Relaxing Positional Alignment in Masked Diffusion Language Models

TL;DR

This work identifies strict position-level supervision as a key bottleneck for open-ended generation in masked diffusion language models, showing that early positional misalignment can cascade under irreversible decoding. It introduces an alignment-flexible training approach that inserts a <slack> token via a CT C-based objective to absorb positional uncertainty without changing the task interface. Across five open-ended generation benchmarks, the CTC-trained model consistently improves over CE-based baselines and demonstrates enhanced robustness to positional shifts. The approach maintains decoding efficiency while enabling flexible alignment, offering a practical path to closing the gap between MDLMs and autoregressive models in open-ended text generation.

Abstract

Masked diffusion language models (MDLMs) have emerged as a promising alternative to dominant autoregressive approaches. Although they achieve competitive performance on several tasks, a substantial gap remains in open-ended text generation. We hypothesize that one cause of this gap is that strict positional prediction makes MDLM decoding highly sensitive to token misalignment, and we show through controlled interventions that a one-position shift can severely disrupt semantics. This observation suggests that enforcing strict positional supervision during training is misaligned with the irreversible denoising dynamics of MDLM decoding. Motivated by this mismatch, we adopt an alignment-flexible supervision strategy during fine-tuning. Specifically, we introduce a special token <slack> via the connectionist temporal classification objective. We apply this approach to the widely used MDLM model and conduct experiments on five open-ended text generation benchmarks. Our method consistently outperforms the original model and improves robustness to positional shifts, indicating that relaxing strict positional supervision is an important factor in improving generation quality in MDLMs.

Figures (14)

• Figure 1: Under standard MDLM decoding, tokens are generated at fixed positions and become immutable once unmasked. A small position error early in decoding can therefore propagate through subsequent steps, leading to cascading semantic degradation (left). Allowing limited alignment flexibility introduces non-semantic buffer positions that absorb local misalignment, preventing error propagation while preserving the irreversible decoding process (right).
• Figure 2: Win rate (%) on Arena-hard as a function of the number of shift boundaries $K$ for LLaDA-8B-Instruct and LLaDA-1.5. Error bands indicate 95% confidence intervals. Win rates are computed against the $K=0$ (no-intervention) outputs of the same model as the reference, yielding a 50% win-rate baseline (dashed line). The Pearson correlation coefficient $r$ is between $K$ and the win rate; larger $r$ indicates greater sensitivity to the intervention.
• Figure 3: Illustration of < slack> token placement in the CTC-trained model. < slack> is observed both within tokenized words ((e.g., D< slack>osa, mas< slack>ala)) and between words, where it functions as mid-text padding. This suggests < slack> can act as a flexible buffer at different locations in the sequence.
• Figure 4: Win rate (%) on Arena-hard as a function of the number of shift boundaries $K$ for the CTC-trained model and the CE-only baseline. Error bands indicate 95% confidence intervals. Win rates are computed against the $K=0$ (no-intervention) outputs of the same model as the reference, yielding a 50% win-rate baseline (dashed line). The Pearson correlation coefficient $r$ is between $K$ and the win rate; $r$ indicates weaker sensitivity to the intervention.
• Figure 5: Tokens per second for the CTC-trained model with and without the merge operation during decoding.