Learnability-Informed Fine-Tuning for Diffusion Language Models

Code for supervised fine tuning and evaluating diffusion LLMs with Learnability-Informed Fine-Tuning (LIFT).

📄 Paper  |  🤗 Hugging Face

---

📰 Updates

• April 2026: Accepted to ICML 2026! 🔥

🗂️ Repository Layout

scripts/   launch scripts
SFT/       fine-tuning and LoRA merge
eval/      evaluation, generation, and scoring
dataset/   local datasets (countdown/sudoku/AIME JSONs)

⚙️ Environment Setup

conda env create -f lift.yml
conda activate lift

🧪 SFT

Run SFT with the root launcher:

bash scripts/sft/run_sft.sh

Merge LoRA adapters for standalone evaluation checkpoints:

python SFT/merge_lora.py \
  --base_model GSAI-ML/LLaDA-8B-Instruct \
  --adapter_path SFT/sft_output/<run>/checkpoint-<step> \
  --output_path SFT/merged_models/<run>/checkpoint-<step> \
  --architecture llada

Use --architecture llada for checkpoints that will be evaluated with this repo's LLaDA evaluation code. The script reads base_model_name_or_path from the adapter config when available; --base_model is used as the fallback.

📊 Evaluation

Run evaluation with:

bash scripts/eval/run_eval.sh

The eval runner prints accuracy when generation finishes and writes the score into each generation JSON.

Supported --dataset keys in eval/eval.py:

gsm8k, math500, countdown, sudoku, aime24, aime25, humaneval, mbpp

📚 Module Docs

• SFT/README.md for training methods, datasets, and SFT scripts
• eval/README.md for evaluation workflow and task details