SSPNet: Leveraging Robust Medication Recommendation with History and Knowledge

1 Introduction

In the healthcare domain, medication recommendation plays a vital role in determining the optimal combination of drugs tailored to the specific health conditions of a patient. The complexity increases significantly in cases of multimorbidity [Skou et al., 2022], where patients suffer from multiple coexisting conditions, making medical decision-making particularly challenging. The current approach to prescribing medications is predominantly manual, which poses a significant challenge in achieving a balance between therapeutic benefits and minimizing associated risks, such as side effects and potential drug interactions. The crafting of an optimal medication regimen is a complex endeavor, even for experienced healthcare professionals, due to the multitude of factors and interactions that must be considered. The extensive adoption of digital health records in modern healthcare settings has significantly enriched electronic health records (EHRs) [Cowie et al., 2017; Evans, 2016], providing detailed and comprehensive patient data. This abundance of healthcare data presents a promising opportunity to develop advanced predictive models that can enhance clinical decision-making processes.

Due to the clinical importance of drug recommendation, numerous promising methods based on deep learning have been proposed. However, these methods often overlook the set-based nature of drug recommendation [Tan et al., 2022]. The predominant existing methods obtain representations of patients' health conditions by aggregating their diagnoses and procedures information into a unified entity [Wang et al., 2021; Mi et al., 2024]. Yet, these methods can lead to a loss of detailed patient health information. Additionally, some methods [Zhang et al., 2017; Wu et al., 2022] treat drug recommendation as a sequential decision-making process rather than the prediction of a set of drugs. These methods implicitly or explicitly establishes an order in the drug recommendation process. However, the process of actual medication recommendation is not governed by a predetermined sequence.

In this paper, we propose Set-to-Set Prediction Net (SSPNet), an encoder-decoder-based framework that predicts appropriate medication combinations in a set-to-set manner. Our model recommends a set of drugs to patients based on the set of their health status. Concurrently, our model is capable of capturing previously used medications, thereby facilitating the provision of personalized drug recommendations. SSPNet 1 consists of a permutation-consistent encoder-decoderbased network and a personalized drug representation module. Two encoders model the dependencies of diagnoses and procedures to obtain a patient's health condition representation set. The personalized drug representation module scales the representation of medications based on the correlation between the patient's health status and historical medication records. The decoder takes into the set of patient health status and drug representations to predict the target set of drugs. Our contributions can be summarized as follows.

1 https://github.com/ResearchGroupHdZhang/SSPNet.git

• We introduce SSPNet, a novel set-to-set prediction framework for automated medication recommendation that handles the complexity of predicting sets of medications. SSPNet encodes patient health condition set from EHRs into robust high-dimensional representations and directly outputs appropriate set of medications. We highlight the robustness of our framework with respect to different orders of the inputs.
• We develop a personalized drug representation module (PDRM) to capture the historical use of medications in patients. This module precisely models the correlation between current health condition and historical medication information. Moreover, the PDRM module also has excellent compatibility and can benefit other baseline approaches.
• We conducted extensive experiments on MIMIC-III, and the results show that SSPNet outperforms current stateof-the-art methods. The effectiveness of SSPNet is validated through both quantitative metrics and qualitative case studies, confirming its practical utility and reliability in real-world clinical settings.

2 RelatedWork

In this section, we present a comprehensive review of the relevant methodologies, categorized according to the types of information utilized.

2.1 Multi-label Classification-basedMedication Recommendation

Existing predominant approaches uniformly regard drug recommendation as a multi-label classification task [Sun et al., 2022; Zhang et al., 2023; Chen et al., 2023; Shang et al., 2019a; Bhoi et al., 2021]. For example, SafeDrug [Yang et al., 2021] and MoleRec [Yang et al., 2023] models a patient's health condition using a Dual-RNN and enhances the efficacy and safety of drug recommendation considering the molecular structures of the drugs. Recently, 4SDrug [Tan et al., 2022] proposed medication combinations by employing a setoriented approach to measure the similarity between symptom sets and individual drugs. These methods assume independence among medications, or equal contribution of all diseases to the medication combinations. Given the complex relationships among drugs, the assumptions are ineffective for medication recommendation. In this paper, inspired by 4SDrug [Tan et al., 2022], we aim to model the dependencies among medications while recommending drugs in parallel.

