Auto-Regressive Masked Diffusion Models

TL;DR

Auto-Regressive Masked Diffusion (ARMD) closes the gap between autoregressive language modeling and diffusion-based discrete models by reframing masked diffusion as a block-wise causal problem and introducing a strictly causal, permutation-equivariant two-stream transformer. This architecture allows fully parallel evaluation of conditional probabilities during training, supports progressive permutation of token orders, and enables strided block-parallel generation for accelerated inference while maintaining global coherence. Empirical results show ARMD achieves state-of-the-art zero-shot perplexities on standard benchmarks with far fewer training iterations than diffusion baselines, and it sets a new benchmark for parallel text generation. The work suggests a practical, flexible path to scalable diffusion-based language modeling that can leverage pre-trained LLMs via fine-tuning and KV caching for efficient decoding.

Abstract

Masked diffusion models (MDMs) have emerged as a promising approach for language modeling, yet they face a performance gap compared to autoregressive models (ARMs) and require more training iterations. In this work, we present the Auto-Regressive Masked Diffusion (ARMD) model, an architecture designed to close this gap by unifying the training efficiency of autoregressive models with the parallel generation capabilities of diffusion-based models. Our key insight is to reframe the masked diffusion process as a block-wise causal model. This perspective allows us to design a strictly causal, permutation-equivariant architecture that computes all conditional probabilities across multiple denoising steps in a single, parallel forward pass. The resulting architecture supports efficient, autoregressive-style decoding and a progressive permutation training scheme, allowing the model to learn both canonical left-to-right and random token orderings. Leveraging this flexibility, we introduce a novel strided parallel generation strategy that accelerates inference by generating tokens in parallel streams while maintaining global coherence. Empirical results demonstrate that ARMD achieves state-of-the-art performance on standard language modeling benchmarks, outperforming established diffusion baselines while requiring significantly fewer training steps. Furthermore, it establishes a new benchmark for parallel text generation, effectively bridging the performance gap between parallel and sequential decoding.

Figures (6)

• Figure 1: (Top Left) An instance of the masked diffusion process on ${\boldsymbol{z}}_{1:N}$ for $T=4$ steps. (Top Right) The causal patterns derived by permuting the sequence ${\boldsymbol{x}}_{1:N} = \pi({\boldsymbol{z}}_{1:N})$. The input sequence is partitioned into blocks according to the reverse masking order: $\mathcal{X}(1)=\{3, 4\}$, $\mathcal{X}(2)=\{1, 6\}$, $\mathcal{X}(3)=\{2, 7\}$, and $\mathcal{X}(4)=\{4\}$. (Bottom Left) The resulting strictly causal attention mask. (Bottom Right) A model instance composed of a single strictly causal layer followed by two causal layers, forming a deep strictly causal architecture.
• Figure 2: Visualization of strided parallel generation. Steps 1--2 represent the sequential generation of the stream heads. Steps 3--5 illustrate the parallel generation of subsequent tokens, where tokens of the same color are generated simultaneously. \ref{['fig:masking_strided_gen']} depicts the corresponding diffusion process and attention mask.
• Figure 3: Schematic of the deep strictly causal architecture. The model is composed of $L^{2s}$ two-stream attention layers (causal and strictly causal streams), followed by $L - L^{2s}$ causal layers. The strictly causal stream (${\bm{G}}^l$) captures rich contextual and positional representations using only past block information, while the causal stream (${\bm{X}}^l$) serves as a shared source of key and value features for the attention layers in both streams. The two streams differ in their query input and masking operations. The final output is taken from the top layer, enabling the model to parametrize the conditional distribution $p_\theta({\boldsymbol{x}}_n \mid \{{\boldsymbol{x}}_i \mid \mathcal{B}(i) < \mathcal{B}(n)\})$ while preserving permutation equivariance with respect to the condition set. Layer-Norm, drop-out, and output projections of the attention layers are dropped for simplicity.
• Figure 4: Illustration of the strided parallel generation, its diffusion process and masking. (Top Left): An instance of masked diffusion process. (Top Right): The causal patterns by permuting the sequence ${\boldsymbol{x}}_{1:N} = \pi({\boldsymbol{z}}_{1:N})$. (Bottom left): The resulting strictly causal attention mask. (Bottom right): The strided parallel generation where steps 1-2 represent the sequential generation of the stream heads, and steps 3-5 show the parallel generation of subsequent tokens. Here, same-colored tokens are generated simultaneously.
• Figure 5: Perplexity vs. training steps for small-size models.

Theorems & Definitions (2)

• Remark 1
• proof