Efficient Autoregressive Inference for Transformer Probabilistic Models

AI summary

Overview

Research area: Machine learning — amortized probabilistic inference, meta-learning, and transformer-based set-conditioning models (neural processes, prior-fitted networks, tabular foundation models).

Technical level: Intermediate to Advanced. The paper assumes familiarity with transformer attention (self-attention, cross-attention, KV caching), autoregressive factorization of joint densities, and the neural process / prior-fitted network literature.

Scope: The paper introduces a "causal autoregressive buffer" that makes joint (multi-target) prediction from set-based transformer probabilistic models asymptotically cheaper — reducing cost from O(K(N+K)²) to O(N²+NK+K²) — and validates it on synthetic functions, EEG time series, a Bayesian model comparison task, and a tabular foundation model.

What This Paper Is About

Transformer models such as neural processes (NPs), prior-fitted networks (PFNs), and tabular foundation models are very good at producing marginal predictions — each target predicted independently — in a single forward pass. But many tasks need coherent joint distributions over several targets at once, and the usual fix (deploying the model autoregressively) forces the model to re-encode the entire context set at every step, which scales poorly. This paper's goal is to keep one-time context encoding while still capturing dependencies between generated targets, so that joint sampling and joint density evaluation become fast and memory-light.

Key Contributions

1. The causal autoregressive buffer. A mechanism that decouples one-time encoding of the static context from lightweight sequential prediction. The context is encoded once and cached; each new prediction attends to the cached context plus previously predicted targets held in a strictly causal buffer. This cuts the cost of autoregressive joint inference from O(K(N+K)²) to O(N²+NK+K²).
2. One-pass joint predictive density evaluation. Because the attention structure is fixed, all K autoregressive conditionals can be computed in a single masked forward pass rather than K sequential passes, and the model recovers exact standard set-based behavior when the buffer is empty (K=1).
3. A unified training strategy. Structured masked attention combined with a buffer-size curriculum lets a single model learn both marginal and autoregressive modes at minimal extra cost. During training, 50% of targets attend only to the context, and 50% attend to the context plus a buffer prefix of length v_m ~ Uniform{1, ..., K}.
4. Broad applicability demonstrated. The buffer is integrated into TNPs/PFNs and into the TabICL tabular foundation model, and is shown to be a general mechanism that also applies to Perceiver-style encoders with P << N latent (pseudo-)tokens, demonstrated on the latent bottlenecked attentive neural process (LBANP).

Main Findings