2.2 Sequential Decision-making inMedication Recommendation

There exist models that regard the drug recommendation process as a series of sequential decision-making processes [Wang et al., 2024]. An early work, known as LEAP [Zhang et al., 2017], uses a recurrent decoder to capture the drugdisease relationship for the current health condition of the patient, allowing sequential drug recommendation. COGNet [Wu et al., 2022] and VITA [Kim et al., 2024] employ an encoder-decoder generative network to predict medicines individually. These methods also mine the relationship between current patient health and past visits via a copy module. However, these methods implicitly or explicitly impose an order on the medication sets [Li et al., 2023], which may lead to sub-optimal recommendations. In this paper, we aim to produce medication combinations with a set prediction method, which models the relationship between diseases and medications while alleviating the ordering restrictions.

2.3 Set Prediction

For a considerable duration, approaches predicting sets from feature vectors have consistently neglected the inherent unordered quality of sets. Recently, several frameworks [Locatello et al., 2020; Kosiorek et al., 2020; Carion et al., 2020] have incorporated permutation-consistent set generators to address this challenge, since they do not require a specific ordering. DSPN [Zhang et al., 2019], a permutationconsistent set generator, employs backpropagation to decode a set. However, generating a set using gradient descent from an initially guessed set can be computationally expensive. The Transformer [Vaswani et al., 2017], which is permutation-consistent, learns to update the elements of the initial set in a joint manner. Therefore, both TSPN [Kosiorek et al., 2020] and DETR [Carion et al., 2020] leverage Transformer architecture to predict the final set in parallel. Inspired by these methods, we design a set recommendation framework based on Transformer architecture, ensuring permutation-consistency for medication recommendations.

3 ProblemFormulation

In this section, we formulate the core task of automated medication recommendation.

3.1 Electrical Health Records (EHR)

Denote D the set of diagnoses, P the set of procedures and M the set of medications. For a patient p, the healthcare information from p's EHR is a sequence V of all historical visits,

where v(i) is the i-th visit and T is the total number of visits for the patient. Each visit v(i) ∈ V is a triplet that contains the diagnosis, procedure, and medication records,

where v(i) d ⊆ D, v(i) p ⊆ P and v(i) m ⊆ M. We denote all EHRs by V, and all patient by P. Note that |V| = |P|. To enhance the accuracy and interpretability of the recommendation, we integrate domain knowledge from the EHR graph and the drug-drug interaction (DDI) graph into the recommendation model. The EHR graph is an undirected graph Gc = ⟨M, εc⟩, where the nodes are all drug M, and the edges are defined as:

If two drugs in D appear in the same medication record of some visit of some patient, they are connected by an EHR graph edge. The DDI graph Gd = ⟨M, εc⟩ is the induced subgraph of the DDI information from database Twosides [Tatonetti et al., 2012] with respect to the given medication space M. Then, we denote Ac,Ad ∈ {0, 1}|M|×|M| the adjacency matrices of the graphs Gc and Gd, respectively.

3.2 Medication Recommendation Problem

For some visit v(t) ∈ V of a given patient p ∈ P, suppose that the diagnosis record v(t) d = {d1, d2, ..., dk}, and the procedure record v(t) p = {p1, p2, ..., pl}. With EHR and DDI graphs Gc and Gd as domain knowledge, we aim to learning a recommendation predication function f(·) : 2D ×2P ↦→ 2M , that generates a proper medication combinations for the visit,

where ˆm(t) ⊆ M is encoded into multi-hot vector in the rest of the paper.

