On the Trainability of Masked Diffusion Language Models via Blockwise Locality

TL;DR

This paper investigates blockwise masked diffusion language models, identifies their stability issues compared to autoregressive models, and proposes locality-aware variants that improve training stability and task performance.

Contribution

It introduces Jigsaw and Scatter models that incorporate locality biases to enhance the stability and effectiveness of blockwise MDMs.

Findings

Jigsaw matches AR-LLM stability on linear regression and Sudoku.

Scatter retains diffusion's planning advantage on path-finding.

Standard random-masking MDMs are suboptimal for ordered generation.

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.