SIGIR2024| RAREMed: No Patient Left Behind - Improving the Accuracy of Drug Recommendations for Patients with Rare Diseases

TLDR: In this paper, we propose a new pre-training-fine-tuning based drug recommendation modeling framework for the problem of low recommendation accuracy of drug recommendation models for patients with rare diseasesRAREMed, and propose two targeted pre-training tasks to improve the model's ability to learn the representation of patient conditions, especially rare disease patients' conditions, and to help the drug recommendation model to improve the accuracy of its recommendations for rare disease patients, thus improving the fairness of the drug recommendation model. Our experimental results on two commonly used public datasets show significant advantages of our approach in terms of accuracy and fairness of drug recommendation models.

Paper Address:/abs/2403.17745

Code Address:/zzhUSTC2016/RAREMed

introductory

With the rapid development of artificial intelligence, its application in the field of healthcare has gained more and more attention, among which drug recommendation is an important task that has achieved rapid development in recent years. The drug recommendation task aims to utilize clinical information such as patients' diseases, surgeries, and medical histories to recommend effective and safe drug combinations for patients or physicians, to reduce physicians' workload, and to reduce the risk of potential medical errors, such as drug misuse and adverse drug-drug interactions.

Existing drug recommendation models focus on improving overall drug recommendation accuracy, but a common problem they face is the fairness issue - the recommendation accuracy for patients with rare diseases is significantly lower than that for other patients, see Figure 1(b). This is mainly due to the fact that rare diseases appear so rarely in the training data that it is difficult for the model to learn accurate representations. As shown in Fig. 1(a), a small number of diseases appear a large number of times, while most of the long-tailed diseases appear a small number of times. On the other hand, existing drug recommendation models model the two key inputs-diseases and procedures-separately, making it difficult for the models to capture the association between the two inputs.

Figure 1: (a) Long-tailed distribution of disease codes on two common publicly available datasets (b) Comparison of recommendation accuracy between patients with rare diseases and patients with common diseases

In this work, in order to address the fairness problem faced by drug recommendation, we utilize the TRANSFORMER architecture and propose two targeted pre-training tasks to improve the model's learning and expressive capabilities. Specifically, the model contains two pre-training tasks: Sequence Matching Prediction (SMP) and Self Reconstruction (SR). The Sequence Matching Prediction (SMP) task allows the model to discriminate whether disease and surgery sequences belong to the same patient, and thus better understand the correlations in the patient's condition information. The Self Reconstruction task allows the model to reconstruct the patient's input code using the learned patient representation, thus capturing the patient's input information more comprehensively.

Definition of the problem

In the drug recommendation task, the input information is usually an EHR (Electronic Health Record, medical health record), which contains information from the patient's medical record\(\mathcal{V}^{(j)} = \{\mathbf{d}^{(j)}, \mathbf{p}^{(j)}, \mathbf{m}^{(j)}\}\)which\(\mathbf{d}^{(j)} = [d_1, d_2, \cdots, d_x] \in \mathcal{D}\)Indicates the disease from which the patient suffers.\(\mathbf{p}^{(j)} = [p_1, p_2, \cdots, p_y] \in \mathcal{P}\)Indicates a procedure performed by the patient.\(\mathbf{m}^{(j)} \in \{0, 1\}^{|\mathcal{M}|}\)Indicates the drug prescribed by the physician. In addition to this, the input information includes Drug-Drug Interaction Graph (DDI Graph, Drug-Drug Interaction Graph) for constraining and evaluating adverse drug interactions in the outcome of drug recommendations.

The drug recommendation task is defined as a sequence of diseases for a given patient\(\mathbf{d}\)and surgical sequences\(\mathbf{p}\)and drug-drug interaction diagrams\(\mathbf{A}\)The goal is to recommend a collection of drugs\(\hat{\mathcal{Y}}\)to maximize prediction accuracy and minimize adverse drug interactions.

The fair drug recommendation task means that, in addition to the above drug recommendation optimization goals, the model needs to have similar recommendation accuracy for both patients with common and rare diseases, reducing the unfairness of the model's recommendation results for patients with rare diseases.

Characterization of patients with rare diseases

