On the Trainability of Masked Diffusion Language Models via Blockwise Locality

Yuxiang Wang, Yu Xiang, Baojian Zhou, Qifang Zhao, Keyue Jiang +2 more
4/27/2026
cs.LGcs.AI

No confident code implementation found for this paper.

Abstract

Masked diffusion language models (MDMs) have recently emerged as a promising alternative to standard autoregressive large language models (AR-LLMs), yet their optimization can be substantially less stable. We study blockwise MDMs and compare them with AR-LLMs on three controlled tasks that stress different aspects of structured generation: in-context linear regression, graph path-finding, and Sudoku solving. We find that standard random-masking MDMs fail to reliably learn linear regression, exhibit high variance training dynamics on graph path-finding, while outperforming AR-LLMs on Sudoku. To mitigate these instabilities, we propose two locality aware blockwise models, namely Jigsaw and Scatter, that inject left-to-right inductive bias by enforcing autoregressive locality within blocks while preserving iterative refinement at the block level. Empirically, Jigsaw matches AR-LLM stability on linear regression and remains strong on Sudoku, while Scatter retains diffusion's planning advantage on path-finding. Our results indicate that standard random-masking MDMs, even with blockwise variants, may be a suboptimal instantiation of diffusion LMs for ordered generation, motivating models beyond random masking.

View on arXivView PDF

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. Any repos shown below are weak matches — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Similar Papers

Cite this paper

@article{wang2026trainability,
  title  = {On the Trainability of Masked Diffusion Language Models via Blockwise Locality},
  author = {Yuxiang Wang and Yu Xiang and Baojian Zhou and Qifang Zhao and Keyue Jiang and Yanghua Xiao and Xiaoxiao Xu},
  year   = {2026},
  eprint = {2604.24832},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2604.24832v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2604.24832.svg)](https://codeofpaper.com/paper/2604.24832)

Discussion