Track new results across pure and applied mathematics.
Search papers, labs, and topics across Lattice.
To resolve the train-inference distribution shift and stateless step dynamics in masked diffusion models, the authors develop PUMBA, a trajectory-aware framework that optimizes consecutive unmasking steps via continuous state passing and multi-step backpropagation through time. Systematic exploration reveals that loose policy alignment avoids the local overfitting of exact matching, while multi-step continuous memory allows the masked denoiser to match the performance of same-sized autoregressive checkpoints. Scaled to LLaDA-8B fine-tuning, the approach shifts the quality-inference Pareto frontier, slashing required function evaluations by up to 26% in block diffusion and 22% in full-canvas generation.
Masked diffusion language models can match autoregressive baselines while cutting sampling steps by over 20%—provided you train across multi-step continuous trajectories rather than isolated random masks.
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.