Fine-grained List-wise Alignment for Generative Medication Recommendation

• Point-wise models evaluate drugs independently and overlook synergistic effects and adverse drug-drug interaction risks that prescriptions must balance.The paper frames recommendation as generating a coherent list that jointly optimizes accuracy and safety.Source: §1 · pp. 1–2
• FLAME formulates drug generation as sequential decisions that add or remove a drug, using potential-based step-wise GRPO for finer-grained feedback.Potential-based signals are combined with list-level outcome reward to assign credit to drug-level decisions.Source: §1 · p. 2, Figure 1 · p. 2
• FLAME injects structured clinical features and embeddings into the LLM token space alongside natural-language inputs to capture textual context and collaborative signals.The framework uses Llama3.1-Aloe-Beta-8B, a domain-specific LLM with enriched medical knowledge.Source: §1 · p. 2
• FLAME achieves state-of-the-art accuracy, controllable safety–accuracy trade-offs, and strong generalization across time and institutions.Experiments include MIMIC-III, MIMIC-IV, and eICU.Source: §1 · p. 2

2.1 Medication Recommendation

• Medication recommendation research includes instance-based methods using single-visit structured features and longitudinal methods modeling hospitalization histories and treatment trajectories.Examples include GameNet, SafeDrug, COGNet, MoleRec, RAREMed, and NLA-MMR; LAMO combines structured diagnoses and procedures with unstructured patient descriptions.Source: §2.1 · p. 3
• Recent LLM methods add natural-language patient context, but many systems still score drugs independently rather than model inter-drug dependencies.FLAME instead formulates recommendation as a list-wise decision process with step-wise GRPO, medical knowledge, collaborative signals, and LLM-based semantic understanding.Source: §2.1 · p. 3

2.2 LLM Fine-Tuning Approaches

• Fine-tuning approaches progress from SFT and RLHF methods such as PPO to DPO and group-relative GRPO, each addressing different training trade-offs.PPO has multi-component overhead, DPO uses static preference data but struggles with dynamic optimization, and GRPO uses group-wise relative comparisons.Source: §2.2 · p. 3
• DAPO and StepGRPO improve aspects of GRPO training, but retain list-level advantage aggregation or depend on ground-truth process data, while step-wise GRPO guides each decision with potential-based rewards.The paper presents this decomposition as a way to model the internal logic of complex tasks such as medication recommendation.Source: §2.2 · p. 3

3.1 Medication Recommendation

• Each EHR visit contains structured features and an unstructured clinical note, while the patient record tracks medications prescribed at each visit.Structured features include demographics, diagnoses, and procedures.Source: §3.1 · p. 4
• The task is to generate a medication set close to the ground-truth prescription while minimizing harmful interactions represented by a symmetric binary DDI graph.An edge indicates a known adverse interaction between two medications.Source: §3.1 · p. 4

3.2 Group Relative Policy Optimization (GRPO)

• GRPO samples candidate outputs and updates the policy using group-relative rewards, but applies the normalized full-response reward uniformly across tokens.This outcome-level credit assignment ignores token-level quality variation and can hinder fine-grained learning and optimization efficiency.Source: §3.2 · p. 4

4 Method: FLAME

• The method combines step-wise GRPO, a two-stage recommendation pipeline, and multi-source knowledge fusion.The section introduces these components in sequence as the core of FLAME.Source: §4 · p. 4

4.1 Step-wise GRPO

• LLM outputs are segmented into medication-level steps and treated as state-action transitions, enabling localized reward signals tied to global output quality.Potential-based shaping adds the difference between consecutive step potentials to the standard GRPO advantage.Source: §4.1 · pp. 4–5

4.2 Fine-grained Alignment for Medication Recommendation

• The step-level potential combines Jaccard correctness, DDI risk, and violations of the candidate-drug constraint.The potential penalizes DDI-based risk and refusal rate while measuring similarity to the ground-truth medication set.Source: §4.2 · p. 5
• Potential differences assign credit to incremental changes in the medication set, and the shaped and outcome-based reward formulations yield the same optimal policy.The potential difference also supports the step-wise GRPO policy objective.Source: §4.2 · p. 6

4.3 Two-stage Recommendation Framework

