Fine-grained List-wise Alignment for Generative Medication Recommendation

This repository provides the official PyTorch implementation for our NeurIPS 2025 Spotlight paper:

Fine-grained List-wise Alignment for Generative Medication Recommendation

📄 Paper: https://openreview.net/pdf?id=Quo3XadYcZ

📌 Poster:

Overview

We propose FLAME, a generative medication recommendation framework that models prescription generation as a drug-by-drug decision process.

The framework follows a two-stage pipeline:

Drug-level filtering (π_cls)
An LLM-based classifier predicts whether each candidate drug should be prescribed for a patient.
List-wise refinement (π_list)
A policy model refines medication sets through add/remove actions, optimizing the overall prescription quality.

To provide fine-grained learning signals, we introduce step-wise GRPO, which assigns rewards at each medication decision step.
This enables the model to consider drug interactions, safety constraints (DDIs), and list-level prescription quality during generation.

If you find this work useful, please consider citing:

@article{fan2026fine,
  title={Fine-grained list-wise alignment for generative medication recommendation},
  author={Fan, Chenxiao and Gao, Chongming and Shi, Wentao and Gong, Yaxin and Zihao, Zhao and Feng, Fuli},
  journal={Advances in neural information processing systems},
  volume={38},
  pages={48037--48061},
  year={2026}
}

Data Preparation

1. MIMIC-III Dataset

You must first obtain access to the MIMIC-III dataset from PhysioNet.

https://physionet.org/content/mimiciii/

After downloading the dataset, place the following files under:

data/data_process/input/mimic-iii

Required files:

ADMISSIONS.csv  
DIAGNOSES_ICD.csv  
PROCEDURES_ICD.csv  
PATIENTS.csv  
PRESCRIPTIONS.csv  
NOTEEVENTS.csv

2. DrugBank Data

Download the following file and place it under:

data/data_process/input

drugbank_drugs_info.csv

Download link:

https://drive.google.com/file/d/1EzIlVeiIR6LFtrBnhzAth4fJt6H_ljxk/view?usp=sharing

Preprocess the Data

Step 1: Filter & process raw MIMIC-III data

python data/process.py

Output:

data/data_process/output/mimic-iii/*

Step 2: Process unstructured clinical notes

This step extracts useful information from clinical notes.

Requirements:

GPT API access

Run the notebook:

data/generate_note.ipynb

Output:

data/data_process/output/mimic-iii/data4LLM_with_note.csv

This step adds a NOTE column to:

data/data_process/output/mimic-iii/data4LLM.csv

and produces:

data/data_process/output/mimic-iii/data4LLM_with_note.csv

Step 3: Generate external embeddings

Run the following notebooks:

data/saved_embedding/get_embed/get_embed_diagpro.ipynb  
data/saved_embedding/get_embed/get_embed_med.ipynb  
data/saved_embedding/get_embed/get_embed_pat.ipynb

Output files:

data/save_embedding/pat_embed_raremed.pkl
data/save_embedding/diag_embed_micron.pkl  
data/save_embedding/pro_embed_micron.pkl  
data/save_embedding/med_embed_molebert.pkl

These scripts require:

data/data_process/output/mimic-iii/voc_final.pkl

as well as pretrained models from:

MICRON
RAREMed
Mole-BERT

GitHub links:

For convenience, we also provide the generated embedding files here:

https://huggingface.co/datasets/cxfan/FLAME_embedding/tree/main

Running the Code

All experiments were conducted on NVIDIA A100 80GB GPUs.

Approximate training cost:

#Admissions	Classifier SFT	List-wise Policy SFT	List-wise Policy Step-wise GRPO
8,741	~200 GPU-hours (8×25h)	~25 GPU-hours (8×3h)	~100 GPU-hours

Considering the high computational cost for reproducing the full training pipeline, we provide the trained model checkpoints, including:

Drug-level classifier π_cls
List-wise policy π_list (after step-wise GRPO training)

Download them from:

https://huggingface.co/cxfan/FLAME_ckpts/tree/main

If you plan to directly evaluate the pretrained models, you only need to run:

Step 1 (part 3)
Step 3 (part 2)

Step 1: Drug-level Classifier (π_cls)

1. Generate training prompts

python ./data/generate_cls_data.py  --output_dir ./data/cls --num_val_visits 100 --save_file_name cls_train_val100

Output:

data/cls/cls_train_val100.json

2. Train π_cls

Download:

Llama3.1-Aloe-Beta-8B
https://huggingface.co/HPAI-BSC/Llama3.1-Aloe-Beta-8B

Place it under:

Models/

Run training:

torchrun --nproc_per_node=X --master_port=XXX ./src/finetune_ddp.py \
    --llm_name "Llama3-Aloe-8B-Alpha" \
    --data_path data/cls/cls_train_val100.json \
    --output_dir outputs/cls \
    --use_pat_embed True \
    --use_dias_embed True \
    --use_pro_embed True \
    --use_drug_embed True \
    --pat_embed_table_path data/saved_embedding/pat_embed_raremed.pkl \
    --dias_embed_table_path data/saved_embedding/diag_embed_micron.pkl \
    --pro_embed_table_path data/saved_embedding/pro_embed_micron.pkl \
    --drug_embed_table_path data/saved_embedding/med_embed_molebert.pkl \
    --eval_epochs 1000 \
    --batch_size 128 \
    --micro_batch_size 1 \
    --num_epochs 1 \
    --learning_rate 5e-4 \
    --early_stopping_patience 20 \
    --med_num 151 \
    --val_set_size 15100

3. Evaluate on test set

Generate test data:

python ./data/generate_cls_data.py  --output_dir ./data/cls --save_file_name cls_test --generate_test True

Output:

data/cls/cls_test.json

Evaluate:

torchrun --nproc_per_node=X --master_port=XXX ./src/evaluate_ddp.py \
                            --model_path outputs/cls/checkpoint-XXXX \
                            --data_path data/cls/cls_test.json \
                            --save_path evaluate_results/cls/test

Output:

evaluate_results/cls/test/evaluation_result.csv

Step 2: List-wise Policy (π_list)

The list-wise policy π_list aims to refine drug combinations by learning from mistakes made by π_cls, through add/remove actions.

1. Generate π_cls predictions on train set

python ./data/generate_cls_data.py  --output_dir ./data/cls --num_val_visits 0 --save_file_name cls_train

torchrun --nproc_per_node=X --master_port=XXX ./src/evaluate_ddp.py \
                            --model_path outputs/cls/checkpoint-XXXX \
                            --data_path data/cls/cls_train.json \
                            --save_path evaluate_results/cls/train

Output:

evaluate_results/cls/train/*

2. Construct training data for π_list

python ./data/generate_list_sft_data.py --evaluate_results_path evaluate_results/cls/train \
                                        --sample_num 100000 \
                                        --output_dir data/list_sft

python ./data/generate_list_sft_data.py --evaluate_results_path evaluate_results/cls/test \
                                        --sample_num 0 \
                                        --output_dir data/list_sft

python ./data/generate_list_grpo_data.py --evaluate_results_path evaluate_results/cls/train \
                                         --sample_num 100000 \
                                         --output_dir data/list_grpo

python ./data/generate_list_grpo_data.py --evaluate_results_path evaluate_results/cls/test \
                                         --sample_num 0 \
                                         --output_dir data/list_grpo

3. Train π_list via SFT

torchrun --nproc_per_node=X --master_port=XXX ./src/finetune_rethink.py \
    --checkpoint_path outputs/cls/checkpoint-XXXX \
    --data_path data/list_sft/100000/mix.json \
    --output_dir outputs/list_sft \
    --save_steps 256 \
    --batch_size 64 \
    --micro_batch_size 1 \
    --num_epochs 1 \
    --learning_rate 5e-4

4. (Optional) Evaluate π_list (SFT version)

python src/evaluate_rethink.py --model_path outputs/list_sft/checkpoint-XXXX --data_path data/list_sft/test/add.json --evaluate_result_path evaluate_results/list_sft --task add

python src/evaluate_rethink.py --model_path outputs/list_sft/checkpoint-XXXX --data_path data/list_sft/test/remove.json --evaluate_result_path evaluate_results/list_sft --task remove

python src/cal_metric.py --results_path evaluate_results/list_sft

Step 3: Fine-tune and Evaluate π_list with Step-wise GRPO (Final Model)

1. Continue from SFT checkpoint

python src/model_process.py --original_model_path outputs/list_sft/checkpoint-XXXX --output_model_path Models/list_sft
python src/finetune_rethink_step_grpo.py \
    --grpo_model_path Models/list_sft \
    --data_path data/list_grpo/mix_100000.json \
    --output_dir outputs/list_grpo \
    --batch_size 32 \
    --gradient_accumulation_steps 2 \
    --num_epochs 1 \
    --num_generations 8 \
    --learning_rate 1e-5 \
    --val_set_size 0 \
    --save_steps 100 \
    --use_vllm True \
    --use_lora True \
    --lora_r 32 \
    --lora_alpha 32 \
    --lora_target_modules "['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']" \
    --alpha 0 \
    --beta 0.2 \
    --step_reward_weight 0.5

2. Evaluate Final Model

python src/evaluate_rethink_grpo.py --base_model_path Models/list_sft \
                                    --lora_path outputs/list_grpo/checkpoint-XXXX \
                                    --data_path data/list_grpo/test/add.json \
                                    --evaluate_result_path evaluate_results/list_GRPO \
                                    --task add
python src/evaluate_rethink_grpo.py --base_model_path Models/list_sft \
                                    --lora_path outputs/list_grpo/checkpoint-XXXX \
                                    --data_path data/list_grpo/test/remove.json \
                                    --evaluate_result_path evaluate_results/list_GRPO \
                                    --task remove
python src/cal_metric.py --results_path evaluate_results/list_GRPO

Name		Name	Last commit message	Last commit date
data		data
src		src
utils		utils
.gitignore		.gitignore
README.md		README.md
nips_poster.png		nips_poster.png

Navigation Menu

Repository files navigation

Fine-grained List-wise Alignment for Generative Medication Recommendation

Overview

Data Preparation

1. MIMIC-III Dataset

2. DrugBank Data

Preprocess the Data

Step 1: Filter & process raw MIMIC-III data

Step 2: Process unstructured clinical notes

Step 3: Generate external embeddings

Running the Code

Step 1: Drug-level Classifier (π_cls)

1. Generate training prompts

2. Train π_cls

3. Evaluate on test set

Step 2: List-wise Policy (π_list)

1. Generate π_cls predictions on train set

2. Construct training data for π_list

3. Train π_list via SFT

4. (Optional) Evaluate π_list (SFT version)

Step 3: Fine-tune and Evaluate π_list with Step-wise GRPO (Final Model)

1. Continue from SFT checkpoint

2. Evaluate Final Model

About

Releases

Packages

Contributors

Languages

Latest commit

History

Folders and files

Repository files navigation

Fine-grained List-wise Alignment for Generative Medication Recommendation

Overview

Data Preparation

1. MIMIC-III Dataset

2. DrugBank Data

Preprocess the Data

Step 1: Filter & process raw MIMIC-III data

Step 2: Process unstructured clinical notes

Step 3: Generate external embeddings

Running the Code

Step 1: Drug-level Classifier (π_cls)

1. Generate training prompts

2. Train π_cls

3. Evaluate on test set

Step 2: List-wise Policy (π_list)

1. Generate π_cls predictions on train set

2. Construct training data for π_list

3. Train π_list via SFT

4. (Optional) Evaluate π_list (SFT version)

Step 3: Fine-tune and Evaluate π_list with Step-wise GRPO (Final Model)

1. Continue from SFT checkpoint

2. Evaluate Final Model

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages