# On Trajectory-Aware Training for Masked Diffusion Language Models

arXiv: https://arxiv.org/abs/2609.37974 · HTML: https://mlsanity.com/paper/2609.37974/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 22 of its 65 references that are in the map.
- First posted on arXiv 29 Sep 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 22 of 22)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [Mercury: Ultra-Fast Language Models Based on Diffusion](https://mlsanity.com/paper/2506.17298/) (2025)
- [Continuous diffusion for categorical data](https://mlsanity.com/paper/2211.15089/) (2022)
- [LLaDA2.0: Scaling Up Diffusion Language Models to 100B](https://mlsanity.com/paper/2512.15745/) (2025)
- [Olmo 3](https://mlsanity.com/paper/2512.13961/) (2025)
- [Path Planning for Masked Diffusion Model Sampling](https://mlsanity.com/paper/2502.03540/) (2025)
- [TinyGSM: achieving >80% on GSM8k with small language models](https://mlsanity.com/paper/2312.09241/) (2023)
- [MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models](https://mlsanity.com/paper/2508.13148/) (2025)
- [Improved Large Language Diffusion Models](https://mlsanity.com/paper/2606.25331/) (2026)
- [Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation](https://mlsanity.com/paper/2606.06712/) (2026)
- [DiffusionGemma Technical Report](https://mlsanity.com/paper/2608.00146/) (2026)
- [Trace-Based On-Policy Distillation for Masked Diffusion Language Models](https://mlsanity.com/paper/2607.16872/) (2026)
- [The Design Space of Tri-Modal Masked Diffusion Models](https://mlsanity.com/paper/2602.21472/) (2026)
- [Limits of Confidence in Diffusion](https://mlsanity.com/paper/2609.20581/) (2026)
- [Learned Relay Representations for Forward-Thinking Discrete Diffusion Models](https://mlsanity.com/paper/2605.22967/) (2026)
- [Temporal Self-Distillation: Faster Inference in Discrete Diffusion Language Models](https://mlsanity.com/paper/2609.15177/) (2026)

## Abstract (arXiv)

Masked diffusion models (MDMs) generate text by unmasking several tokens per step, but they are trained and sampled under different conditions. The model is trained on randomly masked sequences, whereas inference follows a trajectory shaped by the model's own predictions. Additionally, each step has no access to what the previous one computed. Recent methods narrow these limitations from separate angles, leaving open how these choices interact. We introduce PUMBA, a unified framework for trajectory-aware training that trains the denoiser on consecutive steps of policy-induced trajectories, passes information between steps, and optimizes them jointly by backpropagation through time. A controlled study of this design space shows that i) exact train--inference alignment fails due to local overfitting, whereas a looser alignment still brings training masks closer to those seen at inference; ii) passing continuous information outperforms discrete gradient estimators through the commitment at each step; and iii) performance improves as backpropagation through time spans more steps, which we support theoretically. Combined, these components match the best checkpoint of a same-size autoregressive model. Building on these findings, we scale PUMBA to supervised fine-tuning of LLaDA-8B, where it improves the trade-off between performance and number of function evaluations (NFEs) in both full-canvas and block diffusion generation. At matched performance, it needs up to 22% fewer NFEs than standard fine-tuning with twice the budget in full-canvas generation, and up to 26% fewer than standard fine-tuning for the same number of steps in block diffusion.