4 Methodology

Our framework SSPNet consists of four main modules, as shown in Figure 1. The details of these modules are presented in the following subsections.

4.1 Medication RepresentationModule

The co-occurrence and DDI of certain medications in a single prescription is a common practice in human-expert medication recommendation. We fuse the medication's feature vectors with co-occurrence and DDI information in the EHR graph. We first convert all medications into vectors via an embedding matrix Em ∈ R|M|×h , where h is embedding size. Given the drug features M ∈ R|M|×h and the co-occurrence adjacency matrix Ac, we can fuse the co-occurrence information to medication representations by a dual-layer convolutional network (GCN) [Kipf and Welling, 2016] as:

whereWc is the learnable parameter. The ˆAc =Ac+I, and I represents the identity matrix. The diagonal node degree matrix associated with ˆAc is denoted by ˆO. Similarly, we can capture the DDI relations of medications:

The above two parts are aggregated to obtain the representation of medications.

where the λ is a learnable parameter.

4.2 Patient RepresentationModule

Patient's healthcare information includes diagnose records and procedure records. We design two embedding tables, Ed ∈ R|D|×h , Ep ∈ R|P|×h , where each row corresponds to the embedding vectors of the diagnose and the procedure, respectively. Given the diagnose and procedure codes of i-th visit, v(i) d , v(i) p , each diagnose and each procedure are converted into h-dimensional vectors eid and eip using the embedding matrices Eb and Ep, respectively. They constitute the patient's diagnose and procedure embedding matrices D(i) ∈ R|ki|×h and P (i) ∈ R|li|×h . Two encoders based on set attention block (SAB) [Lee et al., 2019; Kosiorek et al., 2020] are utilized to obtain the patient's representation. In the following, we present the details of SABs.

Set Attention Blocks

Set attention blocks aim to model relationships among diagnoses and procedures within the same visit while ensuring permutation-consistency, respectively. Given a query matrix Q ∈ Rn×h and key-value pairs K ∈ Rn×h and V ∈ Rn×h , we can formulate the attention function as follows:

To efficiently capture the interaction information spanning multiple views, SAB leverages multi-head attention by projecting the Q, K, V onto different representation subspaces. Next, a linear layer is used to transform the concatenation of all attention outputs:

Given matrices X,Y ∈ Rn×h , theMultihead Attention Block (MAB) encodes them based on encoder block of the Transformer [Vaswani et al., 2017], without positional encoding:

H = LayerNorm(X+Multihead(X,Y,Y; ω)), (11) Where rFF denotes a feed-forward layer and LayerNorm refers to layer normalization [Ba et al., 2016]. The ω is a parameter. Subsequently, the SAB is capable of being defined in conjunction with the MAB:

To sum up, SAB takes a set as input and models the interactions between each element through self-attention, and ultimately outputs a set of the same dimensionality.

Patient's current health condition encoder

We denote the patient's current visit record as the t-th visit. After converting the diagnose and procedure codes of patient's current visit record into embedding matrices D(t) and P t , we use two SABs, which defined in (12), to encode the diagnoses and procedures:

where ˆ D(t) and ˆ P (t) are the patient's current health condition representations encoded by SABs, respectively.

4.3 Personalized Drug Representations Module

The historical drug treatments can lead to variations in the final prescriptions, even when patients present with similar current health conditions [Shang et al., 2019b; Zheng et al., 2023]. Thus, we have developed a personalized drug representations module. This module leverages historical drug usage information to obtain personalized drug representations.

Pooling by Multihead Attention

The pooling by multihead attention (PMA) [Lee et al., 2019] able to aggregate features while taking into account the importance of different features. Given a set of features X ∈ Rn×h , the PMA can be defined as:

where S ∈ R1×h is a learnable vector.

Scale Drug Representations

We use two PMAs to aggregate diagnosis and procedure representations of patient's all visits.

