MGRec: Structure-Grounded Medication Recommendation via Condition-Aware Molecular Representation Learning

1 Introduction

The rapid digitalization of healthcare has led to an unprecedented accumulation of patient data, significantly transforming clinical decision-making [Senbekov et al., 2020; Agrawal and Prabakaran, 2020]. Among these advancements, medication recommendation systems have emerged as essential tools for optimizing treatment strategies by leveraging vast repositories of Electronic Health Records (EHRs) [Dagliati et al., 2021; Garriga et al., 2022]. Recent advances in artificial intelligence (AI) have enabled the development of predictive models that integrate EHR data to provide more informed and data-driven medication recommendations [Zhang et al., 2017; Shang et al., 2019; Yang et al., 2021b]. These models aim to enhance therapeutic efficacy by identifying optimal drug 1 regimens tailored to patient characteristics.

*: corresponding author

Early medication recommendation models are primarily historical data-driven methods, learning from co-occurrence patterns between patient conditions and discrete medication sets [Choi et al., 2016; Zhang et al., 2017; Shang et al., 2019; Liu et al., 2023; Yang et al., 2021a; Wu et al., 2022], as illustrated in Figure 1(a). While effective in routine clinical scenarios, these methods exhibit limited generalization in sparse settings applied for rare or emerging diseases, where the corresponding medication combinations are rarely observed or entirely absent from historical data. [Jones et al., 2008; Schieppati et al., 2008; Esteva et al., 2019; Faviez et al., 2020]. Recent advancements [Yang et al., 2021b; Yang et al., 2023; Kuang and Xie, 2024] have incorporated molecular substructure to enrich the drug representations, and capture finer-grained associations between patient conditions and pharmacological characteristics, thereby achieving improved predictive performance, as shown in Figure 1(b). However, these methods still fundamentally rely on observed patient-medication data to drive learning, with structural information serving to enhance drug representations rather than providing independent supervision to actively learn clinically plausible treatment associations that are not explicitly present in the data, leading to constrained generalization in data sparse settings.

1 In this paper, "medication" and "drug" are used interchangeably to refer to substances used for the treatment of diseases.

These limitations suggest that when supervision over patient-medication data is scarce, simply enriching drug representations is insufficient, as learning signals remain coupled to observed prescription labels [Kou et al., 2024; Kou et al., 2025]. A more effective strategy is therefore to shift the learning focus away from direct prediction of discrete medication sets to treating molecular structure itself as the modeling target, where the model can be trained to infer the structural and pharmacological characteristics that plausible medications should possess. In this way, structural information directly contributes learning signals through a continuous representation space, rather than relying solely on whether specific prescriptions are observed. Such a reformulation allows the model to exploit smooth relationships induced by both patient similarity and molecular similarity, providing a more suitable foundation for medication recommendation under sparse supervision. Moreover, this formulation naturally enables safety considerations, such as drug-drug interaction (DDI) risks, to be incorporated directly into the representation learning process, rather than imposed as post-hoc constraints on discrete recommendations.

Motivated by this reformulation, we propose the Medication Structure Grounded Recommendation Framework (MGRec), which grounds medication recommendation in a continuous molecular representation space while retaining discrete medications as clinically actionable outputs. As illustrated in Figure 1(c), MGRec integrates patientspecific clinical context with molecular structural knowledge through a Therapeutic-Safety Factorized Conditional Variational Autoencoder (TSF-CVAE). TSF-CVAE performs clinical-molecular context fusion to construct a unified condition representation, and learns a factorized latent space that disentangles therapeutic relevance from safety-sensitive molecular factors. The model infer molecular representations that capture core therapeutic characteristics for the current patient conditions. Final recommended medications labels are instantiated by mapping inferred representations to the closest clinically approved drugs, ensuring practical applicability.

Meanwhile, to enhance the clinical safety of the recommended drug combinations, MGRec incorporates DDIGuided Latent Regularization, which injects pharmacological interaction knowledge into the safety subspace and integrates safety considerations directly into the learned condition-aware molecular representations. MGRec not only improves recommendation accuracy and safety for rare and emerging diseases. In conclusion, the main contributions of this paper can be summarized as follows:

• We propose MGRec, a medication recommendation framework that shifts the learning objective from direct discrete prescription prediction to patient condition-aware modeling
• in a continuous molecular representation space, improving robustness under sparse supervision.
• We introduce a Therapeutic-Safety Factorized Conditional VAE (TSF-CVAE) that disentangles therapeutic and safetyrelated molecular factors, together with a DDI-guided latent regularization mechanism that internalizes pharmacological safety constraints at the representation level.
• Extensive experiments on two real-world benchmark datasets demonstrate that MGRec achieves state-of-the-art performance in both recommendation accuracy and interaction risk reduction, especially in data-sparse settings.

2 RelatedWork

2.1 Medication Recommendation

Existing approaches for medication recommendation can be broadly categorized into historical data-driven methods and molecular substructure-informed methods. Historical datadriven methods rely on patient EHR data to recommend medications, and can be further divided into instance-based methods and longitudinal methods. Instance-based methods such as LEAP [Zhang et al., 2017] focus on the patient information from a specific visit. Longitudinal methods incorporate historical records to capture temporal dependencies. An example is GAMENet [Shang et al., 2019], which integrates memory neural networks with graph convolutional networks to encode longitudinal EHR data. Molecular substructureinformed methods integrate drug molecular structures to enhance recommendation accuracy and safety. These methods include SafeDrug [Yang et al., 2021b], which incorporates molecular structure information to improve medication safety, and MoleRec [Yang et al., 2023], which directly associates patient health conditions with molecular substructures. Notably, RAREMed [Zhao et al., 2024] is the most relevant work to ours, as it explicitly focuses on rare disease medication recommendation. It leverages pre-training to improve medication prediction for rare disease cases, while remaining constrained by its reliance on historical patient records.

2.2 VAE for Condition-Aware Molecular Representation Learning

Variational Autoencoders (VAEs) [Kingma, 2013] provide a principled framework for modeling complex data distributions in a continuous latent space and have been widely adopted for learning molecular representations with semantic structure [Wigh et al., 2022; Zeng et al., 2022; Shi and others, 2025; Ochiai et al., 2023]. In the molecular domain, VAEs have been shown to effectively encode molecular graphs into smooth latent spaces that preserve structural similarity, enabling interpolation and uncertainty-aware inference [G´omez-Bombarelli et al., 2018; Samanta et al., 2020; Jin et al., 2018] for downstream task such as molecule generation [Xue et al., 2019; Pang et al., 2023]. In this work, we leverage VAEs not as generators of new molecules, but as conditional distributional models for learning patient-specific molecular representations to provide continuous supervision and safety control.

3 ProblemFormulation

For a patient x, the EHR is represented as a sequence V (x) = [v(1) , v(2) , . . . , v(Nx) ], where Nx is the total number of visits for the patient, and v(i) represents the details of the i-th visit. Each visit v(i) consists of three main components: v(i) = [v(i) d , v(i) p , v(i) m ]. Here, v(i) d ∈ {0, 1}|D| , v(i) p ∈ {0, 1}|P| and v(i) m ∈ {0, 1}|M| are multi-hot vectors representing the diagnosis, procedures and medications prescribed respectively. At time step t, given the longitudinal diagnosis sequence vtd = [v(1) d , v(2) d , . . . , v(t) d ], procedure sequence: vtp = [v(1) p , v(2) p , . . . , v(t) p ] and medication sequence: vt-1 m = [v(1) m , v(2) m , . . . , v(t-1) m ], our objective is to learn a function f(·) that generates a multi-label output ˆm(t) ∈ {0, 1}|M| . Specifically, ˆm(t) = f(vtd , vtp , vt-1 m ).

4 Methodology

As shown in Figure 2, our proposed MGRec consists of three components: 1) Health condition embedding module processing longitudinal patient medical records to obtain a comprehensive representation of patient health condition, 2) Conditional representation learning module learning ideal molecular representations under various contextualized conditions, 3) Drug inference module inferring personalized drugs according to patient's current health condition.

