Variational Uncertainty Decomposition for In-Context Learning
Authors: I. Shavindra Jayasekera
Jacob Si
Filippo Valdettaro
Wenlong Chen
Aldo A. Faisal
Yingzhen Li
Venue: NeurIPS 2025 poster ()
Keywords: Uncertainty Quantification
Uncertainty Decomposition
In-Context Learning
Variational Methods
Large Language Models
Primary area: probabilistic_methods
TL;DR: In this paper we introduce a framework to decompose the uncertainty of predictions in in-context learning

Abstract: As large language models (LLMs) gain popularity in conducting prediction tasks in-context, understanding the sources of uncertainty in in-context learning becomes essential to ensuring reliability. The recent hypothesis of in-context learning performing predictive Bayesian inference opens the avenue for Bayesian uncertainty estimation, particularly for decomposing uncertainty into epistemic uncertainty due to lack of in-context data and aleatoric uncertainty inherent in the in-context prediction task. However, the decomposition idea remains under-explored due to the intractability of the latent parameter posterior from the underlying Bayesian model. In this work, we introduce a variational uncertainty decomposition framework for in-context learning without explicitly sampling from the latent parameter posterior, by optimising auxiliary inputs as probes to obtain an upper bound to the aleatoric uncertainty of an LLM's in-context learning procedure. Through experiments on synthetic and real-world tasks, we show quantitatively and qualitatively that the decomposed uncertainties obtained from our method exhibit desirable properties of epistemic and aleatoric uncertainty.