• The two-stage framework first uses an LLM binary classifier to filter candidate drugs, then applies a list-wise policy to edit the resulting set.At inference, the policy generates add and remove edits that are applied to form the final recommendation.Source: §4.3 · p. 6
• The list-wise policy is initialized from the classifier and adapted for instruction-driven editing through supervised fine-tuning and step-wise GRPO.Its complementary tasks add missing necessary drugs and remove unnecessary or potentially unsafe drugs.Source: §4.3 · p. 6

4.4 Multi-source Knowledge Fusion

• For each structured entity, knowledge fusion combines a textual embedding with a projected collaborative embedding in the LLM representation space.Collaborative signals cover patient, diagnosis, procedure, and medication levels.Source: §4.4 · p. 7

5.1 Experimental Settings

• MIMIC-III is used for training and evaluation, while MIMIC-IV and eICU test generalization; correctness, DDI safety, and safety–accuracy trade-offs are evaluated.The comparison includes traditional and LLM-based methods, and examines performance under distribution shifts beyond the training domain.Source: §5.1 · p. 7, §5.2 · p. 8

5.2 Overall Performance Comparison

• FLAME outperforms all baselines in Jaccard and F1 scores on MIMIC-III.Table 1 reports the correctness comparison; #Med is average medications per patient, with a ground-truth average of 23.43.Source: §5.2 · p. 8, Table 1 · p. 8
• FLAME achieves a better accuracy–safety balance than safety-controllable baselines and shows strong temporal and external generalization.The evaluations cover safety–accuracy trade-offs on MIMIC-III, temporal transfer to MIMIC-IV, and external transfer to eICU.Source: §5.2 · p. 8, Figure 5 · p. 8

5.3 Ablation Study

• Ablations show that both individualized filtering (πcls) and list-wise modeling (πlist) matter, and that omitting SFT or step-wise preference alignment degrades results.Table 2 reports the component and training-stage comparisons.Source: §5.3 · p. 9, Table 2 · p. 9
• Replacing step-wise GRPO with standard GRPO lowers performance, while step-wise GRPO provides denser feedback, faster learning, and more stable training.Table 2 compares the GRPO variant; Figure 6 shows the training curves.Source: §5.3 · p. 9, Table 2 · p. 9, Figure 6 · p. 9
• Removing structured codes or clinical notes degrades performance, random collaborative embeddings fall short of FLAME, and general LLM backbones perform worse than the domain-specific backbone.The ablations compare structured and unstructured inputs, collaborative embeddings, and Llama3.1-Aloe-Beta-8B against LLaMA2 and LLaMA3.Source: §5.3 · p. 9, Table 2 · p. 9

5.4 In-Depth Analysis

• On MIMIC-III, πcls produces lists averaging 20.83 medications, while πlist makes about 2.03 add/remove edits per patient to address drug dependencies and redundancies.The analysis also reports that πlist can adjust safety through α; it edited 28 Piperacillin-containing lists with an 89% correctness rate, and examples show it correcting redundant and missing predictions.Source: §5.4 · pp. 9–10, Table 3 · p. 10
• 82.1% relative gain: list-wise refinement’s advantage over point-wise classification rises from 0.0028 with standard GRPO to 0.0051 with step-wise GRPO.Step-wise GRPO also shows faster reward convergence and more stable training, attributed to denser intermediate feedback rather than terminal rewards alone.Source: Figure 6 · p. 9, §5.4 · p. 10

6 Conclusion

• FLAME combines structured clinical data, collaborative signals, and unstructured text to support accurate, clinically grounded recommendations.The conclusion describes these sources as hybrid representations that enhance the LLM’s capacity.Source: §6 · p. 10
• Experiments report superior accuracy, controllable safety–accuracy trade-offs, and generalization across temporal and institutional shifts.The conclusion frames these results as evidence of potential for reliable, adaptable clinical decision support.Source: §6 · p. 10
• Future work will target generalization in diverse, evolving healthcare environments and integration of physician decisions for human-in-the-loop recommendations.Source: §6 · p. 10

A Theoretical Analysis

• The outcome-based terminal reward and shaped step-wise reward induce equivalent policies.The analysis relies on potential-based shaping, which does not change the optimal policy.Source: A Theoretical Analysis · p. 13

B.1 Data Preprocessing

