| --- |
| license: apache-2.0 |
| library_name: pytorch |
| tags: |
| - diffusion-language-model |
| - masked-diffusion |
| - unmasking-policy |
| - ontrack |
| - llada |
| - dream |
| base_model: |
| - GSAI-ML/LLaDA-8B-Instruct |
| - Dream-org/Dream-v0-Instruct-7B |
| base_model_relation: adapter |
| --- |
| |
| # OnTrack: learned unmasking policies for LLaDA-8B-Instruct and Dream-v0-Instruct-7B |
|
|
| **OnTrack** policy weights for the paper |
| [What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models](https://arxiv.org/abs/2610.31576) |
| (Jinya Sakurai, Xun Xu, Jaehong Yoon). |
| Code: [github.com/jinya1013/ontrack](https://github.com/jinya1013/ontrack). |
|
|
| Each policy is a small network (about 0.8M parameters) that, at every denoising |
| step of blockwise decoding, decides which masked positions of the active block to |
| commit. OnTrack reads the frozen denoiser's top-1 confidence and its |
| final hidden state. It was trained for one epoch of offline teacher forcing on |
| student-rejoin teacher trajectories built from deterministic Fast-dLLM decoding |
| (threshold 0.9, generation length 256, block length 32) on the GSM8K and MATH |
| training prompts, keeping only trajectories with a correct final answer. |
|
|
| | Folder | Base model | Training config | Teacher samples | |
| |---|---|---|---| |
| | `llada/` | [GSAI-ML/LLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct) | `llada/ontrack.yaml` | 9,169 | |
| | `dream/` | [Dream-org/Dream-v0-Instruct-7B](https://huggingface.co/Dream-org/Dream-v0-Instruct-7B) (rev `05334cb9`) | `dream/ontrack.yaml` | 9,776 | |
|
|
| The base denoisers are not included and stay frozen. |
|
|
| ## Usage |
|
|
| With the [code repository](https://github.com/jinya1013/ontrack): |
|
|
| ```bash |
| hf download jinyas/ontrack --local-dir policies |
| |
| python scripts/sweep.py --output results/ontrack_llada \ |
| --policy-config configs/llada/ontrack.yaml --policy-path policies/llada |
| python scripts/sweep.py --output results/ontrack_dream \ |
| --policy-config configs/dream/ontrack.yaml --policy-path policies/dream |
| ``` |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{sakurai2026unmasking, |
| title = {What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models}, |
| author = {Sakurai, Jinya and Xu, Xun and Yoon, Jaehong}, |
| journal = {arXiv preprint arXiv:2610.31576}, |
| year = {2026}, |
| } |
| ``` |
|
|
| ## License |
|
|
| Apache License 2.0. Use of the base models is subject to their own licenses. |
|
|
