Architectural and Inferential Inductive Biases for Exchangeable Sequence Modeling
Authors: Daksh Mittal
Ang Li
Thomson Yen
C. Daniel Guetta
Hongseok Namkoong
Venue: NeurIPS 2025 poster ()
Keywords: Exchangeability
Inductive Bias
Epistemic uncertainty
Multi-step inference
Transformer architecture
Primary area: deep_learning
TL;DR: We investiage the use of autoregressive models for exchangeable sequences in decision-making, showing multi-step inference improves decision making and standard causal architectures outperform existing custom ones.

Abstract: Autoregressive models have emerged as a powerful framework for modeling exchangeable sequences---i.i.d. observations when conditioned on some latent factor---enabling direct modeling of uncertainty from missing data (rather than a latent). Motivated by the critical role posterior inference plays as a subroutine in decision-making (e.g., active learning, bandits), we study the inferential and architectural inductive biases that are most effective  for exchangeable sequence modeling. For the inference stage, we highlight a fundamental limitation of the prevalent single-step generation approach: its inability to distinguish between epistemic and aleatoric uncertainty. Instead, a long line of works in Bayesian statistics advocates for multi-step autoregressive generation; we demonstrate this "correct approach" enables superior uncertainty quantification that translates into better performance on downstream decision-making tasks. This naturally leads to the next question: which architectures are best suited for multi-step inference? We identify a subtle yet important gap between recently proposed Transformer architectures for exchangeable sequences (Müller et al., 2022; Nguyen & Grover, 2022; Ye & Namkoong, 2024), and prove that they in fact cannot guarantee exchangeability despite introducing significant computational overhead.  Through empirical evaluation, we find that these custom architectures can significantly underperform compared to standard causal masking, highlighting the need for new architectural innovations in Transformer-based modeling of exchangeable sequences.

