Efficient Autoregressive Inference for Transformer Probabilistic Models
Authors: Conor Hassan
Nasrulloh Ratu Bagus Satrio Loka
Cen-You Li
Daolang Huang
Paul Edmund Chang
Yang Yang
Francesco Silvestrin
Samuel Kaski
Luigi Acerbi
Venue: ICLR 2026 Poster (2026)
Keywords: probabilistic machine learning
neural processes
probabilistic meta-learning
amortized inference
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We accelerate autoregressive inference of transformer probabilistic models such as prior-fitted networks and transformer neural processes.

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.

Published as a conference paper at ICLR 2026
EFFICIENT AUTOREGRESSIVE INFERENCE TRANSFORMER PROBABILISTIC MODELS
FOR
Conor Hassan 12, Nasrulloh Loka Cen-You Li Daolang Huang" Paul E. Chang? Yang Yang', Francesco Silvestrin", Samuel 1.55. Luigi Acerbi
*Department of Computer Science, Aalto University, Finland EELLI Institute Finland
*Department of Computer Science, University of Helsinki, Finland *Verda
'Department of Computer Science, University of Manchester, UK conor. hassaneaalto...
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 pass marginal prediction. However, many applications re- quire joint distributions over multiple predictions. Purely autoregressive architec- tures 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 causal aufore- gressive buffer that combines the strengths of both paradigms. The model encodes the context once and caches it; 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 ef- ficient 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 20x faster joint sampling and density evaluation, and up to 7x lower memory usage.
INTRODUCTION
Recent advances in amortized probabilistic inference and meta-learning have produced set-based conditioning models that rapidly adapt to new tasks without retraining. Prominent examples include neural processes (NPs; Garnelo et al. 201 8a; Foong et al. 2020), their wansformer-baseee extensions (Nguyen & Grover, 2022; Chang et al., 2025), prior-fitted networks (PFNs; Müller et al. 2022), and tabular foundation models (Hollmann al., 2023; 2025; Qu et al., 2025; 2026; Zhang al., 2025). These methods share a key architectural principle: they process variable-sized context sets through conmuutationinvvriian encoders that respect the exchangeability of observed data, enabling accurate marginal predictive distributions over new target variables in a single forward pass.
While these models highly efficient for marginal predictions, many applications require co- herent joint distributions over multiple targets. Tasks such as signal interpolation, behavioral data modeling, and multi-column tabular prediction demand capturing dependencies across target vari- ables. The standard approach deploys these models autoregressively at test time (AR; Bruinsma al. 2023): to generate K predictions, the k-th step conditions on the initial context plus all k-1 previous predictions. Since the permutation-invarian encoder uses bidirectional self-attention over all inputs, each appended prediction triggers complete re-encoding of the entire augmented set, yielding O(K (N +K)?) complexity. This severely limits applications with large contexts (N), long target sequences (K), or frequent sampling requirements. Advances in efficient attention (Jacgle al., 1021; Feng et al., 2023) can reduce costs for encoding large static contexts but do not address the repeated recomputation inherent in autoregressive expansion of the conditioning set.
Shared first co-authorship

---

Published as a conference paper at ICLR 2026
Sample generation
Density evaluation
Training step
10
10
Time (s)
Time (s)
Time (s)
10-2
10-3
10-1
10-3
10-1
1.28
912
236
N
NV
N
Memory
Predictive performance
Peak VRAM (GB)
10P
deil diarity tpen
Normalized score
1.0
Ours
TNP-D-AR
10-1
giail
TNP-A
0.5
TNP-D-Ind
10-2
TNP-ND
0.0
128 512
1024
Oun
TNP-D-AR
TNP-4
TNP-D-frd D-tud
TNP-ND
Figure 1: Our method closely matches the predictive performance of expensive joint-prediction methods while delivering significant speedups and lower memory usage. We compare against transformer-basee probabilistic models spanning the expresssiity tradeoff (see Section for details). Our approach reduces the cost of autoregressive joint inference from O(K(N+K)) K))) to yielding up to 20x faster sampling and density evaluation, and 7 x lower peak memory. Panels show wall-clock time (log scale) for (a) sampling (top left), (b) density evaluation (top middle), (c) a training step (top right), (d) peak memory usage versus context points N (bottom left), and (e) normalized predictive performance averaged across six tasks (bottom middle).
To address this limitation, we introduce the causal autoregressive buffer, an architectural mechanism that decouples the expensive encoding of the static context from lightweight sequential prediction. Inspired by the efficiency of purely autoregressive architectures language modeling (Brown al., 2020) and image generation (Chen et al., 2020; Li et al., 2024), our buffer implements causal attention pattern for managing dependencies among generated targets but crucially, it operates alongside the set-based context rather than replacing it. Our approach encodes the initial context once and caches its representation. Targets added to the buffer attend to both the static con- text cache and previously predicted targets through causal masking, capturing dependencies among newly generated samples without requiring context reconic, reducing cost from O(K +K)) K2). When the buffer is empty, the model reduces standard set-based predic-
to O(N2 tor, preserving marginal prediction quality. A unified training strategy using masked attention and buffer-size curriculum allows a single model handle both efficient marginal predictions and ac- celerated autoregressive sampling and joint predictive density evaluation with substantial speedups and comparable predictive accuracy to standard AR deployment (Fig. 1).
Our main contributions are:
1. We introduce the causal autoregressive buffer, which reduces the cost of autoregressive joint inference for transformer-based amortized probabilistic models from to ( K by decoupling one-time context encoding from sequential predic- tion, enabling both efficient autoregressive sampling and one-pass joint predictive density evaluation while recovering exact standard model behavior when the buffer is empty.
2. We propose a unified training strategy, combining masked attention with a buffer-size cur- riculum, that enables single model to learn both modes of operation at minimal cost.
3. We demonstrate applicability across transformer- based probabilistic models including TNPs/PENs (Nguyen & Grover, 2022; Müller et al., 2022) and tabular foundation mod- els (TabICL; Qu et al., 2025), achieving up to 20 faster joint sampling and 7x lower memory while matching the predictive accuracy of full context re-encoding across tasks.
2

---

Published as a conference paper at ICLR 2026
2 PRELIMINARIES
We consider the task of predicting outputs pairs. Given a context set C = {(xn, yn) n=1 with N input-output pairs and a larget set T (x ) we seek to learn predictive distribution (o xi:Mi C) where the model's learnable parameters (Foong et al., 2020). This setup arises in meta-learning, few-shot prediction, and tabular foundation models any setting where single model must condition on variable-sized observed data without task-specific retraining.
inputs, given of observed input-output
Transformer diagonal prediction maps. Transformer architectures (Vaswani al., 2017) are natural fit for this set-based task. Methods such as (diagonal) transformer neural processes (TNPs; Nguyen Grover, 2022) and prior-fitted networks (PFNs; Müller et al., 2022) use two core attention mechanisms. First, the model processes C using multi-head self-attention (MHSA). Then, each target input x queries this summary using multi-head cross-attention (MHCA). This structure leads to an efficient diagonal predictive model where predictions are conditionally independent:
M
(V Xi.MiC)
= II (v m=1
(igi(X rc(C)))
(1)
Here, re(C) is the permutation invvariant summary of the context produced by the MHSA layers', and rig(("-)) :) the final decoding function that produces a parametric prediction for you via MHCA. This may consist of single Gaussian, but more expressive parameterizations include Riemannian distributions (Müller et al., 2022) and mixtures of Gaussians (Uria et al., 2016; Chang et al., 2025). These models are efficiently trained via ciddmmum-iikeihoooo random context-targets data splits.
Autoregressive sampling and predictive density evaluation. Many applications require joint distributions over targets, both for generating coherent joint samples and for evaluating joint pre- dictive densities. While Eq. (1) can be extended with parametric multivariate densities such as a full Gaussian (Markou et al., 2022; Nguyen & Grover, 2022), more flexible alternative applies autoregressive factorization of the same model (Bruinsma et al., 2023). We switch from index m to k (and M to K) emphasize that targets are now processed in specific sequential order:
K
po(vi. C).K:C)=
II
po (
x CU (03,,,)))))))
Crucially, Eq. (2) does not define a new model; a mode of deployment for models described by Eq. (1).2 Each prediction is conditioned on all previously predicted targets, capturing inter- target dependencies. This flexibility comes at a cost: because the permutation-invaaiann encoder uses bidirectional self-attention, adding single predicted target to the conditioning set invalidates all cached representations, foreing a complete re-encoding of the augmented set. The resulting complexity is (KNNKK++*) and generating B parallel sample sequences requires B independent forward- pass streams, each maintaining its own growing context.
Our goal is improve efficiency for both sequential and parallel sampling and predictive density evaluation by encoding the context once and reusing it throughout. Existing autoregressive up- date schemes break this caching: when targets join the conditioning set, the context representation must be recomputed. Our key insight is to separate the roles of initial context C. and predicted tar- gets {(x5 j<k. We preserve permutation invariance for the initial context (encoded once and cached) while handling target dependencies through separate causal mechanism. When needed, the buffer can be merged back into the context to restore full permutation invariance. This selec- tive relaxation in-between fully set-based and purely autoregressive models enables efficient sequential and parallel operations while maintaining the strengths of set-based conditioning.
'See Appendix H.1 for evidence of the impact of permutation invariance for the context set. 2 Eq. (2) imposes specific factorization order While fixing the order can be valid modeling choice under certain circumstances, this breaks permutation invariance. In cases where permutation invariance is required, Monte Carlo approximation he obtained by averaging over multiple target orderings.
3

---

Published as a conference paper at ICLR 2026
Recompute KV cache using self-attention
Autoregressive Buffer (Ours)
Sampling step
((N++)))
KV
KV (Frozes)
concett KV (X))
Densily evaluation
KV Croesss
plall
Atrach buiffirs come
Sorward
RV
RV (roosss
xj
plagl
KV (hosecaa)
p(ocl))
Training mask
K/V
Figure 2: Standard autoregressive (AR) method compared to the proposed causal AR buffer. Top (sampling): Standard AR (left) appends each new prediction (xī. to the context set and re- computes self-attention (SA) over the growing set at each step, with per-step cost ((N + k)))). Our causal AR buffer (right) caches the context KV once and appends each new buffer token to the frozen KV cache, reducing attention cost. Middle (density evaluation): Standard AR requires K forward passes to compute the conditionals p(u because the context set changes at each step (left), whereas the AR buffers (right) enables computing all K conditionals in single pass. Bot- tom (training mask): (left) Standard AR training mask: context self-attention (blue) and target-to- context attention (purple); (right) our training mask: buffer tokens attend the context and one another causally (red). Target tokens are either attend only to the context, or to the context plus a variable number of buffer tokens (purple).
3 EFFICIENT AUTOREGRESSIVE INFERENCE
Core contribution. Our method conditions predictions on static, task-defining context C and dynamic autoregressive buffer B (see Fig. 2). We parameterize the predictive distribution as
K
polyi:x xxixk C)=
II po(yt Fige(xst [rc(C), b1xx-1))), b = re (xt+ vi), [Ie(C), b1:k-1)), k=1
(3) where rg is the buffer encoder, implemented via causal MHSA among buffer tokens and cross- attention to the cached context at each layer; bıst denotes the first & encoded buffer representations; and D1:0 = 0. Crucially, rc(C) computed once and cached. The target decoder rigi cross-attends to the concatenated keys/values from the cached context and the visible buffer prefix, then passes the result through distribution head (e.g., an MLP parameterizing a mixture of Gaussians). When the buffer is empty, Eq. (3) reduces to the standard diagonal prediction map (Eq. (1)).
To couple one- mue-time set-bascd encoding with sequential dependence, the attention must satisfy four requirements: (R1) the context is immutable: encoded once and cached read-only; (R2) the buffer
K denotes the total number of AR targets: the buffer stores up to K -1 previous predictions. Thus, "buffer size 1=16" means at most 15 buffered elements and 16 total predictions. K=1 corresponds to standard AR inference (empty buffer).
4

---