We model the sequential relationship between patient diagnoses and procedures with a Dual-RNN:

Then we concatenate the h(i) d and the h(i) p to obtain each visit's representation V (i) .

We employ an attention mechanism to assess the correlation between patients' historical medical records and their previous health conditions.

Denote the multi-hot vector m(i) ∈ {0, 1}|M| as the i-th historical medications, where m(i) j = 1 represent the patient has been treated with the j-th drug, and m(i) j = 0 signifies that the j-th drug has not been used. Then we aggregate the historical medication records by a(i) :

where the q represents the correlation between the patient's current health condition and the medications previously used.

Then, we scale the drug representations based on the q and use a eed-forward network to obtain the personalized drug representations Z′ :

4.4 Drug Set PredictionModule

The drug set prediction module is designed to recommend appropriate sets of medications based on the patient's health condition and personalized drug representations.

Drug Set Prediction

Inspired by previous works [Zaheer et al., 2017; Carion et al., 2020], we first designed a permutation-consistent decoder. In contrast to sequence generation models [Wu et al., 2022; Kim et al., 2024] that generate drug representations sequentially, the decoder takes as input all candidate drugs and predicts a targeted drug combination, without imposing any irrelevant sequential information. Given the personalized drug representations set Z′ , a multi-head self-attention mechanism defined in (9) is utilized to model the interactions among the medication combinations,

Conditioned on the diagnose representations ˆ D(t) and procedure representations ˆ P (t) of the patient's current visit, we further refine the medication representations,

Then, a fully connected feed-forward network (with a sigmoid function) transforms the representation ˆZ into predicted medication representation ˆM,

Subsequently, we use a feed-forward network FF2 : Rh → R with a sigmoid activation function σ to derive the probability of each drug being recommended,

Finally, we select entries from ˆmwhere values exceed predefined threshold δ, to obtain a multihot prediction vector ˆ o.

Loss Function Design

We first optimize all learnable parameters of SSPNet in training phase. As stated in the previous section, SSPNet gets the predicted medication representation and picks out the predicted medications with drug set prediction module. We consider it as a multi-label binary classification problem. Following [Yang et al., 2021; Yang et al., 2023], we introduce the binary cross-entropy loss Lbec, the multi-label margin loss Lmulti and DDI loss as our loss functions. The details are presented in the Appendix A.

5 Experiments

In this section, we conduct a series of extensive experiments to evaluate the performance of SSPNet, by comparing it with several competing models.

5.1 Dataset

We conducted experiments on the MIMIC-III [Johnson et al., 2016] and MIMIC-IV [Johnson et al., 2023]. Following the setting in [Yang et al., 2021], we split the datasets into training, validation and test as 4 : 1 : 1. The statistical information of the preprocessed dataset is shown in Table 1. Details of data pre-processing are presented in the Appendix B.

5.2 Setting

We implement our method using PyTorch 1.9.0, which is built on Python 3.8.16. We conducted all experiments on an Intel Xeon Platinum 8260 server, comprising 24 CPU cores, 188G memory and a 12GB NVIDIA TITAN V GPU. In our models, we set h = 64 for Ed, Ep and Em. The number of initial medications n = 131. The feed-forward network FF1, FF1 and FF3 are designed as one linear layer. We determined the best hyperparameters based on their validation set performance, setting the threshold δ = 0.5, weight β = 0.95 and Kp = 0.05. The Adam optimizer [Kingma and Ba, 2014] is applied for model training, with a learning rate of 1 × 10-4 .

5.3 Baselines andMetrics

