Self-Supervised Adversarial Distribution Regularization for Medication Recommendation

1 Introduction

Technological innovations in deep learning have achieved great success in various clinical applications[Yue et al., 2020], such as disease diagnosis, onset prediction, and physiological condition monitoring [Chen et al., 2018; Li et al., 2019; Malakouti and Hauskrecht, 2019; Wang et al., 2019; Chen et al., 2019], which enable physicians to better diagnose and treat patients since the beginning of the professional practice of medicine. As one of the hottest research topics, medication recommendation aims to assist physicians in making effective prescriptions according to Electronic Healthcare Records (EHRs) that describe patients in terms of vital signs, diagnoses, and procedures. A comprehensive analysis of EHRs based on deep learning provides physicians with insight regarding patients. It helps caregivers cope with peril brought by the fast-paced way of working in ICUs, significantly improving their quality of life.

∗Contact Author1

†Contact Author2

Deep learning has been adopted for different medication recommendation tasks. Many of these existing methods make great efforts to obtain accurate patient representations to carry out practical medication recommendations, and widely applied approaches include instance-based and longitudinal methods [Lipton et al., 2016; Shang et al., 2019a; Wang et al., 2017]. However, these methods neglect the fatal side effects related to Drug-Drug Interaction (DDI) due to duplication, antagonism, and alternation [Zhang et al., 2017]. Meanwhile, limited researches on this topic apply extra knowledge about DDI to mitigate the problem [Zhang et al., 2017; Ma et al., 2018; Shang et al., 2019b], which requires specific components for knowledge extraction or information fusion, making models complicated and introducing bias that has adverse impacts on the recommendation. Besides, these methods treat patients with different DDI rates, i.e., the fraction of combinations that lead to DDI to the total number of varieties in a set of medications, as a single cohort, which would put strict requirements forwards on the generalization performance of models. However, representations of these patients could negatively affect each other when used to train models due to the different probability distributions they follow.

To address the above problems, we propose a novel SelfSupervised Adversarial Regularization Model for Medication Recommendation (SARMR). Instead of introducing extra knowledge about DDI, SARMR obtains probability distributions of patient representations related to safe medication combinations in the feature space from raw EHRs. Then SARMR applies the knowledge as the true data to adversarially regularize distributions of patient representations with achieving DDI reduction. Patients with different DDI rates are respectively used and regularized as different cohorts, and the adverse impacts on generalization when they are treated as a single cohort could be avoided.

Key challenges for building the model include obtaining probability distributions that reflect low DDI rates and maintaining recommendation performance while shaping distributions. Discovering meaningful embeddings representing desired probability distributions over data is a significant promise of deep learning. The continuous feature space enables the gradient-based optimization to search for distributions that reflect desired properties [G´omez-Bombarelli et al., 2018; Blaschke et al., 2018; Makhzani et al., 2015; Pan et al., 2018]. However, previous methods require extra tasks for the knowledge-based jump in the latent space, which may introduce bias or result in inaccurate search results. In addition, the DDI rate is not a direct output of patient representations but calculated based on the recommended medications. The lack of direct mapping between latent features and DDI rates makes property-oriented methods unavailable. Also, patients suffer from different diseases, and their representations contain other conditions. Focusing on regularizing distributions for desired DDI rates alone may ignore critical information encoded in patients' representations and lead to poor performance on the recommendation.

SARMR overcomes these obstacles in two steps. In the first step, SARMR uses patient records whose DDI rates are lower than a threshold of Drec to recommend medications, and obtains patient representations related to low DDI rates as the true data for the adversarial regularization in the second step. SARMR firstly encodes EHRs with GRUs, and builds a key-value Memory Neural Network (MemNN) [Miller et al., 2016], whose keys are representations of admissions and values are corresponding medications, to model interactions between physicians and patients. Then SARMR uses the representation of the last admission as query to conduct multireading on the MemNN, while Graph Convolutional Network (GCN) [Kipf and Welling, 2016] is presented as the embedding module of the read results. The updated query is used to make the recommendation. Then in the second step, SARMR uses records of all patients regardless of their DDI rates to jointly conducts medication recommendation and adversarial distribution regularization with the Generative Adversarial Network (GAN) [Goodfellow et al., 2014] based on the obtained representations in the first step, so that both effective medication combinations and DDI reduction are achieved.

