Any-Order Any-Subset AR Modeling (A³)

• The paper introduces a generalized framework that extends classical autoregression by modeling arbitrary subsets and orders using rigorous factorization and two-stream attention.
• Its architecture leverages order-agnostic factorization with principled batching and decoupled semantic and structural attention to balance prediction needs.
• Empirical benchmarks reveal that A³ achieves state-of-the-art performance in language, image, and tabular tasks while enabling efficient, parallel, and speculative decoding.

Any-Order Any-Subset Autoregressive Modeling (A³) is a generalized framework for generative modeling of discrete sequences and arrays, unifying and strictly extending classical autoregression, masked modeling, and diffusion-inspired multi-group generation. A³ enables efficient, tractable joint or conditional sampling and density estimation for arbitrary subsets and permutations of variables, thereby supporting flexible tasks such as infilling, outpainting, rewriting, and parallelized decoding. The central architectural and algorithmic advances are based on rigorous order-agnostic factorization, two-stream attention for structural–semantic decoupling, and principled batching strategies for joint prediction. A³ subsumes popular lines such as XLNet, MaskGIT, and discrete diffusion models, while restoring the depth and consistency of autoregressive approaches.

1. Formalism: Order- and Subset-Agnostic Autoregression

A³ models the joint distribution of a sequence $x = (x_1, \dots, x_n)$ by decomposing $x$ into $K$ arbitrary, disjoint subsets $(S_1, \ldots, S_K)$ according to a permutation $\pi$ of indices $[1, n]$. The generalized factorization is: $P(x) = \prod_{k=1}^K P(x_{S_k} \mid x_{S_{<k}})$ where $x_{S_{<k}} = \bigcup_{j=1}^{k-1} x_{S_j}$, and set-prediction within $S_k$ can be performed in parallel (Liu et al., 2024, Du et al., 19 Jan 2026). This formulation recovers standard AR when all $S_k$ are singleton sets in left-to-right order, and masked AR (MAR) or diffusion-style objectives when $x$0 and $x$1 is the set of masked positions. Any permutation and any grouping of indices are permissible, making A³ inherently order- and subset-agnostic (Shih et al., 2022, Pynadath et al., 17 Feb 2026).

This flexibility supports conditional inference for arbitrary masks $x$2: $x$3 where $x$4 is an ordering of $x$5. All AO-/AS-ARM architectures rely on causal masking to enforce this dependency structure.

2. Architectures: Two-Stream Attention and Masking

A³ is operationalized through attention architectures that enforce set-wise and order-wise causality, enabling parallel prediction within subsets while strictly preventing future-token leakage. The two-stream attention mechanism, introduced in XLNet and essential for competitive A³ scaling, tracks both content and query representations per position:

• Content stream ($x$6): Aggregates information from all tokens in current/previous groups.
• Query stream ($x$7): Attends only to content from strictly earlier groups, preventing access to current target tokens (Du et al., 19 Jan 2026, Liu et al., 2024).

Attention masks for both streams enforce the group-wise autoregressive order: $x$8 This separation addresses the structural–semantic tradeoff: $x$9 can summarize all prior generative context, while $K$0 specializes in predicting new tokens, avoiding entanglement that would arise in a single-stream setup (Pynadath et al., 17 Feb 2026). Decoupled rotary position embeddings (RoPE) offer a partial single-stream alternative for short sequences, but degrade with long or highly non-local orderings.

The Fully Masked Transformer (FMT) implements these ideas for images and multimodal data, constructing encoder/decoder masks to enable flexible groupwise decomposition (Liu et al., 2024).

3. Training Objectives and Protocols

Training A³ models proceeds by sampling random permutations $K$1 and partitions $K$2 per minibatch, instantiating the full spectrum from left-to-right AR to masked/diffusion-like objectives. The primary loss is the negative log-likelihood: $K$3 All tokens are always predicted exactly once per training example. Progressive adaptation from a pretrained AR checkpoint is recommended: initialize with singleton groups, expand to blocks of size $K$4, and finally permute orderings within and across groups (Du et al., 19 Jan 2026).

To mitigate redundancy and optimize for actual inference use, recent work recommends minimal edge selection (predicting only the “max index” conditional per mask) and loss reweighting by test-time conditional frequency (Shih et al., 2022). This reduces the number of unique conditionals trained and accelerates convergence.

4. Parallel and Speculative Decoding

A³ architectures support parallel, order-agnostic generation over arbitrary masks—crucial for applications like infilling, outpainting, or low-latency decoding. The Any-Subset Speculative Decoding (ASSD) algorithm (Guo et al., 29 Apr 2025) efficiently samples $K$5 tokens in parallel:

• Draft phase: Propose $K$6 tokens independently under partial conditioning.
• Verification phase: Recompute true joint conditional for each draft.
• Rejection/correction phase: Accept with probability $K$7; otherwise, resample and rerun.