4.1 Health Condition EmbeddingModule

To comprehensively characterize a patient's health condition, we consider two key aspects: the current health condition and the previous health condition. Two embedding layers Ed and Ep are employed to transform the vtd and vtp into vectorized representations:

Subsequently, to capture temporal dependencies across multiple visits, we apply transformer-based encoders [Vaswani et al., 2017], denoted as T(·), to yield contextualized encodings of the diagnosis and procedure sequences:

The current health condition is constructed by concatenating the encoded representations at time t: ec = [h(t) d , h(t) p ] ∈ Rd . Similarly, the previous health condition at time i (where i ≤ t - 1) is obtained by aggregating the corresponding representations from the patient's prior visit (if available): e(i) p = [h(i) d , h(i) p ] ∈ Rd . For cold-start scenarios, where the patient has only one visit (i.e., t = 1), the model exclusively utilizes the current health condition for downstream tasks.

4.2 Conditional Representation LearningModule

Molecular graph encoding via GIN. We firstly employ a Graph Isomorphism Network (GIN) [Xu et al., 2018] to obtain molecular representations of drugs in the given medication sequence vtm. Each drug molecule is treated as a graph G = (V, E), where nodes vi ∈ V represent atoms, and edges eij ∈ E correspond to chemical bonds. The node feature vector zi ∈ Rd encodes atomic properties, while edge features eij ∈ Rd capture bond types, detailed summary of the selected features can be found in Appendix. The messagepassing process at the (l + 1)-th layer is defined as follows:

where h(l) i represents the hidden state of node i at layer l, N(i) denotes the neighboring nodes, and α(l) is a learnable scaling parameter that enhances the expressiveness of the aggregation process. After L layers of message passing, the final node representations are obtained as:

Then we aggregate the learned node features using global sum pooling to derive a holistic molecular representation of the drug: hm = ∑i∈V h(L) i ∈ Rd

Pharma-health contextualized condition construction. After obtaining the structural representation of each drug molecule, we further embed their category labels derived from Anatomical Therapeutic Chemical (ATC) classification system [Schellekens et al., 2011; Garbe et al., 1993]. Specifically, we utilize an embedding layer to transform each drug category label into a dense vector representation V (i) m = Emv(i) m ∈ Rd . Then we compute an adaptive gating vector:

where σ denotes sigmod activation function, Wg ∈ Rd×2d represents the learnable weight matrix, and bg ∈ Rd is the bias term. The final gated pharma-health contextualized condition c which effectively encodes both the categorical pharmacological properties and clinical context of patient's health trajectory is computed as:

where ⊗ denotes the Hadamard product and 1 represents an all-ones vector.

Therapeutic-safety factorized latent modeling. Given drugs prescribed under certain medical contextualized conditions c, we introduce a latent vector z = [zther, zsafety] that factorizes into a therapeutic subspace latent vector zther ∈ Rdther and a safety-related subspace vector zsafety ∈ Rdsafety . These disentangled latent representations are designed to capture the underlying therapeutic substructures that drive prescriptions and the safety patterns, respectively. We assume that the ideal molecular representations for the given c follow the conditional distribution

where hm is the ideal representation and θ denotes the distribution parameters.

Health Condition Embedding Module Conditional Representation Learning Module Drug Inference Module

As the integral is generally intractable, to optimize pθ(hm | c), we approximate the true posterior pθ(z | hm, c) with a variational distribution qϕ(z | hm, c). Applying the variational lower bound [Kingma, 2013; Jensen, 1906], we obtain

Both the conditional prior and posterior are parameterized as diagonal Gaussians over the concatenated latent vector z:

where µθ(c), σ2 θ (c) ∈ Rd and µϕ(hm, c), σ2ϕ (hm, c) ∈ Rd (with d = dther + dsafety) denote the mean and variance of the prior and posterior, respectively. The conditional prior pθ(z | c), parameterized by a multilayer perceptron (MLP) conditioned solely on c, provides a context-dependent latent distribution aligned with the pharma-health context.