We assess the performance of SSPNet through a comparative analysis with the following models: the instance-based standard Logistic Regression (LR), the chain-structured LR classifier model: Ensemble Classifier Chain (ECC) [Read et al., 2011], the dynamic graph-augmented memory-based approach: GAMENet [Shang et al., 2019b], the drug molecule structure-based approaches: SafeDrug [Yang et al., 2021] and MoleRec [Yang et al., 2023]. the sequential decision making approach LEAP: [Zhang et al., 2017], and the conditional generation network-based approaches: COGNet [Wu et al., 2022] and VITA [Kim et al., 2024], Following previous research [Yang et al., 2021; Wu et al., 2022; Yang et al., 2023] on medication recommendations, we employ four widely-used metrics: Jaccard similarity, F1 score, PRAUC and DDI rate, to verify the accuracy and effectiveness of these models. More details of baselines and metrics are demonstrated in the Appendix C.

5.4 Main Result

The Table 2 presents a comparison of our method's performance with all baseline models. Overall, our present SSPNet outperforms the baseline models in Jaccard, F1 score and PRAUC. Our proposed SSPNet recommends drugs in a setto-set manner. Hence, SSPNet can better represent the patient's health condition. Predicting drugs in a set form can avoid the influence of the order of predicted drugs. Furthermore, our PDRM module is capable of capturing the longitudinal historical medical records of patients. Therefore, our method can perform better than other methods.

Traditional classifiers, such as LR and ECC, which do not leverage deep learning, cannot model the complex medication combination recommendation, demonstrate poor performance. The multi-label classification-based medication recommendation methods, like GAMENet, neglect the importance of DDI. Consequently, SafeDrug and MoleRec incorporates the drug molecule structures and control DDI through a loss function, which resulting in improved performance. However, these methods aggregate the health status of patients, leading to the loss of patient health information, and moreover, these methods overlook the drugs previously used by the patients. Considering medication combination recommendation as a sequential decision-making task, LEAP difficult to model long medication decision-making. Using a Transformer, COGNet and VITA show improved performance. However, these methods impose an order on the medications, leading to unsatisfactory performance. VITA relies on patients' historical medical records, which makes it challenging to generalize to situations involving single visits.

5.5 Ablation Study

We conducted the following ablation experiments to validate the effectiveness of each module in SSPNet. The result of different variants of SSPNet is shown in Table 3.

Medication Representation Module

We evaluated the validity of the EHR graph information. The results of SSPNet w/o Gc indicate that the EHR graph information contributes to the final outcome. The EHR graph encompasses co-occurrence information of medications and drug-drug interaction (DDI) information, which can enhance the representation of medications.

Patient Representation Module

We first assessed the validity of the patients' diagnoses and procedures information. We separately removed the information regarding diagnoses and procedures. As indicated by SSPNet w/o D and SSPNet w/o P , removes one of them lead to a diminished model performance.

To verify the advantages of our set encoders, we compared it with methods for modeling patient health conditions using aggregation. As show in Table 3, the SAB-→Dual-RNN [Le et al., 2018] indicates that after replacing the SAB module with the Dual-RNN (the details can referred to [Yang et al., 2021; Yang et al., 2023]), the performance of SSPNet decreased. Applying the set encoder to obtain the representation of patients' health status can better model the information regarding patients' health conditions.

Personalized Drug Representation Module

We have validated the superior performance of our proposed PDRM. As indicated by SSPNet w/o PDRM, the performance of SSPNet decreases when it is not equipped with the PDRM. The proposed PDRM effectively enhances the model's ability to capture historical medication information.

Table 3: Ablation study on MIMIC-III and MIMIC-IV.

Drug Set Prediction Module

To verify the superiority of our ensemble prediction module, we replaced our medication set prediction module with the decoder module of a traditional Transformer, similar to that used in COGNet. The results are presented in Set prediction -→ Sequential prediction. Recommending medications in the form of sets can avoid the sequential recommendation process which introduces specific medication orders, thus enhancing the accuracy of the model.

5.6 Case Study