Published as a conference paper at ICLR 2026
is strictly causal: buffer token may attend only to the cached context and positions j; (R3) information flows out of the context but never back: no edges write into C; and (R4) each target attends to the cached context and visible buffer prefox; targets do not attend to one another.
During training, we enforce (R4)-(4) in a single forward pass using a structured attention mask. We implement this with a transformer backbone that processes context, buffer, and target tokens with distinct role embeddings; buffer tokens carry learned positional embeddings indicating their autoregressive order (see Appendix H.2). This enables parallel loss computation by conditioning each target's 's prediction on the context and a variable-sized, ground-truth buffer prefix. At inference, we use two-stage process: a one-time context encoding followed by prediction as either sam- pling or predictive density evaluation. The total attention cost is O + KN + K²): one-time O(2)) for context self-attention, O K N) for cross-attention reads from the cache, and O (K²) for causal buffer self-attention and target-to-buffe reads. By contrast, naive autoregressive methods cost O(K(N + K)²) due to repeated context recomputation. Architectural details appear in Appendix A.
The model is trained by minimizing the expected negative log-likelihood over a prior distribution of datasets P. Each training task is generated by sampling a dataset DD = {(x, y ~ P. We then split D into three disjoint sets using a random partition x: (1) the context set C = {(x,: yn 1 (2) the buffer ser B {(xx, yk ) and (3) the target ser T = {Xxm }m)]m= with No = N + K + M. We randomly order the buffer set B and compute all predictions for the target set T in single forward pass. A structured attention mask controls whether each target can attend to the buffer, and if so, how many elements: 50% of the targets attend only to the context C; 50% attend to the context plus prefix of the buffer B where Um ~ Uniform (I,, K} (see Fig. 2, bottom). The training objective is:
Training details.
M
||
(()) = ED-P ((C.B.T~<((P)))
log po (Jm Xm: C, (.BE))))
(4)
where Bi.... is the visible portion of the buffer for target m (Um = O for context-only targets). This training curriculum ensures robust performance regardless of buffer state at test time: context- only targets maintain marginal prediction quality, while variable-length buffer exposure teaches the model to flexibly incorporate additional conditioning information. For any fixed conditioning set, minimizing the negative log-likelihood is equivalent to minimizing the KL divergence between the model and the true posterior predictive distribution (Müller et al., 2022: Elsemüller al., 2024); our objective aggregates this across varying amounts of buffer information.
At inference, we support two modes: (i) autoregressive sampling. where the buffer grows incre- mentally with the model's own generated samples (rather than the ground-truth data used during training); and (ii) parallel joint predictive density evaluation, where we pack K buffer tokens con- taining observed target values and K query tokens to evaluate all K conditionals in one pass (see below). The attention structure is identical in both inference modes; only the execution differs (pre- fill followed by sequential updates for sampling, single masked pass for evaluation).
Autoregressive sampling. Given a context C and a sequence of target inputs xj. X we gen- erate samples by first performing one-time prefill of C, caching its keys and values in O(N2) operation. We then decode sequentially following Eq. (3): for each step & = 1,.... , K, we form a target query for input xt. attend to the cached context and causal buffer B sample y from the predictive distribution, and append (x,, v/a) to the buffer with its positional embedding (see Fig. 2, top). Only the buffer's key/value cache is incrementally updated, avoiding context recomputation and yielding ON - NK + K²) total complexity (detailed in Algorithm 1 in Appendix A.3).
Joint predictive density evaluation. Our framework also evaluates the joint predictive density of K K targets, (xt, = in single forward pass. Following the TNP-A variant of Nguyen & Grover (2022), we pack two sets of tokens into the model: (i) buffer tokens (x,, Mikk=1 and (ii) query tokens x for the same inputs. A causal attention mask ensures that each query for x attends to the context C and only the preceding buffer tokens B = {(x], )))) yielding:
K
K
log (v(.. xxixx:C) =
Σ
log logpolui x C, BLA-1).
(5)
5

---

Published as a conference paper at ICLR 2026
This yields the same joint log-density sequential autoregressive evaluation but executes in a sin- gle forward pass with total attention cost +KN*+KN+K*) (see Fig. 2, middle). For a detailed description of the algorithm, see Algorithm 2 in Appendix A.3. Notably, autoregressive predictive density estimates are order-dependent: to recover approximate permutation invariance, we average the predictive densities over multiple buffer orderings (Murphy et al., 2019). See Appendix H.3 for an analysis of how the number of buffer orderings affects estimate stability.
Batched autoregressive sampling. Our method is particularly efficient for autoregressively gener- ating multiple samples in a batch, conditioned on the same context C (e.g., multiple joint predictions for the same observed function values- see Fig. 2, top). naive batched autoregressive approach must re-encode growing context set at every generation step for each of the B samples. To generate B samples of length K, this results in prohibitive total cost of O(BK(N + K)"). In contrast, our approach performs the expensive context prefill (O(N3) only once. This single context cache is then efficiently reused across all B batched generation streams, with only the small, dynamic buffer maintaining separate state for each sample. This reduces the total cost to O(N + B(NK + K3)), making batched sampling practical even for large contexts and batches.
Architectural generality. Our buffer a general mechanism applicable other transformer vari- ants. For instance, Perceiver-style encoder (Jaegle et al., 2021) summarizes the context C into a fixed set of P < N latent tokens, also known as pseudo-toker tokens (Lee et al., 2019; Feng et al., 2023; Lara-Rangel et al., 2025). We can precompute the latent key/value representations once autore- gressive decoding then requires attending only to these P latents and the growing causal buffer. The per-layer attention cost is O(N
for the prefill and (PK+K3) for decoding K samples. P22 To demonstrate this, we applied our buffer
Without our buffer, the cost is to the latent bottlenecked attentive neural process (LBANP: Feng et al., 2023) architecture, TNP variant that encodes the context with pseudo-tokens We found that our buffer yields higher predic- tive densities than standard autoregressive inference with LBANP, likely because the buffer allows the model to condition on both the latent summary and the explicit history of previous points, rather than relying on the compressed latents alone. See Appendix H.4 for results and additional details.
4 RELATED WORK
Transformer probabilistic models. Our method serves modular component within neural processes (NPs; Garnelo et al., 201 8b;a; Bruinsma et al., 2021; Nguyen & Grover, 2022; Dutordoir et al., 2023; Chang et al., 2025) and prior-fitted networks (PFNs; Müller et al., 2022; 2023; Holl- mann et al., 2023). Prior work on efficient NP architectures has focused on improving scalability with respect context set size (Feng et al., 2022; 2023) and reducing memory usage (Feng et al., 2024) for independent predictions. By contrast, our method targets efficiency in autoregressive joint sampling and evaluation a largely unexplored area and is complementary to these architectural improvements. More broadly, cessformer-based probabilistic models increasingly frame Bayesian inference as in-context learning, performing tasks such as posterior approximation and predictive density estimation by conditioning on context observations (Mittal et al., 2023; Gloeckler et al., 2024; Reuter et al., 2025; Chang et al., 2025; Whittle et al., 2025; Mittal et al., 2025). The effective- ness of this paradigm has led to tabular foundation models such as TabPFN (Hollmann et al., 2023; 2025) and TabICI. (Qu et al., 2025; 2026), which perform in-context learning over table rows using the same attention mechanisms standard TNPs and PFNs. Our buffer integrates naturally with all these models.
Autoregressive joint density estimation. Autoregressive approaches are widely used for joint density estimation, from neural autoregressive density estimators (Larochelle & Murray, 2011; 1; Uria et al., 2016; Germain et al., 2015) to normalizing flows (Kingma et al., 2016; Papamakarios 2017; Huang et al., 2018; De Cao et al., 2020; Patacchiola et al., 2024) and order-agnostic autore- gressive models (Uria et al., 2014; Hoogeboom et al., 2022; Liu et al., 2024). Most directly related is the Autoregressive Transformer NP (TNP-A; Nguyen & Grover, 2022), which duplicates targets into queries and observed values for both training and inference. We recognize that this duplication is only needed for predictive density evaluation, not training. Bruinsma et al. (2023) showed that deploying standard NPs autoregressively improves joint predictions but requires expensive context re-encoding at each step. Our buffer combines insights from both approaches: like TNP-A, we en-
al.,

---

Published as a conference paper at ICLR 2026
able parallel predictive density evaluation, and like Bruinsma et al. (2023), model autoregressive dependencies while training on independent targets, but our separate buffer architecture avoids both TNP-A's training overhead and the re-encoding bottleneck.
Connection to other generative modeling techniques. Modern generative models follow two main paradigms: diffusion and flow-matching models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song al., 2021; Lipman et al., 2023) that generate samples through continuous-time dynamics, and autoregressive transformers (Radford et al., 2018; Brown et al., 2020) that generate sequences token-by-token with cached key-value states. Our buffer brings the caching efficiency of autore- gressive transformers to NPs and PFNs: encode the set-based context once (analogous to prompt) and generate efficiently through cached representations, rather than re-encoding the entire context at each autoregressive step. Recent work has shown that these paradigms can be combined (Tang et al., 2025: Arriola et al., 2025; Wu et al., 2025), suggesting future extensions. Recent work has also begun unifying masked diffusion (Lou et al., 2024; Shi et al., 2024; Sahoo et al., 2025) with any-order autoregressive models. Sahoo et al. (2026) exploit the connection between masked dif- fusion models and any-order autoregressive models to introduce KV caching for masked diffusion, enabling sigmificant inference speedups while preserving parallel generation. In a related vein, Par- tition Generative Models (PGMs; Deschenaux et al., 2025) employ self-attention among observed tokens and cross-attention from unmasked to masked tokens mirroring the context-targe target attention pattern in TNPs and PFNs, and pointing to structural convergence between these model families.
EXPERIMENTS
We evaluate the buffer's efficiency gains and predictive fidelity across four domains: regression on synthetic functions, interpolation of real-world EEG data, Bayesian model selection multisen- sory perception model, and pre-training of a tabular foundation model. We first conduct wall-clock and memory benchmarks, then assess predictive performance across these varied domains.
Baselines. We compare against models spanning the efficiency-expressiiee tradeoff, all config- ured with matched parameter counts, same input embeddings and output prediction heads unless noted otherwise (details in Appendix B). TNP-D (Nguyen & Grover, 2022) assumes conditional independence between targets; we evaluate it both with standard parallel decoding (TNP-D-IM) fast but limited) and with autoregressive deployment (7NP-D-AK -D-R, expressive but requires sequential re-encoding). TNP-ND models target dependencies via multivariate Gaussian, enabling one-pass joint predictive density evaluation but limiting expressivity. TNP-A uses causal self-attention for full autoregressive modeling but suffers from slow sequential sampling and high training cost. Ad- sampliiipi ditional task sspcific baselines are introduced as needed. TNP-ND aside, all models use a Gaussian mixture model output head with 20 mixture components unless stated otherwise.
Computational efficiency. Our method trades exact set-based AR updates for efficiency. We benchmark wall-clock time for: (i) autoregressive sampling, (ii) joint predictive density evalua- tion, and (iii) full training step (forward and backward pass), as well as (iv) peak memory usage. All measurements use a unified codebase and run on a single NVIDIA L40S GPU. We optimized all baselines beyond their public versions with KV caching, F (Dao, 2023), and compi- lation, achieving 3 - 10x speedups over original implementations to ensure fair comparison. For our method, we developed custom Triton kernel to optimize memory access during batched sam- pling (details in Appendix C). Benchmarks in Fig. 1 use model architectures matching subsequent experiments with buffer size K 16. For sampling and predictive density evaluation: M 16 targets, batch size B 256. For training: M 256 targets, batch size B = 128.
Our method achieves a superior efficiency profile compared to expressive baselines. For autoregres- sive sampling (Fig. 1, top left), our method is 3 20x faster than the fully autoregressive TNP-A and TNP-D-AR. While TNP-D-Ind and TNP-ND are faster, they cannot capture complex predictive dependencies, as shown later in this section. For predictive density evaluation (Fig. 1, top center), our method's 's speed is on par with the highly parallel TNP-A and is a factor of K x faster than the sequential TNP-D-AR For training speed (Fig. 1, top right), the overhead of our method is minimal, resulting in a training step time comparable to the fastest baselines (TNP-D, TNP-ND) and - 12x faster than TNP-A, which incurs significant computational cost due to its architecture. For memory
7

---

Published as a conference paper at ICLR 2026
Table 1: Average predictive density (1) results on synthetic functions and EEG example. Mean and (SEM) over various functions and context sizes N. for M = 16 targets. See Appendix D.5 for evaluation details and Table A2 for results with larger M. TNP w/ buffer with K=1 tracks the best method; 1-16 (fast) incurs only a slight drop in most cases.
TNP-D
TNP-ND
TNP-A
TNP w/ buffer (ours)
AR
Ind
K=16 (fast)
K=1 (slow)
GP
2.57 (0.020)
2.22 (0.022)
0.80 (0.082)
2.24 (0.018)
2.51 (0.019)
2.56 (0.019)
Sawtooth
1.05 (0.004)
0.94 (0.005)
-0.43 (0.008)
0.98 (0.004)
1.00 (0.005)
1.09 (0.004)
EEG-Int
0.51 (0.013)
0.36 (0.014)
0.46 (0.011)
0.58 (0.014)
0.52 (0.013)
0.54 (0.014)
BEG-For
1.07 (0.004)
-0.74 (0.008)
-0.04 (0.005)
1.23 (0.003)
0.85 (0.004)
1.21 (0.003)
Table 2: Multisensory causal inference model comparison and prediction results. For model selection, we use two metrics: log marginal likelihood root mean -qean- error (LML RMSE) against ground-truth, and difference in LML between م = 4/3 and p 1, reported as RMSE (ALML RMSE). For data prediction, we report average predictive density estimates (Average LL) for M 16 targets, computed using the model selected by the model-selection task. See Ap- pendix D.5 for additional details and evaluations.
TNP-D
TNP-ND
TNP-A
TNP w/ buffer (ours)
AR
Ind
K=16 (fast)
K=1 (slow)
L. RMSE (L) 3.10 (0.005)
86,96 (0.000)
208.51 (0.041)
4.75 (0.012)
3,56 (0.004)
3.47 (0.004)
ALML RMSE (L) 2.44 (0.008) 36.18 (0.000)
25.60 (0.023)
3,29 (0.019)
2.60 (0.010)
2.59 (0.011)
Average LL (1) -2.76 (0.024) -2.77 (0.025)
-3.12 (0.016)
-2.76 (0.024)
-2.76 (0.024)
-2.76 (0.024)
usage (Fig. 1, bottom left), our method requires 7x less VRAM than TNP-D-AR and TNP-A at large context sizes (N = 1024), scaling efficiently due to only needing to cache a single context independent of batch size. For predictive performance (Fig. 1, bottom center), we show normalized scores* averaged across the six tasks presented in Tables 1 and 2; our method closely matches the expressive autoregressive baselines (TNP-D-AR, TNP-A) while substantially outperforming TNP- D-Ind and TNP-ND. We provide additional results, including benchmarks across wider range of batch and target sizes and memory usage comparison, in Appendix C.
Synthetic functions. We consider two prediction tasks: (i) functions drawn from Gaussian pro- cesses (GPs; Rasmussen & Williams, 2006) where the kernel rype is sampled from a set, and its hyperparameters, and (ii) a non--aussiaa sawtooth process with discontinuous derivatives. All mod- els are trained and evaluated on distinct draws from these processes (see Appendix D.2). Results: As shown in Table 1, TNP w/ buffer (K 16) achieves performance comparable to TNP-D-AR while providing substantial speedups (Fig. 1), indicating that the buffer captures relevant target de- pendencies even without full context re-encoding. To verify that buffer training does not degrade AR capability, we deploy our model with K = 1 (empty buffer), matching the performance of TNP-D-AR exactly.
Eleetroencephalogram (EEG) data. Following Markou et al. (2022) and Bruinsma et al. (2023), train TNPs on EEG time series data (Zhang et al., 1995). Each trial contains 256 regularly sampled measurements across 7 correlated channels. See Appendix D.2 for dataset details. We train on an interpolation setting and evaluate on both forecasting and interpolation tasks. Interpolation uses random context/target splits; forecasting uses the first N points as context and the next M as targets (Appendices D.2 and D.5). As shown in Table 1, our method with K=16 is comparable to TNP-D-AR for interpolation but shows a larger gap for forecasting; deploying with K=1 (no buffer) closes this gap. In both settings, our method substantially better than TNP-D-Ind and TNP-ND. Additional results (larger M; permutation effects in forecasting) are in Appendices E.2 and E.3.
*We compute normalized scores for each task by linearly rescaling the average log predictive densities so that the worst-performing method scores O and the best pert-rerforming method scores 1.
See Appendix H.5 for ablations different huffer sizes.
8

---

