Autoregressive Models Rival Diffusion Models at ANY-ORDER Generation

TL;DR

A3 introduces Any-order Any-subset Autoregressive Modeling, a framework that generalizes autoregressive factorization to arbitrary token groups and orders, bridging the gap between autoregressive depth and diffusion-style flexibility. It employs a two-stream attention architecture and a progressive training schedule that starts from standard left-to-right AR and gradually permits multi-token groups and arbitrary orderings, enabling groupwise and dynamic inference strategies. Empirically, A3 outperforms diffusion-language-model baselines on question answering, commonsense reasoning, and story infilling, while requiring substantially less pretraining data and showing favorable scaling with model size. This approach offers a practical, flexible alternative to diffusion models with competitive generation quality and diverse decoding strategies, suggesting a new direction for scalable, efficient language modeling.

Abstract

Diffusion language models enable any-order generation and bidirectional conditioning, offering appealing flexibility for tasks such as infilling, rewriting, and self-correction. However, their formulation-predicting one part of a sequence from another within a single-step dependency-limits modeling depth and often yields lower sample quality and stability than autoregressive (AR) models. To address this, we revisit autoregressive modeling as a foundation and reformulate diffusion-style training into a structured multi-group prediction process. We propose Any-order Any-subset Autoregressive modeling (A3), a generalized framework that extends the standard AR factorization to arbitrary token groups and generation orders. A3 preserves the probabilistic rigor and multi-layer dependency modeling of AR while inheriting diffusion models' flexibility for parallel and bidirectional generation. We implement A3 through a two-stream attention architecture and a progressive adaptation strategy that transitions pretrained AR models toward any-order prediction. Experiments on question answering, commonsense reasoning, and story infilling demonstrate that A3 outperforms diffusion-based models while maintaining flexible decoding. This work offers a unified approach for a flexible, efficient, and novel language modeling paradigm.

Figures (6)

• Figure 1: Architecture of the A3 model. Blue entries in the attention mask denote $0$, and white entries denote $-\infty$. The model employs a two-stream attention module with distinct causal masks. The content stream encodes contextual information and attends to tokens within its own group as well as all preceding groups. The query stream encodes positional conditions and attends only to tokens in preceding groups. The final cross-entropy loss is computed between the input context and the query stream’s output. For illustration, we provide an example grouping with $G_1=\{1,2,3\}, G_2=\{5,6\}, G_3=\{4\}$, showing how the forward process and causal masks are applied.
• Figure 2: Causal masks for content stream and query stream in different stages. Blue for $0$ and white for $-\infty$. Stage 1: AR initialization to reproduce AR factorization. Stage 2: Group expansion by allowing groups of size greater than one. Stage 3: Order permutation with introducing any-order prediction.
• Figure 3: Unconditional generation log(perplexity) and speed using A3-8B. The perplexity is measured by Llama-3.1-8B and we compare several decoding strategies. Dynamic resampling will cost more time but have lower perplexity.
• Figure 4: Performance of A3 with different training data on TriviaQA and perplexity measured by Llama-3.1-8B.
• Figure 5: Model loss of A3 across context lengths, which is stably small.