Autoregressive Models Rival Diffusion Models at ANY-ORDER Generation
Authors: Tianqi Du
Lizhe Fang
Weijie Yang
Chenheng Zhang
Zeming Wei
Yifei Wang
Yisen Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Large language model
permutation language model
Primary area: foundation or frontier models, including LLMs

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. Code is at https://github.com/PKU-ML/Any-order-Any-subset-AR.

Published as a conference paper at ICLR 2026
AUTOREGRESSIVE MODELS RIVAL DIFFUSION MODELS AT ANY-ORDER GENERATION
Tianqi Du
Lizhe Fang'
Weijie Yang²
Chenheng Zhang'
Zeming Wei²
Yifei Wang
Yisen Wang
State Key Lab of General Artificial Intelligence,
School of Intelligence Science and Technology, Peking University
2 School of Mathematical Sciences, Peking University
3 Amazon AGI SF Lab* Institute for Artificial Intelligence, Peking University
ABSTRACT
Diffusion language models enable any-order generation and bidirectional condi- tioning, offering appealing flexibility for tasks such as infilling, rewriting, and self-correction. However, their formulation-prediciing one part of a sequence from another within a single-step dependeeny--limis modeling depth and often yields lower sample quality and stability than autoregressive (AR) models. To address this, we revisit autoregressive modeling foundation and reformulate diffusion-style training into a structured multi-group prediction process. We pro- pose Any-order Any-subset Autoregressive modeling (A3), a generalized frame- work that extends the standard AR factorization to arbitrary token groups and gen- cration 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 two-stream attention archi- tecture and a progressive adaptation strategy that transitions pretrained AR mod- els toward any-order prediction. Experiments on question answering, common- sense 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. Code is at https///ggituub..oo/P////./ /An----- Any-suset--RR,
INTRODUCTION
Diffusion language models have recently emerged as a powerful alternative to autoregressive (AR) modeling for text generation (Li et al., 2022; Gong al., 2025; Gulrajani & Hashimoto, 2023; Sahoo et al., 2024). By iteratively denoising partially masked sequences, they allow any-order generation and bidirectional context conditioning, enabling novel capabilities such as global rewriting, infilling, and self-correction. These properties make diffusion models conceptually appealing for flexible text generation beyond the rigid left-to-right paradigm.
However, diffusion modeling also introduces fundamental limitations. In each denoising step, the model predicts one part of the sequence conditioned on another, forming only a single-layer de- pendency structure. This simplification weakens the model's ability to capture deep, hierarchical dependencies across tokens. As a result, despite their flexibility, diffusion-based models often un- derperform AR models in generation quality and training stability (Kim et al., 2025; Sahoo et al., 2024; Zheng et al., 2024). These challenges highlight an inherent trade-off between the expressive power of AR factorization and the flexibility of diffusion-style generation.
Motivated by this observation, we revisit autoregressive modeling as a foundation for constructing more expressive and stable framework. AR models naturally capture rich, recursive dependen- cies through sequential conditioning but suffer from their strict left-to-right constraint. We therefore
Conresponning Author: Yisen Wang (yisen. wang @ww..epu.ccn..... I This work was completed MIT prior Yifei Wang joining Amazon.

---

Published as a conference paper at ICLR 2026
reformulate diffusion training in more structured way by extending its two-group prediction to multiple groups, so that the model can preserve AR's multi-layer dependency structure while sup- porting arbitrary generation orders. This leads to Any-order Any-subset Autoregressive modeling (A3), a unified framework that generalizes standard AR factorization to allow token groups of ar- bitrary composition and order. A3 retains the probabilistic rigor and stability of AR training while inheriting diffusion models' flexibility for parallel generation and bidirectional context use.
To realize A3 in practice, we design two-stream attention architecture that supports arbitrary group orderings, and develop a progressive training strategy that adapts pretrained AR models to any-order prediction. naturally supports diverse inference strategies. including groupwise AR sampling and dynamic resampling, offering a tunable trade-off between generation speed and quality.
Through comprehensive experiments on question answering (Joshi al., 2017), commonsense rea- soning (Zellers et al., 2019; Sakaguchi et al., 2020; Sap et al., 2019; Bisk et al., 2020), and story infilling tasks (Mostafazadeh et al., 2016), we show that A3 achieves strong performance across diverse benchmarks while enabling flexible and efficient generation. Notably, A3 outperforms state- of-the-art diffusion-basee
models despite using substantially less training data, and demonstrates promising scaling behavior with model size. These results suggest that A3 provides new direction for bridging the gap between AR and parallel generation paradigms.
Our contributions can be summarized as follows:
Conceptually, we propose A3, new sequence modeling paradigm that reformulates diffusion-style group prediction into generalized autoregressive framework, preserving dependency depth while enabling any-order, any-subser generation.
Algorithmicalll, we design a two-stream attention architecture and progressive training recipe that extend pretrained AR models to support arbitrary group orderings.
Practically, we demonstrate that A3 surpasses diffusion-based approaches, offering a flex- ible and scalable foundation for next-generation language modeling.
2 ANY-ORDER ANY-SUBSET AUTOREGRESSIVE MODELING
2.1 FORMULATION
In this section, we will demonstrate the formulation of our proposed Any-order Any-subset Autore- gressive modeling (A3). Masked diffusion language models provide a flexible approach to sequence generation by iteratively predicting masked subsets of tokens conditioned on visible context (He et al., 2023; Nie et al., 2025). Concretely, during training, given sequence TI,N, the index set 1, 1, , N partitioned into two disjoint subsets G₁ and G2. The model predicts the masked tokens fG, conditioned on the visible tokens TG,
P(ae. (G.))
II P(x, xG.). tEGa
This one-step conditional formulation enables parallel prediction and bidirectional context use, sup- porting ony-order generation through repeated resampling of subsets. However, it models only shal- low dependencies-eaae token in G2 depends directly on FG, but not recursively on other predicted tokens. As result, the global dependency structure is limited one layer of conditioning. unlike the multi-step compositional nature of autoregressive (AR) factorization.
To enrich the dependency structure while retaining the flexibility of arbitrary grouping, reformu- late this process through an autoregressive lens. Starting from the AR formulation.
P(ΛIN) II P(x. I<t)) t=1
(2)
we generalize the loken-wise sequence to a group-wise factorization. The index set (I....,NN partitioned into K disjoint groups {G1, G.....,Ga Gx). and the joint probability is expressed as
N is
K
P(IIN) = II P(rga xG
(3)
2

---

Published as a conference paper at ICLR 2026
Each group G3 may contain one more tokens, and the group order is arbitrary. This generalization retains the recursive dependency structure of AR while introducing diffusion-style flexibility. Each factor can be viewed predicting subset of tokens from previously generated subsets, but now organized into multiple layers of dependency.
Training the model randomly sampled group partitions and permutations exposes it to diverse factorization orders and improves robustness to different conditional structures. This resembles the permutation LM objective of XLNet (Yang et al., 2019) but the group level, enabling richer structural modeling. This formulation bridges diffusion's 's parallel prediction and AR's hierarchical modeling, forming the foundation of Any-order Any-subset Autoregressive modeling (A3).
2.2 DISCUSSION WITH PREVIOUS PARADIGMS
Comparison with Masked Diffusion Language Models. Masked diffusion language models (MDLMs) have recently emerged as an alternative to AR decoding, aiming to overcome the se- quential bottleneck of mee-token-at-a-uie sequences. MDLMs can update multiple tokens in parallel and exploit bidirectional context (Austin et al., 2021; Li ct al., 2022; Gong ct al., 2025; Nie et al., 2025). This flcxibility enables controllable, any-order generation and supports tasks such as infilling and global rewriting. However, diffusion approaches face two major limitations. First, inference speed is constrained by the need for many iterative refinement steps, with generation efficiency depending critically on noise schedules and step counts (Gong al., 2025). While large-scale efforts such DiffuGPT and LLaDA demon- strate that diffusion LMs can match or surpass AR models in quality, they still require careful tuning and incur nontrivial decoding costs (Gong et al., 2025; Nie et al., 2025). Second, training stability is less favorable than AR: discrete diffusion objectives require complex forward-revers: processes and additional pretraining or adaptation, making optimization more resource-intensive and sensitive to hyperparameters (He et al., 2023).
generation. By iteratively denoising partially masked
In contrast, A3 preserves the probabilistic rigor and training simplicity of AR modeling while in- troducing flexible groupwise factorization. This allows parallel prediction of token subsets without relying on multi-step denoising schedules, thereby offering both efficiency and stability.
Comparison with AR Multi-Token Prediction. Another line of work seeks to improve AR ef- ficiency by enabling the model to predict multiple future tokens per step (Glocckle et al., 2024; Kou et al., 2024; Zhang et al., 2024). Multi-token objectives and speculative decoding significantly reduce inference latency: for instance, predicting four tokens at once can yield up to threefold speedups while preserving even improving generation quality, particularly in reasoning and code generation tasks (Gloeckle et al., 2024). These methods retain the training stability of standard AR, since the additional objectives can be implemented auxiliary losses with negligible computational overhead. However, multi-token prediction remains bound to a fixed left-to-right ordering, limiting its modeling flexibility. The approach accelerates sequential decoding but does not enable infilling, bidirectional conditioning, or arbitrary ordering of token generation.
A3 generalizes beyond this paradigm by relaxing the strict AR factorization. Through arbitrary group partitions and orderings, A3 supports both sequential and parallel decoding strategies, com- bining the efficiency gains of multi-token prediction with greater structural flexibility. For more discussions with related work, refer to Appendix B.
3 IMPLEMEENTATIONS OF TRAINING AND FLEXIBLE INFERENCE
In this section, we describe the implementation of A3. We begin with the architectural design of A3, where we use a two-stream attention mechanism enable predictions in arbitrary orders. Next, we present our efficient continuous pretraining strategy, which progressively adapts the model from standard AR prediction to group-based prediction. Finally, we introduce the flexible inference strategy of A3, which leverages its general formulation to support diverse decoding modes.
3.1 ARCHITECTURE DESIGN WITH TWO-STREAM ATTENTION
Limitations of Current Architecture. The decoder-only Transformer has become the backbone of modern large language models due its simplicity and effectiveness in next-token prediction.
3

---

Published as a conference paper at ICLR 2026
The onswer
is
14
inches
long
CE Loss
(7 (1)
(3) H
(
(7)
L
Content Stream
Query Stream
19
G1
6
1
2
5
2
6
Q K.
K,
Q
(1)
(2)
(3)
(4)
(5)
(6)
H2(1)
H2)
(3)
(4)
H2( H6
Embedding
The onswer
is
inches
long
G1
6g
62
Figure I: Architecture of the A3 model. Blue entries in the attention mask denote O, and white en- tries denote -o. The model employs 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 the query stream's output. For illustration, we provide an example grouping with G1 21,2,3), G2 (5, G3 {4]. showing how the forward process and causal masks are applied.
operates with a single stream of hidden states, where information flow regulated by causal attention mask. This mask ensures that the representation of the k-th token can only attend to the first tokens, thereby enforcing the AR constraint required for language modeling. While well- suited for the standard left-to-right objective, this design assumes fixed generation order: given the first k tokens, the model trained to treat the (A + 1)-th position as the unique next target. Such rigidity makes it incompatible with any-order prediction, where the next position to be generated need not follow the sequential index.
The encoder-only Transformer, widely used in masked language modeling, represents the opposite design. Rather than causal masking, it processes the full sequence bidirectionally. with missing in- formation represented by mask tokens. Through position embeddings on these masks, the model identifies which locations are to be predicted, allowing arbitrary subsets of tokens to be recon- structed simultaneously. However, this formulation limits dependency modeling: masked positions are predicted in parallel and conditioned only on observed context in single pass. Without recur- sive, multi- layere dependencies across tokens, the encoder-only approach struggles to match the generative fidelity of AR models.
Two-stream Attention Design for Any-order Prediction. To combine the flexibility of encoder- style masking with the dependency modeling strength of autoregression, A3 extend the two-stream attention mechanism proposed by XLNet (Yang et al., 2019). The model maintains two parallel representations for each position: content stream, which encodes semantic and contextual in- formation from the observed tokens, and a query stream, which provides position-awar signals to drive prediction of the next group. This separation allows A3 retain the recursive structure of autoregression while relaxing the generation order constraint of decoder-only models. Figure illustrates the pipeline of the forward process.
4

---

Published as a conference paper at ICLR 2026
Strwom
Query Streom
Content Stream
Quary Streomm
Content Stream
Query Stream
6g 2
6:2
61
12
4
5
5
Stage
Stage 2
Stage
Figure 2: Causal masks for content stream and query stream in different stages. Blue for and white for -cc. Stage 1: AR initialization reproduce AR factorization. Stage 2: Group expansion by allowing groups of size greater than one. Stage 3: Order permutation with introducing any-order prediction.
Formally, let X (#1...., EN) denote the sequence, partitioned into groups {G1..... Gk} In the content stream, the input consists of the observed tokens, embedded and passed through Trans- former layers with designed causal mask. This mask ensures that token at group k can attend to all tokens in groups < k, i.e., both its own group and all groups before it. Thus, the content stream group l aggregates all contextual evidence available up to that point. For group Gk+ the hidden states in the content stream at layer I are computed as:
(t) = Attn (Q= (1) (H=n)
= H G&), V = H (5G)).
(4)
In the query stream, the input shared learnable query vector injected at every position. The key and value matrices are tied to those of the content stream, while the queries are separate. With an appropriately designed causal mask, each query vector at group k can only attend content tokens in groups < k, not including its own group. This forces the query representation to serve as the poseei position- aware predictor for the tokens in group k relying only on prior context rather than future information. Conceptually, the query stream specifies where to predict (positional conditioning), while the content stream provides what to predict (contextual grounding). The hidden states in the query stream at layer I for group G are:
H n = Attn Q= H K= H Gk), V = (<
(5) (1 (1) = w is a learnable query vector shared across positions, and the where the initialization H causal mask ensures the query stream at group k can only access content states from strictly earlier groups.
Finally, the predictive distribution for token 2; Gj is parameterized by:
p(:; (xca)) = Softmax (W. ())
(6)
where L is the final layer and W projects the query hidden state to the vocabulary.
3.2 MULGII-STAGE TRAINING WITH PROGRESSIVE TOKEN GROUPING
Building on the connection between standard AR and A3, we design progressive adaptation strat- egy that smoothly transitions from left-to-right generation to fully flexible any-order prediction. To leverage the stability and strong initialization existing AR models, we begin training A3 from pretrained AR checkpoint and gradually relax its constraints through three stages:
Stage 1: AR Initialization. We align A3 with conventional AR training by setting the two-stream causal masks to exactly reproduce left-to-right factorization (Figure Stage 1). Formally, the sequence øi.N is partitioned into singleton groups:
G {1}, G2 = {2), GN = (N}.
This ensures that P( = I P(æ, xzx), identical to standard AR, providing a stable initialization..
5

---

Published as a conference paper at ICLR 2026
Stage 2: Group Expansion. We expand beyond token-level prediction by allowing groups size greater than one (Figure Stage 2). Concretely, the sequence is partitioned into contiguous segments of fixed size 1, e.g.,
G1 = (1,... 8). G2 = +s+ 2s},
with gradually increased from 1 to 4. This teaches the model to predict multiple tokens jointly within each group while still maintaining AR dependencies across groups.
Stage 3: Order Permutation. We introduce any-order prediction within groups (Figure Stage 3). The group structure G1, G2..... Gx remains sequential, but the token indices assigned to each group are drawn from a random permutation of (V)...,NV. For example, if x is random permutation of indices, then
G = ( (+)), G2 = {T(( 1)...,,, ((2))}
The model therefore learns to predict tokens arbitrary subsets, while still preserving a group-to-group AR factorization:
P(TEN)
II P(ac, xG2e).
(10)
k=1
This exposes the model to diverse intra-group orderings and enables to generalize to arbitrary prediction targets at inference.
By the end of this curriculum, the model is able to predict arbitrary subsets of tokens as coherent groups while preserving the recursive dependency structure of AR. Importantly, every stage of training, each token in the sequence belongs to exactly one group, so all tokens are always predicted, maximizing the learning signal and computational efficiency.
3.3 FLEXIBLE INFERENCE VIA GROUPWISE DECODING
Building the A3 formulation, we propose flexible inference strategies that extend beyond con- ventional AR decoding. The first decoding method we introduce groupwise AR sampling, which generalizes standard left-to-right generation by sampling groups of tokens sequentially rather than strictly one-by-one. Formally, let the token positions of a sequence be partitioned into K groups
= [G1 G2...., Gx where G2 C [I.....n] and U G₁ (l,...,m n... Given a prompt covering groups GI...., Gky the model generates subsequent groups by conditioning on pre- ceding groups:
K
pe G IGk IG
II k=
Pe IGGx).
(11)
Here, FGG denotes the tokens within group Gg, and IG2.. the tokens of all earlier groups. This reduces to the classical AR factorization when G& = for all k but naturally generalizes to larger groups. The procedure is summarized in Algorithm 1. Concretely, several grouping strategies can be applied:
1. Token-wise grouping. Each token treated as its own group, i.e., G3 = {k). The decoding reduces the standard left-to-right AR generation:
(z 2nd) = II (a, Iet).
(12)
2. Fixed-size grouping. Tokens are partitioned into groups of size s, e.g., G5 = {(k - 1)s + 1, 1,..., ks} for {2,4). In this case, the model predicts tokens jointly per step and accelerates decoding.
3. Task-specific grouping, For infilling tasks we allow groups to be arbitrary index subsets and then assign group ids that groups containing masked positions decoded after groups used as context. Concretely, let the sequence be partitioned into left, middle and right index sets L,, M,, R (so ( LU MUMM R). We choose index ko such that

---

Published as a conference paper at ICLR 2026
Algorithm 1 Groupwise AR Sampling
Algorithm 2 Dynamic Resampling
Require: Prompt tokens
F1:m-
Require: Prompt tokens, model fa. criterion
grouping strategy
1: Initialize Fo with prompt tokens, Uo with blank posi- tions
(G1, G2.... Gk model fo
Ensure: Generated sequence Fizz
2: while U, * 10 do
1: Initialize i.mm f- Ii:m
3: for each E U, do
2: Find the last group index ko in the prompt
4:
Compute po(a;; xF.)) 5: end for
3: for k= Ki + to K do
6: Select subset St C Us based on criterion
4: Compute context representa-
7: for each E S, do
tion h fo(
8:
Sample f; ~ po( TF.))
5:
Sample tokens fG,, pol
9: end for
h))
10: Update Fi+1 + F3 St. U141 + Us St
6: end for
11: end while
7:
return Completed sequence dim
12:
return Completed sequence 21.n
every group G3 satisfying G2 n M 0 has h S ko. while every group that contains any masked position satisfies & > ko (groups need not be contiguous and a single group may contain tokens from both T. and R). Under this design, all context groups (those covering I. and R) appear before the masked groups, and the model performs:
pe(xM IG
II po (o IGGx), A=
(13)
which realizes AR dependencies inside each masked group while conditioning both left and right contexts. This flexible assignment enables infilling where context groups are formed from arbitrary subsets of T. U R, and masked spans predicted group-by-group This capability distinguishes A3 from conventional AR models, which cannot directly con- dition on future context during generation.
Dynamic Resampling Inference. Beyond fixed grouping, A3 also supports more adaptive in- ference procedure inspired by iterative refinement (Li et al., 2022; Chen et al., 2024a). Here, the grouping G is not fixed. At each step, the model evaluates all unfinished positions simultaneously, conditioned on the completed tokens. Formally, suppose U, S {l...., ) is the set of unfinished (blank) positions at iteration and is its complement of finished positions. The model computes predictive distributions
po(#; TF)) Vi E U..
(14)
Based on these distributions, we then select subset S, C U, to be committed at this step, according some criterion such maximum confidence, lowest entropy (Kim et al., 2025), or simply random sampling. Once S; is chosen, the tokens at Sy are sampled and added to the finished set:
F1+1 U. UL+1 U, S
(15)
This process repeats until Ur 0, at which point the sequence fully generated. The procedure is summarized in Algorithm 2. The advantage of this dynamic resampling strategy is twofold. First, it allows the model adaptively choose the granularity of generation based on prediction confidence, committing to easy tokens early while deferring more uncertain positions until later. Second, unlike diffusion-style denoising which follows a pre-specified noise schedule, A3 inference directly uses the conditional distributions defined by the AR factorization, ensuring consistency between training and inference.
These inference strategies highlight trade-off between efficiency and flexibility. Fixed-group sam- pling is fast but less adaptive, as performance depends on group alignment with text structure. Dy- namic resampling is slower since all unfinished positions reevaluated at each step, but it yields greater accuracy by adapting token commitment to model confidence. We will compare these strate- gies in the next section real-world tasks.
7

---

Published as a conference paper at ICLR 2026
Table Comprehensive evaluation of different language models. There are types of these models: AR for autoregressive. DD for discrete diffusion, CD for continuous diffusion and A3 for our pro- posed model. For the infilling task, we use ROUGE- 1/2/L score; for other tasks, we use the accuracy (%) metric. refers to the results reported in DiffuLlama (Gong et al., 2025).
QA
CommonSense Reasoning
Infilling
Model
Size
Type
TriQA
HSwag
Wino.
SIQA
PIQA
ROCStorics
Llama-3.1
8B
52.1
76.0
63.9
46.7
80.3
11.7/2.3/10.5
Plaid*
IB
CD
1.2
39.3
51.3
32.3
54.5
12.1/1.1/11.2
Dream
7B
DD
18.3
26.9
51.8
36.6
55.8
11.7/2.3/10.5
Difful.amma
7B
DD
18.5
58.7
56.4
43.2
63.3
23.3555522.2
IB
A3
10.2
40.2
52.8
35.1
64.7
11.8/1.7/11.
A3
3B
15.9
54.3
70.1
11.3/2.3/10.2
8B
A.3
19.4
58.4
60.2
45.2
78.1
19.2/4.6/186
Table 2: Conditional generation quality (measured by perplexity using Llama-3.1 (11ma)311-BB of previous diffusion models with A3 using different dynamic sampling strategy.
Random
Confidence
Entropy
Model
step=512
step=1024
step=512
step=1024
step=512
step=1024
Dream
58.4
46.2
21.3
17.2
18.7
16.4
DiffuLlama
72.3
58.4
24.1
18.3
20.9
14.3
A3-8B
66.4
49.3
20.1
16.8
14.3
11.2
4 EXPERIMENTS
4.1 SETUP
Training Setup. We initialize our models from the LLaMA series, including ILLMA-3.1-88 1-8B. I.L.MA3.2-33 and L.LMMA-3.2-IB (Grattafiori et al., 2024). For training data, we construct mixture of the FineWeb dataset (Penedo et al., 2024) and the SlimPajama dataset (Soboleva et al., 2023), following prior work on DLMs and AR models. From this mixture, we sample 8B tokens and apply sequence packing with maximum context length of 2048. All models are trained with fall-parameter fine-tuning in bf 6. the progressive adaptation recipe, the first two training stages trained for epoch over 20% of the dataset, while the final stage is trained for one epoch over the full dataset. Additional training details are provided in Appendix A.
Evaluation Setup. We adopt the evaluation protocol of Gong al. (2025) compare mod- els against both diffusion and AR baselines. For reading comprehension, we evaluate on TriviaQA (Joshi al., 2017) using exact match accuracy. For commonsense reasoning, we consider Hel- laSwag (Zellers et al., 2019), Winogrande (Sakaguchi et al., 2020), SIQA (Sap et al., 2019), and PIQA (Bisk al., 2020), all assessed by multiple-chot choice accuracy. For story infilling, we use ROC- Stories (Mostafazadeh et al., 2016) and report ROUGE scores (Lin, 2004). We compare against two categories of baselines: (a) the base AR model LLaMA-3.1 1-BB, and (b) recent diffusion language models of varying sizes, including Plaid-1B (Gulrajani & Hashimoto, 2023), Dream-7B (Ye et al., 2025), and DiffuLlama-78 (Gong et al., 2025).
4.2
MAIN RESULTS
The results in Table show that A3 consistently outperforms diffusion- models across QA, commonsense reasoning, and infilling tasks. For example, A3-8B achieves 19.4 accuracy on Triv- iaQA and 78.1 on PIQA, surpassing all the diffusion baselines, while also attaining competitive ROUGE scores for story infilling. These gains are particularly noteworthy given that A3 trained on only 8B tokens, whereas DiffuLlama is trained on 65B. Although A3 still underperforms the AR baseline, this gap is likely attributable to limited training data; with larger-scale pretraining, we expect A3 close the difference further.
8

---

Published as a conference paper at ICLR 2026
log(Perplexity) vs. Group Size
Decoding Time Group Size
Groupwls2 autorogresive
3.6
nanii resampling iconfiencc)
Dynamic resampling (contidenc))
Dynamic resamping lentropy)
Dynomic resampling (entropy)
3.4
3.2
100
Perglesity)
Decoding Time
3.0
2.8
60
2.6
40
2.4
2.2
20
8
10
12
14
16
10
32
14
16
Group
Group size
Figure 3: Unconditional generation log(perplexity) and speed using A3-8B. The perplexity is mea- sured by Llama-3.1-8B and we compare several decoding strategies. Dynamic resampling will cost more time but have lower perplexity.
Table 3: Performance with different training curriculum schedule. We evaulate two variants trained on 2B tokens: 1. original curriculum and, 2. skipping stage 1 and 2 (directly training stage 3: order permutations).
QA TriQA
CommonSense Reasoning
Infilling
Schedule
HSwag
Wino.
SIQA
PIQA
ROCStories
Original
15.6
49.3
56.7
39.6
69.4
13.2/2.3/12.6
Skipping Stage 1 & 2
11.3
44.2
54.1
37.3
64.2
13.1/2.2/12.4
Importantly, A3 demonstrates clear scaling behavior: performance improves steadily from IB to 3B 8B parameters, indicating that the method benefits from larger models in the same way as conventional AR training. Overall, these results confirm that A3 offers a favorable balance between AR and diffusion paradigms, combining strong reasoning accuracy with flexible generation, and holds promise for further improvements under larger-scale training.
Additionally, compare A3-8B with the diffusion models Dream and DiffuLlama conditional generation quality. For each model, we use the same 100 randomly sampled 128-token prefixes from the Pile (Gao al., 2020) conditioning inputs. Each model then generates 2048-token continu- ation, and we evaluate the quality of the generated text using perplexity computed by Llama-3.1 18B.
The results are summarized in Table 2. While Dream achieves slightly lower perplexity under the random sampling schedule, A3-8B consistently outperforms both baselines when using confidence- or entropy- based dynamic sampling, demonstrating clear gains from our sampling strategy.
4.3 ABLATION STUDY
Sampling Strategies. To better understand the trade-off's between the two proposed inference strate- gies in Section 3.3, conducted unconditional generation experiments under the A3 decoding framework. For groupwise AR sampling, we vary the group size from to 4. For dynamic resam- pling, we vary the group size from 16 and implemented two selection criteria: (1) Confidence- based: selecting positions with highest maximum softmax probability. (2) Entropy- positions with minimum output entropy. For each sequence, we sample with temperature of 1.5 and top-p of 0.95. Figure 3 reports the log of perplexity measured by Llama-3.1-88 and the average decoding time for one sequence.
selecting
We observe that dynamic resampling methods consistently achieve lower perplexity than group- wise AR sampling, indicating that they produce higher-quality generations. The confidence- based and entropy- basrop criteria yield very similar performance, with confidence being slightly better smaller group sizes. However, all strategies show a trend of increasing perplexity as group size grows, reflecting the trade-off between decoding granularity and modeling accuracy. We can also

---

Published as a conference paper at ICLR 2026
Table 4: Performance of A3 with different training data
Table 5: Model loss of A3 across context lengths, which is stably small.
on TriviaQA and perplexity measured by Llama-3.1 1-BB..
Model
TriviaQA
log(Perplexity)
Length
Model loss
A3 (6B tokens)
16.2
2.9
256
3.54
A3 (8B tokens)
19.4
2.5
512
3.51
A3 (1OB tokens)
22.5
2.3
1024
3.34
2048
3.23
AR (15T)
52.1
0.8
see that decoding time decreases sharply with larger group sizes, Groupwise AR sampling fastest at the same group size because it only generates the designated group per step, while dynamic re- sampling requires evaluating all unfinished tokens each iteration, making slower. However, as group size increases, dynamic resampling speeds up considerably, nearly matching the efficiency of groupwise sampling at large group sizes.
Overall, these results demonstrate a speed-accuracy trade-off. Groupwise AR sampling is faster but less accurate, while dynamic resampling achieves better perplexity at the cost of slower decod- ing. Importantly, A3 provides the flexibility to choose between these strategies depending the requirements of the application, making it more flexible than conventional AR or diffusion-basee methods.
Curriculum schedule. A3 introduces different causal mask and attention flow from standard AR transformer, and the model must progressively adapt from strict left-to-right prediction to multi- token and eventually arbitrary-order factorization. To assess the sensitivity of the schedule, train two variants on 2B tokens: 1. original curriculum, and 2. skipping stage 1 and 2 (directly training on stage 3: order permutations). Results shown in Table 3. Skipping the early stages consistently hurts performance by 4-6 points on several benchmarks, which proves the importance of such adaptation stage. An adaptive schedule, c.g., based on training loss, may further improve robustness. We plan to investigate this direction in the future work.
Performance with more data. Since the training budget for A3 much less than the baseline (2B for A3, 60B for DiffuLlama and 15T for Llama-3,1-8B)) in order to isolate the architecture effect on the worse performance than the AR baseline, we track how A3 improves under increasing post- training data. We use 6B, 8B (default) and 10B tokens to train A3. The results are shown in Table Performance increases steadily with more data. This confirms that A3 benefits strongly from data scale and that the gap to fully-trained AR models is due to training budget, not limitation of the A3 architecture.
Robustness context length. In order to investigate whether A3 is robust across different context lengths, we input contexts with length 512, 1024 and 2048 to A3 and calculate the loss. The results shown in Table 5. The model loss keeps stable within the training length, indicating the robustness of A3 across different context lengths.
CONCLUSION
We have presented Any-order Any-subset Autoregressive modeling (A3), a novel framework that generalizes traditional autoregressive factorization enable flexible, groupwise generation of to- kens in arbitrary orders. By combining a two-stream attention architecture with progressive train- ing strategy, A3 achieves the dual goals of generation flexibility and modeling stability. Our ap- proach supports wide range of decoding strategies, including groupwise autoregressive sampling and dynamic resampling, offering a tunable trade-off between speed and accuracy. Through com- prehensive experiments, we demonstrate that A3 outperforms diffusion-based models in reasoning, question answering, and infilling tasks. These results highlight A3's ability to balance efficiency, flexibility, and quality, making it promising direction for future sequence modeling. In the fu- ture. we plan to explore scaling A3 to larger models and datasets, as well as applying it to more challenging tasks such long-context reasoning.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
Yisen Wang was supported by National Natural Science Foundation of China (92370129, 62376010), Beijing Major Science and Technology Project under Contract no. 225110000842500, Beijing Nova Program (20230484344, 20240484642), and State Key Laboratory of General Arti- ficial Intelligence. Yifei Wang was supported part by the NSF Al Institute TILOS (NSF CCF- 21 2112665), and an Alexander von Humboldt Professorship. Zeming Wei was supported by Beijing Natural Science Foundation (QY24035).
ETHICS STATEMENT
This work complies with the ICLR Code of Ethics. While our methods are general, they may be applied in contexts with societal implications, including risks related to bias, fairness, and privacy. We encourage responsible use and declare conflicts of interest.
REPRODUCIBILITY STATEMENT
We provide detailed descriptions of our methodology, datasets, model configurations, and eval- uation metries in both the main text and the Appendix. Codes will be released upon accep- tance.
REFERENCES
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models discrete state-spaces. NeurIPS, 2021. 3, 15
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. ACL, 2024. 16, 17
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. AAAI, 2020. 2,
Tianle Cai, Yuhong Li. Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint ar arXiv:2401. 10777477444 2024. 16
Pinzhen Chen, Zhicheng Guo, Barry Haddow, and Kenneth Heafield. Iterative translation refinement with large language models. In EAMT, 2024a. 7
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Lon- gloRA: Efficient fine-tuning of long-context large language models. ICLR, 2024b. 15
Sander Dieleman, Laurent Sartran, Arman Roshannai. Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffu- sion for categorical data. arXiv preprint 1. 2022. 15
Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling? In ICLR, 2025. 15
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. arXiv preprint 24Xi 10171, 2024. 15
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima. et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 100iv:2101. 1r100:2101.00007, 2020. 9

---

Published as a conference paper at ICLR 2026
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction. In ICML, 2024. 3
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. arXiv preprint arXiv:2210.08933, 2022. 15, 16
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al. Scaling diffusion language models via adaptation from 15
autoregressive models. In ICLR, 2025.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, al. The Ilama 3 herd of models. arXiv preprint arXi::2407.2178. 2024 8
Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non-autoregressive
neural machine translation. arXiv preprint arXiv: 171100281,,,, 2017. 15
Jiatao Gu, Changhan Wang, and Junbo Zhao. Levenshtein transformer. NeurIPS, 2019. 16
Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood-based diffusion language models. In NeurIPS, 2023. 1, 8, 15
Junliang Guo, Xu Tan, Linli Xu, Tao Qin, Enhong Chen, and Tie-Yan Yan Liu. Fine-tuning by curriculum learning for non-autoregressive neural machine translation. In AAAI, 2020. 16
Xiaojun Guo, Ang Li, Yifei Wang, Stefanie Jegelka, and Yisen Wang. GI: Teaching ilms to reason on graphs with reinforcement learning. In NeurlPS, 2025. 15
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A Smith. Don't stop pretraining: Adapt language models to domains and tasks. In ACL, 2020. 15
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-Im: Semi-autoregressive simplex- based diffusion language model for text generation and modular control. arXiv preprint arXiv:2210.17432, 2022. 15, 16
Zhengfu He, Tianxiang Sun, Qiong Tang, Kuanning Wang, Xuan-Jing Huang, and Xipeng Qiu. Diffusionbert: Improving generative masked language models with diffusion models. In ACL, 2023. 2,
Yosuke Higuchi, Shinji Watanabe, Nanxin Chen, Tetsuji Ogawa, and Tetsunori Kobayashi. Mask ctc: Non-autoregressive end-to-end asr with etc and mask predict. arXiv preprint arXiv:2005.08700, 2020. 16
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020. 15
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffuuion: Learning categorical distributions. In NeurIPS, 2021. 15
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettleooer.. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In ACL, 2017. 2,8
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu. Continual pre- training of language models. In ICLR, 2023. 15
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham M Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. In ICML, 2025. 1, 1,7
Siqi Kou, Lanxiang Hu, Zhezhi He, Zhijie Deng, and Hao Zhang. Clims: consistency large language models. In ICML,, 2024. 3
Jason Lee, Elman Mansimov, and Kyunghyun Cho. Deterministic non-autoregressive neural se-
quence modeling by iterative refinement. In EMNLP, 2018. 16
12

---

Published as a conference paper at ICLR 2026
Yaniv Leviathan, Matan Kalman, and Yossi Matias, Fast inference from transformers via speculative decoding. In ICML, 2023. 16
Xiang Li. John Thickstun, Ishaan Guirajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion- Im improves controllable text generation. In NeurlPS, 2022.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty. arXiv preprint arXiv:2401.15077 2024. 16
Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint 2tt.::331../683444 2023. 15, 16
Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vander- wende, Pushmeet Kohli, and James Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories. In NAACL, 2016. 2, 8
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, JUN ZHOU, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. NeurlPS, 2025. 2,
Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In ICLR, 2025. 15
Guilherme Penedo, Hynek Kydlfček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. NeurIPS, 2024.
Jack W Rac, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05500 2019. 16
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In NeurlPS, 2024. 1, 15
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. In AAAI, 2020. 2, 8
Mohammad Samragh, Iman Mirzadeh, Keivan Alizadeh Vahid, Fartash Faghri, Minsik Cho, Moin Nabi, Devang Naik, and Mehrdad Farajtabar. Scaling smart: Accelerating large language model pre-training with small model initialization. arXiv preprint arXiv:2409.12903 2024. 15
Maarten Sap. Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. In EMNLP, 2019. EMN1L,,0099.. 2, 8
Tianxiao Shen, Hao Peng. Ruoqi Shen, Yao Fu, Zaid Harchaoui, and Yejin Choi. Film: Fill-in language models for any-order generation. arXiv preprint arXiv:2310.09930, 2023. 16
Daria Soboleva, Faisal Al-Khate,,, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023. URL. https: / //hhugiingfatt co datasets/ cttttttseeeeerbreeeeeeee SS imm j ama- 627B. 8
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICMI,, 2015. 15
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020. 15
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. NewrIPS, 2018. 16
13

---

Published as a conference paper at ICLR 2026
Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit. Insertion transformer: Flexible sequence generation via insertion operations. In ICML, 2019. 16
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309. 16039, 2023. 15
Yiheng Xu, Hongjin SU, Chen Xing. Boyu Mi, Qian Liu, Weijia Shi, Binyuan Hui, Fan Zhou, Yitao Liu, Tianbao Xie, Zhoujun Cheng, Siheng Zhao, Lingpeng Kong, Bailin Wang, Caiming Xiong, and Tao Yu. Lemur: Harmonizing natural language and code for language agents. In ICLR, 2024.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In NeurlPS, 2019. 3,
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508. /5487, 2025.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In ACL, 2019. 2, 8
Qi Zhang, Tianqi Du, Haotian Huang, Yifei Wang, and Yisen Wang. Look ahead or look around? theoretical comparison between autoregressive and masked pretraining. In ICMI,, 2024.
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. In COLM, 2024. 1, 15
14

---

Published as a conference paper at ICLR 2026
A TRAINING HYPERPARAMETERS
We list the hyperparameters in the training stage in Table
Table 6: The hyperparameter list
Hyperparameter
Value
Training
Batch Size
Epoch
[0.2, 0.2, 1]
Optimizer
Adam W
LR
2e-5 (0.9, 0.999)
Betas
Weight Decay
0.01 WarmupLR [50, 50, 50] 2048
LR Schedule
Warmup Iters
Max Sequence Length
Sampling (Section 4.3)
Top-p
0.95
Temperature
1.5
B ADDITIONAL RELATED WORK
Continue Pre-training. Pretraining large language models has been proven to be complex and computationally expensive (Samragh et al., 2024). Consequently, continued pre-training has been proposed as an effective method to adapt existing large language models to specific domains (Ke et al., 2023; Gururangan et al., 2020; Guo et al., 2025) to endow them with new capabilities, such as handling longer contexts (Chen et al., 2024b; Fu et al., 2024; Xiong et al., 2023; Fang et al., 2025) code generation (Xu et al., 2024). Notably, certain continued pre-training efforts, such as those in scaling diffusion language models (Gong et al., 2025), have transcended autoregressive (AR) language modeling by converting large language models into diffusion based architectures, therehy enabling parallel token generation. In contrast, our work retains autoregressive language modeling but innovatively incorporates two-stream architecture and cuni-autorerressvve decod- ing to similarly support parallel prediction of multiple tokens, achieving significant reductions in inference latency compared to both autoregressive and diffusion- baselines.
Diffusion Language Model. Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) have demonstrated remarkable capabilities image generation. Consequently, series of works have sought to extend diffusion models to text generation, which can be roughly divided into the continuous diffusion model and the discrete diffusion model. One straightforward approach involves embedding text data into continuous space and directly applying diffusion models (Li et al., 2022; Gong et al., 2022; Han et al., 2022; Dieleman et al., 2022). However, the scalability of continuous diffusion methods remains challenge, they require substantially greater compu- tational cost compared to AR models to achieve equivalent performance(GGlrrjan & Hashimoto, 2023). To better accommodate the discrete nature of text, an alternative paradigm replaces con- tinuous diffusion with discrete process, introducing an absorbing [MASK] state as noise (Austin al., 2021: Hoogeboom al., 2021: Zheng et al., 2024; Sahoo et al., 2024). Lou al. (2023) demonstrated that masked diffusion models (MDMs) achieve perplexity comparable to even sur- passing that of AR models at the GPT-2 scale. Ou et al. (2025) established foundational theoretical results, affirming the fcasibility of MDMs. In comparison MDMs, our method similarly enables parallel prediction of groups of tokens and leverages bidirectional context. However, by retaining autoregressive modeling, our approach utilizes every token during training, thereby facilitating faster convergence relative to MDMs.
Non-autoregressive Generation. Non-autoregressive (NAR) generation (Gu et al., 2017) accel- erates inference by producing target tokens in parallel, eliminating the dependency on previously generated tokens inherent in traditional AR models. This approach substantially improves genera-
15

---

Published as a conference paper at ICLR 2026
tion speed, but potentially at the cost of some accuracy. Conventional NAR models have focused mainly on machine translation and automatic speech recogmition, emphasizing the trade-off between speed and quality (Higuchi et al., 2020; Lee et al., 30. For instance, Guo et al. (2020) employed labels in translation tasks, training NAR models via a curriculum learning strategy on attention masks. However, early NAR methods in text generation remain constrained by issues of global coherence and diversity. In recent years, insertion-based models(Stern et al., 2019; Gu et al., 2019) and diffusion language models (Gong al., 2022; Lou et al., 2023; Shen al., 2023) have progressively addressed these limitations. For example, SSD-LM (Han et al., 2022) leverages cemi-autoregressvve manner, enabling local bidirectional context updates. Compared to diffusion language models, our method similarly achieves parallel multi-token prediction and bidirectional context modeling, while benefiting from faster training convergence due to the absence of masked tokens during training.
soquencee-o-sequuee
a diffusion- based language model to generate text blocks
Multi-token Prediction. Multi-Token Prediction (MTP) seeks to move beyond strictly stepwise next-token prediction by enabling models to forecast multiple future tokens per autoregressive step, improving inference efficiency with minimal quality degradation. Some methods achieve MTP by modifying the Transformer architecture. For instance, Stern et al. (2018) demonstrated through blockwise decoding that multi-token acceptance feasible when supported by verification mech- anisms. Medusa (Cai al., 2024) add parallel heads predict multi-step continuations with lightweight acceptance rules. Other approaches realize MTP via multi-model collaboration. Specu- lative decoding (Leviathan al., 2023) uses a draft model to propose multi-token candidates that the main model verifies, while EAGLE-like validators (Li et al., 2024) further strengthen multi-token verification. Our method similarly modifies the Transformer structure, but innovatively leverages a two-stream architecture and semi-autoregressive decoding to achieve MTP, significantly reducing generation latency while attaining higher accuracy than baselines.
C THE USE OF LARGE LANGUAGE MODELS (LLMS)
In this work, LLMs are primarily employed for polishing the language of the manuseript to ensure grammatical correctness and coherence. Importantly, all conceptual development, experimental de- sign, and result interpretation are conducted independently by the authors. The use of LLMs is strictly limited to auxiliary tasks, ensuring that the scientific contributions of this paper remain en- tirely unaffected by such tools.
D MORE ANALYSIS EXPERIMENTS
Comparison with speculative decoding. We additionally conduct generation experiments in Figure 3. We use speculative decoding (Leviathan et al., 2023) with Llama-3.2-1B as the draft model and Llama-3.1-81 as the target model, matching A3's group size (max 4). Perplexity is measured by Llama-3. 1-70B. Results are shown in Table 7. Speculative decoding achieves lower perplexity but at higher wall-clock cost, while A3 provides competitive quality with faster decoding in this setting.
Importantly, semi-AR methods such as speculative decoding and multi-token prediction are decoding-leved accelcrators: they maintain the same left-to-right AR factorization and improve effi- ciency through draft-model proposals or multi-step predictions. In contrast, A3 changes the model factorization itself by enabling groupwise and permutation- based prediction. This makes A3 or- thogonal to speculative/MTP: these accelerations can also be applied on top of A3's factorization in principle.
Results on longer contexts. To evaluate whether A3 remains stable under significantly longer con- texts than 2k tokens, finetuned both Llama-3.1-8P and A3-8B 8k-length sequences from PG19 (Rac et al., 2019) using the same training budget (100 steps, batch size 64). We then evalu- ated them the single-document QA task from LongBench v1 (Bai et al., 2024), which requires reasoning over long passages. We used dynamic sampling for A3 with group size and group size 2. The results are shown in Table 8.
A3 (group size 1) improves over the AR baseline with 1.7%, suggesting that the A3 factorization does not degrade long-context modeling and may offer small gains without parallel decoding. A3
16

---

Published as a conference paper at ICLR 2026
Table Comparison between speculative decoding and A3 on generating perplexity and time.
(pepleeely
Time
Speculative decoding
1.9
1.2x
A3
2.1
1x
Table 8: Accuracy and time comparison Single Document QA task of LongBench v1 (Bai al,, 2024).
QA task acc (%)
Average Time
Llama-3. 3--BBB (baseline)
25.4
1.0x
A3 (group size 1)
27.1
1.3x
A3 (group size = 2)
22.5
0.7x
(group size 2) achieves 30% faster decoding, demonstrating that A3's groupwise inference can provide real latency benefits in longer contexts. This result shows that larger groups introduce more parallelism but can slightly reduce accuracy, which consistent with our analyses in shorter contexts.
Our 8k experiments indicate that A3 can scale to significantly longer sequences without degradation and provides decoding-time time advantages via groupwise generation. These results support the poten-
tial of A3 for future long-context extensions and we will explore longer context in future works.
17


Reviews:
[{"review_id": "RWHhQfBZsn", "reviewer": "Reviewer_6bME", "summary": "This paper introduces A³ (Any-order Any-subset Autoregressive modeling), a generalization of standard autoregressive (AR) language modeling. Instead of predicting tokens strictly left-to-right, A³ allows generation in arbitrary orders and subsets while maintaining a valid probabilistic formulation.\n\nThe method builds on a two-stream attention mechanism (content/query streams) similar to XLNet and uses a curriculum training schedule that progressively transitions from left-to-right to multi-token and random-order prediction. At inference, A³ supports both groupwise parallel decoding and dynamic resampling, trading off speed and quality.\n\nExperiments across QA, reasoning, and infilling tasks show that A³ matches or surpasses diffusion-based models (e.g., DiffuLlama, Dream) using significantly less data, and maintains competitive performance to AR models while enabling faster, flexible generation.\n\nOverall it is a practical step toward bridging AR and diffusion paradigms, retaining AR stability while introducing parallel generation flexibility.", "questions": "1. How does A³ fundamentally differ from prior permutation-based autoregressive approaches such as XLNet or Permutation LM beyond the introduction of multi-token subsets? Could you clarify what new modeling capacity A³ enables that these earlier frameworks cannot?\n\n2. The paper emphasizes that A³ enables parallel or groupwise decoding. Could you provide concrete runtime or latency benchmarks (e.g., decoding speedups vs. standard AR and diffusion models) to quantify the practical benefit of this flexibility?\n\n3. The curriculum that transitions from left-to-right to any-order prediction seems important. How sensitive is model performance to the curriculum schedule (e.g., ratio of L2R vs. random-order batches)? Have you explored automatically learned or adaptive curricula?]\n\n4. Since any-order generation can, in principle, apply to structured data (e.g., image or audio tokens), do you anticipate A³ extending naturally to multimodal diffusion–AR hybrids, or would significant architectural adjustments be required?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Insightful formulation. The paper presents a very interesting perspective on any-order, any-subset autoregressive modeling, effectively bridging the strengths of AR and masked diffusion models in a unified probabilistic framework.\n\n2. Clarity and organization. The paper is well written and easy to follow, with clear explanations, sound motivation, and well-structured methodology.\n\n3. Experiments across multiple reasoning and generation benchmarks validate the effectiveness of the proposed approach, showing competitive or superior performance with improved flexibility.", "weaknesses": "1. What's new relative to prior AR generalizations? The author mentioned that the proposed method builds closely on existing ideas from XLNet (permutation-based AR) and masked diffusion modeling, with the main difference being a unified training/inference view. While conceptually elegant, the contribution may feel incremental rather than fundamentally new.\n\n2. Evaluation scope and ablations are limited. Experiments mainly compare against diffusion-style baselines; there is less analysis against modern AR variants (e.g., speculative decoding, parallelized transformers) or ablations on key design choices like grouping strategy and curriculum schedule.\n\n3. Practical speed–quality trade-offs unclear. Although the paper claims parallel generation benefits, detailed runtime comparisons and latency measurements are missing, making it hard to assess the real efficiency gains of “any-order” decoding in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s9QMuaDpFW", "reviewer": "Reviewer_TFMj", "summary": "The paper addresses a challenge in sequence generation: reconciling the training stability/ generation quality of autoregressive (AR) models with the flexibility/ parallelism of diffusion-based models. This problem is theoretically meaningful, as existing paradigms struggle to balance efficiency, quality, and adaptability—key requirements for real-world applications like long-context reasoning, text infilling, and fast content generation.\n\nA3 resolves this by generalizing AR factorization to support \"groupwise token prediction in arbitrary orders,\" with three core designs:\n(1) two-stream attention architecture, (2) progressive training strategy, and (3) flexible inference, which supports groupwise AR sampling (fast, fixed group sizes) and dynamic resampling (high-quality, confidence/entropy-based subset selection).\n\nEmpirically, A3 (trained on 2B tokens) outperforms state-of-the-art diffusion models across QA (TriviaQA), commonsense reasoning (PIQA, HellaSwag), and infilling (ROCStories) tasks (e.g., A3-8B achieves 78.1% PIQA accuracy vs. 63.3% for DiffuLlama-7B). It also surpasses AR models in AR-disadvantaged tasks (e.g., ROCStories ROUGE-L: 18.6 vs. 10.5 for Llama-3.1-8B).", "questions": "See weaknesses.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper focus a critical problem, which addresses a gap in sequence generation—reconciling AR’s stability/quality with diffusion’s flexibility/parallelism—aligned with real-world needs (long-context generation, infilling).\n\n2. Explicitly outperforms AR models in infilling (ROCStories) by utilizing bidirectional context, validating its flexibility.", "weaknesses": "1. Fails to cite or discuss existing NAR research on Block-wise generation (e.g., Block-AR models that split sequences into fixed/masked blocks for parallel prediction) and progressive training (e.g., curriculum-based NAR training that increments block size or relaxes order constraints). This gap obscures A3’s incremental innovation—readers cannot distinguish whether A3’s groupwise/progressive designs are novel or iterative improvements on prior NAR work.\n\n2. No controlled experiment comparing A3 with an AR model trained on the same 2B tokens (e.g., fine-tuned Llama-3.1-8B). The gap with top AR models (e.g., 19.4% vs. 52.1% on TriviaQA) may stem from data scarcity, not A3’s design—undermining assessment of its architectural advantage.\n\n2. Only evaluates sequences of length 2048, with no tests on long contexts (16k/128k tokens). A3’s claimed advantage in solving AR’s long-context inefficiency remains unproven.\n\n3. Figure 3 claims 3–4x faster decoding than AR but lacks direct time data and performances for baseline AR models (same hardware/sequence length) and explicit absolute time values, making efficiency gains hard to quantify.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uptXM429Gh", "reviewer": "Reviewer_yNMv", "summary": "The paper proposes A3 (Any-order Any-subset Autoregressive), a sequence modeling framework that keeps the probabilistic rigor and training stability of standard autoregressive (AR) models but relaxes the decoding order to arbitrary groups of tokens. Concretely, A3 factorizes a sequence into groups and trains the model to predict any group conditioned on previously generated groups, using a two-stream attention layout to separate “what has been generated” from “what is now being predicted.” At inference, A3 can decode group-by-group (for speed) or do dynamic re-sampling of uncertain positions (for quality), so it can handle infilling and other non-left-to-right scenarios in a single model. Empirically, on common LM benchmarks (QA, commonsense, story cloze), A3 outperforms discrete-diffusion / masked-iterative baselines at similar or larger compute, but still trails a comparable AR model that was trained on far more tokens.", "questions": "1. On distributional correctness: your dynamic re-sampling is essentially a confidence- or entropy-driven iterative refinement. How do you ensure that, after several such rounds, the final joint over all tokens is still close to the model’s intended factorization — and how does this compare empirically to ASSD-style “correct-by-construction” decoding for any-subset ARMs? (A small experiment against the scheme in Guo & Ermon 2025 would make the generative claim much harder to dispute.)\n2. On asymptotic scaling vs. plain AR: you attribute the quality gap to “only 2B tokens,” but you also inject extra supervision signals (more factorization patterns, more masked positions). Do you have evidence that, at equal data and wall-clock, A3 does not hurt the base AR’s next-token perplexity, especially on long contexts where order permutations are most disruptive? A curve like “data ↑ → (AR PPL, A3 PPL) → gap ↓” would make the core claim much more convincing.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Clear target: AR-level stability + arbitrary-order generation. The paper correctly identifies a real and currently active gap: classic AR models are stable and likelihood-faithful but order-rigid, while discrete diffusion / iterative masked LM are order-flexible but multi-step, hyperparameter-sensitive, and sometimes harder to train. Proposing a single framework that “looks like AR to the optimizer” but “behaves like an infiller” is a sensible and timely goal.\n2. Groupwise factorization is a clean, composable idea. Writing ($p(x)=\\prod_k p(x_{G_k}\\mid x_{G_{<k}})$) and randomizing the grouping/order during training is conceptually simple, keeps likelihood well-defined, and reuses the well-understood AR training pipeline. Compared with diffusion-style schedules, this reduces the number of design knobs while still giving order flexibility.\n3. Two-stream attention is a reasonable architectural choice. Reusing an XLNet-style separation between “content” and “query” to implement “see only past groups but write to current group” is a sensible reuse of an established mechanism, so the method is not just algorithmic but also implementable on existing Transformer stacks.", "weaknesses": "1. Parallel generation is only *heuristically* correct, not *distributionally* correct. The paper’s decoding story (“decode some groups, resample the uncertain ones”) is an engineering compromise, but it does not give the kind of provable, joint-distribution-correct parallel sampling that very recent any-subset AR work is starting to provide (e.g. ASSD in Guo & Ermon 2025) — those works explicitly address the mismatch between parallel predictions and the target joint, while A3 largely sidesteps it. This makes the “any-subset” claim weaker on the sampling-theoretic dimension.\n2. Comparisons are not load-bearing for the main claim. The paper leans on the fact that it beats discrete/diffusion-style baselines, but those baselines are known to be data/step/hyperparameter hungry; meanwhile A3 is compared to a much better funded AR line only anecdotally (“we used fewer tokens, so we lag”). Without an equal-data, equal-backbone, equal-context comparison against a strong AR infiller (e.g. a repurposed XLNet / permutation LM, or an AR+speculative multi-token decoder), it is hard to measure how much of the gain comes from the grouping idea itself vs. from simply staying in the AR training regime.\n3. Method novelty is incremental over permutation-style / two-stream LMs. At a high level, A3 = permutation/pseudo-permutation LM (XLNet-like) + explicit grouping curriculum + a decoding heuristic. The paper frames it as a “generalization of AR to any order, any subset,” but several of the enabling ingredients — two-stream separation, masked/pseudo-permutation training, curriculum over masks — have been explored in earlier LM or masked-hard-attention work. The paper’s specific combination is tidy, but the conceptual leap is smaller than the title suggests.\n4. Group choice is exogenous and may be the hard part. The model is trained on random / curriculum-defined partitions, but the real deployments that need “any subset” (layout-controlled editing, multi-span infilling, multi-document patching) tend to have structured subsets (spans aligned to discourse / layout). A3 doesn’t learn the grouping policy, and the paper doesn’t show that the learned model is robust to strongly biased or adversarial group patterns. That is exactly where arbitrary-order models tend to break.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7AZsy9xQre", "reviewer": "Reviewer_ky8C", "summary": "The paper proposes A3, a sequence modeling framework that generalizes AR factorization to predict arbitrary groups of tokens in any order, aiming to combine AR’s likelihood-faithfulness with the flexibility/parallelism of diffusion and semi-AR methods. The core is a two-stream attention design (content vs. query streams) plus a progressive curriculum (AR init → group expansion → order permutation). Experiments on QA/commonsense/story-infilling show A3 beats several diffusion LMs and scales with model size, though it still trails an AR baseline at equal parameter counts; the authors attribute this to smaller pretraining budgets.", "questions": "- Compute-normalized comparison. Can you provide training FLOPs and data tokens for A3 vs. AR and diffusion baselines, and where possible train the AR baseline on the same 2B tokens to isolate the architectural effect? \n- Semi-AR baselines. How does A3 compare against modern multi-token-prediction methods under matched setups? A small-scale controlled study would be informative. \n- Curriculum sensitivity. What happens if you skip Stage 2 or change the group-size schedule? Is the two-stream architecture alone sufficient to get most of the gains? \n- Inference consistency. For dynamic resampling, does there exist a well-defined induced factorization (e.g., an ordering over committed sets) that preserves AR semantics, or is this best viewed as a heuristic? Any failure cases?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Neat architectural idea. The two-stream attention extends XLNet-style permutation to groups: content attends to ≤ current group; query attends only to prior groups, yielding a clean predictive head for the next group. The masks and equations are well specified.\n- Thoughtful training curriculum. The three-stage progression (singleton AR → contiguous groups → permuted grouping) offers a practical path from a pretrained AR model to any-order behavior.\n- Flexible inference recipes. The paper describes groupwise AR sampling and dynamic resampling, articulating a speed–quality trade-off and illustrating it empirically.\n- Empirical signal vs. diffusion LMs. On TriviaQA/PIQA/etc., A3 (1B–8B) outperforms diffusion baselines and shows sensible scaling; the table also situates an AR baseline.", "weaknesses": "- Missing baselines from the semi-AR family. The related-work discussion notes multi-token prediction and insertion models, but the experiments do not include strong semi-autoregressive baselines (e.g.,  speculative-with-MTP) that directly target speed-ups while preserving AR quality. This makes it hard to attribute A3’s benefits to grouping vs. other established accelerations.\n- Latency reporting is proxy-based. The speed-quality trade-off uses per-sequence decoding time and Llama-measured perplexity for unconditional generation; end-to-end wall-clock latency under realistic batching and KV-cache reuse is not reported for the task benchmarks. More concrete throughput/latency numbers (tokens/s, ms/sample) would help.\n- Ablation depth. The method depends on group size schedules, permutation schemes, and the two-stream mask design. The paper lacks sensitivity analyses for (i) curriculum choices (stage lengths, s progression), (ii) group selection criteria during dynamic resampling (beyond brief entropy vs. confidence), and (iii) robustness across context lengths.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)