---
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.