• Large speedups with comparable accuracy: Across synthetic functions, EEG, a Bayesian model comparison task, and tabular regression, the method closely matches full context re-encoding while delivering up to 20× faster joint sampling and density evaluation and up to 7× lower memory usage.
• Sampling speed: Autoregressive sampling is 3–20× faster than the fully autoregressive TNP-A and TNP-D-AR (Fig. 1, top left). TNP-D-Ind and TNP-ND are faster but cannot capture complex predictive dependencies.
• Density evaluation speed: Predictive density evaluation is on par with the highly parallel TNP-A and a factor of K× faster than the sequential TNP-D-AR.
• Training overhead is minimal: Training step time is comparable to the fastest baselines (TNP-D, TNP-ND) and 4–12× faster than TNP-A.
• Memory: The method requires 6–7× less VRAM than TNP-D-AR and TNP-A at large context sizes (N = 1024), because only a single context is cached independently of batch size.
• Gaussian process synthetic functions (average predictive density, M = 16 targets): TNP-D-AR 2.57 (0.020), TNP-D-Ind 2.22 (0.022), TNP-ND 0.80 (0.082), TNP-A 2.24 (0.018), TNP w/ buffer K=16 2.51 (0.019), K=1 2.56 (0.019).
• Sawtooth process: TNP-D-AR 1.05 (0.004), TNP-D-Ind 0.94 (0.005), TNP-ND -0.43 (0.008), TNP-A 0.98 (0.004), TNP w/ buffer K=16 1.00 (0.005), K=1 1.09 (0.004).
• EEG interpolation: TNP-D-AR 0.51 (0.013), TNP-D-Ind 0.36 (0.014), TNP-ND 0.46 (0.011), TNP-A 0.58 (0.014), TNP w/ buffer K=16 0.52 (0.013), K=1 0.54 (0.014).
• EEG forecasting: TNP-D-AR 1.07 (0.004), TNP-D-Ind -0.74 (0.008), TNP-ND -0.04 (0.005), TNP-A 1.23 (0.003), TNP w/ buffer K=16 0.85 (0.004), K=1 1.21 (0.003). The paper notes a larger gap versus TNP-D-AR for forecasting, which deploying with K=1 (no buffer) closes.
• Bayesian model comparison (multisensory causal inference, log marginal likelihood RMSE, lower is better): TNP-D-AR 3.10 (0.005), TNP-D-Ind 86.96 (0.000), TNP-ND 208.51 (0.041), TNP-A 4.75 (0.012), TNP w/ buffer K=16 3.56 (0.004), K=1 3.47 (0.004). The ΔLML RMSE values are 2.44 (0.008), 36.18 (0.000), 25.60 (0.023), 3.29 (0.019), 2.60 (0.010), 2.59 (0.011) respectively.
• Bayesian model comparison (average log predictive density, higher is better): all methods except TNP-ND perform similarly — TNP-D-AR -2.76 (0.024), TNP-D-Ind -2.77 (0.025), TNP-ND -3.12 (0.016), TNP-A -2.76 (0.024), TNP w/ buffer K=16 -2.76 (0.024), K=1 -2.76 (0.024).
• Model selection outcome: The model trained with ρ = 4/3 generally achieves higher LML than ρ = 1, aligning with the original finding that participants remap their auditory space to match the visual range. Ground truth was estimated with S-VBMC (Acerbi et al. 2018; Silvestrin et al. 2025).
• LBANP experiment: Applying the buffer to the latent bottlenecked attentive neural process (LBANP) yielded higher predictive densities than standard autoregressive inference with LBANP; the paper attributes this to the buffer allowing conditioning on both the latent summary and the explicit history of previous points.
• Tabular foundation model: A model integrating the buffer into TabICL was pre-trained on 10.24 million synthetic datasets with 1 to 10 features and 8 to 1024 context points, using buffer size K = 32, and evaluated on six UCI and Kaggle datasets with 16 random context/target splits per dataset for N = 16 and N = 1024 context sizes and M = 32 targets. The specific accuracy numbers for this experiment are not present in the content provided.

Methodology in Plain English

The problem in one picture. A set-based model reads a "context" of observed input-output pairs and summarizes it with bidirectional attention. To predict several targets jointly, the standard approach appends each new prediction back into the context and re-runs the whole encoder. Every appended point invalidates every cached representation, so K predictions cost roughly K times the cost of a context of size N+K.

The fix. The authors split the conditioning information into two parts with different rules:

• The context is treated as immutable. It is encoded once, and its keys and values are cached and never recomputed.
• The buffer holds previously predicted targets. It uses a strictly causal attention pattern: a buffer token can see the cached context and only earlier buffer tokens, never later ones. Nothing ever writes back into the context.

Training it. Each training task is a dataset split randomly into three disjoint pieces: a context, a buffer, and a set of targets (N_tot = N + K + M). The buffer is randomly ordered, and a single structured attention mask determines, for each target, whether it sees the context alone or the context plus a variable-length buffer prefix of length v_m drawn uniformly from {1, ..., K}. Half of the targets are context-only. This single forward pass computes the loss, and the training objective is the expected negative log-likelihood under this scheme. For any fixed conditioning set, minimizing that negative log-likelihood is equivalent to minimizing the KL divergence between the model and the true posterior predictive distribution.

Using it. At test time there are two modes. For sampling, the model prefills (encodes) the context once in O(N²), then decodes one target at a time, appending each sampled pair to the buffer and updating only the buffer's key/value cache. For density evaluation, K buffer tokens holding observed target values and K query tokens are packed into one masked forward pass, producing all K conditionals simultaneously — the same joint log-density as sequential evaluation, but in one pass. Because autoregressive densities depend on order, the authors average over multiple buffer orderings to approximate permutation invariance.

Batched sampling. A single context prefill of O(N²) is shared across B parallel sampling streams; only the small buffer keeps separate state per sample. Total cost becomes O(N² + B(NK + K²)) instead of the naive O(BK(N+K)²).

Fair benchmarking. All methods were run in a unified codebase on a single NVIDIA L40S GPU, with baselines optimized beyond their public versions using KV caching, FlashAttention-2 (Dao, 2023), and compilation — yielding 3–10× speedups over original implementations. A custom Triton kernel was written to optimize memory access during batched sampling. Benchmarks used buffer size K = 16; sampling and density evaluation used M = 16 targets and batch size B = 256; training used M = 256 targets and batch size B = 128.

Why This Matters

