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 (Jinya Sakurai, Xun Xu, Jaehong Yoon). Code: 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	llada/ontrack.yaml	9,169
dream/	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:

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

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