Architectural and Inferential Inductive Biases for Exchangeable Sequence Modeling
Daksh Mittal, Ang Li; Thomson Yen; Daniel Guetta, Hongseok Namkoong
Columbia University
{dm3766, a14263, ty2531, crg2133, hn2369) Ocolumba...... edu
Abstract
Autoregressive models have emerged powerful framework for modeling exchangcable sequences-iie -i.d. observations when conditioned on some latent factor-enablin direct modeling of uncertainty from missing data (rather than a latent). Motivated by the critical role posterior inference plays as a subroutine in decision-making (e.g., active learning, bandits), we study the inferential and architectural inductive biases that are most effective for exchangeable sequence modeling. For the inference stage, highlight fundamental limitation of the prevalent single-step generation approach: its inability to distinguish between epistemic and aleatoric uncertainty. Instead, long line of works in Bayesian statistics advocates for multi-step autoregressive generation; we demonstrate this "correct approach" enables superior uncertainty quantification that translates into better performance on downstream recision-making tasks. This naturally leads the next question: which architectures are best suited for multi-step inference? We identify a subtle yet important gap between recently proposed Transformer architectures for exchangeable sequences Müller et [22,, Nguyen and Grover [23, Ye and Namkoong [30), and prove that they in fact cannot guarantee exchange- ability despite introducing significant computational overhead. Through empirical evaluation, we find that these custom architectures can significantly underperform compared to standard causal masking, highlighting the need for new architectural innovations in taaasfoomer-bbseee modeling of exchangeable sequences.
Introduction
Intelligent agents must be able to articulate their own uncertainty about the underlying environment, and sharpen beliefs as it gathers more information. However, uncertainty quantification can be difficult in general without a complete characterization of the nature of the data and sources of uncertainties. Exchangeable sequences represent an important class of data structures on which one can quantify both the nature and degree of uncertainty principally, thereby enabling more robust decision-makinn algorithms. To characterize this concretely, consider a sequence of observations Y1:00 gathered from an unseen environment 0, e.g., noisy answers in math quiz, generated by student's is current proficiency level 0. When marginalized over the latent 0, the joint distribution of the sequence Y1ooo is permutation invariant. This property defines what are known exchangeable sequences, which serve as a fundamental unit of study in uncertainty quantification of latent variables that govern data generation.
Autoregressive sequence modeling has recently gained significant attention as powerful approach for modeling exchangeable sequences Y1..o 22 23
Unlike conventional Bayesian modeling which requires specifying prior over an unobserved latent variable and likelihood for the observed data, often a challenging autoreregeee sequence modeling builds on De Finneti's
30
20
Equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

predictive view of Bayesian inference 6 8 9..
This view directly models the observables Yi: for an exchangeable sequence, epistemic uncertainty the latent variable 0 stems from the unobserved future data 2 4 12]. Viewing future observations as the sole source of epistemic uncertainty in 0 for exchangeable sequences (2 4 [2], autoregressive sequence modeling enables direct prediction of the observables Y1:00- offering a conceptually elegant and practical alternative to conventional Bayesian approaches.
Transformers have emerged as the dominant architecture for autoregressive sequence modeling 120 22
17 15 29,
owing to their remarkable performance in natural language and vision applications [5] 10]. As Transformers are increasingly employed to meta-learn probabilistic models for large-scale tabular datasets [32 20), they offer unique opportunity to move beyond traditional prediction tasks or merely replicating supervised algorithms- an area that has been the primary focus so far. Instead, following De Finetti's 's perspective, when meta-trained on tabular datasets, these models can effectively quantify epistemic uncertainty, which powers peccision-making and active exploration across diverse domains, including recommendation systems, adaptive experimentation, and active learning B1. For instance, we can train sequence models collection of tables, each representing different disease diagnosis setting-akin to applying meta-learning in the context of disease diagnosis. These models can then be leveraged to actively gather additional data in a previously unseen disease diagnosis setting to enhance model's predictive performance (see Figure I for illustration).
However, using Transformers to model exchangeable sequences for decision makiing presents its own challenges. Since the existing literature has primarily focused on traditional prediction tasks or the replication of supervised algorithms rather than decision-making. it has overlooked the perspective that epistemic (reducible) uncertainty stems from missing data. Accurately distinguishing between epistemic (reducible) and aleatoric (irreducible) uncertainty is crucial for decision-orienten applications. significant limitation in the current literature is its predominant focus on one-step predictive uncertainty 122, 23 one-step predictions [18), which fails to differentiate between epistemic and aleatoric uncertainty. contrast, long line of work in Bayesian statistics [28] 24
advocates multistep inference, offering more robust framework for differentiating between epistemic and aleatoric uncertainty.
Our first key contribution is to empir- ically and theoretically demonstrate the limitations of one-step inference in
Collection of Tables [e.g. each representing different disease diagnosis sett rg)
Contribution 1: How do interence/ make decisions using sequence models?
sequence models for uncertainty quan- tification (Section 3 Specifically, we show that one-step inference leads to
112
Irain
1-1
loss of information (Theorem 2 which degrades uncertainty estimation (Figure 3) and results in suboptimal performance in downstream decision- making tasks such as multi-armed ban- dits and active learning (Theorem 3 Figure 3). We address this by using multi-step autoregressive generation in sequence models. Our findings are consistent with those of Zhang et al. 31.. who also advocate for multi- step inference to improve active ex- ploration via Thompson sampling.
3.0
Ours [Multi stepl
Autoregressive sequence model
0.0
Prior work (Ore-step)
5.8
Inferance
Active Learning (decision-making task) (e.g. which individuals (X) should we conduct expensive test and get abels r)
10 Numrerr Rewedeed Label
Contribution 2: What Transformer
architectures good for training this data?
Figure 1: (Left) Meta-learned sequence models be used for de- cision making for e.g. active learning (Right) Our contributions (1) multi-step outperforms one-step (2) Investigate which transformer architectures are good for training exchangeable data
This brings us to the next key question:
What kind of architectures should we use for modeling exchangeable sequences, particularly when performing multi-step inference? Architectural choice is a crucial source of inductive bias, and existing approaches have attempted to incorporate exchangeability as an inductive bias through specialized masking strategies [22] | 30 However, a critical gap in the literature remains: these prior works conflate exchangeability with the invariance properties enforced by their masking schemes that don't necessarily guarantee valid probabilistic inference.
Another key contribution of our work is in articulating this gap. We introduce conditional per- mutation invariance (Property I as a way to formalize the invariance enforced by these masking
2

---

strategies. Correcting previous works 22 23 30 that implicitly assume their architectures achieve exchangeability, we show that enforcing conditional permutation invariance alone insufficient to guarantee full exchangeability. Specifically, existing masking-bass based approaches do not ensure the conditionally identically distributed (c.i.d.) property (Property 2 which essential for the validity of probabiliitic inference in exchangeable sequence models. As a result, despite their intended design, such models may still violate exchangeability when used in practice (see Section 4 By clearly distinguishing between exchangeability and the weaker notion of conditional permutation invariance, our work not only clarifies existing misconceptions in the literature but also establishes more rigorous foundation for exchangeable sequence models.
Moreover, a significant drawback of this masking scheme, as discussed in Section 4 is that introduces computational overhead without yielding any tangible improvements in model performance. To evaluate its impact, we empirically assess the effectiveness of enforcing conditional permuttttonono invariance (Property[ [T) within the model architecture, comparing it to standard causal masking, which does not enforce permutation invariance. Surprisingly our results show that enforeing Property [] not only fails to provide any performance benefits but actually performs worse than causal masking (Figure 5). These findings underscore the need for new research directions to develop more effective inductive biases for Transformers in exchangeable sequence modeling.
2 Conceptual background
In this section, we first review De Finetti's 's predictive view of Bayesian inference 6 8 9 autoregressive sequence modeling of exchangeable sequences can power it. For most of our discussion, we focus on the setting where observations are given by Y1:0c. However, this framework can be easily extended to contextual settings where observations are (X1:s0+ Y1:s)), with X serving as the context (Section C). We start by reviewing the conventional Bayesian modeling paradigm, where the modeler posits a latent parameter 0, along with a prior u(0), and likelihood (0 0) PY The joint probability of observations Yiso is then expressed as PY V):,>)) - Po,Y, yt. (d(0)
and how
20
Given any observable data Yı., the posterior over latent parameter expressed as u(-lYi.e). Note that pu-| iit) represents the epistemic (reducible) uncertainty, which gets resolved as more data is collected while the likelihood P (0) represents the alcatoric uncertainty and comes due to inherent randomness in the data.
Predictive view of uncertainty. Instead of positing explicit priors and likelihoods a proposed la- tent parameter space, we consider a different probabilistic modeling approach where we directly model the observable Y1.00 without explicitly relying on any latent parameter. This view heavily relies on the infinite exchangeability of the sequence Yisoc. defined as: IP(Y1,... Y,) = P(Yx(1), (()) Ya(n)) ) for any and permutation æ. De Finetti's 's theorem states that if an infinite sequence is exchangeable
then the sequence can be represented as a mixture of i.idd random variables.
Theorem (De Finetri's theorem). If sequence Yuo is infinitely exchangeable then there exists latent parameter and unique measure pd.) over 8 such that, for any n
P(Y = =Vi.n)
P(Y = =v.0)p(d0)) 1=1
In addition justifying conventional Bayesian modeling, De Finetti's theorem also establishes that for infinitely exchangeable sequence Y1.0o. the epistemic uncertainty in the latent parameter 0 in Theorem aaiise solely from the unobserved Y1.o [ 13 [6] In other words, epistemic uncertainty in is the same as predictive uncertainty in Y1:0o- De Finetti [8) Hewitt and Savage [19] in fact show that the latent parameter 0 in Equation ( is entirely function of the Y1.0o, that is,
0 = =((Yr.)). To illustrate, consider coin B from Figure and suppose it is tossed repeatedly. Let y. denote the outcome of t-th toss of coin B where y,(²) = 1 for heads, and yr(B) 0 for tails.
The sequence (T) ..... (B)) } is exchangeable, since its probability distribution depends only on the outcomes of the tosses, not their order. In this case, De Finetti's Theorem holds trivially, and the
latent parameter represents the probability of obtaining 'heads' when the coin B is flipped. Further,
this parameter 0 can be estimated as 0 lim 1 Ei=,Y, x(())
3

---

By abusing notation and writing = P(-10) where 8= = f(Yise)), we can interpret P(- (-)2 |1:00)
itself as the latent parameter 0. Therefore, given some observation 1 gezzrating I P(: is equivalent to sampling 0 ~ pu-([1x:)- This indicates we can do equivalent Bayesian inference 12 using P(.))..).
Differentiating epistemic and aleatoric uncertainty. Since epistemic uncertainty is equivalent the predictive uncertainty of future observations, it can be reduced with additional observations. In contrast, aleatoric uncertainty refers to uncertainty that remains irreducible, even with more observations. For example, in the coin toss scenario (Figure 2 the uncertainty in Coin A is aleatoric,
arising from inherent randomness of fair coin toss. No matter how many times coin A is flipped,
there will always be uncertainty about the outcome of the next coin A toss. Coin B, on the other hand, has epistemic uncertainty which can be eliminated by flipping the coin once.
To quantify epistemic uncertainty and distinguish from alcatoric uncertainty, we must autoregressively generate 11 P-)). As
Decision problem
Which colin (A 8) should we toss first if we want to maximize total H (heads) in M coin tosses
Decision using one-step Inference
Wrong answer
Coin H)
shown in Figure - consid- ering the entire sequence of future coin tosses is crucial-examining only a single toss results in iden- tical predictive uncertainty for both coins, making them indistinguishable. For- mally, since P(Y]).)) (P()))))(0)))..) one- step prediction/generation fails to distinguish between epistemic and aleatoric uncertainty. We explore this further in Section
H]
Select any coin (A or B] randomly for the first toss as the same predictive uncertainty.
Coin A: Fair coin with one side and other side
Decision using multi-step inference
Correct
PI Y1. Ysr) 0,5
Both side H th 0.5
Coin A:
for all (""""") E (H, 73
M
Aleatoric uncertainty
Coin
- Mr 0.5
Epistemic uncortainty
Both side T with 0.5
P/Y T, TT"",, 0.5
Coin B: Biased coin but we don't tt know whether both sides both sides T
Choose coin for the first toss, If results keep tossing it otherwise snift to the coin A
Figure 2: [Impact Multi-step inference v/s One-step inference in deci- sion making] Coins A and are considered identical by single-step inference because both have the same level of predictive uncertainty in their rewards. However, multi-step inference highlights key difference: for Coin B, the uncertainty can be reduced (epistemic) by performing single toss, whereas for Coin A. all the uncertainty is irreducible (aleatoric) and arises from inherent randomness a fair coin toss. Consequently, multi-step inference prioritizes tossing Coin B first to reduce epistemic uncertainty. (Illustration is adapted from (24).
Learning probabilistic model from data. We now shift our focus to di- rectly learning P(Y.T WIrT). We represent the transformer-basee autore- gressive sequence model, parametrized by é, as Pa. Rewrite (Y VI.T)
as an autoregressive product: IT'- Po(Y:.1 Ye+1 YI. = yis), where p( Pix)) rep- resents one-step predictive distribution, often referred to posterior predictive in conventional Bayesian terminology. The transformer is trained to minimize the KL-divergence between the true data generating process P' and the model Po
DRi [
E Ui P [oo (Po (Ki+1 1 = v1:) =]]
1=0
Given a dataset of sequences (vi. 1 S N} generated from the true data generating process P()), we train the autoregressive sequence model (transformer) on this data using the following objective: - N E ET- log Po = v Y = - For simplicity. denote Po = v(b1.) p((+1) (v). This training procedure enables the model to learn the single-step predictive distributions that collectively define the full sequence likelihood. Once trained, given any observed data Vi.t the transformer sequence model can be used generate future samples using one-step inference or multi-step step inference. We now formally define one-step inference and multi-step inference. Suppose we want to predict Y1+1- given y1ai

---

Definition 1 (One-step inference) Given f1:t. we generate each Y, ~ Po(-IBra) iid from the sequence model. That is, the probability of Y++1.T 1++1.7 is equal I Po(Y vilvi:). We denote this distribution by Po (U--IT))
Definition 2 (Multi-step inference) Given yla, we generate Y++I:T autoregressively from the sequence model. That is the probability of Y = 1:+1.T is equal to I Po v(l We denote this distribution by PM (y(41:))
3 Inferential inductive biases for decision making
Much of the literature on autoregressive sequence modeling focuses on one-step inference or predic- tion. For example, recent work by Hegselmann et al. [88] empirically demonstrates the scalability of sequence models for training tabular data, but their study emphasizes one-step prediction. Similarly, while Müller al. [22], Nguyen and Grover [23], Hollmann et al. (20 discuss predic- tive uncertainty, their focus remains exclusively one-step predictive uncertainty. As mentioned carlier in Section 2 and elaborated on later, one-step inference does not adequately differentiate between epistemic and aleatoric uncertainty. In contrast, multi-step inference provide more ef- fective means of quantifying epistemic uncertainty, which is crucial for sequential decision making applications [28 (4]. This distinction was illustrated informally in the coin toss example in Figure To clarify this further, we present a more formal example in Section D.I
3.1 Theoretical characterization of impact on decision-making
We emphasize the significance of the gap between single-step and multi-step inference by examining its impact on decision-making Recall Figure 2 where we needed to decide whether to flip coin A or B first to maximize the total number of heads over M coin tosses. One-step inference fails to distinguish between the two coins, leading to a suboptimal decision of selecting a coin randomly, In contrast, multi-step inference addresses this issue, enabling optimal decision-making
Characterizing information loss in one-step inference v/s multi-step inference: We analyze how our uncertainty quantification and inference in inneral-deterioratee when relying solely on one-step inference instead multi-step inference. Let Yt 1:T IP((-1) be some data from the data generating process. Further, let 50: (1++1T) = I1 Pf = w1l1i.)) be the one-step inference model and let PM (1++1.) 1I Po 11 be the multi-step inference model, where
we generate Y, autoregressively. Then the expected difference between the log- log-likelihood of the multi-step inference model and single-step inference model is given by the following result.
Theorem 2. Assuming Po PP, then the difference Eloog[Ph (u++1r) log PO (1++1=))) is
equal to Σi= Tyy 1141a-1lVia), where T(A; BJC) is the mutual information between A and B conditional on C and expectation is w.r.t. V1. P(-).
This demonstrates that relying solely on one-step inference results the loss mutual information among 3t+ Consequently, the expected likelihood under the multi-step inference model is higher than that under the one-step inference model and single-step inference inherently less effective than multi-step inference. We further analyze expression Σi=++ T(y:: 1t+1s- 11/1=/ in a specific setting (Example and show that as o² (epistemic uncertainty) increases, the performance gap between one-step and multi-step inference widens, with one-step inference becoming increasingly suboptimal.
Example 1. Assuming Y 0 where N(O, 7?) and 0 ~ N(p, Let Ve+1:1 IP((V11))) and (a')2 = (22 + - -1 then expression (Ei I(y; 1111)) equal to 2 log((1 " T 2 log(1 + (T -
Characterizing impact on decision making: We now examine how one-step inference affects downstream decision-making tasks. Specifically, we consider a one-armed bandit problem where the reward of the first arm given by y()) N(0 with 0 N(H,,²). While the second arm has constant reward (2) 0. It well known that Thompson sampling suffers O(VT) Bayesian regret in multi armed bandits [25] We now consider the performance of Thompson sampling implemented
5

---

3.0
Multi-step inference
6-
Multi step Inferencc
1.3
Mult step interence
2.5
OOncstop inference
inference
1.2
One-step Inference
Bayestaan Regret
1.1
2.0
1.0
0.9
0.8
1.0
0.7
10
20
20 40
60
20
30
40
50
Context length
Number of (T)
Number of Revealed Labels
(a) Uncertainty Quantification
(b) Multi-armed Bandits
(c) Active Learning
Figure 3: Comparing one-step inference and multi-step inference (lower is better): (a) Uncer- tainty Quantification: Comparing multi-step log-loss for one-step and multi-step inference (Train horizon: 100, Target Length: 50) (b) Multi-armed Bandits: Comparing Bayesian regret under Thomp- sampling algorithm using one-step and multi-step inference. (c) Active Learning: Comparing log-loss under uncertainty sampling algorithm using one-step and multi-step inference.
based just the one-step inference one-step inference (predictive uncertainty), with proof in Section
Theorem 3. There exists one-armed bandit scenarios in which Thompson sampling incurs O(T)
Bayesian regret if it relies solely on one-step predictions from autoregressive sequence models.
Generalizing to the contextual setting: We extend our analysis the contextual setting in Section D where the context X Widd Px- Additionally, we characterize the information loss associated with one-step inference in Bayesian linear regression and Gaussian processes in the same section.
3.2 Empirical investigation
In this section, we empirically evaluate the impact single-step and multi-step inference on uncer- tainty quantification, well on downstream optimization tasks such multi-armed bandits and active learning. We find that multi-step inference significantly outperforms one-step inference, being up to 60% more efficient in bandit settings and requiring up to 10 times less data in active learning to achieve the same predictive performance.
3.2.1 Uncertainty quantification (UQ)
We evaluate one-step and multi-step inference by generating datasets using Gaussian Processes (GP), a common choice in prior works 23] Specifically, we employ a GP with an RBF kernel: GP(m,, where X'XXX')) of exp (-| X'|l/ 2(2). Additionally, Gaussian noise (O, o²) is added the outputs. The input X drawn i.i.d. from Px. compare the performance of the two inference strategies, we use the multi-step log-loss metric. Further details on the metrics and experimental setup can be found in Section Figure illustrates the comparison multi-step log-loss performance between one-step and multi-step inference. Consistent with our theoretical results (Theorems [ and 4), the results demonstrate that one-step inference performs worse than multi-step inference.
3.2.2 Multi-armed Bandits
Problem: We consider a two-armed Bayesian bandit setting with arms (C, DJ and T rounds during which the arms are pulled. In each round i, based on the information collected so far, ((A,, A < < t 1}, arm A, is selected, and reward Y: is observed. For each arm a E (C, the rewards are distributed as y(a) iid N(Bla), (T())) where the mean reward 1:2 g(a) follows prior distribution g(a) ~ N (f(a)) )). The objective is to determine which arm A, E D} to pull in each round t in order to minimize the Bayesian regret, defined as:
OOrr code repository is available
at:
https://github com wtt.//oo///oooeee -amkoong-aaa.
Inductive- siteee-biiss--eeee bww.oaaeeeeccccaaaaleeeeeeeeeee- wwchangeable-sequeeeee

---

( y(di)), where expectation is taken over the randomness in the rewards (Y), the actions/policy (At), and the means (#).
B[maxxe DD [Toka)} (y Σt=
Training and Evaluation: We train two transformers, one for each arm (C and D). To compare the two inference strategies, we first sample 0* is for each arm. We then implement Thompson Sampling (using the respective inference strategy) with the trained transformers to acquire rewards over a horizon T. The regret is evaluated as |max,e (Tg*(a)} - E Y, r,(A)), Finally, we average the regret over 1000 different samples of 0", with each run consisting of T' = 100 steps, to compute the Bayesian regret. Additional details about training and algorithm are in Section
Results: Our results are summarized in Figure B). As expected, cumulative regret increases with the number of pulls for both inference strategies. The figure also demonstrates that multi-step inference significantly outperforms one-step inference having upto 60% less regret.
3.2.3 Active Learning
Problem: In active learning, the goal is to adaptively collect labels Y for inputs to maximize the performance of model v)()). We focus on pool-based setting, where pool of data points X is given, and the objective is to sequentially query labels Y for X E X We consider a regression setting in which inputs X lid Px, and outcomes are generated from an unknown function f", such that Y = (X) + 6x. where the noise 6x ~ NO,, T is heteroscedastic. Additionally, the aata-generating function f drawn from a distribution Ja.
Training and Evaluation: We consider meta learning setup where we train the sequence model (transformer) data {(1 1:N: Y j [1, M]] generated from the original data generating process. To evaluate the two inference strategies, we first sample f ~ μ and generate a dataset X J = D. which contains both pool dataset (Drool) and test dataset (Dtest), ). Using the trained transformer and the respective inference strategy for uncertainty sampling, we sequentially select X Xppool for which labels are queried. At each time step i, given the collected data Dd Deoot we evaluate the transformer model's performance as ExxY)) P D=)]
The results are summarized in Figure For both inference strategies, prediction accuracy improves, and loss decreases as more data points are acquired. However, multi-step inference sigmificantly outperforms single-step inference, achieving the same performance level with nearly 10 times fewer samples.
Results:
Now that have established the importance of multi-step inference, the next key question arises: What architectures are best suited for modeling exchangeable sequences, especially when performing multi-step inference?
4 Architectural inductive biases
Ensuring that the sequence model Ps is infinitely changeable enables robust performance and reliable statis- tical inference [30]. Several prior works have proposed ar- chitectural approaches to enforce exchangeability. Müller al. [22] introduced masking scheme designed to force exchangeability (Figure In this scheme, all con- text points attend to one another, allowing the model to condition on the entire context set without any pre- defined ordering constraints. This design ensures that the model's predictions remain invariant to the order of context points, aligning with the principle that exchange- able sequences should not depend on the specific order in which observations are presented. Figure illlsttates the masking scheme corresponding to this architecture, where (x1, 11,.9....... (x3, Ma) denote the context points, and (xa, 0) represents the target point for which 54 is to be predicted. Nguyen and Grover (33 proposed similar architecture with modifications to improve training effi-
(A. (r Mak (5,
(x,00
Figure 4: A representative attention mechanism and masking scheme widely used in prior literature to enforce ex- changeability. However, it only en- sures the conditionally permutation- invariant property.
7

---

ciency; however, their design contained an error where the target point did not attend to itself. This issue was later corrected by Ye and Namkoong [3)]
Despite these efforts, the exchangeability guarantees of these architectures remain uncertain. Previous works implicitly assume that their masking schemes enforce exchangeability but fail to provide formal characterization of what is actually achieved. To address this gap, we introduce a precise definition that explicitly captures the invariance imposed by these architectures. Our analysis reveals that all transformer-basee architectures in the literature [22 23 [0] aim to enforce a "conditional permutation invariance property." which we formally define below.
Property 1. Pg is conditionally permutation invariant if (M,
. ,,a))) =
This property ensures that the predictive uncertainty in Ve+11 remains the same under any permutation of the context V1:1- We refer the reader to Section for the corresponding definitions in contextual settings. Additionally, in Section 4.1 we provide detailed discussion of the transformer architecture with the conditional permutation invariance property, including its efficient training procedure and the computational requirements for inference.
However, while conditional permutation invariance is necessary characteristic of exchangeability. is not sufficient. Enforcing this property alone does not guarantee full exchangeability- critical oversight prior work. For example, all exchangeable sequence models must also satisfy another crucial property called the conditionally identically distributed (c.i.d.) property, also known as the martingale property:
Property 2. Recalling (5(() A Po = 11.-1), Ps is conditionally Identically Distributed (c.i.d.) if E[ (b)
[ (v)
This property ensures that the expected predictive distribution at time & + 1, given past observations (Yuz-1), is consistent with the predictive distribution at time t. The importance of c.i.d. property in exchangeable sequence models what powers their ability to quantify epistemic uncertainty, emphasized by previous work in Bayesian statistics [2] 4 [2 []]. However, an autoregressive sequence model that satisfies Property does not necessarily satisfy Property 2 We demonstrate this with concrete example provided in Appendix D.2
Although we have established that C-permutation invariant architectures not achieve full exchange- ability, it is still important to assess whether incorporating Property I into transformer architectures offers any advantages. To investigate this, we compare it to the standard causal transformer architec- ture (Section 4.2). which does not exhibit this property. Specifically, architecture and (2) the standard causal masking architecture.
analyze two architectures: (1)
the conditionally sermutation-invarian We describe their respective masking schemes, efficient training procedures, and computational requirements for inference.
4.1 Conditionally permutation invariant architecture
As we have established, achieving conditional permutation invariance requires ensuring that all context points can attend to one another. We adopt the same masking scheme shown in Figure 4 which, as previously discussed, has also been utilized in prior works such as 22 20).
efficient training procedure: Training directly a naive manner can inefficient, as would require processing T separate sequences, each of length i E 21, .T) for an sequence data of length T. A more efficient approach is to fix a context length i and train the transformer on multiple target points simultancously. This process is then repeated for all possible context lengths i. The corresponding masking scheme is shown in Figure Importantly, this method is equivalent to training for P(y| V1:4, z) across multiple values of in parallel, solely to optimize training efficiency. However, during inference, a multi-step (autoregressive) prediction approach should be adopted, described in Section 3
Predicting a sequence of length T requires approximately ((T3) computational effort. A major drawback of this masking scheme is that, even with KV caching, the compute cannot be reduced to (T2). This limitation arises because, at each inference step, all outputs from the attention heads must be recomputed every point attends to the newly added context point (sec
Inference compute:
8

---

0.42
Standard p-eomutttion invariant
0.9
Standard
Standard
0.40
C- permutation invariant
C perrmutttion invariant
0.38
0.8
0.40
Loos
0.36
0.7
0.38
Log
10g
0.36
0.32
0.34
0.30
0.32
0.4
20
40
50
80
10
40
20
40
50
Context length
Epochs
Context length
(a) In-training horizon perfor-
(b) Training/data efficiency of (c) Out-of-training horizon per- two architectures
mance of two architectures
formance of two architectures
Figure 5: Comparing C-permutation invariant architecture and Standard causal architecture (lower is better): (a) In in-training horizon performance (b) Training/Date efficiency (c) Out-of traaiinn horizon performance
Figure 9 for details). This behavior contrasts with the typical causal masking approach, where such recomputation is avoided.
4.2 Standard causal architecture
In this scheme, each context point attends only to the previous context points and not to any future
points. Specifically, the 3t context point attends only the context points within i 1. In the interest of space, the masking scheme for this architecture is illustrated in Figure (Section A
An efficient training procedure: As before, efficient training procedure involves fixing context length i and training the transformer on multiple target points simultancously. The corresponding masking scheme is shown in Figure g
Inference compute: Predicting a sequence of length T requires approx. O(T3) computational effort. However, with KV caching, this can be reduced to O(T²) (see Figure 10)
4.3 Empirical investigation
In this section, we compare the contextual permutation invariance architecture and standard causal architectures. As is common in the literature [22 23 24], we use the standard log-loss metric to compare these architectures and defer the analysis of downstream performance to Section E Our findings indicate that the in-training horizon performance of the C-permutation-invariin architecture is comparable to that of the standard causal architecture. However, the standard causal architecture horizon by approximately 10%
outperforms the C-permutation-ing invariann model on the out-of-trainin and demonstrates greater training and data efficiency by up to 20%. Furthermore, as noted earlier in Section 4 the standard causal architecture requires significantly less inference computation due to same data generating process in
benefits from KV caching. To evaluate these architectures
Section 3.2.1
Evaluation Metric:
To compare these two architectures, we use two metrics one-step log-loss and multi-step log-loss. These metrics are described in detail Section
Both masking schemes are trained on data generated from Gaussian Processes. Additional details about the architectures and training process are provided in Section B For brevity, we present only the results based on multi-step log-loss in the main body, while results for one-step log-loss are deferred Section [
Results: We evaluate these architectures across three dimensions performance on sequence lengths within the training horizon, and beyond the training horizon, training/data efficiency.
In-training horizon performance: Figure 5)a) shows that the performance of both masking schemes is comparable, with no evident advantage of enforcing conditional-permmuttio invariance. Furthermore, additional ablation studies (see Section Fi indicate that standard-causal masking may outperform cewiiitionnl-eemmtatattee invariance masking.

---

2. Data/Training efficiency: To evaluate the data/training efficiency of these architectures, we compare their performance across various training epochs. Figure(Jjb indicates that standard causal masking exhibits superior training efficiency. Similar findings were consistently observed in the ablation studies (Section
3. Outof-training horizon performance: To assess this performance, we train both architectures to horizon m 15 and evaluate their performance on horizons beyond m 15. Figure (f)), along with additional experiments (see Section Db, suggests that the standard causal masking may achieve better performance on out-of -ot---in-ninngggg horizons.
5 Conclusion and future work
We empirically and theoretically demonstrate that one-step inference using sequence models, which has been the primary focus of the literature thus far, is insufficient for distinguishing between epistemic and alcatoric uncertainty, ultimately leading to suboptimal decision-manl making. In contrast, multi-step autoregressive generation effectively overcomes this limitation, as evidenced by empirical results in multi-armed bandit and active learning settings. On the architectural side, much of the existing work has focused on enforcing a masking scheme that, as we identify, only ensures conditional permutation invariance in transformers rather than full exchangeability. Empirically, find that this approach performs worse than standard causal masking and significantly increases the computational cost of multi-step inference, as it cannot leverage KV caching. A limitation of work is that were unable to identify a transformer architecture that achieves both full exchangeability and computational efficiency for decision-making and active exploration. Developing such architecture remains important direction for future research.
References
[1] C. C. Aggarwal, X. Kong, Q. Gu, J. Han, and S. Y. Philip. Active learning: A survey. In Data classification, pages
Chapman and Hall/CRC, 2014.
[2] P. Berti, I.. Pratelli, and P. Rigo. Limit theorems for class of identically distributed random variables. Annals of Probability, 32(3):2029 2052, 2004.
[3] Berti, E. Dreassi, L. Pratelli, and Rigo. A class of models for bayesian predictive inference. Bernoulli, 27(1):702 726, 2021.
[4] Berti, E. Dreassi, F. Leisen, L. Pratelli, and Rigo. Bayesian predictive inference without a prior. Statistica Sinica, 34(1), 2022.
[5] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, and A. Askell. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33, 2020.
[6] D. M. Cifarelli and E. Regazzini. De Finetti's contribution to probability and statistics. Statistical Science, 1144):253-222, 1996.
[7] B. De Finetti. Classi di numeri aleatori equivalenti. la legge dei grandi numeri nel caso dei numeri alcatori equivalenti. sulla legge di distribuzione dei valori in una successione di numeri aleatori equivalenti. R. Accad. Naz. Lincei, Rf S 6a, 18:107-110, 1933.
[8] B. De Finetti. La prévision: ses lois logiques, ses sources subjectives. In Annales de I l'institut Henri Poincaré, volume 7, pages 1-68, 1937.
[9] B. De Finetti. Theory of probability: critical introductory treatment, volume 6. John Wiley & Sons, 2017.
[10] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL httppp:///pppperrereerrr//rrrrre
10

---

[11] F. Falck, Z. Wang, and C. Holmes. Is in-context learning in large language models bayesian? martingale perspective, 2024. BKLLitpps///axx....o/////////// 00793
[12] E. Fong, C. Holmes, and S. G. Walker. Martingale posterior distributions. Journal of the Royal Statistical Society, Series B, 2023.
[13] S. Fortini and S. Petrone. Predictive distribution (de f'inett'ss view). Wiley StatsRef: Statistics Reference Online, pages 1-9, 2014.
[14] S. Fortini, L. Ladelli, and E. Regazzini. Exchangeability, predictive distributions and parametric models. Sankhyā: The Indian Journal of Statistics, Series A, pages 86-109, 2000.
[15] J. P. Gardner, J. C. Perdomo, and L. Schmidt. Large scale transfer learning for tabular data via language modeling. In The eiihty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[16] P. R. Hahn, R. Martin, and S. G. Walker. On recursive bayesian predictive distributions. Journal of the American Statistical Association,
2018.
[17] S. Han, J. Yoon, S. O. Arik, and T. Pfister. Large language models can automatically engineer features for few-shot tabular learning. In Proceedings of the 41st International Conference on Machine Learning, 2024.
[18] S. Hegselmann, A. Buendia, H. Lang, M. Agrawal, X. Jiang, and D. Sontag. Tabllm: Few-shot classification of tabular data with large language models. In F. Ruiz, J. Dy, and J.-W. van de Meent, editors, Proceedings of the 26 International Conference Statistics, volume 206 of Proceedings of Machine Learning Research, pages 5549-5581. PMLR, 25-27 Apr 2023. URL.tttss//proceedd.ng mlr pares//v06//heeee http///00///eeeeaaaaaeeeee html
Artificial Intelligence and
[19] E. Hewitt and L. J. Savage. Symmetric measures on cartesian products. Trunsactions of the American Matemattiaa Society, 80(2):470-501, 1955.
[20] N. Hollmann, S. Müller, L. Purucker, A. Krishnakumar, M. Körfer, S. B. Hoo, R. T. Schirmeis- ter, and F. Hutter. Accurate predictions on small data with tabular foundation model. Nature, 637(8045)3319-22
2025.
[21] E. Kauffannn O. Cappe, and A. Garivier. On bayesian upper confidence bounds for bandit problems. In N. D. Lawrence and M. Girolami, editors, Proceedings of the Fifteenth Interna- tional Conference Artificial Intelligence and Statistics, volume 22 of Proceedings of Machine Learning Research, pages 592-600, La Palma, Canary Islands, 21-25 Apr 2012. PMLR.
[22] Müller, N. Hollmann, S. P. Arango, J. Grabocka, and F. Hutter. Transformers can do bayesian inference. In Proceedings of the Tenth International Conference on Learning Representations, 2022.
[23] T. Nguyen and A. Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In Proceedings of the 39th International Conference on Machine Learning, 2022.
[24] I. Osband, Z. Wen, S. M. Asghari. V. Dwaracherla, X. Lu, M. Ibrahimi, D. Lawson, B. Hao, B. O'Donoghue, and B. V. Roy. The Neural Testbed: Evaluating Joint Predictions. In Advances in Neural Information Processing Systems, 2022.
[25] D. J. Russo, B. Van Roy, A. Kazerouni, I. Osband, and Z. Wen. A tutorial on thompson sampling. Foundations and Trends0 in Machine Learning, 11(1):1-96, 2018.
[26] B. Settles. Active learning literature survey. 2009.
[27] R. Weber. On the gittins index for multiarmed bandits. The Annals of Applied Probubility, pages 1024-1033, 1992.
[28] Z. Wen, I. Osband, C. Qin, X. Lu, M. Ibrahimi, V. Dwaracherla. M. Asghari, and B. V. Roy. From predictions to decisions: The importance of joint predictive distributions. arXiv: 10070024 [cs.LG), 2022.
11

---

[29] J. Yan, B. Zheng, H. Xu, Y. Zhu, D. Chen. J. Sun, J. Wu, and J. Chen. Making pre-trained language models great on tabular prediction. In The Twelfth International Conference on Learning Representations, 2024.
[30] N. Ye and H. Namkoong. Exchangeable sequence models quantify uncertainty over latent concepts, 2024. URL http:::///// wwtpxiii..rg/bb//////0/ hwtpg//bb///0....//// 040.2033077
[31] K. W. Zhang, H. Namkoong, D. Russo, et al. Posterior sampling via autoregressive generation. 1xXx: [cs.LG), 2024.
[32] Z. Zhao, Birke, and L. Chen. Tabula: Harnessing language models for tabular data synthesis. arXiv preprint 1rX1:2210.12746, 2023.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Main results in Section 3 and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation in Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
13

---

Justification:
Section
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section
B
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification:
Please refer to the supplementary materials for the link.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnps.cm
httpii///giidess/deeSSbbbbmt wesssnnnoiccyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ cpipp/.ccppullc//uuuii wwwtc//ggue////eeeeceoocc CoodeSubeiiieeee ssi wsionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
All the graphs include error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
15

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section
B
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpppp//eerrpp....ccbbbbb.biiie
Answer: [Yes]
Justification: Yes we conform with NeurIPS Code of Ethics and there is no potential harmful effect of this research.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: here appear be no direct pathways harmful negative applications arising from this work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All the creators or original owners of assets (e.g.. code, data, models) are credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt
package should be provided. For popular datasets, paperswi thcode. cow/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Thee are no new assests created using this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: Core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
19

---

A
Additional Demonstrations
1, (rs (x,, 0)
(4,3)
(X..%)
(x5.3s)
(aod) (5.3)) (sat)) (4.0)
(x,,0)
Figure 6: Standard causal transformer architecture: attention mechanism and masking scheme
(u0t)) lavis Drundl (y00)
(euA))
(599
04
06.39
057
06
0470
04001
0x0
6-40
Figure 7: Efficient training procedure for condi- tionally permutation invariant transformer
Efficient training procedure for stan- dard causal transformer architecture
Figure 8:
()
(x,,)) (e (x))P)) (4,.0)
(x (Xg.)s (x) (x. fel (xg
(5.))
(463))
Predict
Predict
(N).0)
Yo
9x)
94
9x)
(x).0)
(r,.94
(s..0)
Figure 9: Conditional permutation invariant architecture masking scheme at inference stage
20

---

Boos) (1a))t)) (tyy
(5 (fa,3s (ty. Fs) (s4,
(Rx (1a.2)) (Xy Pa) (ta Pal (At. 0)
(Ry)))
(%.3))
(ss)1)
(Rx(5)
Prodict
Predict
Ya
9,
(3.00)
(ss.9)
(s,00)
(x4.9%)
(x,
Figure 10: Standard causal architecture masking scheme at inference stage
B Experiments Details
In this section, we present the experimental details from Sections 3.2 and 4.3
B.1 Details experiments in Section 3.2.1 and
Data generating process: As previously mentioned, we generate data synthetically using Gaus- sian processes. Specifically, we employ a Gaussian Process (GP) with a Radial Basis Function (RBF) kernel: f GP(m, K.), where m(X) represents the mean function, and (CX,X')
of exp - 2
represents the covariance function. Additionally, Gaussian noise (0, o²) is added to the outputs. The input X drawn i.i.d. from Px. Unless stated otherwise, the parameters are set as follows: m(X) = 0, X 2 2.0], of 1.0, l 1.0, 0.1.
Evaluation Metric: To compare these two architectures, we employ two metrics: one-step log-loss and multi-step log-loss. A detailed description of these metrics is provided below.
1. One-step log-loss: Given (#1sm, Flam Fm+1),, we generate ym++1 from the true data gener- ating process, i.e., Vm++1 PP- [2m Fims V1:m). The one-step log-loss for the sequence model is then calculated
log Po(ym+1 24, Timms 31.)) This value is further aver- aged over multiple instances of (I1:ms Blam 1m+1).
2. Multi-step log-loss: Given (I1:m, I1:ms Im+1):), we generate 4m++1:7 from the true data- generating process, i.c., Vm++1:7 PP 1:T, 2 Tim: V1.m). The multi-step log-loss for the sequence model is computed
-log P (3m+1:7 [m+1:T Zimm {1m)). which can
be expressed -E-1 log (a Ii 1. 3 Fimm Dlomy Monss) This value is then averaged over multiple instances of (Fiams Mism: Fm+1r)).
We refer to t as the context-length and (T - t) as the target length.
To compare the conditionally permutation- invariant architecture with the standard causal architecture, we use a decoder-only transformer with the following parameters. Both architectures share the same parameters, differing only their masking schemes. The model parameters are as follows:
Transfromer architecture and training details:
Model dimension: 64
Feedforward dimension: 256
Number of attention heads: 4
Number of transformer layers: 4
Dropout: 0.1
21

---

Activation function: GELU
For embedding (æt, y),
neural network with two layers of sizes [256, 64]. Additionally, final linear layer is used to predict the mean μ and standard deviation o the output distribution, modeled as Y N (H, o²). For training the transformers, we use the Adam optimizer with default parameters, and the learning rate is adjusted using a cosine scheduler. The training parameters are follows: Warmup ratio is 0.03, minimium learning rate 3.0e-9, learning rate is 0.0003, weight decay to 0.01 and batch size is 64. For all the experiments we train the transformer for 400 epochs.
Computational resources: We use NVIDIA A100-SXM4-800 for training models. For the standard-causal architecture takes 4hr, while for C-permutation variant architecture takes 17hr to train the model.
Evaluation: We evaluate the trained transformers, corresponding each architecture, using one- step log-loss and multi-step log-loss. For each experiment, we use 8192 test samples and conduct evaluations across five different random seeds. This process includes retraining the models on different training datasets and evaluating them on distinct test datasets.
B.2 Experimental details for Section 3.2 Rayesian Bandits
Data generating process: Recall that, for each anm [C, the rewards distributed as yla) N (o()), (r(al)²). where the mean reward ((a) follows prior distribution (la)
1:T
N (ula), (ola))²). Where for Arm C, set (C)) 0, o (C) 0.5 and ((C) 0.5. While for Arm
D, we set (D) =0, (D)) 0.9 and ((D) 0.1.
Exploration Algorithm: Numerous algorithms have been proposed in the literature to address this trade-off, such Thompson Sampling [25). Bayes-UCB [21] and Gittin's 's index (27). However, to
compare the two inference strategies, we fix the algorithm to Thompson Sampling and implement using the respective strategy.
Transformer Architecture and Training: As specified earlier for each arm a E {C, D}, we train
separate transformer (decoder-only). The model parameters and training details are specified below:
Model dimension: 64
Feedforward dimension: 256
Number of attention heads: 4
Number of transformer layers: 4
Dropout: 0.1
Activation function: GELU
For embedding (z, y).
neural network with two layers of sizes [256, 64]. Note that, in this case there is no context, we set x as O. Additionally, a final linear layer used to predict the mean l and standard deviation of the output distribution, modeled Y N(u, o²). For training the transformers, we use the Adam optimizer with default parameters, and the learning rate is adjusted using cosine scheduler. The training parameters learning rate is 3.0e learning rate is 0.0003, weight decay to 0.01 and batch size is 64. For all the experiments we train the transformer for 400 epochs.
follows: Warmup ratio is 0.03, minimium
Evaluation Details: The results presented in Figure ()) are averaged over 1000 different experi- ments, with each experiment run for 100 steps.
One-step Multi-step inference (Thompson sampling): In Algorithm I we outline the imple- mentation of Thompson sampling algorithm for one-step and multi-step inference using a trained transformer. Algorithm| setting
corresponds to one-step inference, while choosing J > enables multi-step inference. In our experiments, we set J = 100 for multi-step inference. For refer- ence, we also provide the standard Thompson sampling algorithm for Gaussian- multi-armed bandits (see Algorithm 2) In our experiments, we set J 100 for multi-step inference.
22

---

One-step and Multi-step inference (Thompson sampling) using sequence models (transformers) in multi armed bandits setting
Algorithm 1
Require: Trained transformer (P)a) for each of the Arm e [1, K], horizon T, Number of
autoregressive generations in each iteration (J): one-step inference J 1, while for multi- step inference J L.
Initialization:
2: for 1 + K do
3: }(o) - 0
|/ Observed rewards for each arm number of pulls for arm
4: n²()) +
5: end for
6: for t= do
7: Autoregressively generate J rewards for each arm, conditioned on observed rewards
8: for = do
9:
(a) tn ~ Pa(-D ((a)) (,(0)+11
{(
I/ In (a) +1 f indicates time-step and (a)) + indicates conditioning on [y(a)] (1
observations and generations.
10: end for
(a)
11:
Select A. arg max Isask
Pull arm A, and observe reward (((A)) Update the collection of observations for arm Ar:
12:
13:
(Ai) t (Ad) +1, ydAd t y(Ad) u (r((Ai).
end for
B.3 Experimental details for Section 3.2
Active Learning
Data generating process: Recall that data generating process is as follows features X iid and outcomes are generated from an unknown function f³, such that Y = fX) + ex, with noise ex N (0, xxx being heteroscedastic and the data-generating function f drawn from distribution μ. Our Px consists of 100 non-overlapping clusters. Further, within each cluster f (X) is highly correlated, while across clusters, the correlation is low. Further
Px
clusters the noise ex is high, while on others it is low. This setting introduces a setting where is necessary to differentiate between alcatoric and epistemic noise for efficiently querying the labels for model improvement.
There are various active learning query strategies, such as uncertainty sam- pling techniques (e.g., margin-sampling entropy) and Bayesian Active Learning by Disagreement (BALD) 26 I. For this study, we focus on Uncertainty Sampling adapted to the regression setting.
Exploration Algorithm:
Transformer Architecture and Training: We train a decoder-onjy transformer on sequential- data {(xf)) I. y / E [1, M]} generated from the original data generating process. The model parameters and training details are specified below.
Model dimension: 64
Feedforward dimension: 256
Number of attention heads: 4
Number of transformer layers: 4
Dropout: 0.1
Activation function: GELU
For embedding (x, y), we use a neural network with two layers of sizes [256, 64). Additionally, final linear layer is used to predict the mean μ and standard deviation of the output distribution.
23

---

Algorithm 2 Thompson Sampling for Multi Armed Bandits (Gaussian-Gauseee setting)
Require: Number of arms K, prior mean p(a), prior variance (o!a))² 2 known reward variance
(r(a)) horizon T.
1: Initialization:
2: for + K do
3: (2)) +
number of pulls for arm a
4: 5(w)
// sum of rewards for arm
5: (1(0) t p((a)
// posterior mean for arm
6: (ō(a)) 2 + (ola))
posterior variance for arm
7: end for
8: fort= 1 T do
9:
Sample mean from each arm's posterior:
(a) Mr
N
( (ol.))
for 1, I,...,,K. K.
(a)
Select arm A, arg max
Pi ISask
11: Pull arm A, and observe reward Y, (Ac)
Update the sufficient statistics for arm As:
n (Ae) t (d)) + 5(A,) t s(Ac) + y((c).
13: Update posterior for arm A₁-
Posterior mean: ((A)) (a(4e)) / (Ar) + sta0
+ m(de)
(olds))²
:2
+ n((A))
(o(4.))²
1
Posterior variance:
(oldi)))
72
14: end for
modeled as Y ~ (N,, For training the transformers, parameters, and the learning rate is adjusted using a cosine scheduler. The training parameters are follows: Warmup ratio is 0.03, minimium learning rate is 3.0e-5 learning rate is 0.0003, weight decay to 0.01 and batch size is 64. For all the experiments we train the transformer for 400 epochs.
the Adam optimizer with default
Evaluation Details: Results shown experiment run for 50 steps.
in Figure 3 c) are averaged over 50000 experiments with each
One-step / Multi-step inference (Uncertainty sampling): In Algorithm 3 we describe the imple- mentation of one-step and multi-step inference based uncertainty sampling using trained transformers. In our multi-step inference experiments, we set J = 20 and / = 20.
C Contextual setting definitions
Consider contextual setting, where the context lid Px.
Exchangeability definition: An infinite sequence (X (Xioo0 Y1:00) ) is exchangeable if for any n and permutation
P((1,, Yi), ",,,,, Y.)) = P((( (()), Ya()), . (Xu(n), Ya(n))).
24

---

One-step and Multi-step inference (Uncertainty sampling) using sequence models (transformers) in active learning setting
Algorithm 3
Require: Trained transformer (Pa), Horizon T, Initial data available D Pool choose from 2 Number of autoregressive generations in each iteration (J): In one-step inference J = 1, while for multi-step inference J > 1; Number of generation paths over which variance taken I.
1: for = + T do
2: for E XX do
3: Generate trajectories for each sample X and collect the mean output on each trajectory in
list 4.x.
4: Initialize 4,x 0
5: for
7 do
6:
Autoregressively generate sample outputs for X for each arm, conditioned on the available data.
7:
for j do
8:
(1.x) ~ Po(-1D'-1, (X, (x,) (x, (x
9:
end for
10:
l(.x)
< (Xx) u U5E 10,))
11:
end for
12:
Estimate variance of the mean output across I trajectories V(x,)) Varianceyeluuxx (Y)
end for Select X = arg max Vi.x pont Query X and get the true label/output Y*.
14:
15:
16:
Update D' + Dr-1 (X*,Y* Y
17:
end for
Transformer training:
N T-1
min
log Ps ( = v
Vi. = V 1:::
j=1 6=0
Conditional permutation invariance property:
Po(Yeti((1,,1)) X441)
Ya(i)), (Xx(()) Yau)), Xet1).
Conditionally identically distributed property: In the presence of covariates, where ~ Px independently, the c.i.d. property extends the sequence model as follows:
p1+1 (y x)) (1tt-1 = (1 x).
where
p. (-,= = Xx Y1s- X1s1-1)
p
x)
and
Po (Ymm = Xt+1 x, 11.,, =: (2)
D Examples and Proofs
D.1 Example differentiating epistemic and aleatoric uncertainty
Example 2. Consider a setting where the observation Y Pa N(0 -²) follows a normal distribution with mean 0 l N(a Given observations yı.t, the posterior distribution (()l)r:)) N(a, of) represents the epistemic (reducible) uncertainty, while y NO, 7²) captures the alcatoric (irreducible) uncertainty caused by inherent randomness in the data.
25

---

One-step inference corresponds to P(Y(vi:)), which is equivalent to Po(Y)u N (at, of +2)). In this case, the predictive uncertainty for (characterized by V 7²), combines both epistemic (aî) and aleatoric (r²) uncertainties. Importantly, this approach does not allow us to separate the two components.
Multi step inference, on the other hand involves F(Y (1:)). which is (by De Finneti's s) equivalent to [[ Po (o ( ddo. Furthermore, as T - oo, we have limy-+o Σi= Y, = where 0 pu-l = N (at: of). Hence, the variance of the long-term average + Σ 2i=4+1 Y is approximately of isolates the epistemic uncertainty. Therefore, multi-step inference enables the differentiation between epistemic (of) and aleatoric (7²) uncertainties.
D.2 Example demonstrating Property does not lead to Property And exchangeability
Example 3. Suppose (0 (0) 1.- Further, conditioned on first observation, assume that (1( Pa(1) (0) Po(O)) P3(1) == (1 //F Observe that Y1 and Y2 are ii.d. Finally, conditioned on first two observations, assume that (0) Pa(O0 F22 8+
and (1) (o(l)),, 1 -Yi
Y2) Here, it is evident that Po(u|Y Y=(2)) (u 12) satisfying Property However, we observe that ((P) Yi)= ([((1 (2 Y2/2) 81) / + 1/3 is not equal to (0) = 1/3. Thus, E ( (2) f * (v). indicating
that the sequence model not exchangeable.
D.3 Proof of Theorem
Recall that 9t+1:T P(())1))) is the data generating process. Further, po (B(+1T) = Ih Pa = vllvaa) is the one-step inference model and PM (BI-1:) = IT PolP
-1 is the multi-step inference model, where Y; generated autoregressively.
Theorem
2
Assuming
P. then the difference E (log [PM (((+1:)))
[log
po (w-11r)) is equal to
(y
where T(A; C) is the mutual information between A and B conditional on C' and expectation is w.r.t V1.T (P-). Proof As data is generated from P(-). Therefore, under the assumption that Po. = PP we have that
T
Po = wolon..)))
log
PM (WL+1r)) E [og po (W++1r)) [(a) E log
II i+++1
II Ps volenas)) Li=t+1
E
log
T
P(Y = yi volvia-1)) 1)
(b) E log
II 1
log II IP(Y wi vellene) 1) -s=++1
I1 P(Y ys yxl81:111)
log
| P(Y = volvr:e)
26

---

T
IP y 111:11111
log
P(Y = villn
P(Y = y) 1=lva--1)
II'
log
PP(Y vil311.)
i=
(c)
Ty,, 1e+1a-118ha))
Here (a) follows from the definition of PM Po follows from the following identity
(b) follows from the assumption Pg. = PP (c)
I(X:Y Z) I[in
P(X,Y Y 2)
P(X Y. Y. Biin
P(Y X, Z) P(Y Z)
E In
P(X Z) P(Y Z)
P(X 2)
Also note that it is equal to En (Dul (F(U+1IT/BL:)IIT
P(y P(W1BB1.)))
D.4 Proof of Example I
Proof
Recall that, Y; = + E, &~~N(0 ~(0,7))
0 o(),
We first estimate KL divergence
T
DKL P(1 via))
||
II ==1+1
P(y: b1=)).
Because Y; 0 ~ NN(,,²), 72), therefore the posterior P(O
y1:)) is normal with mean = Further conditional on 0, the future
Eis Mi and variance is of (++ a2
1t+1..... vr are ii.d NO0, Therefore we get, Vt+1:T Vist A(pp 1, of
where 1 is the (T - t) dimenssoonal vector of all ones, and / the (T- x (T - t) identity.
For the one-step inference Denote this covariance by Dا = (72 + + 1.
Ply: | V1:t). each ys has distribution y: Via ~ N 7² +2-0?)
Now, we want the KL divergence between two (T- dimensional Gaussians with P = MMu1, Ep) with Σp = 72 + of 11 and Q NN,,1, Σo) with Σq = (r2 + of) T. For two K-dimensional Gaussians N(H,, Ev) and N N(H1, E1), the KL divergence is
DKL. DKI.(N(uo, Eo) [[, En)) = (IT(TT'E)) + (w) pO) (w) - mo) - K In Σ
Σg
Therefore we get
DkL()) = K ln 1 of - In( 1
}
72
As the expression independent of y1:., taking expectation over yıu our final expression remains the same.
27

---

D.5 Generalization of Theorem to the contextual setting
Loss of information in one-step v/s multi-step: Let y1-1:T 1 F1.T) be some data from true data generating process. Further, let P8 131- ((+1T) I Po = y: Ti:es 2;)
and PM (IL+I:T, (++1:)) I[T=ett Po(Y; - 1,
Theorem 4. Assume Pa = P. The difference (log P. (00-1..)+++10)))
Pg (ve+1ar, ( me+1r)) is
T(y Ue
lid ~ Px. where expectation is V1. ~ P(( and F1T
D.5.1 Proof of Theorem
Recall that V1+1.T ~ P(-1W1a,41: is generated from the true data generating process. Fur- ((41. (1++11,
ther, po @
= Il P8 (1,, and PM (1++1.T =
I Po = 1 11,-1,
Theorem
4
Assuming Po = P. then the difference log PM (yL+1:1, (30+111,0+111)) met1er))
CP (11411, is equal to
Mys ye+ 1:- (91,,
id where expectation is V1.r ~ P(1 and ø1.7 ~ Px.
Proof We can follow exactly same procedure as in the proof of Theorem for this proof.
D.6 Characterizing ( for Bayesian linear regression and Gaussian processes
Example 4. Bayesian Linear Regression Assuming Y = 0 X + where (0,, and (u, )). Inputs X are drawn ii.d. from Px. Let D. = (Zin and the posterior
0 D. N(J,, E'), then expression is equal to
|diag(XX) XX'X
ΣX,
log
X
where X is the matrix of F++1, 4(+2 .T-
Example 5. Gaussian Processes- Assuming Y = f(X) + where ~~ GP(m, KC), where m mean and XCX,, X') is the covariance. Additionally, Gaussian noise (0, o²) is added to the outputs. The imput X is drawn i..dd from Px. Let D, = (æ1., (1a).. Suppose, under multi-step inference P(y++1-T Y1.ts X1.T) N(y+411 pp, Kp). Further, under single-step inference Q(y++1.7) = I1 N( μ4, o,) = NV(+4 Mq: Eq). Then, expression 2
is equal to
[log
Eg Krp
Ex, DEp.Eq.MOuPP
- - -T + tr (K Eo) + (MQ Mp)) TKp (MQ MP))
where d = -)-
and the posterior DDD, N(L,EE) then expression is equal to
diag XX'X"
Ex,y
log
X
where X is the matrix of Tt+1, 2+42 ... TT.
28

---

D.6.1 Proof of Example 4 and
Proof
Recall that, under multi-step inference 1 Yias (1:)) Nyy411 Mp, P)). Further,
under single-step inference Q(y++1:7) 1 M(y pr; of = NNy+i: HQ: EQ). Now,
M
KL Divergence between two multivariate Gaussian distributions: N(y Hpp Kp) and Q = Ny
Mq: Eq) is give by
EQ
Du (PIQ)) log
d Eq) + (1q 1p))
ppp) "Kppppp (Mq up)
Kr
where: d is the dimensionality of (i.e., d = T - t); the trace of matrix. Therefore, we get
denotes the determinant and tr(.) denotes
-+ of K
Daa (PIQ) = 2 log
- ) + (Kp' Eq) + (uQ 1p)) "Kp' "Kp up)
In Bayesian linear regression the means remain the same, i.e., 1Q pp. further the expression reduces to the following:
DRI (M(X XX'X" (X) +
that is,
XE'XT
Da (PIQ)
log diag XX'X")
Final expressions follows from proof of Theorem
in Section D.3
D.7 Proof of Theorem 3
Recall that our setting was such that the reward of first arm is generated as (()) N(0, 72) with 0 Nu,, o²). While the second arm has a constant reward (2) O.
Theorem 3
There exists one-armed bandit scenarios which Thompson sampling incurs O(T)
Bayesian regret if it relies solely one-step predictions from autoregressive sequence models.
Consider bandit setting, where a 0 and μ < 0. Implementing Thompson sam- pling using the one-step inference, does not differentiate between aleatoric and epistemic un- certainty. This means whenever we will sample
Proof
> we will choose the wrong arm to pull as pamean reward from arm 1) < O(mean reward from Arm 2). As Y (2, therefore, P(Y > 0) = कb (#). Hence,
average we will suffer (T (#)) regret over horizon T.
Similarly, when 0 and μ 0, whenever will sample Y pull. Therefore in this case, it suffers T (=#) Bayesian regret over horizon T. On the other hand, the multi-step inference (assuming it done till oo) will have regret these two case.
will choose the wrong
E Additional Experiments
In this section, we present ablation studies to evaluate the performance of different architectures, specifically comparing conditionally permutation-invariann and standard causal masking schemes. In
29

---

0.450
Standard causal
0.46
Standard causal
Negative log probability
0.425
C-permutation invariant
Negative log probability
0.44
C-permutation invariant
0.42
0.40
0.375
00.38
0.350
0.36
0.325
0.34
0.300
0.32
0.30
20
40
60
80
100
20
40
60
80
100
Context length
Context length
(a) In-training horizon performance
(b) Out-of-training horizon performance
Figure 11: One-step log-loss: Comparing the two architectures on one-step log-loss metric [Dimen-
sion: 1].
Section we provide the corresponding results from Section 4.3 using one-step log-loss as the evaluation metric. Section E.2 explores the effect of varying the input dimension X, while Section E.3 examines the impact of different noise levels in the observed output
E.1
One-step log-loss figures corresponding to Section 4.3
In Figures aTaa and (T),, we present the in-training horizon and out-of-training horizon performance B[a) a) and55c) The results exhibit similar trends, where both masking schemes perform equally well in-distribution. However, beyond the training horizon, the causal masking approach appears to outperform the conditionally permutation-invatian masking.
for the corresponding settings shown in Figures
E.2 Ablations on dimensions
In this section, we present the results of our ablation study the dimension d of the context, where X U1-2,22
In-training horizon performance: Figure 12
shows the in in-training horizon performance for both architectures across different dimensions. Our findings align with previous observations. Notably, for d = 16, the provided context appears insufficient to improve prediction log-loss.
Training/Data efficiency: Figure
13 illustrates the training and data efficiency across different
dimensions.
Out-of-training horizon performance: We only analyze out-of-training efficieny for the dimension 4 in addition to dimension 1. As shown in Figure 14 the results remain consistent with our earlier findings.
Multi-step v/s One-step inference: Figure 15 present a comparison of multi-step and one-step inference across different dimensions.
E.3
Ablations on noise
In-training horizon performance: Figures 16 illustrates the in-training horizon performance for both architectures under varying levels of observation noise. Our findings remain consistent with previous observations.
Training/Data efficiency: Figures 17 shows the training and data efficiency across different levels of observation noise. The results align with earlier findings.
Multi-step v/s One-step inference: Figures present comparison of multi-step and one-step inference under different observation noise levels.
30

---

E.4
Downstream performance of the two architectures
Figure compares the performance of the conditionally permutation invariant architecture and the standard causal architecture in an active learning setting. Our results indicate that the standard causal architecture outperforms the conditionally permutation-ines invarriat architecture.
1.30
C-permutation invariant
Standard causal
1,470
Standard causal
C-pernutation invariant
1.475
1.25
1.465
1.470
1.00
1.460
1.15
Log 1.465
1.455
1.10
1.450
1.460
1.05
Standard causal
1.445
C. perm utation inveriant
1.00
1.455
20
30
40
50
60
70
80
90
20
30
40
50
60
70
80
90
20 30 40
50 60 70 80 90
Context longth
Context langth
Context langth
(a) Dimension:
(b) Dimension: 8
(c) Dimension: 16
Figure 12: Ablation dimension (In-training horizon performance): Comparing two architec- tures [Training horizon: 100, Metric: Multi-step log-loss, Target length: 10].
1.470
causal
1.470
Standard
E-permutation Invariant
C-permutation Invariant
1.469
1.468
Standard causal
1.468
Log
C-permutation invariant
1.467
1.466
1.466
1.466
1.25
1.465
1.465
10
10
20
30
40
Epochs
Epachs
Epachs
(a) Dimension: 4
(b) Dimension: 8
(c) Dimension: 16
Figure 13: Ablation on dimension (Training/Data efficiency): Comparing two architectures [Training horizon: 100, Metric: Multi-step log-loss, Target length: 10].
Standard causal
1.35
Standard causal C-permutation invariant
1.30
C-permutation invariant
Negative log probability
1.30
1.25
Loss
11.5
9e
1.20
1.20
1.15
1.15
1.10
1.10
30
50 60 70
90
20
60
100
Context length
Context length
(a) Dimension: 4, Multi-step log-loss
(b) Dimension: 4, One-step log-loss
Figure 14: Ablation on dimension (Out-of-training horizon performance): Comparing two
architectures [Training horizon: 15, Target length: 10].
31

---

1.49
step inference
1.400
Multistep inference
One inference
1.475
One step inference
1.48
Log Loss
1.455
Log Legs
1.47
1.460
M
1.455
1.46
1.450
Multi-step Inference
144
One-step inference
1.440
10
50
10
40
Context length
Context length
Context length
(a) Dimension: 4
(b) Dimension:8
(c) Dimension: 16
Figure 15: Ablation on dimension (Uncertainty Quantification): Comparing one-step inference and multi-step inference [Train horizon: 100, Metrie: Multi-step log-loss, Target Length: 50].
0.42
causal permmutttionn invariant
0.74
0.08
permutation invariant
C. permutation invariant
0.06
0.72
Log Loss
0.04
Log Loss
0.70
0.02
0.36
Log
0.68
0.00
0.34
0.02
0.66
-0.04
0.30
0.64
20
40
60
80
40
70
90
30
60
80
Context length
Context length
Context length
(a) Noise: 0.05
(b) Noise: 0.1
(c) Noise: 0.2
1.18
1.51
1.84
Standard causal
Standard causal
Standard
C-permutation Invariant
C-permutation Invarlant
1.83
permatation invariant
1.16
1.49
1.82
1.48
1.47
Log
112
1.46
1.80
1.45
1.10
1.44
1.78
1.43
90
Context longth
Context longth
Context length
(d) Noise: 0.5
(e) Noise: 1.0
(I) Noise: 2.0
Figure 16: Ablation on noise (In-training horizon performance): Comparing two architectures [Training horizon: 100, Metric: Multi-step loss, Target length: 10).
32

---

1.10
0.8
Standard
0.9
Standard
Standard causal
permutation invariant
permutation invariant
1.05
C-pemuutation Invarlant
0.8
1.00
Log Loss
0.7
Lost
0.95
0.90
0.4
0.3
0.2
0.5
0.80
0.75
0.1
0.4
20
40
20
30
50
Epochs
Epochs
Epochs
(a) Noise: 0.05
(b) Noise: 0.
(c) Noise: 0.2
1.64
Standard causal
Standard causal
Standard
1.35
C-permutation Invariant
C-permutation Invariant
1.92
C-permutation Invarfant
1.60
1.90
1.58
1.88
1.56
1.225
1.54
3
1.86
1.20
1.52
1.84
1.50
1.82
1.15
1.48
Epochs
Epochs
Epochs
(d) Noise: 0.5
(e) Noise: 1.0
(1) Noise: 2.0
Figure 17: Ablation on noise (Training/Data efficiency): Comparing two architectures (Training horizon: 100, Metrie; Multi-step log-loss, Target length: 10].
33

---

2.2
Multi-step Inference
Multi-step Inference
Multi-step inference
One-step inference
2.5
One-step inference
2.0
One step inference
1.8
Log Los5
2.0
1.5
1.2
1.0-
1.0
0
10
20
30
40
50
0
10
20
30
40
50
0
20
30
40
50
Context length
Contaxt length
Context length
(a) Noise: 0.05
(b) Noise: 0.1
(c) Noise: 0.2
1.7
Multi-step Inference
Multi-step Inference
1.975
Multi-step Inference
step
step
1.950
step
1.925
Lo63
1.5
1.000
L19
1..4
1.60
1.875
1.55
1.850
1.825
1.2
1.50
1.800
1.45
1.775
10
20
30
40
50
0
10
20
30
40
50
10
20
30
40
50
Context length
Context length
Context length
(d) Noise: 0.5
(e) Noise: 1.0
(f) Noise: 2.0
Figure 18: Ablation on noise (Uncertainty Quantification): Comparing One-step inference and multi-step inference [Train horizon: 100, Metric: Multi-step log-loss, Target Length: 50].
1.3
Standard
Less NNegetive top-kkllbeed)
1.2
Opprmutation invariant
1.1
1.0
0.9
0.8
0.7
0.6
0.5
10
20
30
40
50
Number Revsaled Labels
Figure 19: Active Learning: Comparing C-permutation invariant and standard causal architecture for the active learning setting.
34


Reviews:
[{"review_id": "AjQL3hQaf1", "reviewer": "Reviewer_dFZM", "summary": "Autoregressive models such as Transformer-based model architectures are increasingly used for exchangeable sequence modeling, but common approaches rely on **one-step inference**, which fails to distinguish between epistemic uncertainty and aleatoric uncertainty.\n\nThis limitation negatively impacts decision-making tasks.\n\nThe paper did multiple analyses to show that \n- multi-step inference in training is essential for reliable uncertainty quantification and downstream task performance;\n- existing architectural approaches exchangeability are flawed and computationally expensive without providing benefits;\n- designing transformer architectures that achieve full exchangeability while remaining computationally efficient remains an open challenge.", "questions": "n/a", "limitations": "n/a", "rating": 3, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Reviewer Expertise Disclaimer\nThis work falls outside my core area of expertise. While I have some familiarity with sequence modeling and uncertainty quantification, I encourage the area chairs to weigh this review accordingly, especially regarding technical novelty and empirical depth.\n\n## Strengths\n1. **Theoretical and Empirical Analysis**: The paper provides a thorough theoretical critique, in the context of uncertainty quantification, of one-step inference and supports its arguments with empirical results demonstrating the advantages of multi-step inference in specific tasks.\n2. **Architectural Insights**: It offers a detailed examination of transformer architectures concerning **exchangeability**, contributing to the understanding of their limitations and guiding future architectural designs.\n\n## Limitations and Concerns\n1. **Limited scope of novelty**: The paper's core critique seems novel but only in a narrow context. Broad limitations of single-step inference and standard Transformer architectures have been explored in broader reasoning and planning contexts (e.g., Shojaee et al., 2025; Valmeekam et al., 2023; Creswell et al., 2022). While the paper’s domain (exchangeable sequences) is more narrow, the broader concern has precedent.\n2. **Uncertainty quantification is conceptually important but underused in practice, especially in sequential decision making**: While uncertainty quantification is theoretically valuable, the claim that it is a prerequisite for intelligent agents is overstated. While uncertainty quantification is indeed valuable for certain types of decision-making, it is not a universally accepted prerequisite for intelligent behavior. Classical control systems, expert systems, and many successful RL algorithms function without explicit UQ mechanisms, despite the presence of partial observability and noisy environments.\n3. **Limited generality of the scope**: Although the paper frames its contributions in terms of “decision-making,” the empirical evaluations are limited to multi-armed bandits and active learning, and especially in exchangeable settings. These are important, but relatively constrained, problem settings.\n\nCombined, the scope of novelty and impact of the paper seems relatively limited for a NeurIPS publication.\n\n## Recommendation\nDue to my distance from the specific subdomain, I do not make a definitive accept/reject recommendation. I encourage a review from a domain expert who can more fully evaluate the novelty and technical depth of the paper.", "was_revised": true, "final_justification": "I keep my ratings and have explained in more details in the response to the authors.", "rebuttal": "## **General response for Reviewers/AC/SAC**\n\nWe thank the reviewers for their positive and constructive feedback. We are glad that reviewers appreciate the clarity and significance of our theoretical analysis on exchangeability and multi-step inference vs. one-step inference (dFZM, UKPT, Sm7a), and recognize the value and timeliness of our empirical result on failure of architectural inductive biases (dFZM, UKPT, Sm7a, JZjB). At the same time, we acknowledge the valid concerns raised about specific aspects of our contribution. In light of these, we emphasize the broader significance of our contributions and contextualize it further.\n\nThe machine learning community faces a fundamental challenge in uncertainty quantification (UQ). Modern ML systems are powered by representation learning, excelled through minimizing loss functions and benchmarking on standard datasets. Most existing UQ approaches fall outside this paradigm—requiring hand-crafted priors, complex variational approximations, or ensemble methods that don't scale to modern ML architectures.\n\nThis methodological gap creates significant inefficiencies in downstream decision-making tasks. For example, Thompson sampling—a popular algorithm for active exploration—typically relies on hand-crafted priors and closed-form posterior updates. This limits practical deployment of principled decision-making algorithms in modern ML systems.\n\n**Autoregressive Sequence Modeling: A Paradigm Shift**\n\nWhat is exciting about autoregressive sequence modeling for exchangeable sequences is that it brings UQ back into ML's comfort zone. Instead of wrestling with intractable posteriors or hand-crafted priors, we directly model observables through De Finetti's predictive view. This approach:\n- Leverages the same optimization machinery (gradient descent on likelihood) that powers modern ML.\n- Enables direct benchmarking through standard metrics (log-likelihood, downstream task performance).\n- Scales naturally to large datasets.\n- Provides principled UQ without additional hyperparameters.\n\nThis approach represents a fundamental shift in how we do UQ, making it compatible with tools and practices that have made modern ML successful. However, this promising direction involves important subtleties concerning inferential and architectural inductive biases that the community has largely overlooked. A central contribution of our work is the systematic identification and clarification of these subtleties both theoretically and empirically.\n\n**Summary of Contributions**\n\nOur work addresses a fundamental gap in autoregressive sequence modeling for exchangeable sequences, making three key contributions:\n1. Analysis of One-Step vs Multi-Step Inference (also highlighted by dFZM, UKPT and JZjB)\n    - We prove that one-step inference fundamentally loses mutual information (Theorem 2), leading to degraded UQ.\n    - We demonstrate that this limitation causes one-step inference to suffer $O(T)$ regret in bandits setting (Theorem 3).\n     - Empirically, we show multi-step inference achieves up to 60% better performance in bandits and requires 10× fewer samples in active learning.\n2. Formal Characterization of Architectural Inductive Biases (also highlighted by dFZM, UKPT, Sm7A)\n     - We introduce conditional permutation invariance (Property 1) to formally characterize what existing Transformer architectures actually enforce.\n     - We prove that conditional permutation invariance alone is insufficient for exchangeability, as it fails to guarantee the conditionally identically distributed (c.i.d.) property (Property 2).\n     - This corrects misconceptions in prior work (Müller et al., 2022; Nguyen & Grover, 2022; Ye & Namkoong, 2024) that implicitly assumed their architectures achieved full exchangeability.\n3. Computational and Performance Analysis of Architectural Choices (also highlighted by dFZM, UKPT, Sm7A)\n    - Standard causal architectures outperform specialized \"exchangeable\" architectures by ~10\\% on out-of-training horizons with 20% better training and data efficiency.\n    - Additionally, enforcing conditional permutation invariance introduces significant computational overhead at inference ($O(T^3)$ without KV caching benefits).\n\n**Summary of responses to other concerns raised by reviewers**\n1. Reviewers (Sm7a, JZjB) point out that we do not propose a new architecture to address the limitations we identified. We acknowledged this limitation in our paper (lines 406-409). This represents a fundamental challenge in the field that our work has identified but not fully resolved (see details in response to reviewers Sm7a, JZjB). Our empirical findings suggest that standard causal masking may offer a reasonable intermediate solution. However, developing architectures that satisfy both properties while maintaining computational efficiency remains an important open research direction.\n2. Reviewers (UKPT, Sm7a) raised concerns about computational overhead for multi-step inference in practice. In many practical settings where our approach would be useful, the benefits of better decision-making far outweigh the computational cost. For example, in clinical settings, the cost of data acquisition far exceeds computational costs. Additionally, one can determine an optimal lookahead horizon to prevent wasteful use of computational resources. Exploring this optimal horizon represents an interesting future direction—we also discuss the factors that impact the lookahead horizon with details in response to reviewers UKPT, Sm7a.\n3. Based on feedback from reviewers (UKPT, Sm7a), we will include a real-data experiment in the camera-ready version demonstrating the benefits of multi-step inference over one-step inference.\n\n## **Response to Reviewer dFZM**\n\nWe thank the reviewer for their thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised.\n\n>**Uncertainty quantification (UQ) is ... underused in practice**\n\nWe agree that UQ remains underutilized in practice despite its theoretical importance. **Our work directly aims to address this gap.**\n\nAs the reviewer notes, the theoretical need for reliable UQ in efficient exploration has been well-established (Osband et al., 2016; Charpentier et al., 2022; Deisenroth et al., 2011), yet practical deployment remains limited to specific domains like online recommendation systems (Li et al., 2019), clinical decision-making (Lopez et al., 2025), and hyperparameter optimization (Golovin et al., 2017).\n\n**The core issue is scalability.** Modern ML systems are powered by representation learning, excelled through minimizing loss functions and benchmarking on standard datasets. Most existing UQ approaches fall outside this paradigm—requiring hand-crafted priors, complex variational approximations, or ensemble methods that don't scale to modern ML architectures. This creates fundamental tension: while representation learning drives modern ML's success, most principled UQ methods cannot effectively integrate with learned representations.\n\nThis scalability challenge directly impacts deployment in downstream decision-making tasks. Consider Thompson sampling, a cornerstone algorithm for sequential decision-making. Traditional implementations rely on closed-form posterior updates and hand-crafted priors, making them incompatible with learned representations that power modern ML systems. This incompatibility creates a critical bottleneck that limits practical UQ deployment.\n\n**Autoregressive sequence modeling for exchangeable sequences offers a solution**—By directly modeling observables through De Finetti's predictive view, it brings UQ back into ML's comfort zone—no intractable posteriors or hand-crafted priors required. This approach is inherently scalable and compatible with modern deep learning paradigms.\n\nHowever, this promising direction involves important subtleties that the community has largely overlooked, which we address in this paper:\n- **Inferential inductive biases:** The critical distinction between one-step and multi-step inference.\n- **Architectural inductive biases:** Proper enforcement of exchangeability and its empirical implications.\n\nWe demonstrate the importance of these subtleties by focusing on data-acquisition tasks where agents must decide what data to collect next. In these settings, separating epistemic from aleatoric uncertainty is necessary for efficient performance, as shown by our theoretical results (Theorems 2-3) and empirical validation (Section 3).\n\n>**Limited scope of novelty** and **Limited generality of the scope**\n\nWe thank the reviewer for the references to broader planning literature. While those works demonstrate advantages of multi-step planning in general reasoning tasks, our contribution is fundamentally different:\n\n1. **Theoretical characterization:** We provide a formal analysis of why one-step inference fails in downstream decision-making tasks as compared to multi-step inference, quantifying the information loss (Theorem 2) and downstream impact (Theorem 3).\n\n2. **Exchangeable data focus:** Our work targets sequence models meta-trained on tabular/exchangeable datasets—a rapidly growing area with applications in recommendation systems, clinical trials, and adaptive experimentation. This differs from the general reasoning contexts addressed in above mentioned works.\n\n3. **Architectural insights:** We identify and formalize a critical gap in existing \"exchangeable\" architectures—they enforce only conditional permutation invariance, not full exchangeability, while being computationally expensive without performance benefits.\n\nMulti-armed bandits and active learning are canonical examples of decision-making on exchangeable data, providing analytical tractability while covering important practical applications. To our knowledge, both the theoretical analysis of multi-step vs. one-step inference and the architectural results for exchangeable sequence modeling are novel contributions."}, {"review_id": "tzLaMIVwVp", "reviewer": "Reviewer_UKPT", "summary": "This paper reveals a potential flaw of using Transformers trained with one-instance generation loss for posterior inference. The work shows that the current approaches to building and using sequence models for exchangeable data on regression tasks are often theoretically unsound and practically suboptimal.", "questions": "(Q1) The distinction between epistemic and aleatoric uncertainty is a key justification for multi-step inference. However, this approach introduces practical challenges. Could you provide an analysis of the computational overhead of this method? How does the inference time scale with the \"imagination horizon\" (the length of generated sequences) and the size of the model? Following this, how should one determine the optimal imagination horizon? Is there a risk that a poorly chosen horizon could lead to noisy or biased uncertainty estimates, and how sensitive are the downstream task improvements to this choice?\n\n(Q2) Following (W3), the experiments mainly demonstrate your points on synthetic data. How do you expect these findings to translate to large-scale, noisy, real-world tabular datasets where features may have complex dependencies? Would the performance gap between single-step and multi-step inference become more or less pronounced in such settings?\n\n(Q3) The paper convincingly demonstrates the limitations of applying autoregressive models to exchangeable data. This raises a fundamental question: what is the rationale for focusing on adapting Transformers for this task, rather than exploring architectures inherently designed for permutation-invariance, such as Deep Sets, or other graph-based neural networks?\n\n(Q4) What would be equivalent problem/strategy when number of steps becomes large?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength\n\n(S1) The paper's primary strength is its formal and clear-headed analysis. The distinction drawn between \"conditional permutation invariance\" and true \"exchangeability\" is a particularly insightful and valuable contribution that clarifies a subtle but critical misunderstanding in prior work.\n\n(S2) The paper provides compelling evidence that specialized attention mechanisms, despite being designed to handle exchangeability, can fail in practice. This finding serves as an important and actionable \"negative result\" for researchers, cautioning against the pursuit of complex architectural fixes that may not be theoretically sound.\n\n(S3) By demonstrating the fundamental mismatch between autoregressive models and exchangeable data, the work implicitly questions the community's default choice of Transformers for this problem class. It forces a necessary and critical re-evaluation of whether adapting ordered models is a more fruitful path than developing architectures inherently suited to permutation-invariance.\n\nWeakness\n\n(W1) The paper's motivation is well-grounded in showing that current Transformer-based approaches are flawed. However, it does not sufficiently justify why the community should continue to focus on adapting Transformers for this task at all, rather than exploring more natural alternatives like Deep Sets or other architectures explicitly designed for permutation-invariant data.\n\n(W2) The paper's critique, while sharp, is focused on a small and specific set of recent works. To broaden its impact and appeal to a general audience, the paper would benefit from contextualizing its findings within a wider survey of methods for modeling exchangeable data, thereby making a stronger case for the novelty and significance of its specific contributions.\n\n(W3) While using Gaussian Processes as a case study is a reasonable starting point, the paper's claims about modeling general tabular data would be substantially more convincing with a broader empirical evaluation. Demonstrating the identified issues and the benefits of multi-step inference on a variety of well-known, real-world tabular datasets is essential to validate the practical relevance of the findings.", "was_revised": true, "final_justification": "Thanks for the addressing my comments and answering my questions. It makes sense to limit the scope to Transformers. I still believe experiments on real-world data/applications would increase the value of the work.", "rebuttal": "We thank the reviewer for their thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised. We will incorporate relevant portions of this discussion into the camera-ready version of the paper and thank the reviewer for their constructive feedback. Additionally, we have summarized the broader significance of this work and our main contributions in our response to the first reviewer (dFZM) and kindly request the reviewer to refer to that response as well.\n\n\n>**W1: Why Focus on Transformers Rather Than Natural Alternatives?** and **Q3: Rationale for Transformer Focus**\n\nThanks for raising an excellent question about exploring alternatives like Deep Sets. We focus on Transformers for three key reasons: \n\n1. **Scalability and Practical Impact:** Transformers have demonstrated unprecedented scalability and are the current SOTA for sequence modeling. Given their widespread adoption in practice, understanding their inductive biases in the context of  exchangeable data has immediate practical implications.\n\n2. **Meta-learning Capabilities:** Unlike traditional methods that typically model single datasets with parametric distributions, Transformer-based sequence models can meta-learn across multiple datasets simultaneously.\n\n3. **Foundation for Future Work:** Our theoretical analysis establishes fundamental principles for modeling exchangeable data that transcend specific architectural choices. These findings provide critical guidance for future research, particularly in understanding: (1) the trade-offs between multi-step and one-step inference approaches, and (2) the distinction between conditional permutation invariance and the martingale property.\n\nWe agree that exploring inherently permutation-invariant architectures (e.g., State Space Models, Set Transformers, Deep Sets) represents important future work, and our findings will help inform these research directions as well. We will add this discussion to clarify our scope.\n\n>**W2: Broader Context and Survey**\n\nWe thank the reviewer for this valuable suggestion. We will expand our related work to better contextualize our contributions within the broader uncertainty quantification landscape.\n\n**The Significance of This Work:** Despite decades of research, the ML community still faces fundamental challenges in uncertainty quantification. Most existing approaches—Bayesian neural networks, ensemble methods, variational inference—fall outside ML's core paradigm of optimizing loss functions and benchmarking. They require hand-crafted priors, complex approximations, or computationally prohibitive ensembles that don't scale to modern architectures. This methodological gap creates a fundamental tension: while representation learning drives modern ML's success, most principled UQ methods cannot effectively integrate with learned representations\n\nThis scalability challenge directly impacts deployment in downstream decision-making tasks. Consider Thompson sampling, a cornerstone algorithm for sequential decision-making. Traditional implementations rely on closed-form posterior updates and hand-crafted priors, making them incompatible with the learned representations that power modern ML systems. This incompatibility has created a critical bottleneck that limits practical UQ deployment.\n\n**Autoregressive Sequence Modeling as a Paradigm Shift:** What makes this approach exciting is that it brings uncertainty quantification back into the realm of what ML does well. Through De Finetti's predictive view, we can directly model observables using standard optimization machinery, enabling principled uncertainty quantification without additional hyperparameters or architectural complexity.\n\n**Critical Gaps We Address:** However, this promising direction has important subtleties the community has overlooked: (1) one-step inference cannot distinguish epistemic from aleatoric uncertainty, and (2) specialized \"exchangeable\" architectures are based on informal intuitions rather than rigorous analysis. We address these subtleties in the paper leading to stronger foundation, and also demonstrate their effectiveness in downstream tasks.\n\nIn addition, we will expand our survey to include:\n\n- Classical Bayesian approaches and their scalability limitations\n- Non-parametric methods (Dirichlet Processs)\n- Neural approaches (Deep Sets, Set Transformers, Neural Processes)\n- The positioning of autoregressive sequence modeling within this landscape\n\nThis contextualization will clarify why our theoretical contributions—distinguishing conditional permutation invariance from true exchangeability, and demonstrating the fundamental importance of multi-step inference—represent critical advances for making uncertainty quantification practical in modern ML systems.\n\n\n>**W3: Real-World Validation** and **Q2: Real-World Dataset Performance**\n\nThank you for the helpful feedback. In the camera-ready version, we will include an additional experiment on a real-data example to illustrate the benefits of multi-step inference compared to one-step inference in the active learning setting.\n\nMeanwhile, we note that the performance gap in real settings depends on the following dataset characteristics:\n\n- **Signal-to-noise ratio:** Higher noise requires longer horizons to distinguish epistemic from aleatoric uncertainty. For a fixed horizon, the gap between multi-step and single-step inference diminishes with increased noise levels.\n\n- **Epistemic vs. aleatoric ratio:** Higher epistemic uncertainty leads to larger performance gaps between multi-step and single-step inference.\n\n- **Dimension and heterogeneity:** Greater noise heterogeneity increases the performance gap between multi-step and single-step inference. Higher dimensionality may also amplify this gap.\n\n\n>**Q1: Computational Overhead and Horizon Selection**  \n\n**Computational Complexity:** Multi-step inference scales as O(H×B×T²) where H is the imagination horizon (multi-step inference window), B is the number of trajectories, and T is the sequence length.  This complexity is H times greater than one-step inference, which requires O(B×T²) operations.\n\nWhile longer horizons generally yield better performance, for a fixed computational budget, there exists a fundamental trade-off between horizon length and the number of trajectories used for uncertainty estimation.\n\n**Optimal Horizon Selection:** The optimal horizon depends on dataset characteristics:\n- **Signal-to-noise ratio:** Higher noise requires longer horizons to distinguish epistemic from aleatoric uncertainty\n- **Epistemic vs. aleatoric ratio:** When aleatoric uncertainty dominates, the performance gap between inference methods diminishes, and requires longer horizons.\n- **Dimension and heterogeneity:** Complex dependencies may require longer horizons\n\nTheoretically characterizing this dependence represents an interesting avenue for future research.\n\n**Practical Guidelines:** We recommend a validation-based approach - incrementally increase horizon length until the empirical mean converges and variance estimates stabilize, then balance horizon against trajectory count based on computational constraints.\n\n>**Q4: Large Step Regime**\n\nFor very long horizons, we might employ a heuristic approach of batching along with a diversity criteria.  \n\n- **Batching:** Sampling multiple points based on uncertainty estimates from single step.\n- **Diversity Criteria:** Introduces diversity between sampling points."}, {"review_id": "SPhhXDp5hz", "reviewer": "Reviewer_Sm7a", "summary": "This paper investigates the optimal inferential and architectural inductive biases for modeling exchangeable sequences with autoregressive models, particularly for downstream decision-making tasks. The authors make two primary contributions: \n- First, they argue that the prevalent use of one-step inference is insufficient as it fails to distinguish between epistemic and aleatoric uncertainty. They advocate for multi-step autoregressive generation, demonstrating both theoretically and empirically that it provides superior uncertainty quantification, leading to better performance in active learning and bandit tasks. \n- Second, the paper critically examines existing Transformer architectures designed to enforce exchangeability. The authors prove that these architectures only guarantee a weaker property of \"conditional permutation invariance\" and not full exchangeability. Through empirical evaluation, they show these specialized architectures are computationally inefficient and perform worse than a standard Transformer architecture.", "questions": "1. Given that standard causal models outperform specialized ones, what is the path forward for incorporating exchangeability as an inductive bias? \n2. Multi-step inference is more computationally expensive. How should practitioners manage the trade-off between its superior performance and its higher computational cost in real-world applications?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper makes a clear and significant argument for using multi-step inference in decision-making contexts, which is supported using both theory and experiments.\n2. It provides a valuable and timely critique of prior architectural work, demonstrating that simpler causal models are more efficient and effective than complex \"permutation-invariant\" transformer designs that don't fully guarantee exchangeability\n\nWeaknesses:\n1. While the paper critiques existing architectures, it does not propose a novel, efficient architecture that can guarantee full exchangeability.\n2. The experiments, while effective, are limited to synthetic data; validation on real-world datasets would further strengthen the claims.", "was_revised": true, "final_justification": "Authors have agreed to add more empirical results in the camera-ready version. I would like to keep my \"Accept\" rating.", "rebuttal": "We thank the reviewer for their thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised. We will incorporate relevant portions of this discussion into the camera-ready version of the paper and thank the reviewer for their constructive feedback. Additionally, we have summarized the broader significance of this work and our main contributions in our response to the first reviewer (dFZM) and kindly request the reviewer to refer to that response as well.\n\n\n>**While the paper critiques existing architectures, it does not propose a novel, efficient architecture that can guarantee full exchangeability.** and **Given that standard causal models outperform specialized ones, what is the path forward for incorporating exchangeability as an inductive bias?**\n\nEnsuring full exchangeability represents a central challenge highlighted by our work (lines 406-409). Our analysis demonstrates that merely enforcing Property 1 (conditional permutation invariance) is insufficient to guarantee full exchangeability. While Property 1 can be readily enforced through attention masking schemes (as demonstrated in prior work), Property 2 (the conditionally identically distributed/martingale property) poses a fundamental challenge.\n\nThe difficulty arises because enforcing the martingale property typically necessitates imposing specific parametric constraints on the underlying distributions—constraints specifically designed to satisfy the martingale condition. However, such parametric restrictions significantly limit representational capacity, thereby undermining the core promise of Transformers and deep learning more broadly. This creates a fundamental tension between architectural expressivity and theoretical guarantees.\n\nLooking forward, our empirical findings suggest that **standard causal masking** may offer a reasonable intermediate solution for practitioners seeking to perform multi-step inference on exchangeable sequences, despite not enforcing strict permutation invariance. Additionally, we identify several promising research directions:\n\n1. **Hybrid approaches:** Combining autoregressive models with architectural components that partially enforce exchangeability while maintaining computational efficiency\n\n2. **Alternative architectures:** Exploring state-space models or Set Transformers that may offer better trade-offs between exchangeability guarantees and computational efficiency\n\n3. **Regularization-based approaches:** Developing training objectives that encourage exchangeable behavior without hard architectural constraints\n\n4. **Theoretical advances:** Characterizing when approximate exchangeability is sufficient for decision-making tasks\n\nOur finding that standard causal masking outperforms specialized architectures suggests that the inductive bias of exchangeability may be better incorporated through training procedures or regularization rather than architectural constraints.\n\n\n\n> **The experiments, while effective, are limited to synthetic data; validation on real-world datasets would further strengthen the claims.**\n\nWe appreciate this feedback. While our main paper focuses on controlled synthetic experiments to clearly demonstrate the theoretical principles, we will include an additional experiment with real-data example in the camera-ready version, showcasing the benefits of multi-step inference over the one-step inference. \n\n\n\n> **How should practitioners manage the trade-off between its superior performance and its higher computational cost in real-world applications?**\n\nThe computational overhead of multi-step inference scales linearly with the lookahead horizon (multi-step inference window) $H$, requiring $H$ times more computation than single-step inference. However, practitioners can effectively manage this computational trade-off through:\n\n**Optimal Horizon Selection:**  For any fixed computational budget, an optimal trade-off exists between lookahead horizon length and the number of generated trajectories for uncertainty estimation. This optimum is fundamentally determined by dataset characteristics, including:\n- Signal-to-noise ratio\n- Relative magnitudes of epistemic versus aleatoric uncertainty\n- Data dimensionality and complexity\n\n**Practical Guidelines:** We recommend adopting a validation-based approach whereby practitioners incrementally increase horizon length until the empirical mean converges and variance estimates stabilize. Subsequently, the horizon should be balanced against trajectory count based on available computational constraints.\n\nImportantly, in many decision-making applications (e.g., active learning, clinical trials), the cost of acquiring new data substantially exceeds computational costs, thereby justifying the additional computational overhead.\n\n**Theoretical Characterization:** Developing theoretical frameworks to predict optimal horizon lengths given specific dataset characteristics remains an important open problem and represents a promising direction for future research."}, {"review_id": "s84XBwuiv5", "reviewer": "Reviewer_JZjB", "summary": "The paper tackles both the inferential and architectural inductive biases that matter when using Transformer sequence models on exchangeable data such as i.i.d. rows in a table.\nOn the inference side, it proves that the common one-step prediction heuristic systematically discards information.\nOn the architecture side, the authors show that recently proposed “exchangeable” masking schemes for Transformers is suboptimal.", "questions": "1. Is C-permutation invariant the same as conditional permutation invariant?\n2. Based on my understanding, the proposed theory does not apply to most time series data due to exchangeable property, is that correct?\n3. Can the authors provide some concrete example on how one apply either one-step or multi-step inference? A simple real world example would help a lot for understanding the paper.\n4. Can the authors discuss how strong is the assumption $\\hat{P_\\phi} = \\mathbb{P}$ ?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strength\n1. This paper tackles sequential modeling from a unique perspective. They show how the inference routine affects performance. I think this is a nice and meaningful perspective for the ML community.\n2. The paper is well organized, it is easy to find information and follow.\n\n## Weakness\n1. Typo: line 223 (undertainty), appendix D.2 (exchangeablity), line 15 (infintely).\n2. After proving existing masks insufficient, the paper did not provide a corresponding solution that satisfies both Property 1 and Property 2.", "was_revised": true, "final_justification": "The authors have addressed my concerns, i've decided to remain my score.", "rebuttal": "We thank the reviewer for their thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised. We will incorporate relevant portions of this discussion into the camera-ready version of the paper and thank the reviewer for their constructive feedback. Additionally, we have summarized the broader significance of this work and our main contributions in our response to the first reviewer (dFZM) and kindly request the reviewer to refer to that response as well.\n\n\n>**After proving existing masks insufficient, the paper did not provide a corresponding solution that satisfies both Property 1 and Property 2.**\n\nWe acknowledge this limitation highlighted by the reviewer, which we also noted in our paper (lines 406-409). This represents a fundamental challenge in the field that our work has identified but not fully resolved. Specifically:\n\n1. **Property 1 (conditional permutation invariance)** is relatively straightforward to enforce through architectural modifications such as Set Transformers, permutation-equivariant networks, or the masking schemes we analyzed.\n\n2. **Property 2 (conditionally identically distributed/martingale property)** is significantly more challenging to impose architecturally. Traditional approaches typically assume parametric distributional forms, which limits representational power and undermines the flexibility that makes Transformers appealing for exchangeable sequence modeling.\n\nOur empirical findings suggest that standard causal masking may offer a reasonable intermediate solution for practitioners seeking to perform multi-step inference on exchangeable sequences, despite not enforcing permutation invariance. However, developing architectures that satisfy both properties while maintaining computational efficiency remains an important open research direction for future work.\n\n>**Can the authors discuss how strong is the assumption $\\hat P_\\phi =\\mathbb P$ ?**\n\nThis assumption is made primarily for theoretical clarity in Theorem 2 to isolate the fundamental information loss inherent in one-step versus multi-step inference. In practice, this assumption will not hold exactly due to finite data and model capacity limitations.\nWhen $\\hat P_\\phi \\neq \\mathbb P$, the information loss would include additional term proportional to KL divergence  $D_{KL}(\\mathbb P \\|\\| \\hat P_\\phi)$. However, our key theoretical insight—that one-step inference systematically discards mutual information between future observations—remains valid regardless of model accuracy.\n\n>**Can the authors provide some concrete example on how one apply either one-step or multi-step inference? A simple real world example would help a lot for understanding the paper.**\n\nWe provide detailed algorithmic descriptions for multi-step inference in active learning and multi-armed bandit settings in the appendices (Algorithms 2 and 3).\n\nTo clarify this with a simple example, consider the coin scenario illustrated in Figure 2, which demonstrates how multi-step inference enables the distinction between epistemic and aleatoric uncertainty—a prerequisite for effective decision-making—while one-step inference fails to do so.\n\n**The Setup:** We have two coins (A and B) and must decide which to flip first to maximize expected rewards over multiple tosses.\n\n**One-step inference approach:**\n          \n- Predicts only the immediate next outcome for each coin\n\n- Finds that both coins A and B have identical predictive uncertainty for their next toss\n\n- Cannot distinguish between the coins, leading to random selection\n\n**Multi-step inference approach:**\n\n- Simulates multiple future tosses for each coin to understand the uncertainty structure\n\n- Reveals a crucial difference:\n\n    -  **Coin B (epistemic uncertainty):** Multi-step simulation shows that after one toss, all future outcomes become deterministic. If the first toss is heads, all subsequent tosses will be heads; if tails, all will be tails. The uncertainty can be completely resolved with a single observation.\n\n   -  **Coin A (aleatoric uncertainty):** Multi-step simulation reveals that even after observing many tosses, each future toss remains uncertain due to the inherent randomness of a fair coin. The uncertainty is irreducible.\n\n- Multi-step inference correctly identifies that Coin B should be chosen first because its uncertainty can be resolved (providing valuable information), while Coin A's uncertainty cannot be reduced regardless of how many times it's observed\n\nThis example illustrates the core advantage: multi-step inference captures the reducibility of uncertainty, which is essential for optimal sequential decision-making in applications like active learning and bandit problems.\n\n\n>**Is C-permutation invariant the same as conditional permutation invariant?**\n\nYes, \"C-permutation invariant\" is our shorthand notation for \"conditional permutation invariant.\" We will clarify this notation in the revision to avoid confusion.\n\n\n>**Based on my understanding, the proposed theory does not apply to most time series data due to exchangeable property, is that correct?**\n\nThat is correct, our current theoretical framework specifically applies to exchangeable data and does not extend to time series data where temporal order matters. The exchangeability assumption is fundamental to our approach—time series data violates this assumption due to temporal dependencies.\n\n\n>**Typo: line 223 (undertainty), appendix D.2 (exchangeablity), line 15 (infintely).**\n\nThank you for catching these errors. We will correct all typos in the camera-ready version."}]

Author rebuttal:

Meta-review: The paper proposes to revisit decision making along the following lines:
*  decision making, as opposed to mainstream classification, must consider uncertainty quantification (UQ, under-used in ML practice) and thus distinguish among epistemic and aleatoric uncertainty;
* this distinction requires to consider multiple-step inference (only enabling to reduce the epistemic uncertainty) instead of single-step inference.
* the literature handles exchangeable sequences (sequential events are independent conditionally to the distribution) through specialized masking strategies;
* still, "the invariance properties enforced by such masking schemes don’t necessarily guarantee valid probabilistic inference". 

Basically, the paper introduces the property of conditional permutation invariance; shows that this property is enforced by transformer architectures, and that this property does not guarantee exchangeability, entailing a loss of performance. 

Rev UPKT appreciates the clear-headed analysis. They suggest that broader experiments are essential to appreciate whether the benefits are worth the extra-cost of the approach, and that the critique should not focus on the only transformer architectures. Authors see their main contributions as i) distinguishing conditional permutation invariance from true exchangeability; ii) demonstrating the fundamental importance of multi-step inference, and state that these represent critical advances for making uncertainty quantification practical in modern ML systems.
They agree to extend the scope and promise an experiment on real-world data in the revised version.

Rev Sm7a appreciates the valuable critique of the considered transformer architectures, with regrets that the authors do not propose another architecture and only provide evidence on synthetic data. Authors state that standard causal masking may offer a reasonable intermediate solution.
Decision: Accept (poster)