Auto-Regressive Masked Diffusion Models

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.

• The paper demonstrates that ARMD bridges autoregressive and masked diffusion models via block-wise causal reformulation to enable efficient likelihood evaluation and fast inference.
• It leverages a strictly causal, permutation-equivariant attention architecture to ensure robust training and coherent, parallel generation.
• Experimental results indicate that ARMD achieves lower perplexity and faster decoding compared to state-of-the-art diffusion and autoregressive baselines.

Auto-Regressive Masked Diffusion Models: Bridging Causal Modeling and Parallel Generation

---

Introduction and Motivation

The Auto-Regressive Masked Diffusion (ARMD) model (2601.16971) addresses core inefficiencies in language modeling for discrete domains, reconciling the strengths of autoregressive models (ARMs) and masked diffusion models (MDMs). ARMs, characterized by their sequential token generation and efficient parallel training, have dominated language modeling benchmarks but suffer from slow sampling and limited support for bidirectional or non-canonical generation orders. MDMs, in contrast, enable parallel sampling but have lagged behind ARMs in sample quality and training efficiency, particularly for text data. ARMD reformulates MDMs as block-wise causal models, facilitating parallel likelihood evaluation and enabling progressive permutation-based training.

---

Block-wise Causal Reformulation of MDMs

The central technical contribution reinterprets the masked diffusion process as a block-wise autoregressive model. In standard MDMs, a discrete sequence is corrupted by iterative random masking, and the reconstruction is carried out via sequential denoising steps. ARMD introduces a permutation $\pi$ that reorders tokens by their masking timesteps, inducing a partition into blocks $[\mathcal{X}^{(1)}, \ldots, \mathcal{X}^{(T)}]$ (Figure 1). In this reordering, each block is autoregressively generated conditioned only on prior blocks, effectively transforming the original masked diffusion objective into a single forward-pass causal likelihood across all timesteps. This reformulation enables both efficient parallel training and compatibility with key-value caching for fast decoding.

Figure 1: Visualization of the masked diffusion process, block-wise permutation ordering, strictly causal attention mask, and deep strictly causal model composition.

---

Strictly Causal Permutation-Equivariant Attention Architectures

To instantiate the block-wise causality, ARMD introduces a strictly causal, permutation-equivariant attention architecture. Whereas typical causal attention allows intra-block interaction, strictly causal attention masks eliminate any access to the current block, enforcing that each token can only attend to preceding block representations. This masking enables the model to flexibly support arbitrary block orderings and underlies the progressive permutation training regime.

Query representations for strictly causal attention are constructed via prefix aggregation over past blocks, utilizing parameter-efficient schemes (e.g., linear dot-product-based pooling of positional embeddings) to maintain expressivity without prohibitive computational overhead. The model is designed as a two-stream stack (Figure 2), with each layer updating both a standard causal and a strictly causal hidden state, sharing key-value projections and differing only in masking and query input composition.

Figure 2: Schematic of the deep strictly causal architecture with two-stream attention layers followed by causal layers, supporting both block-wise causality and permutation equivariance.

---

Strided Parallel Generation: Efficient Inference

ARMD leverages its permutation-equivariant structure for strided block-parallel (SBP) generation, achieving fast and coherent inference. The sequence is partitioned into $S$ parallel streams, and a permutation interleaves their heads, allowing simultaneous token generation at large physical separation. This respects conditional independence assumptions needed for diffusion-based parallel sampling, and the block-wise masking ensures global coherence (Figure 3).

Figure 3: Illustration of strided parallel generation: initial stream heads generated sequentially, followed by parallel generation of blocks at equal relative positions, with strictly causal masking ensuring coherence.

---

Experimental Results and Empirical Analysis

ARMD is evaluated against the strongest MDM baselines (D3PM, PLAID, SEDD, RADD), and autoregressive transformers (GPT-2), across LAMBADA, WikiText2/103, PTB, and One Billion Words. Notable findings include:

• On small and medium model scales, ARMD matches or exceeds GPT-2 performance on most datasets, achieving lower perplexity than all diffusion baselines while requiring significantly fewer training steps.
• ARMD achieves best diffusion-model perplexity on LM1B (22.36 with 1M steps), closing the gap to the best AR transformer result (22.83).
• Progressive permutation training improves robustness to non-canonical orderings without degrading left-to-right performance.
• Strided parallel generation maintains high output entropy and coherent language, achieving inference speedups proportional to block parallelization factor $S$; increasing SBP fine-tuning steps further bridges the remaining gap to sequential AR sampling.

---

Theoretical and Practical Implications

ARMD's causal-permutation reformulation and two-stream attention provide crucial advances:

• Parallel Training Efficiency: Unlike MDMs and OA-ARMs, all denoising conditionals are evaluated in a single forward pass, reducing training wall-clock time and total FLOPs.
• Flexible Generation Order: The permutation-equivariant architecture generalizes ARMs and supports arbitrary token unmasking/reordering, necessary for non-sequential or conditional generative scenarios.
• Inference Scalability: Strided parallel generation, leveraging strictly causal masking, enables substantial reductions in inference latency without severe degradation in sample quality.
• Adaptability: Because ARMD is architecturally close to ARMs, pre-trained AR LLMs could be fine-tuned for diffusion-style sampling with simple modifications, offering a pathway to scale DLMs further.

---

Speculation and Future Directions

ARMD establishes masked diffusion as a viable, scalable alternative to ARMs for text modeling. Its design principles suggest promising future avenues:

• Diffusion Finetuning for Large LMs: Pretrained LLMs could be adapted for block-wise diffusion inference, potentially inheriting their sample quality and controllability.
• Hybrid Planning and Conditional Generation: The ability to modulate sampling order and block structure may enhance planning, control, and selective context synthesis in AI systems.
• Scaling Laws and Memory-Efficient Training: ARMD's efficiency may enable the scaling of DLMs to trillion-parameter regimes, especially leveraging architectural sharding and model parallelism.

---

Conclusion

Auto-Regressive Masked Diffusion Models (ARMD) bridge the long-standing gap between sequential efficiency and parallelism in discrete language modeling. By reframing masked diffusion as block-wise autoregression and introducing strictly causal, permutation-equivariant attention, ARMD achieves state-of-the-art empirical performance with dramatic reductions in training cost vs. prior diffusion models, supports fast, coherent parallel generation, and offers strong architectural compatibility for further scaling and adaptation. This synthesis paves the way for next-generation diffusion LMs with both performance and flexibility competitive with autoregressive paradigms.