Correctness is guaranteed by induction, and the number of model calls is $K$8, never exceeding the number of tokens. This speculative parallelism yields up to $K$9 speed-up in wall-clock generation and strictly preserves sample quality (e.g., identical perplexity to sequential decoding on WikiText infilling tasks) (Guo et al., 29 Apr 2025).

Fully Masked Transformer and related frameworks also enable exploitation of modern hardware through batch prediction and caching, supporting both few-step and AR-efficient regimes (Liu et al., 2024).

5. Empirical Results and Benchmarks

A³ models match or outperform strong discrete diffusion and conditionally masked baselines across domains:

• Language modeling and infilling: On ROCStories infilling, A³-8B achieves ROUGE-1/2/L of 19.2/4.6/18.6 versus DiffuLlama-7B's 23.3/5.5/21.2; PIQA commonsense accuracy improves by 14.8 points over diffusion (Du et al., 19 Jan 2026). ASSD-augmented XLNet matches or exceeds much larger diffusion models on code completion.
• Image generation: FMT-trained A³ models obtain FID = 3.72 (FMT-L, AR order, ImageNet 256) and adapt robustly between regimes. Random-order training ensures generalization to unseen inference orders, while rigidly ordered models collapse under permutation shifts (Liu et al., 2024).
• Tabular and continuous data: A³ (MAC protocol) matches or exceeds previous bests (ARDM, ACE) in joint and marginal bits-per-dimension on CIFAR-10 (e.g., joint 2.81, marginal 1.81) and similar for ImageNet32 and UCI datasets (Shih et al., 2022).

The flexibility of partition size $(S_1, \ldots, S_K)$0 allows interpolation between highly parallel inference (few sets), which benefits efficiency, and fine-grained AR (many sets), which maximizes sample quality.

6. The Structural–Semantic Tradeoff and Design Guidelines

A³ exposes a fundamental tension in any-order modeling: prediction accuracy benefits from attending semantically to neighboring tokens (semantic locality), while global summarization (for full-context reasoning) demands structural recency (structural locality). For arbitrary permutations, these axes can diverge substantially. Two-stream attention resolves this tradeoff by specializing streams for each, whereas single-stream alternatives (including decoupled RoPE) fail to scale for long sequences (Pynadath et al., 17 Feb 2026).

For practical deployment:

• Use two-stream attention on long sequences, multimodal data, or when strict correctness of both local and global statistics is critical.
• For lightweight or short-context tasks, decoupled RoPE reduces overhead with minimal loss under moderate permutation locality.
• Adaptive grouping schedules ($(S_1, \ldots, S_K)$1 selection, dynamic partitioning) can further improve tradeoff control (Liu et al., 2024).

7. Applications, Limitations, and Extensions

A³ supports diverse conditional and joint inference applications: sequence infilling, visual inpainting, multimodal alignment, and conditional control—all without retraining or loss of tractability. Moreover, practitioners can choose the mask and partition distributions in training to align with expected deployment queries.

Limitations remain:

• Extreme scaling ($(S_1, \ldots, S_K)$2B parameters, $(S_1, \ldots, S_K)$3 masked inference) is largely unexplored (Guo et al., 29 Apr 2025).
• Some intermediate states (few large sets) underperform the very best AR or MAR regimes on specific FID/ROUGE axes, suggesting A³ is not universally optimal for all snapshot choices (Liu et al., 2024).
• Efficient KV caching for arbitrary order remains an open architectural frontier.

Potential extensions include adaptive on-the-fly partitioning, hybrid continuous–discrete A³ via diffusion latents, and cross-modal A³ frameworks that intermingle text, image, and audio tokens (Liu et al., 2024).

---

Summary Table: A³ Key Aspects and Comparisons

Aspect	A³ (Any-Order Any-Subset AR)	Classical AR / MAR	Discrete Diffusion
Order/Subset Flexibility	Arbitrary permutation, arbitrary subsets	Fixed (AR), 2-group (MAR)	Arbitrary, but poor depth
Attention Masking	Groupwise causal (+two streams)	Causal (AR), full (MAR)	Full / masked
Parallelism	Full within each group, batched AR steps	1-step (AR), global (MAR)	Full, but conditional
Core Objective	Grouped NLL, progressive adaptation	Next-token, masked	Masked LR, diffusion loss
Decoding	Speculative, parallel, provably exact	Sequential, iterative	Diffusion-step, imprecise
Empirical Performance	SOTA or parity language/image/tabular	High (task-specific)	Often outperformed by A³

A³ represents a unified, extensible paradigm bridging probabilistic rigor, expressive dependency capture, and versatile conditional generation (Du et al., 19 Jan 2026, Guo et al., 29 Apr 2025, Liu et al., 2024, Shih et al., 2022, Pynadath et al., 17 Feb 2026).