Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding

TL;DR

This paper addresses the challenge of generating tokens in arbitrary orders while preserving the true joint distribution. It introduces AS-ARMs and a principled speculative decoding scheme (ASSD) that enables parallel generation with a provable guarantee of fidelity, bounded by the number of tokens predicted. The authors provide a training paradigm based on an XLNet-like architecture with a binary lattice mask to realize 2^N joint-factorizations, along with a density-estimation oracle and a draft-then-verify workflow. Empirically, AS-ARMs with ASSD achieve speedups over sequential decoding and match or surpass stronger baselines on infilling benchmarks and code-generation tasks with substantially smaller models. Overall, the work revives AS-ARMs as a promising approach for efficient, high-quality language generation.

Abstract

In arbitrary-order language models, it is an open question how to sample tokens in parallel from the correct joint distribution. With discrete diffusion models, the more tokens they generate in parallel, the less their predicted distributions adhere to the originally learned data distribution, as they rely on a conditional independence assumption that only works with infinitesimally small timesteps. We find that a different class of models, any-subset autoregressive models (AS-ARMs), holds the solution. As implied by the name, AS-ARMs can generate tokens in any order, and in parallel. Moreover, AS-ARMs support parallelized joint probability density estimation, allowing them to correct their own parallel-generated token distributions, via our Any-Subset Speculative Decoding (ASSD) algorithm. ASSD provably enables generation of tokens from the correct joint distribution, with the number of neural network calls upper bounded by the number of tokens predicted. We empirically verify that ASSD speeds up language generation, without sacrificing quality. Furthermore, we provide a mathematically justified scheme for training AS-ARMs for generation, and show that AS-ARMs achieve state-of-the-art performance among sub-200M parameter models on infilling benchmark tasks, and nearly match the performance of models 50X larger on code generation. Our theoretical and empirical results indicate that the once-forgotten AS-ARMs are a promising direction of language modeling.

Figures (4)

• Figure 1: Attention Masks in AS-ARMs: Red means that attention is not allowed, while green means that attention is allowed from the query to the key/value. The numbers on each token represent the generation order, where lower-numbered tokens come first. Panel \ref{['subfig:parallel_sample_attention']} shows how, by attending to the same conditioning ("This", "is"), the tokens "machine" and "learning" can be generated in parallel. They do not attend to themselves nor each other. Panel \ref{['subfig:density_estimation_attention']} shows how we can conduct one-step density estimation on a sequence, with a permuted causal-like attention mask. The mask enforces the order "is", "This", "learning", "machine", where each token only attends to those decoded before it.
• Figure 2: Probabilistic Graphical Model: Shows the discrete-time Markov chain for the forward noising process, and its time reversal (i.e., data generation). This justifies Equation \ref{['eqn:loss']}.
• Figure 3: Fixed (Recursive Binary Lattice) Versus Any Permutation Mask Decomposition: Validation loop metrics on generated sequences from each training strategy. The curves shown are for models trained with an effective batch size of $96$ across four NVIDIA RTX A4000 devices. Each validation iteration has $36$ sequences of length $512$ tokens.
• Figure 4: Narrow ($\mathbf{1\%\rightarrow10\%}$) Versus Wide ($\mathbf{1\%\rightarrow85\%}$) Prompting Rates: Validation loop metrics on generated sequences from each training strategy, as it relates to the distribution of prompt lengths in the train set. The curves shown are for models trained with an effective batch size of $320$ across five NVIDIA RTX 6000 Ada devices. Each validation iteration has $64$ sequences of length $512$ tokens from OpenWebText, where the task is to infill $95\%$ of the masked sequence given a $5\%$ prompt.

Theorems & Definitions (9)

• lemma 1
• Theorem 1
• Theorem 2
• proof
• proof
• proof
• proof
• Theorem 3
• proof