Figure 2: The horizontal coordinate is the number of occurrences of the rarest disease in a patient, and the vertical coordinate is the (a) number of diseases, (b) number of surgeries, (c) number of drugs, and (d) drug prevalence corresponding to the group of patients suffering from different prevalence, all averaged

We divided the patients in the MIMIC-IV dataset into groups based on the prevalence of the rarest disease suffered and counted the characteristics of the patients within each group, as shown in Figure 2, and we have the following two observations:

Observation 1: Patients with rare diseases have more complex conditions. As shown in Figure 2(a)(b), the number of diseases suffered by patients and the number of surgeries performed by patients within each patient group decreased as the prevalence of the patient's disease increased.

Observation 2: Treatment options for patients with rare diseases are more complex and individualized. As shown in Figure 2(c)(d), patients with rare diseases often require more and rarer drugs for treatment.

All of these present greater challenges for drug recommendations for patients with rare diseases. Providing more accurate drug recommendation results for these patients requires capturing their conditions more fully and providing more comprehensive and personalized drug recommendation results.

methodologies

In order to provide more accurate and fair drug recommendations, we propose Robust and Accurate REcommendations for Medication (RAREMed) model, as shown in Figure 3. The model is divided into three main parts: patient coding, pre-training and fine-tuning.

Figure 3: RAREMed model diagram

Patient representation (Encoder)

In order to capture more comprehensive information about the patient's condition, given the patient's disease\(\mathbf{d}\)Surgery\(\mathbf{p}\)and other inputs, we concatenate these codes as a unified sequence input to the transformer model:

\[\begin{align} \label{eq:1} input = [\text{CLS}] \oplus \mathbf{d} \oplus [\text{SEP}] \oplus \mathbf{p}, \end{align} \]

Where [CLS] and [SEP] are special encodings that denote the sequence opener and separator, respectively.\(\oplus\)Indicates sequence linkage symbols. The patient's disease and procedure codes are sorted in order of importance in relation to the patient's current hospitalization, and the importance has been labeled in the dataset by a medical professional.

On this basis, we designed three embedding layers, where theToken Embedding layer (Token Embedding)represents the semantics of each token.Relevance Embedding layer (Relevance Embedding)used to encode the importance of each token in relation to the location only.Segment Embedding layer (Segment Embedding)Used to encode the category of the input token, identifying whether each token belongs to a disease or a procedure.

Finally, these embedding go through a transformer encoder that generates the patient representation, encoded in the output layer using [CLS] notation:

\[\mathbf{r} = \text{Encoder}(E_{\text{tok}}(input) + E_{\text{seg}}(input) + E_{\text{rel}}(input))[0]。 \]

Pre-training

In order to enhance the representation learning capability of the model, we designed two pre-training tasks for the drug recommendation task:

Task #1: Sequence Matching Prediction Sequence Matching Prediction (SMP) ：The goal of the SMP task is to enable the model to predict whether two sequences of input diseases and surgeries belong to the same patient. We provide a model for each real patient's (\(\mathbf{d_i}\), \(\mathbf{p_i}\)) A positive sample pair matches a negative sample pair (\(\mathbf{d_i}\), \(\mathbf{p_j}\)), of which\(\mathbf{p_j}\)surgical sequences from another patient sampled at random. We then use Binary Cross-Entropy (BCE) loss to optimize the model parameters:

\[L_{SMP} = - \log(\hat{y}_i) + \log(1-\hat{y}_j), \]

included among these\(\hat{y}_i=\sigma(W_1\mathbf{r}_i + b_1)\in\mathcal{R}\) denotes the predicted probability of a positive sample pair, the\(\hat{y}_j\)denotes the predicted probability of a negative sample pair. Here.\(\sigma\)denotes the sigmoid function.\(W_1\in\mathbb{R}^{dim}\)cap (a poem)\(b_1\in\mathbb{R}\)are trainable parameters.

**Task #2: Self Reconstruction Self Reconstruction (SR): ** The goal of the SR task is to make it possible for the model to reconstruct the input sequence from the patient representation. This task will encourage RAREMed to capture and preserve as much information as possible from the input encoding. The loss function is defined as follows:

\[L_{SR} = -\sum_{j=1}^{|\mathcal{D}|+|\mathcal{P}|} \left[ \mathbf{c}_j \log(\hat{\mathbf{c}}_j) + (1-\mathbf{c}_j) \log(1-\hat{\mathbf{c}}_j) \right], \]

included among these\(\hat{\mathbf{c}} = \sigma(W_2\mathbf{r}+b_2)\in[0,1]^{|\mathcal{D}|+|\mathcal{P}|}\)denotes the probability of all diseases and surgeries reconstructed by RAREMed.\(W_2\in\mathbb{R}^{(|\mathcal{D}|+|\mathcal{P}|)\times dim}\)cap (a poem)\(b_2\in\mathbb{R}^{|\mathcal{D}|+|\mathcal{P}|}\)are learnable parameters. Here.\(\mathbf{c}\in\{0,1\}^{|\mathcal{D}|+|\mathcal{P}|}\)Indicates the true label. Only if the corresponding label is present in the input sequence will the\(\mathbf{c}_j\)Set to 1.

Fine-tuning and Inference (Fine-tune and Inference)

After pre-training, we fine-tune RAREMed on the drug recommendation task to adapt it to the downstream drug recommendation task. To predict drugs, we integrate a multi-label classification layer and utilize patient representations as inputs:.

\[\hat{\mathbf{o}} = \sigma(W_3 \mathbf{r} + b_3), \]

where \(\hat{\mathbf{o}}\in[0,1]^{|\mathcal{M}|}\)is the probability that the drug is recommended.\(W_3\in\mathbb{R}^{|\mathcal{M}|\times dim}\)cap (a poem)\(b_3\in\mathbb{R}^{|\mathcal{M}|}\)are learnable parameters.

We optimize the model parameters using the following loss function:

\[L_{bce} =-\sum_{i=1}^{|\mathcal{M}|} \left[ \mathbf{m}_i \log(\hat{\mathbf{o}}_i) + (1-\mathbf{m}_i) \log(1-\hat{\mathbf{o}}_i) \right].\\ L_{multi} = \sum_{i,j: \mathbf{m}_i=1, \mathbf{m}_j=0} \frac{\text{max}(0, 1-(\hat{\mathbf{o}}_i-\hat{\mathbf{o}}_j))}{|\mathcal{M}|}.\\ L_{ddi} = \sum_{i=1}^{|\mathcal{M}|} \sum_{j=1}^{|\mathcal{M}|} \mathbf{A}_{ij} \cdot \hat{\mathbf{o}}_i \cdot \hat{\mathbf{o}}_j. \\ L = (1-\beta)((1-\alpha) L_{bce} + \alpha L_{multi}) + \beta L_{ddi}, \]

where 𝛼 and 𝛽 are hyperparameters that balance the different loss contributions.

In the reasoning process, we recommend to the patient that the probability is greater than the threshold\(\delta = 0.5\)of drugs. Thus, the final set of recommended drugs\(\hat{\mathcal{Y}}\)It can be defined as.

\[\hat{\mathcal{Y}} = \{i | \hat{\mathbf{o}}_i > 0.5, 1\leq i \leq |\mathcal{M}|\}. \]

test

Experimental setup

We used two EHR datasets, MIMIC-III and MIMIC-IV, and modeled the screening and processing of the data after previous work.

We selected Jaccard, PRAUC, F1, DDI rate and #Med as the rubrics to evaluate the recommendation accuracy and safety of each model. Among them, the higher the Jaccard, PRAUC, and F1, the more accurate the drug recommendation; the lower the DDI rate, the fewer the drug interactions and the safer the recommendation results; and the closer the #Med is to the average number of drugs prescribed by the doctor, the more reasonable the model prediction is.

Results

RAREMed is more accurate and safer than existing drug recommendation methods.

RAREMed produces fairer drug recommendation results, with significantly higher drug recommendation accuracy for patients with rare diseases than existing models
The pre-training task, the uniform sequence encoding, and the two additional embedding layers designed all have a positive impact on recommendation accuracy and fairness

Impact of hyperparameters on the results of RAREMed recommendations