• Structured diagnoses and procedures are paired with concise ICD descriptions, while GPT-4o compresses verbose dictionary titles.The approach follows prior preprocessing work and LAMO [10] to preserve semantic content and model compatibility.Source: §B.1 · p. 13
• Medication codes are mapped from National Drug Codes to DrugBank IDs, and selected clinical-note sections are parsed into patient condition descriptions.Extracted sections include history of present illness, past medical history, allergies, and medications on admission.Source: §B.1 · p. 13

B.2 Dataset Statistics

• Table 4 summarizes statistics for the processed MIMIC-III data.Source: Table 4 · p. 14, §B.2 · p. 14

C.1 Training Details and Hyperparameters

• Training runs on NVIDIA A100-SXM4-80GB GPUs with bf16 quantization, LoRA, and the adamw_torch optimizer.The environment uses Python 3.10 and PyTorch 2.5.1.Source: §C.1 · p. 14
• Drug-level classifier SFT uses Llama3-Aloe-8B-Alpha, four two-layer GELU projectors, learning rate 5e-4, batch size 128, and one epoch.The projectors encode patient, diagnosis, procedure, and medication information.Source: §C.1 · p. 14
• List-wise policy SFT initializes from the classifier’s model and projector weights, using learning rate 5e-4, batch size 64, and one epoch.Source: §C.1 · p. 14
• Step-wise GRPO initializes from the list-wise SFT weights and uses learning rate 1e-5, batch size 16, eight generations, and one epoch.The DDI-related hyperparameter α is selected from [0, 2, 5, 10, 20, 30, 40, 50], while β is 0.5 and λ is 5.Source: §C.1 · p. 14

C.2 Baseline Methods and Evaluation Metrics

• LEAP uses current-visit data, while GAMENet and SafeDrug incorporate DDI graphs or molecular encoders; COGNet, MICRON, and MoleRec model prescription history, medication changes, or drug substructures.These baseline descriptions span single-visit prediction, safety-aware models, and approaches that incorporate historical prescriptions or fine-grained drug structure.Source: §C.2 · pp. 14–15
• RAREMed uses a Transformer-based patient encoder, NLA-MMR aligns patient and medication representations across modalities, and LAMO incorporates unstructured clinical notes.The descriptions identify different representation strategies used by the baselines.Source: §C.2 · p. 15
• Jaccard Score and F1 Score assess prediction accuracy, while DDI rate measures harmful interactions among predicted medications.The DDI rate is defined as the ratio of harmful drug-drug interactions among predicted medications.Source: §C.2 · p. 15

C.3 Prompt Templates

• The drug-level classifier evaluates each candidate drug independently for appropriateness using the patient representation and a textual description of that drug.Its prompt requests a yes-or-no judgment about whether the candidate is effective and safe for the patient.Source: §C.3 · p. 15
• The list-level policy refines the classifier’s predicted set by adding missing clinically necessary drugs and removing unnecessary or potentially unsafe drugs.The two tasks use separate prompts to revise the pre-predicted list in light of the patient’s condition and medical history.Source: §C.3 · p. 16
• Patient prompts combine current and, when available, previous hospitalization information, including structured diagnoses, procedures, medications, demographics, and unstructured clinical narratives.The current representation includes information such as allergies, history of present illness, and past medical history.Source: §C.3 · p. 16
• Hybrid representations pair textual descriptions of structured codes with collaborative embeddings projected into the LLM embedding space by a learnable adapter.The embeddings are produced by domain-specific encoders and supplement, rather than replace, the textual token sequence.Source: §C.3 · p. 17

D Limitations

• LLM training and inference overhead can constrain deployment in resource-limited clinical settings, particularly for real-time or on-device use.The authors identify computational cost and deployment challenges as limitations warranting further investigation.Source: D Limitations · p. 17
• FLAME lacks end-to-end automated multimodal processing, and manually invoking external APIs for patient narratives limits scalability and consistency.This limitation applies despite the framework’s use of both structured and unstructured patient data.Source: D Limitations · p. 17
• Training and evaluation on static EHR datasets do not fully capture dynamic clinical workflows, leaving adaptation to continual learning or live decision support for further investigation.Source: D Limitations · p. 17

E Broader Impacts

• Training-data bias may produce disparities in recommendation quality, especially for patient populations underrepresented in historical records.Source: E Broader Impacts · p. 18
• Deployment without proper clinical oversight risks harmful or misleading recommendations, including in unauthorized or nonmedical settings.The authors call for responsible deployment practices, transparency, and safeguards.Source: E Broader Impacts · p. 18