Our contributions could be summarized as follows:

• We propose a self-supervised strategy to shape distributions of patient representations for safe medication recommendation, which obtains prior distribution reflecting low DDI rates from raw EHRs and conducts adversarial distribution regularization with GAN.
• A key-value MemNN is constructed to learn interactions between doctors and patients by carrying out multi-hop reading on the MemNN, so that contexts in historical EHRs are derived for informative patient features.
• Experimental results demonstrate that SARMR outperforms all the baseline methods, and it can make effective medication recommendation with an F1 at 0.6608, while it also achieves a DDI reduction at -2.72%.

2 RelatedWork

2.1 Distribution Regularization

Discovering informative embeddings representing meaningful distributions over data is a significant promise of deep learning. The continuous latent space enables models to search for distributions that reflect desired properties. For example, [Lim et al., 2018] incorporate molecular properties into latent features when using conditional variational autoencoder to generatemolecules with desired properties. [G´omezBombarelli et al., 2018] add a regression task to autoencoder (AE) for property prediction to connect latent features with target properties, searching for desired attribute in the feature space. And [Blaschke et al., 2018] extend the idea by applying Bayesian optimization on the obtained latent space to find new molecular structures with target properties. [De Cao and Kipf, 2018] further utilize reinforcement learning to carry out optimization towards desired chemical properties.

Distribution Regularization for desired properties has also been discussed before. [Makhzani et al., 2015] match the aggregated posterior of hidden vectors of AE with an arbitrary prior distribution to generate meaningful samples. [Kadurin et al., 2017] apply the strategy to drug discovery, and introduce a neuron responsible for reflecting desired properties to develop new molecules. Meanwhile, [Pan et al., 2018] use adversarial regularization to match latent representations of graphs to match a prior distribution, so that meaningful graph embeddings in a continuous vector space are achieved.

SARMR differs from these methods since no extra task is required to guide the search in the latent space. Instead, a prior distribution directly relating to low DDI rates is obtained from raw EHRs for adversarial regularization.

2.2 Medication Recommendation

To model relations covering multiple input views in EHRs, [Le et al., 2018] use memory augmented neural networks to achieve better performance in drug prescription task and disease progression task. But the method does not take relationships between drugs into account, so [Shang et al., 2019b] build a graph based on the co-occurrence of drugs in EHRs, and uses GNN to obtain embeddings for those drugs. Meanwhile, [Zhang et al., 2017] address medication recommendation from the view of sequential decision-making, and use a recurrent decoder to model drug dependency, while drug-todisease mapping is modeled by content-based attention. But these methods concentrate on visit-level temporal information, so [Choi et al., 2016] present a two-level neural attention model to model both visit-level and variable-level sequential information, which also provides a detailed interpretation of the prediction results. Most of these proposed methods ignored EHRs related to patients with a single visit, so [Shang et al., 2019a] pre-train its transformer-based visit encoder on EHR data from patients with a single visit, then the model is fine-tuned on EHRs of patients with multiple re-admissions.

Compared with previous methods, SARMR applies a selfsupervised strategy to obtain target distribution from EHRs for the adversarial distribution regularization. And SARMR conducts multi-hop reading on the key-value MemNN to contextual information for medication recommendations.

3 ProposedMethod

Details of SARMR are described in this section. SARMR firstly selects patient records by a threshold Drec for DDI rate, and uses the data to train the Medication Recommendation (MedRec) module that consists of a encoder and a predictor. Then SARMR carries out adversarial regularization on the distributions of all patient representations based on the informative patterns obtained by the encoder to achieve DDI reduction. The overview is demonstrated in Figure 1.

3.1 Notations

EHRs describe patients in terms of different views, and each patient is recorded as a sequence of multivariate observations X = [x1, x2, ..., xT ], where T is the number of admissions, and xt represents records of the tth admission. xt consists of three sets of medical codes, including diagnoses dt , procedures pt , and medications mt , which belong to the medical code sets Sd, Sp, and Sm respectively, and the corresponding numbers of distinct codes in EHRs are |S∗|. These records are transformed into multi-hot vectors std ∈ R|Sd| , stp ∈ R|Sp| and stm ∈ R|Sm| to act as inputs of SARMR.

