Relaxing Positional Alignment in Masked Diffusion Language Models

TL;DR

This paper proposes a flexible positional supervision method for masked diffusion language models, improving their robustness and performance in open-ended text generation by reducing sensitivity to token misalignments.

Contribution

It introduces a <slack> token and a TCTC-based training strategy to relax positional constraints in MDLMs, enhancing generation quality and robustness.

Findings

Outperforms original MDLM on five benchmarks

Increases robustness to token misalignments

Improves open-ended text generation quality

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.