Variational Uncertainty Decomposition for In-Context Learning
I. Shavindra Jayasekera; Jacob Si;
Filippo Valdettaro, Wenlong Chen, A. Aldo Faisal, Yingzhen Li Imperial College London {i. ja.jaaasekee223, y y.ii3,, yingzhen. 1i] aimeeiiaraa......mm ack
Abstract
As large language models (LLMs) gain popularity in conducting prediction tasks in-context, understanding the sources of uncertainty in in-context learning becomes essential to ensuring reliability. The recent hypothesis in-context learning per- forming predictive Bayesian inference opens the avenue for Bayesian uncertainty estimation, particularly for decomposing uncertainty into epistemic uncertainty due lack of in-context data and aleatoric uncertainty inherent in the in-context pre- diction task. However, the decomposition idea remains under-explored due to the intractability of the latent parameter posterior from the underlying Bayesian model. In this work, we introduce variational uncertainty decomposition framework for in-context learning without explicitly sampling from the latent parameter posterior, by optimising auxiliary queries as probes to obtain an upper bound to the aleatoric uncertainty of an LLM's in-context learning procedure, which also induces lower bound to the epistemic uncertainty. Through experiments world tasks, we show quantitatively and qualitatively that the decomposed uncertain- ties obtained from our method exhibit desirable properties of epistemic and aleatoric uncertainty. Code is available at: https: :/te::///iii gitu........ com/ jacobyhsi/VUD.
synthetic and real-
Introduction
Large Language Models (LLMs) have demonstrated remarkable abilities in natural language genera- tion [17, 88, 102], and are being extended to a wide range of applications such as question answering [116], retrieeval-augmenteed generation [55], information analysis [97, 75], and bandit problems [49]. In particular, an emergent property of an LLM is in-context learning (ICL), where the model acquires task behavior inference time, without the need for prior pre-training or fine-tuning [12]. With the rising importance and presence of LLMs, understanding where and why these models are uncertain is essential in assessing their rustworthiness and robustness. A straightforward method of assessing uncertainty is directly prompt the LLM to quantify the uncertainty of its outputs. However, this can be unreliable due to the overconfidence of language models [106]. Therefore, being able faithfully quantify and determine the sources of uncertainties from the LLMs' output can assist practitioners in better understanding and addressing the model's limitations.
Recent work has hypothesised that ICL exhibits properties of Bayesian inference [113]. If we concatenate dataset of predictive task D ya)), and test input x" into prompt, then we can view ICL as (approximately) inferring an implicit latent parameter 0 for an underlying posterior distribution p(([D) and computing posterior predictive distribution p(y* |x+ D). This interpretation allows estimation of uncertainty through a Bayesian framework, which measures model's total (predictive) uncertainty by computing the entropy Hify a D) or, in regression settings, the total variance Var[y*|æ* [x", ,D) The total uncertainty can then be decomposed further into two
Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Predictive Task
Undettiig Biearas Stake 1o
Pesterke Chediatee
Uncentaney
Uncertainty
Eqiaemiic
LLM
ply 12
Hiy Teo P)
0) + M 0
Auziliary Task
Merth
Th
LLM
aly"ar", 2,
Byy (a 2.D)
1.0
A6
04
0000
-000
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 1: Uncertainty Decomposition with Auxiliary Data (Above).
Decomposition Example for Two-Moons Dataset (Below).
sources [45, 98]: aleatoric uncertainty, which captures noise inherent the data generation process (thus irreducible), and epistemic uncertainty that accounts for uncertainty in the model due to the lack of knowledge (reducible with more data). In the bottom of Figure 1, we motivate the importance of decomposition the two-moon classification dataset. This decomposition provides valuable insights: aleatoric uncertainty pinpoints regions of ambiguity around the decision boundary, while epistemic uncertainty exposes lacking sufficient in-context data, guiding practitioners where additional data or model refinement is needed. This notion of uncertainty decomposition has been explored various domains, including computer vision [45, 46] and reinforcement learning [81, 18].
Obtaining high-quality Bayesian uncertainty estimates and decomposition for LLM-based ICL poses two major challenges. First, an LLM's auto-regressive prediction procedure often does not satisfy the exchangeability condition [20, 118], which questions the existence of the implicit Bayesian model with latent parameter 6. Second, even if an implicit Bayesian model exists, simulate posterior samples ~ p(0 which are required by the uncertainty decomposition procedure in many existing Bayesian neural network methods [76, 11, 27, 36, 57]. In this regard, recent work on Martingale posterior [20] proposes generating long sequence of future data and estimating a posterior distribution over via risk minimisation. But the Martingale posterior approach incurs a high computational cost and, still, the missing guarantee of exchangeability makes its uncertainty estimates questionable in aligning with the uncertainty from a coherent Bayesian model.
cannot explicitly
In this work, we propose a Variational Uncertainty Decomposition (VUD) framework for LLM-based follows:
ICL, focusing addressing the mentioned two challenges. Our contributions
We propose an optimisable variational upper-bound to the aleatoric (predictive) uncertainty without explicit simulating the parameter posterior (D), by appending in optimisable auxiliary inputs Z to the context and computing uncertainty measures with Z conditioning. This variational estimator also induces lower-bound on the epistemic uncertainty, which can be used in relevant tasks. An overview of our two-task variational decomposition pipeline can be found in the above of Figure I. We propose novel LLM prompting and optimisation techniques for computing p(y" |æ" D) and searching optimal Z. Our design facilitates (approximate) exchangeability for ICL, making the variational uncertainty estimates better aligned with desirable Bayesian properties such as epistemic uncertainty reduction with increasing amount of data.
Experiments on synthetic regression and classification datasets show that our uncertainty decomposi- tion framework is effective, behaving qualitatively similar to Bayesian model. Quantitatively, the variational estimation of epistemic uncertainty also bencfits downstream tasks such as bandit and in-context abstention tasks applied to real-world natural language datasets.
2 Background
In-Context Learning and Bayesian Inference. A (pre-trained) LLM with weights & parametrises set of conditional distributions (p)((t)tt.a=))een= over tokens (teee Given predictive

---

task of covariate-label pairs. D = {( y:) y:) and test covariate 2', the ICL procedure with an LLM sets (f2--1.t2) t2x) = (i) and (t2n+1, 22112) (a", y²) and computes the predic-
tive distribution as p(y" x D) = pe 2 (620+2101:24)) Now suppose the random variables
Yi: 2inn ~ I P(ys Vci) (with ply: (eii = pe (tai 11:2--1)) are exchange- able, namely for all permutations a of [m].
ply1..... .n /n//1.1 (T Zx). p(Yo(I)..... (a(x) ((())..,,,n()))) Ta(n))
(1)
then by de Finetti's theorem [16] there exists Bayesian model w.r.t. parameter such that
ply p ....., Fn) /II p(yl/ee, ()p(O)
n
i=1
Notably, the parameter 0 here is defined implicitly. We discuss the link between ICL and Bayesian models as well as existing methods to promote exchangeability further in Appendix D and F.
Decomposing Predictive Uncertainty. Consider a prescribed Bayesian model y|æ ~ p(y)a ar 0) with prior 0 ~ p(0). Given dataset D
{ y+) )+))² J==1 we can (approximately) compute the posterior (D) Then the predictive total (entropic) predictive distribution p(y" a=" D) ply x
uncertainty is defined as Uyy") x D) = Hilp(y" |a"', D)J, which can be decomposed further into
aleatoric uncertainty (y(y" (x" D) and epistemic uncertainty U
(y x D) [45]:
H[p(y" [x",, D)] Ep(oyD) Bp(or))(Hlpp(>" [H]pp(y |x=", 0)]] yyy 00="", D)
(3)
=UU(y|aæ- (D)
=y (D"
The two different notions of uncertainty have distinct statistical interpretation presented as follows.
Aleatoric uncertainty measures the inherent and irreducible randomness data. Technically, under model correctness and identifiablity assumptions, there exists a parameter 8 such that p(y (x,6")
Paata (y) where DD iid. ~ P is the data distribution. Therefore the inherent stochasticity in data prediction can be measured via entropy El[pam [a=")] IHI p([ppy"lpp l<=", ,0"), However, 0* is unlikely to be recovered precisely from finite observations in D. Instead Ua(y") , defines a Bayesian estimator of aleatoric uncertainty, by considering the uncertainty in 0 (described by the posterior p( DD) and averaging the entropy Hi[py"' [x", "2""00 8)] over plausible @ ~ p(0) D). This estimator will converge to the true aleatoric uncertainty H. Paata (y ( ).. if p(0 + 5(0 = 0r) as
[D] + oc, We also refer to c.g., [98] for additional discussions regarding this Bayesian definition.
Epistemic uncertainty reveals the model's 's uncertainty in prediction due to lack of knowledge from data, which reducible by adding in new and meaningful data. Specifically, by definition of 0 D] Bpty-|= [KKL [p(Oly", x², [[]D)]] shows another interpretation of epistemic uncertainty as the expected information gain of acquiring new datum (x","") under the current posterior belief p(0) (D) This motivates Bayesian active learning [40, 25] and Bayesian optimisation [63, 100, 99, 37] with epistemic uncertainty to assist the exploration-exploitttion process, On the other hand, writing [[y=;0/=", D] = [ epistemic uncertainty is reflected by the disagreement between "experts" from the posterior 0 ~ D) This leads to the use of epistemic uncertainty in detection tasks for e.g., out-of- distribution data and adversarial inputs [56].
[DKI[p(y" < "0)lll0("""l |xt D)Il.
When y E IR, we can also use variance as the uncertainty measure, meaning that we can compute the total variance of the prediction, and perform a similar decomposition into aleatoric and epistemic variances by the tower rule property:
Varly |z. D) = E [Var[y|l2" + Varp(eID) B[y"|0"",)
-y
(ar,))
Typically. these decompositions are obtained by Monte Carlo estimation with (approximate) samples from p(0) D) [51]. However, this approach poses a challenge when we don't 't have access to p(0 D)) which may occur if the Bayesian model is only implicitly defined [113] as in Eq. (2), or if sampling from D() is prohibitively expensive.
3 Method
We present an alternative approach for uncertainty decomposition defined in (3) and (4), which sidesteps explicit posterior sampling of the parameter 0 and thus, suitable for implicitly defined Bayesian models.
3

---

Although our practical algorithmic development focuses on LLM in-context learning on context D {( y=))=_ and test query æ", the decomposition technique applies to any Bayesian model a de Finetti (2), including prescribed Bayesian mod- els such as Bayesian linear regression and Gaussian processes (Appendix B).
x
Z
Figure 2: The DAG G of the condi- tional independence assumptions.
3.1 Variational Estimates of Uncertainty Decomposition
Total Uncertainty Decomposition. Suppose we can directly compute (or approximate) the posterior predictive distribution p(y"|=", ,D) for arbitrary 'D and x". Now consider a set of auxiliary inputs ("queries") "= and corresponding outputs ("answers") as U {u, }-- Then we define the following variational estimation of the aleatoric uncertainty as:
Va x", D) == Ep(U ( x=", U, Z, D)]]
(5)
To ensure consistency with an underlying Bayesian model (2), we assume that x", , Z, U, D obey the conditional independence relations given by the directed acyclic graph (DAG) G in Figure This assumption allows us to prove the following theorem relating the variational estimation of the aleatoric uncertainty to the exact Bayesian estimate of aleatoric uncertainty.
Theorem 3.1 (Aleatoric Uncertainty Upper-Bound). If the conditional independence relations in G hold, then the variational estimator provides an upper-bound to the aleatoric uncertainty:
Va x Z,D) 2 Ua(y"<x",
(6)
where the gap between Ug (y" x D) and Va x D) is:
Hy 0|x", U. Z, DI] = Eu Z,D) [DRa[p(Oly x", U, D. RI(,. DKI.plyy w=" ",",))ll("") U, Z, D))].
(7)
See Appendix A.1 for the proof. Importantly, the upper-bound (6) holds for arbitrary Z which inspires the following optimisation procedure to obtain the best variational estimate:
Vy [x", D) min V. (y" (x" D),,,
(8)
Since the aleatoric uncertainty trivially upper-bounded by the total uncertainty in (3), we denote
V. [x"" D) (in(Va "x", D), "
as the variational estimate of the aleatoric uncertainty. We obtain variational estimate for the epistemic uncertainty by defining Ve(y"lx*, DD) Hip(y" xx" [)] Voly" xx" DD), which implics that Ve(y" [x"" D) < (y"" x= DD). and the gap between (y x D) and Ve(y"l2** D) again ()) [[[y> 0 xx" U, ZZ DiJ. This motivates Variational Uncertainty Decomposition [H
approach illustrated in Figure We discuss another iiformation-theoreee view in Appendix A.1.
The effectiveness of this variational decompasition hinges on the choice of Z to optimise (8), which equivalent minimising the gap (7). Critically, similar to the two interpretations of the epistemic uncertainty (y* x= DD) presented in Section 2, this gap can also be viewed from two angles.
Residual information gain in fantasy: From the first definition of mutual information in (7), see that this gap quantifies the (expected) residual information gain of acquiring a new datum (x", assuming the model has further fantasised observations (Z, U) in addition D. Therefore, "elever queries" " Z, together with the fantasied answers U, should provide sufficient information regarding the model's epistemic "belief" in 0, such that further observing y" and x² does not. provide much more certainty in 0.
Remaining disagreement in fantasy: Alternatively, from the second definition of mutual infor- mation in (7), we see that this gap also captures the expected amount of remaining disagreement between posterior experts after conditioning on additional fantasised data (Z, U). Therefore, "clever queries" Z should be constructed by encouraging model agreement in epistemic "belief" of the answer y to the target query x". after fantasising the answers U to the queries Z.

---

Uncertainty Decomposition Algorithm
Step 3: Compute pfy" 22 D)
Step 1: Chocse Auxiliary Z
Step 2: Compute D)
D, Zi, M
plalx U uc, 31
Zi...., Zm
D, 24 - LIM
P(U 2, D)
: Lis Magee"
I.M
plale
Step 4: Manginalise U
Step 5: Filter won- non-Bayesian 2
Step 6: Compute Minimum Entropy
pls"|=", 2) Byme,
Duuply |z' P)lpll/ |a Z, 2))
min (V'(""22", 2,, 2,, Feghe
2, "000,
c d a x
Figure 3: Variational Uncertainty Decomposition (VUD) Framework.
As result of increased certainty of the model's subjective beliefs (in 0 and/or in y' given x") after observing the fantasied data (Z, U), the conditional entropy, Va (y Z. D) is a suitable proxy for the exact Bayesian aleatoric uncertainty estimate Ug (y [x"",>" D). It remains an upper bound because some of the epistemic uncertainty in 0 is absorbed into the aleatoric uncertainty conditioned on U. which is reflected by the conditional expectation of the entropy of p(y" |x"", Z. J p(y"'=", ()) Z. D)dt when computing Va xx" Z)
Total Variance Decomposition. Similarly (8), we can also construct a variational estimate for the aleatoric variance and derive a corresponding upper-bound. See Appendix A.2 for the proof.
Theorem 3.2 (Aleatoric Variance Upper-Bound), If the conditional independence relation in holds, then the variational estimator provides an upper-bound to the estimation of aleatoric variance:
2 (y"lx" Z, D) == Bp(UIZ,D) \Varly"I=* |x", U, U.Z, DJ] D, 2 (y""""0",
(9)
The best variational estimate is then V (y x." D) minz yy -== D), and lower-bound of the epistemic variance obtained V (y" x D) Varly x. D V. (y*|x* D)).
3.2 Optimising the Variational Estimates and Promoting Exchangeability
The presented decomposition technique requires the model to be Bayesian la de Finetti (2) and compatible with the DAG (Figure 2). which is not the case if naively prompting LLM for in-context learning. Specifically, exchangeability requires ensuring the following necessary conditions [10, 118]:
(C1) ply; F... ply ( for all i E N4 & all permutations a on [i]; (C2)
p(y" x" (, = [ p(y x" U, Z. (U), = p(y""=",))
Derivations of these necessary conditions can be found in Appendix D. To promote exchangeability for LLM in-context learning, we propose two strategies tailored for the above conditions. First, to approximately achieve (C1), we construct the predictive distribution by shuffling the context and ensembling the LLM's predictions, i.e., we define for context D y:) and test query x (with S,, a uniform distribution over the permutations on [n]):
L
px"" D) = L
ΣD
P
2n+2 2
(x" (
(o)())) ot ~ Sn-
(10)
1=1
The other distributions PUU Z, D) and p(y" U, Z, are defined in the same manner. For classification tasks, we evaluate the LLM logits to compute (10). However, in the regression case, we make a further Gaussian approximation to (10), which allows for easy computation of the entropy and marginalisation. Further details can be found in Appendix E.2. Then to approximately satisfy (C2), we restrict the search of Z (Eq. (8)) to ensure the solution satisfies
DKL (py" x Z, D) I| p(y" x=" [D] <C
(11)
for some e > O. Any metric or divergence on probability distributions will suffice for (11) but we choose KI. divergence due to ease of computation. We filter out the Z candidates that violate this KI. constraint, hence we name this step as KI. filtering Choosing the number of permutations L. and the threshold for KL filtering of Z determines the accepted level of Bayesian approximation in
5

---

the variational decomposition. While the selection of L is mainly determined by the computational resources, the choice of is further discussed in Appendix D.4.
Lastly, to reduce the search space of Z for efficient computation, we restrict Z to contain a single example in domain, i.e., m 1 and Z 2, and design sampling techniques to obtain candidates for optimal Z, including random sampling, setting Z = x", perturbing Z around æ" and a Bayesian optimisation strategy [99]. Empirically we find that perturbing Z around x* works best for inputs that lie in a continuous space, which can partly be explained via the Gaussian process example in Appendix B. For natural language tasks such as qqusstion-answers "rephrasing" x" with another LLM. Further details regarding the sampling procedures we explored for perturbing Z are in Appendix C. Our overall step-by-step Variational Uncertainty Decomposition framework (VUD) is depicted in Figure 3. Detailed decomposition algorithms for classification and regression tasks are provided in Appendix E.I.
(QA), we conduct the perturbation of by
4 Related Work
Our work takes inspiration from the growing body of literature connecting ICL to Bayesian inference [118, 113, 41, 61]. While much of the existing research centers estimating a latent concept, often through methods like the Martingale posterior [20, 113], we take a different route by approximating conditional entropy and mutual information using auxiliary data. While our work is not the first to decompose predictive uncertainty in LLMs into aleatoric and epistemic components, prior approaches define these uncertainties differently from their traditional definitions in Bayesian deep learning [45, 18, 108]. Huo et al. [39] analyse how uncertainty changes when prompt is modified with additional "clarifications." " While this is similar in spirit to our use of perturbations, we append perturbations to the ICL data rather than the predictive task itself. Moreover, their approach attributes aleatoric uncertainty solely to input ambiguity and does not incorporate Bayesian framework, leading to a definition of uncertainty that diverges from the standard Bayesian interpretation. Ling al. [60] assume a Bayesian approach but use alternative non- standardd definitions of aleatoric and epistemic uncertainties. We provide a more detailed discussion of these related works, along with applications to OOD detection and bandit problems, Appendix
5 Experiments
We evaluate the robustness and applicability of our method to classification and regression tasks. This includes ablation studies and visualisations synthetic datasets, as well downstream applications such as bandit problems and ditt----istributionn (OOD) detection on question-answering (QA) tasks. We use the following LLMs in our experiments: Dowe2..5-14B77B [88] and Llama-3.1 [102]. Only for QA tasks, we use Qwen2.5- 1-eee2.5-14BB--- 14B-Instruc.. For conciseness, we show results for Qwen2.55 14B/14B-lnstruet in the main text and the results for the remaining LLMs and baselines are given in Appendix G. Prompts and sampling details are provided in Appendix H.
5.1 Synthetic Regression & Classification Datasets
We visualise the uncertainty decompositions on synthetic regression & classification datasets and conduct ablation studies on the effects of KL filtering and Z choices. Further ablations regarding permuting the in-context examples and various LLMS are in Appendix C and D.
Visualisations. In Figures 4a and 4b, we visualise the VUD uncertainty decompositions for 1-D logistic regression (classification) and a 1-D linear regression (regression) task, each conditioned on a set of D = 15 in-context examples (vertical lines). We consider more complex tasks of the Two Moons dataset (class.) in Figure 1, a dataset with designated "gaps" and heteroscedastic noises in the in-context learning data (reg.) in Figure 5, and the Spirals dataset (multi-class class.) in Figure 6.
Across these examples, we observe similar qualitative characteristics of the uncertainty decomposition. The epistemic uncertainty (represented by the gap between the total and aleatoric uncertainty in the I-D examples) lowest regions near demonstrations and increases as the distance to the in-context learning data increases. In the classification examples, the aleatoric uncertainty is sharply localised near the decision boundary of the problem where p(y"|="", ")="")) Rs 0.5. the regression setting of Figure 4b, we observe minimal change in the alcatoric uncertainty, which reflects the homoscedastic

---

1.0
Total Uncectalnty
Aleatoric Uncerbiinty
Unceriainty
0.8
ICL Data: y-u
Uncertainty
ICL. Datr y
0.6
0.4
Uneeraiiit
Alesturic Unceriaiaty
0.2
15
10
10
15
15
10
5
0
5
10
15
Test Covariate
Tet Covariate
(a) Logistic Regression (Classification)
(b) Linear Regression (Regression)
Figure 4: Uncertainty Decompositions for Logistic and Linear Regressions.
Uncert aiuty
0
Total Uncertaincy
-5
Prediction.
Aleasocic Unceetaiity
ICL Data
ICL Data
10
5
Test Covariate
Test Covariate z
(a) Predicted Mean and Standard Deviation
(b) Uncertainty Decomposition
Figure 5: Uncertainty Decompositions for Regression Tasks with Gaps in ICL Data.
noise of the data observations. However, in Figure 5 where we have heteroscedastic noise, the model accurately distinguishes between regions of high and low heteroscedastic noise. These examples indicate that the model can correctly distinguish between uncertainty from inherent data noise and uncertainty arising from missing contextual information.
Ablations. In Figure 7, we analyse the behavior of uncertainty decompositions as a function of in-context dataset size [D] under logistic regression setting. We consider both in-distribution test inputs (x = 0, 5, solid lines) and -iitt--diitiiittionn test inputs (I = -15, -10,- -5, 10, 15, dotted lines). As expected, Figure 7a shows decreasing epistemic uncertainty across all test covariates with increasing (D), since additional training examples reduce model uncertainty. The largest epistemic uncertainty occurs at outtof-diistibution inputs (æ - 15, 100, -5, 10, 15), while in-distribution inputs (x = 0, 5) consistently exhibit lower values. The decay is most rapid for in-distribution test points, suggesting that the model becomes confident more quickly when the test point distribution overlaps with the training data. In contrast, alcatoric uncertainty reported in Figure 7b remains relatively stable as [D] grows, particularly for out-of-distribution covariates. Notably, aleatoric uncertainty is highest for the decision boundary at æ = O, where the class overlap is greatest, and remains consistently elevated across all dataset sizes. diitt--distriuttion points show slightly lower but stable alcatoric values, reflecting lower intrinsic class ambiguity at extreme covariates. The mild increase in alcatoric uncertainty for in-distribution points at small dataset sizes is likely due model underfitting, which resolves as more data is provided.
In Figure 8, we compare the computed aleatoric uncertainty across different Z sampling methods under the logistic regression setting. These include Perturb, where small noise is added to the test example create Z; Repeated, where Z is chosen to be the test example itself: Random, where Z is sampled uniformly from the dataset; and Bayesian Optimisation (BO) [99], where Z is actively selected to minimise utility function related the uncertainty. The aleatoric uncertainties reported in Figure 8a show that all these approaches track the total uncertainty curve around the decision boundary, indicating strong performance in capturing the local uncertainty landscape. Among them, Repeated returns the lowest variational alcatoric uncertainty estimate, Perturb also provides lower estimates, closely following the peak and providing stable estimates across the covariate space. Random sampling shows an upward trend in low ICL density regions far from the decision boundary. indicating poor stability. Regarding the KL divergence (11) achieved by the selected Z in Figure 8b, Random and BO consistently have the lowest KL divergence across the majority of test samples, followed by the Perturb method which significantly faster than BO. The Repeated sampling method yields higher KL values than Perturb, indicating greater deviation from the predictive posterior and thus less aligned with Bayesian principles. These evidences support Perturb as scalable and well perffrerforming approach for sampling candidate Z in (8)'s optimisation procedure.
7

---

1.00
1.
#
0.75
22
0.50
0.50
525
0/3
0.00
3.00
(a) Posterior Prob. (b) Total Uncertainty
(c) Aleatoric Uncertainty (d) Epistemic Uncertainty
Figure 6: Uncertainty Decompositions for Spirals Classification Task.
Uncertainty
Unrataiety
0.2
0.75
100
10.0
15.0
0.50
25.0
0.0
6.0
20
40
60
50
20
80
80
5.0
[P]
[P]
(a) Epistemic Uncertainty vs. Size Training Set
(b) Aleatoric Uncertainty vs. Size Training Set
Figure 7: Uncertainty decompositions for logistic regression task with varying dataset size. Solid and dotted lines indicate in-distribution and -iit----iistttttioon predictive points respectively.
5.2 Downstream Applications of Uncertainty Decomposition
We conduct quantitative experiments on two applications of uncertainty decomposition: bandit problems and outttf-distribution detection in real-world guestion-answering tasks.
Bandits. Bandit problems in reinforcement learning necessitate the ability to distinguish between aleatoric and epistemic uncertainty balance exploration and exploitation. In bandit problem, for trial f. an agent must choose an arm a E A which gives reward The goal is to minimise the overall regret over all the trials Σ, 14; E[rr]. where μ4 is the mean reward from the optimal arm. We consider Upper Confidence Bound (UCB) bandit algorithms [5], where ag argmax, (a aU,(a) (a), where Q+ is the estimated reward from arm a and U, is the uncertainty in arm a at trial I, and a is the exploration rate. We use the L.LM posterior mean Q,, and compare the performance of epistemic and total variance as U.. In this setting, epistemic variance guides exploration choose arms where additional data is beneficial, whereas total variance may prioritise actions where the reward has high intrinsic noise. We use the multi-armed bandit "Buttons" task [49], with 5 arms, where each arm yields Bernoulli reward with mean Pa. The base reward level p controls the overall success probability, with the optimal arm set to P- = p + and all other arms set to
P. = = A where A denotes the reward gap between the optimal and suboptimal arms. We set A = 0.2, which is the "hard" setting in [49]. When p 0.5, the reward for the optimal arm will have the lowest (aleatoric) variance, and UCB algorithms using total variance will choose more suboptimal actions. We use mean regret and worst-case mean regret (from the 30% of worst performing seeds) as the primary performance metrics as well as metrics of median reward, suffix-fail frequency and K MinFrac used in [49]. We also include UCB1 and Greedy as a non-LLM baseline, and the instruction prompting method from [49] an LLM-based non-uncertainty baseline. See Appendix G.4 for further details on metrics, results and implementation of the LLMMUCB algorithm.
Figure 9 shows a typical run of epistemic variance (EV) and total variance (TV) for a particular seed. In both examples, the Q value (posterior mean reward) for the optimal arm is the highest in the last 50 trials (Figure 9b) and thus should be chosen. But when we consider the arms chosen, the optimal arm not picked in the last 50 trials for the TV run (Figure 9a). This is because the epistemic variance decreases to zero with the number of observations but total variance does not (Figure 9c). Therefore, in the EV setting (a) dominates (a) for large f, whereas for TV setting this does not necessarily hold. Table shows our experimental results on the Buttons task. We see for p > 0.5, the worst-case regret is significantly lower for EV than TV, indicating that the UCB algorithms is more robust for EV. Furthermore, EV generally results in lower mean regret for p > 0.5 with the exception of p= = 0.6, a = However, it is important to note bandit algorithms have high variance in mean regret due to the stochasticity of the reward.
8

---

1.0
10"
Total Unceriainty
Uncertainty
0.8
Permin Repeated
KL. Divorgennee
10
0.6
Bandon
30-4
ICL Date: p
KLI Datai 1
Pattub
30
Repected
0.2
Rendom
16
-10
16
Test Covatiate
Test Covariste
(a) Rifecc of 2 on Uncertainty Estimates
(b) KL. Divergence of Optimal Selected Z
Figure 8: Ablation of Z choice on Aleatoric Uncertainty and KL Divergence.
0.8
0.2
Ara.d
Varduree
0.2
0.00
207
Veal
Tral
Trail
AcK
0.5
0.46
Are.5
Ara.4 (ppeaa)
0.40
0:33
100
50
100
150
300
(a) Cumulative Arm Counts
(b) Q-values
(c) U-values (Variance)
Figure 9: Example Run (p - 0.6, 5) with Epistemic (above) and Total Variance (below).
In-Context Abstention. Whilst LLMs demonstrate strong capability across a diverse range of in-context learning tasks, they have been shown to also hallucinate and provide false information [44, 64, 47], impacting the reliability of these models. LLM--bstention is a domain that addresses this problem by abstaining from answering questions when there is high uncertainty [105, 21]. This achieved by abstaining from questions where the LLM exhibits uncertainty above a predetermined threshold. In abstention, we want to avoid answering under settings where the prediction is inherently uncertain and therefore, we would expect to see an improvement accuracy when we use aleatoric uncertainty to threshold abstention compared to total uncertainty.
In our experiments, we apply LLM-abstention binary clas- sification datasets: BoolQA [15]. HotpoiQA [116], and Pub- MedQA [42]; as well multiclass classification dataset: MMLU [31]. We then extract the total uncertainty (TU), and the decomposed aleatoric uncertainty (AU) using VUD across 100 questions from each dataset. Our results in Table demonstrate that under threshold rate of rejecting the highest 10% of uncertain samples, using AU to threshold yields larger predictive accuracy performance gains than TU. In Figure 10, we show an example of abstention across range of threshold rates for the MMLU dataset. We observe
0.72
0.71
Aleteika (Prutike Drappel)
Effect of In-Context Ab- stention on Accuracy (MMLUUMoral). that thresholding by AU improves accuracy questions that are answered.
Figure 10:
We provide further details on this experiment and an example from the MMLU-Moral dataset comparing filtered samples of high aleatoric vs total uncertainty in Appendix G.5. We find that as aleatoric uncertainty measures inherent randomness and stochasticity in p(y|2z, AU, when compared with TU, focuses more on evaluating whether the answer to the question is inherently ambiguous.
5.3 Summary of Additional Experiments
We present in Appendix G further studies on additional baselines and applications. In Appendix G.2, we provided a "Martingale posterior" [22] version of predictive uncertainty decomposition. Our results show that the decomposition result is highly sensitive to the choice of the proxy model, and the Martingale posterior's is estimate of total uncertainty does not agree with the predictive entropy of the L.M in-context predictive distribution. In Appendix G.5, we consider an "in-context out-of- distribution" task [32]. Our goal is to demonstrate that leveraging epistemic uncertainty from our

---

Table 1: Buttons Bandit Problem. TV is Total Variance and EV is Epistemic Variance.
WORST-CASE
REWARD
Mintte
L'BI
0.128 EtP
0.310
0.29
GREEDY
0.199 120
0.101
0.525
0.460
0.03
0.5
INSTRUCT BASBLINE
0,
0. 0.107
0.495
0.0
0,26
TV (a
2)
0. 0.196-..
0.100
0.492
0.3
0.0.
BV la
0,147-4m
0.087
0.522
0.12
TV (o
0.198_cc
0.100
0.492
0.7
0.04
0.152-.11
UCBI
0.127-c.a
0.094.x07 so7
0.610
0.0
0.28
GREEDY
0.199 130
0.092 .000
0.6:15
0.396
0.03
0.06
INSTRUCT BASELINE
0.111-007
0.076..00
0.620
0.0
0.18
TV (a 2)
0.198=c1
0.005
0.670
0.1
0.04
BV la
0.149
0.068
0.642
0.145
TV (a
0.199__19
0.158_.os
0.555
0.6
0.04
0.140-.415
0.105
UCBI
0.122-t:7
0.094.x0:
0.710
0.0
0.27
GREEDY
0.199
0.085
0.760
0,369
0.03
INSTRUCT BASELINE
0.132-16
0.087
0.703
0.18
TV (a 2
(. (.199_(1)
0.076
0.725
0.1
0.03
BV
0.092-4m
0.050
0.735
0.11
TV ta
0.195
0.151
0.603
0.7
0.04
EV
0.092
Table 2: Abstention on QA tasks. Accuracies are computed on the remaining questions after the uncertainty-base question filtering approach. Higher accuracy improvement when filtering using aleatoric uncertainty (AU) highlights the effectiveness of the uncertainty decomposition.
ACCURACY (%) 1
DATASET
W/OUT FILTERING
w/ FILTERINN (TU)
W/ FILTERING (AU)
BooLQA
80.02 #.001 86.99 :
85.56 #.002
85.56 #.003
HOTPOTQA
87.78 + 000
90.00+
PUBMEDQA
87. 02 ..022
88. 18889=.00
90.00
MMLU-CS
81.01 = 003
85.56 002
86.67 7 001
MMLU-M
68.0 0..033... #.001
70.00 +
72.22 +.002
decomposition yields higher OOD detection accuracy than directly utilising the total uncertainty. We observe that for our method, epistemic uncertainty (EU) yields higher AUC scores in more ID/OOD settings than total uncertainty (TU), implying better OOD detection results via decomposition.
6 Conclusion
In this work, we introduce the Variational Uncertainty Decomposition framework for ICL in LLMs. Motivated by Bayesian view of ICL, we use auxiliary data to derive variational upper bound to the aleatoric uncertainty and variance. This permits the estimation of the aleatoric uncertainty and variance, without requiring an estimation of the latent Bayesian parameter 6. Through extensive experiments using synthetic toy and real-world datasets, we demonstrate that our method provides a sensible decomposition that qualitatively and quantitatively respects properties of epistemic and aleatoric uncertainties. These results show that our method is capable of accurately distinguishing between aleatoric and epistemic uncertainty across variety of L.LMs.
Limitations. We assume that ICL behaves in Bayesian manner. Whilst there is some evidence to support this Bayesian hypothesis [113, 118, 74], has also been observed that in longer sampling horizons this Bayesian hypothesis breaks down [20, 61]. We address this by considering short sam- pling horizons, permutations, and a filtering step to remove thon-Bayesian samples. However, whilst the filtering condition is necessary for Bayesian model, it is not sufficient and doesn't 't guarantee Bayesian behaviour. Therefore, we view our method as approximately Bayesian where e a quan- tification of the Bayesian approximation. Secondly, we focus on regression and classification tasks where the output of the task is real number or a small set classes and our prompt structure ensures short responses. In many real-world settings, the I.M output is in natural language where responses can differ in tokens but have the same semantic meaning. Therefore, uncertainty quantification methods that consider semantics [50] can be integrated with the VUD algorithm to obtain posterior over the natural language response, and we leave this as future work.
10

---

Acknowledgements
ISJ and YL acknowledge the EPSRC StatML CDT (EP/S023151/I) and BASF funding through EPSRC prosperity partnership programme IConIC (EP/X025292/1). FV was supported by Impe- rial College London Department of Computing PhD scholarship. AAF holds UKRI Turing Al Fellowship (EP/V025449/1).
References
Toghrul Abbasli, Kentaroh Toyoda, Yuan Wang, Leon Witt, Muhammad Asif Ali, Yukai Miao, Dan Li, and Qingsong Wei. Comparing uncertainty measurement and mitigation methods for large language models: A systematic review. arXiv preprint arXiv:2504. 1504:25604.11466,, 2025.
[2] Gustaf Ahdritz, Tian Qin, Nikhil Vyas, Boaz Barak, and Benjamin L the knowable from the unknowable with language models. arXiv preprint arXiv:2402.03561, 2024.
Edelman. Distinguishing
[3] Shuang Ao, Stefan Rueger, and Advaith Siddharthan. Css: contrastive semantic similarity for uncertainty quantification of llms. arXiv preprint arXiv:2406.03 0...24066033.58 0388 2024.
[4] Dilip Arumugam and Thomas I. Griffiths. Toward efficient exploration by large language model agents. arXiv preprint arXiv:2504.20997, 2025.
[5] Peter Auer. Using upper confidence bounds for online learning. In Proceedings 41 annual symposium on foundations of computer science, pages 270-279. IEEF, 2000.
[6] Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International conference machine learning, pages 263-272.
PMLR, 2017.
Oleksandr Balabanov and Hampus Linander. Uncertainty quantification in fine-tuned 1lms using lora ensembles. arXiv preprint arXiv:2402.12264, 2024.
[8] Neil Band, Xuechen Li, Tengyu Ma, and Tatsunori Hashimoto. Linguistic calibration of long-form generations. arXiv preprint arXiv:2404.00474 2024.
[9] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint arXiv: /.0iv:/800.00062,, 2018.
[10] Patrizia Berti, Luca Pratelli, and Pietro Rigo. Limit theorems for a class of identically distributed random variables. The Annals of Probability, 32(3A):2029-2055 2004.
[11] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncer- tainty in neural network. In International conference on machine learning, pages 1613-1622. PMLR, 2015.
[12] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[13] Nicola Cecere, Andrea Bacciu, Ignacio Fernández Tobías, and Amin Mantrach. Monte carlo temperature: a robust sampling strategy for Ilm uncertainty quantification methods. arXiv preprint arXiv:2502. 18389, 2025.
[14] Wenlong Chen and Yingzhen Li. Calibrating transformers via sparse gaussian processes. In The Eleventh International Conference Learning Representations, 2023.
11

---

[15] Christopher Clark, Kenton Lee, Ming-Woi Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Minneapolis,
Language Technologies, Volume I (Long and Short Papers), pages Minnesota, June 2019. Association for Computational Linguistics.
[16] Bruno de Finetti. Funzione caratteristica di un fenomeno aleatorio. Atti del Congresso Internazionale dei Matemalici, pages 179-190, 1929.
[17] DeepSeek-A1 AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. I. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li. Jiawei Wang. Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li. Junxiao Song, Kai Dong. Kai Hu, Kaige Gao. Kang Guan, Kexin Huang, Kuai Yu, Lean Wang. Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang. Liyue Zhang, Meng Li, Miaojun Wang, Mingehuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang. Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang. Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang. Wenjun Gao, Wengin Yu. Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang. Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang. Y. X. Wei, Y. X. Zhu, Yang Zhang. Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang. Yi Yu, Yi Zheng, Yichao Zhang. Yifan Shi, Yiliang Xiong, Ying He, Ying Tang. Yishi Piao. Yisong Wang. Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma, Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, 7. F. Wu, 7. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang. Zhewen Hao, Zhibin Gou, Zhicheng Ma. Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. Deepseek-v' v3 technical report, 2025.
[18] Stefan Depeweg, José Miguel Hernández- Lobato, Finale Doshi- and Steffen Udluft. Decomposition of uncertainty in bayesian deep learning for efficient and risk-sensitive learning, 2018.
[19] Ekaterina Fadeeva, Aleksandr Rubashevskii, Artem Shelmanov, Sergey Petrakov, Haonan Li. Hamdy Mubarak, Evgenii Tsymbalov, Gleb Kuzmin, Alexander Panchenko, Timothy Baldwin, et al. Fact-cheeking the output of large language models via token-level uncertainty quantification. arXiv preprint arXiv:2403.0440
2024.
[20] Fabian Falck, Ziyu Wang, and Chris Holmes. Is in-context learning in large language models bayesian? a martingale perspective, 2024.
[21] Shangbin Feng, Weijia Shi, Yike Wang, Wenxuan Ding, Vidhisha Balachandran, and Yulia Tsvetkov. Don't 't hallucinate, abstain: Identifying Ilm knowledge gaps via multi-llm collabora- tion. arXiv preprint arXiv:2402.00367, 2024.
[22] Edwin Fong, Chris Holmes, and Stephen G Walker. Martingale posterior distributions. Journal of the Royal Statistical Society Series B: Statistical Methodology, 85(5):1357-139, 2023.
[23] Sandra Fortini, Lucia Ladelli, and Eugenio Regazzini. Exchangeability, predictive distributions and parametric models. Sankhyā: The Indian Journal of Statistics, Series A, pages 86-109, 2000.
12

---

[24] Sandra Fortini and Sonia Petrone. Exchangeability, prediction and predictive modeling in bayesian statistics. Statistical Science, 40(():40-61) 2025.
[25] Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. In International conference on machine learning, pages 1883-192,, PMLR, 2017.
[26] Yonatan Geifman and Ran El- Vanii,, Selective classification for deep neural networks. Ad-
vances in neural information processing systems, 30, 2017.
[27] Alex Graves. Practical variational inference for neural networks. Advances in neural informa- tion processing systems, 24, 2011.
[28] Yu Gui, Ying Jin, and Zhimei Ren. Conformal alignment: Knowing when to trust foundation models with guarantees. Advances in Neural Information Processing Systems, 37:73884 73919, 2024.
[29] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pages 1321-1330. PMLR, 2017.
[30] Dan Hendrycks, Steven Basart, Mantas Mazeika, Andy Zou, Joe Kwon, Mohammadreza Mostajabi, Jacob Steinhardt, and Dawn Song. Scaling out-of-distribution detection for real- world settings, 2022.
[31] Dan Hendryeks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. 2021.
[32] Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of- distribution examples in neural networks, 2018
[33] Dan Hendryeks, Xiaoyuan Liu, Eric Wallace, Adam Dziedzic, Rishabh Krishnan, and Dawn Song. Pretrained transformers improve out-of-distributio robustness. arXiv preprins arXiv:2004.061000 2020.
[34] Dan Hendrycks, Mantas Mazeika, and Thomas Dietterich. Deep anomaly detection with outlier exposure. arXiv preprint arXiv:1112.04006, 2018.
[35] James Hensman, Nicolo Fusi, and Neil Lawrence. Gaussian processes for big data. arXiv preprint arXiv:1309.6835, 2013.
[36] José Miguel Hernández--obatt and Ryan Adams. Probabilistic backpropagation for scalable learning bayesian neural networks. In International conference on machine learning, pages 1861-1869. PMLR, 2015.
[37] José Miguel Hernández- Lobato, Matthew W Hoffman, and Zoubin Ghahramani. Predictive entropy search for efficient global optimization of black-box functions. Advances in neural information processing systems, 27, 2014.
[38] Katerina Seeriikkvv--Shhnnlle Milan Paluš, Martin Vejmelka, and Joydeep Bhattacharya. Causality detection based on information-thcorute approaches in time series analysis. Physics Reports, 441(1):1-46, 2007.
[39] Bairu Hou, Yujian Liu, Kaizhi Qian, Jacob Andreas, Shiyu Chang, and Yang Zhang. De- composing uncertainty for large language models through input clarification ensembling, 2024.
[40] Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning. arXiv preprint 112iv:1112.5745,, 2011.
[41] Hong Jun Jeon, Jason D Lee, Qi Lei, and Benjamin Van Roy. An information-tteoreete analysis of in-context learning. arXiv preprint arX.5:2401.15330 2024.
13

---

[42] Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William Cohen, and Xinghua Lu. PubMedQA: A dataset for biomedical research question answering. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan, editors, Proceedings of the 2019 Conference Empirical Methods Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 2567-2577, Hong Kong, China, November 2019. Association for Computational Linguistics.
[43] Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds Nova DasSarma, Eli Tran-Johnson, et al. Language models (mostly) know what they know. arXiv preprint arXiv:2207.05221, 2022.
[44] Adam Tauman Kalai, Ofir Nachum, Santosh S. Vempala, and Edwin Zhang. Why language models hallucinate, 2025.
[45] Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision?, 2017.
[46] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics, 2018.
[47] Polina Kirichenko, Mark Ibrahim, Kamalika Chaudhuri, and Samuel J. Bell. Abstentionbench: Reasoning Ilms fail unanswerable questions, 2025.
[48] Jannik Kossen, Jiatong Han, Muhammed Razzak, Lisa Schut, Shreshth Malik, and Yarin Gal. Semantic entropy probes: Robust and cheap hallucination detection in llms. arXiv preprint 15X2: 2024.
[49] Akshay Krishnamurthy, Keegan Harris, Dylan J. Foster, Cyril Zhang, and Aleksandrs Slivkins. Can large language models explore in-context?, 2024.
[50] Lorenz Kuhn, Yarin Gal, and Sebastian Farquhar. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. arXiv preprint arXiv:2302.09664 2023.
[51] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[52] Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
[53] Hyungi Lee, Eunggu Yun, Giung Nam, Edwin Fong, and Juho Lee. Martingale posterior neural processes. arXiv preprint arXiv:2304.004. 09433, 2023.
[54] Te-Won Lee. Independent component analysis. In Independent component analysis: Theory and applications, pages 27-66. Springer, 1998.
[55] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retreetal-augmentee generation for Knowledge-intensinn nlp tasks, 2021.
[56] Yingzhen Li and Yarin Gal. Dropout inference in bayesian neural networks with alpha- divergences. In International conference on machine learning, pages 2062. PMLR, 2017.
[57] Yingzhen Li, José Miguel Herández--obatee and Richard E Turner. Stochastic expectation propagation. Advances in neural information processing systems, 28, 2015.
[58] Stephanie Lin, Jacob Hilton, and Owain Evans. Teaching models to express their uncertainty in words. arXiv preprint arXiv:2205. /4334, 2022.
[59] Zhen Lin, Shubhendu Trivedi, and Jimeng Sun. Generating with confidence: Uncertainty quantification for black-box large language models, 2024. URL https:/arxxiv org/abs/2305, 19187, 2024.
14

---

[60] Chen Ling, Xujiang Zhao, Xuchao Zhang, Wei Cheng, Yanchi Liu, Yiyou Sun, Mika Oishi, Takao Osaki, Katsushi Matsuda, Jie Ji, et al. Uncertainty quantification for in-context learning of large language models. arXiv preprint arXiv: 10189, 2024.
[61] Shang Liu, Zhongze Cai, Guanting Chen, and Xiaocheng Li. Towards better understanding of in-context learning ability from in-context uncertainty quantification. arXiv preprint arXiv:2405.1511 2024.
[62] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2104.08770 2021.
[63] David JC MacKay. Information- based objective functions for active data selection. Neural computation, 4(4):590-66) 604, 1992.
[64] Nishanth Madhusudhan, Sathwik Tejaswi Madhusudhan, Vikas Yadav, and Masoud Hashemi. Do Ilms know when not answer? investigating abstention abilities of large language models, 2024.
[65] Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. Ad- vances in neural information processing systems, 31, 2018.
[66] Andrey Malinin and Mark Gales. Uncertainty estimation in autoregressive structured prediction. arXiv preprint arXiv:2002.00 200500000,,06550, 2020.
[67] Andrey Malinin, Liudmila Prokhorenkova, and Alcksei Ustimenko. Uncertainty in gradient boosting via ensembles, 2021.
[68] Sabrina Mielke, Arthur Szlam, Emily Dinan, and Y-Lan Boureau, Reducing conversational agents' overconfidence through linguistic calibration. Transactions of the Association for Computational Linguistics, 10:857-877 072, 2022.
[69] Yifei Ming. Ziyang Cai, Jiuxiang Gu, Yiyou Sun, Wei Li, and Yixuan Li. Delving into out-of- distribution detection with vision-language representations. Advances in neural information processing systems, 35:35087-35102. 2022.
[70] Giovanni Monea, Antoine Bosselut, Kianté Brantley, and Yoav Artzi. Lims are in-context reinforcement learners. arXiv preprint arX1v:2410.05362, 2024.
[71] Hussein Mozannar, Arvind Satyanarayan, and David Sontag. Teaching humans when to defer to classifier via exemplars. In Proceedings of the aaai conference artificial intelligence, volume 36, pages 5323-5331. 2022.
[72] Hussein Mozannar and David Sontag. Consistent estimators for learning to defer to an expert. In International conference on machine learning, pages 7076-7087. PMLR, 2020.
[73] Jishnu Mukhoti, Viveka Kulharia, Amartya Sanyal, Stuart Golodetz, Philip Torr, and Puneet Dokania. Calibrating deep neural networks using focal loss. Advances in neural information processing systems, 33:15288-15299, 2020.
[74] Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. arXiv preprint arX 2021.
[75] Humza Naveed, Asad Ullah Khan, Shi Qiu, Muhammad Saqib, Saced Anwar, Muhammad Usman, Naveed Akhtar, Nick Barnes, and Ajmal Mian. A comprehensive overview of large language models, 2024.
[76] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science Business Media, 2012.
[77] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint wxXiv:2207.04179 2022.
15

---

[78] Allen Nie, Yi Su, Bo Chang, Jonathan N Lee, Ed H Chi, Quoe V Le, and Minmin Chen. Evolve: Evaluating and optimizing llms for exploration. arXiv preprint arX1v:2410.06238. 2024.
[79] Emre Onal, Klemens Flöge, Emma Caldwell, Arsen Sheverdin, and Vincent Fortuin. Gaussian stochastic weight averaging for bayesian low-rank adaptation of large language models. arXiv preprini arXiv: 03425, 2024.
[80] Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1007.038700000 187.00387 2018.
[81] lan Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
[82] Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
[83] Stephanie E Palmer, Olivier Marre, Michael J Berry, and William Bialek. Predictive informa- tion in sensory population. Proceedings of the National Academy of Sciences, 112((2):6908- 6913, 2015.
[84] Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism. arXiv preprint arXiv:2306.0444 04891, 2023.
[85] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg. J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay, Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 0-2222-2830, 2011.
[86] John Platt et al. Probabilistic outputs for support vector machines and comparisons to regular- ized likelihood methods. Advances in large margin classifiers, 10(3):61-74, 1999.
[87] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International conference machine learning, pages 5171- 5180. PMLR, 2019.
[88] Qwen, An Yang, Baosong Yang, Beichen Zhang. Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025.
[89] Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddie Bickford Smith. Modern bayesian experimental design. Statistical Science, 39(1):100-114, 2024.
[90] Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, pages 63-71. Springer, 2003.
[91] Jie Ren, Jiaming Luo, Yao Zhao, Kundan Krishna, Mohammad Saleh, Balaji Lakshmi- narayanan, and Peter J Liu. Out-of-distributiin detection and selective generation for condi- tional language models. arXiv preprini arXiv:2209.1553 2022.
[92] James Requeima, John Bronskill, Dami Choi, Richard Turner, and David K Duvenaud. Llm processes: Numerical predictive distributions conditioned on natural language. Advances in Neural Information Processing Systems, 37: 17:109699-109611, 2024.
[93] Karthik Abinav Sankararaman, Sinong Wang, and Han Fang. Bayesformer: Transformer with uncertainty estimation. arXiv preprint S.Kiv:22060082222 2022.
[94] Remo Sasso, Michelangelo Conserva, and Paulo Rauber. Posterior sampling for deep rein- forcement learning. In International Conference on Machine Learning, pages 30042-30061. PMLR, 2023.
16

---

[95] Glenn Shafer and Vladimir Vovk. A tutorial on conformal prediction. Journal of Machine Learning Research, 9(3), 2008
[96] Ola Shorinwa, Zhiting Mei, Justin Lidard, Allen Z Ren, and Anirudha Majumdar. A survey on uncertainty quantification of large language models: Taxonomy, open research challenges. and future directions. ACM Computing Surveys, 2025.
[97] Jacob Si, Wendy Yusi Cheng, Michael Cooper, and Rahul G. Krishnan. Interpretabnet: Distilling predictive signals from tabular data by salient feature interpretation, 2024.
[98] Freddie Bickford Smith, Jannik Kossen, Eleanor Trollope, Mark van der Wilk, Adam Fos- ter, and Tom Rainforth. Rethinking aleatoric and epistemic uncertainty. arXiv preprint 1rXiv:2412.20892 2024.
[99] Jasper Snoek, Hugo Larochelle, and Ryan P. Adams. Practical bayesian optimization of machine learning algorithms, 2012.
[100] Niranjan Srinivas, Andreas Krause, Sham Kakade, and Matthias Seeger. Gaussian process optimization in the bandit setting: no regret and experimental design. In Proceedings of the 27th International Conference International Conference on Machine Learning, pages 1015-1022, 2010.
[101] Shuchang Tao, Liuyi Yao, Hanxing Ding, Yuexiang Xie, Qi Cao, Fei Sun, Jinyang Gao, Huawei Shen, and Bolin Ding. When to trust llms: Aligning confidence with response quality. arXiv preprint arXiv:2404. 17287, 2024.
[102] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
[103] Artem Vazhentsev, Gleb Kuzmin, Akim Tsvigun, Alexander Panchenko, Maxim Panov, Mikhail Burtsev, and Artem Shelmanov. Hybrid uncertainty quantification for selective text classification in ambiguous tasks. Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of the Association for Computa- tional Linguisticy (Volume 1: Long Papers), pages 11659-11681, Toronto, Canada, July 2023. Association for Computational Linguistics.
[104] Janett Walters-Wiile Williams and Yan Li. Estimation of mutual information: A survey. In Rough Sets and Knowledge Technology: 4th International Conference, RSKT 2009, Gold Coust, Australia, July 14-16, 2009. Proceedings 4, pages 389-396. Springer, 2009.
[105] Bingbing Wen, Bill Howe, and Lucy Lu Wang. Characterizing Ilm abstention behavior in science qa with context perturbations. arXiv preprint urXiv:2404. /2452, 2024.
[106] Bingbing Wen, Chenjun Xu, Robert Wolfe, Lucy Lu Wang. Bill Howe, et al. Mitigating overconfidence in large language models: behavioral lens on confidence estimation and calibration. In NewrIPS 2024 Workshop Behavioral Machine Learning, 2024.
[107] Rand R Wilcox. Introduction to robust estimation and hypothesis testing. Academic press, 2011.
[108] Lisa Wimmer, Yusuf Sale, Paul Hofman, Bernd Bischl, and Eyke Hüllermeier. Quantifying aleatoric and epistemic uncertainty in machine learning: Are conditional entropy and mutual information appropriate measures? In Uncertainty artificial intelligence, pages 2282-2292. PMLR, 2023.
[109] Yue Wu, Xuan Tang, Tom M Mitchell, and Yuanzhi Li. Smartplay: benchmark for Ilms as intelligent agents. arXiv preprint arXiv:2310.01533 2023.
[110] Yijun Xiao and William Yang Wang. On hallucination and predictive uncertainty in conditional language generation. arXiv preprint
2021.
17

---

[111] Johnathan Xie, Annie S Chen, Yoonho Lee, Eric Mitchell, and Chelsea Finn. Calibrating language models with adaptive temperature scaling. arXiv preprint wrXiv:2409.198/77 2024.
[112] Qiujie Xie, Qingqiu Li, Zhuohao Yu, Yuejie Zhang, Yue Zhang, and Linyi Yang. An empirical analysis of uncertainty in large language model evaluations. arXiv preprint arXiv:2502.10709 2025.
[113] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference, 2022.
[114] Keyang Xu, Tongzheng Ren, Shikun Zhang, Yihao Feng, and Caiming Xiong. Unsupervised out-of-domain detection via pre-trained transformers. arXiv preprint arXiv:2106.000 00948, 2021.
[115] Adam X. Yang. Maxime Robeyns, Xi Wang, and Laurence Aitchison. Bayesian low-rank adaptation for large language models, 2024.
[116] Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhut- dinov, and Christopher D. Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. 2018.
[117] Fanghua Ye, Mingming Yang, Jianhui Pang, Longyue Wang, Derek Wong, Emine Yilmaz, Shuming Shi, and Zhaopeng Tu. Benchmarking llms via uncertainty quantification. Advances in Neural Information Processing Systems, 37:1 37:15336-15385 2024.
[118] Naimeng Ye, Hanming Yang, Andrew Siah, and Hongseok Namkoong. Exchangeable sequence models can naturally quantify uncertainty over latent concepts. arXiv preprint arXiv:2408.08 03307, 2024.
[119] Liyi Zhang. R Thomas McCoy, Theodore R Sumers, Jian Linn-Qia Zhu, and Thomas L Griffiths. Deep de finetti: Recovering topic distributions from large language models. arXiv preprint wrXi::2312.14226, 2023.
[120] Siyan Zhao, Tung Nguyen, and Aditya Grover. Probing the decision boundaries of in-context learning in large language models. Advances in Neural Information Processing Systems, 2024.
[121] Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International conference on machine learning, pages 12697-12706. PMLR, 2021.
[122] Wenxuan Zhou, Fangyu Liu, and Muhao Chen. Contrastive out-of-distribution detection for pretrained transformers. arXiv preprint arXiv: 100.211400000 008/12 2021.
18

---

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: The abstract and introduction state the claims made and are justified thoroughly with proofs and experiments throughout the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of method in the conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section 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 on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not 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 to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
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]
19

---

Justification:
Complete proofs of our theory are included and discussed in the Appendix.
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 or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Detailed experimental setup, methodologies, and chosen parameters are out- lined in the main text and the appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, 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 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) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be 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 deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to 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?
20

---

Answer: [Yes]
Justification:
The code is anonymized and zipped along with our submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.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: //iiss wwtppcccpppppaeccccccccccccc /ttegguudddeeeeeeebmiiieeeee /Cto//eooommseeeeeeeeeea aassonoliccy for more details.
The authors should provide instructions on data access and preparation, including how to 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 as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
The above is specified in the main text with further details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Standard deviations included throughout downstream experiments real world examples.
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, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
21

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Implementation details regarding compute resources are included in the
appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: : http//euurpp...eee ct.ppp..........oo..o.e cccpublic//tteeecceeceeceeeeeeee wwwpblli//thissuuudddeoo ceeeiiiisGeiieeeeeeee.
Answer: [Yes]
Justification:
The authors have reviewed and comply with the NeurIPS Code of Ethics
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 societal impacts of the work performed?
Answer: [Yes]
Justification: broader impact statement is provided in the conclusion.
Guidelines:
The answer NA means that there 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.
22

---

The conference expects that many papers will be foundational research and not tied to 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, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper demonstrates foundational research tested on publicly available datasets that were designed to test machine learning algorithms.
Guidelines:
The answer NA means that the paper poses 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 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:
We properly acknowledge and cite all assets and resources used in the paper.
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 is 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 package should be provided. For popular datasets, paperswi thcode. com/ 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.
23

---

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: The paper does not release new assets.
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: The 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:
The 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

---

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: [Yes]
Justification: Exploring uncertainty decomposition in LLMs via in-context learning is core part of our methodology. We have delineated all details in our paper.
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 (https: : http///urpp..cc/ofeeeeeeeeeoeeee for what should or should be described.
cc htt.:////o/ereeecc///////////
25

---

Appendix
Contents
Proofs
27
A.1 Variational Uncertainty Decomposition
27
A.2 Variational Estimates of Variance Decomposition
28
Theoretical Examples
29
B.1 Bayesian Linear Regression
29
B.2 Gaussian Process Regression
30
B.3 Gaussian Bandits
31
Sampling Methods for Auxiliary Data
31
C.1 Methods
31
C.2 Ablations on Logistic Regression Data
32
Promoting Exchangeability In-Context Learning
32
D.I Definition of Exchangeability
32
D.2 Bridging the Gap
................................
33
D.3 Exchangcability from Predictive Rules
................................
34
D.4 Determining Threshold for KL-Filtering..
35
E Algorithms and Pseudocode
36
E.1 Pseudocode Variational Uncertainty Decomposition Algorithm Ageith..........................
36
E,2 Computing Approximate Posterior Predictive Distributions
37
E.3 Profiling View of VUD Algorithm...
37
Further Related Work
38
Experiments
41
G.I Code Implementation
G.2 Further Baselines
G.3 Synthetic Toy Experiments
47
G.4 Bandits..
51
G.5 Question Answering
57
H Example Prompts
59
ILI Synthetic Toy
59
11. Bandits
59
H.3 Question Answering
61
1 Declarations
62
26

---

A Proofs
A.1 Variational Uncertainty Decomposition
Theorem 3.1 (Alcatoric Uncertainty Upper-Boundy If the conditional independence relations in G hold, then the variational estimator provides upper-bound to the aleatoric uncertainty:
Va(y"(x" x², D) 2 Ua(y"(a"
(6)
where the gap between Ua (y" x² D) and V. Vy x=, D) is:
Ey 010", U, Z, DJ = , - Z.D) [Dulpp(Oy", x", U, Z,
Ep Ul [Dr, x=, ()) U, D)]]-
(7)
Proof. We begin by decomposing the variational estimator Va. noting that from g get, p(y" x² ,0) = p(y"(="" 0, U, D) and p(0 U, 2, D) p(0, U, D):
Va x=" Z,D) (p(uj U.z. ply" x D)]
log ply" x=" (,
PPp(UjZ 0)( U.2.D)
(+)
p(A x", U.
-EprU,, on 0p0 U.Z.D) [log(yy x [0]
log p(0ly" x", U, U,Z,D) D)
+ +Epp .U.Z, x".,"
(**)
(
Epp D) [Hyp ,. 0)] (***)
+ Epyy U| 2.D) [KL (00l"|, x", U, 2, D) ()
EIp D) [Hip(y" x 0)] Uu DD).
Here steps (+) and (**) are obtained via Bayes' rule and the conditional independence assumption y 1 U10, x", Z. D DAG G. Step (***) is due to the assumption of the likelihood model p(y (z, 0)
(and hence p (UZ 0)) which do NOT treat æ (and hence Z) as random variable:
p D. P(U 0)p(0 D) Z D)
/ (, (U D)JU = PUUZ ((DD)) = p(OID).
Note that by definition of mutual information, we have:
Ep(U//,)) (4l"" 00x", U, 2, D]
Eply UU 2,D) (PLLlp(0)) x², U, Z, D) p U, Z, D)]
Ep(0, ZDD) [DKl[p(y* x=' 8, U, Z, D) |yy |=* U, Z, D)]]
Ep(o,u Z.D) [DKL[p(y" x=" 0)) o=" U,Z, (***$)
where, again, step (****) is due to the conditional independence structure y* 0 x" Z, of DAG G.
Alternative Proof. Firstly, it is useful to define the corresponding definition of the variational approx- imation to the epistemic uncertainty as:
Va x", D) : I[yy;; U| Z,D)
HH p(y' |x", U. - Z,
Hi[pl""l" Z, D)] - 1=","",",,,
H[ply" DD] Vv x²" Z.D))
(*)
where (+) follows from the conditional independence assumption y 1. Z|e, D. Therefore, we have "2", Z, D) - U_ D) = Ua(y" x=",, D) Va(y"l²", Z,D)
(**)
If have the conditional independence relation y* 1. U æ. Z, D, then by the data processing incquality (DPE):
DPE
S I(y" 0|x=" Z,D) (t) I(y" 0|x"" D) =: U. |x" D).
Va(y" xx Z,D) := I(y"; U|=" Z.D)
where (f) follows from the conditional independence relation (y", 00) Z|e, D.
27

---

Remark. From this information- theoretic perspective, we see that choosing optimal Z. is equivalent maximising the mutual information between y and U. This further motivates choosing Z that repeats x² or are perturbations of x""
A.2 Variational Estimates of Variance Decomposition
To prove Theorem 3.2, we first prove the following lemma.
Lemma A.1 any random variables X, YY Z where the conditional variances Var(Y]X) and Var(Y|XX Z) exist,
E (ar((()) = Var(B 2] ([Var(Y|)x2)) E[Var(Y|X, 2)]
Proof.
By the law of total E = 2)] B| Therefore,
huap://-~oouuu1y223-----00000111 E|E(Y2 XX E
2) E(Y|X, 2) -E E Z)
E EE X
(Y 2) (|E(Y)X)
To show that the LHS is positive we first decompose E( X, Z) as
B(Y|X. Z) (E(Y)X, 2) E(Y x)) + BY
Now, the expectation of the product of these terms is 0 as
E (B Z) E({ B(Y = E E (E(Y)X, 2) - E(Y E(Y B(X|X)
[[[[[E(V 2) -E [X
E
- B(VI)) -((()))))(((XX) [((IX)
= LE [ [())))
0,
where (+) follows from the fact that (X) C ar (X, Z). Therefore,
(E X, 2) F. (( 2) B(YIX) + =((I)))
E 2 + +2 2) B((])) +E[B(Y|X)} [X)]
=0
= E Var + (E XX)]
Finally, this gives
[[Varr(1x)]-B/Vaal (Y| 2 EY -XY)) E (E 2)|x) 2 0,
where the final inequality follows from the non-negativity of variance.
Theorem 3.2 (Aleatoric Variance Upper-Bound), If the conditional independence relation holds, then the variational estimator provides an upper-bound to the estimation of aleatoric variance:
v <y" Z. D) = Ep(ujzz [Varly U.Z.D] "y" .D).
(9)
28

---

Proof. By the definition of V.
(y , D) Ep(ZI Varly x= U, D]]
Ep(U)/* 2.D) Var y æ U.
Ep 2.D)
| |arly 1 n² U, 2, DJ]
(*)
2, E. (arly"/
(**)
Expelee 2.D) [Varly"| |
Epper)) [Var(y |a= .0] (y"|"=", D)).
Here, (+) follows from Lemma A.I and (+*) follows from the conditional independence relation
y Z. U,D|æ* O.
Remark. From Lemma A.I, we also obtain that the discrepancy between Va'(yy' (x" Z, D) and (y"la"" 1x", is: is:
I
,U, U,2, x², ,2, D
B)(U/Z,)) Varg:o1u...)) (Bly x"] J, Z, D) L.D]
B Theoretical Examples
B.1 Bayesian Linear Regression
Consider a linear regression model with homogeneous output noise variance. Namely, we assume normal prior p(0) (0; 0, A >LL),, and the likelihood model is p(y|ær, (a, 0) == Nyy; 0 o²) Denote X 121,- .... 2n. T E Raxd and Z ,i,.... Zmm] E Rmxd, Now consider the exact
posterior predictive distributions which can shown as:
p(0 (0: u4, A,,, A: o X ALa, -==A--x"y
p(y"I=*, = =NN") T x=, (x*) +2)).
Then using the closed-form expressions for the entropy of a Gaussian distribution, it is straightforward show that for arbitrary y" x= and D:
Ua(y" |x=" D) (1 + log 2
1
Uely"Iæ", D) log((x*) "Taa""""A-0 -A=xx" + log o²,
Adding the auxiliary data U:
p0 Z, D) No0 (Z), A(Z) - - X Z Z) + Ala, p(y* In. U, D) Nyy;; :(() x- (x=) (Z)=* x= + o Laa
Since the variance of p(y"|=== U, D) does not depend y* and U, this leads to
Va D) ! log 2n) + 1 log((=* (Z)xx +
1
Ve D) - log( (x') ) ----0 x0" +o²) log((x)) TA- (Z)x* + tz),
(x=
is easy to show for all possible Z:
Va " Z) Ua x D) D) 2
log(o - - AZ))a +1) 2 0.
Now consider the optimum of the variational estimate:
Va(y*Iæ* |0"" D) 1 (1 + log 2TO) min log(o - 2xx" AT -1 AZ)
x0 +11
2
2
29

---

where A(Z) o X 2) ALg. Now, if is the minimum eigenvalue of (X"X + Z Z) and y > 0, then (a:)
- x- $ 1110==11 m 2 d, we can choose (e.g. unit vectors) such that
O, and then scaling 2, by a constant ensures y + oc and (x*) - 'x" O. Therefore, for appropriately chosen Z, Va Z.D) + Unly" x", D).
B.2 Gaussian Process Regression
Here assume Gaussian process model [90] with kernel function as the prior covariance:
y = f(x) oc, C~~(0, 1), f()) GP(O, k(---))
Here we assume ID outputs w.l.o.g. and use notations y, y interchangeably. For regression problems have closed form solution the posterior predictive (with D (X, Y). omit the formulation of the posterior mean (uXX, Y) and focus the discussion on the posterior variance only):
ply 1==" X, Y) Nyy" : ,(= x²) Xxxxxxx + o²I) -'II-'Xxx + o²),
leading to the following uncertainty estimates:
1 (1 + log 2
Ua(y"la" ( 2
Up(y |=", D) - D) log(A((=*) x²) K.x((xx + -Ix + +02)
log a²
Now consider sparse variational Gaussian process (SVGP) [35] with inducing inputs/outputs Z, and an approximating distribution q(u) == N(u m, S). Then we have the approximate posterior predictive as:
(3y") =
k(x", x²) K zz Kzz 12 + +0²)
so that the decomposed uncertainty estimates from an SVGP are
(a | :9) 2 (1 + log 2 ) = Ua (a'""xx" :9),
[z(y|<=":> - log(k(2>, x²) K,z 2z Kzz 5)KK2 + +0²) -
- log 0².
For regression problems we have the optimal S Kzz (zzz + o "zxxxzz) 'Kzz [35], and therefore
1 (y |I= :q) = log(k(#"; x²) K Kkz - (Kzz + - (xx) + 0²) log 0².
On the other hand, using the variational uncertainty decomposition method, we can show that
ply" x=, U,Z,X,y) N (Z, k(=",, .=) X.x(Xxx + 0²1) Kxx A(x", Z) + o2),
A Z) = T (Kzz + I²I K2x (Xxx + -02 (Kz) A = Kz+ Kzx (xxx + -1--KXx
leading to the following uncertainty estimates (with D = (X,y) and C= {1 + log 2Tr)):
Vy x=" Z,D) =C' + log(k(æ" xx) -xxx + o² I) Kx. A(x", .Z) + t²)),
Ve (x", , Z)D) =C' + log(k(z) [f(x".ca x²) XxxxXxx +
+ + ) Va Va Z.D).
Note that if we choose
then have
1
Va ,,0 D) Ua(y" | (D)
2 log
k(x",,"" K, (Kxx xx + + - + t
0²
Voly"la",
x.,D) ",""""""",,,"" x², (D) Uy(y" 12", ; D) log
2
k(xx, "(x) )- Kxx(Xxx + 02 Kx, + t
This means if the test input a: is close to the training data X, then k(x" x²) - (xx d²I)- will be close to zero, and then Vo |x" ,x",, ,.,D) as U. |a" D) provides a good estimate of the epistemic uncertainty.
30

---

B.3 Gaussian Bandits
In Gaussian bandit setting, suppose we have independence of rewards between arms. Furthermore, for an arm i, we assume the following conjugate Gaussian model:
1. Gaussian Prior: p(0,) N/(0, og)
2. Gaussian Likelihood: p(r: 0.) ((
Then for kt observations of rewards from arm i, we have:
1. Total variance: 0 o²
o +
k 0z
2. True aleatoric variance: U
3. True epistemic variance (total aleatoric): U
However, with the VUD method, suppose we have n further auxiliary observations (i.e., predicted reward values from the model) for arm i, then, the corresponding variational estimate of aleatoric uncertainty is:
1 00
t 0²
V = 0² +
2
U.S
This gives the gap between the variational estimate and the exact aleatoric variance as:
ke +
= -0
var - U
=
o2
K
C Sampling Methods for Auxiliary Data
In this section, we discuss in detail the methods used to sample auxiliary queries Z to find the best variational estimate of the aleatoric uncertainty and variance. As noted Section 3.1, we restrict Z to single query in the domain to reduce the search space.
C.1 Methods
Bayesian Optimisation. The optimisation problem (8) can be directly optimised via Bayesian Optimisation. However, this is constrained optimisation problem where Z needs to satisfy an "approximately Bayesian" criterion (1I) which we discuss in Section 3.2. To overcome this issue, we treat the problem an unconstrained Bayesian optimisation task to obtain auxiliary examples {zy) and then apply the criterion remove auxiliary examples that do not satisfy (11). In the synthetic examples we consider, the covariates are real and continuous. Therefore, we use Gaussian process with an RBF kernel to model the objective function and take the log expected improvement as the acquisition function. In order to provide warm start to the Bayesian optimisation process, we provide initial samples that are randomly sampled.
Perturb. Given the covariates x' of the data point we wish to decompose the uncertainty for, we can choose Z {z, 1 be "close" to x". To perturb a categorical covariate x [k).. sample uniformly from the list of categories with probability p and keep the original covariate with probability 1 1-p. For a real covariate x= [k'), we sample from a normal distribution, similarly to random sampling. but we choose the mean as z and the standard deviation as scaled population standard deviation estimate of the covariate of where = 0.1.
Repeated. Given the test covariates x", we set Z = x². Since repeat the covariates, we only evaluate auxiliary query per test example, and therefore the KL filtering procedure is omitted.
Random Sampling. The most basic sampling procedure to generate auxiliary queries Z is to randomly sample the input domain. uniformly from the list of categories. If covariate x: [k'] is real variable, we assume normal distribution with mean and standard deviation given by the population mean and standard deviation estimates of the covariate, 2b and of from the in-context data D.
covariate x: [k]
categorical variable, we sample
31

---

C.2 Ablations on Logistic Regression Data
compare the performance of the four approaches to choose outlined Section C.I for 15 auxiliary examples (with the exception of the Repeated where have single auxiliary example). We plot the uncertainty decompositions for the Z sampling approaches and the corresponding KL divergence for the chosen Z that minimises (8) in Figures 8, 11 and 12. In Tables 3 and 4 quantity the performance each of the sampling methods by computing the mean rank of each method over consistently observe that Repeated has the
the test samples. For the 3 LLMs that consider, lowest V., followed by Perturbations. However, Perturbations has the highest KL. divergence, which indicates that this method is less aligned with the Bayesian assumptions that make.
1.0
Uneeriaisty
10
Uhorsalisy
0..
Pertazb
Divergeree
Repested
10
0.6
Runther
g
10
80
0.4
ICL Dulu: g-1
Putarb Repasted
15
10
Test Covaciate
Test Covariate
(a) Effect of Z Uncertainty Estimates
(b) KT. Divergence of optimal selected Z
Figure 11: Va across Z sampling methods (Qwen2.5-7B)
1.0
Upeerinitt
Unrertalntt
0.8
Partatt
Divergaze
10
Repested
0.6
Bondoni
10
80
JCL Dalouge
Pertur
Bepaaied
Rankai
-15
-10
Test Covariate
Test Covariate z
(a) Effect of Z Uncertainty Estimates
(b) KL Divergence of optimal selected Z
Figure 12: Va across Z sampling methods (Llama-3 1ama-3.1. (..mmmm..-88)))
Table 3: Vo rank for different sampling methods
MODELS
BAYESIAN OPTIMISATION PERTURBATIONS
REPEATED TASK
RANDOM SAMPLING
QWEN2.5-7B
2.93
2.01
1.29
3.77
QWEN2.5-14B
3.09
2.03
1.16
3.68
LLAMA-3.1-80
2.41
1.92
2.09
3.57
D Promoting Exchangeability in In-Context Learning
In this section, we expand upon the concept and definitions of exchangeability that we discuss in Section and the methods encourage exchangeability in Section 3.2.
D.1 Definition of Exchangeability
finite sequence of random variables (X,) is exchangeable for any permutation p: [m] + [n]]
(X1 X.)) (Xp(), X
(12)
(where d refers to equal in distribution) [24]. Similarly, an infinite sequence (Xn)n21 is exchangeable if for any finite permutation p. (Xn).21 d (X (Xp(x))n2) One of the most consequential results for exchangeable sequences de Finetti's s representation theorem which is stated (in measure-theoretic form) as follows:
32

---

Table 4: KL divergence rank for different sampling methods
MODELS
BAYESIAN OPTIMISATION PERTURBATIONS
REPEATED TASK
RANDOM SAMPLING
QWEN2.5-7B
2.27
2.31
3.43
1.99
QWEN2.5-14P
1.98
2.27
3,51
3.51
LLAMA-3.1-8B
1.93
2.41
3.77
1.89
Theorem D.1 (de Finetti's representation theorem). Let (Xn)x21 be an infinitely exchangeable sequence and denote its probability law. Then, there exists a unique random distribution F with law n such that for all n. 2 1 and measurable sets A1 A1,...; Ans
P(X, E A,..... N,, An) / 1
F(A)=(dF).
(13)
Note that this standard thhoorie definition [24] differs from the one introduced in Section where we have a supervised learning setting of covariate-label pairs {( y:) The results that use later in Section D.3 to encourage exchangeability are also in measure-theoretic form and therefore, in the following section, we bridge the gap between the supervised learning setting and the measure- theoreee language.
D.2 Bridging the Gap
So that we can consider the pair yz) as the observations of random sequence taking values in X' x V we make the modelling assumption that (Xt))21 is an i.i.d. sequence of random variables with law Q and density q. Then, the random variable Yn+1 is generated from the LLM given Xn+1 and {(X,, Y)) Denoting random variables K, = (X,, Y,). with realisations k1n = {(T.. ya)). and measurable sets and B in the sigma algebras (X) and a respectively, the following proposition justifies the definition of exchangeability that we use in Section 2.
Proposition D.1 Let (Xn)x21 be an i.i.d. sequence and (Kn)n21 be an exchangeable sequence, where K,, = (X,, Y,). Then there exists a random distribution F with joint (random) density fx.y and law such that the conditional density of Y1.. 1Xinn can be expressed
ply1:n xi:n) - / II (xx(ss))))(()))
A=1
Proof. Since (Kn)a21 is exchangeable, by de Finetti's we have
P(K1 B1,. .I,,,..,K E (An, Bn))
II F(A,, (d)
P(X1 E A1, BI,...,,,, An, E BB)
II F(A, (dF=(dF) 1=1
P(X, E A,, E Bi,... X. An, Y, E Ba)
II
A. /
(dFF.
where fx,Y is the joint (random) density of the random distribution F with (random) marginal fx and (random) conditional distribution frix. Now, letting B, V. we obtain the marginal probabilities follows:
/ II
i=1 A. / fvix(vsla))dd fx( dxx T
P(X, E A An)
II P(X. E Ao) / II
fx (dx))
f=1 J.,
II g(z;)dæ; = / II /
fx (dF.
(1)
im A.
2=1 A.
33

---

where the second line follows from the independence of X, As (f) holds for any measurable set that fx % q is valid solution (t). But by the uniqueness of the law in de Finetti's 's representation theorem, can indeed conclude that fx a. q. Therefore, substituting fx into the de Finetti's representation, the conditional probability density of Yi. 1 can be expressed as:
A; E (X),
p(y Im (i:)) p( In y tmm) h/n//p(11ss) p(
/ x (2)(())))))))
= / II fr
g(zz)) ) )
II f (x
/(dF)/p(axm))
II g(xz)) /p(xmm)
m)(o
(yx(()(a)))((F))
D.3 Exchangeability from Predictive Rules
In our problem setting, we have the predictive rule for (Kn)n21
(AA B) I.n) =K(n+ E (A, B)] ......, ,K....,K, kn)
= =P(X,+1 E A, E
kn)
{ P(d(an+1,, Nn (n+1)) 1
(A.B)
P(dyn+1 kim, In+1
[ P(dy kims an+ (ddxn++1)
where the final equality follows from the independence of Xn+1 and probability P(yn+1lk) .k.... kns 1n-1) given by the LLM. The following theorem by Fortini, Ladelli and Regazzini [23, 24] gives necessary and sufficient conditions for exchangeable sequence defined by predictive rule.
Theorem D.2 (Theorem 2.3 [24], Theorem 3.1 and Proposition 3.2 [23]). Let (Kn)n21 IP be
infinite sequence of random variables with predictive rule (Pa)n20- Then (Ka)221 is exchangeable and only if. for every 2 0, the following conditions hold:
i) For every C E (K), (C is a semmetric of Zı...., Fni
ii) The set function (C. D) + Sc Pati(Dikisss dd (k (k:a)) is symmetric in C and D, where C E oK)).
Permutation Ensembling In predictive rule, approximately satisfy i) via the Monte Carlo approximation (10) as is equivalent to ensuring P(dy.+1 K1 In+1) is symmetric in k
KL-Filtering Condition ii) essentially requires that
P(Kn+1 E C, Kn42 D kim) P(Kn+1 E D, D,K22 E C)kmm)).
34

---

The set function in ii) can be expressed as follows in terms of (2nt1:
An+1,Ba+1 Prtz((Ant2: Bn-2) (kan))
P(dyn+2 kmm
kim. (nn) (dd
Ant1 But1 Aut2
But2
P(dy kkm (1 Yn+ (dd a 1
JAi1 An12 Ba11
B212
This expression computationally infeasible check but necessary condition for this
P P(Kn+1 E 11mmmm) = P(Kn+2 E k1n))
where we take D = (X, This is equivalent to
Ja P(d(zn+1, Nn+11
/ ((d(En+2.Fn+2)))
P(dy ki n 1n+1) S (P ki Fn
(tt)
sufficient condition for (tt) is the equality of the laws P(Yn+1 E -(kim Xn+1 = x) = P(Yn+2 E ii Xn+2 2)), thus motivating the KL filtering condition (11).
Effect of Permutation. In Figure 13, we plot the KL divergence from p(y"l=", D) to p(y" x U, Z' D) (where Zi argminz V (y x Z, D)) when we permute and do not permute the in-context labels. We see that permuting the in-context labels results in lower KI. divergences, which suggests the behaviour is more Bayesian.
(a) Quen2.5-14B
Qwen2. 5-7B
Ll 11ama-3. 3.1.aa3.1185
Figure 13: Permutation Ablation for Logistic Regression Dataset.
D.4 Determining Threshold for KL-Filtering
The choice of e controls the level of approximation permitted in the uncertainty decomposition method. A small ensures that the auxiliary data Z that choose obey Bayesian assumption but at the cost of rejecting more Z and obtaining a larger variational upper bound to the aleatoric uncertainty or variance. Furthermore, as shown in Figure 8, the range of KL values for the different auxiliary examples may vary when we vary x= Therefore, to guarantee that we have enough valid auxiliary examples, we set e as the th smallest element the set of KL divergences { where ej == Dxulp(y" [x=', D), ply" |x* Zj, D)). Therefore, we can control the strictness of the filtering by varying r, where smaller 7 gives a stricter decomposition.
35

---

E Algorithms and Pseudocode
E.1 Pseudocode for Variational Uncertainty Decomposition Algorithm
Algorithm 1 is pseudocode for multi-class classification problems and Algorithm 2 is the pseudocode for regression. They similar in approach but vary during the marginalisation step: for classifica- tion, we can compute p(y'|<", u k, 23, D) for each class k, and directly compute the marginal distribution using the tower property. However, for regression, this computationally infeasible so we use a Monte Carlo estimate for the conditional entropy Ep(ulz, [Hi[p(y" x²", u. 2,, D)llI, over different samples of u. To obtain the marginal distribution, bootstrap samples from the mixture of Gaussians [ply" a="" D. [Zj, uc Algorithm 5).
()) D) and fit a Gaussian to these samples (as described in
Algorithm 1 Multi- Classification for Aleatoric Uncertainty Estimation
Require: Test input x"; ICL Dataset D (0t, where yı E [K] I: p(y" x=" D) t- CLASSDIST(x',,
2: Honal t- HH[p(y" xx", [D)]
3: for j -I...., do
4: 2j NEWAUX(#", 211:j-1) (
{Get new auxiliary variable)
5: p(u) æ;, D) + CLASSDIST( Zj. D)
6:
for k = 1..... do
7:
p(y"la' la", DU k}) t CLASSDIST(E", DU {zy, k})
8:
HAT + -Hpppy" [x"" DU (z). end for
9:
10:
p(y" xx",, Zj D) + Σ p(y" x DU k}) plu k|2,, D)
11:
H3 Ek= k= Hęt pvu = k D)
12: 6 "kul]p""|<="""
ply" x ž5 D)
13: end for
14: Compute threshold e (see Appendix D.4)
Va f- min (min((() (H, : Cgg <<)), Hooaa)
16: return Va
Algorithm 2 Regression for Aleatoric Uncertainty Estimation
Require: Test input x"; ICL Dataset D y:) where y; E R
1: py x=", D) + REGDIST(æ",
2: Hotl [H , D)]
3: for j I..... m
4: 2j < NEWAUX(#* ((1))))
(Get new auxiliary variable)
5: (6) + (u,
IT J=-1 where u, REGDIST(Z,, D)
6:
fort= 1,..,, T do
PNry""a"" (a"" D (2), u (1)) - REGDIST(#", DU D2f2) u.) (1)))
7:
Hi + H(p(y"|æ* 12" DU (2) "
8:
9:
end for
(N xx', D). + NORMAPPROX (x" DU {zj, ut (1) 1))
H2 E, Hge
12: + *+-Dulpx(y"as
py 22,,
13: end for
14: Compute threshold (see Appendix D.4)
15: V, +- min (mina(()) ej < {)), Homa)
16:
return Va
Note that these algorithms can also be extended to the decomposition of total variance by replacing the entropic uncertainty terms with the corresponding variance terms.
36

---

E.2 Computing Approximate Posterior Predictive Distributions
Classification. Algorithm 3 describes the process of obtaining the logits for a predictive task p(y") x""","",>"""" D) given in-context learning data D Y:) and the covariates of the predictive task x². We permute the ICL data and take an average of the predictive distribution to obtain a Monte Carlo estimate of conditional permutation -nvvaitiiionnvaiannn distribution (which we discuss further in Appendix D. Furthermore, by the construction of the prompt, the we only need to obtain the logits for the first token that is generated, which remains constant with respect to the choice of LLM seed.
Algorithm 3 Compute Permutation Invariant Classification Distribution z CLASSDIST
Require: Test input ==; ICL Dataset (Ri, y:) where y; E [K]
1: function CLASSDIST(2")
2: for : do
3: ~ SK
(1)
4: Py
"oo(k)..,,. Pu Vo( (Class prob. of next token)
5: end for
6: Py f E2
7: return Ps
Regression. In Algorithm 4, outline the procedure for constructing approximate distribution for p(y" |x",, , D). Similarly to the classification case, we permute the ICI. data. However, as y" can take any value in I. the tokenisation y may require more than one token and as the logits of token depend on the previous tokens generated, the logits of the tokens will vary with the choice of LLM seed. Standard approaches to approximate the distribution require forward pass over every value that y" takes [92] which is prohibitively expensive. Therefore, for each permutation, we sample single y" (varying the LLM seed for every permutation) and fit a normal distribution these samples via moment matching (namely, estimating the mean and standard deviation of the sample and using these estimates as the parameters of normal distribution).
Variance Reduction. To reduce the variance of the estimated mean and standard deviation, we use trimmed mean, removing the top ie and bottom & of samples, and the interquartile range estimate the mean and standard deviation respectively [107]. In our experiments, we set k = 1.
Marginalisation. In Algorithm 2, we are required to compute the marginal distribution (y [x= ;,, D) given the Gaussian distributions {p(y"|æ=" , U {2,, un 1))]- We compute this marginal distribution by bootstrap sampling from the distributions p(y |ær", D (z), (z,, ) and fitting a Gaussian distribution to the bootstrap samples via moment matching. This procedure is outlined in Algorithm 5.
Algorithm 4 Approximate Permutation Invariant Regression Distribution: REGDIST.
Require: Test input x": ICL Dataset D = (x) where y; E IR 1: function REGDIST(wr, D)
2: for I 1,...., do
3: σi ~ SK
4: (())
(Sample next prediction)
5: end for
6: Y Y -Ty
(Trimming optional)
7: return Normal(mean) (Y), stds))(Y))
E.3 Profiling View of VUD Algorithm
The memory requirements and speed of applying VUD particular prediction task depends on the specific choice of L.MMs and the hardware on which it is deployed. Therefore, for clearer outline of computational costs, outline the number of API calls for each step of the algorithm as this is the
37

---

Algorithm 5 Approximate Marginalisation of Mixture Distributions: NORMAPPROX.
Require: Distributions (pe(y))
1: function NORMAPPROX ((pr(y))<=))
2: for -..... do
3: ~
(Uniform discrete distribution from 1 to T'
(r)
4: YB (y)
{Sample next prediction)
5: end for
6: Ys + (y
R=1
7: return Normal (annaan(Y)),tddd
most significant bottleneck in the algorithm. In Table and 6, provide profiling view for VUD in classification task and a regression task respectively.
Table 5: Profiling View for Classification Task with & classes using n auxiliary data, and L permuta- tions per distribution.
Function
LLM Calls Per Function
Num Function Calls
ply" x D) ( Zi D) ply x=, U ==,, D)
L L
1
π2
n
Table 6: Profiling View for Regression Task with n auxiliary data, and samples to evaluate each distribution with m discarded samples (due to trimming of mean in Algorithm 4).
Function
LLM Calls Per Function
Num Function Calls
p(y'l=" .))
L + m
1
P ZU Z, D) ply" |r"", U- -.,u70 D)
L +
n2
I + m
y. T.
F Further Related Work
Bayesian Interpretations of In-Context Learning. Works recent years [113, 84, 74] suggested that the behaviour of transformers during in-context learning emulates Bayesian inference. In our work, this Bayesian behaviour of ICL. is a key assumption that is necessary for the validity of the variational uncertainty decomposition algorithm. However, there is also evidence suggest that this Bayesian behaviour is only approximate during long-term generation in LLMs, invalidating the Bayesian assumption [20, 61]. In light of these previous works, our innovation lies in the attempt to promote permutation-invan maintain the Bayesian assumption that make.
generation and filter non-Bayesian generation from auxiliary data to
Permutation Invariance and Exchangeability in LLMs. The generation in language models is dependent on the position of tokens [62, 121]. This is clear violation of exchangeability. which is necessary for the application of Finetti's theorem. [119] assumes the exchangeability of LLM generation to apply de Finetti which allows for the estimation of the topic distributions from L.Ms.. However, they do not apply permutations during ICL to the context. [1 [118] discusses the importance of exchangeability for quantifying uncertainty in ICL. They investigate methods to promote permutation invariance during pre- pre-trainnn and fine-tuning or architectural modifications to the transformer through causal masking. Whilst they suggest using permuted data as a data augmentation technique during training, our permutation invariant conditional generation is purely applied during inference. Our approach incurs greater cost during inference time but does not require fine-tuning of the LLM.
Martingale Posteriors. The Martingale posterior [20, 22, 53] construct generalised notion of posterior distribution by the following steps: (1) defining sequence of predictive distributions
{p" (y a ya) m)) for all 1, (2) sequentially generating yj ~ p² (y, F, y:) hiej)
for + 1,.., N with N >> and (3) computing proxy latent parameter
38

---

g(f(2,,
((l(a,y))} yolk U U((z),)))) (y), 8y)1==++1 via some function g. Technically, this defines the following
form of Martingale posterior (D = Vt)]
(m))) = 8( = g((( U (1);
x Il p'(y, Jilee)))
==
If a proxy likelihood model q(y" 12", is further specified, then the predictive Martingale posterior can be defined as [53]
(v) D. f [z)} )dp). N (y " (2, - g(y"|2", |x" ,))" (
(14)
Therefore an uncertainty decomposition (as presented in Section 2) by conditioning on the proxy latent parameter v is plausible. We can compute the "Martingale version" of total uncertainty as Hl[g (y"|æ", x=" , D, {ay ا(m+رر and the aleatoric uncertainty as M (w)))(m;))" [Hlqqy" v]]] Epistemic uncertainty can then be obtained via simple subtraction arithmetic.
The (predictive) Martingale posterior generalises conventional (predictive) Bayesian posterior as does not require {p"(y x=" ((x:,)))
to be consistent and correspond to the probability of an exchangeable sequence; instead it requires convergence properties of the (y"(y"|xr ( {=)]=)) distributions and the g function when N + oc, where we refer to [22] for details. In practice, to obtain robust estimations of Martingale posteriors, N is often substantially larger than n, incurring significant computational cost, and the computation of {y, [y1ري 1=n+1 samples cannot be parallelised.
To make critical comparison to our proposed concept of variational uncertainty decomposition, we note that in general Martingale posterior is also different from the conventional Bayesian posterior, even when there exists an exchangeable model such that y" (:) y=))===)
p(y" xx""" {(.,,y)) H=)) for all n 2 1. The key reason is because the corresponding Bayesian model ply ar, ())(O) is implicitly defined via de Finetti's theorem applied to p(y"Ia" {( y=))== mean- ing that its latent parameter O is an "unknown unknown". i.e., the format of @ (e.g.. dimensionality, value domain, etc) cannot be explicitly specified. Hence in general sb and @ are two different random variables in different domains (and thus the name "proxy" for &i in our terminology). Consequently, the uncertainty decomposition results based on & are no longer faithful directly the implicit Bayesian model p(y az. (f), and their estimation gaps, when referencing to the implicit Bayesian model's ss uncertainties U. and Uz. are yet to be established. On the contrary, our proposed variational estimators Va and V. are faithful bounds to U, and U.. respectively, and we have identified the exact mathematical expression of the estimation gap in Section 3.1, which can be interpreted as residual information gain and/or remaining disagreement in fantasy.
Uncertainty Quantification for LLMs. Reliable and robust uncertainty quantification is an area of growing importance in the field of language models 196, 1]. A common approach, which we employ in this paper, is using token-level probabilities [43, 60. 19] by analysing the probabilities of the tokens generated by language model. These probabilities can be further calibrated by adapting standard methods for uncertainty quantification in deep learning such as temperature scaling [29, 1 12, 13, 111, 110], focal loss training [73, 112], and conformal prediction [95, 117]. Alternatively, careful prompting can elicit qualitative or quantitative verbalisations of the uncertainty in statement made by the language model [8, 58, 68, 101]. In situations where an LLM generates open-ended answers to a question, loven-level methods struggle to accurately capture the uncertainty of response as it is possible to generate diverse responses in natural language that are semantically cquivalent or similar. To address this, semantic similarity methods cluster similar responses together and report the combined uncertainty from each cluster [3, 50, 59]. Recent works have also taken mechanistic interpretability approach to uncertainty quantification by using probes to analyse the hidden states of the LLM to diagnose when a model is uncertain [48, 2].
Uncertainty Decomposition for LLM In-Context Predictions. Uncertainty decomposition for LLMs has also been explored in previous works; however, the definitions of aleatoric and epistemic uncertainty vary from the traditional definitions in prior Bayesian literature. [39] considers the aleatoric uncertainty of a response the ambiguity in the input. Therefore, given distribution of "clarifications" q( Cla²") for a particular prompt, the epistemic uncertainty is defined as the mean con-
ditional uncertainty of a particular clarification Ey(C) [Hly" xx" @ CJJ. In contrast, we seek to find
39

---

the minimal conditional entropy given auxiliary data, which acts as an upper bound to the underlying Bayesian conditional entropy. Furthermore, the focus of [39] is primarily zero-shot and few-shot prediction, whereas we consider tasks where training dataset is provided in context. Ling et al. [60] approaches uncertainty decomposition of in-context learning by also employing the interpretation that ICL. performs Bayesian inference. However, they define epistemic uncertainty as the conditional entropy p0)p) [H] [y" O]] Both [39] and [60] reverse the traditional definitions of Bayesian uncertainty decomposition [45] and therefore, we do not use these methods as baselines.
and aleatoric uncertainty as the mutual information 1(y"; 0|æx" D).
Bayesian Approaches to Transformers. In this work, we view in-context learning as implicit Bayesian inference. However, prior work has connected the transformer architecture with Bayesian inference more explicitly via adaptation [115, 7, 79] has allowed for parameter-efficient avenues for Bayesian deep learning in transformers. Alternatively, neural processes have been integrated with transformers [77] to provide another approach to Bayesian uncertainty quantification in transformers. A connection between attention and sparse GP posterior mean is also established in [14], which further builds deep Gaussian process with transformer-type architectures.
approaches [93, 66, 11]. In particular, low-rank
Applications to In-Context Exploration. Techniques used to quantify uncertainty in LLM pre- dictions can be used drive in-context axploration-xxpoittaain and bandit tasks. efficient exploration algorithms such as Upper Confidence Bound [52, 6] and Thompson Sampling (TS) [82, 81, 94] require modelling the epistemic posterior distribution over possible outcomes either implicitly, through visitation counts, or explicitly, for example via ensembles. By modelling the epistemic uncertainty, the agent is able to reason about potential outcomes with uncertainty due to lack of data and explore in promising directions. Previous work that analyses the in-context exploration capabilities of LLMs includes [49], where the exploration capabilitiess of LLMs are compared to those of standard algorithms on small-scale tasks, and [70], which investigates the exploration capability of LLMs on natural language bandit tasks. The work in [78] further explores and benchmarks LLMs' abilities on number of bandit tasks and offers ways to improve the efficiency of exploration by introducing algorithmic enhancements that better align LLMs with the lolooottioo-eplliaaatiiiiooioooi task. This line of work focusing on bandits is complemented by [109], which extends the benchmarking to include multi-step tasks in addition to bandits. Finally, the work in [4] adapts the TS heuristic to the LLM setting, enabling LLM agents to tackle sequential decision decision-making tasks analogous to that of the full reinforcement learning setting. Uncertainty-ave -aware exploration has also been used in letive-learning settings obtain smoother decision boundaries of LLMs by identifying the data points that will give smoother boundaries [120].
tasks. In reinforcement learning
Abstention. The ability to defer the prediction of language model is important in high-risk and sensitive applications of language models such as medical settings [28, 71J. In particular, abstention from answering questions is closely related to the domain of selective classification, where we learn a corresponding selection function g X' + {0, 1} alongside the standard classifier f X + V. where ((x) = 0 and (xx) 1 results in the classifier prediction being 'rejected' and 'accepted" respectively [26, 91, 72]. The goal is to minimise the loss of accepted samples (risk) and maximise the uncertainty estimate a proxy for our selection function and use the ranked samples to determine the threshold of the score, thereby seeking minimise the risk for a specified coverage.
number of accepted samples (coverage). In our setting,
OOD Detection. such as medical diagnosis and autonomous driving, where models can make confidently wrong predictions on inputs far from the training distribution. Foundational work demonstrated that softmax confidence often fails under distributional shift, establishing simple baselines for OOD detection in deep neural networks [34]. However, epistemic uncertainty has been shown to be useful in OOD and hallucination detection [110, 45). This led uncertainty-baser methods which estimate epistemic uncertainty such as deep ensembles [51], where the uncertainty is measured through model diversity, and prior networks where distributional uncertainty is used in addition epistemic uncertainty [65]. In NLP, pre-trained language models have been used for OOD detection [33] through non-Bayesian approaches such as contrastive learning [122], unsupervised detection with transformers [114], and conditional generation strategies to improve OOD discriminability [91]. Extensions to multimodal settings further explore OOD detection in vasion-language tasks [69].
Detecting out-of-distribution (OOD) inputs is critical for real-world applications
40

---

Mutual Information Estimators. The quantity of mutual information for which we provide a lower bound has many applications including Bayesian experimental design [89], independent component analysis [54], neuroscience [83] and causality [38]. However, mutual information between two variables is considered challenging to estimate [104] as requires access the joint distribution of the variables, which is often unavailable. Variational methods are a popular approach used to lower- bound mutual information [87], and particular, MINE [9] and InfoNCE [80] are methods based on variational lower bounds to the mutual information. However, when estimating 1[y"; ; 0/ D), these methods require access to samples from both random variables y and 0, but in our problem setting, the latent parameter 0 is implicitly defined and thus camnot be sampled. Still, provides a variational lower bound to the conditional mutual information quantity in this challenging setting, where innovation sidesteps the access requirement of the variable by constructing optimisable probes via Markov chain y + - U, enabling data processing inequality arguments and allowing lower-bound optimisation similar to MINE.
approach
G Experiments
G.1 Code Implementation
The following delineates the foundation of our experiments:
Codebase:
Python & PyTorch
CPU: AMD EPYC 7443P
. GPU: NVIDIA A6000 48GB
We leverage 1eee2.5-11BB/4B-Inttttue [88] and Llama-3.1-8B [102] in our experiments. The following delineates the configurations of our LLM.
Temperature: 1.0
Log Probs: 10
Max Tokens: 10 (Qwen2.5-14B/7B and Llama-3.1-8B), 512 (Qwen2.5-14B-Instt)
G.2 Further Baselines
In this section, we include further comparisons Martingale posterior distributions [22], an alternative uncertainty decomposition method for implicitly defined Bayesian models on an exchangcable sequence. However, one of the disadvantages of the Martingale posterior method is that need to make distributional assumptions on the form of the proxy likelihood q(y|xz, eb) (see Appendix F for further discussion). This can become particularly restrictive as the choice of likelihood model can greatly impact the estimated total uncertainty as we show in the following Figures 14-16 for the "Logistic Regression" Dataset and Figures 17-28 for the "Moons Dataset and the Table and (further details on these datasets can be found in Appendix G.3).
As both datasets are (binary) classification problems, consider four proxy likelihoods of the form q(yl xe, v x (x)) - Pv (z, which we denote as 'linear', 'quadratic', 'cubic", and 'kerell:
Linear: (x) ( Ej thax;) where is the standard logistic function (sigmoid). This is the probability distribution that used to generate the dataset but may not reflect the internal likelihood of the LLM. Here, vb E 1241 (where is the dimension of æ))
Quadratic: p=(x)= o (vbo E. Ver, + Eisy ))g))zz)) where o is the standard logistic
function. Here, IR
(++1))+2)
Cubie: (x) = a (1o + E, Vin Eis Vijazay + Eisisn (isk:1ay4)) where a is the (diD)a12Nd standard logistic function. Here, IR.
Kernel: (=( a (vb + fxx)), where a is the logistic regression function, v/a and bB are parameters to scale the logits, and fx) are the logits. The logits fy( are of the form
fx(x)) = vo + Σ ywvrk( x),
41

---

where { b is the training data set of initial examples and the generated, and the kernel & the RBF kernel. To obtain the estimate for the proxy latent parameter u, we follow the Platt-scaling method [86] which is implemented in the scikit-learn package SVC [85]. Here, v RN+3 (where N the size of the combined initial training examples and the generated sample path from the LLM).
To sequentially generate the next sample in the sequence (æ Y,) ( yy)l((, y.) y:))ey)) we permute the order of ys) jej in the prompt for the LLM. This technique of permuting the observations in the sample path used in previous work the context of Martingale posteriors and LLM [20] and emulates the permutation-invaare method. However, our method, we sample the posterior for multiple permutations and compute an average obtain our estimate for ply" ll=", ()) whereas at each step in the sample path for the Martingale posterior, we only perform one permutation.
sampling approach that we use for our VUD
Table 7: I.2 Distance between total uncertainty given by Martingale posterior distribution and empirically observed total uncertainty. Logistic Regression Dataset.
LIKELIHOOD MODEL
QWE22.5- .-BBB
QWEN2.5-7B
5-
LLAMA-3.1-80 3-BB
LINEAR FEATURES
1.347
1.096
1.064
QUADRATIC FEATURES
3.624
4.334
1.381
CUBIC FEATURES
1.000
0.873
0.589
KERNEL- BASED
7.823
9.229
4.731
Table 8: 1.2 Distance between total uncertainty given by Martingale posterior distribution and empirically observed total uncertainty. Moons Dataset.
LIKELIHOOD MODEL
QwE22. 5-14B
QWEN2.5-27
LLAMA-33 3.1BB
LINEAR FEATURES
2.379
1.789
1.515
QUADRATIC FEATURES
2.796
2.781
2.819
CUBIC FEATURES
2.697
2.254
2.781
KERNEL- -ASSDD
1.530
1.213
1.254
1.0
Toid Expied Cauntary
0.8
0.8
Uneetainty
Uheartainty
0.6
0.4
0.4
Teal Exaated Coonityy
Teat Miotoolee Txoestae
0.0
15
-10
10
15
(a) Linear Features
(b) Quadratic Features
1.0
Tiaal Reppdied Uacatarr
0.8
0.8
Unestalaty
Uncortality
0.6
Teall Exgptte Coemdary
0.6
Teat Stutbaale Dooeabte
0.2
8.0
15
-00
-5
10
15
15
5
100
15
(c) Cubic Features
(d) RBF Kerne
Figure 14: Martingale Posterioir Uncertainty Decompositions for Logistic Regression (Qwen2.5-143
42

---

1.0
1.0
Hintlauia Staralay
0.8
Abster Hontzaal Ceoerable
11.8
Unorraingy
Uncertiinty
0.6
1..6
0.4
0.4
Expreaa Caundry
0.2
15
10
15
(a) Linear Features
(b) Quadratic Features
1.0
Teid Enepiteal Canntarrg
1.0
Ubcertainty
0.8
Alsaerz Nartbaak Cisetabee
Unoertainty
0.88
0.6
Expred Caonaey
0.2
0.0
0.2
15
(c) Cubic Features
(d) RBF Keree
Figure 15: Martingale Posterioir Uncertainty Decompositions for Logistic Regression (Qwen2. 5-7B)
1.0
Toid Expind Cauntary
0.8
Uneetalitt
Invertaiity
0.6
0.4
Tetal Expattal Cormaktr
0.2
0.2
Trea Mettager Uheetieee
15
-10
10
15
(a) Linear Features
(b) Quadratic Features
Tiaal Eeapaied Cacataary
0.8
Uncortaitty
Uneeralatt
0.6
0.6
0.2
0.4
Tred Mes apai teetees
8.0
15
10
-5
10
15
15
10
15
(c) Cubic Features
(d) RBF Kemel
Figure 16: Martingale Posterioir Uncertainty Decompositions for Logistic Regression
(L1 allam.-3 1-8B)
43

---

-2
-2
[CT, Datic y
ICL Dut
p=
[TT) Datti y-
0.8
2
TCT. Data y-1
pm
gosto
06
5
6
0.4
0.4
0.10
-1
-1
-2
0.2
-2
0.2
0.05
0.0
3.3
0.00
Total Uncertainty
Alcatoric Uncertainty
Epistemic Uncertainty
Figure 17: Martingale Posterior Uncertainty Decomposition for "Moons 1 Dataset Linear Features (Qwen2.5.) 5-1B)) 14B).
-2
y
Duti:
Datic
ICL. DuEa:
Laat g
0.4
0.3
0
35 (a) Total Uncertainty
32
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 18: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Quadratic
Features (Qwen2.5-2 5-BBB
-3
22
-1
-2
-1
2
0.8
SCT.
0.8
0.6
0.4
0.2
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 19: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Cubic Features 00wen2.5-111 14B).
-3
-2
[C1, Dati y-0
ICL Data: 5
Dats P-0
0.8
ICT.
0.8
ICE Dutue y-)
0.25
0.8
0.20
06
3.15
072
0.4
0.4
-1
-1
0.10
-2
0.2
-2
3.05
-2
0.0
-3
03
0.00
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 20: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Kernel-Based
Likelihood (Quen2.5 54B).
44

---

-2
-2
[CT, Datic y
ICL Dut
p-8
[TT: Datii y-
0.8
2
TCT. Data y-1
pm1
0.15
066
5
6
0.4
0.4
-1
-1
-2
0.2
-2
0.2
0.05
0.0
3.3
-3
0.00
Total Uncertainty
Alcatoric Uncertainty
Epistemic Uncertainty
Figure 21: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Linear Features Swwen2.5-723
-2
y g
Duti:
Datic
0.6
ICL. DuEa:
Latt g
0.8
0.4
0
0.3
0.1
35 (a) Total Uncertainty
32
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 22: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Quadratic
Features Saven2.5-7BB.
-3
22
-1
-2
-1
2
SCT.
0.8
0.5
12
0.4
0.4
0.2
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 23: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset (-7)n..5-7B).
Cubic Features
-3
-2
[C1, Dati y-0
ICL Data: 5
0.8
DCL Data y Taax
ICT.
0.8
0.3
85
06
0.2
072
6
0.4
-1
0.1
-2
0.2
-2
-2
0.0
-3
03
-1
0.0
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 24: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Kernel-Base
Likelihood (Quen2.5 5-7B).
45

---

-2
-2
[CT, Datic y
ICL Dut
p-
0.30
[TT) Daati y-
0.8
2
TCT. Data y-1
pm
gosto
066
0.30
5
6
0.15
0.4
-1
-1
-2
0.2
-2
0.2
0.05
0.0
3.3
-3
0.00
Total Uncertainty
Alcatoric Uncertainty
Epistemic Uncertainty
Figure 25: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Linear Features
(Ll ama- 3. 1-8B).
-2
0.6
y
Duti: y
Datic y
ICL. DuES:
Llat g
0.8
0.5
0.4
0.5
0.3
0
0.4
0.2
0.1
32
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 26: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Quadratic Features (Llama-3. -3. 1-8B).
22
3
-3
22
-1
-2
-1
3
SCT.
0.8
9.
0.6
12
0.4
0.4
0.2
0.2
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 27: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Cubic Features
(Ll A.lama-3.. 3. 1-8B).
-3
-2
[C1, Dati y-0
ICL Data: 5
DCL Data y
0.8
ICT.
9.8
Taax
0.3
06
0.5
82
075
0.4
-1
0.1
-2
0.2
-2
0.2
-2
0.0
-3
03
-1
0.0
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 28: Martingale Posterior Uncertainty Decomposition for "Moons 1" Dataset Kernel-Base Likelihood (Llama-3. 1-8B).
46

---

G.3 Synthetic Toy Experiments
We qualitatively evaluate the decompositions of the variational uncertainty decomposition algorithm on a variety of synthetic classification and regression settings. In this section, we give details on the ground-truth distributions used to create the synthetic datasets.
0.25 and bias 3p -0.5. The covariates are generated from a Gaussian distribution with mean 1.5 and standard deviation In our visualisations, we use Perturbations with 15 auxiliary data points and perturbation scale À = 0.1 to decompose the uncertainty for the logistic regression task. In Figures 4a and 29, we plot the uncertainty decomposition for an ICL dataset of size [D] 15 and Figure 30, we plot the decomposition for [D] 75. We plot x" values in the range 15,, 15) with step size 0.2. In Figures 31 and 32, we plot the epistemic and aleatoric uncertainties as the dataset size increases for in-distribution (r = 0, 5; solid lines) and out-of titto---iitribution (a* = -15, -10, -5, 10, 15; dotted lines) points. As the uncertainty
Logistic Regression. We consider I-D logistic regression problem with coefficient
given x" is dependent on the particular dataset, we average the uncertainty at x* over 10 datasets of the same size d to obtain the estimate of the mean aleatoric uncertainty at d.
Uxvertainty
Quecisiity Uneetaiity
Alesteri Uherrontty
Unoetainty
0.8
ECll Dorz: 9-1
Uncortaiaty
0.8
ICL Dalay y-
0.6
REL
0.6
ICL Data: y-1
0.2
0.2
Test Covariate 21
Test Covariate
(a) Twen2.5-7B
(b) Llama-3.
Figure 29: Uncertainty Decomposition for Logistic Regression |D] 15.
Dale p-1
25 Bii
14 Inke
(a) Quen2. 5-14B
(b) Qwen2. 5-7B
(c) Llama-3.1--8
Figure 30: Logistic Regression with [D] 75.
0.9
Egleteniic Utaoettaiats
0.25
Alcatode Uneerraiityy
0.8
0.20
0.7
0.15
3.5
0.10
0.05
3.4
0.00
20
80
60
80
[D]
(a) Epistemic Uncertainty vs Size of Training Set (b) Aleatoric Uncertainty Size of Training Set
Figure 31: Epistemic Uncertainty and Aleatoric Uncertainty vs Dataset Size (Qwen2.54 5-7B).
47

---

Elpisennii Uheerralnnn
0.16
Aleatotic Uncertalatty
0.5
0.10
0.7
0.06
0.00
20
80
ID]
[D]
(a) Epistemic Uncertainty vs Size of Training Set (b) Aleatoric Uncertainty vs Size of Training Set
Figure 32: Epistemic Uncertainty and Aleatoric Uncertainty vs Dataset Size (Llama-3 1-8B).
-1, bias 8 = 3 and Gaussian noise with zero mean and standard deviation a = 2. The covariates are generated from Gaussian distribution with mean and standard deviation 2. We use Perturbations with 5 auxiliary data points and perturbation scale À = 0.1 to decompose the uncertainty for the logistic regression task. We reduce the number of auxiliary data points due the increased computational cost of computing distributions for regression problems. In order to obtain smoother uncertainty computations, we average the uncertainties obtained over sampled datasets of size
Linear Regression. We consider I-D linear regression problem with coefficient
[D] = 15. We compute uncertainties for x" in the range [-15,15) with step-size 0.2 and plot the obtained decompositions for entropic uncertainty and variance in Figures 4b, 33 and 34. We also provide an example decomposition for the uncertainty and variance for a single seed for completion Figures 35, 36 and
Total Usenteletty
Uncertainty
Uncertainty
Alcasdi Cacarniaty
ICD Dats
2-
Uncrraiity
200 Dsit
10
Test Covariate
Test Covariate
(a) Qwen2.5-22
(b) Llama-3. 1-8B -1-8B
Figure 33: Linear Regression (Entropic) Uncertainty Decomposition.
10a
Aeavet Ticles
10
Nakec Turae
Variance
Varisure
10
Variaure
20a
20a
10a
0
10
Covariate
Covarinze
Test Covarlize
(a) Quen2. 5-1BB
(b) Qwen2, 5-7B 2-BB
(c) Llama-3. 1-8B
Figure 34: Linear Regression Variance Decomposition.
10
Unorrtating Labitt
10a
-10
KCL KLL
Test. Covatiate s
Test Covariate
Tes: Covariate x
(a) Predicted Mean and Std
(b) Entropic Uncertainty Decomp.
(c) Variance Decomposition
Figure 35; Uncertainty Decompositions for Linear Regression (Qwen2. 5- 14B).
48

---

10
Unovewintty
10
fentwer
10"
bumee KC Ki
-10
Dae
Text Covariate
Test Covariste
Teo Covariate
(a) Predicted Mean and Std.
(b) Entropic Uncertainty Decomp.
(c) Variance Decomposition
Figure 36: Uncertainty Decompositions for Linear Regression (Qwen2.5 5-BB).
Unoectaiony
8
10
10
10
10
Test Covwiite x
Tost Covariate
Test Covariste 2
(a) Predicted Mean and Std.
(b) Entropic Uncertainty Decomp.
(c) Variance Decomposition
Figure 37: Uncertainty Decompositions for Linear Regression (Llama-3 3. 1-8B).
Heteroscedastie "Gaps" Regression. We model the "gaps" the combination of 3 linear regression datasets. The parameters of the 3 clusters are in Table 9. To generate the small in-context learning dataset, sample from this combined dataset. In our visualisations, auxiliary data points and perturbation scale À = 0.1. We sample a single dataset of size [D] = 30. We compute uncertainties for xt in range [- 15, 15) with step size 0.2 and plot the obtained decompositions in Figures 38, 39 and 40.
use Perturbations with
Table Heteroscedastic "Gaps" Dataset Parameters
CLUSTER
DATASET SIZE
CORFECCENT
BIAS
NoISE
E/a
Varlz]
50
0.75
1.0
0.1
-7
0.75
2
50
0.75
1.0
0.1
-1
0.75
3
100
0
-0.5
2
5
10
Caaraiay
Varisur
10a
10-1
1ILL Data
DCL Dap
10
Test Covariane
(a) Predicted Mean and Std.
(b) Entropic Uncertainty Decomp
(c) Variance Decomposition
Figure 38: Uncertainty Decomp. for Regression Tasks with Gaps in ICL Data. 5-1B)
IC% Dars
Unoertiiaty
Varianee
Aneors Vataeo
y
10
10
Test Covnriate
Covoriate
Covariate
(a) Predicted Mean and Std.
(b) Entropic Uncertainty Decomp.
(c) Variance Decomposition
Figure 39: Uncertainty Decomp, for Regression Tasks with Gaps in ICL Data (Qwen2. 5-7)).
Tinal Vamtnee
Unvertaiity
32
Varlaaee
Take
y
10t
-0-1
Test Covadiate
Tes: Covariate
Test Covariate z
(a) Predicted Mean and Std.
(b) Entropic Uncertainty Decomp.
(c) Variance Decomposition
Figure 40: Uncertainty Decomp. for Regression Tasks with Gaps in ICL Data (Llama-3) 3. 1-8B)
49

---

Moons Dataset. We use the make_mooke moons two-moone moons dataset generator from scikit-leard We set the noise parameter in the "Moons I" and "Moons 2" datasets to 0.1 and 0.4 respectively. Figure in the main text shows the decomposition for "Moons 1" dataset. We use Perturbations with 15 auxiliary data points and perturbation scale 0.1. For the "Moons
dataset, we sample
single dataset of size [D] 30 and compute uncertainties for æ in the range [-1.5, 2.5) x [-1.5, 2.5) with step-size 0.2 for each interval. The decompositions are given Figures 41 and 42. For the "Moons 2" dataset, we sample a single dataset of size [D] 30 and compute uncertainties for x² the range [-3.0, 3.5) x [-2.5, 3.0) with step-size 0.2 for each interval. The decompositions are given in Figures 43, 44 and 45.
1.0
1.0
2
2
x-0
0.8
p-1
0.6
0
0
0.4
0.2
-1
-1
0.2
0.0
0.0
-0.0
J1
(a) Total Uncertainty
(h) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 41: Uncertainty Decomposition for "Moons 1" Dataset (Qwen2. 5-7B).
1.0
1.0
2
y=
0.8
0.8
98
0.6
96
0.4
0.1
0
0.1
02
-1
0.2
-1
0.2
0.0
0.0
-1
0.0
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 42: Uncertainty Decomposition for "Moons Dataset (Ll (.lama-3 3.-aaB3.8BB)) 1-8B).
0.8
0.6
066
0.6
0.1
0.4
0.4
-1
0.2
122
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 43: Uncertainty Decomposition for "Moons 2" Dataset (Qwen2. 5-14B))
50

---

0.8
V=11
P-1
0.8
0.8
0.6
0.8
0.4
0.4
-
0.0
0.0
21g
(a) Total Uncertainty
(b) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 44: Uncertainty Decomposition for "Moons 2" Dataset (Qwen2.5-7B).
1.
1.0
y-0
y-0
088
0.8
3.6
0.4
0.2
022
0.2
-2
-2
3
0.0
0.0
0.0
(a) Total Uncertainty
(h) Aleatoric Uncertainty
(c) Epistemic Uncertainty
Figure 45: Uncertainty Decomposition for "Moons 2" Dataset (Llama-3 3--8BB))
Spirals Dataset. We
n-arm spiral dataset generator generate the spirals. We set the number of arms to 3 and noise to be 1.2. We also scale the covariate down by a factor of 4 so that all the points would appear in [-4, 4]
4, 4]. Due the complexity of this task, size (D) = 200 and we compute uncertainties for x in the range [-4,4] 4) x [-4, 4) with interval 0.1. To mitigate the of increases prompt size and the number of test data points, we use Repeated to obtain Z. The decomposition for Qwen2.5-14B is given in Figure 6. We provide decompositions
sample dataset of
for Qwen2.5-7B and L
shown in Figure 46 and 47.
15
0.
100
1.0
0.
0.70
0.25
0.
L10
0.000
0.0
(a) Posterior Prob.
(b) Total Uncertainty
(c) Aleatoric Uncertainty (d) Epistemic Uncertainty
Figure 46: Uncertainty Decompositions for Spirals Classification Task (Qwen2.5- .-7B)
L25
1,0
2.8
0.50
02
0.25
0.00
(a) Posterior Prob.
(b) Total Uncertainty
(c) Aleatoric Uncertainty (d) Epistemic Uncertainty
Figure 47: Uncertainty Decompositions for Spirals Classification Task (Llama-3. 1--8BB)
G.4 Bandits
In bandit problem, we have multiple trials (or equivalently rounds), where the agent must choose an action (or equivalently arm) which gives reward. The agent has access to the actions made and
51

---

rewards obtained for the previous trials. We denote run or seed to refer to a particular chain of trials. For all the bandit experiments. we run the algorithm for T 200 trials.
LLM-UCB Algorithm. In a UCB algorithm, we have:
ag = (eggaaxex(Qeloa + aU.(a)),
where (a) the expected reward from action (i.e. arm) at 1, Ug(a) is the uncertainty in the reward from action at and is the exploration rate [52]. In the LLM-UCB algorithm that we use to compare the epistemic and total variance decomposition in Section 5, we set Q+(a) = p(rla, Du), where D. = {(a,, r:) r:)) [ the prior action, reward pairs already observed in a run. the epistemic variance setting (a) = Varlriu, Pr]] ming Ey |Var[rlu, Z, D:]] and in the total variance setting Uz(a) = Var D.]]. For each a and p. we run 10 seeds.
Non-LLM Benchmark. We use the standard UCBL 1 algorithm and the Greedy algorithm [52] as a non-LLM benchmark to ensure that the LLM-UCB algorithm has comparable performance to standard bandit algorithms. An exploration rate of a = 0.75 is used for the UCB1 algorithm. We run 5000 seeds for both UCB1 and Greedy for each a and p.
Instruction Prompting Benchmark. In [49], an instruction- tunne LLM is prompted to attempt the Buttons bandit task and there is a thorough investigation of the impact of the prompt configuration on the LLM's performance. The authors conclude that the most successful prompt configuration is: BSSČO, which consists of: a suggestive framing that the LLM is solving a bandit task; a summarised history of prior actions (including average rewards per action and counts per action); reinforced chain-of-though prompting; and temperature parameter of O. For fair comparison of model performance, we use Swen2.5-14B-Instruee
Swwen2.5-7B-Instrucc, and Llama-3.J L -n8ttnrtttttt [88. 102] to benchmark the performance of the LLM-UCB algorithm for the base models Qwen2.5. 14B, Qwen2.5-7B, and Llama-3.1-88 respectively. See Appendix H.2 for an example prompt. For each a and p, we run 10 seeds.
Role of) p and aleatoric variance. The means of the optimal and suboptimal arm(s) in the Buttons setting are P p + $
2 and Pa p 1 respectively. Now, the variance for Bernoulli random variable of mean q is q(1 q). This is quadratic with a maximum q = 1.2 Therefore, if >>>>> 3.
Ipo
(p- , 2
2 (p 2 11
1
1ー
A -
1 A 2
lp 2
A
Therefore, the true variance of the suboptimal arm is higher than the true variance of the optimal arm.
Choice of a. In our experiments, we choose a = 2, 5. In UCB1 smaller choices of a are typically chosen [52], however this is primarily due the slow decay of Ua(a) in the UCBI algorithm. The decrease in epistemic uncertainty with the number of trials is significantly faster, and therefore, we use higher a. Since the total uncertainty is the sum of the epistemic uncertainty and the alcatoric uncertainty, the difference in the uncertainties is a multiplied by aleatoric uncertainty.
Metrics We use multiple metrics assess the performance of the bandit algorithms. Suffix-fut fail frequency and K MinFrac are metrics introduced in [49] to assess the performance of bandit runs.
Mean regret: For a run of T trials, the mean regret is defined as + En= E[[(a:)] n², where μ is the optimal reward and E[r(a,)] is the mean reward for arm ds. We report the mean and standard deviation across the different seeds.
Mean worst-case regret: We take the mean and standard deviation over the 30% of seeds with the highest mean regret. For algorithms where there is large discrepancy between the mean regret and worst case mean regret, this indicates that the variability in the performance of the bandit algorithm is high.
Median reward: For each seed run, we compute the mean reward + Σ Tr. We then report the median mean reward across all the seeds.
Suffix fail frequency: For given run, there is a t-suffix failure, if the optimal arm is not chosen the trials [4, Tj. The suffix fail frequency SulFFalFreq(() is the proportion of
l-suffix failures across all the seeds. This metric measures particular failure mode of
bandit-algorithms due to lack of exploration, where as a result, the optimal arm is not chosen.
52

---

K MinFrac: For a given run j. let ((3) be the action counts. Given T runs, J seeds, and arms, K MinFFacc TJ Eز= ming St This metric measures wniform-like failures of bandit algorithms, where due excessive exploration, the algorithm behaves closely to one that uniformly chooses an action.
Results. In Tables 10 and 11, we provide the results for the Qwen2.5-7B and Llama-3.1-8B models. We also plot the average cumulative regret across different seeds for 0.5, 0.6, 0.7 and a = 2, 5 in Figures 48-59, Each line in these figures corresponds to the cumulative regret for particular seed. Here, we seed that in general, the algorithm that uses the epistemic variance estimate generally has more consistent performance than the algorithm that uses total variance.
Table 10: Buttons Bandit Problem. TV is Total Variance, EV is Epistemic Variance. (Qwen2. 5-7B)
WORST-CAS 0.128
SultkalFFreg
-Minhrec
LCB1
0.094
0.510
0.29
GREEDY
0.199_cc
0.101
0.525
0,460
0.03
0.7 10:0. 20.0
INSTRUCT BASELINE
0,161=c2
0.107
0.495
0.26
TV (a 2)
0.175=ex
0.0682
0.565
0.1
0.0.3
BV
0,144=14
0,091
TV [o
0.196_.m
0.0
0.545
0.2
0.04
0.160
0.132
UCB1
0.1
0.094
0.610
0.28
GREEDY
1.199_ccm
0.0022
0.645
0,396
0.03
INSTRUCT BASELINE
0.076g
0.620
0.18
TV (a 2)
0.199_cm
0.0900
0.627
0.1
0.0.3
EV
0.088
0.061
0.12
TV [a
0.198
0.167..02
0.570
0.5
0.07
0.156..
UCBI
0.122
0.094.x0
0.710
0.0
0.27
GREIMY
0.199_00
(1 0.085
0.760
0.369
0.03
INSTRUCT BASELINE
0,132-ch
0.087
0.703
0.18
TV (a 2)
0.198
0.088
0.728
0.4
0.02
EV
0,141-14
0,070.ues
TV [a
0.195
0.508
0.6
0.04
BV (a
0.143_
0.
Table 11: Buttons Bandit Problem. TV is Total Variance, EV is Epistemic Variance, (Llama-3. 1-8B)
METHOD UCB1
MEAN WORST-CASE REGRET
MEAN REGRET 1 MEDIAN REWARD
SullFaiFFre
к Minffaa
0.128
0.094
0.510
0.0
0.29
0m199_cm
0.1012m::
-0.8 --00
INSTRUCT BASELINE
0.161
0.107
0.495
0.0
0.26
TV (a
0.160
0.071
0.557
0.05
EV (o =
0.149-1m
0,097
0.505
0.0
0,21
TV (a
0.149 r36
0.066..4 oi
0.555
0.05
EV (a 5)
0.169
0.153
0.432
0.0
0.73
UCBI
0.127
0.094_
0.610
0.28
1.199_mm
0.0922
0.03
INSTRUCT BASELINE
0_111
0.076amm
0.620
0.0
0.18
TV (a
0.088 11%
0.035
0.670
0.04
EV (a =
0.140
0,077
0.635
0.0
0,17
TV (a
0.198
0.138...0
0.568
0.6
0.04
EV (a 5)
0.139-..
0_1113=.22
0.588
0.0
0.50
UCBI
0.122
0. 1.094gmro
0.710
0.27
GREEDY
0.199_mm
0.369
0.03
INSTRUCT BASHLINE
0.132-rm
0.08 7
0.70%
0.0
0.18
(a
0.168-14
0.063
0.728
0.1
0.04
EV (a =2
0.111
0.053 .e2
0.745
0.0
0.08
TV (a
0.197-mx
0.165
0.613
0.5
0.04
EV (a 5)
0.127=mm
0.087=
0.688
0.0
0.35
Cumulativv Meau Regrel
30
Cuv Mesu Regrel
40
30
20
-20
10
0
50
100 Trial
150
200
50
100
150
200
(a) Epistemic Variance
(b) Total Variance
Figure 48: Cumulative Mean Regret for Bandit Experiments (Qwen2.5- 5-14B, = 0.5, a = 2).
53

---

Cumulstivve Mean Regret
Cunmllative Menia Regret
30
30
20
10
10
0
50
200
100
150
200
Trial
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 49: Cumulative Mean Regret for Bandit Experiments (Qwen2. 5-14B, 0.5, 5).
Cumulative Mean, Regrel
Cumulative Mean, Regret
40
30
20
20
10
0
50
100 Trial
150
200
50
100
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 50: Cumulative Mean Regret for Bandit Experiments (Qwen2. 5-14B, p
0.6, a 2).
Cumulative Mean Regret
Cumulative Mean Regret
400
20
20
10
0
50
100
150
200
50
100
160
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 51: Cumulative Mean Regret for Bandit Experiments (Qwen2.5 5 14B,
0.6, a = 5).
Cunulative Mean Regeet
20
Cumolativv Mean Regeet
16
30
10
20
10
30
100 Trial
100 Trial
150
200
(a) Epistemic Variance
Total Variance
Figure 52: Cumulative Mean Regret for Bandit Experiments (Quen2. 1ree2.5-1493 5-14B, 0.7, a 2).
Cumulitivv Meau Regret
Curnalative Mesu Regrel
40
30
20
20
10
0
50
100 Trial
150
200
50
100
150
200
(a) Epistemic Variance
(b) Total Variance
Figure 53: Cumulative Mean Regret for Bandit Experiments (Qwen2.5- 5-14B, = 0.7, a = 5).
54

---

Camulativv Mesu Regret
40
Cunualat Menia Regret
30
30
20
10
10
0
50
100 Trial
200
100
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 54: Cumulative Mean Regret for Bandit Experiments (Qwen2. 5-7B, p 0.5, 2).
Cumulative Meau Regrel
Cunnlative Meau Regret
40
30
20
20
10
0
50
100 Trial
150
200
50
100
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 55: Cumulative Mean Regret for Bandit Experiments (Qwen2. 5-7B, p 0.5, a 5).
Cumulative Mean Regrel
Cunnlativv Mean Regrel
40
15
30
10
20
0
50
100 Trial
150
200
50
100
150
200
(a) Epistemic Variance
(b) Total Variance
Figure 56: Cumulative Mean Regret for Bandit Experiments (Qwen2 5-7B, = 0.6, a = 2).
Comulstive Mean Regeet
Counulative Mean Regret
30
30
20
20
10
10
100
100
200
Trial
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 57: Cumulative Mean Regret for Bandit Experiments (Qwen2.5 5-7B, 0.6, a 5).
Cumulative Mesu Regret
Curnulative Mesu Regrel
40
30
30
20
-0
10
0
50
100 Trial
150
200
50
100
150
200
(a) Epistemic Variance
(b) Total Variance
Figure 58: Cumulative Mean Regret for Bandit Experiments (Qwen2.5.m 5-7B, = 0.7, a 2).
55

---

Cumulativve Mesa Regret
Cumulaat Menia Regret
30
30
20
10
10
0
50
200
100 Trial
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 59: Cumulative Mean Regret for Bandit Experiments (Qwen2.5- p 0.7, 5).
Comulativv Mean Regret
Connultivve Mean Regret
40
30
30
20
20
10
10
0
50
100
2010
100
150
200
Trial
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 60: Cumulative Mean Regret for Bandit Experiments (Llama-3. 1-8B, p 0.5, a 2).
Cumulative Meau Regrel
Cunulative Meau Regret
40
30
20
20
10
6
50
100 Trial
150
200
50
100
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 61: Cumulative Mean Regret for Bandit Experiments (Ll 1 1-8B, p 0.5, a 5).
Counllative Menn Regret
Counulative Mean Regrett
40
30
20
20
10
0
100
100
200
Trial
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 62: Cumulative Mean Regret for Bandit Experiments (Ll a.lama-3... 1.1...1.. 1-8B, 0.6, a 2).
Cumulative Meau Regret
30
Cunnlatiive Menu Regrel
40
20
30
-20
10
0
0
50
100 Trial
150
200
50
100
150
200
(a) Epistemic Variance
(b) Total Variance
Figure 63: Cumulative Mean Regret for Bandit Experiments (L) a.lama-3 1-8B, p = 0.6, a = 5).
56

---

Oumulat Mesia Regret
Cumulattiv Menia Regret
20
30
20
10
10
0
50
100 Trial
200
100
150
200
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 64: Cumulative Mean Regret for Bandit Experiments (Llama- 1-88, p 0.7, 2).
Counllative Mean Regret
30
Cumnuaativ Mean Regret
40
30
20
20
10
10
50
100
200
100
150
200
Trial
Trial
(a) Epistemic Variance
(b) Total Variance
Figure 65: Cumulative Mean Regret for Bandit Experiments (Llama-3.1 1..ma-3.11 1-8B, p 0.7, 5).
G.5
Question Answering
Datasets. In our experiments, leverage binary classification datasets including BooIQA [15].
HotpotQA [116], and PubMedQA [42] as well as multiclass classification dataset MMLU [31].
BoolQA is reading comprehension dataset that studies yes/no questions. HotporQA is dataset with Wikipedia- based questions that contain complex reasoning explanations for answers. PubMedQA biomedical question answering dataset collected from PubMed abstracts to answer research questions with yes/no//aybe. MMLU is a massive multitask test consisting of multiple-choice questions from various branches of knowledge. For the binary classification datasets, we preprocess them by extracting the "yes/no" questions, followed by formulating each sample in "Question:... Context:..." " format and mapping its labels into integers: ("no":0, "y"s":1)". For MMLU, each sample is formulated in "Question:... Choices:... format.
In-Context Out-of-Distribution Detection. We apply VUD question answering tasks. We first examine out-of-din distiiibibiiitoo (OOD) detection via area under the ROC curve (AUC) [32]. Our goal demonstrate that leveraging epistemic uncertainty from our decomposition yields higher OOD detection accuracy than directly utilising the total uncertainty. This enables practitioners to identify unreliable model predictions on unfamiliar inputs, improving the robustness and trustworthiness of deployed QA systems. In our main experiments, we leverage BoolQA [15], HotpotQA [1 [116], and PubMedQA [42] interchangeably of equivalent sample size as the in-distribution cnttof-ddisributtinn (OOD) datasets [67]. We formulate these datasets as binary classification tasks (yes/no). For our reference baseline, we extend the Deep Ensembles framework [39] to our OOD detection task by ensembling the output distributions of multiple different in-context example sets. For both methods, we leverage a training set size of [D] = 15 ICL samples and a test set size of |ai + Foool 120 for our ID and OOD samples and average our experimental results across 3 seeds. For our method, we generate |Z| = 20 perturbations by prompting the LLM to rephrase with relevant context from the test sample. For Deep Ensembles, we leverage 5 different in-context learning sets.
(ID) and
Before our discussion, note that OOD detection from an ICL perspective can be particularly challenging. Traditionally, OOD detection leverages the entire training set to train the model [30, 32]. However, in the ICL setting, we are limited by the context length and quality of the LLM. Another issue that persists guaranteeing that the QA datasets are semantically different enough where their distribution differs. Despite the difficulties, in Table 12, we observe that for our method, epistemic uncertainty (EU) yields higher AUC scores more ID/OOD settings than total uncertainty (TU), implying better OOD detection results via decomposition. When compared to Deep Ensembles,
57

---

Table 12: -iitr--Diitribttion Detection AUC scores on QA tasks. Higher AUC values for epistemic
uncertainty (EU) highlights the effectiveness of the uncertainty decomposition.
AUC (DEEP ENSEMBLES)
AUC (OURS)
ID/OOD
BOolQA
HOTPOTQA
PUBMEDQA
BooLQA
HOTPOTQA
PUBMEDQA
TU
0.04
0.604 .00
0.355 +.00
0.570
BOoLQA
EU
0.347+.001
0.619: + =0002
0.600
0.395 1.395-000
TU 0.677
0.684 :.0
0.712
0.754 002
HOTPOTQA
BU 0.659 2.00
0.638:
0.780
0.775
TU 0.6662.00
0.360
0.679:..04
0.382
PUBMEDQA
BU
0.606
0.329
0.471
0.4833..01
we notice that 1) the AUC scores for EU are considerably lower and 2) the AUC of the decomposed EU often underperforms when compared to its own TU.
In-Context Abstention. We provide the following example from MMLU-Moral where extract the highest decomposed aleatoric uncertainty sample (filtered out by AU) that is not present in the total uncertainty filtered samples (not filtered out by TU).
"Question: For which of these two scenarios does the main character (who uses l/me/my) do something clearly morally wrong, according to ordinary moral standards in the US as of 2020?
Scenario protected the catered food from my cold
Scenario 2 I was pole dancing during the erotic show.
Choices:
O: Wrong, Wrong
Wrong, Not wrong
2: Not wrong, Wrong
3: Not wrong, Not wrong"
To interpret the question, we find that the example of "Scenario protected the catered food from my cold" could be morally "Wrong" or "Not Wrong". ". If the character did not attend the event to protect the food, it is "Not Wrong". the character did attend but protected the food by staying further away from it, it could be morally "Wrong" since the individual could still pass the germs onto the food. This task is an example the broader class of selective classification problems [26], which has previously shown that AU can be more effective than total uncertainty thresholding metric for selective classification [103].
In Figures 66 and 67, we depict that across multiple thresholds, AU samples achieve mostly superior results for datasets BoolQA, HotpotQA, PubMedQA, and MMLU-CS.
0.031
0.93
Total (TU) Alnatorie Uneertsinty (Va)
Uncetailll
Alastoric Unentiiiity (Va)
0.92
0.85
0.01
0.85
0.00
0.89
0.85
0.80
0.87
030
0.31
0.51
1.00
0.25
Drigped)
Drigee)
(a) BoolQA
(b) HotpotQA
Figure 66: Effect of In-Context Abstention on Accuracy. BooIQA and HotpotQA Datasets.
58

---

0.91
Tutal Unoutaatsty (TU)
Total Uncurtastt (TU)
Alestotic Uneatiialy (Va)
Aleatatic Unattisialy (Va)
0.93
0.87
Accuary
Ascuracc
0.89
0.85
0.84
0.81
0.35 0,94
0.90
0.82
0,50
1.00
0.58
0.50
0.38
3.85
Alstention (Praction Doupped)
Drupped)
(a) PubMedQA
(b) MMLU-CS
Figure 67: Effect of In-Context Abstention on Accuracy. PubMedQA and MMLU-CS Datasets.
H
Example Prompts
H.1 Synthetic Toy
Synthetic Classification Experiments
x1 -1.75; x2
x1 -0.16; 0.57 <output>ocloutput>
-0. 21 <output>i<houtpur>
x1 0.4; x2
-0.05 05 <output>i<\outpurs
x1 0.2; x2 0.4 <output>
Synthetic Regression Experiments
-0.7 <output> 4.9 <loutput>
x -1.1 <output> 3.7 <\output>
= 4.8 <output> -1.6 <\output>
x 0.2 <output>
H.2 Bandits
Bandit Classification Experiments (LLM-UCB Algorithm)
action <revard>1<<reward>
action Creward>0</rewardd
action <reward>i<\reeenn action <reward>
59

---

Bandit Classification Experiments (Instruct Baseline)
</sytemm
You are a bandit algorithm in a room with 5 buttons labeled blue, green, red, yellow, purple. Each button is associated with a Bernoulli distribution with a fixed but unknown mean; the means for the buttons could different. For each button, when you press it, you will get reward that is sampled from the button's associated distribution.
have 200 time steps and, on each time step, you can choose any button and receive the reward. Your goal is to maximize the total reward over
time steps.
At each time step, will show you summary of your past choices and rewards. Then you must make the next choice, which must be exactly one of blue, green, red, yellow, purple. Let's think step by step make sure make good choice. You must provide your final answer within the tags <Answer>COLOR</Anseee / where COLOR is one of blue, green, red, yellow, purple. </user>>
So far you have played 7 times with your past choices and rewards summarized as follows:
blue button: pressed 3 times with average reward 0.67
green button: pressed 2 times with average reward 0.50
red button: pressed 0 times
yellow button: pressed 1 times with average reward 0.00
purple button:
pressed times with average reward 1.00
Which button will you choose next? Remember, YOU MUST provide your final answer within the tags <hnwer>>COLOR</Ansver where COLOR is one of blue, green, red, yellow, purple. Let's think step by step to make sure we make good choice. <lassistant|>
60

---

H.3
Question Answering
Downstream Prediction
You are given a set of in-context examples a new input. Your task is to predict the label of the new input.
Please carefully review the following examples and their labels inside <ottput><Dabllsk/outtttt tags:
Question: is marley from...
Context: when john senses...
<output>1<<output>
Question: are all the...
Context: following the unsuccessful...
<output>o<<outpuu>>
Now, predict the label for this new input:
Question: did the titans... Context: despite bertier's paralysis...
IMPORTANT: Output ONLY the label inside <output></output> tags. Do not add any explanation, text, formatting.
Your response must strictly follow this format:
Z Perturbations (Binary Classification)
Please rephrase the following:
Question: do the titans Context: while celebrating
While rephrasing the above, incorporate context from the following and make sure its intertwined/nteercomecce
Question: did zz top play
Context: 'sdoubleback" a song
Use the following format when rephrasing:
<rep> Question: (Rephrased Question]?
Context: (Rephrased Context). </rep>
61

---

Z Perturbations (Multiclass Classification)
Please rephrase the following:
Question: A scientist, using electrodes...
Choices:
0: Depolarization
1: Repolarization
2: Hyperpolarization
3: Resting potential
While rephrasing the above, you must incorporate context from the following and make sure it's teterrtinne//nteerconneeee
Question: During exercise, adrenaline secretion...
Choices:
0: increased plasma glucose.
1: increased plasma fatty acids.
2: increased plasma ACTH.
3: increased sympathetic nerve activity.
Use the following format when rephrasing:
<rep> Question:... Choices... </rep>
I Declarations
Use of Generative AI. The experimental data is collected from open-sourced L.MMs declared in the relevant experiment sections.
Broader Impact. This work aims to improve the reliability of LLMs through principled uncertainty quantification but may also amplify risks if used without safeguards for fairness and transparency.
62


Reviews:
[{"review_id": "6tuSNV8G2j", "reviewer": "Reviewer_aVVb", "summary": "This paper tackles the problem of uncertainty decomposition in large language models for the in-context learning (ICL) task, based on the recent hypothesis that performing ICL using large language models (LLMs) is akin to generating from the posterior predictive distribution. The authors provide a theoretical framework to try and tackle two main components required for uncertainty decomposition of this Bayesian posterior predictive: exchangeability and posterior sampling. The authors introduce a theoretical framework by introducing additional optimizable parameters that help tackle the first problem approximately and provide an optimizable upper bound on aleatoric uncertainty. The authors also present techniques to generate the necessary auxiliary parameters to improve the tightness of the bound to generate desirable properties such as well calibrated epistemic uncertainty. The authors also present experimental results on synthetic regression and classification datasets.", "questions": "- Can you comment on how we would improve the practical adoption of this approach from a perspective of test time scaling? Existing field of Bayesian deep learning has multiple approaches such as knowledge distillation [1], etc. to help, but wondering what would you suggest to improve the test time computation.\n\n- The paper would also benefit from comparison with additional approaches.  If we were to borrow from the literature of Bayesian Deep Learning again, there are techniques such as MC Dropout[2], BatchEnsembles [3], etc. I am also open to hearing from the authors on what other baselines they have in mind.\n\n- Can the authors comment on the difficulty on optimizing the new parameters introduced in this approach? It'll be great to also have a profiling view of the different steps involved per experiment and model. \n\n- Can the authors comment on how we might extend this approach to new datasets/tasks (such as MMLU)? \n\nReferences\n\n[1] Vadera, M., Jalaian, B. and Marlin, B., 2020, August. Generalized bayesian posterior expectation distillation for deep neural networks. In Conference on Uncertainty in Artificial Intelligence (pp. 719-728). PMLR.\n\n[2] Gal, Y. and Ghahramani, Z., 2016, June. Dropout as a Bayesian approximation: representing model uncertainty in deep learning. In Proceedings of the 33rd International Conference on International Conference on Machine Learning-Volume 48 (pp. 1050-1059).\n\n[3] Wen, Y., Tran, D., & Ba, J. BatchEnsemble: an Alternative Approach to Efficient Ensemble and Lifelong Learning. In International Conference on Learning Representations.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The paper provides a principled approach to uncertainty decomposition in ICL by introducing variational bounds that alleviate the intractable posterior sampling problem. I really liked the fact that the authors have built their approach \n\n- The paper is also well positioned and motivated given the interest in understanding the uncertainties arising from deploying LLMs to critical applications \n\n- The authors have done a good job describing their work in the context of existing literature. I also appreciate the fact that the authors have clearly mentioned the underlying assumptions as they relate to the interpretation of ICL as computing posterior predictive distribution.\n\n\nWeaknesses\n\n- I think the paper would greatly benefit if it were to include experiments from additional tasks (such as answering on MMLU dataset for instance). Having experimental results on synthetic datasets is useful for visualization and initial understanding, but the paper definitely needs additional experiments to support the application perspective. \n\n- While I appreciate the fact that the authors have clearly written down the assumptions under which the Bayesian framework can be attributed to ICL, this can be considered a potential shortcoming of the overall idea of Bayesian framework as an approximation to ICL. While this is not necessarily a critique of this existing paper, but it's a critique of the broader idea and as a result it does impact the significance of this work.\n\n- When it comes to practical application, this method would add additional complexity in terms of finding the appropriate $Z$, and performing multiple forward passes.", "was_revised": true, "final_justification": "The authors engaged thoughtfully in the rebuttal process and provided add'l experiment results that show benefits of their approach on existing benchmarks of MMLU-Moral and Computer science.", "rebuttal": "### Summary Response:\n\nWe sincerely thank the reviewers for their constructive feedback. We are encouraged that reviewers appreciated the ***novel*** (**bTbz**, **esK3**, **y4yc**) ***principled approach*** (**aVVb**) and ***sound theoretical framework*** (**s63H**) to the ***important and underexplored*** (**y4yc**) problem of uncertainty decomposition in in-context learning.\n\nOur work addresses the fundamental challenge of computing uncertainty decompositions on an intractable and implicitly defined Bayesian model of ICL in LLMs. Specifically:\n\n- We derive an upper bound to the aleatoric uncertainty (conditional entropy/variance) of the Bayesian model.\n- To our knowledge, we are the first achieve this without sampling from the Bayesian model posterior or assumptions on the form of the likelihood, via probing the model with auxiliary inputs.\n- We provide a clear statement of our assumptions and employ techniques such as permutation-based sampling and KL filtering to approximate exchangeability to ensure the validity of our method.\n- Whilst our method is particularly suited to LLMs, it can be applied to any model where there is access to the posterior predictive distribution.\n- We demonstrate both quantitative and qualitative performance of VUD across a range of synthetic datasets and real-world downstream tasks.\n\nIn response to reviewer suggestions, we have expanded our experimental analysis. New results include:\n\n- Table A (see **aVVb**): Additional experimental OOD tasks for natural language on the MMLU dataset\n- Table B (see **aVVb**): New baseline comparisons using Martingale posterior distributions\n- Table C (see **bTbz**): Analysis of the sample efficiency in the number of samples $Z$\n\nWe provide detailed responses to the reviewer comments and questions below.\n\n---\n\n### Individual Response:\n\n**Comments: MMLU**\n\n> Can the authors comment on how we might extend this approach to new datasets/tasks (such as MMLU)?\n> \n\nWe conduct an out-of-distribution experiment using Qwen/Qwen2.5-14B-Instruct between MMLU-CS computer science (\"college_computer_science\", \"high_school_computer_science”) and MMLU-M moral (”moral_scenarios”) questions from the MMLU dataset.\n\n**Table A: OOD on MMLU dataset**\n\n|  |  | Deep Ensembles (AUC) |  | Ours (AUC) |  |\n| --- | --- | --- | --- | --- | --- |\n| ID/OOD |  | MMLU-CS | MMLU-M | MMLU-CS | MMLU-M |\n| MMLU-CS | TU | - | **0.302 (.001)** | - | 0.337 (.000) |\n|  | EU | - | 0.178 (.000) | - | **0.521 (.001)** |\n| MMLU-M | TU | 0.376 (.000) | - | 0.397 (.000) | - |\n|  | EU | **0.391 (.001)** | - | **0.502 (.001)** | - |\n\nVUD outperforms Deep Ensembles considerably. We aim to include further MMLU subtopic comparisons in our paper.\n\n**Comments: Bayesian framework assumptions**\n\n> While I appreciate the fact that the authors have clearly written down the assumptions under which the Bayesian framework can be attributed to ICL, this can be considered a potential shortcoming of the overall idea of Bayesian framework as an approximation to ICL. While this is not necessarily a critique of this existing paper, but it's a critique of the broader idea and as a result it does impact the significance of this work.\n> \n\nIn this work, we choose to build upon the intuition of a Bayesian framework for ICL [1]. In particular, we explicitly state assumptions as a rigorous way to provide a solid foundation to promote further Bayesian methods in this domain. However, to encourage practical adoption of this method we propose techniques to promote exchangeability and bridge the gap between theory and practice. In this way, our contribution lies both in articulating a principled Bayesian theory of ICL and in developing practical methods to make that interpretation more applicable in real-world settings.\n\n**Comments: Test Time Scaling**\n\n> Can you comment on how we would improve the practical adoption of this approach from a perspective of test time scaling?\n> \n\nTo leverage test-time scaling [3], variational estimates for aleatoric uncertainty (steps 2-4 in Fig. 3) in can be computed in parallel across a range of auxiliary variables $Z_i$. This may allow to further optimisation of the variational lower bound, for instance by considering pairs of auxiliary variables $(Z_i, Z_j)$. An example of a greedy procedure is outlined below:\n\n1. **Parallel Evaluation**: compute $V_a (y^* |x^*, \\mathcal{D} ,Z_i)$ for multiple $Z_i$, where each $Z_i$ is a single datapoint.\n2. **Initial Filtering**: Discard candidates that fail to meet a KL-divergence threshold\n3. **Combinatorial Expansion**: Form combinations (e.g. ordered pairs $(Z_j, Z_k)$) from the remaining candidates and compute $V_a (y^* | x^* , \\mathcal{D} ,(Z_j,Z_k))$\n4. **Secondary Filtering**: Apply the KL condition again to filter out weak combinations.\n5. **Selection**: Choose the optimal set based on the resulting $V_a$ score\n\nThis simple scheme highlights how increased test-time compute can be used to improve the quality of the uncertainty estimates, a potential direction for future work.\n\n**Comments: Test Time Compute**\n\n> Existing field of Bayesian deep learning has multiple approaches such as knowledge distillation, etc. to help, but wondering what would you suggest to improve the test time computation.\n> \n\nWe agree that test-time scaling is an important consideration for practical adoption. One avenue to improve the test-time scaling of this method is that, when computing decompositions for a large set of predictions $X$ given the same in-content learning dataset $\\mathcal{D}$, it is possible to re-use the auxiliary variables $Z$ for different predictions. This can reduce the number of forward passes and in particular, for the binary classification scenario, reduce the number of forward passes by 33%.\n\n**Comments: Additional Baseline Comparisons**\n\n> The paper would also benefit from comparison with additional approaches. If we were to borrow from the literature of Bayesian Deep Learning again, there are techniques such as MC Dropout, BatchEnsembles, etc. I am also open to hearing from the authors on what other baselines they have in mind.\n> \n\nIn our method, the Bayesian model is obtained via de Finetti’s theorem and thus defined implicitly. Therefore, we do not know the form of $\\theta$, including the likelihood or prior so the ground truth aleatoric and epistemic uncertainties are unobservable. However, standard Bayesian DL methods assume $\\theta$  to be the LLM parameters so are not comparable baselines in this setting.\n\nNevertheless, if we assume a particular likelihood model, the Martingale posterior method [2] can be used to decompose total uncertainty (as it similarly uses de Finetti’s theorem). However, this approach is highly sensitive to the choice of likelihood model. Indeed, as we show in Table B, across various tasks and likelihood choices, the L2 divergence between the Martingale total uncertainty and empirically observed uncertainty remains significant, highlighting the difficulty of selecting a suitable likelihood model in practice.\n\n**Table B: L2 Divergence between uncertainty from Martingale Posterior Total Uncertainty and Empirically Observed Uncertainty**\n\n| Likelihood (Row); Model (Column) | Qwen7B | Qwen14B | Llama8B |\n| --- | --- | --- | --- |\n| Logistic Regression (Linear Features) | 1.023 | 0.804 | 2.158 |\n| Probit Regression (Linear Features) | 1.518 | 0.941 | 2.680 |\n| Logistic Regression (Quadratic Features) | 3.398 | 2.556 | 3.369 |\n| Probit Regression (Quadratic Features) | 3.319 | 2.441 | 2.429 |\n\n**Comments: Optimizing Parameters**\n\n> Can the authors comment on the difficulty on optimizing the new parameters introduced in this approach? It'll be great to also have a profiling view of the different steps involved per experiment and model.\n> \n\nThe key hyperparameters in the method are the KL-threshold $\\epsilon$ and the choice of auxiliary data $Z_i$.\n\nTo ensure that all the auxiliary variables aren’t filtered out by the KL-filtering, we rank the KL divergences $\\epsilon_i$ by size and only keep the $r$ auxiliary variables with the smallest KL divergence. This ensures that the method guarantees a decomposition whilst $r$ becomes a proxy for the level of approximation, which is left to the discretion of the user. See Appendix D.3 for further details.\n\nThe hyperparameter tuning for auxiliary variable choice is dependent on sampling method. We discuss 4 strategies for sampling $Z$ in the main text and the best performing methods are:\n\n- Perturbing about the predictive task $x^*$: For continuous features, we add Gaussian noise (std = 0.1 $\\times$ std of the covariate in $\\mathcal{D}$). For categorical inputs, we perturb uniformly across classes. For natural language, we use LLM-based paraphrasing. We find that 5-10 perturbations suffice in most cases; beyond that, improvements in the lower bound are minimal.\n- Repeating $x^*$. This does not require hyperparameter tuning\n\nOverall, these two methods do not require significant hyperparameter tuning and therefore, straightforward to implement.\n\nRegarding profiling, we agree this would provide a useful insight. We provide profiling for the logistic regression problem in the table below as an example.\n\n**Table: Profiling View**\n\n| **Function** | **LLM Calls Per Function** | **Num Function Calls** |\n| --- | --- | --- |\n| $p(y\\|x,D)$ | 5 | 1 |\n| $p(u\\|Z_i,D)$ | 5 | 5 |\n| $p(x\\|u,Z_i,D)$ | 10 | 5 |\n\n---\n\n[1] Xie, Sang Michael, et al. \"An explanation of in-context learning as implicit bayesian inference.\" *arXiv preprint arXiv:2111.02080* (2021).\n\n[2] Fong, Edwin, Chris Holmes, and Stephen G. Walker. \"Martingale posterior distributions.\" *Journal of the Royal Statistical Society Series B: Statistical Methodology* 85.5 (2023): 1357-1391.\n\n[3] Muennighoff, Niklas, et al. \"s1: Simple test-time scaling.\" *arXiv preprint arXiv:2501.19393* (2025)."}, {"review_id": "a4lod2pScV", "reviewer": "Reviewer_s63h", "summary": "The paper introduces a uncertainty decomposition framework for in-context learning (ICL) with LLMs, proposing a method to estimate and decompose predictive uncertainty into aleatoric and epistemic components without requiring posterior sampling over latent parameters. The authors optimize over auxiliary inputs to upper-bound aleatoric uncertainty and derive epistemic uncertainty as the residual. The method is theoretically grounded and designed to approximate Bayesian behavior under certain assumptions. Experimental results on synthetic datasets, bandit problems, and question-answering tasks show that the decomposed uncertainties align well with intuitive expectations.", "questions": "1. Since the framework allows auxiliary input and LLM inference is often costly, sample complexity becomes an important consideration. It would be helpful if the authors could discuss how the number of auxiliary samples affects the quality of uncertainty quantification. Specifically, how does performance scale with more auxiliary inputs, and is there any theoretical or empirical guidance on the minimal number of points needed to achieve reliable decomposition?\n\n\n2. In the bandit setting, what auxiliary information is used to support the uncertainty decomposition? I checked Appendix F.3 but couldn’t find a clear description. Additionally, while the paper presents relevant metrics and graphs, it remains unclear how the proposed uncertainty quantification framework is concretely integrated into the bandit decision-making process. Could the authors clarify how the uncertainty estimates are used to guide action selection in practice?", "limitations": "N/A", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength: This paper is well-written and presents a sound theoretical framework. The underlying problem is timely and important in the current AI landscape.\n\nWeakness: The major weakness is that, upon careful reading of the theoretical framework, it seems the role of the LLM is not tightly integrated into the core uncertainty quantification method. With access to auxiliary information for $U$ and $Z$, the proposed approach could be applied to a broad class of Bayesian models, not necessarily tied to in-context learning or LLMs. Another concern is that the assumed access to auxiliary data may be overly strong. In the experiments, the model is allowed to query points and responses near the test input $x$ and target $y$ which could enable many alternative methods to perform effective uncertainty quantification as well. The work would benefit from a more nuanced discussion of why the proposed method is non-trivial.", "was_revised": true, "final_justification": "The authors addressed the most part of my concern, therefore I decide to increase my rating.", "rebuttal": "### Summary Response\n\nWe sincerely thank the reviewers for their constructive feedback. We are encouraged that reviewers appreciated the ***novel*** (**bTbz**, **esK3**, **y4yc**) ***principled approach*** (**aVVb**) and ***sound theoretical framework*** (**s63H**) to the ***important and underexplored*** (**y4yc**) problem of uncertainty decomposition in in-context learning.\n\nOur work addresses the fundamental challenge of computing uncertainty decompositions on an intractable and implicitly defined Bayesian model of ICL in LLMs. Specifically:\n\n- We derive an upper bound to the aleatoric uncertainty (conditional entropy/variance) of the Bayesian model.\n- To our knowledge, we are the first achieve this without sampling from the Bayesian model posterior or assumptions on the form of the likelihood, via probing the model with auxiliary inputs.\n- We provide a clear statement of our assumptions and employ techniques such as permutation-based sampling and KL filtering to approximate exchangeability to ensure the validity of our method.\n- Whilst our method is particularly suited to LLMs, it can be applied to any model where there is access to the posterior predictive distribution.\n- We demonstrate both quantitative and qualitative performance of VUD across a range of synthetic datasets and real-world downstream tasks.\n\nIn response to reviewer suggestions, we have expanded our experimental analysis. New results include:\n\n- Table A (see **aVVb**): Additional experimental OOD tasks for natural language on the MMLU dataset\n- Table B (see **aVVb**): New baseline comparisons using Martingale posterior distributions\n- Table C (see **bTbz**): Analysis of the sample efficiency in the number of samples $Z$\n\nWe provide detailed responses to individual reviewer comments and questions below.\n\n---\n\n### Individual Response\n\n**Comments: Significance of Method**\n\n> The major weakness is that, upon careful reading of the theoretical framework, it seems the role of the LLM is not tightly integrated into the core uncertainty quantification method. With access to auxiliary information for $U$ and $Z$, the proposed approach could be applied to a broad class of Bayesian models, not necessarily tied to in-context learning or LLMs.\n\n> In the experiments, the model is allowed to query points and responses near the test input $x$ and target $y$ which could enable many alternative methods to perform effective uncertainty quantification as well. The work would benefit from a more nuanced discussion of why the proposed method is non-trivial.\n\n\nWe agree that the proposed method is general and can be applied to a broad class of Bayesian models and indeed, in Appendix B, we theoretically verify the validity of our approach on Bayesian linear regression and Gaussian processes. However, our method is motivated in-context learning in LLMs, where the Bayesian model is implicitly defined and we **do not have access to** the likelihood $p(x|\\theta )$ or posterior $p(\\theta|\\mathcal{D})$ or samples of $\\theta$, but we **do have access to** the posterior predictive distribution $p(y^* |x^* , \\mathcal{D} )$. To our knowledge, we are the only method to provide a bound on the conditional entropy of an implicitly defined Bayesian model given these restrictions.\n\nFurthermore, we tailor our approach to LLMs by using practical techniques to encourage exchangeability in model generations which are grounded by clear theoretical motivations.\n\nFinally, the most popular LLMs such as ChatGPT, Claude and Gemini are proprietary models. These models are closed-source and in a black-box setting which prevents practitioners from accessing their weights to apply traditional Bayesian deep learning techniques. However, our ICL based UQ method provides a mechanism to elicit and quantify  Bayesian uncertainty, since these black-box models like ChatGPT and Claude provide access to output logits and the query of points near the test input, allowing our method to be applied to a broad class of proprietary models.\n\n**Comments: Sample Complexity of Auxiliary Variables**\n\n> How does performance scale with more auxiliary inputs, and is there any theoretical or empirical guidance on the minimal number of points needed to achieve reliable decomposition?\n> \n\nIn our experiments, we use up to 15 auxiliary variables, if not fewer, for more computationally intensive tasks such as regression. We find empirically that 5-10 auxiliary variables is sufficient for a reliable decomposition with diminishing results beyond that. We include an example of bounds per number of auxiliary examples in Table C below.\n\n**Table C: Sample Complexity (Logistic Regression)**\n\n| Row ($x$ value); Col (num $Z$ ) | 5 | 10 | 15 | 20 | 30 |\n| --- | --- | --- | --- | --- | --- |\n| -10 | 0.177 (0.033) | **0.158 (0.024)** | 0.164 (0.002) | 0.201 (0.021) | 0.209 (0.000) |\n| -5 | 0.227 (0.033) | **0.210 (0.034)** | 0.249 (0.001) | 0.250 (0.001) | 0.252 (0.002) |\n| 0 | 0.838 (0.020) | **0.826 (0.016)** | 0.849 (0.010) | 0.860 (0.008) | 0.868 (0.003) |\n| 5 | 0.235 (0.016) | 0.224 (0.018) | 0.218 (0.018) | 0.212 (0.016) | **0.208 (0.014)** |\n| 10 | 0.173 (0.019) | **0.164 (0.019)** | 0.173 (0.018) | 0.182 (0.007) | 0.185 (0.001) |\n\nDeveloping theoretical guarantees for sample complexity is challenging without further assumptions on the particular prediction tasks as we don’t have access to the ground truth or Bayesian model. Nevertheless, under certain assumptions on the Bayesian model we prove for the multi-armed bandit setting that the tightness of the bound is $\\mathcal{O}(n^{-1})$ where $n^{-1}$ is the number of auxiliary examples (via an application of the Bernstein-von-Mises theorem). Please find the following for the proof for the Gaussian bandit setting:\n\nIn Gaussian bandit setting, assume independence of rewards between arms. Furthermore, for an arm $i$, we assume the following model:\n- prior: $p(\\theta_i) = \\mathcal{N}(0, \\sigma_0^2)$\n- likelihood: $p(r_i \\mid \\theta_i) = \\mathcal{N}(\\theta_i, \\sigma^2)$\n\nThen for $k_t$ observations of rewards from arm $i$, we have:\n- total variance: $\\sigma^2 + \\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t}{\\sigma^2} \\right]^{-1}$\n- true aleatoric variance: $\\sigma^2$\n- true epistemic variance = total - aleatoric: $\\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t}{\\sigma^2} \\right]^{-1}$\n\nHowever, with the VUD method, suppose we have $n$ further auxiliary observations, then, the corresponding variational estimate of aleatoric uncertainty is: $\\sigma^2 + \\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t + n}{\\sigma^2} \\right]^{-1}$.\n\nThis gives the gap between the variational estimate and the exact aleatoric variance as:\n$$\\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t + n}{\\sigma^2} \\right]^{-1} = \\mathcal{O}\\left(\\frac{1}{k_t + n}\\right) = \\mathcal{O}\\left(\\frac{1}{n}\\right).$$\n\nNote: For non-Gaussian bandits we can apply Bernstein-von Mises to obtain a similar bound.\n\n**Comments: Bandit Clarifications**\n\n> In the bandit setting, what auxiliary information is used to support the uncertainty decomposition? …  Could the authors clarify how the uncertainty estimates are used to guide action selection in practice?\n> \n\nFor the “Buttons” multi-armed bandit task, the task is to predict the reward $r_t(a)$ for an action $a$ and prior action reward pairs $ \\mathcal{D}_ t = \\{ (a_i , r_i ) \\}_{ i = 1 }^{ t - 1 }$. In this setting $x$ is $a$ and $y$ is $r_t(a)$. For computational efficiency, we use the repeated $Z$ strategy: so for a particular action $a$, the auxiliary data would be $Z=a$ (so there is only 1 auxiliary data point per action). We have to do uncertainty decompositions for each action in the action space $\\mathcal{A}$.\n\nOnce we have uncertainty estimates $U_t(a)$ for each action and expected reward estimates $Q_t (a) \\approx \\mathbb{E} [r_t |a, \\mathcal{D}_ t ]$, we compute the UCB score for each action, $\\mathrm{UCB}_ t(a)=Q_t (a)+c U_t (a)$ where $c$ is the exploration constant (a hyperparameter we can choose). Then, the next action is the action that has the minimal UCB score $a_t =\\mathrm{ argmin } _{ a \\in \\mathcal{A} } \\{ \\mathrm{ UCB } _ t (a) \\} $. (In the situation that multiple actions have the same UCB score, we choose uniformly from these minimising actions)."}, {"review_id": "nazs1onhHK", "reviewer": "Reviewer_bTbz", "summary": "The paper introduces a uncertainty decomposition method, variational uncertainty decomposition (VUD), that leverages in-context examples for uncertainty decomposition of LLMs. Motivated by a Bayesian interpretation of ICL, the authors propose a method that estimates aleatoric and epistemic uncertainties without requiring explicit posterior sampling over latent parameters. Instead, the approach uses optimisable auxiliary inputs (Z) to provide a variational upper bound on aleatoric uncertainty, with the residual treated as epistemic uncertainty.", "questions": "N/A", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n- The proposed method is very interesting and novel. The experiment on QA tasks are especially interesting to me. However, I am not an expert in math/bandit and cannot justify the validity of the proposed method. But the experiement results in Table 2 looks very promising.\n\n\n\nWeakness:\n\n- On the definition of in-distribution and out-of-distribution examples in QA experiments. These QA datasets are mainly about world knowledge. However, existing LLMs undergo large-scale pre-training and already store sufficient knowledge for most of these questions. In such a context, how you define out-of-distribution? It is unclear how in-context examples can help train these models or help them acquire additional knowledge such that some testing questions are in-distribution with respect to these in-context examples. The only in-distribution knowledge is the format, i.e., the model needs to output 0 or 1 to represent no or yes.\n- How Deep Ensembles is implemented in the QA setting? From my understanding, each set of in-context examples will give us a output distribution of the model, and by taking the average of that distribution over all sets of in-context examples, we can compute the entropy of the mean; then, we compute the mean entropy of each individual output distribution, which will be the aleatoric uncertainty; by subtract the aleatoric uncertainty from the entropy of mean, we get the epistemic uncertainty. Is that the implementation of Deep Ensembles?\n- Another concerns is on the decomposition of uncertainty and the source of aleatoric uncertainty. For these QA tasks, what is the source of aleatoric uncertainty? The only explanation of aleatoric uncertainty in these QA tasks, from my perspective, is the question ambiguity. But does there really exist ambigous questions in the test set, since these datasets should be considered as high-quality datasets and should be clear enough? If not, what is the source of the two types of uncertainty?\n- Finally, is in-context learning-based UQ method still useful for QA tasks? To the best of my knowledge, existing LLMs after SFT/post-training (not the pre-trained ones) can already follow human instructions and does not need in-context examples to answer such knowledge-intensive tasks. Even sometimes we provide them, these examples are only used for format clarification, instead of train these LLMs or equip them with additional world knowledge. Will the proposed methods be too \"toy\" for the actual applications of existing LLMs?\n\nAgain, the reviewer find this paper interesting and impressive. Once the other reviewers that are more experienced in the math deriviations can help verify the correctness and validity of this paper, I would like to see this paper being accpected.", "was_revised": true, "final_justification": "This paper already undergoes sufficient discussion. After actively engaging in the discussion with the authors and reading the opinion from other reviewer, I am pretty confidence that my judgement is correct and recommend the acceptance of this paper.", "rebuttal": "### Summary Response\n\nWe sincerely thank the reviewers for their constructive feedback. We are encouraged that reviewers appreciated the ***novel*** (**bTbz**, **esK3**, **y4yc**) ***principled approach*** (**aVVb**) and ***sound theoretical framework*** (**s63H**) to the ***important and underexplored*** (**y4yc**) problem of uncertainty decomposition in in-context learning.\n\nOur work addresses the fundamental challenge of computing uncertainty decompositions on an intractable and implicitly defined Bayesian model of ICL in LLMs. Specifically:\n\n- We derive an upper bound to the aleatoric uncertainty (conditional entropy/variance) of the Bayesian model.\n- To our knowledge, we are the first achieve this without sampling from the Bayesian model posterior or assumptions on the form of the likelihood, via probing the model with auxiliary inputs.\n- We provide a clear statement of our assumptions and employ techniques such as permutation-based sampling and KL filtering to approximate exchangeability to ensure the validity of our method.\n- Whilst our method is particularly suited to LLMs, it can be applied to any model where there is access to the posterior predictive distribution.\n- We demonstrate both quantitative and qualitative performance of VUD across a range of synthetic datasets and real-world downstream tasks.\n\nIn response to reviewer suggestions, we have expanded our experimental analysis. New results include:\n\n- Table A (see **aVVb**): Additional experimental OOD tasks for natural language on the MMLU dataset\n- Table B (see **aVVb**): New baseline comparisons using Martingale posterior distributions\n- Table C (see **bTbz**): Analysis of the sample efficiency in the number of samples $Z$\n\nWe provide detailed responses to individual reviewer comments and questions below.\n\n---\n\n### Individual Response\n\n**Comments: LLM’s Pretrained Knowledge and Motivation behind ICL based UQ**\n\n> On the definition of in-distribution and out-of-distribution examples in QA experiments. These QA datasets are mainly about world knowledge. However, existing LLMs undergo large-scale pre-training and already store sufficient knowledge for most of these questions. In such a context, how you define out-of-distribution? It is unclear how in-context examples can help train these models or help them acquire additional knowledge such that some testing questions are in-distribution with respect to these in-context examples. The only in-distribution knowledge is the format, i.e., the model needs to output 0 or 1 to represent no or yes.\n\n> Finally, is in-context learning-based UQ method still useful for QA tasks? To the best of my knowledge, existing LLMs after SFT/post-training (not the pre-trained ones) can already follow human instructions and does not need in-context examples to answer such knowledge-intensive tasks. Even sometimes we provide them, these examples are only used for format clarification, instead of train these LLMs or equip them with additional world knowledge. Will the proposed methods be too \"toy\" for the actual applications of existing LLMs?\n\n\nThe following outlines the motivation of ICL based UQ:\n\n- Existing LLMs undergo large-scale pre-training. However, there exists private, sensitive, inaccessible data such as electronic health records.\n- Serving the top performing open-sourced LLMs [1] have incredibly high compute requirements, necessitating at least 8 80GB A100 GPUs. Likewise for LLM-fine tuning. For cheaper LLMs, they often have much poorer performance [2] (~65% on MMLU for Llama3-8B). Thus, ICL alleviates users with computational constraints.\n- The most popular LLMs such as ChatGPT, Claude and Gemini are proprietary models. These models are closed-source and in a black-box setting which prevents practitioners from accessing their weights to apply traditional Bayesian deep learning techniques. Our ICL based UQ method provides a mechanism to elicit and quantify uncertainty since these proprietary models like ChatGPT and Claude provide access to output logits.\n- We find that formatting the output from LLMs is crucial. Even if the model knows the answer, it needs to express it in the expected form: \"yes\" vs. \"no\", \"0\" vs. \"1\", \"The answer is: True\" vs. \"False\". In our QA experiments, the model is expected to output 0/1 labels. If not prompted clearly, it may hallucinate and output varying text, leading to inconsistent output logits.\n\n**Comments: Deep Ensembles Implementation**\n\n> How Deep Ensembles is implemented in the QA setting? From my understanding, each set of in-context examples will give us a output distribution of the model, and by taking the average of that distribution over all sets of in-context examples, we can compute the entropy of the mean; then, we compute the mean entropy of each individual output distribution, which will be the aleatoric uncertainty; by subtract the aleatoric uncertainty from the entropy of mean, we get the epistemic uncertainty. Is that the implementation of Deep Ensembles?\n> \n\nYes, this interpretation is correct. To further clarify, instead of training multiple neural networks (traditional Deep Ensembles), we: \n\n1. Sample different in-context example sets ($D_k$) for each ensemble member. \n\n- Each $D_k$ contains randomly chosen few-shot examples from the training data.\n\n2. Query the same LLM with each $D_k$ to get predictions $p_k(y|x, D_k)$. \n\n- The LLM generates a probability distribution over labels for each prompt.\n\n3. Aggregate predictions i.e. taking the mean across $D_k$ and computing the entropy (total uncertainty) as well as the mean entropy of each ensemble distribution (aleatoric uncertainty) to estimate epistemic uncertainty.\n\n- **Total Uncertainty (TU)**: Entropy of the mean prediction $\\left(H\\left(\\frac{1}{K}\\sum_k p_k\\right)\\right)$.\n- **Aleatoric Uncertainty (Va)**: Mean entropy of individual predictions $\\left(\\frac{1}{K}\\sum_k=H(p_k)\\right)$\n- **Epistemic Uncertainty (Ve)**: $TU - Va$ (disagreement due to $D_k$ variability).\n\n**Comments: QA Task Source of Aleatoric Uncertainty**\n\n> Another concerns is on the decomposition of uncertainty and the source of aleatoric uncertainty. For these QA tasks, what is the source of aleatoric uncertainty? The only explanation of aleatoric uncertainty in these QA tasks, from my perspective, is the question ambiguity. But does there really exist ambiguous questions in the test set, since these datasets should be considered as high-quality datasets and should be clear enough? If not, what is the source of the two types of uncertainty?\n\nHigh-quality datasets can still contain intrinsic aleatoric uncertainty. As quoted from BoolQ [3] “the authors labelled 110 randomly chosen examples. If there was a disagreement, the authors conferred and selected a single answer by mutual agreement. We call the resulting labels “gold-standard” labels. On the 110 selected examples, the answer annotations reached 90% accuracy compared to the gold-standard labels”. Similarly, the same can be said about HotpotQA [4] where there is a **~**13 EM points gap between the Human UB and Human gap in Table 8, reflecting aleatoric uncertainty in the questions. Lastly, PubMedQA [5] includes \"maybe\" answers to model uncertainty that accounts for 11% of the PQA-L(abeled) dataset.\n\n---\n\n[1] Chiang, Wei-Lin, et al. \"Chatbot arena: An open platform for evaluating llms by human preference.\" *Forty-first International Conference on Machine Learning*. 2024.\n\n[2] Yang, An, et al. \"Qwen3 technical report.\" arXiv preprint arXiv:2505.09388 (2025).\n\n[3] Clark, Christopher, et al. \"Boolq: Exploring the surprising difficulty of natural yes/no questions.\" *arXiv preprint arXiv:1905.10044* (2019).\n\n[4] Yang, Zhilin, et al. \"HotpotQA: A dataset for diverse, explainable multi-hop question answering.\" *arXiv preprint arXiv:1809.09600* (2018).\n\n[5] Jin, Qiao, et al. \"Pubmedqa: A dataset for biomedical research question answering.\" *arXiv preprint arXiv:1909.06146* (2019)."}, {"review_id": "sYCEqBf79k", "reviewer": "Reviewer_y4yc", "summary": "This paper proposes Variational Uncertainty Decomposition (VUD), a method for separating aleatoric and epistemic uncertainty in in-context learning (ICL) with large language models, where model parameters are inaccessible. The approach uses a variational upper bound on aleatoric uncertainty via optimizable auxiliary inputs, and defines epistemic uncertainty as the residual from total predictive uncertainty. VUD operates without needing posterior sampling, bootstrapping, or fine-tuning, making it scalable and compatible with black-box models. The method is validated through theoretical analysis and empirical results across regression, classification, decision-making, and out-of-distribution detection tasks.", "questions": "1. Clarify the practical implementation and optimization of auxiliary inputs\n\nThe proposed method relies on optimizing auxiliary inputs $Z$ to estimate a variational upper bound on aleatoric uncertainty. Could the authors provide more detail on how this optimization is performed in practice? Specifically, how are these inputs initialized, what constraints or regularization are applied, and how sensitive is the method to their configuration? Clearer documentation of this component would improve reproducibility and help assess robustness.\n\n2. Explain the assumptions and limitations of approximate exchangeability in prompting\n\nThe method assumes approximate exchangeability in in-context examples to justify a Bayesian interpretation of ICL. Could the authors clarify how this assumption holds in practical prompting setups, especially when support examples are ordered or semantically diverse? It would be helpful to understand under what conditions this assumption breaks down, and how the method behaves in such cases.\n\n3. Provide empirical validation of the variational upper bound’s tightness\n\nTheoretically, the method guarantees a variational upper bound on aleatoric uncertainty, but the paper lacks a clear empirical assessment of how tight or accurate this bound is in practice. Could the authors include a quantitative analysis or visualizations to show how well the estimated uncertainty components match ground-truth uncertainty (when available) or known benchmarks?\n\n4. Discuss generalizability across model scales and prompt configurations\n\nThe experiments focus on a fixed set of tasks and models. Could the authors comment on how QUD performs across LLMs of different sizes or under prompt variations (e.g., number of in-context examples, style of formatting)? Insight into the method’s generalizability would help evaluate its practical utility and deployment readiness.\n\n5. Expand on downstream value and decision-making use cases\n\nWhile the method is evaluated on bandit decision-making and out-of-distribution detection, the practical impact of decomposing uncertainty is not deeply explored. Can the authors elaborate on how this decomposition improves real-world outcomes, such as more calibrated decisions or better sample efficiency? A more application-driven perspective would strengthen the significance and practical relevance of the contribution.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. The paper proposes a new framework, Variational Uncertainty Decomposition (VUD), that enables decomposition of total predictive uncertainty into aleatoric and epistemic components in in-context learning (ICL) scenarios, where traditional Bayesian inference is infeasible due to the black-box nature of large language models (LLMs).\n\n2. A key innovation is introducing optimizable auxiliary inputs (Z) to estimate an upper bound on aleatoric uncertainty through a variational formulation. This bypasses the need for access to latent parameter posteriors and enables uncertainty quantification directly from black-box LLMs.\n\n3. The paper defines epistemic uncertainty as the difference between total predictive uncertainty and the variational estimate of aleatoric uncertainty, offering a principled and interpretable decomposition without requiring multiple model runs or explicit posterior sampling.\n\n4. To support the Bayesian interpretation of ICL, the authors introduce prompting and optimization techniques that approximate exchangeability, a key assumption for Bayesian analysis. This makes the method more scalable and applicable in real-world LLM settings without modifying model internals.\n\n5. The method is evaluated across regression, classification, bandit decision-making, and OOD detection tasks, with both synthetic and real-world datasets. The paper also provides theoretical guarantees on the variational upper bound and the validity of the decomposition under assumptions, combining empirical rigor with formal soundness.\n\nWeaknesses:\n\n1. The proposed method is theoretically grounded and empirically validated across diverse tasks. However, the work relies on assumptions such as approximate exchangeability and the optimization of auxiliary inputs, which may not generalize well to all prompting formats or model architectures. The paper does not thoroughly analyze the sensitivity or robustness of the method with respect to prompt length, model size, or task type. Also, while theoretical guarantees are provided, empirical validation of these guarantees (e.g., tightness of the variational bound) is limited.\n\n2. The paper is technically sound but occasionally dense, particularly in its presentation of the variational objective and its connection to Bayesian foundations. The role of auxiliary input $Z$ and the practical steps for its optimization are central but underexplained for a general audience. A clearer explanation of the prompting strategy used to achieve approximate exchangeability, along with visual or intuitive illustrations, would significantly improve accessibility.\n\n3. The work addresses an important and underexplored problem: how to estimate and decompose uncertainty in in-context learning with large language models. However, its impact is somewhat constrained by a lack of practical analysis on how this uncertainty decomposition improves downstream decision-making beyond synthetic settings. Broader applications and more user-centric evaluation (e.g., real-world deployment scenarios) would enhance the significance of the contribution.\n\n4. The idea of decomposing uncertainty in ICL via a variational approach is novel, especially in black-box LLM settings where traditional Bayesian methods are inapplicable. However, the use of variational bounds and auxiliary optimization is conceptually inspired by prior work in approximate Bayesian inference. The novelty lies primarily in adapting these tools to the ICL setting rather than introducing entirely new inference techniques.", "was_revised": true, "final_justification": "The authors provided thoughtful responses that clarify several aspects of their method, including practical strategies for optimizing auxiliary inputs, theoretical justification for approximate exchangeability, and promising downstream applications. However, their explanation lacks key details that limit the overall strength of the rebuttal. Specifically, there is insufficient discussion of regularization and sensitivity in the optimization process, no empirical validation or visual analysis of the variational bound’s tightness, and a lack of systematic evaluation across model sizes and prompt styles. While the theoretical contributions are sound and the potential impact is clear, the absence of concrete empirical support and reproducibility guidance tempers the confidence in the method's robustness and generalizability. Given these considerations, a slight score increase (e.g., to a weak accept) could be justified if the missing components can be reasonably addressed in a camera-ready revision. However, I believe the current weaknesses are significant enough that I will stay with my original borderline accept score.", "rebuttal": "### Summary Response\n\nWe sincerely thank the reviewers for their constructive feedback. We are encouraged that reviewers appreciated the ***novel*** (**bTbz**, **esK3**, **y4yc**) ***principled approach*** (**aVVb**) and ***sound theoretical framework*** (**s63H**) to the ***important and underexplored*** (**y4yc**) problem of uncertainty decomposition in in-context learning.\n\nOur work addresses the fundamental challenge of computing uncertainty decompositions on an intractable and implicitly defined Bayesian model of ICL in LLMs. Specifically:\n\n- We derive an upper bound to the aleatoric uncertainty (conditional entropy/variance) of the Bayesian model.\n- To our knowledge, we are the first achieve this without sampling from the Bayesian model posterior or assumptions on the form of the likelihood, via probing the model with auxiliary inputs.\n- We provide a clear statement of our assumptions and employ techniques such as permutation-based sampling and KL filtering to approximate exchangeability to ensure the validity of our method.\n- Whilst our method is particularly suited to LLMs, it can be applied to any model where there is access to the posterior predictive distribution (p.p.d.).\n- We show quantitative and qualitative performance of VUD across a range of synthetic datasets and real-world downstream tasks.\n\nIn response to reviewer suggestions, we have expanded our experimental analysis. New results include:\n\n- Table A (see **aVVb**): Additional experimental OOD tasks for natural language on the MMLU dataset\n- Table B (see **aVVb**): New baseline comparisons using Martingale posterior distributions\n- Table C (see **bTbz**): Analysis of the sample efficiency in the number of samples $Z$\n\n---\n### Individual Response\n\n**Question 1**\n\nWe investigate 4 optimisation methods in the paper and empirically observe that choosing $Z$ close to or equal to $x^*$ optimises this expression. This is theoretically motivated via the proof of Theorem 3.1 and 3.2 in Appendix A, and analysis on tractable Bayesian models in Appendix B.\n\nIn practice we perform one of two strategies for choosing $Z$:\n\n- Perturbing about the predictive task $ x^* $ : For continuous features, we add Gaussian noise (std = 0.1 × std of the covariate in $\\mathcal{D} $ and mean= $ x^* $). For categorical inputs, we perturb uniformly across classes. For natural language, we use LLM-based paraphrasing. We find that 5–10 perturbations suffice in most cases; beyond that, improvements are minimal.\n- Repeating $x^*$. This does not require hyperparameter tuning\n\nWe also investigated Bayesian optimisation (BO) but found that the number of auxiliary variables required were significantly higher than the perturbing or repeating. Furthermore, the optimal $Z$ for the BO process almost always was a small perturbation from $x^*$, motivating the simpler approaches.\n\nTo choose the threshold for KL filtering we set $\\epsilon$ by ranking the KL divergences $\\epsilon_i$ for each auxiliary datapoint $Z_i$ and choose the $r$-th smallest divergence. Therefore, we guarantee  a decomposition and $r$ becomes a proxy for the level of approximation, which is left to the discretion of the user. We found this allowed for a more robust method because the scale of KL divergences can vary between different tasks. \n\n**Question 2**\n\nBandit problems are an example of a dataset where the in-context examples have an implicit ordering and the reward variables are not exchangeable (as later rewards have a higher mean). However, since the reward model $r_t(a)$ is stationary, the reward distribution when we restrict to the rewards with action $a$ is exchangeable. This is sufficient for the VUD method to still be applicable for ordered examples. This suggests that the theoretical assumptions for VUD can be loosened, which we would like to further investigate in future work (e.g. non-stationary reward models).\n\nIn our natural language experiments, we do not directly measure the semantic diversity [1] and compare the effect of the diversity on the exchangeability of the data. However, we hope to investigate this further in future work, in particular, integrating VUD with methods with methods that measure semantic entropy [2] to decompose uncertainty on predictions with natural language output.\n\n\n**Question 3**\n\nThe true aleatoric uncertainty depends on the (unknown) Bayesian prior and likelihood governing the LLM, which are inaccessible in practice. As such, exact ground-truth comparisons are not possible. \n\nHowever, we can approximate this by comparing against the uncertainty in the known data-generating process in our synthetic and bandit settings. While this serves only as a heuristic, since the LLM’s internal modelling may diverge from the true generative model, it still offers useful insights. For instance, in homoscedastic generative settings, we generally observe that VUD estimates reflect this structure.\n\nFurthermore, under certain assumptions on the Bayesian model we prove the multi-armed bandit (MAB) setting that the tightness of the bound is $\\mathcal{O}(n^{-1})$ where $n^{-1}$ is the number of auxiliary examples. Please find the following for the proof:\n\nIn Gaussian bandit setting, assume independence of rewards between arms. Furthermore, for an arm $i$, we assume the following model:\n- prior: $p(\\theta_i) = \\mathcal{N}(0, \\sigma_0^2)$\n- likelihood: $p(r_i \\mid \\theta_i) = \\mathcal{N}(\\theta_i, \\sigma^2)$\n\nThen for $k_t$ observations of rewards from arm $i$, we have:\n- total variance: $\\sigma^2 + \\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t}{\\sigma^2} \\right]^{-1}$\n- true aleatoric variance: $\\sigma^2$\n- true epistemic variance = total - aleatoric: $\\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t}{\\sigma^2} \\right]^{-1}$\n\nHowever, with the VUD method, suppose we have $n$ further auxiliary observations, then, the corresponding variational estimate of aleatoric uncertainty is: $\\sigma^2 + \\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t + n}{\\sigma^2} \\right]^{-1}$.\n\nThis gives the gap between the variational estimate and the exact aleatoric variance as:\n$$\\left[\\frac{1}{\\sigma_0^2} + \\frac{k_t + n}{\\sigma^2} \\right]^{-1} = \\mathcal{O}\\left(\\frac{1}{k_t + n}\\right) = \\mathcal{O}\\left(\\frac{1}{n}\\right).$$\n\nNote: For non-Gaussian bandits we can apply Bernstein-von Mises to obtain a similar bound.\n\n**Question 4**\n\nWe observe that increasing the number of in-context examples leads to a consistent trend: **epistemic uncertainty decreases**, while **aleatoric uncertainty stabilises**, as shown in Figures 7, 9c, 15, and 16 (logistic regression and bandit tasks), which aligns with the intuition for these quantities.\n\nWe adopt the standard prompt format used in prior work on Bayesian views of ICL [7]. However, we note the prompt wording can have a conditioning effect on the p.p.d., effectively changing the prior of the LLM, similar to behaviour observed in LLM processes [6]. We believe a systematic analysis of the prompt style effect on model prior is beyond the scope of this project but we agree this is a valuable direction for future work.\n\nDue to computational constraints, we focus on experiments with 7B, 8B and 14B parameter models but we observe that VUD provides sensible uncertainty decompositions across larger model sizes as well. However, if the p.p.d. is poor (e.g. with very small models), the total uncertainty (which is a function of the p.p.d.) can be unreliable, leading to less meaningful decompositions.\n\n**Question 5**\n\nDecomposing uncertainty in LLMs provides greater transparency: this distinction helps practitioners decide when to trust model outputs or seek more information.\n\nPractically, VUD enables more informed data acquisition in settings like Bayesian experimental design (BED), active learning [3] and bandit problems. In these applications, maximizing mutual information (MI) guides which unlabeled samples to label next. VUD offers a tractable lower bound on MI, making it a practical tool for sample-efficient learning. In particular, VUD allows for applications of BED, active learning and bandits in natural language settings [4].\n\nWe agree that applying VUD to real-world scenarios is an important next step, and we are pursuing such directions. One such real-world deployment is tabular data. There exists an abundance of tabular electronic health records and financial statements where rows can be serialized into natural language [5]. Thus VUD can be used for OOD detection for ICL where there are limitations on LLM-finetuning budget. Our goal is to allow healthcare practitioners (e.g. Abridge/OpenEvidence) and financial analysts (e.g. BloombergGPT) to identify unreliable model predictions on unfamiliar inputs, improving the trustworthiness of deployed QA systems.\n\n**Other Comments: Visual Illustrations**\n\n> A clearer explanation of the prompting strategy used to achieve approximate exchangeability, along with visual or intuitive illustrations, would significantly improve accessibility.\n\nWe will include further visual illustrations of the method and practical advice for optimisation to improve the accessibility for general audience if accepted.\n\n\n---\n\n[1] Hoffman et al. \"Semantic diversity: A measure of semantic ambiguity based on variability in the contextual usage of words.\" (2013).\n\n[2] Kuhn et al. \"Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation.\" (2023).\n\n[3] Rainforth et al. \"Modern Bayesian experimental design.\" (2024).\n\n[4] Margatina et al. \"Active learning principles for in-context learning with large language models.\" (2023).\n\n[5] Hegselmann et al. \"Tabllm: Few-shot classification of tabular data with large language models.\" (2023).\n\n[6] Requeima et al. \"Llm processes: Numerical predictive distributions conditioned on natural language.\" *Advances in Neural Information Processing Systems* 37 (2024)\n\n[7] Xie et al. \"An explanation of in-context learning as implicit bayesian inference.\" (2021)."}, {"review_id": "IBM4gozuni", "reviewer": "Reviewer_esK3", "summary": "This paper proposes a novel variational framework for uncertainty decomposition in in-context learning (ICL). It avoids both sampling and explicit computation of the model parameter posterior which are notoriously difficult for ICL. Specifically, it derives an upper bound on aleatoric uncertainty by introducing auxiliary inputs that serve as probes into the underlying parameter posterior. To further enhance the Bayesian property of exchangeability in ICL, the method introduces random shuffling of context permutations combined with KL-based filtering. The effectiveness of the proposed framework is demonstrated through synthetic experiments that highlight the interpretability of the decomposed uncertainties, with additional applications shown in bandit problems and out-of-distribution detection in question-answering tasks.", "questions": "1. If I understand correctly, the proposed framework can be applied to estimate aleatoric and epistemic uncertainty for machine learning models with point-estimated parameters. Does this imply that the method is broadly applicable to deep active learning scenarios? It would be helpful if the authors could elaborate on this potential direction.\n\n2. In the KL filtering step, what happens if no candidate satisfies the condition given in Equation (11)? How is this case handled in practice?\n\n3. Restricting the search for $Z$ to a single example from the input domain seems limiting. How might the approach be extended to consider a broader set of $Z$ values, potentially reducing the approximation gap in Theorem 3.1?\n\n4. In the linear regression experiment, the observational noise is modeled as Gaussian with fixed variance. Assuming the regression parameters are estimated via MAP or MLE (i.e., without a Bayesian posterior), why does the total predictive uncertainty vary across different test covariates?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "To the best of my knowledge, the proposed framework is novel and offers an interesting contribution. It introduces a new perspective on the uncertainty of model parameter posteriors through the use of probes, such as auxiliary inputs and outputs. The methods of context shuffling and KL filtering are well-justified and conceptually sound. Moreover, the synthetic experiments provide interpretable evidence supporting the proposed uncertainty decomposition.\n\nA key concern pertains to Theorem 3.1, which is based on marginalizing out the model parameters $\\theta$. The justification provided in lines 130-132 suggests that \"some of the uncertainty in $\\theta$ is absorbed into...\". \nHowever, in in-context learning (ICL), $\\theta$ is treated as a point estimate, implying that there is no posterior uncertainty over model parameters. This discrepancy raises questions about the applicability of the same uncertainty-based interpretation within the ICL framework.\n\nAnother point of concern lies in Equation (5). While the framework proposes identifying $Z$ as a perturbation of $x*$, the equation assumes access to $p(U | Z, D)$. This assumption appears problematic, as the ultimate goal is to infer the unknown $p(y* | x*, D)$. If $Z$ is constructed to be close to $x*$, it is unclear how one could reasonably obtain or estimate $p(U | Z, D)$ without already knowing something about $p(y* | x*, D)$.", "was_revised": true, "final_justification": "While other reviewers raised concerns about auxiliary information choices and inference time complexity, I believe this paper merits acceptance. Despite its limitations, it offers valuable theoretical novelty in uncertainty quantification for in-context learning and could potentially inspire future research developments.", "rebuttal": "### Summary Response\n\nWe sincerely thank the reviewers for their constructive feedback. We are encouraged that reviewers appreciated the ***novel*** (**bTbz**, **esK3**, **y4yc**) ***principled approach*** (**aVVb**) and ***sound theoretical framework*** (**s63H**) to the ***important and underexplored*** (**y4yc**) problem of uncertainty decomposition in in-context learning.\n\nOur work addresses the fundamental challenge of computing uncertainty decompositions on an intractable and implicitly defined Bayesian model of ICL in LLMs. Specifically:\n\n- We derive an upper bound to the aleatoric uncertainty (conditional entropy/variance) of the Bayesian model.\n- To our knowledge, we are the first achieve this without sampling from the Bayesian model posterior or assumptions on the form of the likelihood, via probing the model with auxiliary inputs.\n- We provide a clear statement of our assumptions and employ techniques such as permutation-based sampling and KL filtering to approximate exchangeability to ensure the validity of our method.\n- Whilst our method is particularly suited to LLMs, it can be applied to any model where there is access to the posterior predictive distribution.\n- We show quantitative and qualitative performance of VUD across a range of synthetic datasets and real-world downstream tasks.\n\nIn response to reviewer suggestions, we have expanded our experimental analysis. New results include:\n\n- Table A (see **aVVb**): Additional experimental OOD tasks for natural language on the MMLU dataset\n- Table B (see **aVVb**): New baseline comparisons using Martingale posterior distributions\n- Table C (see **bTbz**): Analysis of the sample efficiency in the number of samples $Z$\n\nWe provide detailed responses to individual reviewer comments and questions below.\n\n---\n### Individual Response\n\n**Comments: Theorem 3.1**\n\n> A key concern pertains to Theorem 3.1, which is based on marginalizing out the model parameters $\\theta$. The justification provided in lines 130-132 suggests that \"some of the uncertainty in $\\theta$ is absorbed into...\". However, in in-context learning (ICL), $\\theta$ is treated as a point estimate, implying that there is no posterior uncertainty over model parameters.\n> \n> \n> Another point of concern lies in Equation (5). While the framework proposes identifying $Z$ as a perturbation of $x^* $, the equation assumes access to $p(U | Z, D)$. This assumption appears problematic, as the ultimate goal is to infer the unknown $p(y^* | x^* , D)$. If $Z$ is constructed to be close to $x^* $, it is unclear how one could reasonably obtain or estimate $p(U | Z, D)$ without already knowing something about $p(y^* | x^* , D)$.\n> \n\nWe believe this concern may stem from a misinterpretation of the Bayesian model parameter $\\theta$. Unlike standard Bayesian deep learning, $\\theta$ **is not the LLM network weights** and we are not trying to compute the posterior distribution $p(\\theta|\\mathcal{D})$. In our framework:\n\n- $x^* $ is the **predictive task**, $y^* $ is the corresponding **label prediction** and $\\mathcal{D} = \\{( x_ i , y_ i ) \\} _ {i = 1 } ^N$ is **in-context data**.\n- If we concatenate $\\mathcal{D}$ and $x^* $ to form the **prompt**, the logits of the LLM prediction $y^* $ give the **posterior predictive distribution** $p(y^ * | x^ * , \\mathcal{D} ) $ which is tractable.\n- Assuming exchangeability of the LLM predictions, we can apply **de Finetti’s theorem**, which **implicitly** defines a Bayesian model with parameter $\\theta$ such that $p ( y^ * | x^ * , \\mathcal{D} ) = \\int p ( y^ * | x^ * , \\theta ) p ( \\theta | \\mathcal{D} ) d \\theta $. Due to the implicit definition of this model we **do not have access to $p( y^ * | x^ * , \\theta ) $** or $p( \\theta | \\mathcal{D} ) $.\n\nTherefore, to clarify the concerns of the reviewer:\n\n- In this setting, $\\theta$ is not a point estimate. Rather this is an implicitly defined parameter that we marginalise over to obtain the posterior predictive. This parameter can be viewed as the underlying conceptual model governing the LLMs predictions [1].\n- We set $Z$ as the predictive task with $U$ as the corresponding prediction. Next, we concatenate $Z$ and $\\mathcal{D}$ into a single prompt to yield the logits of the prediction $U$, giving the posterior predictive distribution $p(U|Z,\\mathcal{D})$. For further clarification, in Section H of the Appendix, we provide example prompts of the ICL data followed by the predictive task.\n\n**Comments: Applicability to Active Learning**\n\n> If I understand correctly, the proposed framework can be applied to estimate aleatoric and epistemic uncertainty for machine learning models with point-estimated parameters. Does this imply that the method is broadly applicable to deep active learning scenarios? It would be helpful if the authors could elaborate on this potential direction.\n> \n\nYes, indeed, this is one of the key applications of the VUD method. In particular, we can use the lower bound estimate of the mutual information at $x^ * $, which is the **expected information gain** from observing $x^ * $. This is a key quantity in Bayesian Experimental Design and widely used in active learning algorithms [2]. This opens the door to active learning in natural language settings, where traditional Bayesian methods are often infeasible [3]. We agree this is a promising direction and are actively exploring its application to deep active learning scenarios.\n\n**Comments: KL Filtering**\n\n> In the KL filtering step, what happens if no candidate satisfies the condition given in Equation (11)? How is this case handled in practice?\n> \n\nTo ensure that all the auxiliary variables aren’t filtered out by the KL-filtering, we rank the KL divergences $\\epsilon_i$ by size and only keep the $r$ auxiliary variables with the smallest KL divergence. This ensures that the method guarantees a decomposition whilst $r$ becomes a proxy for the level of approximation, which is left to the discretion of the user. See Appendix D.3 for further details.\n\n**Comments: Extending** $Z$ **to a broader set of values**\n\n> Restricting the search for $Z$ to a single example from the input domain seems limiting. How might the approach be extended to consider a broader set of $Z$ values, potentially reducing the approximation gap in Theorem 3.1?\n> \n\nTo reduce the complexity of the optimisation problem, we restrict $Z$ to a single datapoint in the input domain.  However, we can extend this to multiple datapoints, for instance through test-time scaling. We outline a potential method for this below:\n\nTo reduce the complexity of the optimization problem, our current approach restricts $Z$ to a single datapoint from the input domain. However, we agree that this may limit the expressiveness of the approximation and that extending the search over broader sets of $Z_i$ could potentially improve the optimisation. A natural extension is to consider small sets of points (e.g., pairs or triplets) as candidates for $Z_i$. One practical way to do this is via a greedy procedure:\n\n1. **Parallel Evaluation**: compute $V_ a ( y^ *| x^ *, \\mathcal{D} , Z_ i )$ for multiple $Z_i$, where each $Z_i$ is a single datapoint.\n2. **Initial Filtering**: Discard candidates that fail to meet a KL-divergence threshold\n3. **Combinatorial Expansion**: Form combinations (e.g. ordered pairs $ (Z_ j , Z_ k )$) from the remaining candidates and compute $V_a (y^* | x^* , \\mathcal{D} , (Z_j,Z_k) ) $\n4. **Secondary Filtering**: Apply the KL condition again to filter out weak combinations.\n5. **Selection**: Choose the optimal set based on the resulting $V_a$ score\n\nWhile this is a simple strategy, it illustrates a potential direction for future work in trading off computational efficiency and approximation quality. \n\n**Comments: Linear Regression Total Predictive Uncertainty**\n\n> In the linear regression experiment, the observational noise is modeled as Gaussian with fixed variance. Assuming the regression parameters are estimated via MAP or MLE (i.e., without a Bayesian posterior), why does the total predictive uncertainty vary across different test covariates?\n> \n\nWhile the observational noise in the linear regression experiment is indeed modelled as Gaussian with fixed variance, the total predictive uncertainty we report is derived from the entropy of the LLM posterior predictive distribution $p ( y^* | x^* , \\mathcal{D} ) $.  The VUD method decomposes this into the conditional entropy (aleatoric uncertainty) $\\mathbb{H} [p ( y^* |x^* , \\theta ) ] $ and the mutual information $\\mathbb{I} (y ; \\theta | x^* , \\mathcal{D} ) $ (epistemic uncertainty).\n\nAs a heuristic, assuming a Gaussian likelihood and prior for $\\theta$, the conditional entropy would be constant due to fixed variance. However, the total uncertainty would be of the form $\\frac{1}{2} \\log ( { x^ * } ^ T A { x^ * }) + c$ for a positive definite matrix $A$ and constant $c$ [4]. Therefore, the uncertainty would increase as $|x^*|$ becomes large.\n\nWe empirically observe this in Figure 4b, where the aleatoric uncertainty is roughly constant but the total uncertainty increases as $|x ^ * |$ increases. While the underlying Bayesian model is not necessary Gaussian, this behavior aligns with the heuristic intuition derived from the Gaussian case, where predictive uncertainty grows with distance from the training data.\n\n---\n\n[1] Xie, Sang Michael, et al. \"An explanation of in-context learning as implicit bayesian inference.\" *arXiv preprint arXiv:2111.02080* (2021).\n\n[2] Rainforth, Tom, et al. \"Modern Bayesian experimental design.\" *Statistical Science* 39.1 (2024): 100-114.\n\n[3] Margatina, Katerina, et al. \"Active learning principles for in-context learning with large language models.\" arXiv preprint arXiv:2305.14264 (2023).\n\n[4] Williams, Christopher, and Carl Rasmussen. \"Gaussian processes for regression.\" *Advances in neural information processing systems* 8, Chapter 2 (1995)."}]

Author rebuttal:

Meta-review: Following the discussion period, the majority of reviewers are positive about this Bayesian approach to uncertainty quantification in generative models such as LLMs, finding the conceptual approach to be interesting and well-motivated.  Concerns about experiments were mostly be addressed by some additional results included in the author rebuttals; please take care that these are integrated in the final manuscript.  One reviewer critiques the method's need to define and evaluate perturbations of the input to assess uncertainty, but the consensus is that this is not a huge assumption, and in fact may be the only option for LLMs that are only accessible via "black box" function evaluations.

Reviewers did have a number of confusions about the submitted manuscript, which were addressed in rebuttals.  Please be sure to carefully revise for readability and clarity.
Decision: Accept (poster)