Improve model card: Add pipeline tag, library, links, and sample usage

#1
by nielsr HF Staff - opened
Files changed (1) hide show
  1. README.md +65 -4
README.md CHANGED
@@ -1,9 +1,70 @@
1
  ---
2
- license: mit
3
- datasets:
4
- - openai/gsm8k
5
  base_model:
6
  - GSAI-ML/LLaDA-8B-Instruct
 
 
 
 
 
7
  ---
8
 
9
- Post-Training Lora models on gsm8k task based on LLaDA-8B-Instruct for the paper Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
 
 
 
2
  base_model:
3
  - GSAI-ML/LLaDA-8B-Instruct
4
+ datasets:
5
+ - openai/gsm8k
6
+ license: mit
7
+ pipeline_tag: text-generation
8
+ library_name: transformers
9
  ---
10
 
11
+ # ESPO: Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective
12
+
13
+ This repository provides the official implementation and checkpoints for the paper:
14
+ [**Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective**](https://huggingface.co/papers/2512.03759)
15
+
16
+ **Project Page:** [https://jingyangou.github.io/ESPO-Demo/](https://jingyangou.github.io/ESPO-Demo/)
17
+ **Code Repository:** [https://github.com/ML-GSAI/ESPO](https://github.com/ML-GSAI/ESPO)
18
+
19
+ We introduce **ESPO (ELBO-based Sequence-level Policy Optimization)**, a principled reinforcement learning framework that adapts to diffusion large language models (dLLMs). This framework addresses the fundamental challenges of applying traditional RL methods to dLLMs by viewing the **entire sequence generation as a single action** and leveraging the **ELBO** as a tractable proxy for sequence-level likelihood.
20
+
21
+ ESPO introduces:
22
+ - **Sequence-level optimization** for diffusion LLMs via the ELBO objective.
23
+ - **Per-token normalized ratio estimation** and **robust KL regularization** for stable large-scale training.
24
+ - **Consistent gains** across math, coding, and planning benchmarks, significantly outperforming token-level baselines.
25
+
26
+ ## Model Details
27
+ This specific model is an ESPO-fine-tuned LoRA adapter built on `GSAI-ML/LLaDA-8B-Instruct` for the GSM8K task. ESPO-GSM8K, ESPO-Math, ESPO-Countdown, and ESPO-Sudoku are provided as LoRA adapters, which can be loaded on top of the base LLaDA-8B-Instruct model for lightweight and efficient fine-tuning.
28
+
29
+ ## Sample Usage (Quick Start)
30
+ You can load and use the ESPO-fine-tuned LoRA adapters with the `transformers` and `peft` libraries. Here's an example for inference using an ESPO adapter (e.g., ESPO-Math):
31
+
32
+ ```python
33
+ from transformers import AutoModelForCausalLM, AutoTokenizer
34
+ from peft import PeftModel
35
+ from eval.generate_utils import generate # Note: This `generate` function is a custom utility from the ESPO codebase.
36
+ import torch
37
+
38
+ base_model_path = 'GSAI-ML/LLaDA-8B-Instruct'
39
+ peft_model_path = 'GSAI-ML/ESPO-Math' # Use 'GSAI-ML/ESPO-GSM8K' for this model, or other ESPO adapters.
40
+
41
+ tokenizer = AutoTokenizer.from_pretrained(base_model_path)
42
+ model = AutoModelForCausalLM.from_pretrained(
43
+ base_model_path, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda")
44
+
45
+ peft_model = PeftModel.from_pretrained(model, peft_model_path, device_map="cuda")
46
+
47
+ prompt = "The point $(0,0)$ is reflected over the vertical line $x=1$. When its image is then reflected over the line $y=2$, what is the resulting point?\
48
+ \
49
+ Write your answer in the form $(x, y)$ where $x$ and $y$ are real numbers."
50
+ messages = [{"role": "user", "content": prompt}]
51
+
52
+ input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
53
+
54
+ output_ids = generate(peft_model, input_ids, tokenizer, steps=128, gen_length=256, temperature=0.9, remasking="low_confidence")
55
+
56
+ output_text = tokenizer.batch_decode(output_ids[:, input_ids.shape[1]:], skip_special_tokens=True)[0]
57
+ print(output_text)
58
+ ```
59
+
60
+ ## Citation
61
+ If you find ESPO useful in your research, please consider citing our paper:
62
+
63
+ ```bibtex
64
+ @article{ou2025principledrldiffusionllms,
65
+ title={Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective},
66
+ author={Jingyang Ou and Jiaqi Han and Minkai Xu and Shaoxuan Xu and Jianwen Xie and Stefano Ermon and Yi Wu and Chongxuan Li},
67
+ journal={arXiv preprint arXiv:2512.03759},
68
+ year={2025},
69
+ }
70
+ ```