The reconstruction term in Eq. (8) encourages the decoded representation to match the prescribed drug, while the KL term aligns the approximate posterior with the conditionaware prior. Let qϕ = qϕ(z | hm, c) and pθ = pθ(z | c). The overall objective is:

To parameterize the approximate posterior qϕ(z | hm, c) and the likelihood pθ(hm | z, c), we employ two MLPs as the conditional encoder and conditional decoder, respectively. The encoder outputs the posterior parameters

where µϕ, σ2ϕ ∈ Rd . We similarly partition the posterior parameters into therapeutic and safety components, µϕ = [µther ϕ ; µsafety ϕ ] and σϕ = [σther ϕ ; σsafety ϕ ], and apply the reparameterization trick [Kingma, 2013]:

The conditional decoder reconstructs the pharma-health contextualized molecular representation h′m from the disentangled latent vector z and condition c:

The training loss of this module consists of a reconstruction term and a KL regularization term. The reconstruction loss is implemented as the mean squared error (MSE) between hm and h′m:

The KL divergence term aligns the posterior with the condition-aware prior. Let µϕ,k, σ2ϕ,k and µθ,k, σ2 θ,k denote the k-th elements of the posterior and prior parameters, respectively. Then

4.3 Drug Inference Module

Therapeutic matching and category filtering. We first determine the necessary drug categories based on the current patient health condition. Specifically, we construct a filtering mask as a constraint on the learned drug latent space.

To achieve this, we integrate all therapeutic latent vectors with the current patient condition by an MLP to obtain a therapeutic matching score for each drug category:

where s represents the matching scores for N drug categories, σ(·) is the sigmoid activation function. Subsequently, we apply a threshold τ = 0.5 to construct a filtering mask:

where I(·) is the indicator function, returning 1 if s > 0.5 and 0 otherwise, mask ∈ {0, 1}|M| is the filtering mask applied to the drug categories.

Latent sampling for condition-aware inference. For all drug categories where the mask value equals 1, we utilize their corresponding drug category label embeddings with the patient's current status ec to obtain current contextualized conditions c′ . Subsequently, we obtain latent vectors from the condition-aware prior and use the decoder to draw latent vectors from

The latent vector is then used to decode a therapeutic molecular representation:

To obtain the final prescribed drug labels, we evaluate the embedding similarity between each generated molecular structure in ˆh and all drugs within its corresponding category. Specifically, we categorize each generated molecule as the most similar drug within its category, thereby obtaining the final recommended drug set ˆm:

where Mc represents the set of candidate drugs within the category, and Sim(·, ·) denotes the cosine similarity of molecular structural embeddings.

To ensure clinical safety, we introduce a DDI-guided latent regularization mechanism driven by the disentangled safety vector zsafety, which plays two roles: (i) learning a latent representation that reflects the structure of the DDI graph, and (ii) regulating the recommended drug set to mitigate DDIs.

Safety subspace learning via the DDI graph. To ensure that the latent space captures clinically meaningful interaction patterns, we impose a representation-level DDI supervision on zsafety. For each drug pair (i, j), we compute a latent interaction score:

where fDDI is a lightweight MLP. Given the DDI adjacency matrix D ∈ R|M|×|M| , the safety subspace is trained to reproduce the interaction structure via

where wij is an optional normalization weight to address class imbalance.

Safety regularization on the recommended set. LDDI-repr shapes the safety subspace, but it does not directly constrain the recommended drug set. We therefore introduce a representation-level safety regularizer:

where si and sj are the therapeutic matching scores assigned to drugs i and j. This loss discourages the model from simultaneously assigning high confidence to drug pairs that the safety subspace predicts as high-risk, thereby coupling safety awareness with therapeutic relevance recommendation.

Following previous work [Shang et al., 2019; Yang et al., 2021b], we treat the prediction of each drug categorical labels as an independent task and use the BCE loss for optimization:

where si is the correspond therapeutic matching score as the recommendation probability.

Overall objective. The overall training loss is:

where α, β, γ1, and γ2 are hyperparameters that respectively control the contributions of reconstruction quality, latent regularization, DDI-aware representation shaping, and safety regularization.

5 Experiments

In this section, we conduct extensive experiments to make a comprehensive evaluation of our proposed method and answer the following four questions:

• RQ1: How does the performance of the proposed MGRec compare to that of existing methods?
• RQ2: Does MGRec effectively provide reliable recommendations for patients with rare and emerging diseases?
• RQ3: How do the different components and hyperparameters of MGRec affect its performance?

5.1 Experiment Setup

Dataset. Electronic health record (EHR) data from two realworld datasets are utilized, specifically MIMIC-III [Johnson et al., 2016] and MIMIC-IV [Johnson et al., 2023]. In line with prior studies [Shang et al., 2019; Yang et al., 2021b], the datasets are processed and randomly split into training, validation, and testing sets with a ratio of 4:1:1. The statistics of the processed datasets are provided in Appendix.

Evaluation Metrics. Four commonly adopted metrics in medication recommendation [Shang et al., 2019; Yang et al., 2021b; Yang et al., 2023] are used as the evaluation metrics: Drug-Drug-Interaction Rate (DDI), Jaccard Similarity Score (Jaccard), F1-score, and Precision-Recall Area Under Curve (PRAUC). Detailed formulation can be found in Appendix.

Baseline. We compare MGRec against ten state-of-the-art medication recommendation models, covering both historical data-driven approaches and molecular substructure-informed approaches. Detailed introduction can be found in Appendix.

Experimental settings. The hyperparameters for all baseline models are selected based on their performance on the validation set. Model optimization is performed using the AdamW optimizer [Loshchilov and Hutter, 2017] with a learning rate of 1e-4 and weight decay of 1e-3. The implementation is based on PyTorch and runs on an NVIDIA Tesla V100 (16GB) GPU. All of the experiment are conducted eight times, and the mean and standard deviation of the results are reported. The hyperparameters could be found in Appendix.

5.2 Overall Performance Comparison (RQ1)

Table 1 summarizes the overall performance of all methods. MGRec achieves the best Jaccard/F1/PRAUC and the lowest DDI on both datasets. Compared with history-driven and structure-informed baselines, MGRec better balances effectiveness and safety by inferring condition-aware molecular representations and incorporating DDI supervision in the latent safety subspace. These gains persist across MIMIC-III and MIMIC-IV, indicating robust generalization under sparse supervision.

Figure 3 illustrates the recommendation accuracy of four models across patients with varying numbers of visits in the MIMIC-III dataset. MGRec consistently outperforms the baselines in overall recommendation accuracy. Notably, even for patients with only one visit, where data sparsity poses a significant challenge, MGRec maintains a comparable level of accuracy. This demonstrates MGRec's strong generalization ability to provide reliable recommendations even in datasparse scenarios.

5.3 Case Study (RQ2)

To intuitively demonstrate the advantage of MGRec in providing reliable medication recommendations for patients with emerging and rare diseases, we sample several example cases and conduct a detailed analysis of their recommendation results. Table 2 presents two representative patients from the MIMIC-III dataset: Patient E with an emerging disease and Patient R with a rare condition.

Patient E suffers from postconcussion syndrome, a condition that appears only once in the dataset. To address this unique case, MGRec specifically recommends Sultiame, an anticonvulsant and antiepileptic drug [Ben-Zeev et al., 2004], to manage postconcussion syndrome-related symptoms [Ryan and Warden, 2003; Gorman and Shahwan, 2016; Bresnahan et al., 2019]. Additionally, the model recommends Furosemide to alleviate retention of urine [Brater, 1998], demonstrating its ability to provide clinically relevant treatment strategies even for unseen conditions.

Patient R has been diagnosed with poisoning by aromatic analgesics, a condition with an occurrence rate of only 0.004% in the dataset. To treat this patient, MGRec recommends Acetylcysteine to eplenish glutathione levels and prevents liver damage [Heard, 2008], as well as Heparin, an anticoagulant, to manage coagulation defects [Oduah et al., 2016]. These recommendations align with clinical practices for poisoning cases, showcasing MGRec's capability to make precise and condition-aware medication suggestions even when handling extremely low-frequency diseases.