We present a sample patient with historical visits to illustrate our model's drug recommendation. Figure 2 shows the visualization of drug representations through t-SNE [Van der Maaten and Hinton, 2008] under the initial, personalized, and predictive phases, corresponding to Equations (7), (20) and (22), respectively. In the initial phase, the drug representations have identical distributions for each patient. Many existing methods recommend drugs based on drug representations in the initial phase, which leads to suboptimal results. In the personalized phase, our PDMR module personalizes the drug representations according to the drugs the patient has previously used. Compared with the initial phase, the target drug representations are more separated from other drugs. Drug recommendations based on personalized drug representations can effectively improve the model accuracy. In the predictive phase, our permutation-invariant framework predicts the representations of recommended drugs based on the patient's health status set and the personalized drug representations.

5.7 The Stability of Set-to-Set Framework

We further assessed the performance of SSPNet, GAMENet, SafeDrug, MoleRec, Leap, COGNet and VITA under following five different input drug sequences:

• No process preserves the inherent sequence of drugs, without performing any manipulation.
• Frequent is sorted the order of the drugs in descending order based on their frequency.
• Rare first is arranged in ascending order by frequency.
• Vocabulary is sorted of drugs in alphabetical order.
• Random is randomly shuffled the order of the drugs.

The results are shown in Figure 3. Our model appropriately treats patient health information and medications as sets, and employs permutation-invariant encoders and decoders to predict drug prescriptions. Thus, our model exhibits stable performance and possesses superior robustness. MoleRec, SafeDrug, and GAMENet models aggregate patient diagnostic and surgical information to obtain patient health status, rather than treating patient health information as a set. This approach leads to a slight deficiency in the robustness of these models. VITA, COGNet and Leap which regard drug recommendation as a sequential decision-making process are sensitive to the medical implications of drug order.

5.8 Compatibility of PDRM

We validated the capability of our proposed PDRM module to capture historical medication information and its orthogonality. We applied the personalized drug representation model to MoleRec, SafeDrug and GAMENet. The performance is shown in Table 4. The performance of these models is enhanced after the integration of our PDRMmodule. We further evaluated the performance with varying numbers of historical visits. The results shown in Figure 4 as shown. Since most of the patients in the MIMIC-III data set have less than five visits [Wu et al., 2022], we conducted experiments in patients with fewer than five visits. MoleRec and SafeDrug have only taken into account the current and historical health conditions of patients, neglecting the drugs that the patient has previously used. Upon the incorporation of the personal medication module, they are able to account for historical medication information. GAMENet has individually considered the relationship between patients' health status and medication as well as their historical drug use, leading to suboptimal outcomes. When GAMENet is equipped with our PDRM module, it can jointly consider the relationship between patients' health status, medication, and historical medication information. Our proposed PDRM not only shows exhibits flexibility in its application in various models, but also excellent performance across varying numbers of historical visits.

6 Conclusion

In this paper, we introduce SSPNet, a novel end-to-end framework designed to tackle the intricate task of automated medication recommendation. Unlike existing methods, which often overlook the essential set-based requirements, SSPNet employs a set-based approach to encode comprehensive patient health condition and predicts appropriate prescription through a permutation-consistent decoder. Furthermore, we have designed a personalized drug representation module with excellent compatibility to capture the medications previously used by patients. Our extensive evaluations on the MIMIC-III and MIMIC-IV datasets demonstrate that SSPNet outperforms existing state-of-the-art models.

Ethical Statement

The research presented in this paper utilizes the MIMIC dataset, which is a deidentified data set that is compliant with the Health Insurance Portability and Accountability Act (HIPAA). Under HIPAA regulations, the use of deidentified data for research purposes does not constitute human subjects research. The authors of this study have taken all necessary precautions to ensure that the research is conducted in an ethical manner, respecting the privacy and rights of the individuals represented in the dataset.

Acknowledgments

This work is partially supported by the Guangdong Provincial Key Lab of Integrated Communication, Sensing and Computation for Ubiquitous Internet of Things (No.2023B1212010007). Liang-Jie Zhang and Lin Ma are the corresponding authors.