---
title: Any-Order Any-Subset AR Modeling (A³)
url: https://www.emergentmind.com/topics/any-order-any-subset-autoregressive-modeling-a
type: topic
---

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

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 [2410.10511], [2601.13228]. 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 $K = 2$ and $S_2$ is the set of masked positions. Any permutation and any grouping of indices are permissible, making A³ inherently order- and subset-agnostic [2205.13554], [2602.16092].

This flexibility supports conditional inference for arbitrary masks $S$:
\[
p(x_S \mid x_{\bar S}) = \prod_{t=1}^{|S|} p(x_{\sigma(t)} \mid x_{\bar S}, x_{\sigma(<t)})
\]
where $\sigma$ is an ordering of $S$. 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 ($H_c$):** Aggregates information from all tokens in current/previous groups.
- **Query stream ($H_q$):** Attends only to content from strictly earlier groups, preventing access to current target tokens [2601.13228], [2410.10511].

Attention masks for both streams enforce the group-wise autoregressive order:
\[
A_{i, j} =
\begin{cases}
0, & i \leq j \text{ (allow past/self-attend)} \\
-\infty, & i > j \text{ (prohibit future-attend)}
\end{cases}
\]
This separation addresses the structural–semantic tradeoff: $H_c$ can summarize all prior generative context, while $H_q$ specializes in predicting new tokens, avoiding entanglement that would arise in a single-stream setup [2602.16092]. 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 [2410.10511].

## 3. Training Objectives and Protocols

Training A³ models proceeds by sampling random permutations $\pi$ and partitions $(S_1, \ldots, S_K)$ per minibatch, instantiating the full spectrum from left-to-right AR to masked/diffusion-like objectives. The primary loss is the negative log-likelihood:
\[
\mathcal{L}(\theta) = - \mathbb{E}_{\pi, \{S_k\}} \sum_{k=1}^K \sum_{i \in S_k} \log p_\theta(x_i \mid x_{S_{<k}})
\]
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 $s > 1$, and finally permute orderings within and across groups [2601.13228].

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 [2205.13554]. 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 [2504.20456] efficiently samples $k$ tokens in parallel:
- **Draft phase:** Propose $k$ tokens independently under partial conditioning.
- **Verification phase:** Recompute true joint conditional for each draft.
- **Rejection/correction phase:** Accept with probability $\min(1, q_i/p_i)$; otherwise, resample and rerun.

Correctness is guaranteed by induction, and the number of model calls is $O(N)$, never exceeding the number of tokens. This speculative parallelism yields up to $10\%$ speed-up in wall-clock generation and strictly preserves sample quality (e.g., identical perplexity to sequential decoding on WikiText infilling tasks) [2504.20456].

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 [2410.10511].

## 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 [2601.13228]. 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 [2410.10511].
- **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 [2205.13554].

The flexibility of partition size $K$ 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 [2602.16092].

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 ($K$ selection, dynamic partitioning) can further improve tradeoff control [2410.10511].

## 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 ($\gg$B parameters, $90\%+$ masked inference) is largely unexplored [2504.20456].
- 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 [2410.10511].
- 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 [2410.10511].

---

**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 [2601.13228], [2504.20456], [2410.10511], [2205.13554], [2602.16092].

Source: https://www.emergentmind.com/topics/any-order-any-subset-autoregressive-modeling-a