A drug graph DG ∈ R|Sm|∗|Sm| is constructed based on all medication combinations m∗ to indicate whether two drugs i and j have been prescribed to the same patient, if so, DGi,j = DGj,i = 1, otherwise DGi,j = DGj,i = 0. Similarly, a DDI graph IG ∈ R|Sm|∗|Sm| is constructed to indicate polypharmacy interactions with drug pairs or higher-order drug combinations, and IGi,j = IGj,i = 1 if medication i and j would cause side effects, otherwise IGi,j = IGj,i = 0. Given IG, the DDI rate Dr(mt ) ∈ [0, 1] for mt could be calculated as Eq.(1) shows.

SARMR uses a threshold Drec for DDI rate to select patients Prec with Dr(mT ) ≤ Drec, and uses their records to train MedRec to obtain the Gaussian distribution p(z) for adversarial regularization. Given the medical records of a patient xt(t < T) as well as dT and pT , the goal of SARMR is to predict medications mT .

3.2 Medication Recommendation

The encoder uses two GRUs to obtain temporal information from std and stp , where t ≤ T. SARMR firstly embeds std and stp into continuous representations with embedding matrices Wd ∈ R|Sd|×dim and Wp ∈ R|Sp|×dim, where dim is the embedding size that determined with Gaussian Process as a hyper-parameter. Then the model obtains hidden states for each time stamp as Eq.(2) and Eq.(3) show, and combines them with a linear embedding layer to generate the representation qt for the tth admission.

After encoding patient representations, SARMR carries out multi-hop reading on a key-value MemNN KV to obtain contextual information from interactions between physicians and patients, and uses GCN to transform the results into continuous embeddings.

The MemNN KV uses qt (t < T ) as keys and corresponding ground truth medications stm(t < T) as values. For each hop k, SARMR calculates an attention weight ak between uk-1 and KV T-1 key as Eq.(4) shows, where uk-1 is the output of last hop, q∗ ∈ KV T-1 keys , and Wa is a weight matrix. Then SARMR reads weighted memories ok from KV as Eq.(5) shows, where Z k are medication embeddings generated by GCN in the k hop, and s∗m are values in KV T-1 value . Given ok , the query uk-1 is updated to uk according to Eq.(6). The reading process would be repeated for K hops to get uK .

For the calculation of Z k , the drug graph DG is processed following GCN procedures as Eq.(7) shows [Kipf and Welling, 2016], where ˆD is a diagonal matrix such that ˆDii = ∑j DGij and I is a identity matrix. SARMR applies a two-layer GCN on DG to obtain embeddings for drugs as Eq.(8) shows, where We is the medication embeddings from DG, and W1 is the hidden weight parameter matrix. Finally, the correlation between qT and Z K is calculated according to Eq.(9), where We is a weight matrix for the attention.

SARMR predicts the final result as Eq.(10) shows. The equation takes the concatenation of patient representation qT , multi-hop reading result uK , and weighted embeddings of medication e as inputs, where S(.) is a sigmoid function and f(.) is a fully connected layer.

3.3 Self-Supervised Adversarial Regularization

The key idea of SARMR is to adversarially match latent features of patients to a prior distribution reflecting low DDI rates, which is obtained in a self-supervised way that requires no extra knowledge, so that patients would be prescribed safe medication combinations. To to so, given the threshold Drec, patients whose DDI rates Dr(mT ) ≤ Drec are used to train the MedRec module formed by the encoder and predictor described previously, and SARMR uses the Gaussian distribution p(z) that qT r follows [Kadurin et al., 2017; Pan et al., 2018], which represents the distribution over low DDI rates, as the source of real data in the adversarial regularization. Specifically, the mean and covariance matrix of all qT r are calculated to get the fitted Gaussian distribution p(z). Then all patients regardless of their DDI rates are used to train the model, and their representations qT a contains two parts: qT f related to patients whose Dr(mT ) ∈ (Drec, 1], which follows pf(z), and q′Tr that follows p′ (z), which comes from qT r but has been affected by qT f .

