|
|
--- |
|
|
base_model: |
|
|
- GSAI-ML/LLaDA-8B-Instruct |
|
|
datasets: |
|
|
- Jiayi-Pan/Countdown-Tasks-3to4 |
|
|
license: mit |
|
|
pipeline_tag: text-generation |
|
|
library_name: transformers |
|
|
--- |
|
|
|
|
|
# ESPO-Countdown-LLaDA-8B-Instruct-LoRA |
|
|
|
|
|
This repository contains a LoRA adapter for the `GSAI-ML/LLaDA-8B-Instruct` model, fine-tuned on the Countdown task using the **ELBO-based Sequence-level Policy Optimization (ESPO)** framework, as described in the paper [Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective](https://huggingface.co/papers/2512.03759). |
|
|
|
|
|
**ESPO** is a principled reinforcement learning framework designed for diffusion large language models (dLLMs). It addresses the fundamental challenges of adapting RL methods to dLLMs by treating entire sequence generation as a single action and leveraging the ELBO (Evidence Lower Bound) as a tractable sequence-level likelihood proxy. This approach resolves the mismatch between RL and the non-autoregressive nature of dLLMs, leading to significant performance improvements on mathematical reasoning, coding, and planning tasks. |
|
|
|
|
|
- 📚 Paper: [Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective](https://huggingface.co/papers/2512.03759) |
|
|
- 🌐 Project Page: [ESPO Demo](https://jingyangou.github.io/ESPO-Demo/) |
|
|
- 💻 Code: [ML-GSAI/ESPO](https://github.com/ML-GSAI/ESPO) |
|
|
|
|
|
## Sample Usage |
|
|
|
|
|
You can load and use this ESPO-fine-tuned LoRA adapter on top of the base LLaDA-8B-Instruct model. Below is a quick start example demonstrating how to perform inference. |
|
|
|
|
|
```python |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
from peft import PeftModel |
|
|
from eval.generate_utils import generate # Note: `eval.generate_utils` is a custom module from the GitHub repository and needs to be accessible. |
|
|
|
|
|
base_model_path = 'GSAI-ML/LLaDA-8B-Instruct' |
|
|
peft_model_path = 'GSAI-ML/ESPO-Math' # Replace with 'GSAI-ML/ESPO-Countdown' for this specific model |
|
|
tokenizer = AutoTokenizer.from_pretrained(base_model_path) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
base_model_path, trust_remote_code=True,torch_dtype="bfloat16", device_map="cuda") |
|
|
peft_model = PeftModel.from_pretrained(model, peft_model_path, device_map="cuda") |
|
|
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?\ |
|
|
\ |
|
|
Write your answer in the form $(x, y)$ where $x$ and $y$ are real numbers." |
|
|
messages = [{"role": "user", "content": prompt}] |
|
|
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda") |
|
|
output_ids = generate(peft_model, input_ids,tokenizer, steps=128, gen_length=256, temperature=0.9,remasking="low_confidence",) |
|
|
output_text = tokenizer.batch_decode(output_ids[:, input_ids.shape[1]:], skip_special_tokens=True)[0] |
|
|
print(output_text) |
|
|
``` |
|
|
|
|
|
## Citation |
|
|
|
|
|
If you find ESPO useful in your research, please consider citing our paper: |
|
|
|
|
|
```bibtex |
|
|
@article{ou2025principledrldiffusionllms, |
|
|
title={Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective}, |
|
|
author={Jingyang Ou and Jiaqi Han and Minkai Xu and Shaoxuan Xu and Jianwen Xie and Stefano Ermon and Yi Wu and Chongxuan Li}, |
|
|
journal={arXiv preprint arXiv:2512.03759}, |
|
|
year={2025}, |
|
|
} |
|
|
``` |