5.4 Ablation Study (RQ3)

To examine the contribution of each key component in MGRec, we design the following ablation variants. For clarity, we use concise abbreviations to denote each removed component.

• w/o Dr removes the DDI-guided latent regularization, including both representation-level DDI supervision and setlevel safety regularization, disabling explicit safety modeling in the latent space.
• w/o Fz removes the therapeutic-safety factorization by replacing the disentangled latent representation with a single unified latent variable.
• w/o Te replaces the Transformer-based temporal encoder in the health condition embedding module with a simple embedding layer.
• w/o Ge replaces the GIN-based molecular graph encoder with a drug embedding layer, removing explicit molecular structural reasoning.

Table 3 reports the ablation results on the MIMIC-III dataset. Removing any component leads to consistent performance degradation, confirming the necessity of each design choice. In particular, w/o Dr results in a substantial increase in DDI rate, demonstrating the importance of incorporating pharmacological interaction knowledge directly into the latent space. The variant w/o Fz exhibits both reduced accuracy and increased DDI risk, indicating that explicitly disentangling therapeutic relevance from safety-sensitive factors enables better control over the effectiveness-safety trade-off.

Moreover, removing the Transformer encoder (w/o Te) or the molecular graph encoder (w/o Ge) degrades predictive performance, highlighting the importance of modeling longitudinal patient conditions and molecular structural information, respectively. Notably, the performance drop caused by w/o Ge is also accompanied by a higher DDI rate, underscoring the role of molecular structure in mitigating adverse drug interactions.

5.5 Hyperparameter Study (RQ3)

To understand the influence of hyperparameters on the efficacy of MGRec, we considered four hyperparameters: the reconstruction weight (α), KL regularization weight β, representation-level DDI regularization weight γ1 and generation-level safety regularization weight γ2.

Fig. 4 illustrates the impact of key hyperparameters on MGRec. As shown in Fig. 4(a), the reconstruction weight

α controls the trade-off between preserving molecular structure fidelity and learning abstract therapeutic representations, where moderate values yield the best accuracy and safety balance. Fig. 4(b) shows that the KL regularization weight β stabilizes the condition-aware latent space and improves generalization under sparse supervision, while overly large values lead to latent collapse and reduced performance.

Fig. 4(c) and Fig. 4(d) examine the two DDI-related weights γ1 and γ2. Increasing γ1 strengthens the alignment between the safety subspace and the DDI graph, effectively reducing interaction risk, whereas excessive emphasis may interfere with therapeutic representation learning. Similarly, γ2 directly penalizes high-risk co-recommendations at the set level, leading to lower DDI rates but potentially conservative recommendations when overly large. All hyperparameters are selected based on validation performance to balance effectiveness and safety.

6 Conclusion

In this work, we present MGRec, a structure-grounded medication recommendation framework that reframes the learning objective from discrete prescription prediction to conditionaware inference in a continuous molecular representation space. Rather than relying on observed co-occurrence patterns, MGRec models the pharmacological characteristics that plausible medications should possess given a patient's clinical context, enabling more robust generalization under sparse supervision. The proposed Therapeutic-Safety Factorized CVAE disentangles therapeutic relevance from safetysensitive molecular factors in a unified latent space, while DDI-guided latent regularization embeds pharmacological interaction constraints directly into the representation learning process. Extensive experiments on two real-world EHR benchmarks demonstrate that MGRec consistently improves recommendation accuracy while reducing drug-drug interaction risk, particularly in data-sparse scenarios.

Acknowledgments

This paper is partially supported by the National Natural Science Foundation of China (No.62502491; No.12227901). The AI-driven experiments, simulations and model training were performed on the robotic AI-Scientist platform of Chinese Academy of Sciences.Anhui Science Foundation for Distinguished Young Scholars (No.1908085J24).