Patient representations qT a obtained by the encoder are adversarially shaped to match pf (z) to p(z), while p′ (z) is also corrected back to p(z), so that they would all follow the same distribution and present the desired low DDI rates. To achieve the goal, a discriminator is attached on top of encoder to form a GAN model, and qT a will act as fake data while samples from p(z) would act as real data. Meanwhile, patients suffer from different diseases and their representations contain information of various conditions as well as medications. It is critical to guarantee that the information is mapped to accurate medication combinations when representations are regularized. Thus, SARMR jointly train the GAN model and the MedRec within each mini-batch in two phases: the prediction phase and the regularization phase.

In the prediction phase, the MedRec module is updated to minimize the loss between predicted medications and ground truth medications, so that the prediction performance is improved. In the regularization phrase, SARMR updates the discriminator to distinguish real data generated by the target distribution from features generated by the encoder. Once the training procedure is done, patient representations qT a are regularized to follow p(z), and these representations could still be mapped to corresponding medications to finish the recommendation by the predictor. Thus, both the expected low DDI rate and prediction performance are guaranteed.

3.4 Training and Inference

Following the joint training strategy, there are two loss functions for SARMR. The GAN regularization process makes use of the typical loss function for GAN model [Goodfellow et al., 2014]. The MedRec module is updated to minimize the weighted loss for medication recommendation shown in Eq.(11) [Shang et al., 2019b], which consists of the binary cross entropy loss Lb and multi-label margin loss Lm. Lm is used to make the predicated probability of ground truth medications have at least 1 margin larger than others, which would benefit the predictive performance of SARMR.

Here, y ∈ R|Sm| is sTm, i.e., the ground truth of medications for the final admission, ˆ y ∈ R|Sm| is the predict result in the form of probability, and ˆ Y is the predict medication set, so ˆ yi and ˆ y[ ˆ Yj ] are the probabilities that the ith and ˆ Y th j medication would be prescribed to the patient. The constraints for the weights α1 and α2 are that α1 > 0, α2 > 0 and α1 +α2 = 1.

4 Experiment

In this section, SARMR is compared with different baseline methods on the real-world clinical dataset MIMIC-III v1.4 [Johnson et al., 2016]. The model is implemented with PyTorch and trained on a NVIDIA TITAN Xp GPU, and more information about source code could be found at Github 1 .

4.1 Dataset

Patients with at least two admissions to hospitals are selected for experiments, and all their diagnoses and procedures are mapped to fixed vocabularies to act as inputs, while medications prescribed during the first 24-hour in hospitals are used as ground truth for the recommendation [Shang et al., 2019b]. Totally 6350 patients are selected, and numbers of diagnoses, procedures and medications contained in their records are 1960, 1432, and 153. The top-100 severe DDI types are used to build IG, and Kernel Density Estimate on DDI rates indicates a Gaussian distribution whose mean is 0.4, so Drec=0.4.

1 https://github.com/yanda-wang/SARMR

4.2 Baselines

The following methods are compared with SARMR.

• Leap [Zhang et al., 2017]: Leap addresses medication recommendation as a sequential decision-making task, and uses reinforcement learning to improve accuracy.
• RETAIN [Choi et al., 2016]: RETAIN uses a two-level attention model to detects influential past visits for sequential prediction.
• DMNC [Le et al., 2018]: DMNC models multi-view interactions and long-term dependencies via memory augmented neural network to recommend medications.
• GAMNet [Shang et al., 2019b]: GAMENet uses a dynamic memory network to model historical EHRs, and applies GCN on DDI graph to achieve DDI reduction.
• MedRec: MedRec contains the encoder and predictor of SARMR, but without adversarial regularization.

4.3 Metrics

The performance of SARMR is evaluated using Jaccard Similarity (Jaccard), Precision Recall AUC (PRAUC), Average F1 (F1), and changes of average DDI rate (△DDI). Among these metrics, △DDI, as shown in Eq.(14), indicates the difference between DDI rates of predicted results and ground truth, which shows whether DDI is reduced.

4.4 Evaluation

As shown in Table 1, SARMR achieves the best △DDI at -2.72% and succeeds in attaining DDI reduction. Meanwhile, the figure for MedRec is 3.04%, showing that the GAN model has successfully regularized the distribution of patient representations to reduce the DDI rate. GAMENet uses extra DDI knowledge as a memory component to reduce DDI, and the △DDI is 3.06%. The comparison proves that regularizing distributions of representations could be a more effective strategy for DDI reduction. The remaining methods have much higher DDI rates and fail to predict safe medications.