Joint predictive distributions are what make a model's outputs coherent rather than a bag of disconnected guesses. This paper shows that the standard autoregressive route to coherence — re-encoding the context at every step — is unnecessary. By keeping the context immutable and confining sequential dependency to a small causal buffer, it preserves the flexibility of set-based conditioning while borrowing the key-value caching efficiency that made autoregressive language models practical. It also removes a real deployment barrier: batched joint sampling for many samples from the same context now shares one context cache.

Real-world applications (all drawn from the paper's own experiments):

• EEG time series — interpolation and forecasting across 7 correlated channels, where joint coherence across channels and time matters.
• Tabular regression — pre-training a small-scale tabular foundation model on synthetic structural-causal-model data for multi-column prediction.
• Computational neuroscience model comparison — computing log marginal likelihoods over 400 experimental trials per dataset for an audio-visual localization model with 7 free parameters, to distinguish between two model variants (ρ = 1 vs ρ = 4/3).
• Environmental and infrastructure monitoring — the paper's tabular evaluation includes the Individual Household Electric Power Consumption, Gas Turbine CO and NOx Emission, Bike Sharing, Jena Climate, Power Consumption of Tetouan City, and California Housing Prices datasets.

Industry relevance. Any deployment that needs many joint samples or joint density evaluations per context — probabilistic forecasting, uncertainty-aware tabular prediction, posterior approximation — benefits directly from 3–20× faster sampling and 6–7× lower VRAM at large context sizes. The mechanism is architecture-agnostic: the paper applies it not only to TNPs/PFNs but also to tabular foundation models (TabICL) and Perceiver-style latent-bottleneck encoders (LBANP).

Future Directions

• Combining with diffusion and flow-matching. The paper notes that generative modeling's two main paradigms — diffusion/flow matching and autoregressive transformers — have recently been combined (Tang et al. 2025; Arriola et al. 2025; Wu et al. 2025) and that this suggests future extensions of the buffer. Related work on KV caching for masked diffusion (Sahoo et al. 2026) is cited as a parallel direction.
• Order-agnostic and permutation-invariant variants. Because the autoregressive factorization fixes an order and breaks permutation invariance, the authors average over multiple buffer orderings. Section H.3 analyzes how the number of orderings affects estimate stability, but the general problem of efficient any-order joint inference remains open.
• Scaling to larger foundation models. The tabular demonstration pre-trained on 10.24 million synthetic datasets with a buffer size of K = 32 and a network size comparable to the original TabPFN. Whether the same buffer mechanics hold at larger pretraining scales is not established in the paper.
• Improving rather than merely matching full re-encoding. The LBANP result — where the buffer outperformed standard autoregressive inference — hints that conditioning on an explicit history plus a latent summary can be better as well as faster. Understanding when the buffer is not just an approximation but an improvement is left open.

Target Audience

Researchers and practitioners working on neural processes, prior-fitted networks, and tabular foundation models who need joint rather than marginal predictions. It is also relevant to anyone building amortized Bayesian inference or meta-learning systems where the same context must be queried repeatedly for many joint samples or density evaluations, and to engineers concerned with inference latency and GPU memory at large context sizes. Readers should already be comfortable with transformer attention mechanics and the autoregressive factorization of joint densities; the asymptotic and training-mask details assume that background.

Accuracy note: The provided paper content includes the TabICL tabular experiments' setup but is truncated before their numeric results table, so no tabular accuracy figures are reported here.

Authors’ abstract

Set-based transformer models for amortized probabilistic inference and meta-learning, such as neural processes, prior-fitted networks, and tabular foundation models, excel at single-pass marginal prediction. However, many applications require joint distributions over multiple predictions. Purely autoregressive architectures generate these efficiently but sacrifice flexible set-conditioning. Obtaining joint distributions from set-based models requires re-encoding the entire context at each autoregressive step, which scales poorly. We introduce a causal autoregressive buffer that combines the strengths of both paradigms. The model encodes the context once and caches it; a lightweight causal buffer captures dependencies among generated targets, with each new prediction attending to both the cached context and all previously predicted targets added to the buffer. This enables efficient batched autoregressive sampling and joint predictive density evaluation. Training integrates set-based and autoregressive modes through masked attention at minimal overhead. Across synthetic functions, EEG time series, a Bayesian model comparison task, and tabular regression, our method closely matches the performance of full context re-encoding while delivering up to $20\times$ faster joint sampling and density evaluation, and up to $7\times$ lower memory usage.