Published as a conference paper at ICLR 2026
Multisensory causal inference model comparison and data prediction. We evaluate method on a popular computational neuroscience model that determines how the brain combines sensory stimuli from different sources (Körding et al., 2007). Using publicly available data from an audio-visual localization experiment (Liu et al., 2025), we consider two model variants, each with 7 free parameters, differing only in their auditory recalibration parameter p E {1, 4/a), and evaluate two tasks: (1) Model selection. For each method, we train two models on simulators with = I and p 4/3. We use the trained models for the challenging task of computing the log marginal likelihood (LML) of real experimental data, which requires evaluating the joint likelihood (Murphy, 2012):
LML (OypP(NNNNNN
log p(ul
(6)
which is inherently an autoregressive prediction task, as each prediction conditions all pre- vious data points, so it is perfectly suited for our models. For each dataset, we estimate the ground-truth LML for both p 1 and p 1/s using S-VBMC, a method proven effective similar problems (Acerbi et al., 2018; Silvestrin et al., 2025). We report LML RMSE and ALML RMSE (the difference between model metries, useful for model comparison) in Table 2. (2) Data prediction. Using the model selected in (1), we predict outputs on the real dataset and report average log-predictive densities (Table 2). See Appendix D.3 for experimental details and Appendix D.5 for evaluation settings.
Results. We evaluate our method using data from the 15 participants of the original study, extracting two non-overlapping subsets of 400 experimental trials each (400 data points), re-
UML (TNP w/ buffer)
LML (TNP w/ buffer)
-914
30
R² 1.00
0.92
-1085
sulting in a total of 30 datasets. The model
pm4/3
trained with p 4/3 generally achieves bet-
p=1
ter (higher) LML than p = 1, aligning with the original finding that participants
-1256
30
-1256
-1085
-914
-30
remap- ping their auditory space to match the visual range (Liu al., 2025). Fig. shows that the LML and ALML approximations obtained with our method are remarkably close the ground-truth. Furthermore, our method per- forms on par with TNP-D-AR and outperforms all other baselines on model comparison (Ta- ble 2). All models except TNP-ND perform similarly on the data prediction task. For additional results, Appendix F.
LML (True)
UML (True)
Figure Multisensory causal inference model comparison versus ground-truth. (Left) Log marginal likelihood (LML) comparison for both and p 4/3. (Right) LML difference (p = 4/3 p = 1) comparison. Our method
closely aligns with the ground-truth.
Small-scale tabular foundation model. We integrate our autoregressive buffer into the TabICI. foundation model architecture (Qu al., 2025). While the original work focused on classification, we pre-train our model from scratch for regression tasks. We reuse TabICL's set encoder to compute feature embeddings upfront and focus modifications on the final dataset-wise in-context learning transformer. Our core methodological contribution is the buffer mechanism, implemented by structured attention mask. This allows the model to condition on its recent predictions by storing them in a dynamic buffer. We pre-train this architecture model (SCM) prior (Hollmann et al., 2023; Qu et al., 2025), where each training instance is formed by partitioning datasets into distinct sets of context, buffer, and target points. Our network size and training scale are comparable to the original TabPFN (Hollmann et al., 2023); the model is pre- trained on 10.24 million synthetic datasets containing 1 to 10 features and 8 to 1024 context points, with a buffer size of K = 32. Full details are provided in Appendix D.4.
synthetic data from a structural causal
Results. We evaluate on six UCI and Kaggle d Individual Household Electric Power Con- sumption, Gas Turbine CO and NOx Emission, Bike Sharing, Jena Climate, Power Consumption of Tetouan City, and California Housing Prices. We form 16 random context/target splits per dataset for both N = 16 and N = 1024 context sizes and M-32 targets under interpolation (Int) and forecast- ing (For) tasks; the latter for time-series datasets only (all excluding California Housing Prices). We compare three inference modes from the same pre-trained model: "Ind" (independent predictions),
'UCI: WtCl:httts:///chhhi ics.uci edu/. Kaggle: https: //wp://wwwwww/////// kaggl e com/.

---

Published as a conference paper at ICLR 2026
Table 3: Average log-predictive density (1) results on UCI datasets with TabICL. We evaluate our AR buffer integrated into TabICL foundation model against independent and standard AR baselines. Performance is measured on interpolation (Int) and forecasting (For) tasks across six real-world datasets. Results are reported as mean and standard error over 16 randomly sampled mini-datasets (M- = 32) low context (N = 16) and high context (N = 1024) settings.
Low CONTEXT REGIME (N 16)
Electric Cons.
Gas Turhine
Bike
Sharing For
Tetouan
Jena
Cali. Int
For
For
Int
Independent 0.38 maz -2.87
-0.65
w.s
-1.46
0.84 6:10)
-0.04 out)
-0.59
-4.71
0.10
-4.53
-1.31 (007)
Standard AR 0.88 mis -0.97 (15t
-0.48
-1.00
m00
1.43 10:0
1.02
am
-0.09
0:0
-2.39
0030
0.64
am -2.15
(030 -1.17 m00
AR w/ buffer 0.78 m.. -1.00 asi
-0.48
6..8
-0.98
1.31 (8.0)
0.94
0. -0,10
0.15)
-2.41
(0.3))
0.55
2..15 (334)
-1.16 (009)
HIGH CONTEXT REGIME (N 1024)
Electric Cons.
Gas Turbine
Bike
Sharing For
Tetouan
Jena
Cali. Int
Int
For
Int
For
Int
Int
For
Int
For
Independent 1.78 ono 1.64 1 -0.01 6.e -0.60 mes
2.54 mos:
2.32 mo
0.36 mo
-1.12 0035)
2.01 0.00)
1.56 19 -0.44 2006
Standard AR 1.78 0.06:
1.70 0.15)
-0.01
6.10
-0.47
m2m:
2.54
m09
2,40 mum
0.36
momh
-0.08
(:23)
2.01
(.00))
1.80
(.13)
-0.44 (.06)
AR buffer 1.79 auee) 1.70 0.15) -0.01
(310
-0.48
m22
2.53
mas
2.39
mos:
0.36
mosi
-0.12
0122)
2.01
0006)
1.64
005)
-0.44 (.0))
"Standard AR" (step-by-step autoregression, K=1 equivalent), and "AR w/ buffer" (ours, K =32)- Results in Table 3 show that both AR modes consistently outperform independent predictions in low context settings and in high context forecasting. Crucially, AR w/ buffer matches standard AR within standard errors, demonstrating that the buffer preserves dependencies while enabling efficient autoregressive inference. In Appendix G, we provide results for the N = 256 setting.
6 DISCUSSION
We introduce a causal autoregressive buffer that decouples one-time context encoding from lightweight sequential updates in transformer-basee probabilistic models. By caching context keys/values and routing target-to-target dependencies through a causal buffer, we reduce the at- tention cost from O [ ([([N+K)3) to O(N2 + N K2). Across synthetic functions, EEG interpo- lation, multisensory modeling, and tabular prediction, method matches autoregressive baselines while achieving up to 20x faster joint sampling with minimal additional training cost over standard models and up 10x lower training cost than autoregressive-specif baselines.
There are several limitations. First, cost grows with K. Runtime and memory include an O(K)) term from causal self- self-attention in the buffer, and currently learn a fixed set of buffer positional embeddings. Scaling to longer horizons without growing training complexity may be possible via rotary position embeddings (RoPE; Su al., 2024) or attention biasing (ALiBi; Press al., 2022). Second, for long buffers, quality can drift relative to exact AR that re-encodes the context at each step. Exploring the draft-verify process from speculative decoding (Leviathan et al., 2023; Chen et al., 2023) could enable adaptive inference strategies using the buffer for improved performance.
A practical strength of our method is its plug-and-play applicability: the buffer is implemented via attention masks and token roles, demonstrated by its integration into the TabICL architecture. While we currently perform joint training of the base model and buffer, our method could be directly applied to pretrained NPs/PFNs. Parameter-effiiiet fine-tuning (Houlsby et al., 2019; Hu et al., 2022) could offer a direct path to enable buffered inference without full retraining. We also leave to future work deeper exploration of alternative attention backbones (e.g.. Jaegle et al., 2021; see Appendix H.4) and broader inference tasks (Chang et al., 2025).
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
This work was supported by the Research Council of Finland (Flagship programme: Finnish Center for Artificial Intelligence, FCAI) and ELISE Networks of Excellence Centres (EU Horizon:2020 grant agreement 951847). The project is also supported by the EuroHPC Joint Undertaking and its members including top-up funding by Ministry of Education and Culture. CH and SK were sup- ported by Business Finland Virtuall.ab4Pharmaa grant agreement 3597/31/2023) and the European Union (Horizon Europe, grant agreement 101214398, ELLIOT). NL was funded by Business Fin- land (project 3576/31/2023 and LUMI AI Factory (EU Horizon Europe Joint Undertaking and its members including uop-up funding by Ministry of Education and Culture). YY was supported by the Ministry of Education and Culture's Doctoral Education Pilot under Decision No. VN/3137/2024- OKM-6 (The Finnish Doctoral Program Network Artificial Intelligence, AI-DOC). SK was sup- ported by UKRI Turing AI World-Leading Researcher Fellowship (EP/W0029731), LA sup- ported by Research Council of Finland grants 356498 and 358980, the latter also supporting FS. The authors also acknowledge the research environment provided by ELLIS Institute Finland.
We acknowledge Verda for providing computational resources. Additionally, we acknowledge CSC -T Center for Science, Finland, for computational resources provided by the LUMI supercomputer, owned by the EuroHPC Joint Undertaking and hosted by CSC and the LUMI consortium (LUMI projects 462000943 and 462000874). Access was provided through the Finnish LUMI-OKM allo- cation. We also acknowledge the computational resources provided by the Aalto Science-IT Project from Computer Science IT.
Funded by the European Union. Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those the European Union or the granting authority. Neither the European Union nor the granting authority can be held responsible for them.
Co-funded by the European Union
EuroHPC Joint Undertaling
ETHICS STATEMENT
This work uses only publicly available datasets and synthetic simulators, with no sensitive data involved. The methods are for research purposes and pose no foreseeable ethical risks. We have followed the ICLR Code of Ethics.
REPRODUCIBILITY STATEMENT
We will release the code, including the training and evaluation pipelines, well configuration files, in the repository linked below.? All experiments use public datasets or, when applicable, simulator for synthetic data. Algorithmic details are presented in Algorithms 1 and 2, and all hyperparameters and training schedules are specified in the configuration files and documented in the appendix. Ablation studies are also reported in the appendix. We do not release pretrained weights, and no special data licenses or usage constraints apply.
REFERENCES
Luigi Acerbi. Variational Bayesian Monte Carlo. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2018
Luigi Acerbi. Variational Bayesian Monte Carlo with noisy likelihoods. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2020.
Luigi Acerbi, Kalpana Dokka, Dora Angelaki, and Wei Ji Ma. Bayesian comparison explicit and implicit causal inference strategies in multisensory heading perception. PLoS Computational Biology, 14(7)cc006110, 2018.
hhttps://gi thub. com acerbi a

---

Published as a conference paper at ICLR 2026
Mariannne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block Diffusion: Interpolating between autoregres- sive and diffusion language models. In International Conference Learning Representations, 2025.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisti- cians. Journal of the American Statistical Association, 112(518):859-87) 2017.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in Neural Information Processing Systems, 2020.
Wessel P Bruinsma, James Requeima, Andrew YK Foong, Jonathan Gordon, and Richard E Turner. The Gaussian neural process. In 3rd Symposium on Advances in Approximate Bayesian Inference, 2021.
Wessel P Bruinsma, Stratis Markou, James Requeima, Andrew YK Foong, Tom Andersson, Anna Vaughan, Anthony Buonomo, J Scott Hosking, and Richard E Turner Autoregressive conditional neural processes. In International Conference on Learning Reprexentations, 2023.
Paul E Chang, Nasrulloh Loka, Daolang Huang, Ulpu Remes, Samuel Kaski, and Luigi Acerbi. Amortized probabilistic conditioning for optimization, simulation and inference. In International Conference on Artificial Intelligence and Statistics. PMLR, 2025.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean- Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning. PMLR, 2020. Tri Dao. FlashAttention-22: Faster attention with better parallelism and work partitioning. In Inter- national Conference on Learning Representations, 2023.
Nicola De Cao, Wilker Aziz, and Ivan Titov. Block neural autoregressive flow. In Uncertainty in artificial intelligence. PMLR, 2020.
Justin Deschenaux, Lan Tran, and Caglar Gulcehre. Partition generative modeling: Masked model- ing without masks, 2025. URL http:: http://axxivvv..o org/ abs 2505. ...0.... 18883.
Vincent Dutordoir, Alan Saul, Zoubin Ghahramani, and Fergus Simpson. Neural diffusion pro- cesses. In International Conference on Machine Learning. PMLR, 2023.
Lasse Elsemüller, Hans Olischläger, Marvin Schmitt, Paul-Christian Bürkner, Ullrich Koethe, and Stefan T. Radev. Sensitivity amortized bayesian inference. Transactions on Machine Learning Research, 2024.
Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Efficient queries transformer neural processes. In NeurlPS 2022 Workshop on Meta-Learning, 2022.
Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent bot- tlenecked attentive neural processes. In International Conference on Learning Representations, 2023.
Leo Feng, Frederick Tung, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Memory efficient neural processes via constant memory attention block. In International Confer- ence on Machine Learning. PMLR, 2024.
Andrew YK Foong, Wessel P Bruinsma, Jonathan Gordon, Yann Dubois, James Requeima, and Richard E Turner. Meta-learning stationary stochastic process prediction with convolutional neu- ral processes. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2020.
12

---

Published as a conference paper at ICLR 2026
Marta Garnelo, Dan Rosenbaum, Chris J Maddison, Tiago Ramalho, David Saxton, Murray Shana- han, Yee Whye Teh, Danilo J Rezende, and SM Ali Eslami. Conditional neural processes. In International Conference Machine Learning. PMLR, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo Rezende, SM Ali Eslami, and Yee Whye Teh. Neural processes. In ICML 2018 Workshop Theoretical Foundations and Applications of Deep Generative Models, 2018b.
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. International Conference Machine Learning. PMLR, 2015.
Manuel Gloeckler, Michael Deistler, Christian Weilbach, Frank Wood, and Jakob Macke. All- in-one simulation-based inference. In International Conference on Machine Learning, PMLR,
2024.
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models with- out positional encodings still learn positional information. Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1382-1390, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2020.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. TabPFN: A transformer that solves small tabular classification problems a second. International Conference on Learning Representations, 2023.
Noah Hollmann, Samuel Müller, Lennart Purucker, Arjun Krishnakumar, Max Körfer, Shi Bin Hoo, Robin Tibor Schirrmeister, and Frank Hutter. Accurate predictions on small data with a tabular foundation model. Nature, 637(8045):319 326, 2025.
Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Rep- resentations, 2022.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficien transfer learning for nlp. In International Conference on Machine Learning. PMLR, 2019.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations. PMLR, 2022.
Chin-Was Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Learning. PMLR, 2018.
Bobby Huggins, Chengkun Li, Marlon Tobaben. Mikko J. Aaros, and Luigi Acerbi. PyVBMC Efficient Bayesian inference in Python. Journal of Open Source Software, 8(86):5428, 2023. doi: 10.21 10//joss..5542...
Kazuki Irie. Why are positional encodings nonessential for deep autoregressive transformers? re- visiting petroglyph. Findings of the Association for Computational Linguistics: ACI. 2025, pp- 551-559, 2024.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International Conference on Machine Learning. PMLR, 2021.
Durk P Kingma, Tim Salimans, Rafal Jozefowiez, Xi Chen, llya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2016.
David C Knill and Alexandre Pouget. The Bayesian brain: The role of uncertainty in neural coding and computation. Trends in Nearosciences, 27(12):122-71, 2004.
13

---

Published as a conference paper at ICLR 2026
Konrad P Körding, Ulrik Beierholm, Wei Ji Ma, Steven Quartz, Joshua B Tenenbaum, and Ladan Shams. Causal inference in multisensory perception. PLoS One, 2(9):e943, 2007.
Jose Lara-Rangel, Nanze Chen, and Fengzhe Zhang. Exploring pseudo-token approaches in trans- former neural processes. arXiv preprint arXiv:2504 14416, 2025.
Hugo Larochelle and lain Murray. The neural autoregressive distribution estimator. In international conference on artificial intelligence and statistics. PMLR, 2011.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set Transformer: A framework for attention-based permutation-iivvaiiat neural networks. In Inter- national conference on machine learning. PMLR, 2019.
Yaniv Leviathan, Matan Kalman, and Yossi Matias, Fast inference from transformers via speculative decoding. In International Conference on Machine Learning. PMLR. 2023.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. Advances in Neural Information Processing Systems, 2024.
Yaron Lipman, Ricky Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference on Learning Representations, 2023.
Shuze Liu, Trevor Holland, Wei Ji Ma, and Luigi Acerbi. Distilling noise characteristics and prior expectations in multisensory causal inference. 2025.
Sulin Liu, Peter J Ramadge, and Ryan P Adams. Generative marginalization models. In Interna- tional Conference on Machine Learning. PMLR, 2024.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In International Conference Machine Learning. 2024.
Stratis Markou, James Requeima, Wessel P Bruinsma, Anna Vaughan, and Richard E Turner. Practi- cal conditional neural processes via tractable dependent predictions. In International Conference Learning Representations, 2022.
Sarthak Mittal, Niels Leif Bracher, Guillaume Lajoie, Priyank Jaini, and Mareus A Brubaker. Ex- ploring exchangeable dataset amortization for bayesian posterior inference. In JCML. Work- shop on Structured Probabilistic Inference and Generative Modeling, 2023.
Sarthak Mittal, Niels Leif Bracher, Guillaume Lajoie, Priyank Jaini, and Marcus Brubaker. Amor- tized in-context Bayesian posterior estimation. arXiv preprint arXiv:2502.06601, 2025.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Trans- formers can do Bayesian inference. International Conference on Learning Representations, 2022.
Samuel Müller, Matthias Feurer, Noah Hollmann, and Frank Hutter. PFNs4BO: In-context learning for bayesian optimization. In International Conference on Machine Learning. PMLR, 2023.
Kevin P. Murphy. Machine Learning: A Probabilistic Perspective. MIT Press, 2012.
Kevin P Murphy. Probabilistic Machine Learning: Advanced Topics. MIT press, 2023.
Ryan L. Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Janossy pooling: Learning deep permutation-invarice functions for variable-size inputs. In International Confer- ence on Learning Representations, 2019.
Tung Nguyen and Aditya Grover. Transformer Neural Processes: Uncertainty-awar meta learning
via sequence modeling. In International Conference Machine Learning. PMLR, 2022.
George Papamakarios, Theo Paylakou, and Iain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2017.
14

---

Published as a conference paper at ICLR 2026
Massimiliano Patacchiola, Aliaksandra Shysheya, Katja Hofmann, and Richard E Turne.. Trans- former neural autoregressive flows. In ICML 2024 Workshop on Structured Probabilistic Infer- ence & Generative Modeling, 2024.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference Learning Representations. PMLR, 2022.
Jingang Qu, David Holzmüller, Gaël Varoquaux, and Marine Le Morvan. TabICI: A tabular foun- dation model for in-context learning on large data. In International Conference on Machine Learning. PMLR, 2025.
Jingang Qu, David Holzmüller, Gaäl Varoquaux, and Marine Le Morvan. Tabielv2: A better, faster, scalable, and open tabular foundation model. arXiv preprint arXiv:2602. 2026.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018. URL https: //oppa:///ppaaiiii wto. nddx/ language-unn www.aaseeeuuuuuriieeeeee
Carl Edward Rasmussen and Christopher KI Williams. Gaussian Processes for Machine Learning. MIT Press, 2006.
Arik Reuter, Tim GJ Rudner, Vincent Fortuin, and David Rügamer. Can transformers learn full Bayesian inference in context? International Conference on Machine Learning, 2025.
Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin T Chiu, and Volodymyr Kuleshov. The diffusion duality. In International Conference on Machine Learning, 2025.
Subham Sekhar Sahoo, Zhihan Yang, Yash Akhauri, Johnna Liu, Deepansha Singh, Zhoujun Cheng, Zhengzhong Liu, Eric Xing, John Thickstun, and Arash Vahdat. Esoteric language models: Bridging autoregressive and masked diffusion lims, 2026. URL https:://axxtii org/ / 2506. 01928.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances in Neural Information Processing Systems, 2024.
Francesco Silvestrin, Chengkun Li, and Luigi Acerbi. Stacking Variational Bayesian Monte Carlo. Transactions Machine Learning Research, 2025.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. PMLR, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. Interna- tional Conference on Learning Representations, 2021.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En-
hanced transformer with rotary position embedding. Neurocomputingg 568:12706: 2024.
Haotian Tang, Yecheng Wu, Shang Yang, Enze Xie, Junsong Chen, Junyu Chen, Zhuoyang Zhang, Han Cai, Yao Lu. and Song Han. HART: Efficient visual generation with Hybrid AutoRegressive Transformer. In International Conference on Learning Representations, 2025.
Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator. In Inter- national Conference on Machine Learning. PMLR, 2014.
Benigno Uria, Mare-Alexandre Côté, Karol Gregor, lain Murray, and Hugo Larochelle. Neural autoregressive distribution estimation. Journal of Machine Learning Research, 17(205):1-37, 2016.
15

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems. Curran Associates, Inc., 2017.
George Whittle, Juliusz Ziomek, Jacob Rawling, and Michael A Osbore.. Distribution trans- formers: Fast approximate Bayesian inference with on-the-fly prior adaptation. arXiv preprint arXiv:2502.02463, 2025.
Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dLLM: Training-free acceleration of diffusion LLM by enabling KV cache and parallel decoding. arXiv preprint arXiv:2505.2200 226788, 2025.
Xiao Lei Zhang, Henri Begleiter, Bernice Porjesz, Wenyu Wang, and Ann Litke. Event related potentials during object recognition tasks. Brain Research Bulletin, 38(6):531-538 1995.
Xiyuan Zhang, Danielle C Maddix. Junming Yin, Nick Erickson, Abdul Fatir Ansari, Boran Han, Shuai Zhang, Leman Akoglu, Christos Faloutsos, Michael W Mahoney, et al. Mitra: Mixed synthetic priors for enhancing tabular foundation models. In Advances Neural Information Processing Systems, 2025.
Chunsheng Zuo, Pavel Guerzhoy, and Michael Guerzhoy. Position information emerges in causal transformers without positional encodings via similarity of nearby embeddings, In Proceedings of the s1 International Conference Computational Linguistics, pp. 9418 943000 2025.
16

---

Published as a conference paper at ICLR 2026
Table of Contents
A Method Details
A.1 Modules and notation
18
A.2 Training mask that implements (R4)
18
A.3 Algorithms for autoregressive sampling and log-likelihood evaluation
19
B Transformer Neural Process Baselines Details
20
B.I TNP-D
20
TNP-ND
20
TNP-A
20
C Computational Efficiency Details
21
C.I Scaling with Batch Size
21
C.2 Impact of Custom Triton Kerne
22
C.3 Comparison to Open- Sperce Baselines
...............
22
C.4 Training Time Scaling
................................
23
Impact of Attention Patterns on Training .pe.............................
24
C.6 Memory Usage
25
D Experimental Details
26
D.1 Model Configuration
26
D.2 Datasets
27
D.3 Multisensory causal inference model and experiment details
28
Tabular model details
31
D.5
Evaluation Details
32
E
Additional Log-Predietive Densitty Results Synthetic and EEG Tasks
34
E.1
Predictive Power of Different Heads
34
E.2 Results of Larger M
34
E.3 EEG Forecasting w/ and w/o Target Permutation
35
F Additional Multisensory Causal Inference Model Results
G Additional Tabular Foundation Model Results
35
H Ablations and Extra Experiments
36
H.1 Comparison to Con-Perrutaaion-llvaiinnn Transformers
36
H.2 Positional Embeddings Ablation
36
H.3 Number Samples Order Averaging Ablation
36
H.4 Extension Latent Bottlenecked Attentive Neural Processes Model
H.5 Buffer Size Ablation
38
I Use of Large Language Models
38
17

---

Published as a conference paper at ICLR 2026
A
METHOD DETAILS
This appendix spells out the modules used in Eq. (3), the single block-sparse attention mask that implements requirements (RI)-R4), and the exact procedures for autoregressive sampling and one- pass joint log-likelihood evaluation.
A.1 MODULES AND NOTATION
Our method uses three sets tokens: context C, buffer B, and targets of sizes N, K, M, respec- tively. Throughout this paper, let
Eg X
Ey :V-R", a K} +R
denote learned embeddings for inputs, outputs, and buffer positions. In addition, we introduce role embeddings that indicate token type, denoted by roke sake and eigi for context, buffer, and target tokens, respectively.
Context encoder rg. Given context pairs C {(x, you J.N construct context tokens: e (x,) + (Un) e process them with bidirectional MHSA (no positional embeddings), and cache per-layer keys/values:
(KV = rc(C) (computed once; immutable).
Buffer encoder rg. For buffer prefix B1:k {(x], , yJ ) form tokens cjaf (x())) + (aj)))+(+++++++))))))))) +e then apply strictly causal MHSA on {e Jo so that each token restricted to attend only to earlier tokens in the sequence, and in addition, each token performs cross-attention to the cached context ( This yields per-layer KV 8 that we update incrementally at inference:
[KV BL 1
rB (c())).
Target decoder Figt and prediction head. For target input x,, create query token lut Cm (xY,, + etole. The target decoder Figr performs single cross-attention from to the concatenated keys/values the context cache (KVE} and the visible buffer prefix (KV) B B }
lgt
followed by normalization and an MLP:
[(KVC), (KV (KVB.....))
1))
hm = Figs em
Φn (hhm)
where v is the distribution head (e.g., s., the mixture-of-Gaausian head).
A.2 TRAINING MASK THAT IMPLEMENTS (R1)-(R4)
We concatenate tokens as [C, B, T] with sizes N, K, and M, respectively, and use one block-sparse attention mask consisting of the following five unmasked sections (everything else is masked):
(1) Self-attention within context. Context tokens attend bidirectionally other context tokens. Context never attends to buffer or targets (context is read-only outside this block).
(2) Buffer reads context (cross-attention Each buffer token can read (attend to) all context tokens. This lets the buffer incorporate task information from the cached context while keeping the context cache immutable.
(3) Causal self-attention within the buffer. Within the buffer itself, attention is strictly causal: a buffer token at position j can only read earlier buffer positions <j (no future reads). This encodes the autoregressive dependency among realized targets.
(4) Targets read context (cross-attention). Each target query read the entire cached context. There are no edges between targets.
(5) Targets read buffer (prefix only, cross-attention). Each target query can read only visible prefi of the buffer. The visible prefix length for target m is en training (teacher forcing): we = for 50% of targets and sample Um (a.fform1,...... for the rest (the curriculum);
18

---

Published as a conference paper at ICLR 2026
Transformer diagonal prediction map training mask
Transformerr diagonal prediction map training mask w/ AR buffer
KIV
KIV Target cross-attn
Context self-attn
Buffer-context cross-attn
Buffer self-atin
Target-context cross-atin
Figure A1: Block-sparse attention masks with and without an autoregressive buffer. Left: diagonal prediction-map transformer (e-g., TNP/PFN): context attends to itself and each target reads the entire context. Right: our buffered variant inserts autoregressive memory B between context and targets, adding three blocks: (i) buffer reads context (ii) causal self-attention within buffer (iii) target reads varying number of elements from start of buffer, depending curriculum.
Algorithm Autoregressive sample generation for K targets
Require: Context C {on, Vn (n))) > target inputs {
1: (KV2) < rc(C)
O(N2); cached b empty buffer cache
2: Initialize (KV* B
3: for k = 1 to K do
4: h& + Figt ( [KVV)) (KV B 1))
5: Sample v po(: po(:; (hx))
6: Append (sit vt); update (KVg Bra) (strictly causal)
7: end for
8: return {
sampling: at step k. the active query sees the realized prefix k-1; one-pass joint log-likelihood: packed queries use Um= to recover the autoregressive chain in single forward pass.
All other connections are masked: context never reads buffer or targets; targets never read targets; and buffer never reads targets. This single pattern implements the four requirements from the main text- inteimlllable context, strictly causal buffer, unidirectional flow out of context, and target access (context + visible buffer). See Fig. Al for the diagram.
Complexity. Under this mask, full prediction pass costs O(N2+NK' K22 attention operations per layer: one-time ((N2) for C, ((NK) for reads from C, and O 2(2) for causal buffer self- attention. This replaces the (KK(N+))2) cost of naive AR re-encoding. Packing / target orders in parallel (for order averaging) isolates the B buffer sets while sharing the context cache, yielding O(22 B(NK + K2)).
A.3 ALGORITIMS FOR AUTOREGRESSIVE SAMPLING AND LOG-LIKELIIOOD EVALUATION
We include here the pseudocode for the main procedures used in our method. Algorithm details the autoregressive sampling procedure, and Algorithm 2 presents the joint likelihood evaluation.
19

---

Published as a conference paper at ICLR 2026
Algorithm 2 Joint log-likelihood evaluation for K targets
Require: Context C = {(rn Bn)) ordered targets ((r W
1: (KV2) < re(C)
O(N2); cached
2: Build all buffer tokens; compute (KV (..x) under causal mask
B.ix)
3: Build target queries + k=
4: Mask: target & sees Bis-- and all C
5: Compute (log Jk
6: return Σ log
B TRANSFORMER NEURAL PROCESS BASELINES DETAILS
We summarize the baseline transformer neural process (TNP) variants used in our comparisons, following Nguyen & Grover (2022). Architectural hyperparameters appear in Appendix D.I.
B.1
TNP-D
This model takes as input a context set {(x,: > and a target set {x 1M /m=1 Similar to Ap-
pendix A, the context embeddings etth are processed with bidirectional MHSA with no positional encodings. Each target is decoded by:
Figl [ig(eltt dlest m rc(C)).
h
Φm = -))
where & is the distribution head (Gaussian in the original paper: we primarily use a mixture of Gaussians). The left panel of Fig. Al shows the training mask for TNP-D. This model is trained via maximum likelihood estimation of independent targets given fixed context set.
At deployment, the decoding can be independent or autoregressive, yielding TNP-D-Ind and TNP- D-AR methods. INP decodes all targets independently in single pass. It is fast (context and targets encoded once), but cannot capture dependencies between targets.
TNP-D-AR decodes targets sequentially, appending each sampled (x,, V/)) to the context. This captures joint structure but requires re-encoding the growing set at each step. TNP-D-Ind is invari- ant target order; TNP-D-AR is order-sensitive, we approximate the predictive distribution by averaging over multiple target orderings.
B.2 TNP-ND
This model encodes the context once and decodes all targets simultaneously by parameterizing joint multivariate Gaussian distribution over the outputs. The embedder and transformer backbone identical to those of TNP-D-Ind:
re(C) rc()
hm
Kigs
Then the joint distribution is obtained via
= AiND (h has),
where uND is the multivariate Gaussian head that outputs both mean vector and valid covariance matrix. The mean is produced per target, and a lightweight self-attention head over the of tar- gets yields fixed-width embeddings that are transformed into valid covariance factor. This design supports a variable number targets and invariant to target order.
The training optimizes the joint multivariate Gaussian likelihood of the target points. At inference, the joint samples and log- ligeliihl are computed a single pass. This model invariant to the order of target points.
B.3 TNP-A
The key difference between this model and TNP-D is the attention mechanism on the target set. M This model processes three sets: the context ((xn, yn vn) J the target (xt, m= and the observed
20

---

Published as a conference paper at ICLR 2026
(m)mm= I== by {em J- Similar to TNP-D, the context embeddings attend to each other. For the target set, each em attends to the context and the previous observed target embeddings e U-igt iigm" Likewise, the observed target embeddings attend to context and previous target embeddings (Fig. of Nguyen & Grover 2022).
target {(x,, Vm
( m= To differentiate, denote the embeddings of {(x), (
igt
The target causal mask allows TNP-A to model the joint likelihood simultaneously in one single pass, assuming the observations are given (c.g.. for training and test log- logeliiel evaluations). For prediction generation, however, each sampled target needs to be re-encoded and attended for the generation next targets, yielding a sequential re-encoding procedure. The causal mask on the target set is sensitive to the target order, and thus the final likelihood is average dom permutations. Note that this model processes duplicated target set-[x;,, 1.M
multiple ran- J Jm and an observed sequence { (( yo J m= this creates significant computational overhead in both the training and the inference, particularly when the target set is large (see e.g. Appendix C and Figs. A7 to A9).
Compared to our method, TNP-A can be viewed as TNP-D with 'frozen buffer' { Dim 3 m= of
size K M containing the observed targets. For likelihood evaluation where all sets are processed one shot, the behavior of TNP-A and our approach are analogous, with the {(xx, y ) JM m= serving role similar to our buffer. However, for AR sampling, TNP-A repeatedly re-encodes the full context and target sets after each sampled Vo whereas our method dynamically adapts to new samples. Moreover, since TNP-A does not afford a dynamic-size target set decoupled from the 'in-contett' targets, training is much more expensive than our method (see Fig. the main text).
C COMPUTATIONAL EFFICIENCY DETAILS
This section provides additional empirical results to support the efficiency claims made in the main paper. We present an analysis of performance scaling with batch size, ablation study of our custom kernel, comparison against unoptimized open-source baselines, and further ablations on training time. In all subsequent plots, the absence of data point for given method indicates that the experiment failed due to an out-of-memory (OOM) error for that specific configuration.
C.I SCALING WITH BATCH SIZE
To analyze the effect of batch size B, provide expanded results for autoregressive sampling and joint log-likelihood evaluation in Fig. A2 and Fig. A3, respectively. These plots show the wall-clock time function of the number of context points N for various batch sizes. The results confirm that our method's 's performance advantage over autoregressive baselines like TNP-A is consistent and often widens as the context and batch size increase.
Sample generation time (M 16)
1024
Time (5)
10-1
Bolk kgl
10-2
256
256
Figure A2: Autoregressive sampling time (log scale) versus context size N for an expanded range of batch sizes B.
21

---

Published as a conference paper at ICLR 2026
Log likelihood evaluation time (M 16)
B 128
B 256
B 512
8 1024
Time (6)
10-1
belk
10-2
128
256
1024
256
1024
125
250
1024
128
250
1024
Ours
TNP-D-Ind
INF-DAR
TNPA
TNPNN
Figure A3: Joint log-likelihood evaluation time (log scale) versus context size N for an expanded range of batch sizes B.
C.2 IMPACT OF CUSTOM TRITON KERNEL.
To isolate the contribution of our custom attention kernel, we compare the sampling time of our method with and without this optimization. The kernel is designed to accelerate key computational step: the cross-attention between the batched target embeddings (batch size B) and the concatenation of batched buffer cache with shared context cache (batch size 1). A naive implementation would explicitly expand the context cache tensor B times to match the batch dimension before the attention operation. This "expand" operation is memory- seemory-bandwidhe bandwidth intensive and creates large, redundant intermediate tensor.
Our custom Triton kernel avoids this bottleneck by fusing the expansion and attention computations. The kernel loads the single context cache into fast SRAM and reuses it for each item the batch, calculating the attention on-the-fly without ever materializing the full expanded tensor in slower global memory. As shown in Fig. A4, this memory-centric optimization provides a substantial speedup that grows with the batch size B.
Sample generation time (M 16)
128
B = 256
B = 512
8 = 1024
Time (5)
10-1
6x
32
64
128
256
512
1024
32
64
128
256
512
1024
32
64
128
256
512
1024
32
64
128
256
512
1024
N
N
N
N
-
Ours (w/ Trilon)
Ours (w/o Triton)
Figure A4: Ablation study for autoregressive sampling, comparing method with and without the custom Triton kernel across different context and batch sizes.
C.3 COMPARISON TO SPEN-SOURCE BASELINES
To demonstrate the fairness of primary comparisons, benchmark our optimized baseline implementations against their standard, publicly available versions. The results for sampling and likelihood evaluation
shown in Fig. A5 and Fig. A6. Our optimized baselines are consistently 3 10x faster than their standard counterparts. This confirms that our method's 's performance gains due to fundamental algorithmic advantages, not unfair comparison against unoptimized code.
22

---

Published as a conference paper at ICLR 2026
Sample generation time (M 16)
B = 128
B 256
512
B = 1024
(6)
10
Emik Bx
Ged
10-2
10-3
255
256 512
256
512
1024
64
128
512
1024
TNP-D-Ind
(buse)
(bune)
TNP-ND (buse)
TNP-ND (sumpiled)
DNP-D-Ind (compiled)
D-R (compiled)
TNPA (compiled)
Figure A5: Comparison of our optimized baseline implementations against standard open-source versions for autoregressive sampling.
Log- likelhool evaluation time (M 16)
B 128
B = 256
B 512
B 1024
10
107
Time (6)
10--1
10-2
10"3
64
128
256
512
1024
64
125
250
512
1024
64
125
250
512
1824
64
3.25
250
512
1024
N
TNP-A (base)
TNPA (compiled)
TNP ND (base)
TNP-ND (compplle)
Figure A6: Comparison of optimized baseline implementations against standard open-source versions for joint likelloodd evaluation.
C.4 TRAINING TIME SCALING
We further analyze the scaling of training step time with respect to the number of target points M for different batch sizes. Figs. A7 A9 show this relationship for batch sizes of 64, 128, and 256, respectively. The results show that as the context, target, or batch size increases, TNP-A becomes increasingly expensive to train relative to all other methods.
Training step time (forward + backward, B = 64)
M 128
M = 256
M 512
100
Time (s)
10-1
32
64
128
256
512
1024
32
64
128
256
512
1024
32
64
128
256
512
1024
N
-
Ours
TNP-D-Ind
TNP-
TNP-ND
Figure A7: Training step time number of target points M for batch size B 64.
23

---

Published as a conference paper at ICLR 2026
Training step time (forward backward, 128)
128
256
512
109
Time (s)
10
Sammi
64
128
256
1024
32
128
512 1024
64
128
256
512
1024
N
N
N
Ours
TNP-D-Ind
TNP-
TNP-ND
Figure A8: Training step time vs. number target points M for batch size B 128.
Training step time (forward backward, 256)
M 128
M 512
10P
Time (s)
10-1
64
128
256 N
512
1024
32
128
256
512
1024
64
128
256
512
1024
N
Ours
TNP-D-Ind
TNP-A
TNP-ND
Figure A9: Training step time vs. number target points M for batch size B 256.
C.5 IMPACT OF ATTENTION PATTERNS ON TRAINING SPEED
key difference between the baseline models their compatibility with modern, efficient attention implementations. The causal attention mask required by TNP-A during training is incompatible with kernels like FlashAttention, forcing it to use Py PyTorhh's 's standard, but slower, "math" attention backend. In contrast, models like TNP-D and ours can leverage these faster kernels.
In Appendix B, we discussed the duplicated processing of TNP-A on the target set, which incurs significant computational overhead. To determine if TNP-A's slow training is fundamental to its architecture or merely an artifact of this kernel incompatibility, we conduct a controlled ablation. We disable Fttentntention for all methods, forcing fair comparison the same standard PyTorch attention backend. The results, shown in Figs. A0 to A12, are unequivocal. Even on level playing field, TNP-A's training time is orders of magnitude slower than all other methods. This confirms that its high computational cost is an inherent consequence of its autoregressive design, not just an implementation detail.
24

---

Published as a conference paper at ICLR 2026
Training step time (forward backward, 64)
256
M 512
10P
Time
10:1
32
64
128
256
512
1024
32
64
128
255
512
1024
32
64
128
256
512
1024
N
N
N
TNP-D-Ind
TNP-
TNP-ND
Figure A10: Training step time vs. number of target points M using the standard PyTorch attention backend (FlashAttention disabled). Batch size B 64.
Training step time (forward + buckward, 8 128)
M = 128
M 256
M 512
10
Time (S)
10-1
32
64
128
256
512
1024
32
64
128
255
512
1024
32
64
128
256
512
1024
N
N
---
Ours
TNP-D-Ind
T
TNP-A
TNP-ND
Figure A11: Training step time number of target points M using the standard Py Torch attention backend AttashAttention disabled). Batch size B 128.
Training step time (forward + backward, B = 256)
M = 128
M = 256
M = 512
10P
Time
10-1
32
64
128
256
512
1024
32
64
128
256
512
1024
32
128
256
512
1024
N
Ours
TNP-D-Ind
TNP-A
TNP-ND
Figure A12: Training step time number of target points M using the standard PyTorch attention backend (FlashAAtentionn disabled). Batch size B 256.
C.6 MEMORY USAGE
Figure A13 reports peak GPU memory consumption during autoregressive sampling as a function of context size N across different batch sizes B. Our method maintains consistently low memory usage across all configurations. requiring 6-7x less VRAM than TNP-D-AR and TNP-A at large context sizes (N 1024). This efficiency stems from our fixed-size buffer mechanism: while autoregressive baselines must cache representations that grow with context size and batch size, our method only caches buffer representations of size K, independent of the batch. TNP--DIInd and
25

---

Published as a conference paper at ICLR 2026
TNP-ND show lower memory usage but, demonstrated in the main text, cannot capture complex predictive dependencies.
Peak memory usage (M 16)
B 128
= 256
512
1024
10P
Peak VRAM (GR)
let bodl
10-1
and
gill
10-2
32
04
128
256
512
1024
32
64
128
256
512
1024
32
64
128
250
512
1024
32
64
128
250
512
1024
N
N
N
Ours
INP-D-Ind
INP-DAR
TNPA
INP-ND
Figure A13: Peak GPU memory usage during sampling as function of context points for differ- ent batch sizes B. Our method scales efficiently due to its fixed buffer size, using substantially less memory than expressive autoregressive baselines.
D EXPERIMENTAL DETAILS
D.1 MODEL CONFIGURATION
In our paper, we use MLP to map context pairs, buffer pairs, target points to tokens. Then trans- former is applied to the sequence of tokens. We used Eexture-of-Gaussian (GMM) head as our main head distribution (more expressive than single Gaussian head, demonstrated in Appendix E). In general, we train all models (except the tabular model; see Appendix D.4 for details) with the following settings.
Training configurations.
Optimizer: Adam with learning rate 1 10 (unless stated otherwise). 3 (0.9, 0.999), no weight decay. For TNP w/ buffer, we use the same settings, but apply weight decay of 0.01 for stability.
Scheduler: Cosine schedule with warmup; warmup ratio 0.1 for all experiments. For TNP w/ buffer, we use warmup ratio 0.05.
Training loop: 32 epochs.
Embedder. We use a 3-layer MLP with 256 hidden layer dimension and 128 output dimension. There is skip connection between the first linear layer and the MLP output.
Transformer backbone. This has 6 layers of transformer encoder modules, each with multi- head attention 4 heads and dimension 128 followed by MLP feedforward of layers, dimension 128 + 256 - 128. This is the transformer attending context, buffer, and target set (Appendix A and Appendix B).
Prediction head. Note first that different distribution heads involve individual parameterization structures. Therefore, another layer of distribution-specinh NNs required to process the above transformer outputs. This NN module is considered part of the distribution head (the & in Ap- pendix A and Appendix B).
For our method, TNP-D, and TNP-A, the head consists of layers of MLP with dimension 128 256 + 3 Dy * Neooponents- where D, is the output dimension of the problem and Ncomponeuts is the number of Gaussian components. The MLP output is then chunked into weights, means, and standard deviations (of the same shape) which parameterize the GMM, and the outputs are sampled in parallel for D, > We set N Noompoments 20 for all tasks except for EEG where N Naomponents 8.
26

---

Published as a conference paper at ICLR 2026
the setting from Nguyen & Grover (2022), where the targets mapped to a mean and a Cholesky matrix, which parameterize the multivariate Gaussian. The mean of each target is mapped by an MLP with dimension 128 256 Dy- The Cholesky matrix requires two steps: (i) the target tokens (conditioned on context via the above transformer backbone) are first decoded into H E RM by another 3-layer transformer (no positional encoding, 4 heads, each layer with dimension 128 and MLP 128 + 256 + 128, no mask) and then an MLP projector (128- + 256 + 20); (ii) the Cholesky matrix is taken as L lo
For TNP-ND,
Trained model selection. We track the loss value in each epoch as we train the models. The parameters with the best loss value on the validation set are selected for the evaluations on separate test set.
D.2 DATASETS
Gaussian Process (GP) Functions. As first toy case, we test on GP functions (see Rasmussen & Williams 2006 for details of GPs). In this example, batch contains 128 functions of one di- mensional inputs (D = 1) and one dimensional observations (D, = 1). The inputs are sampled from interval 2-2, 2] using the Sobol sequence. For each batch, we first sample kernel class from enaoredexppoeential (RBF), Matérn-3/2, Matérn-5/2 with probabilities 0.4, 0.3, and 0.3, respectively. Conditional on the chosen class, each function receives its own kernel hyperparameters: the variance of Uniform (0.5, 1.5] and the lengthscale & ~ Uniform (0.1, 1/, broadly covering diverse classes of functions of amplitude around 1. We then sample functions from GP k)). where ls represents the sampled kernels, and add i.i.d. Gaussian observation noise with variance 10-5, The resulting values are randomly partitioned into context, buffer, and target sets. Note that within batch the kernel class is fixed, whereas the hyperparameters are sampled independently for each function.
During training. we sample the context set size N between 4 and 192 with a maximum buffer size of 16.
Sawtooth Functions. The second example is the non- Gaussian sawtooth functions (Bruinsma et al., 2023). In this example, batch contains 128 functions of one dimensional inputs (D = 1) and one dimensional observations (D, 1). The inputs are sampled from interval [-2, using the Sobol sequence. An input and output follows:
y(x) Ynonoise (x) + C,
Vnoncise (x) (w((u, - Φ)) mod
where IR2 is direction sampled uniformly from the unit sphere via - g/llg/la with g (O Ip); w, &. and denote the frequency, phase offset, and additive noise, respectively; and the parameters are drawn independently as w Uniform[3, 5), Uniform [0, 1], and e ~ NN o²) with noise scale a ~ Uniform(0.05. 0.1).
During the training, we sample N between 8 and 128 and the maximum number of buffer is 16.
Elletrrenceepalloramm (EEG). The dataset contains 11, 520 trials of 122 subjects from 7 corre- lated channels with 256 time points each. The output channels are individually standardized to zero mean and unit variance. We randomly select 10 for the test set, reserve 10 for cross-validation, and the remaining for the train set. This leaves 7802 trials for the training and 896 for testing.
During the training, the trials are replicated for 200 times and shuffled. Each batch contains 32 trials sampled from the shuffled set. We select between 4 and 192 of the 256 time points to be context points, 32 buffer points, with the remaining being target points. Each batch has fixed size of context set.
We evaluate on both interpolation (random masking) and forecasting (temporal masking) tasks using the test subjects. The test set splits the 256 time points into context and target. For interpolation, we sample the specified number of context and target points from the full time sequence (Appendix E). For forecasting, we take the first N points as context set and the consecutive M points as target set. Forecasting with N- 192 context and M 64 target involves the full sequence.
27

---

Published as a conference paper at ICLR 2026
Multisensory causal inference model dataset. In the last example, adopt one of the multisen- sory causal inference models described in Liu et al. (2025) to build a simulator, which we then use to generate training data (full setup and generation procedure, iment, are provided in Appendix D.3). The inputs x correspond to the experimentally manipulated variables of the study, namely Tlype: SA; Sy and Vievel, where Tyyee denotes the task type (auditory vs visual localization), 9 and sv are the true locations of auditory and visual cues presented to human participants, and Vievel the level of noise applied to visual cues. We first generate sets of input points for the simulator to obtain the outputs y, which represent the predicted responses.
well description of the exper-
For training, we construct two datasets from the simulator with different values of p. a variable of the model regulating the level of recalibration of the auditory perceptual range (with p = 1 representing no recalibration and p = 4/3 representing full recalibration to the visual range, see Appendix D.. for more details), and train two separate models for each setting. We sample N between 0 and 400, fix N₁ = 256, and set the buffer size to a maximum of 16. For the zero-context case, we introduce one "dummy point", to indicate the absence of context to the model. During evaluation, we use the publicly available dataset obtained from the experiment described Liu et al. (2025). 8 For each of the 15 participants in the study, we extract two overlalllg subsets of experimental data of 400 trials each. We do so by stratifying on the joint levels of Vievel E {0, 1 and "yyee E {0, 1} (more details on these variables below), and extracting the two sets such that (i) within each split the six (3 x 2) strata are represented as evenly as possible, and (ii) the per-stratum counts are matched between splits. This yields 30 datasets overall (2 per participant).
For details of the real experiments and the complete data generation setup in the simulator, see Appendix D.3.
Tabular foundation model. We pretrain a tagk-sgi tabular model on synthetic data and eval- uate it on three UCI datasets. This larger model uses dedicated training procedure; architectural and training details are in Appendix D.4, and the evaluation protocol is in Appendix D.5.
D.3 MULTISENSORY CAUSAL INFERENCE MODEL AND EXPERIMENT DETAILS
To probe our method's 's suitability for Bayesian model comparison, we consider a computational neuroscience study investigating multisensory causal inference, described in Liu et al. (2025).
D.3.1 ORIGINAL NEUROSCIENCE EXPERIMENT
Stimuli and procedure. In this work, we take into account subset of the experimental data obtained from 15 human participants who, at each experimental trial, were asked to perform one of two localization tasks, which the authors refer to as bisensory visual (BV) and bisensory auditory (BA) localization. In both cases they were presented with an auditory cue, located at an angle uniformly sampled among {-15--- - 10 -5a 0%, 5", 10 15º} from the participant, and a visual one, either at the same location as the auditory one (*= 1/2 of trials) or at an angle uniformly sampled between -20º and 20. They were either asked report the location of the visual (BV) or the auditory (BA) stimulus on a screen. Here we call those locations 8y and SA. respectively. The level of noise Vievel associated with the visual stimulus location was experimentally manipulated by modifying the size of the stimulus itself. In practice, this meant presenting a small ( (Vieel = 0; ss 1/3 of trials), medium (Veeel = 1; 1/3 of trials) large (Vevel 2; ~ 1/3 of trials) visual stimulus.
Each participant completed a total of 1000 trials.
Cognitive models. Here we focus on two versions of the "vanilla" model described in the original paper (Liu et al., 2025). On each trial, the participant is assumed to believe the two stimuli could come from either a common (C = 1) or different (C = 2) source, assigning a fixed prior probability 1) Psame to the former case. Regardless of this, the participant has Gaussian priors over stimuli locations p(*4) A (* 0, o and p(sy) = NVsvv 00 oz).
A key assumption of the model is that participants do not have direct access to the true location of the stimuli, but only to noisy auditory and visual percepts, a common feature in Bayesian models of perception (Knill & Pouget, 2004). These percepts are modeled as A p(s and zy = s stey
"https://gi thub. com/ 1.8x2001 /Audi evi sua 1 -inffrence
28

---

Published as a conference paper at ICLR 2026
respectively in case of common source, and ZA p(sA EA) and zv sv + Ey in case of separate sources. Here = SA = Sy represents their common location when C' = 1. while EΛ (0,,0% and Ev (V o represent the auditory and visual perceptual noise. While (med) (high) based on the a is assumed to be fixed, σv can assume three separate values (ov -v .o (experimentally manipulated) size of the visual stimulus Vievet. Finally, p represents "recalibration" factor to account for the fact that the range of auditory stimuli (30") is different from that of visual ones (40%.. In our experiment, this is the factor that differentiates the two models we set out to compare: in the first, we set 1; in the second, we set p 4/a (thus re-mapping auditory percepts to the same scale as visual ones).
Here we describe a BA trial, but the following is easily generalizable to BV ones. When asked about the location of the auditory stimulus, participants are assumed to consider both scenarios (common vs different sources) by evaluating
p(s C =) p(s [A
( C 2) =p(sA (A,
as well
PCC IA.
Psame).
The final estimate A of the location then inferred by weighting the two hypotheses (common vs separate sources) by their posterior probability, so
A =pccc=11 FA TV. OV. oS, Psame)
s. p(s C== 11ds+
PCC 2 FA FV, AA. TV. TS. Prame)
SA p(SA C =2)dsA-
Finally, the response of the participant is modeled as (SA. with a probability of X, and y Uniform[-45 45] with probability of A. Here À represents the "lapse rate", the probability of a participant being distracted/disengggeed and giving a random answer (which we fix at 0.02),
while σы represents motor noise.
Both models thus have free parameters, which we re- re-paaametrize as logo (low) (low) logo
(med) (mee))
logo (hght
(high)
logo logos, logoıs as and logitpsame for the purposes of simulation and model-fitting.
D.3.2 SIMULATION
For training all models, we produce 11.5 millions synthetic datasets. what follows we go through the simulation of single trial. As trials are independent from one another, generating more of them simply involves repeating this process.
Stimuli. Following the setup used in Liu et al. (2025), we sample $ Uniform(-15, - 10, -5, 0, 5, 10, 15) and C ~ (1i 8y Uniform[- 200 20 as a continuous variable (if C 2) we set sv (if C = 1). We then sample Vievel ~ Uniform (0, 1, 2}. representing the perceptual noise associated with sy. This (med)
Then we either sample
(low) regulates whether σv ov ov oi or σy o o (high)
Finally, we sample Tiype ~ Uniform (0, 1}. representing the task (BV if Ttype = 0, BA Ttype = 1).
Parameters. For each synthetic dataset, the prior generative distributions for the 7 free param- eters are Gaussians truncated at two standard deviations above and below the mean. We use Nuruncated (p, o²) to denote such distributions, with μ being the mean and the standard deviation. Similarly to empirical Bayes approaches (Murphy, 2023), we use wakiimum-likelhoon estimates of the individual participants' parameters from Liu et al. (2025) guide for setting these priors, so
29

---

Published as a conference paper at ICLR 2026
to generate realistic parameter ranges. The parameter distributions use in this work are:
logov
(low)
~Ninnaated (0, 1.5");
V
logo (med)
~ Vurueatte (ogoo (low)
1, 12);
V
r Iruncated logo (med) + 0.75, 0.5²);
(high)
logo
V
logo
A
Auuneated 1.75, 0.5?);
logas - 2.5, )²;;
logo ~ Muncattd (0, 0.5²);
logitpsame ~ N (1.,, 1.5²).
Note that logov logov (med) and logo (high) (high)
(low)
(med)
are not independent from each other, but carry the (med)
assumption that in most cases logor logo((ow)
(low)
logo (med) logo (high) which reflects the intent of the
experimental manipulation of Vievel-
Responses. Here we describe scenario in which Tiype = (BA trial), but the process is the same for Tyye O. In simulating the responses, we follow the hierarchical structure specified by the model. First we computed the sensory percepts a = p(8 + EA) and zy = 8y - Ev by sampling Ea NO, o and εy ot) We then evaluate 8 (recall we considering a BA trial) as in Eq. (7), and sample the final response as either ~ AS, dM) y Uniform 45). with probability regulated by the lapse rate À (which set 0.02, see above).
D.3.3 GROUND-TRUTH ACQUISITION
Here we describe how we obtained our log marginal likelihood (LML) estimates (in the form of lower bounds, see below), which we then use as ground-truth to compare our approach to baselines.
Problem setting. Fitting the cognitive model to a dataset involves finding the posterior over model parameters given empirical data and model
ply 0, X, k,P)((0)
p(0 y, X.
(8)
ply x)
where
(logov
(low) (low)
logo
(med)
(med)
loga
(high)
logo A, logas, logo loge
{8
(t) Veeer (t) 1400
Jtyp
and
44000
It=l-
Here / represents the trial number within the dataset (recall we are using data splits of 400 trials each, see Appendix D.2), and we set p(0) to the truncated Gaussians we use for sampling parameters in our simulation (see Appendix D.3.2), with probability density of values beyond the truncation boundaries set to a "floor value" " of N55 0).
While the posterior over parameters is often instrumental in answering scientific questions, the cru- cial quantity we are interested in estimating is the model evidence (also called marginal likelihood) p(y X, p) (i.e., the denominator in Eq. (8)). it represents straightforward metric for model selection. In fact, assuming a flat prior over models p(p = 1) p(p 4/3) 0.5, the model evidence as function of represents the unnormalized posterior over models.
Stacking Variational Bayesian Monte Carlo. To compute reliable estimate of the marginal likelihood to use as our ground-truth, we use Stacking Variational Bayesian Monte Carlo (S-VBMC, Silvestrin et al., 2025). This is principled approach to merge ("stack") approximate posteriors generated by a set of independent runs of its parent algorithm, Variational Bayesian Monte Carlo (VBMC, Acerbi, 2018; 2020). This is done in simple post-processing step, which has been shown to greatly improve the approximate posterior quality in a variety of challenging settings. In addition a posterior distribution, S-VBMC outputs estimate of the evidence lower bound (ELBO),
30

---

Published as a conference paper at ICLR 2026
which, the name suggests, is lower bound the (log) model evidence (Blei al., 2017), the quantity we are interested in for model comparison. As the approximation of the posterior approaches the true one, this quantity gets closer to the true model evidence, with equality when the approximation is perfect. As S-VBMC proved very effective computational neuroscience problems (Silvestrin al., 2025), including one very similar the considered here (Acerbi et al., 2018), we deem it a suitable method for estimating lower bound on model evidence to use as ground-truth.
While an in-depth description of S-VBMC and VBMC is beyond the scope of this work (an inter- ested reader should refer the original papers cited above), in the following paragraphs we briefly report details of our implementation of both.
VBMC implementation details. To obtain an approximate posterior, the Python implementation of VBMC (Huggins et al., 2023) requires absolute and plausible upper and lower bounds for each parameter. We use the sampling bounds defined in Appendix D.3.2 as absolute bounds, and replicate the process considering 1.5 standard deviations (as opposed to 2) from the mean to establish the plausible ones.
Another required input is a target density function (i.e., the unnormalized posterior), for which we use the numerator of Eq. (8), ply 0, x,X,,i p))(). We do this both with p = and p = 4/3, representing the two models set out to compare.
Finally, VBMC requires a starting point in the parameter space, which we uniformly sample between plausible bounds independently for each inference run.
S-VBMC implementation details. After obtaining converging VBMC runs for each of our 30 datasets (2 for each of the 15 participants, Appendix D.2) for both models, stack the resulting posteriors with S-VBMC. We maintain the default settings, therefore the only inputs required are the VBMC runs themselves. With this, obtain total of 60 "stacked" ELBOs (two per each dataset, corresponding to our two competing models) to use as ground-truth.
D.4 TABULAR MODEL DETAILS
This section describes the TabICL model and explains how the training dataset was generated. No- tably, the base architecture used for this tabular data example is different from the one used in the other experiments, highlighting the broad applicability of our method.
D.4.1 ARCHITECTURE
Set encoder. We reuse the first two stages of TabICL (Qu et al., 2025) without modification: the distribution-aware column processor (TFeol, implemented with induced self-attention blocks) followed by the context-aware row-wise transformer (TFm.) with RoPE. Scalars are mapped by a 1 - 128 linear layer; each column is then processed across rows by an ISAB stack (Lee et al., 2019) with three blocks, four heads, 128 inducing points, feed-forward hidden dimension of 256. The row-wise encoder has three layers with four heads, feed-forward hidden dimension of 256, and RoPE base 100,000, We prepend two [ (CLS) tokens per row and concatenate their outputs, yielding 128). We use most ten features per table.
a diddeinnnsioall row embedding
Tokenization and additive target encoding. The set encoder produces one row token per sample for context, buffer, and target rows (dimension 128; only selects the subset of the vector correspond- ing to the [CLS] token dimensions). Context and buffer tokens receive the target value additively via a small target encoder (linear 1212 Buffer tokens also receive a learned positional embedding indicating their autoregressive index (up to 32 positions). This keeps labels additive, lets us compute the set encoder once, and makes the buffer explicit at the token level.
Dataset-wise ICL with a buffered mask. On top these tokens we run a dataset-wise trans- former with twelve layers and four heads, model width 128, and feed-forward size 256. The atten- tion mask is the only architectural change relative to TabICL: context attends bidirectionally and is read-only at inference; the buffer uses strictly causal self-attention; target queries attend to the cached context and to the causal prefix of the buffer; there are no edges into context from buffer or targets. The maximum buffer size is 32 tokens and query 512 targets per task.
31

---

Published as a conference paper at ICLR 2026
Prediction head. Predictions GMM head with 20 components and minimum standard deviation of 10-3
Caching. The column and row set encoder is computed once for all rows. During autoregressive decoding we cache keys/values for the context once and update only the buffer cache, the same context cache is reused across parallel generations.
D.4.2 DATA GENERATION AND PREPROCESSING
SCM prior and task family. We generate datasets with the MLP-based structured causal model (SCM) prior in the style of Hollmann et al. (2023), following the dataset-wise set-encoded regime of TabICL. (Qu et al., 2025). Concretely, we first sample a DAG with layered (MIP-sttlle) connectivity and then define each variable
c f(Pa(c) E, where Pa(c) its parents, f is small MLP with nonlinearity, and e is independent noise. Unless stated otherwise, we sample the feature dimension d E [1, 10]. and per-task context sizes N E [8, 1024]; targets are continuous responses with datasct-ppecific noise levels. The cause sampler follows the TabPFN prior (including mixed marginals); the SCM therefore yields columns that may be non-Gaussian or discrete at source, which we handle with the TabICL preprocessing described below.
Sampling of task partitions. For each generated dataset draw random partition (C, T) with N ~ Uniform 8,..., 1024), buffer capacity fixed at K = 32, and target count M = 512. Per batch, we fix (d, N, M) across tasks avoid padding and stack samples directly.
Preprocessing. We adopt the TabICL PreprocessingPipellie and it on context features only. The fitted transform is then applied to context, buffer, and target features. Regression targets are standardized using context statistics, i.., 1 = (y Hy,c)/oy, /,,.. and the same (H, or) are used for buffer and targets. No missing values are synthesized by the SCM generator.
Summary of preprocessing pipeline. We use three-stage, per-column pipeline following Qu et al. (2025): (i) standard scaling: (ii) normalization (power, i... Yeo-Johnson); and (iii) outlier handling via a 2-score threshold 4.0. At transform time, values outside the fitted range are clipped to the training (context) min/max before normalization, mirroring TabICL's behavior.
D.4.3 TRAINING PROCEDURE
We train with AdamW (learning rate 10-4, 3= 0.95), weight decay 0.0). batch size 64 datasets per step, gradient clipping at 0.5, and dropout 0.0 throughout the backbone. Mixed-precision training uses AMP with bfloat16 6. All runs use float32 tensors at the data interface. A cosine schedule with warmup is used
warmup_steps 2000 takes precedence over the nominal warmup_ratis = 0.20; num_cycles = 1. Auto- matic mixed precision is enabled with of 64 independent tasks (datasets) with feature dimension d sampled from {1,.... 10] and context size N from (8. ....... 1024]; buffer size and target count are fixed at K K232 and M=512. Training is capped at max steps 160,000, i.e., one epuch effective duration. This corresponds to approxi- mately 64 x 160,000 10.24 million synthetic tasks seen during pretraining. The global data seed 123. We trained the model on a single NVIDIA A100 80 GB GPU for approximately 3 days.
ww..dyypeeblllatt....... 6. Each training step draws batch
D.5 EVALUATION DETAILS
In this paper log-likelihood values are always averaged (LL divided by the number of target points M).
GP & Sawtooth functions. We evaluate likelihood values over 1024 functions, each repeated 4
times with models trained on different seeds and context sizes N = 8, 16, 32, 64, 128 (statistics of 1024 x 4 x 5 evaluations). Each likelihood evaluation is average of 128 permutations (log averaged likelihood). In other words, we have 1024 x x averaged likelihoods,, and each averaged value merges 128 orders of the target set.
32

---

Published as a conference paper at ICLR 2026
Table AI: Head comparison on synthetic function. We compare average log-likelihood (f) results on our main GMM head and on standard Gaussian distribution head.
TNP-D
TNP w/ buffer K=4
AR
Ind
K=16
K=1
GP (M 16)
2.57 (0.020)
2.22 (0.022)
2.51 (0.019)
2.55 (0.019)
2.56 (0.019)
GP (M = 128)
3.29 (0.013)
2.15 (0.022)
3.27 (0.013)
3,28 (0.013)
3.29 (0.013)
Sawtooth (M 16)
1.05 (0.004)
0.94 (0.005)
1.00 (0.005)
1.08 (0.004)
1.09 (0.004)
Sawtooth (M 128)
1.15 (0.003)
1.16 (0.003)
1.15 (0.003)
1.16 (0.003)
1.16 (0.003)
TNP-D- DausD-Ga
TNP Gaussian w/ bufter
Ind
K=16
K=4
K=1
GP (M 16)
2.50 (0.019)
2.13 (0.023)
2.48 (0.019)
2.53 (0.019)
2.53 (0.019)
GP (M 128)
3.23 (0.013)
2.06 (0.023)
3.25 (0.013)
3.27 (0.013)
3.27 (0.013)
Sawtooth (M = 16)
0.96 (0.004)
0.82 (0.006)
0.85 (0.006)
0.98 (0.004)
0.99 (0.004)
Sawtooth (M = 128)
1.10 (0.003)
0.82 (0.005)
1.10 (0003)
1.11 (0.003)
1.11 (0.003)
EEG data. We train each model once with fixed seed; the evaluations are over 896 trials from 20 subjects held during training, each repeated with N 8, 16, 32, 64, 128, 192. For the EEG forecasting, the target set consists of time points immediately after context points, and, in the main results (Table 1), the target set permutations are applied, as done in Bruinsma al. (2023). We additionally demonstrate in appendix Table A3 that forecasting with permuted target set outperforms fixed sorted target. The number of permutations we apply 128.
Multisensory causal inference model. We train model for each setting of p (p and = 4/3). In the model selection scenario, the full 400-point dataset from each the 30 batches is used the target, and evaluate the LML across all cases. This procedure is repeated 5 times, with 128 different sequence permutations per run. the data prediction scenario, we first select the winning model from the model selection stage, and then compute log-likelihoods on the same 30 batches, each repeated with N = 8, 16, 32, 64, 128, 256. The results of both experiments are summarized in Table 2. Here also 128 permutation for all batches.
Tabular foundation model. We pretrain task-agnostic tabular model on synthetic data (Ap- pendix D.4) and evaluate on three UCI datasets: Individual Household Electric Power Consump- tion". Gas Turbine CO and NOx Emission Da Bike Sharing" Jena climate dataset !2 Power Consumption of Tetouan 1 and California Housing Price
For each dataset, evaluate likelihood values over 16 randomly sampled subsets. The context and target sets are set to N = 128, . 32. Each likelihood evaluation is an average of 128 permutations.
33

---

Published as a conference paper at ICLR 2026
Table A2: Average Log-likelihood (1) results on synthetic functions and EEG example. Supple- mentary results of Table on larger target set and various deployed K. When M >> K, we evaluate every K targets once and perform AR for M/K steps.
TNP-D
TNP-ND
TNP-A
Ind
GP (M 16) GP (M 128)
2.57 (0.020)
2.22 (0.022)
0.80 (0.082)
2.24 (0.018)
3.29 (0.013)
2.15 (0.022)
(0.023)
3.10 (0.012)
Sawtooth (M 16)
1.05 (0.004)
0.94 (0.005)
-0.43 (0.008)
0.98 (0.004)
Sawtooth (M 128)
1.14 (0.003)
0.94 (0.005)
0.39 (0.005)
1.12 (0.003)
EEG-Int (M 16)
0.51 (0.013)
0.36 (0.014)
0.46 (0.011)
0.58 (0.014)
EEG-Int (M 64)
0.88 (0.011)
0.35 (0.014)
0.50 (0.010)
0.95 (0.012)
EEG-For (M 16)
1,07 (0.004)
-0.74 (0.008)
-0.04 (0.005)
1.23 (0.003)
EEG-For 64)
1.12 (0.003)
-1.08 (0.007)
-0.23 (0.004)
1.20 (0.003)
TNP w/ buffer K=4
K=16
K=1
GP (M 16)
2.51 (0.019)
2.55 (0.019)
2.56 (0.019)
GP (M 128)
3.27 (0.013)
3.28 (0.013)
3.29 (0.013)
Sawtooth (M 16)
1.00 (0.005)
1.08 (0.004)
1.09 (0.004)
Sawtooth (M 128)
1.15 (0.003)
1.16 (0.003)
1.16 (0.003)
EEG-Int (M 166
0.52 (0.013)
0.54 (0.014)
0.54 (0.014)
EEG-Int (M 64)
0.90 (0.011)
0.91 (0.011)
0.91 (0.011)
EEG-For (M 16)
0,85 (0.004)
1.17 (0.003)
1.21 (0.003)
EEG-For (M 64)
1.12 (0.003)
1.18 (0.003)
1.19 (0.003)
E ADDITIONAL LOG-PREDICTIVE DENSUITY RESULTS ON SYNTHETIC AND
EEG TASKS
E.1 PREDICTIVE POWER OF DIFFERENT HEADS
In this paper. GMM prediction head. We compare the predictive performance of GMM to standard Gaussian distribution head. In Table A1, GMM is able to achieve better predictive performance, particularly the non-Gaussian Sawtooth functions.
E.2 RESULTS OF LARGER M
As supplementary results of Table 1, we evaluate log- values on larger target set. For TNP w/ buffer, we evaluate K points per Algorithm 2 and proceed to the next target subsets by conditioning the context and evaluated points. This requires M/ K steps of evaluations. The results are reported in Table A2. As we decrease the number of buffer targets K 15, the performance of our TNP w/ buffer becomes stronger, while more iterations (and thus computational time) are required.
pnttps://archive. ics. power4consumption http://hww...akooccccc///s//ssss "
emisssionrdtaHset
Pnttps://www, http://ww.a..oooaaaaaaaac//aaeea www.eiicss..........oo.
the buffer http:///ww....aaaaa..a.////////p
city
14 "https: :/tp::///wwww. kaggle e. com dat.satts/ ttlasetss/cammugee//aaaa - i innnnriii pprices ces
Note that when K = 1, our method is equivalent to standard TNP-D AR, as the actual number of points in
34

---

Published as a conference paper at ICLR 2026
Table A3: EEG forecasting w/ and w/o target set permutation. The target set of EEG forecasting are the points immediately after the context set. Our main paper applies permutation to the target set while this table compares against forecasting of fixed temporal order (sorted).
TNP-D
TNP-ND
TNP-A
AR
EEG-For (M 16)
1.07 (0.004)
-0.74 (0.008)
-0.04 (0.005)
1.23 (0.003)
EEG-For 16, sorted)
0.85 (0.005)
-0.74 (0.008)
-0.004 (0.005)
1.14 (0.004)
EEG-For (M 64)
1.12 (0.003)
-1.08 (0.007)
-0.23 (0.004)
1.20 (0.003)
BEG-For (M = 64, sorted)
0.89 (0.005)
-1.08 (0.007)
-0.23 (0.004)
1.16 (0.003)
TNP w/ buffer K=4
K=16
K=1
EEG-For (M 16)
0.85 (0.004)
1.17 (0.003)
1,21 (0.003)
EEG-For (M = 16, sorted)
0.76 (0.006)
0.87 (0.005)
1.09 (0.004)
EEG-For (M 64)
1.12 (0.003)
1.18 (0.003)
1.19 (0.003)
REG-For (M 64, sorted)
0.78 (0.005)
0.89 (0.004)
1.11 (0.004)
Table A4: Multisensory causal inference model selection extra results. Supplement for Table on model comparison case with extra evaluation on K = 4 and R2 metrics for LMI and ALMI.
TNP-D
TNP-ND
TNP-A
AR
Ind
LML RMSE (L)
3.10 (0.005)
86.96 (0.000)
208.51 (0.041)
4.75 (0.012)
ALML RMSE (1)
2.44 (0.008)
36.18 (0.000)
25.60 (0.023)
3.29 (0.019)
LML R (1)
1.00 (0.000)
-0.43 (0.000)
-7.22 (0.003)
1.00 (0.000)
ALML R² (1)
0.93 (0.001)
-1.44 (0.000)
-6.74 (0.014)
0.87 (0.002)
TNP w/ buffer K-4
K=16
K=1
LML RMSE (U)
3,56 (0.004)
3.48 (0.002)
3.47 (0.004)
ALML RMSE (1)
2.60 (0.010)
2.59 (0.009)
2.59 (0.011)
LML H² (1)
1.00 1.00 (0.000)
1.00 (0.000)
1.00 (0.000)
ALML R² (1)
(..2000001))
0.92 (0.001)
0.92 (0.001)
E.3 EEG FORECASTING w/ AND w/o TARGET PERMUTATION
In our main paper, the EEG forecasting task is evaluated with the permuted target set, following the procedure Bruinsma al. (2023). We repeat the experiment by forecasting the target of fixed temporal order. In Table A3, we show that averaging over random target order as done in the paper, provides better overall performance.
F ADDITIONAL MULTISENSORY CAUSAL INFERENCE MODEL RESULTS
As supplementary results to Table 2, we include additional metrics and evaluation settings. Specif- ically, for the model comparison task, report the coefficient of determination (R²) for both the LML and ALML with respect to the ground-truth (see Table A4). For the data prediction task, we present results with larger target size of M 128 (see Table A5). In addition, for completeness, we evaluate both the model comparison and data prediction tasks with K = 4. With varying K, we observe little to almost no performance degradation compared to TNP-D AR, especially for the data prediction case.
G ADDITIONAL TABULAR FOUNDATION MODEL RESULTS
We report results for intermediate context size (N 256) in Table A6. Consistent with main findings, AR w/ buffer matches standard AR within standard errors across all tasks, while both autoregressive methods outperform independent predictions forecasting tasks.
35

---

Published as a conference paper at ICLR 2026
Table A5: Multisensory causal inference model data prediction task normalized log-likelihood (1) results. Supplementary results of Table 2, with extra evaluation on K 4 and on larger target set M 128.
TNP-D
TNP-ND
TNP-A
Ind
Pred LL (M 16)
-2.76 (0.021)
-2.77 (0.025)
3.11
-2.66 (0.024)
Pred LL (M 128)
-2.71 (0.015)
-2.74 (0.016)
-3.17 (0.012)
-2.11 (0.015)
TNP w/ buffer K=4
K=16
K=1
Pred LL (M 16)
-276 (0.024)
-2.76 (0.024)
-2.76 (0.024)
Pred LL (M 128)
-2.11 (0.015) -2.71 (0.015)
-2.71 (0.015)
Table A6: Average log-predictivg density (t) results on UCI datasets with TabICL (N 256)..
Results are reported as mean and standard error over 16 randomly sampled mini-datasets (M 32) for interpolation (Int) and forecasting (For) tasks.
INTERMEDIATE CONTEXT REGIME (N 256)
Electric Cons.
Gas Turbine
Bike
Sharing For
Tetouan
Jena
Cali.
Int
For
Int
For
Int
Int
For
Int
For
Int
Independent
1.65 ans 1.21 0:3
-0.44 (.10
-1.06 mm
2.38 moni 1.98 min
0.56 mon
-1.45 0119)
2.03 (10 0.59 0139) -0.55 ()
Standard 1.67 0016 1.57 mxy
-0.44 (.16
-0.73 mm 2.39 mm) 2.24 mm 0.57 mmn
0.39 1m
2.03 mer
1.45 (r)5)
-0.54 pm
buffer 1.67 14 1.56
6.10
-0.76 m:::
m05
mom
mo
0.28
0.20
m.0
())))
-0.54 0.
H ABLATIONS AND EXTRA EXPERIMENTS
H.1 COMPARISON TO AON-PERMUTATIOON-NAARRA TRANSFORMERS
To isolate the effect of permutation invariance in the context set, we replace our model with plain autoregressive decoder Transformer that treats the context fixed input sequence. This sequential baseline performs substantially worse than our method the GP task and across context sizes (see Fig, A14), indicating that explicitly maintaining permutation invariance over the context set or at least part of critical for performance.
H.2 POSITIONAL EMBEDDINGS ABLATION
We also trained our method without positional embeddings in the buffer and performed evaluations with M K 16, shown in Table A7, and observed statistically significant difference in predictive performance. This aligns with findings in causal-transformer work showing that models can infer positional structure without explicit encodings (Haviv et al., 2022; Irie, 2024; Zuo al., 2025). While not strictly required, positional embeddings may still support future extensions, such as scaling to larger buffer sizes via ALiBi (Press al., 2022) or RoPE (Su et al., 2024).
Table A7: Average joint predictive log- degsity (t) for positional embedding ablation on the GP task; reported mean (SEM).
TNP-D-AR TNP-D-Ind 2.57 (0.02) 2.22 (0.02) 0.80 (0.08) 2.24 (0.02)
TNP-ND
TNP-A
Ours w/ pos. emb
Ours w/o pos. emb
2.51 (0.02)
2.51 (0.02)
H.3 NUMBER OF SAMPLES ORDER AVERAGING ABLATION
We study the effect of the number of sequence samples (permutations) used for order averaging. We report results the multisensory causal inference task, where method (with buffer size K = 16) is used to compute the LML of a dataset by averaging over multiple permutations. As shown in Fig. A15, increasing the number of permutations reduces the RMSE of the estimated joint
36

---

Published as a conference paper at ICLR 2026
M
M 4
M = 16
3.2
Joint log- log-density
2.4
1.6
0.8
0.0
128
32
Number of context points N
TNP buffer, K=M +- TNP-D-AR Decader Transformer
TNP bufler, K=1
Figure A14: Average joint predictive log-density (1) on the GP regression task comparing the De- coder Transformer models with and gold standard TNP-D-AR on varying number of context points N and number of targets M = 1, 16.
LML
ALML
RMSE
16
32
128
16
32
64
128
N. of sequence samples
of sequence samples
Oues
TNP-A
TNP-D-AR
Figure A15: Average RMSE (4) on the LML (left) and ALML (right) estimation in the multisensory causal inference task for different numbers of sample permutations.
density relative to the true joint at a rate comparable to existing autoregressive baselines, while our method remains significantly faster. This indicates that for proposed method, order averaging does not introduce additional performance degradation relative to the gold standard baselines for given number of permutations.
H.4 EXTENSION TO LATENT BOTTLENECKED ATTENTIVE NEURAL PROCESSES MODEL
To assess the generality of the proposed autoregressive buffer, we integrate it into a perceiver-style Latent Bottlenecked Attentive Neural Processes (LBANP) architecture (Jaegle et al., 2021; Feng et al., 2023). The context set is first encoded into fixed-size latent array, and the autoregres- sive buffer operates over targets on top of this latent array bottleneck. We evaluate this BNP with buffer model on the GP regression task with 4 and 16 targets. As shown in Fig. A16, the LBANP equipped with our autoregressive buffer matches or slightly outperforms standard autoregressive deployment of the Perceiver architecture (when K 1). This result is likely due to the fact that the buffer allows the model to explicitly represent the recent history of targets, bypassing the compressed representation of the context for immediate short-term dependencies, thus slightly enhancing pre- dictive performance. These results demonstrate that our method extends naturally to bottlenecked / perceiver-style architectures, supporting its generality beyond full-attentic models.
37

---

Published as a conference paper at ICLR 2026
M
16
2.4
Joint log degg-denttyy
1.6
0.8
0.0
32
64
32
Number context points
LBANP w/ butter, K=M
LBANP
LBANP
Figure A16: Autoregressive buffer extension Latent Bottlenecked Attentive Neural Processes (LBANP) model. Average joint predictive log-density (t) on GP with varying number of context N and number of targets M 4, 16.
N=#
N= 8
N= 16
Joint log donsity
0
16
16
64
16
32
64
N=
N= 128
3
Number of targets (M)
w/ buffer, K-Targets
TNP.D.AR
Figure A17: Average joint predictive log-density (f) on the GP task across number of targets M up to 64, with buffer size K set equal to M, for varying numbers of context points N from 4 to 128.
H.5 BUFFER SIZE ABLATION
We evaluate the effect of the buffer size K on GP regression task, training a model with maximum buffer size K = 64. As shown in Fig. A17, the performance of our method remains stable across this range and does not degrade relative to the autoregressive bascline, indicating that increasing K up 64 does not harm predictive quality.
I USE OF LARGE LANGUAGE MODELS
Idea generation and exploration. We used Large Language Models (LLMs) the early stages of this work support idea generation, brainstorming, and the exploration of possible methodological
38

---

Published as a conference paper at ICLR 2026
directions. LLMs were also employed for tasks such identifying related work through web search and summarization, which helped us gain an initial overview of relevant literature.
Coding assistant. LLMs provided assistance with coding, primarily by generating boilerplate components of the codebase, visualization scripts, and test codes. They were also used for drafting parts of the implementation in PyTorch. All code produced or suggested by LLMs was carefully reviewed, verified, and modified where necessary to ensure correctness and reliability.
Writing assistant. Finally, LLMs were used in preparing the manuscript, particularly for refining clarity, conciseness, and grammatical correctness. They supported rephrasing and restructuring of text, helping us to communicate ideas more effectively while maintaining the accuracy and integrity of the content.
39


Reviews:
[{"review_id": "rkIpRXGsJk", "reviewer": "Reviewer_jkX4", "summary": "Neural processes are finding applications across a range of areas, including in tabular foundation models like TabPFN and in weather modelling e.g. in Aardvark Weather. In many of these applications, it’s important to be able to model joint predictive distributions that capture the dependencies between target variables at multiple locations e.g. when imputing multiple missing values into a row of a tabular dataset or when predicting precipitation across a region to assess flood risk in weather forecasting.  Autoregressive transformer neural processes are arguably the go-to method for solving such tasks.\n\nThis paper makes a neat contribution: it addresses one of the central limitations of autoregressive transformer neural processes and sequence models -- their high computational cost at inference time. It introduces a lightweight causal buffering mechanism that preserves the expressive power and calibration benefits of full autoregression while dramatically reducing computation and latency. This makes autoregressive probabilistic models viable for real-time and large-scale applications, from time-series forecasting to neural operator learning. In doing so, the paper bridges the gap between expressive but slow autoregressive methods and fast but less flexible parallel inference, providing a practical path toward scalable, uncertainty-aware inference in scientific and foundation-model settings.\n\nI liked the paper and would vote for acceptance.", "questions": "Small point, but I’m not a huge fan of the phrase “joint likelihood” since in many contexts in the paper the primary focus is really producing a \"joint predictive density” rather than the likelihood of the parameters need for e.g. learning. Indeed, if you trained a model only using likelihoods derived from univariate predictives, you could then immediately use this to produce joint predictive densities (c.f. the AR CNP paper). So I'd place the emphasis on joint distributions most of the time, rather than on the likelihood functions.\n\nline 49: \"However, this breaks the set-based structure”  I think that, at this stage in the paper, this is a bit ambiguous. E.g. in standard mode, adding the generated targets back into the context and recomputing everything is arguably as good as it gets in terms of permutation invariance. Perhaps say “However, this involves significant computational overhead…”\n\nline 161: \"In practice, Eq. (2) is not exact for likelihood evaluation as it breaks permutation invariance of the model.” I don’t agree with this. This is really a modelling choice, or something that is determined by the task itself, rather than being wrong or right. You can choose to model the order as latent (and therefore something that you need to average over) or known and fixed (in which case you don’t). For time-series, for example, there is often an implied order over the targets. For spatial data, there might not be. \n\nline 291: you could point out here that discrete diffusion is really an any order AR model in disguise and is very closely related to a NP with a discrete input set (arguably it is an NP) and as such could leverage your ideas.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper has at its core a simple but neat idea that works well. I liked the fact that it had a wealth of experimental results including\n\n1. Synthetic function modelling (Gaussian Process (GP) functions and sawtooth synthetic functions.)\n2. EEG forecasting and interpolation including real EEG recordings from 20 human subjects\n3. Multisensory causal inference modelling involving simulated and real human behavioural data from a multisensory integration study (Liu et al., 2025).\n4. Tabular foundation model experiments\n5. Further Efficiency and ablation studies\n\nThe results were strong -- the method really works.\n\nThe presentation and the writing were very clear and well polished. Figures were very well presented. Generally, I thought it was a well executed piece of work.", "weaknesses": "The idea is quite simple which could be viewed as a weakness, but I actually view this as a strength for multiple reasons: it's easy to understand, it's simple to implement, and because of this it could be deployed widely. \n\nThe idea might also seem niche as AR-TNPs are not super well-known, but this is also not correct. TNPs in the sim2real setting were rebranded by the impactful TabPFN line of work and the contributions here are directly aligned with this breed of tabular foundation model as the final experiment shows. Moreover, NPs have been deployed in another Nature paper in Aardvark weather, which has also been impactful and which could leverage these results to produce scalable weather forecasts. I expect more and more examples of applied NPs will emerge over the coming years.", "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": "n1YDqaQXb4", "reviewer": "Reviewer_G5io", "summary": "The paper proposes a “causal autoregressive buffer” to accelerate joint inference in transformer-style probabilistic models such as Transformer Neural Processes (TNPs) and Prior-Fitted Networks (PFNs). The idea is: instead of repeatedly re-encoding the entire context set at every autoregressive step, the model encodes the context once, caches it, and then maintains a lightweight causal buffer of generated targets. This is claimed to both (i) avoid recomputing attention over the full growing set, and (ii) still model target-to-target dependencies during sequential rollout. The authors argue this reduces complexity from  \n$O(K (N+K)^2) \\) to \\( O(N^2 + NK + K^2) $ \nand yields up to 20× faster joint sampling while matching predictive quality of stronger autoregressive baselines, across synthetic regression, EEG forecasting, and tabular regression tasks.", "questions": "see in weaknesses", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "### 1 Practical inference speedup for probabilistic transformers  \nThe work directly targets a real bottleneck: autoregressive joint inference in set-conditioned probabilistic models is extremely expensive if you have to repeatedly re-encode the entire (context ∪ generated targets) set at every step. The proposed causal autoregressive buffer reduces the need to rerun the full encoder every time by caching the context representation once and then only extending a lightweight causal buffer of generated targets. This reduces the stated complexity from \\(O(K (N+K)^2)\\) to \\(O(N^2 + NK + K^2)\\) for K rollout steps, which the authors claim translates into 3×–20× wall-clock speedups in their experiments.\nThis is a meaningful systems win for anyone who wants to deploy TNP/PFN-style models in real-time or interactive settings.\n\n### 2 Keeps target-to-target dependency modeling (vs. fully independent heads)  \nA common cheap fallback in meta-regression / neural process style models is to just predict each target independently conditioned on the same context encoding. That is fast but discards correlations across targets.  \nHere, the buffer still allows each new target to attend causally to previously generated targets in the rollout, so you capture sequential statistical structure across predicted outputs without paying the full re-encode cost. This is especially relevant for tasks like trajectory / time-series forecasting (e.g. EEG forecasting) where temporal consistency matters.  \nIn other words, you get something closer to coherent joint samples instead of just pointwise fits.\n\n### 3 Unification: one model that can act both “independently” and “autoregressively”  \nThe paper proposes a training curriculum in which, during training, some targets are forced to condition only on the static context block (independent prediction mode), and others are allowed to attend to a prefix of previously generated targets through the buffer (autoregressive mode). The same set of parameters is trained to handle both regimes.  \nIf this actually works robustly, it’s attractive from an engineering perspective: you don’t need to maintain two separate heads (one batchable + independent, one fully autoregressive). You get a single model that can cheaply do fast batched prediction or slower-but-more-coherent sequential rollout.\n\n### 4 Conceptual bridge between permutation-invariant context encoders and standard AR decoding  \nArchitecturally, the paper is trying to marry two worlds:\n- permutation-invariant / set-encoder style context processing (as in Neural Processes, PFNs, etc.), and\n- standard causal decoding with KV caching (as in language models).", "weaknesses": "## 2. Major Concerns\n\n### 2.1 Claimed novelty vs. standard KV caching  \nThe proposed “causal autoregressive buffer” is essentially a block-structured attention mask plus caching of keys/values:\n\n- A *static* context block (permutation-invariant encoder over the observed context set) is computed once and cached.\n- A *growing* autoregressive buffer stores keys/values for previously generated targets and is updated step by step with a causal mask.\n- During generation, new target tokens attend to (a) the frozen context block and (b) previous targets in the buffer, but the buffer cannot “write back” into the context block.\n\nThis is very close in spirit to ordinary autoregressive Transformer inference with KV caching: you prefill a prefix once, then incrementally extend the cache with each new token under a causal mask. \n\nThe paper repeatedly describes this as a general architectural mechanism that “decouples context encoding from sequential prediction,” but standard decoder-only Transformers already do this: prefix context is cached, new tokens causally attend to the prefix and prior tokens. The only real twist here is that the “prefix” came from a permutation-invariant encoder instead of from plain tokens, and that the mask enforces a hard separation between context and generated targets (“R1–R4”).\nRight now, the paper oversells this as conceptually new. It reads more like applying known KV-caching style inference to set-conditioned meta-learners. The authors need to justify why this is more than “cache the context encodings and use a causal mask for the generated block.”\n\n---\n\n### 2.2 Theory is bookkeeping, not analysis of approximation error  \nThe main “theoretical” contribution is a complexity argument: naive autoregressive deployment of a TNP-like model requires re-running attention over the enlarged context-plus-buffer at every step, costing \\( O(K (N+K)^2) \\), while the buffer method pays \\( O(N^2) \\) once for the context and then \\( O(NK + K^2) \\) to roll out K targets.\n\nThis is essentially computational accounting, not a theoretical guarantee about inference quality. The paper makes a strong claim that the buffered rollout “preserves model quality,” because each new prediction can still attend to all prior targets via the buffer and to the cached context. But there is no formal analysis of when this approximation matches the behavior of the “true” fully-updated model that would have re-encoded the augmented context set after each new target is added.\n\nIn fact, once you stop re-encoding, you’ve thrown away strict permutation invariance over the *union* of context and generated targets. After generation begins, the model is no longer allowed to revisit and re-symmetrize the combined set. The paper acknowledges degradation when the number of autoregressive targets exceeds the buffer size K used during training.\nThis is a core limitation, but it’s treated as an implementation detail rather than a fundamental modeling gap.\n\n---\n\n### 2.3 Empirical evaluation seems arranged to flatter the method  \nMost of the storytelling is about runtime: joint sampling speed, likelihood evaluation speed, etc. The new buffer method is reported as 3–20× faster than autoregressive baselines like TNP-A or TNP-D-AR, which repeatedly re-encode context at each step, especially at large N.\n\nConcerns:\n\n- The paper admits their implementation uses heavy engineering (FlashAttention-2, Triton kernels, KV caching), and also says baselines were “optimized beyond their public versions.” But we never see an ablation that isolates the architectural contribution from pure systems tuning. We can’t tell if the 20× number is “new idea” vs. “better CUDA.”\n- There is no memory comparison. The buffer maintains per-layer KV caches for both the frozen context block and the autoregressive buffer tokens, which can blow up for long rollouts and large batch sizes. Only wall-clock time is emphasized, which is convenient but incomplete.\n- When you actually look at predictive quality (e.g. EEG forecasting), the buffered model with K=16 can underperform the slow full autoregressive baseline (TNP-D-AR), and sometimes is even worse than extremely small-buffer variants (effectively K=1).\n  So the tradeoff is not “same accuracy, way faster.” It’s “sometimes noticeably worse, but faster.”\n\nDespite that, the abstract and intro still assert that the method “matches predictive accuracy … while delivering up to 20× faster joint sampling,” which is too strong given these cases.\n\nIn short, the evaluation is narrated as “little to no quality loss, huge speedup,” but the actual numbers show a real quality vs. speed tradeoff.\n\n---\n\n### 2.4 Generality claims are speculative and weakly supported  \nThe paper claims broad applicability: Perceiver-like architectures with pseudo-tokens, probabilistic neuroscience modeling, and “tabular foundation models” such as PFN-style inference.\n\nBut:\n\n- The Perceiver extension is hypothetical; there is no experiment demonstrating the buffer with learned pseudo-tokens.\n- The “tabular foundation model” experiment is actually quite small. The model is trained from scratch on synthetic structural causal model data, then evaluated on a few UCI-like tabular tasks with N=128 context / M=32 targets.\n  That’s nowhere near the scale implied by the phrase “foundation model.”\n- In those tabular results, “AR w/ buffer (K=32)” performs roughly on par with a standard AR baseline (K=1) and modestly better than fully independent predictions.\n  That mostly shows that causal conditioning across targets helps, which is expected. It does not prove that the specific buffering trick is uniquely enabling.\n\nSo the claim that this mechanism is a general, broadly applicable inference upgrade feels speculative. The evidence is narrow, and in some cases purely suggestive.\n\n---\n\n### 2.5 Missing baselines / missing ablations  \nThe baselines are primarily TNP variants (TNP-D, TNP-D-AR, TNP-A, etc.).\nHowever, to really argue “our method is necessary,” I would expect at least:\n\n1. **A simple two-tower baseline**  \n   - Tower A encodes the context set once (frozen).\n   - Tower B is a causal decoder over targets that cross-attends into Tower A, with standard KV caching.\n   - This is, in spirit, what the proposed model is.  \n   Without ablations on masking structure and curriculum, we don’t know if the fancy buffer rules (R1–R4) are actually critical or if a trivial cross-attend-decoder would do the same.\n\n2. **A plain autoregressive Transformer treating the context as a prefix**  \n   - Just linearize the context into a token sequence, feed it as a prefix prompt, and then autoregress over targets exactly like a language model, ignoring permutation invariance entirely.\n   - The paper insists that permutation invariance of the context block is crucial, but it never quantifies how much you lose if you drop that constraint and just go full decoder-style.", "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": "HB7FycKLHd", "reviewer": "Reviewer_Jq2K", "summary": "The paper addresses a conflict between two types of models: set-based probabilistic models and autoregressive models. When models (like Neural Processes) try to generate a sequence of predictions (e.g., filling in a signal), they have to re-process all the original data and all the new predictions at every single step. The core idea is to decouple the static context from the dynamic predictions by introducing a new architectural component called the \"causal autoregressive buffer\". The advantage is (i) on the computational advantage: The expensive $\\mathcal{O}(N^2)$ cost (encoding the context) is paid only once at the beginning and (ii) preserved accuracy as shown in the experiments section.", "questions": "1. What happens when the target sequence length M is larger than the buffer size? I saw a discussion on Appendix E.2 and also the limitation in the final discussion section, which suggests that to achieve better performance, we need more computational time. I'm wondering what the final complexity is if including the operation \"we evaluate every $K$ targets once and perform AR for $M/K$ steps\". Will this still lead to the 20 times speed up? How often does the situation \"target count exceeds training bounds of the buffer\" happen in real-world analysis?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and novel. The problem of repeated, expensive re-encoding is visualized in Figure 1, which provides an intuitive understanding of the entire paper's premise. The proposed solution is an elegant architectural fix that decouples the static context from the dynamic predictions, supported by a clear complexity analysis ($\\mathcal{O}(N^{2}+NK+K^{2})$). \n\n2. The empirical evidence is convincing. It delivers a speedup (up to 20x) with small additional training cost and without sacrificing predictive accuracy.", "weaknesses": "I'm not an expert in this field (neural processes, probabilistic meta-learning, amortized inference), so I'll leave my comments in the following Questions sections.", "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": "qm9439MPnw", "reviewer": "Reviewer_agRA", "summary": "The paper introduces a causal autoregressive buffer for transformer-based amortized probabilistic models (NPs, PFNs, tabular FMs). The key idea is to encode the context once and keep it in a static cache; generated targets enter a causal buffer that attends to the context cache and prior buffer tokens, enabling joint dependencies without re-encoding the augmented set each step. Complexity drops from (O(K(N!+!K)^2)) to (O(N^2 + KN + K^2)), supporting batched AR sampling and single-pass joint likelihood evaluation. Experiments across synthetic, EEG, cognitive, and tabular tasks match standard AR accuracy while achieving speedups.", "questions": "- Is it possible to have tests on larger (K) regimes and very large contexts to map out failure boundaries, plus ablation on positional embeddings inside the buffer and ordering effects when approximating permutation invariance via order-averaging?\n - Is it possible to quantify order-averaging required to stabilize joint likelihoods with a metric measuring their cost/benefit?\n- How do the authors decide on the size of (N)s? Is it there a experiment for emprically showing for large values of N, proposed method has compounding advantages?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This is a practical architectural improvement that preserves set-conditioning benefits while unlocking efficient joint predictions typical of AR models. The idea of a role-aware attention mask (immutable context, causal buffer, no writes back to context) feels simple yet powerful, and widely applicable to TNPs/PFNs/tabular FMs. The potential for large efficiency gains in meta-learning scenarios with repeated sampling is significant.\n\n- The factorization and masking constraints (R1–R4) are clearly stated; buffering semantics are sound. The complexity analysis matches the masking structure.\n- Training uses a buffer-size curriculum with structured masks that let targets attend to context or a variable prefix of the buffer, aligning training and inference. The link to minimizing KL to posterior predictive under varying conditioning sets is consistent with PFN theory.\n- The claims of accuracy comparison with standard AR and speedups are supported on small/medium contexts and (K< 32) (as per reported settings). The authors also acknowledge a degradation risk when target counts exceed the trained buffer size.", "weaknesses": "In the experiments, they are rather limited to moderate sizes of K and small/medium-sized contexts. I have a feeling that this may be limiting the characterization of the approach. Also, residual order dependence for likelihood evaluation requires averaging; I think an analysis on the order sensitivity is missing.", "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)