Besides, SARMR outperforms all the baselines with the highest Jaccard, PRAUC, and F1 at 0.5039, 0.7688, and 0.6608 respectively. Among those baselines, DMNC uses an encoder similar to SARMR, and its sub-optimal performance indicates that key-value MemNN plays an important role in modeling EHRs. GAMENet uses a similar key-value MemNN without multi-hop reading, and its lower figures on metrics show that reasoning over interactions between patients and doctors are essential.

To further illustrate the effectiveness of adversarial regularization, the kernel density estimate is conducted on DDI rates of predict results and ground truth, which are illustrated in Figure 2. As shown in Figure 2.(a) and Figure 2.(b), density curves of DDI rates related to Leap and DMNC shift to the right compared with that of ground truth, which explains why these two methods lead to high DDI rates in the predict results. Similarly, the curve for RETAIN in Figure 2.(c) also shifts to the right, but the density has much higher values at high DDI rates with small variance, and that is why RETAIN has the highest △DDI. The results of RETAIN, GAMENet, and MedRec show similar trends, and the densities of DDI rates in their predict results have small variance while the means are more extensive than that of the ground truth. Meanwhile, the density of SARMR has the highest values in low DDI rates, indicating that the regularization has successfully matched distributions of patient representations to the desired Gaussian distribution to achieve DDI reduction.

Besides, how the number of DDI types affects different methods is explained to show the robustness of adversarial regularization on DDI reduction. In addition to the top-100 DDI types, the top-40, top-60, and top-80 types are considered. The results are shown in Table 2. SARMR can achieve DDI reduction regardless of the number of DDI types. Meanwhile, when using the top-40 DDI types, GAMENet succeeds in reducing DDI rates and achieves a △DDI at -1.08%, but the method fails to do so when more types are considered. The figure for Leap drops dramatically from 7.80% to -17.80% as the number of DDI types decreases from 100 to 40, while DMNC and RETAIN cannot achieve DDI reduction in any case. The results demonstrate the robustness of SARMR since it is the only method that leads to a DDI reduction when a different number of DDI types are used.

Ablation Study

To evaluate the effectiveness and necessity of each component of SARMR, the model is deconstructed by replacing or removing these components to build variants as follows:

• S-GRU: The variant replaces the two GRUs in SARMR with a Single GRU to model patient representations.
• S-Encoder: The variant removes the MemNN and uses the encoder's output to recommend medications directly.
• No-GCN: The GCN that embeds medications in SARMR is replaced with an embedding matrix.
• MedRec: MedRec removes the discriminator in SARMR and it is not adversarially regularized.

The results in Table 3 and Figure 3 shows that SARMR achieves the best performance. △DDI for MedRec is 3.04%, showing that the GAN model is essential to reduce the DDI rate. Otherwise, the density curve would shift to the right as shown in Figure 3.(d). No-GCN has the second-best performance on Jaccard and F1 as well as the most similar density curve with SARMR, indicating that the clinically meaningful embeddings of medications obtained by GCN could assist SARMR in achieving better performance. Meanwhile, S-Encoder has a dramatically high △DDI at 5.34%, and the density has higher values in high DDI rate, which proves that without an appropriate decoder to interpret information from the encoder, the regularization on patient representations may even lead to an adverse impact on the performance. S-GRU has the lowest Jaccard and F1. The results show that temporal patterns in diagnoses and procedures should be modeled separately and then combined to obtain a comprehensive representation rather than directly treated as a whole.

5 Conclusion

In this paper, we propose a novel self-supervised adversarial distribution regularization strategy SARMR for safe medication recommendation. Existing methods treat distributions related to different DDI rates as a single cohort, and requires extra tasks to apply knowledge for DDI reduction. SARMR conducts multi-hop reading on MemNN to derive contextual information from EHRs, and adversarially regularizes patient representations based on desired distributions obtained from raw EHRs to achieve DDI reduction. For now, SARMR applies a fixed number of hops for all patients, while their EHRs contain different amounts of information. Given the potential for facilitating accurate information extraction, we expect adaptively determining the number of hops in future work.

Acknowledgments

This research has been supported by the NSFC under Grant No.61972384 and CPSF under Grant No. 2017M621192.