SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Suman Bhoi
Mong Li Lee
Wynne Hsu
Ngiap Chuan Tan
$^{1}$ National University of Singapore  $^{2}$ SingHealth Polyclinics sumanbhoi@u.nus.edu, {leeml,whsu} @comp.nus.edu.sg, tan.ngiap.chuan@singhealth.com.sg
Patients with co-morbidities often require multiple medications to manage their conditions. However, existing medication recommendation systems only offer class-level medications and regard all interactions among drugs to have the same level of severity. This limits their ability to provide personalized and safe recommendations tailored to individual needs. In this work, we introduce a deep learning-based fine-grained medication recommendation system called REFINE, which is designed to improve treatment outcomes and minimize adverse drug interactions. In order to better characterize patient's health conditions, we model the trend in medication dosage titrations and lab test responses, and adapt the vision transformer to obtain effective patient representations. We also model drug interaction severity levels as weighted graphs to learn safe drug combinations and design a balanced loss function to avoid overly conservative recommendations and miss medications that might be needed for certain conditions. Extensive experiments on two real-world datasets show that REFINE outperforms state-of-the-art techniques.
The widespread adoption of digital patient records in hospitals has led to the creation of a rich knowledge base in the form of Electronic Health Records (EHR). This valuable resource can be utilized for various tasks in the medical domain, such as mortality prediction [19], outcome prediction [14], and medication recommendation [16, 2]. Concurrently, advances in drug therapy have resulted in growing interest in research aimed at enabling personalized and safe treatment recommendations for managing patient conditions. Current medication recommendation systems are coarse-grained, providing recommendation at the class level rather than individual drugs, and are unable to differentiate between patients who require single versus multiple drugs from the same class. For example, while diabetic patients are often prescribed Metformin to manage their blood sugar levels, patients with uncontrolled blood sugar level may be given additional blood glucose lowering drug such as Glipizide which belong to the same medication class  $\mathrm{A10B^1}$  as Metformin. As such, the recommended treatment regime is the same for both types of patients which is suboptimal. Thus, it is essential to develop a fine-grained recommender system capable of providing recommendations at the drug level.
Existing recommender systems that consider drug interactions assume that all the interactions have the same severity. This is not true in practice as resources on drug information such as the DDIter database [23] record different levels of interaction severity, ranging from mild to moderate to severe. By assuming all interactions are equally severe, existing recommender systems tend to avoid any combination of drugs with known interactions, even when the benefits outweigh the potential risks.
1https://www.whocc.no/atc_ddd_index/?code=A10B&showdescription=no
For example, in the benchmark MIMIC-IV [10], Metformin and Lisinopril are prescribed together  $36\%$  of the time for patients with diabetes and hypertension to better manage their conditions, despite having a moderate level of interaction severity.
In this work, we propose a fine-grained medication recommender system that takes into account the severity of drug interactions and leverages information gathered over past visits to provide safe and personalized drug treatment regimes. We represent the varying severity of drug interactions as a weighted graph and employ a graph attention network to obtain drug interaction representations, enabling the system to learn safe drug pair combinations. Moreover, to avoid overly conservative recommendations, we design a balanced loss function that weighs the benefits of a pair of drugs against the severity of any potential drug interaction. This approach allows for more effective and personalized treatment recommendations while minimizing the risk of adverse drug interactions. Additionally, since clinicians often titrate treatment regimes for patients based on their past lab test responses and medication dosages, we incorporate such trend information and adapt the vision transformer [7] (ViT) to learn an effective patient representation.
Our contributions can be summarized as follows: (1) We design a novel deep learning system that combines knowledge from patient information over multiple visits and severity-based drug interaction database to provide personalized, safe, and accurate fine-grained medication recommendation. (2) We propose to extract and model trend information from the medication dosage titrations and lab test responses over multiple visits to obtain effective patient representation. This enables the system to adjust the treatment regime as the patient condition changes. (3) We design a balanced DDI loss function that weighs the benefits of a pair of drugs against the severity of any potential drug interaction. (4) We provide comprehensive quantitative and qualitative evaluations on a benchmark medical dataset MIMIC-IV [10] and a proprietary outpatient dataset to demonstrate the capability and effectiveness of our proposed approach in providing fine-grained medication recommendations while considering drug interaction severity and patient history.
Early works in medication recommendation focus on learning a collection of rules from EHR. Solt and Tikk [17] extract rules from discharge summaries, while Lakkaraju et al. [11] use Markov Decision Processes to learn the mapping between the patient characteristics and treatments. However, this approach may introduce conflicting rules, and is difficult to generalize and scale. Subsequent works employ recurrent neural networks to model sequential dependency in patient's past visits to improve the performance of medication recommendation [1, 12]. However, they do not consider drug interactions, which are crucial to minimize adverse drug reactions.
LEAP [27] incorporates an external drug-drug interaction (DDI) database and uses a recurrent decoder to model the drug-disease relationship in the current visit for recommending a set of medications. GameNet [16] uses dynamic memory and Graph Convolutional networks (GCN) to personalize the medication recommendations based on patients' longitudinal visit history and drug interactions. CompNet [21] is another GCN based recommender system that models relation among patient's past visits as well as the drug interactions by combining reinforcement learning with relational GCN.
PREMIER [2] introduces a two-stage recommender system that employs dual attention-based recurrent neural networks to model patient visits and adapts graph attention networks to model drug co-occurrences and drug interactions. SafeDrug [25] uses RNN based patient visit representation and incorporates drug molecular structure information to enhance the safety of the recommended drugs. They introduce a threshold during training to penalize the recommended set of medications if the total drug interaction in the set is above the threshold. MICRON [24] uses recurrent residual networks to capture the medication changes over the visits, and incorporate drug interaction information from an external database in the design of their DDI loss function.
COGNet [22] introduces a copy module, which generates the probability of a medication being repeated from past visits. They model the drug interaction information with the help of a GCN and a DDI loss function. MoleRec [26] introduces a hierarchical architecture to model the association between drug substructures at the molecular level and the target disease to provide personalized medication recommendation. [18] proposes a causal inference based medication recommendation model. They model patient visit information as causal graphs and DDI as propositional satisfiability
Figure 1: Overview of REFINE.
problem. However, they use additional information such as patient symptoms which are not available in our datasets, making a direct comparison difficult.
Our approach differs from existing works in that we provide fine-grained medication recommendation by explicitly modeling drug names, and the trend information in the medication dosages as well as the lab test responses gathered over multiple visits. We also ensure the safety of the recommendations by modeling the severity of drug interactions to reduce the risk of harmful combinations.
Figure 1 shows the framework of the proposed system REFINE for REcommendation of FINE-grained medication. The input is a list of patient visits. The  $i^{th}$  visit is represented as  $\boldsymbol{x}_i = [d_i, l_i, m_i]$  where  $d_i$  is a multi-hot vector of diagnosis,  $l_i$  is a vector of lab tests and responses,  $m_i$  is the vector of medications and dosages, and  $i \in [1, T]$ . All continuous values are normalized using min-max normalization [9]. The goal is to predict a set of medications  $y_T$  for the  $T^{th}$  visit, taking into account the patient history, the trend in the medication dosages and lab test responses, as well as the severity of drug interactions.
3.1 Trend Information Extraction
We observe that the treatment regimes for patients with the same disease vary depending on whether their conditions are well managed. This is typically reflected in the fluctuations in the lab test responses over the various visits. Further, clinicians often titrate the dosage of first-line medications before introducing additional medications to achieve the optimal treatment outcome. As such, we extract trend information from medication dosages and lab test responses over the past visits in the form of slope and variance from [15].
Let  $\mathbf{var}_i^m$  and  $\mathbf{slope}_i^m$  be the vectors of dosage variance and rate of dosage change for the medications across visits 1 to  $i$  respectively. The dosage variance and rate of dosage change for the  $k^{th}$  medication can be computed as follows:
$$
\operatorname {v a r} _ {i} ^ {m} [ k ] = \frac {\sum_ {j = 1} ^ {i} \left(m _ {j} [ k ] - \mu_ {i} ^ {m} [ k ]\right) ^ {2}}{(i - 1)} \tag {1}
$$
$$
\operatorname {s l o p e} _ {i} ^ {m} [ k ] = \frac {\sum_ {j = 1} ^ {i} \left(t _ {j} - \bar {t} _ {i}\right) \left(m _ {j} [ k ] - \mu_ {i} ^ {m} [ k ]\right)}{\sum_ {j = 1} ^ {i} \left(t _ {j} - \bar {t} _ {i}\right) ^ {2}} \tag {2}
$$
Figure 2: Patient representation obtained from transformer encoders.
where  $m_j[k]$  is the dosage of the  $k^{th}$  medication at the  $j^{th}$  visit,  $\mu_i^m [k]$  is the mean of the dosages of the  $k^{th}$  medication across visits 1 to  $i$ ,  $t_j$  is the interval (in days) between the  $1^{st}$  visit and the  $j^{th}$  visit,  $\bar{t}_i$  is the average  $t_j$  for  $j\in [1,i]$ . The variance  $\mathbf{var}_i^l$  and slope  $\mathbf{slope}_i^l$  for the lab test responses are obtained in the similar manner.
We stack the variance  $\mathbf{var}_i^m$  and slope  $\mathbf{slope}_i^m$  and medication dosages  $m_{i}$  to obtain a 2D array  $\in \mathbb{R}^{3\times N_m}$  where  $N_{m}$  is the total number of medications. We flatten the 2D array into a 1D array sequence, similar to the approach used in Vision Transformer (ViT) [7]. To do this, we concatenate the three values (medication dosage, slope, and variance) for every medication in a visit to obtain a 1D array sequence  $s_i$ . We repeat the process for the lab test response information to obtain a 1D array sequence  $r_i$ .
3.2 Patient Representation
To obtain a comprehensive representation of a patient's evolving health condition, it is crucial to capture and combine the sequential dependencies present in their EHR data. Specifically, we utilize three separate transformers [7, 20] to model the sequential dependencies in the patient's diagnosis, lab test response, and medication information over multiple visits. These transformer modules can be used individually or in combination to derive a patient representation even when some information sources are missing from the EHR data. The embeddings  $\boldsymbol{d}_i^e$ ,  $\boldsymbol{r}_i^e$ , and  $\boldsymbol{s}_i^e$  for diagnosis, lab test response, and medication dosage respectively are passed to their respective transformers to generate the patient representation  $\boldsymbol{q}_T$  as shown in Figure 2.
Each transformer is composed of multiple layers of encoders, and each encoder has a multi-head self-attention block followed by a position-wise feed-forward network. For the  $i^{th}$  visit, the encoder in the first layer of the transformer takes as input the concatenation of the respective input embeddings from visits 1 to  $i$ . The input to the subsequent encoders is the concatenation of the outputs from the encoder of the previous layer. This helps to preserve and process the information from the patient's visit history. Residual connections are employed around self-attention mechanism to facilitate the propagation of lower layer features to higher layers. To stabilize and accelerate neural network training, we apply layer normalization after the multi-head self-attention block and the position-wise feed-forward network. We also apply dropout to avoid over-fitting. The outputs of the transformers for the  $i^{th}$  visit are  $h_i^d$ ,  $h_i^r$ , and  $h_i^s$  for diagnosis, lab test response, and medication dosage respectively.
Since our goal is to predict the set of medications for the current visit  $T$ , we fuse the encoded sequential representation for the diagnosis and lab test response up to the current visit, and the encoded representation for the medication dosage information up to the previous visit to obtain the
Figure 3: Details of modeling drug information from the EHR as well as the drug interaction database.
final patient representation  $\pmb{q}_T$  as follows:
$$
\boldsymbol {q} _ {T} = \boldsymbol {h} _ {T} ^ {d} + \boldsymbol {h} _ {T} ^ {r} + \boldsymbol {h} _ {T - 1} ^ {s} \tag {3}
$$
3.3 Drug Information Representation
We observe that clinicians take into account the severity levels of interacting drug pairs and the severity of drug interactions can vary significantly. This information is crucial when prescribing medications to ensure patient safety. In addition, the frequency of co-occurrence of drug pairs observed in the EHR can provide valuable insights into commonly or rarely prescribed combinations, further guiding the decision-making process.
We represent drug co-occurrence information from the EHR and drug interaction from the DDI database DDInter [23] as two separate graphs  $G^{C} = (V, E^{C})$  and  $G^{D} = (V, E^{D})$  respectively. Each node  $n_{i} \in V$  depicts a drug  $i$ . The weighted edges in  $E^{C}$  and  $E^{D}$  depict the co-occurrence frequency and interaction severity between drug pairs, respectively. For co-occurrence, a weighted edge  $(\pmb{n}_i, \pmb{n}_j, w_{ij}^C) \in E^C$  gives the co-occurrence of drug  $i$  and drug  $j$  with a frequency of  $w_{ij}^{C}$ . The co-occurrence frequency is calculated by counting the number of occurrences of a drug pair in the prescribed medication sets in the EHR. These values are normalized to range between 0 and 1.
For drug interactions, a weighted edge  $(\pmb{n}_i,\pmb{n}_j,w_{ij}^D)\in E^D$  gives the interaction between drug  $i$  and drug  $j$  with a severity of  $w_{ij}^{D}$ . The severity information is given as Major, Moderate, Minor, or Unknown/No interaction in the DDIter database. In this work, we represent these severity levels numerically as 1, 0.66, 0.33, and 0, respectively.
We adapt GATv2 [5] to learn node representations from the weighted drug interaction graph  $G^{D}$ . GATv2 has two layers with  $z$  attention heads in the first layer, and 1 attention head in the second layer (see Figure 3). For the first layer, the attention weight between nodes  $\pmb{n}_j$  and  $\pmb{n}_k$  at the  $b^{th}$  attention head is calculated as follows:
$$
\delta_ {j k} ^ {b} = \frac {w _ {j k} ^ {D} \times \exp \left(\boldsymbol {a} ^ {\prime} \text {L e a k y R e L U} \left(\boldsymbol {E} ^ {b} \cdot \left[ \boldsymbol {n} _ {j} \| \boldsymbol {n} _ {k} \right]\right)\right)}{\sum_ {\boldsymbol {n} _ {i} \in S _ {j}} w _ {i j} ^ {D} \times \exp \left(\boldsymbol {a} ^ {\prime} \text {L e a k y R e L U} \left(\boldsymbol {E} ^ {b} \cdot \left[ \boldsymbol {n} _ {j} \| \boldsymbol {n} _ {i} \right]\right)\right)} \tag {4}
$$
where  $S_{j}$  is the set of drugs that interact with drug  $j$ ,  $a'$  is the transpose of the weight vector of a single layer feed-forward neural network,  $E^{b^*}$  is the embedding matrix for the  $b^{th}$  attention head with  $E^{b} = [E^{b^*} \| E^{b^*}]$  [5],  $w_{jk}^{D}$  is the edge weight depicting the severity of interaction between drugs  $j$  and  $k$ , and  $\| \cdot \|$  denotes the concatenation operation. The output from the first layer for node  $n_j$  is given by:
$$
\boldsymbol {o} _ {j} ^ {D} = \left\| _ {b = 1} ^ {z} \sigma \left(\sum_ {\boldsymbol {n} _ {k} \in S _ {j}} \delta_ {j k} ^ {b} \boldsymbol {E} ^ {b ^ {*}} \cdot \boldsymbol {n} _ {k}\right) \right. \tag {5}
$$
where  $\sigma$  is the sigmoid function. With this, the output from the second layer for node  $n_j$  is given by:
$$
\boldsymbol {u} _ {j} ^ {D} = \sigma \left(\sum_ {\boldsymbol {n} _ {k} \in \mathcal {S} _ {j}} \delta_ {j k} \boldsymbol {E} \cdot \boldsymbol {o} _ {k} ^ {D}\right) \tag {6}
$$
where  $\pmb{E}$  is the embedding matrix for the second layer, and  $\delta_{jk}$  is obtained by applying Equation 4 on the output of the first layer for nodes  $j$  and  $k$  with  $w_{jk}^{D}$  updated to the average of  $\delta_{jk}^{b}$  for  $1\leq b\leq z$ .
We obtain the drug interaction representation  $Z_{D}$  by stacking  $\boldsymbol{u}_{j}^{D}$ ,  $j \in \{1,2,\dots,N_{m}\}$  and  $N_{m}$  is the total number of drugs. Similarly, we can obtain a drug co-occurrence representation  $Z_{C}$  by applying another GATv2 on the drug co-occurrence graph  $G^{C}$  and stacking the node representation  $\boldsymbol{u}_{j}^{C}, j \in \{1,2,\dots,N_{m}\}$ . The final drug information representation  $Z$  is given by:
$$
\boldsymbol {Z} = \boldsymbol {Z} _ {C} + w \boldsymbol {Z} _ {D} \tag {7}
$$
where  $w$  is a parameter that controls the relative importance of the two representations.
The final output vector  $\mathbf{y}_T$  is obtained by applying attention-based aggregation on the patient representation and the weighted drug information representation as follows:
$$
\boldsymbol {y} _ {T} = \sigma (\boldsymbol {W} \cdot \left(\boldsymbol {q} _ {T} + \boldsymbol {Z} \cdot \boldsymbol {\lambda}\right)) \tag {8}
$$
where  $\lambda = \mathrm{softmax}(\mathrm{transpose}(Z)\cdot q_T)$  and  $W$  is the gradient matrix of the transformation function in the linear layer.
We consider a medication is recommended if  $y_{T}[k] > 0.5$ ,  $k \in \{1,2,\dots,N_{m}\}$ . To understand the factors influencing the model's recommendations, Appendix A shows how we can derive the contributions of the different information sources.
We formulate the recommendation task as a multi-label prediction task and employ the widely used binary cross entropy loss  $\mathrm{L_{bce}}$  and hinge loss  $\mathrm{L_{hinge}}$  [16, 3]. To capture the trade-off between the benefits and potential interactions of a drug pair, we propose a balanced drug interaction loss function  $\mathrm{L_{bdi}}$  defined as follows:
$$
\mathrm {L} _ {\mathrm {b d i}} = \sum_ {i = 1} ^ {T} \sum_ {j = 1} ^ {N _ {m}} \sum_ {k = 1} ^ {N _ {m}} \left(\beta w _ {j k} ^ {D} - (1 - \beta) w _ {j k} ^ {C}\right) \cdot y _ {i} [ j ] \cdot y _ {i} [ k ] \tag {9}
$$
where  $w_{jk}^{D}$  and  $w_{jk}^{C}$  depict the severity of drug interaction and the frequency of co-occurrence obtained from the graphs  $G^{D}$  and  $G^{C}$  respectively,  $\beta$  is a controllable factor which provides a balance between drug interaction and drug co-occurrence information.
The final objective function is given by:
$$
\mathrm {L} _ {\text {t o t a l}} = \gamma_ {1} * \mathrm {L} _ {\text {b c e}} + \gamma_ {2} * \mathrm {L} _ {\text {h i n g e}} + (1 - \gamma_ {1} - \gamma_ {2}) * \mathrm {L} _ {\text {b d i}} \tag {10}
$$
where  $\gamma_{1}$  and  $\gamma_{1}$  are hyperparameters.
We implement REFINE in PyTorch and carry out the training on two NVIDIA Titan RTX GPU. We use two datasets, MIMIC-IV [10] and PRIVATE, in our experiments. MIMIC-IV is a publicly available inpatient dataset consisting of 299,712 patient (mostly ICU patients) data collected between 2008 and 2019, whereas PRIVATE is a proprietary outpatient dataset from 6 primary care clinics over a span of ten years. Table 1 gives the statistics. Patients with fewer than two visits are filtered out (see details in Appendix B).
We evaluate the accuracy as well as the level of drug interactions in the recommended medication set. For accuracy, we use the metrics precision-recall area under the curve (AUC), F1, and Jaccard
Table 1: Characteristics of datasets.
Attribute	MIMIC-IV	PRIVATE
Number of patients	61310	85039
Number of diagnosis	2000	20
Number of lab tests	573	18
Number of medications	735	50
Avg number of visits per patient	4.0	6.9
Avg number of diagnosis per visit	12.5	1.8
Avg number of lab tests per visit	35.6	5.6
Avg number of medications per visit	8.9	4.1
Table 2: Results of comparative study.
Methods	MIMIC-IV	PRIVATE
AUC	F1	wDDI	AUC	F1	wDDI
LR	0.632 ± 0.006	0.526 ± 0.004	0.055 ± 0.0009	0.639 ± 0.007	0.561 ± 0.002	0.146 ± 0.0004
RF	0.635 ± 0.004	0.531 ± 0.003	0.053 ± 0.0005	0.642 ± 0.003	0.565 ± 0.005	0.140 ± 0.0003
BART	0.637 ± 0.002	0.540 ± 0.002	0.056 ± 0.0006	0.646 ± 0.004	0.571 ± 0.004	0.142 ± 0.0005
LEAP	0.625 ± 0.003	0.514 ± 0.005	0.047 ± 0.0003	0.633 ± 0.005	0.558 ± 0.007	0.117 ± 0.0001
DMNC	0.633 ± 0.006	0.528 ± 0.002	0.056 ± 0.0002	0.640 ± 0.007	0.566 ± 0.002	0.158 ± 0.0005
GameNet	0.641 ± 0.003	0.543 ± 0.005	0.053 ± 0.0005	0.647 ± 0.001	0.575 ± 0.004	0.131 ± 0.0003
CompNet	0.638 ± 0.002	0.546 ± 0.006	0.050 ± 0.0002	0.645 ± 0.005	0.582 ± 0.001	0.134 ± 0.0005
PREMIER	0.655 ± 0.001	0.565 ± 0.002	0.047 ± 0.0001	0.669 ± 0.001	0.607 ± 0.003	0.122 ± 0.0001
MICRON	0.658 ± 0.003	0.567 ± 0.006	0.047 ± 0.0006	0.662 ± 0.002	0.597 ± 0.006	0.120 ± 0.0001
SafeDrug	0.652 ± 0.004	0.564 ± 0.002	0.045 ± 0.0001	0.671 ± 0.003	0.608 ± 0.006	0.113 ± 0.0002
COGNet	0.661 ± 0.002	0.573 ± 0.001	0.054 ± 0.0006	0.682 ± 0.005	0.616 ± 0.002	0.129 ± 0.0004
MoleRec	0.670 ± 0.004	0.579 ± 0.003	0.044 ± 0.0001	0.688 ± 0.002	0.623 ± 0.006	0.112 ± 0.0001
REFINE	0.708 ± 0.001*	0.635 ± 0.003*	0.043 ± 0.0001*	0.729 ± 0.004*	0.656 ± 0.003*	0.111 ± 0.0002*
* indicates that the result is statistically significant when compared to the second best with p-value  $< {0.05}$  .
similarity [2, 25]. For the level of drug interactions, we introduce a weighted drug interaction score (wDDI) defined as follows:
$$
\mathrm {w D D I} = \frac {1}{T} \sum_ {i = 1} ^ {T} \left(\frac {2}{| M _ {i} | | M _ {i} - 1 |} \sum_ {j \in M _ {i}} \sum_ {k \neq j \in M _ {i}} w _ {j k} ^ {D}\right) \tag {11}
$$
where  $M_{i}$  is the set of recommended medications at the  $i^{th}$  visit for a patient,  $|M_{i}|$  is the cardinality of  $M_{i}$ ,  $w_{jk}^{D}$  is the severity of interaction between drug  $j$  and drug  $k$  obtained from the graph  $G^{D}$ . The values of the various metrics are averaged over all the patients in the test set.
We split the datasets into training, validation and test sets in the ratio 4:1:1, and report the performance on the test set. Appendix C provides details of the training and hyperparameter settings. Sensitivity experiments and additional comparative and ablation studies can be found in Appendix D.
4.1 Comparative Analysis
We compare the performance REFINE with baseline methods such as Logistic Regression (LR) [13], Random Forest (RF) [4], BART [6] LEAP [27], DMNC [12], GameNet [16], CompNet [21], PREMIER [2], MICRON [24], SafeDrug [25], COGNet [22], and MoleRec [26]. We perform 10 rounds of bootstrapping sampling and report the mean and standard deviation.