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