Contextual Similarity Distillation: Ensemble Uncertainties with a Single Model
Authors: Moritz Akiya Zanger
Pascal R. Van der Vaart
Wendelin Boehmer
Matthijs T. J. Spaan
Venue: ICLR 2026 Poster (2026)
Keywords: Uncertainty Quantification
Epistemic Uncertainty
Reinforcement Learning
Deep Ensembles
Exploration
Neural Tangent Kernel
Primary area: reinforcement learning
TL;DR: We introduce contextual similarity distillation, a scalable and principled method for uncertainty quantification and exploration.

Abstract: Uncertainty quantification is a critical aspect of reinforcement learning and deep learning, with numerous applications ranging from efficient exploration and stable offline reinforcement learning to outlier detection in medical diagnostics. The scale of modern neural networks, however, complicates the use of many theoretically well-motivated approaches such as full Bayesian inference. Approximate methods like deep ensembles can provide reliable uncertainty estimates but still remain computationally expensive. In this work, we propose contextual similarity distillation, a novel approach that explicitly estimates the variance of an ensemble of deep neural networks with a single model, without ever learning or evaluating such an ensemble in the first place. Our method builds on the predictable learning dynamics of wide neural networks, governed by the neural tangent kernel, to derive an efficient approximation of the predictive variance of an infinite ensemble. Specifically, we reinterpret the computation of ensemble variance as a supervised regression problem with kernel similarities as regression targets. The resulting model can estimate predictive variance at inference time with a single forward pass, and can make use of unlabeled target-domain data or data augmentations to refine its uncertainty estimates. We empirically validate our method across a variety of out-of-distribution detection benchmarks and sparse-reward reinforcement learning environments. We find that our single-model method performs competitively and sometimes superior to ensemble-based baselines and serves as a reliable signal for efficient exploration. These results, we believe, position contextual similarity distillation as a principled and scalable alternative for uncertainty quantification in reinforcement learning and general deep learning.

Published as a conference paper at ICLR 2026
CONTEXTUAL SIMILARITY DISTILLATION: ENSEMBLE UNCERTAINTIES WITH A SINGLE MODEL
Moritz A. Zanger, Paseal R. Van der Vaart, Wendelin Böhmer & Matthijs T. J. Spaan Department of Intelligent Systems, Delft University of Technology Delft, 2628 XE, The Netherlands
Correspondence to: m. a 2 zange r tudel ft.nl
ABSTRACT
Uncertainty quantification is critical aspect of reinforcement learning and deep learning, with numerous applications ranging from efficient exploration and stable offline reinforcement learning to outlier detection in medical diagnostics. The scale of modern neural networks, however, complicates the use of many theoret- ically well-motivated approaches such as full Bayesian inference. Approximate methods like deep ensembles can provide reliable uncertainty estimates but still remain computationally expensive. In this work, we propose contextual similarity distillation, a novel approach that explicitly estimates the variance of an ensemble of deep neural networks with single model, without ever learning or evaluating such an ensemble in the first place. Our method builds on the predictable learning dynamics of wide neural networks, governed by the neural tangent kernel, to de- rive an efficient approximation of the predictive variance of an infinite ensemble. Specifically, we reinterpret the computation of ensemble variance regression problem with kernel similarities as regression targets. The resulting model can estimate predictive variance at inference time with single forward pass, and can make use of unlabeled target-domain data data augmentations to refine its uncertainty estimates. We empirically validate our method across a variety dut-of-distribution detection benchmarks and sparse-reward reinforcement learning enviromments. We find that our single-model method performs competitively and sometimes superior to ensemble-base for efficient exploration. These results, we believe, position contextual similarity distillation principled and scalable alternative for uncertainty quantification reinforcement learning and general deep learning.
supervised
baselines and serves reliable signal
INTRODUCTION
With the deployment of increasingly large deep learning systems to real-world applications, efficient uncertainty quantification has become an essential challenge of modern deep learning. Assessing the reliability in predictions is crucial in applications ranging from out-of out-o-diistriutton (OOD) detection to deep reinforcement learning (RL), where uncertainty estimation is used to drive exploration, stabilize offline learning, increase data efficiency, or to design cautious, safety-aware agents. A necessary condition for designing and deploying such agents is their ability to quantify uncertainty reliably and efficiently.
Bayesian methods for deep neural networks address this challenge with solid theoretical footing (Goan and Fookes 2020 Pearce et al. 2020 [zmailov et al. 2021) but often require coarse approx- imations or costly sampling from a complex posterior: To this end, deep ensembles from random initializations/ Lakshminarayanan et al. (2017): Osband et al. (2016); Qin al. (2022) have emerged as a simple but reliable method for estimating predictive uncertainty in neural networks. While usually more efficient than full Bayesian inference, the computational cost of training several models remains a burden, particularly with increasing parameter spaces.
In this paper, introduce contextual similarity distillation (CSD), novel single-model approach that directly estimates the variance of a random initialization ensemble of deep NNs without ever training evaluating such ensemble in the first place. The theoretical motivation for our approach

---

Published as a conference paper at ICLR 2026
derived from recent work characterizing the learning dynamics of wide neural networks through the Neural Tangent Kernel (NTK, Jacot et al. 2018; et al. 2020). Under some conditions, this setting allows us to describe deep ensembles and in particular their predictive variance by the NTK Gaussian Process (NTK GP, He et al. 2020), providing an analytical expression for ensemble uncertainties. Although one can in principle solve these analytical expressions explicitly without requiring training of an ensemble of models, these computations quickly become infeasible when considering large models or datasets, as frequently encountered in the field of RL.
In contrast, we devise novel method called contextual similarity distillation (CSD) that is amenable to regular training pipelines based on gradient descent and approximates predictive ensemble variance with a single forward pass. We derive our method from the insight that ensemble variance can be obtained as the result of structured supervised regression problem, where labels correspond to kernel similarities between training points and test point 2;. As result, one can obtain the predictive variance of a deep ensemble for known query point 2; by training a single NN on a regression task using gradient descent and a carefully designed label function dependent on Z₁. We then extend this "single-query" approach to work efficiently for arbitrary queries Ty by formulating a contextualized regression model that involves regression tasks with family of context-dependent label functions. This formulation moreover enables CSD to refine its uncertainty estimates by leveraging unlabeled data, for example from target domain of interest or from data augmentation techniques, approach that has proven extraordinarily successful in the field of self-supervised and representation learning
(Chen et al. 2020 Guo al. 2022 Caron et al. 2021p.
We analyze the practical effectiveness of CSD through an empirical evaluation on a variety of distribu- tion shift detection tasks (Van Amersfoort et al. 2020) using the MNSHIO MNIST, KMNIST,
and NOTMNIST datasets (Xiao et al.) 221 Deng 2012, Clanuwat et al. 2018). We moreover use
CSD to generate an exploration signal on sparse- reward reinforcement learning problems from the visual RL benchmark VizDOOM (Kempka et al. 2016). Empirically, CSD consistently achieves competitive and sometimes superior uncertainty estimation to finite deep ensembles and other baseline methods while maintaining lower computational cost. We helieve these results establish CSD) as both principled and scalable alternative to ensemble-based uncertainty quantification and exploration methods.
BACKGROUND
For our default framework, we consider a finite Markov Decision Process (MDP, Bellman 1957
of the tuple (S, A. R, r, p), with state space S, action space A, immediate reward distribution S x A + A (R). discount 7 E [0, 1]. transition kernel P: A (S)). and the start state distribution μ E (S). Here, (P( indicates the space of probability distributions over some space Z and random variables are denoted with uppercase letters. Given state S, at time f, agents choose action A, from a stochastic policy и S + (AA) and subsequently receive the immediate reward R2 ~ RR(-SS,,A) A:) and observe next state S ~ ((I(S,,)). The expected discounted sum of future rewards, conditioned on a particular state s and action is known as the state-action value and is given by QTT a) = Ep, 20 7 Rr |So = An a). This value function adheres to a temporal consistency condition described by the Bellman equation (Bellman
1957
Qa1(s, Ep,, [Ro (S, Ai) s, Au = a]
where Ep [-] indicates that S₁ and A1 are drawn from and respectively. The expected return of policy π can compactly be expressed through the state-action value and the starting state distribution through
J(T) Eso~H.Anrr (So. Ao)].
(2)
The objective of reinforcement learning is to find an optimal policy n" that maximizes the above equation x- arg max (TT)
2.1 EXPLORATION IN REINFORCEMENT LEARNING
A fundamental challenge in attaining an optimal policy x² lies in the sxploration-exploitaaee tradeoff: an agent must decide whether to exploit its current knowledge to maximize returns or whether to explore novel actions in order to discover better strategies. Efficient exploration is particularly erucial
2

---

Published as a conference paper at ICLR 2026
in dinenssinnii or sparse-reward settings, where naive strategies such random exploration require prohibitive amounts of interactions.
A widely used approach to exploration is optimism in the face of uncertainty (Auer et al., 2008 Auer 2002), where agents prioritize actions with high epistemic uncertainty in value estimates. In the context of model-free RI,, provably efficient algorithms often rely on the construction of upper confidence bound (UCB) that overestimates the true optimal value Q (s.a) with high probability (Jin et al., 2018: 2020; Neustroev and de Weerdt (020). This may be implemented by adding a well-chosen exploration bonus b(s, a) to value estimates according
QTh a) Q* (s, a) b(s,a).
In small state-action spaces, such bonuses can be derived from count-based concentration inequal- ities (Bellemare et al., 2016 Jin et al.) 2020). whereas high-dimensional, continuous domains
usually require function approximation, significantly complicating efficient uncertainty estimation
(Ghavamzadeh et al. 2015 Osband et al.) 2016 Lakshminarayanan et al. 2017 Burda et al. 2019).
With the widespread use of deep neural networks, deep ensembles LLakshminarayaaan et al. 2017)
based on random initialization have become a dominant tool for quantifying epistemic uncertainty in high-dimensional continuous spaces (Chen et al. 2017 Osband al., 2019) He et al. 2020). An informal intuition behind the effectiivnness of ensembles is the tendency of randomly initialized NNs to converge to diverse minima in the training loss landscape (Fort et al. (020). leading to higher prediction diversity for unseen inputs. The variance among ensemble members can then be used to measure the model's uncertainty for a specific input.
2.2 NEURAL TANGENT KERNEL GAUSSIAN PROCESSES
In order to better understand the properties of deep ensembles and to design better exploration algorithms, an analytical description of deep neural networks and their learning dynamics is desirable. While general framework remains elusive, significant progress has been made in the field of deep learning theory. In particular, seminal works by bacot et al.) (2018) and Lee et al. (2020) have shown that wide neural networks trained by gradient descent are well-described by their linearized training dynamics and thus predictable.
For this, let neural networks be parametrized functions fæx 0,) R + R and denote training data X = {a, E R E ND and training labels ) = {w E Rli E {l....,Ny, ND We assume training is performed using gradient descent with infinitesimal step sizes, also referred to as gradient flow. The initialization weights Gg are drawn i..dd from a normal distribution 8, ~ N, and deep ensembles are formed by training multiple independently initialized neural network functions. We furthermore assume so-called NTK- PTKKpprrametiization which scales forward and backward passes in proportion to layer widths (see Jacot al. 2018 Lee et al. 2020 for details).
key result by Lee et al. 2020) is that in the limit of infinite layer widths, the training dynamics of deep networks are described exactly by a Taylor expansion around the parameter initialization 0g In this setting, the NTK ((r, x²)= R + IR, first described by Jacot et al. (2018) emerges as the defining function governing learning dynamics:
Oo(2, x') = Vaf(z, 0o) T (2f(('
The NTK can be interpreted as a similarity measure between inputs based on gradient representations of the inputs x and x'. Crucially, Jacot et al. (2018) find that in the limit of infinite layer width, Ozz, x²) becomes deterministic despite random weight initialization Oo(2,z²) O x²) and remains constant throughout training, inducing analytically solvable training dynamics. As result, the post-training NN function [(, 0zw) can be characterized a deterministic function of the random initialization [(x 0o) through
f(x, 0x.) = f(x, 0o) + Oz,
(5)
Here, we have overloaded notation to indicate the vectorization 0 X) E RIXND, O X') E N No, and forth. The matrix @ X)) is also known the training Gram matrix, will refer to it. Further extending this framework, He et al. (2020) demonstrate that by introducing suitable function priors on [xx On), akin to the well-known randomized prior functions by Osband al.
3

---

Published as a conference paper at ICLR 2026
(2119) , the post-training function is described by Gaussian Process (GP. Rasmussen and Williams 2006):
~ ((C, ((
(6)
(ovlf(2t,
where Xe is an arbitrary test data set. An outline of the derivation of Equations 5 and 6 is provided in
Appendix [A] Consequently, the variance of an ensemble over infinite random initializations is given by
(x) -000
The above expression provides with theoretical footing for understanding the behavior and uncertainty estimates deep ensembles. In the following sections we will describe our approach for estimating Eq. 7 not as the result of training several random models but deterministically with single model.
3 CONTEXTUAL SIMILARITY DISTILLATION
We now proceed to describe our approach, contextual similarity distillation (CSD), The main objective of our method is to approximate the variance of an infinite deep ensemble, as described by Eq. directly with single model.
3.1 ENSEMBLE VARIANCE PREDICTIONS FOR A PRIORI QUERIES
We introduce the underlying idea of CSD in the simplified setting of a priori known test points. Given a test query point 2t, it is our goal to estimate the variance (1f( of an ensemble of independently initialized NNs, trained on dataset X. It is important note that could in principle obtain this variance via the NTK GP by solving Eq. 7 This, however, requires inversion of the potentially very large Gram matrix @ X),, which becomes computationally prohibitive for most datasets and models of interest, including RL applications where sample sizes can go into the
billions.
Instead of solving Eq. 7 directly, we leverage an alternative perspective that arises naturally from the learning dynamics of wide neural networks. Specifically, we begin with the simple observation that the variance of wide ensemble at test point 2, can be computed efficiently as the solution to regular supervised regression problem of a single model with a particular label function. For this, let g 0.) be a NN such that its NTK is equal to f with (y(x, x²) (xx, (i.e., with the same architecture and initial weight distribution). Recall that the post-training NN function g(æ, 0oo) with squared loss y is given by
g(x, 0xx = g(x, 0o) + (10,
g(X 0o))
(8)
is straightforward to see that for small function initialization 09(2, 0b) = O, Vr the r.h.s. this expression, when choosing the label function Vi, (X) O(X 22), simplifies to
Ix, (x, 0xx)) = O(x,
where we used the subscript v, to indicate the function's 's dependence on the label function Vz,- This identity now recovers exactly the problematic right term of Eq. containing the Gram inversion ( X1- Note that 9r, (I, 0zx)) is obtained "naturally" the result of gradient-based regression, without requiring explicit inversion of ( 2 X') or training of large ensemble at any point. The ensemble variance in a query point 2, can be obtained as
[If(rz 0xx ) ((z), Iv. (2e, 0oo),
(10)
which can be computed efficiently. Fig. I illustrates the above-described process of obtaining expression 10 geometrrcclll.. While simple, we believe this formulation provides a crucial insight: uncertainty estimation for NN can be phrased as a singular prediction problem of kernel similarities.
For example, small function initialization can simply he obtained by redefining f(z, 0t) = f(z, 0c)
faz 0g).
4

---

Published as a conference paper at ICLR 2026
Training Data
NIK-PP Variane (x
Kenne Price ((a,a)
Similerity Regrossio si,tt, Similerity-Relabbeeed Duta
Figure/ 1: Illustration of regression tasks with query-dependent NTK similarities labels. The difference between the kernel prior function (xx, x) (dotted line) and the post-training regression function 9N, (x, foe) matches exactly ensemble variance in Se note that we shifted curves black by constant offset for each subplot illustrate this equality). Plots from left to right depict the same principle, but for different query points Tt.
3.2 ENSEMBLE VARIANCE ESTIMATION FOR ARBITRARY QUERY POINTS
In the above derivation, we outlined an efficient method for obtaining ensemble variances at specific test query point Ft known a priori. An obvious limitation of this approach, however, that the used labeling function Vz, (X) @(X, æ²) and by extension the model 92, (x. 00)) is inherently dependent on the test point 2; and not usable for arbitrary queries.
To overcome this limitation, we now formulate contextualized regression model g(æ, c, 0e), where e serves as a context variable that determines the label function used during training the function g(x, c, 0t). Specifically, instead of defining a label function that depends on a single fixed test query Tt. we construct family of label functions parameterized by the context c, ()( = O c).
This means that for set of context data C {cz R {1,.... Ne) ). the model g(r, c, 0.) is optimized to solve a supervised regression problem associated with labels (x()).
Intuitively, this approach can be interpreted as attempt to interpolate between multiple regression solutions that were trained on the same dataset X' but with different label functions (x()). Geomet- rically, this corresponds conjoining the functions Ja, in Fig. [ aalong new dimension c. So long as
g(x, c, 0z)) maintains the approximate dynamics of (r, az)). ) this model can be evaluated quickly for arbitrary test points by setting c in
g(z, c, 0x)) Olz, x)0(X,
(11)
This generalization accordingly enables ensemble variance estimation across arbitrary points without requiring a separate regression solution for each individual query by computing
(x 0(x,2) g(x, x, 0oo).
(12)
An intuitive interpretation of the function g(z, x, (xo) is that it captures an ensemble's confidence gained through observing the training data X, weighted by its similarity to æ. The resulting variance of Eq. [2] can then be understood as the difference between a prior uncertainty term O 2)) and the confidence term g(x, x, 0.)). One should note at this point, that the evaluation of g(x, 000) for contexts c & C not used during training requires g to generalize to novel c. Furthermore, the introduction of the context variable may influence the training dynamics of g, putting this approach into the realm of approximate algorithms. We have added a section to Appendix B. that discusses and summarizes used approximations and their implications for practical settings.
Finetuning Variance Estimates with Context Data. Before proceeding to describe our practical setup, we outline a property of contextualized similarity distillation that emerges through the above- deseribed modeling choices. Our theoretical motivation highlights that exact ensemble variances (in the NTK regime) can be obtained when the test point 2, is known a priori. The implication of the subsequent formulation as contextualized regression problem is that, when available, one
5

---

Published as a conference paper at ICLR 2026
Figure 2: Top Row: Variance an ensemble of 100 randomly initialized neural networks on a 2D toy regression task. Red dots training points. Bottom Row: Variance prediction by contextual similarity distillation (CSD) with single model on the same regression task.
can include unlabeled context data C during training to obtain better uncertainty estimates in the domain of interest, as we will show later the experimental section. This property also opens up improve uncertainty estimation, approach that has proven extraordinarily successful in the field of self-supervised and representation learning
possibility of using unlabeled data augmentations
(Chen al. 2020; Guo al. [0022] Caron al., (222)) and not easily incorporated with standard approaches for uncertainty estimation (Lakshminarayanan et al. 2017 Gal and Ghahramani 2016
Burda 2019
3.3 CONTEXTUALIZED SIMILARITY DISTILLATION WITH DEEP NEURAL NETWORKS
Building on this theoretical basis, we proceed to describe a setting for contextualized similarity distillation with deep neural networks. This section outlines algorithmic design choices we found to be computationally efficient while maintaining the approach's 's thcoretical motivation.
First, we parameterize the contextualized regression model g(æ,, c, 0.x)) as an inner product between a feature vector (αx, Ieam) and context vector (c, fesa) as
g(r, 0o)) = Φ(x, ( (c, fcm))
(13)
Conceptually, this parametrization can be thought of as introducing a context-dependent final layer of weights, represented by (e, 0err), to the regression model g. Computationally, this inner product parametrization bears the advantage that g(X',
, C, 0x)) E R Ne can be evaluated quickly without requiring explicit forward passes for each pairing ( X, Cj C).
Second, we approximate the NTK prior (xx, with partial gradients. Given that O(x, x²) is not involved in backward gradient computations, computing the full analytical empirical prior kernel functions (ax, x²) is often not computationally prohibitive, but can pose burden for models with large parameter spaces. We find that gradients with respect only the last layer weights of sufficient in practice and further accelerate computation. Assuming, the last layer of f is a dense layer such that (ax 0o) = sp(x, 00 01:L-111111 of, we have
(x = Vog fxx 0o) Vag ((zf((100) p(z, 1T ( 10.
(14)
The resulting training pipeline for g(x, c, 0,) involves a simple supervised regression task with
minimization of the squared loss, where (Is, ci) sampled randomly from X' and C
1 Σ (glan, 0.) cz), ((O.) N
(15)

---

Published as a conference paper at ICLR 2026
Table 1: Distribution Shift Detection. Test accuracy and average OOD detection metrics across
MNIST, FashionMNIST, KMNIST, NotMNIST. OOD metrics are evaluated for each ID dataset against the remaining OOD datasets and a perturbed version of the ID dataset.
Method
Acc.
AUROC
AUPR-IN
AUPR-OUT
MCD
94.39 + 0.10
85.67 0.21
81.73 0.34
86.44 0.20
BNN-MCMC
87.70 = 0.38
83.17
62.65
82.28 + 0.71
BNN-Laplace
90.86 = 0.62
81.38
79.43
#.66
RND
96.18 + 0.05
94.17
94.01 0.31
ENS(3)
96.91 + 0.04
92.30 0.09
92.83
91.37 + 0.11
ENS(15)
97.18 = 0.03
94.00 0.07
94.70
92.99 0.06
CSD
96.29 = 0.07
96.63 = 0.35
96.94 + 0.39
96.19 + 0.32
CSD-Aug.
96,28 = 0.06
98.22 = 0.14
98.51 + 0.13
97.80 + 0.17
CSD-OOD,
96.30 - 0.06
98.57 = 0.14
98.86 + 0.12
98.19 + 0.15
Lastly, we propose several choices for the context data C. We find that the arguably simplest choice, that is to reuse the training set ~ X, works well in practice and is easily implemented. addition, apply data augmentations to the training samples X' when using context data. For this, we employ the well-established set of augmentations from the contrastive learning literature (Chen et al., 2020). We note here, that designing novel data augmentation techniques for the purpose of uncertainty quantification a promising avenue (see for example works by Wen et al. (2020) and Wu and Williamson (2024)). Unlike contrastive learning and many other self- supe-susee methods, our approach does not require data augmentations to preserve the nature of the original label and can in principle use any unlabeled data. Finally, when available, unlabeled data from the test distribution of interest can be used and often provides an additional improvement in uncertainty estimation, as we will show empirically.
it is possible
EMPIRICAL EVALUATION
Our empirical evaluation aims to provide us with a better understanding of contextual similarity distillation in practice. Given that our approach introduces approximations beyond the theoretical framework, we investigate whether CSD maintains its theoretically motivated properties in practice with high-dimensional problem and parameter spaces. Specifically, we aim to assess whether CSD provides scalable alternative to deep ensembles and other established methods in uncertainty quantification, including Monte Carlo dropout (Gal and Ghahramani] (016)) Bayesian NN based on Markov chain Monte Carlo sampling (BNN MCMC, Garriga-Ali Alonso and Fortuin 202 a Laplace approximated Bayesian NN (BNN Laplace, [mmer et al. 2021). deep ensembles of sizes 3 and 15 (ENS. Lakshminarayanan et al. 2017) and random network distillation (RND, Burda et al. 2019). Furthermore,
analyze how algorithmic design choices, such as the choice of context data, influence uncertainty estimates. Lastly, we seek to evaluate our approach's efficacy as an exploration signal for deep reinforcement learning agents on sparse-reward visual exploration tasks from the VizDoom (Kempka et al. 2016) suite.
Code for experimental reprodccttono available at http:///////aoooooooooooooto/aaa
github. cwm/ http://ooconaaayoooy//otteeuuuue wtt.::/wwwuusssmsmaaaaa.m.-m-mmo
wee.vviidooo zdoom
4.1 DISTRIBUTION SHIFT DETECTION
Following prior work (Van Amersfoort et al. 2020 Immer et al., [202] Rudner et al. 2022). evaluate uncertainty estimates in image classification under distribution shift, where a model trained an in-distribution dataset is evaluated on inputs from shifted distribution.
In particular, we train models on one of the FashionMNIST, MNIST, KMNIST, NotMNIST datasets and evaluate uncertainty estimates the other, shifted datasets and perturbed version of the in-distribution dataset. Well-calibrated epistemic uncertainty estimates will correlate with dataset shift, such that -fitt--iiittiuution samples are likely to be rated more uncertain than in-distribution
7

---

Published as a conference paper at ICLR 2026
VoDDoom Eavinonment
Donse
Syurse
Very Sparse
0.5
Optial CSD
DON
BDON-P
-IS--C51
MyWayllome
0.0
0.3
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Environment internetions (in 10²)
Figure 3: (Left)): Visual observation in the VizDoom environment (Kempka et al. Second Left to Right): Mean learning curves in variations VizDoom MyWayHome. Shaded regions 90% Student's 's confidence intervals from 10 seeds.
2016). (From
samples. To compare methods quantitatively, we use the thresholddindepeeddent area under the receiver operating characteristic curve (AUROC) metric, as well as the area under the precision-rechl curve for in-distribution (AURR-IN) and ditto--distribution (AUPR-OUT) samples. The AUROC metric can be interpreted as the likelihood an OOD sample receiving higher uncertainty than an ID sample, while AUPR-IN and AUPR-OUT provide additional sensitivity to dataset size and the choice of the positive class. For these metrics, Table repprts the average and standard deviation over 10 seeds, averaged over all permutations of ID and OOD datasets, along with average accuracy. Full detailed results are provided in the supplementary material.
To analyze the role of the used context data, evaluate three versions of CSD: baseline that only uses training data (CSD), a variant incorporating data augmentations to training samples (CSD-Aug.). and a model using context data from the evaluation distribution (CSD-OOD). Even in the basic version, CSD demonstrates highly effective distribution shift detection, surpassing baseline methods on variety of datasets while requiring only a single model. Our results furthermore suggest that incorporating data augmentations and target-distribution context data indeed significantly improves performance.
4.2 EXPI.ORATION VIZDOOM
We now evaluate CSD in reinforcement learning task with nigh-dimensional observation spaces and sparse rewards. For this, we consider visual navigation tasks in the VizDOOM environment, where agents explore a 3D maze-like environment with ego-perspective image observations, The agent is tasked with reaching a goal while receiving a minimal constant negative reward except upon successful completion, where reward of 1 is given. We consider three variations of the task, where agents are initialized at increasing distances from the goal, defining progressively harder exploration tasks (details provided in Appendix
We use DQN agent (Mnih et al.) 2015) as base algorithm and include uncertainty estimates by CSD as intrinsic reward (full details provided in Appendix C). For comparative evaluation, we compare the performance of RSDD-aased exploration with several baseline algorithms, including deep networks (DQN, Mnni et al. 2015). random network distillation (RND, |Burda et al., 2019). bootstrapped Q Qentwwrkss (BDQN+P, Osband et al. 2019), and information -irreceedd sampling (IDS, Nikolov et al., 2019). Fig. 3 shows mean learning curves across 10 random seeds. Interestingly, sparse version of the environment appears to be the hardest, a circumstance we believe is due to the spawning point lying in a sidearm of the maze map. Of the tested methods, only CSD was able to find the goal across all seeds and environments, with RND performing most competitively.
5 RELATED WORK
Our work builds the extensive body of literature in the field of uncertainty quantification in deep learning and reinforcement learning. Ensemble learning (Dietterich) (0000) has emerged as on the most effective and reliable approaches uncertainty estimation (lakshminarayanan et al. 2017)
8

---

Published as a conference paper at ICLR 2026
and has been widely adopted in the deep reinforcement learning literature. In particular, ensembles can be used for efficient exploration by sampling random models (Osband et al. 2016 Qin et al.)
2022] OSband and Van Roy 2017). by constructing upper confidence bounds for exploration bonuses (Chen et al. 2017] Donoghue et al. 2018) or by estimating information gain (Nikolov et al. 2019). Several works moreover rely
deep ensembles reduce overestimation and improve learning
stability (Fujimoto et al.
2018
Haarnoja et al. 2018
Chen et al. 2022 extending to the challenging
offline setting (An et
2021
Agarwal et al.
2020 Smit al.
2021
A number of previous works have focused on reducing ensemble size, notably by disaligning the
Jacobian of networks (An al., 202T), adding repulsive loss terms (Sheikh et al. 2022). or through
architectural diversification (Osband et al. 2019; Zanger et al., (202)) Notably, varios works aim to quantify epistemic uncertainty with single model (Pathak et al. 2017 Burda al. 2019 Filos et al., 2021; Guo et al., 2022, Lahlou et al. (02T), often by measuring prediction errors. To the best of our knowledge, few single-mode methods in the field offer an interpretation as ensemble or posterior uncertainty.
In broader sense, ensembles have been studied extensively from a Bayesian perspective (Hoffmann) and Elster 2021 D' Angelo and Fortuin 2021). In particular, some of our work relies on the NTK GP characterization of deep ensembles by He et al. (2020). who, in turn, rely on seminal work by seminal work on the NTK by Jacot et al. (2018) and Lee et al. (020). Subsequent analysis has used the NTK to disentangle ensemble variance (Kobayashi et al. 2022). Recent works Wilson et al. (2025) rely on NTK theory to derive sampling-based uncertainty estimator, while Whivo-Calvo--ee Calvo-Ordoñer et al.) (2024) construct uncertainty estimates using several regression models. In contrast to the latter, our method uses coatextualized regression model that allows for single-model uncertainty estimates in deep learning setting.
6 CONCLUSION
This work introduced contextual similarity distillation (CSD), a novel single-model approach for uncertainty quantification that estimates the predictive variance of an ensemble with single model and forward pass. By reframing ensemble variance estimation as a structured regression problem, CSD enables efficient uncertainty estimation without requiring the training of multiple models, stochastic forward passes, or explicit kernel matrix inversion. Instead, phrasing predictive variance estimation as a contextualized regression problem is amenable to standard training pipelines with deep NNs and gradient descent.
We implemented CSD deep learning setting and performed a comparative evaluation on a variety of distribution shift detection and reinforcement learning tasks. Empirically, found that CSD provides uncertainty estimates competitive and sometimes superior to deep ensembles and other alternatives all tasks. This makes CSD attractive option for guiding exploration in RL, as experiments on high-dimensional exploration tasks confirmed. Our results furthermore confirmed that our approach can leverage unlabeled target domain data and data augmentations to further refine uncertainty estimates. We believe our work opens up several avenues for future research. Due to its conceptual similarity to contrastive learning approaches, we believe refining the generation and incorporation of contextual data through augmentation is an exciting avenue for research that is currently not commouplace in the context of uncertainty quantification. Moreover, natural extension of our approach could aim to include an explicit quantification of aleatoric uncertainty so as to provide a complete separation of aleatoric and epistemic uncertainties within one model. Although our current derivations to not consider the learning dynamics of such probabilistic models, we believe such an extension to be feasible and valuable. Lastly, our approach could be leveraged to drive exploration in various hard exploration tasks or drive stability in offline RL.
Our findings, we believe, position CSD as a scalable alternative to deep ensembles, offering a principled and computationally efficient method for uncertainty quantification in deep learning.
9

---

Published as a conference paper at ICLR 2026
7 ACKNOWLEDGEMENTS
This project has received funding from the EU Horizon 2020 programme Epistemic Al under grant number 964505 and the Dutch Research Council (NWO) project Reliable unttof-Distributtonn Gener- alization in Deep Reinforcement Learning with project number OCENW.M.21 ,244 Computational resources for experimental studies were provided by the Delft High Performance Computing Centre (DHCC) and the Delft Artificial Intelligence Cluster (AIC).
10

---

Published as a conference paper at ICLR 2026
REFERENCES
R. Agarwal, D. Schuurmans, and M. Norouzi. An optimistic perspective on offline reinforcement learning. International conference on machine learning. pages 104-114, PMLR, 2020.
G. An, S. Moon, J.-H. Kim, and H. O. Song. Uncertainty-base offline reinforcement learning with diversified g-ensemble. Advances in neural information processing systems, 34.7436-7447, 2021.
Auer. Using confidence bounds for explottaiion-exploraiio trade-offs. Journal of machine learning research, 3, 2002.
Auer, T. Jaksch, and R. Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems, 21, 2008.
M. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
R. Bellman. A Markovian decision process. Journal of mathematics and mechanics, 6, 1957.
Y. Burda, H. Edwards, A. J. Storkey, and O. Klimov. Exploration by random network distillation. In International conference on learning representations, ICLR, 2019.
S. Calvo-Ordoñez, K. Palla, and K. Ciosek. Epistemic uncertainty and observation noise with the neural tangent kernel. arXiv preprint MKX::2409.039555 2024.
M. Caprio, S. Dutta, K. J. Jang, V. Lin, R. Ivanov, O. Sokolsky, and Lee. Credal bayesian deep learning. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL
net forum?ire 4NHF9AC5U1
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. Proceedings of the IEEE/CVF international conference on computer vision, pages 9650-9660, 2021.
R. Chen, S. Sidor, Abbeel, and J. Schulman. UCB exploration via Q-ensembles. arXiv preprint 17X.vv1706.011502 2017.
T. Chen, S. Kormblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
X. Chen, C. Wang, Z. Zhou, and K. Ross. Randomized ensembled double Q-learning: Learning fast without model. arXiv preprint arXhv:2101.05982, 2021.
T. Clanuwat, M. Bober-Irizar, A. Kitamoto, A. Lamb, K. Yamamoto, and D. Ha. Deep learning for classical japanese literature. CoRR, abs/
2018.
J. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent neural networks typically occurs at the edge stability. In International conference on learning representations, 2021. URL https::/oppenevvll
net/ forum? ?id= jh- rTtvkGeM
Cuzzolin. The Geometry of Uncertainty: The Geometry of Imprecise Probabilities. Artificial Intelligence: Foundations, Theory, and Algorithms. Springer International Publishing, Cham, 2021.
F. D' Angelo and V. Fortuin. Repulsive deep ensembles are bayesian. Advances in Neural Information Processing Systems, 34:3451-3465. 2021.
A. Daniely, R. Frostig, and Y. Singer. Toward deeper understanding neural networks: The power of initialization and dual view expressivity. Advances in neural information processing systems, 29, 2016.
Delft Artificial Intelligence Cluster (DAIC), 2024.
Delft High Performance Computing Centre (DHPC). DelftBlue Supercomputer (Phase 1), 2022.
11

---

Published as a conference paper at ICLR 2026
L. Deng. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 2966):111-14, 2012.
T. G. Dietterich. Ensemble methods in machine learning. In Multiple classifier systems: First international workshop, MCS. Springer, 2000.
D'Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville. Sample- efficient reinforcement learning by breaking the replay ratio barrier. In International Conference on Learning Representations, ICLR, 2023.
L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, Dunning, al. Impala: Scalable distributed deep-RL with importance weighted actor-leamer architectures. In International conference on machine learning. PMLR, 2018.
A. Filos, E. Vértes, Z. Marinho, G. Farquhar, Borsa, A. Friesen, F. Behbahani, T. Schaul, A. Barreto, and S. Osindero. Model-value inconsistency as a signal for epistemic uncertainty. arXiv preprint arXiv:2112.04153, 2021.
S. Fort, H. Hu. and B. Lakshminarayanan. Deep Ensembles: A Loss Landscape Perspective. 1rXiv:1912.02757 [cs, stat]. June 2020.
S. Fujimoto, H. Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning. PMLR, 2018.
Y. Gal and Z. Ghahramani. Dropout as Bayesian approximation: Representing model uncertainty in deep learning. In International conference on machine learning, pages 1050-1059, PMLR, 2016.
A. Garriga-Alonso and Fortuin. Exact langevin dynamics with stochastic gradients. arXiv preprint arXiv:2102.01691, 2021.
M. Ghavamzadeh, S. Mannor, J. Pineau, and A. Tamar, Bayesian reinforcement learning: A survey. Foundations and trends in machine learning. 8, 2015.
E. Goan and C. Fookes. Bayesian neural networks: An introduction and survey. Case studies in applied Bayesian data science: CIRM Jean-Morlet chair, Fall 2018, pages 45-87, 2020.
Gogianu, T. Berariu, M. C. Rosca, C. Clopath, L. Busoniu, and Pascanu. Spectral normalisation for deep reinforcement learning: an optimisation perspective. In International conference machine learning, pages 3734-3744. PMLR, 2021.
Z. Guo, S. Thakoor, M. Pislar, B. Avila Pires, F. Altché, C. Tallec, A. Saade, D. Calandriello, J.-B. Grill, Y. Tang, et al. B Exploro Exploration by bootstrapped prediction. Advances in neural information processing systems, 35:31855-31870, 2022.
T. Haamoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
B. Hanin and M. Nica. Finite depth and width corrections the neural tangent kernel. arXiv preprint 159,0999, 2019.
H. Hasselt. Double Q-learning. Advances neural information processing systems, 23, 2010.
B. He, B. Lakshminarayanan. and Y. W. Teh. Bayesian deep ensembles via the neural tangent kemnl.. Advances in neural information processing systems, 33, 2020.
K. He, X. Zhang, S. Ren, and J. Sun. Delving deep into rectifiers: Surpassing human-level per- formance on Imagenet classification. Proceedings of the IEEE international conference on computer vision, 2015.
K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
12

---

Published as a conference paper at ICLR 2026
M. Hessel, J. Modayil, H. Van Hasselt, T. Schaul, G. Ostrovski, W. Dabney, D. Horgan, B. Piot, M. Azar, and D. Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAl conference artificial intelligence, volume 32, 2018.
L. Hoffmann and C. Elster. Deep ensembles from bayesian perspective. arXiv preprint arXiv:2105.13283, 2021.
E. Hüllermeier and W. Waegeman. Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods. arXiv:1910.0944 [cs, stat), Sept. 2020.
A. Immer, M. Korzepa, and M. Bauer. Improving predictions of Bayesian neural nets via local linearization. In International conference artificial intelligence and statistics, pages 703-711. PMLR, 2021.
P. Izmailov. S. Vikram, M. D. Hoffman, and A. G. G. Wilson. What are Bayesian neural network posteriors really like? International conference on machine learning, pages 0129-4640. PMLR, 2021.
A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
C. Jin, Z. Allen-Zhu, S. Bubeck, and M. I. Jordan. Is Q-learning provably efficient? Advances in neural information processing systems, 31, 2018.
C. Jin, Z. Yang, Z. Wang, and M. I. Jordan. Provably efficient reinforcement learning with linear function approximation. In Proceedings of Thirty Third Conference on Learning Theory, volume 125. PMLR, 2020.
M. Kempka, M. Wydmuch, G. Runc, J. Toczek, and W. Jaskowski. ViZDoom: Doom-based AI research platform for visual reinforcement learning. In IEEE Conference on computational intelligence and games. IEEE, 2016.
S. Kobayashi, P. Vilimelis Aceituno, and J. Von Oswald. Disentangling the predictive variance of deep ensembles through the neural tangent kernel, Advances Neural Information Processing Systems, 35:25335-25344 2022.
S. Lahlou, M. Jain, H. Nekoei, I. Butoi, P. Bertin, J. Rector-Brooks, M. Korablyov, and Y. Bengio. Deup: Direct epistemic uncertainty prediction. arXiv preprint arXiv:2102.0850,, 2021.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Le Cam. Asymptotic methods in statistical decision theory. Springer Science Business Media, 2012.
J. Lee, J. Sohl-dickstein, J. Pennington, R. Novak, S. Schoenholz, and Y. Bahri. Deep neural networks as gaussian processes. In International conference on learning representations, 2018.
J. Lee, L. Xiao, S. S. Schoenholz. Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. Journal of Statistical Mechanics: Theory and Experiment, 2020, Dec. 2020.
A. Lewkowycz, Y. Bahri, E. Dyer, J. Sohl-Dickstein, and G. Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/6 2020.
V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Ried- miller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep reinforcement learning, Nature, 518, 2015.
R. M. Neal. Bayesian Learning for Neural Networks. Lecture Notes in Statistics. Springer- 1996. URL /twpp:://wwwww../p http:://wwwww..nge..........
G. Neustroev and M. M. de Weerdt. Generalized optimistic Q-Lerrning with provable efficiency. In International conference autonomous agents and multi-agent systems, pages 913-921, 2020.
13

---

Published as a conference paper at ICLR 2026
N. Nikolov, J. Kirschner, F. Berkenkamp, and A. Krause. Information-directed exploration for deep reinforcement learning. In International conference on learning representations, ICLR, 2019.
I. Osband and B. Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In International conference on machine learning, pages 2701-2710. PMLR, 2017.
I. Osband, C. Blundell, A. Pritzel, and B. Van Roy. Deep exploration via bootstrapped DQN. Advances in neural information processing systems, 29, 2016.
I. Osband, B. Van Roy, D. Russo. Z. Wen, et al. Deep exploration via randomized value functions. Journal of machine learning research, 20, 2019.
B. O' O'Donoghue, I. Osband, R. Munos, and Mnih. The uncertainty Bellman equation and explo- ration. In International conference machine learning. PMLR, 2018.
D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity-driven exploration by self-supervised prediction, In International conference on machine learning. PMLR, 2017.
T. Pearce, F, Leibfried, and A. Brintrup. Uncertainty in neural networks: Approximately Bayesian ensembling. In International conference on artificial intelligence and statistics, pages 234-26 244. PMLR, 2020.
C. Qin, Z. Wen, X. Lu, and B. Van Roy. An analysis of ensemble sampling. Advances in Neural Information Processing Systems, 35:21602-21614 2022.
Rashid, B. Peng, W. Böhmer, and Whiteson. Optimistic exploration even with a pessimistic initialisation. Proceedings ICLR 2020, 2020.
C. E. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning. MIT Press, 2006.
T. G. Rudner, Z. Chen, Y. W. Teh, and Y. Gal. Tractable function-space variational inference in bayesian neural networks. Advances in neural information processing systems, 11-22686-22695 2022.
A. M. Saxe, I.. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 1312,6100 2013.
T. Schaul, J. Quan, I. Antonoglou, and D. Silver. Prioritized experience replay. In Y. Bengio and Y. LeCun, editors, International conference on learning representations, ICLR, 2016.
D. Schmidt and T. Schmied. Fast and data- eata-efficieet training of rainbow: An experimental study on Atari. arXiv preprint arXiv:2111.1024 2021.
M. Seleznova and G. Kutyniok. Neural tangent kernel beyond the infinite-width limit: Effects of depth and initialization. In International Conference on Machine Learning, pages 19522-19560. PMLR, 2022.
G. Shafer. A Mathematical Theory of Evidence. Princeton University Press, Apr. 1976.
H. Sheikh, M. Phielipp, and L. Boloni. Maximizing ensemble diversity in deep reinforcement learning. In International conference on learning representations, 2022.
J. Smit, C. Ponnambalam, M. T. Spaan, and F. A. Oliehoek. PEBL: Pessimistic ensembles for offline deep reinforcement learning. In Robust and reliable autonomy the wild workshop the 30th international joint conference of artificial intelligence, 2021.
J. Van Amersfoort, L. Smith, Y. W. Teh, and Y. Gal. Uncertainty estimation using single deep deterministic neural network. International conference on machine learning, pages 9690-9700. PMLR, 2020.
P. Walley. Statistical Reasoning with Imprecise Probabilities. Chapman & Hall, 1991.
Y. Wen, G. Jerfel, R. Muller, M. W. Dusenberry, J. Snoek, B. Lakshminarayanan, and D. Tran. Combining ensembles and data augmentation can harm your calibration. arXiv preprint arXiv:2010.09875. 2020.
14

---

Published as a conference paper at ICLR 2026
J. Wilson, C. van der Heide, L. Hodgkinson, and F. Roosta. Uncertainty quantification with the empirical neural tangent kernel. arXiv preprint arXiv:2502.02870, 2025.
L. Wu and S. A. Williamson. Posterior uncertainty quantification in neural networks using data augmentation. In International Conference on Artificial Intelligence and Statistics, pages 3376- 3384. PMLR, 2024.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: A novel image dataset for benchmarking machine learning algorithms. CoRR, 0tos/108.00744 2017.
G. Yang and E. J. Hu. Tensor programs iv: Feature learning infinite-width neural networks. In International conference machine learning, pages 11727-11737. PMLR, 2021.
M. A. Zanger, W. Boehmer, and M. T. Spaan. Diverse projection ensembles for distributional reinforcement learning. In International conference on learning representations, 2024.
15

---

Published as a conference paper at ICLR 2026
A LINEARIZED NEURAL NETWORK LEARNING DYNAMICS
For completeness, we briefly outline sketch for how the GP interpretation of wide neural networks governed by NTK dynamics described in Expression can be obtained. This section largely follows the seminal works by Jacot et al. (2018), Lee et al. (2020) and He et al. (2020), to whom we refer readers interested in further details.
We begin by constructing first-order Taylor expansion of the neural network function f(r, 0n)
around its initialization parameters 0p:
fun (z. 0,) f(x, + Vef(I, -0)
(16)
When trained on and y with the squared error loss L fil (X;; 0,) 211 gradient flow with learning rate induces an evolution of 0g according to
d 8, aV = -a o (X (x,,,)) V.) (.
(17)
dt
In function space, this evolution translates to the expression
d 01) Vofun(z, 06) d
dt fim (z; (z;
9: = a0.( (f)m -)).
(18)
dt
where x'( = Vof(æ, Bo) Vof( (o) is the (empirical) tangent kernel of fun (2. 0z). Since this linearization has constant gradients Vof( 0o), the resulting differential equation is linear and solvable. For the substitution i = (fiin (X;; V), we obtain the training error dynamics d de U. - (A, to which an exponential ansatz yields the solution
fun(XX: 0,) - y = xx (f(X: -)),
(19)
where the matrix exponential e- was used. Plugging Eq. 19 back into Eq. at the identity
18 arrives
d fri (1;; -000=) (l)x ());;0)) y)-
(20)
dt
This differential expression is explicit in its terms such that we can obtain solution by integration through
fii (n; (r; = f(x, 0o) + (/' d
de fin )x)
(21)
= f( (0) Ou(z, (x) (1000)))) I) -0)))))))))))))))))))))))))))- -)),
(22)
which recovers Eq. 50.55for n + co. A central result by Pacot et all 2018) and extended in the linearized setting by Lee et al. (2020) that, as layer widths of the neural network go to infinity, the NTK Ov(x,22) becomes deterministic and constant and the linear approximation [i;; 0,) becomes exact w.r.t. the original function lim Jíin(x: 0,) f(z 02).
Rewriting the (infinite width) post-training test and training functions as affine transformation of the initialization yields
(1W (xi,O0))) (122,0m) (f
O( X102, (f(x2 +
(12::3) (0(x..))8(2288-1))
(23)
For the earlier described parametrization of f. the set of initial predictions is known to follow multivariate Gaussian distribution (Lee al.) 2018) described by the neural network Gaussian process (NNGP) f(X, 0o) ~ (O, x()')(') (and analogously for X2),, where
((, Es, [f (x),
(24)
Affine transformations of multivariate Gaussian random variables ~ N Ex) with y =
BX are, in turn, multivariate Gaussian random variables with distribution Y Nía
16

---

Published as a conference paper at ICLR 2026
Bpx, BExB") We here omit explicit derivations and rearrangements for brevity. As con- sequence, Eq. [3 with initialization covariance from Eq. 24 is also described by a multivariate Gaussian with mean and covariance given by
Ea. (X,, 0x))] -0,, ())
Cov(f(E, -10.
(25)
-
h.c.).
where h.c. refers to the Hermitian conjugate of the preceding term. He et al. (2020) then introduce constant "correction" terms to the function initialization described in Eq. 24 in particular such that <(x, (zx z²). This simplifies Expression 25 significantly and permits a Gaussian process interpretation with the final expression given by Eq.
B FURTHER DISCUSSIONS
Below, we further discuss the approximate nature of our method and provide a more general discussion of the terminology used in the context of this paper and the broader field of uncertainty quantification.
B.1 DISCUSSION ON APPROXIMATIONS
As our method relies on several approximations, we include a discussion that aims to provide an overview of the approximate nature of our method and in which settings it is exact or where deviations may be more likely.
The first central approximation make is to model neural networks with dynamics governed by deterministic and constant NTK. [Jacot et al. (2018) show that this is the case for fully connected NNs with NTK parametrization trained on squared loss. The implied dynamics are solved assuming gradient flow, that is with infinitesimal step sizes and full-batch gradients. Jacot et al. (2018) and Lee al. (2020) moreover show that convergence and final generalization behavior empirically well-described by wide but finite architectures including fully connected NNs, convolutional NNs and residual architectures, trained with stochastic gradient descent. The function initialization scheme proposed by He et al. (2020) allows for Gaussian process interpretation of NNs from random initialization and largely relies on the same assumptions as the above-described works.
Our theoretical motivation, outlined in Sections 3.1 and 3.2 relies the GP description of deep ensembles and the implied assumptions. Given this setting, that assuming NTK parametrization with infinite widths, function initialization according to He et al. (2020), and gradient flow with squared loss, the derivation for single-quest query ensemble variances in Section 3.1 is exact. In our contextualized model deseribed in Section 3.22
introduce an additional approximation through the introduction of an explicit context variable c, which may interfere with the training dynamics of g(x, c, 0). Let training tuples be x = (x, c) and X = {aq, ..., xN, } and let the NTK be (1r, c), (a', d')) Vag(a, c, 0o) Vag( e', 0o). The analogous regression solution to the function g(Ø, c, 8) by minimizing the loss in Eq. 15 becomes
g(æ,,, c, 02) (r(x", X O ( -(("))
(26)
A natural setting in which these training dynamics recover Eq. I is when gradients are independent between context, that is ((, c), (z, c')) 0 if + c' and maintain the gradient structure of O(z, x²) with (z((2, c), cz) = 2'), Vc E C. However, this setting would hardly permit meaningful interpolations and extrapolations between different contexts e, such that one engages in trade off between generalization capability towards general contexts e and interference in the training dynamics.
Beyond this, our practical setting approximates the NTK prior function with partial gradients as outlined in Eq. [4 of Section 3.3 The influence of this approximation choice generally depends on architecture, but we found it to perform well in our experiments using deep convolutional and residual architectures. Lastly, the RL exploration setting involves data streams rather than fixed datasets X, further deviating from the earlier delineated dynamics. Understanding the influence of this non stationarity on training dynamics is an open problem, and we believe countermeasures like promising avenue for future research.
periodic resets Oro al. 2023)
17

---

Published as a conference paper at ICLR 2026
While alleviating and quantifying the assumption stated above largely open problems in deep learning theory (Hamin and Nical 2019; 2019: Seleznova and Kutyniok] 2022 Yang and Hu [202]] Cohen et al. [021]] Lewkowycz et al. 2022). various approaches exist that aim to quantify errors w.L.t more realistic NN behavior, e.g. in the finite-width regime or with discrete gradient descent. We outline one such direction, following work by [ee al. (2020), to quantify approximation errors of the linearized NN dynamics by assuming
1. NNs of depth L to have equivalent layer-widths n2 = 122 =... = nL = n
2. a full-rank analytical NTK, i.c. Amin (O(X", (0)))) 0, Amax ( ())) < 00
3. maximum learning rate of cko I Oterit Amaa
4. the contextualized training set X is contained in the unit ball, i.e. ||0 S 1 for all
x E X² with distinct elements.
5. nonlinearities (x) to satisfy
(z)
o o'(x') x
[o(0),, I|
sup axst
< oc
(27)
x²|
6. block diagonality of the contextualized NTK as stated in Eq. (26) such that the infinite-width limit recovers g(x, c, a0) = O(z, (12,)))0(4 )0(", "
Under these conditions, we have that the wnddh NN g(x, c, 2x)) behaves as g(Ø, 50) = (zx, OX x .) and one can show that (Lee et al. 2020), with high probability over random initialization, a linearized NN (x, e, 0z) solution approximates a finite-width NN (2,, e, ax))) trained with gradient descent with (non-critical) step size ao with (1; c, 0oo) - greal 022 = O (t)) Daniely et al. (016) furthermore show that the empirical tangent kernel
at initialization O, concentrates at the same A rate, such that we have ||@u(z", x²) - O(x >>) =
O
Taken together, we can conclude that the finite-width approximation error as occurring in Eq. under the stated conditions, scales as
[V((((z,0x)) (0, g00n (z, 0oo)) (
(28)
We refer interested readers to Lee et al. 2020), Daniely et al. (0066), and Hanin and Nica 2019) for detailed derivations and further results going beyond strict NTK-regimes.
B.2 DISCUSSION ON TERMINOLOGY
While there is broad agreement that it is important to distinguish different sources of uncertainty in machine learning, there remains debate about how these notions should be formally captured and which terminology is appropriate in different contexts. A widely used conceptual distinction is between epistemic uncertainty - uncertainty arising from limited knowledge about the true or optimal model parameters and aleatoric uncertainty uncertainty inherent to the stochasticity in the data-generte seattaggerraingg process (Hüllermeier and Waegeman, 2020). Together, these sources of uncertainty are considered to constitute total uncertainty, though alternative frameworks exist that depart from such an additive decomposition (Shafer 1976] Cuzzolin [002I). Several mathematical frameworks aim formalize these notions: Bayesian inference, arguably the most prominent, capture epistemic uncertainty in the form a posterior distribution over plausible models, given a prior distribution (Neal 1996). However, alternative frameworks (c.g., frequentist statistics (Le Cam 2012) or imprecise probability (Shafer 1976] Walley| 1991] Caprio et al., (024h) provide alternative perspectives, in part to alleviate the often restrictive requirement of well-specified priors. Our work is situated within the Bayesian viewpoint in the sense that we use the term epistemic uncertainty to refer to variability induced by the posterior distribution over functions compatible with the observed data. Within this framework, Gaussian processes offer nonparametric model class that enables analytical Bayesian inference, and we use the variance of the GP posterior predictive distribution as measure of epistemic uncertainty. is important to note, however, that for probabilistic models that explicitly model observation noise (i.c., aleatoric uncertainty), the posterior predictive variance conflates epistemic and aleatoric components and therefore reflects total predictive uncertainty rather
18

---

Published as a conference paper at ICLR 2026
than epistemic uncertainty alone (Hüllermeier and Waegeman 2020). In setting, however, assume a deterministic GP model without an observation noise term, and thus interpret the posterior predictive variance as epistemic uncertainty.
C EXPERIMENTAL DETAILS
In the following, we outline details on our experimental setup. This includes hyperparameter settings, hyperparameter search procedures, algorithmic and experimental details, and dataprocessing details.
C.I HYPERPARAMETER SETTINGS
In order to facilitate comparable results, experiments conducted using central codebase and follow similar modeling choices such as architectures, optimizer, etc. where sensible. All experiments resnet-based model (He al. 2016) following the IMPALA architecture by [Epehhlt al. (2018). We optimized essential and coorithm-specife hyperparameters through a search on a selected subset of experiments.
Distribution shift detection. In the supervised distribution shift detection, architecture with 2 residual blocks and channels widths 32 and 64. Hyperparameters were searched on the FasHionMNIST dataset as the in-distribution sets. Each dataset is normalized to zero-mean and standard deviation 1 using the training set statistics. For the main classifier we apply random horizontal flips (p=0.5), random vertical flips (p=0.5) and random sized crops (zoom range between 1.0 and 1.3) to training data in all experiments. Learning rate and algorithm-specific hyperparameters were optimized independently, meaning we first performed a search for learning rates, which we used in the (if applicable) subsequent algorithm-specific parameter searches. Table 2 contains lists of searched parameters, with parenthesis indicating algorithm- specific parameters and italics indicating the parameter used during the learning rate search. The final hyperparameters were chosen based on the average AUROC metric and are reported in Table
the IMPALA
and the remaining datasets as -f-iiiiibiuutionn
VizDoom. In the RL experiments, we conducted full grid search on the My HayoepSerrsre variation of the environment and chose parameters based on performance after 5 10th steps. Our basic network architecture is based on the rainbow (Hessel et al. (0188) network proposed by SSchmitt) IMPALA (Espeholt al. 2018) (see also Fig. 4) We use residual blocks with channel widths according to Table 6 Detailed final hyperparameter settings are given in Table 5 We use the same network architecture for value functions and RND/CSD networks (up to output dimensions). All agents furthermore use a data preprocessing pipeline as outlined Table
and Schmied (2021) who in turn base their architecture
C.2 IMPLEMENTATION DETAILS
In this section, we briefly outline implementation details concerning CSD and the tested baselines.
Data augmentations For both the distribution shift detection experiments (CSD-Aug.) and the VizDoom experiments, we add data augmentation to obtain additional context variables in CSD. In both experiments, we apply augmentations with probability of p 0.25 and specific augmentations are listed in T
Data and context sampling. To compute the loss 15
we sample minibatches Xmb from a buffer or data set. Context minibatches C.b either simply reuse Xmb, are generated by applying data augmentations as outlines above, by sampling from context data set. We compute inner products over all pairings of the two batches with p(Xmb, 0tem) TT(CCm, Ocm) E RRNax Nms and compute loss 15 elementwise. Finally, we sum the average diagonal loss and the average off-diagonal loss.
Normalization.
During training, we normalize prior features ((x,0))' 1):1-1) s0(x, 10:--1)/11 /1/(, 1122 feature vectors Φ(x, Oten) context vectors (c, 0am)) (bc,, 0tm))//l///(e. lcm))l22. When computing predictive variances
ò(I, 0reat)/ll((z Oie and
19

---

Published as a conference paper at ICLR 2026
Table 2: Searched hyperparameters for distribution shift experiments.
Hyperparameter
Values
Learning rate (All)
3- 100" 10 3 10-2, 10-1
Dropout probability (MCD)
0.05. 0.1, 0.15, 0.25, 0.5
RND Learning rate (RND)
10-1
CSD Learning rate (CSD)
[00 1001:3::0000,,,:::33:0100011
10 10-1
Table 3: Searched hyperparameters for VizDoom
Hyperparameter
Values
Learning rate (all)
[1.25 10 1,0 20 10.
5- 1-10-44 6,25 10..,
Loss (all)
Huber, C51]
Prior function scale (BDQN+P, IDS)
1.0,3.0,550)
Initial bonus 3 (RND, CSD)
[0.05, L.,, 10.0]
RND Learning rate (RND)
1:25-100::25--1- 10-4
5-10:
CSD Learning rate (CSD)
3.2
5 5-10:2:225-1--775--100
Table 4: Hyperparameter settings for distribution shift experiments.
Hyperparameter
MCMC
Laplace
MCD
ENS
RND
CSD
Main Classifier Network
Learning rate
10-3
10-3
3 10-4
10-3
10-3
10-3
MLP hidden layers
2
MLP layer width
256
Channel Widths
32, 64
RND/CSD Network
Learning rate
n/a
10-3
10-2
MLP hidden layers
n/a
2
2
MLP layer width
n/a
256
256
Channel Widths
n/a
16
32
Target hidden layers
n/a
Output dimensions
n/a
256
256
Ensemble size
n/a
n/a
n/a
3, 15
n/a
n/a
Dropout rate
n/a
n/a
0.1
n/a
Prior Precision
100
n/a
n/a
Posterior Temperature
1.0
1.0
n/a
n/a
Posterior Samples
30
30
100
n/a
Epochs per sample
2
n/a
n/a
n/a
Burn-In Epochs
10
n/a
n/a
n/a
Adam epsilon
n/a
1
10
10
Learning rate anneal
Linear
Batch size
256
Initialization
Orthogonal Saxe et al.
2013)
20

---

Published as a conference paper at ICLR 2026
FC 256
ReLU
N Layers
FC 256
Residual Block
ReLU
Conv. 3 x stride 1
Residual Block
ReLU
Residual Block
x3
Conv. x stride
(32, 64, 64] ch.
Max. 3 3, stride 2
ReLU
Conv. x 3, stride
/55
42 x 42
Figure 4: Illustration of the architecture for VizDoom environments. Based on the architecture used by Espeholt et al. (2018).
Room 15 suarse
Goal
Room 17 verg spores
Figure 5: Map for the VizDoom MyWayHoee environment. Agents spawned in the sparse and very sparse locations to vary the exploration difficulty.
21

---

Published as a conference paper at ICLR 2026
Table 5: Hyperparameter settings for VizDoom experiments.
Hyperparameter
DQN
BDQN+P
RND
IDS
CSD
Adam Learning rate
2.5
10-4
2.5
10
6.25
10
2.5
10
6.25 10-4
Prior function scale
n/a
1.0
n/a
1.0
n/a
Heads K
1
101
1/101
101/101
Ensemble size
n/a
10
n/a
10/1
n/a
Initial bonus Siii
n/a
0.1
0.1
Final bonus Oriaal
n/a n/a
0.01
0.01
0.01
Bonus decay frames
n/a Huber Huber
n/a
3.3 106
3.3 10
3 106
Loss function
C51
Huber/C51
C51
Channel Widths
32, 32, 64
MLP hidden layers
1
MLP layer width
256
RND / CSD Network Parameters
Adam Learning rate
n/a
n/a
2.5
10
n/a
10-4
Channel Widths
n/a n/a
16, 16,
32
n/a
16,
16, 32
MLP hidden layers n/a n/a
3
n/a
3
(main)
MLP hidden layers n/a n/a
n/a
n/a
1
(context)
MI.P hidden layers (prior)
n/a
n/a
MLP layer width
n/a
n/a
256
n/a
256
Target hidden layers n/a
n/a
1
n/a
1
Output dimensions
n/a n/a
256
n/a
256
Initial in e-greedy
1.0
Final é in e-greedy
0.01
decay frames Training starts
500, 000
100, 000
Discount
0.997
Buffer size
1, 0000,, 000
Batch size
256
Parallel Envs
16
Adam epsilon
0.005 /btch size
Initialization
He uniform
He et al. 2015
Gradient clip norm
10
Regularization
spectral normalization (Gogiauu et al.
2021
Double DQN
Yes (Hassell 2010
Update frequency
Target lambda
1.0
Target frequency
8000
PER 3o
0.45 (Schaul et 2016
n-step returns
Table 6: VizDoom Preprocessing
Parameter
Value
Grayscale
Yes
Framc-skipping
No
Frame-stacking
Resolution
42 x 42
Max. Episode Length
2100
22

---

Published as a conference paper at ICLR 2026
Table 7: Data augmentations for context data.
Distribution Shift
VizDoom
RandomHoorizontllllipp = 0.25)
RandommPerspectivvepp 0.5)
Random Random Verticallllicc 0.25)
RandommHorizoonallPippo 0.5)
Rotate(p = 0.25)
RandomResizeeCCoppt = [0.75, 1.0)
GaussianBlur(o 1.0, p p= 0.25)
RandomResszzeCropp = [0.75, 1.0)
Raigtmmrrightnesstrr -1.0, p= 0.5)
RandomContrast(f = [-1.0, 1.0), p= 0.5)
inference time, we rescale by
(if (x)) || (0 01 (F (0:0)) (,,)) 012,0001111 00:1111) - Ö(z, 0 T (c, dcas)), (29)
to obtain predictions in the original scale again.
Small function initialization. While our theoretical suggests using small function initialization with g( 0g) = 0, Vr, preliminary experiments with a reparametrization ÿ(x, 0e) = gxx, 0²) gx,, 0o) showed no significant differences. In our main implementation we thus refrain from using this reparametrization in favor of simplicity.
Environment Details. We conduct experiments on three variations of the VizDoom VizDoom environment My MyWayHoee. A top-down view of environment map shown in Fig. [5] In the dense setting, at the beginning of each episode agents are spawned in random positions of the map, such that the goal position is encountered stochastically without requiring coordinated exploration. The sparsity of the problem is increased by changing the agents spawning location deterministically to a room further from the goal position, that is Room 13 for the sparse setting and Room 17 for the very sparse setting. As described in Section 4 the reward function is sparse. A constant reward of -1 * 10-1 given every timestep and a reward of is given for reaching the goal. Episodes are limited length of 2100 timesteps.
Reinforcement Learning Implementation. We outline the basic implementation of our DQN- based RL agent. The agent follows the same algorithmic flow as the established DQN-algorithe (Mnih et al. 2015) and subsequent variations (Hessel et al. 2018: Schmidt and Schmied, (021). The agent maintains a replay buffer transitions, from which we sample minibatches of transition Noo
Xmh = {8i, as, s/, Tep where T; are terminations. Q-networks are then updated at a fixed frequency using the sampled minibatch. As is established, use target networks with slow-moving parameters for value learning.
We provide intrinsic rewards as generated by CSD to the DQN agent to incentivize exploration. For all our experiments including intrinsic rewards (CSD and RND), we use separate value functions for the intrinsic reward and employ intrinsic reward priors, mechanism suggested by Zanger et al. (2024) which includes intrinsic rewards to the forward pass of the value network. This addresses a common issue with intrinsic reward learning as described previously by Rashid al.] (2020) by preventing underestimation of unseen actions. Specifically, intrinsic reward priors redefine the forward pass of the intrinsic Q-function according to
(,,(s, 0, 0m) Qin(s, 0) + i a, Onn),
where (s, a, 0m) denotes an intrinsic reward term, in experiments generated by either RND or CSD with parameters 0m. The altered function Qin a, 0..) is then used drop-in replacement for the Q-function in the used algorithm.
Pseudocode for Reinforcement Learning Experiments. We provide pseudocode for DQN agent with CSD in Algorithm [ For clarity, omit standard algorithmic details such double Q-learning distributional value functions, prioritized experience replay, separate value functions for intrinsic reward, and intrinsic reward priors.
23

---

Published as a conference paper at ICLR 2026
Algorithm 1 CSD-DQN
1: initialize CSD model g(8, a, 8c; Cc. 0.) = (58 a, 0e) (b(8e, ae 0e) with 6o-
2: initialize CSD prior (s, a, Sc: Oc, 0p) = (0 0p) (s(8, dor Op) with Op.
3: initialize Q-function Q(s, a, 0e) with 0g and target parameters 0o.
4: sample initial state So from the environment.
.......'' 6: take action < arg max A d')3
5: for = 1,...,77 7' do
7: obtain observations (st, TT) from the environment.
8:
store samples (, 0e-1, re, St, Tt).
9:
sample transition tuple {5t, a,, rs, s Ti B from buffer
10:
sample context tuple {8;, â;, fi, 8 Tile
Nob ~ from buffer
11:
generate intrinsic reward Fin ( (s, s,, 02) g( 84, a, 0z). generate next action a == arg max, {Q(s) a' a', 0,)).
12:
13:
update Q-function 0, 02 Vo, C(,,) with
Nmb
[(0.) 2 Σ
(ra + B ri + Q(si, ai, B.) Q66 d/ 0,)))
14:
update CSD model 0e t 0, VE, ((0.) with
Nan
((O,) 2N > (g(sis ais Si, ai, 0.) 0* ais Sis ai 0p))²
15:
if 1. % freq 0 then
16:
update target parameters 0, + 10g + (1 x)0c
end if
18: end for
Figure 6: Left: Original Image. Right: Perturbed OOD Image.
C.3 ADDITIONAL EXPERIMENTAL RESULTS
We report the detailed results of our distribution shift detection tasks. Tables 8 to show OOD detection metries for the datasets FashionMNIST. MNIST, NotMNIST, and KMNIST. Each table shows the test accuracy and average AUROC, AUPR-IN and AUPR-OUT scores against the remaining three training datasets and an additional perturbed dataset. The perturbed dataset constructed by applying data augmentations the ID dataset. In our experiments,
random brightness changes (p 1.0, r = [-1.0, 1.0)), random contrast changes(p 1.0, = [-1.0, 1.0)), , and randomly set patches of an image to zero (p- 1.0, r = 1.0, 1.0)). Fig. shows example of this.
Runtime analysis We report runtime comparisons for all algorithms used in the distribution shift detection tasks in Fig. 7 All efficient uncertainty estimation methods, including ours, run faster than the smallest ensemble tested (ENS(3)), while incurring some overhead compared to a single-model baseline (ENS(1)), which does not come with built-in uncertainty quantification capability. Among these, CSD has slightly higher runtime, which is consistent with the algorithmic structure of our method that comprises separate feature and context model. We note that this runtime comparison is not entirely one-to-one: different methods vary their application mechanism, with Laplace and MCD being mainly post-hoc methods, while others devise dedicated learning algorithms.
24

---

Published as a conference paper at ICLR 2026
Table 8: Distribution Shift Detection. FashionMNIST as ID dataset.
Method
Acc.
AUROC
AUPR-IN
AUPR-OUT
MCD
89.24 + 0.36
82.23 + 0.48
79.88 + 0.75
83.01 + 0.34
BNN-MCMC
85.73 0.24
85.01 0.62
85.16 0.68
83.38 + 0.62
BNN-Laplace
88,57 = 0.80
86.50
86.32
85.95
RND
91.90 = 0.15
93.93 - 0.73
93.45
93.64
ENS(3)
92.90 = 0.09
88.90 = 0.20
89.63
88.16 = 0.20
ENS(15)
93.33 = 0.06
91.93 = 0.12
92.83
91.09 0.12
CSD
91.93 = 0.17
96.18 = 0.67
96.49 + 0.74
95.74 0.62
CSD-Aug.
91.92 0.16
97.84 0.30
98.24 0.27
97.34 0.31
CSD-OOD.
91.96 + 0.13
97.35 + 0.50
97.87 + 0.45
96.72 + 0.56
Table 9: Distribution Shift Detection. MNIST as ID dataset.
Method
Acc.
AUROC
AUPR-IN
AUPR-OUT
MCD
98.97 = 0.06
90.03 + 0.23
87.70 + 0.38
89.01 + 0.32
BNN-MCMC
94.29 + 0.39
80.24
80.20 2.05
77.33 2.56
BNN-Laplace
94.17 + 1.01
74.05 1.70
1.90
74.39
RND
99.85 = 0.02
93.8
04.25:0..3
ENS(3)
99.95 = 0.01
+.22
92.32 0.31
ENS(15)
99.97 = 0.00
95.33 = 0.06
96.31 0.06
03.79
CSD
99.88 + 0.01
96.78 + 0.58
96.96 + 0.72
96.25 + 0.57
CSD-Aug.
99.87 0.02
98.39 0.17
98.63 0.20
97.94 0.19
CSD-OOD.
99.87 = 0.02
99.37 = 0.08
99.51 + 0.07
99.14
Table 10: Distribution Shift Detection. NotMNIST as ID dataset.
Method
Acc.
AUROC
AUPR-IN
AUPR-OUT
MCD
95.17 0.14
83.21 0.45
75.86 0.89
85.73 + 0.18
BNN-MCMC
90.20 0.44
87.05
85.93
87.68 0.63
BNN-Laplace
95.29 - 0.52
86.38
8
$7.55
RND
96.25
0.12 0.12
95.49 0.82
95.81
95.23 1.74
ENS(3)
92.37
92.11 0.30
91.93 - 0.27
ENS(15)
97.47 0.05
94.26 1 0.17
93.29 0.17
CSD
96.48 + 0.08
96.98 - 0.41
97.26 0.44
96.86 + 0.36
CSD-Aug.
96.45 0.09
98.51 0.22
98.70 0.24
98.31 0.21
CSD-OOD.
96.49 + 0.10
98.49 + 0.35
98.78 + 0.29
98.21 + 0.42
Table 11: Distribution Shift Detection. KMNIST as ID dataset.
Method
Acc.
AUROC
AUPR-IN
AUPR-OUT
MCD
94.18 + 0.26
87.22 0.75
83.48 + 0.74
88.00 + 0.77
BNN-MCMC
80.57 1.29
80.40
79.31
80.75
BNN-Laplace
85.39 + 1.79
78.58 2.66
3
79.47 + 2.49
RND
96.73 = 0.21
93.50 + 1.17
92.93 1.05
ENS(3)
97.68 0.10
93.88 0.24
94.49 0.26
93.05 0.24
ENS(15)
97.96 = 0.06
94.68 = 0.11
95.39 0.12
13.81+0.1
CSD
96.89 + 0.13
96.57 + 0.73
97.05 + 0.74
95.90 + 0.74
CSD-Aug.
96.90 0.19
98.12 0.46
98.45 0.41
97.61 + 0.53
CSD-OOD,
96.86 + 0.12
99.06 - 0.19
99.30 0.14
98.71 0.25
25

---

Published as a conference paper at ICLR 2026
Runtime Analysis Distribution Shift Detection
40000
Total Runtine (sec)
30000
20000
10000
ENS(35)
ENS(15)
ENS(3)
CSD
RND BNN-MCMC
MCD
BNN-Laplace
ENS(1)
Figure 7: Runtime all algorithms on the Distribution Shift Detection task. Runtimes are reported in seconds for one seed completion on single Nvidia RTX 3060 GPU.
USE OF LARGE LANGUAGE MODELS
Large language models (LL.Ms) used assist in the preparation of this paper. Their usage was limited to refining sentence structure and verifying grammar, punctuation, and general language usage. No content or substantive research contributions were generated by LLMs.
26


Reviews:
[{"review_id": "PJ3BWd9PYn", "reviewer": "Reviewer_8hVE", "summary": "This paper presents an approach to estimate the uncertainty of a neural network prediction through contextual similarity distillation. The method is built upon the Neural Tangent Kernel theoretical results, which express a probability distribution of a neural network function after training based on the distance between a test sample and the training sample as measured by a gradient-based kernel (NTK). Using a clever choice of labels, they train a neural network to predict the uncertainty that would result of an ensemble.", "questions": "Questions\n* In section 3.1, it is said that $g$ has the same architecture as $f$, therefore $\\Theta_g(x, x’) = \\Theta(x, x’)$. In the text earlier, there is no justification for this: the closer argument I saw was that $\\Theta(x, x’)$ is equal for every stage of training of weights $\\theta_f$ under the assumption of gradient flow. Is the two networks having the same architecture sufficient for this equality to hold? If so, could you introduce this somewhere in section 2.2?\n\n* Figure 1, the legend describes the “Kernel Prior $\\Theta(x,x)$”, yet the function changes across the three subfigures, where $x_t$ is the element that changes. Should it be $\\Theta(x, x_t)$?\n* Section 3.3: is the feature vector $\\phi(x, \\tilde{\\theta}_f)$ related to parameterized function $f$? If so, what exactly from $f$ does it refer to?\n* Equation 14, is the RHS calculated on the trained $f$ or on the initial $f$ parameters? (This also applies to equation 13 too if $\\phi(x, \\tilde{\\theta}_f)$ is also based on $f$). \n* If it is on the trained parameters, how does that work when it is estimated during a continuous training process, such as in RL?\n\nComments:\n* Introduction, Section 3.1: It is said often in the paper that, in the field of RL, there are often large models or datasets. Depending on the RL application, this may not be true: most applied RL uses small neural networks, and datasets are created from experience (which can be sparse/expensive). Maybe these claims could be more nuanced?\n* Section 2: The objective of finding the optimal policy is not RL specific (this is the case for every MDP solving method, s.a. dynamic programming or model-predictive control). The RL specific challenge is to do so through learning on trajectories acquired in the environment.\n* Section 2.1: The fundamental challenge of exploration is specific to the case of online RL (to differentiate from offline RL)  \n* Insufficient upper margin for eq. (3)\n* End of section 2.1: Due to their initialization, ensemble members are not great at the beginning of the RL process to determine the model’s uncertainty: that’s where randomized prior functions (RPF, Osband et al., 2019) are helpful.\n\nTypos:\n* Introduction: “distribution shift detection tasks(Van Amersfoort” and “NOTMNIST datasets(Xiao et al.”: space missing before the “(“\n* Section 2: “agents […] subsequently receiveS the immediate reward […] and observeS the next state”: there should not be an S as agents is plural\n* Section 3.1: “it is our goal IS to estimate”: correct the sentence", "limitations": "", "rating": 10, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper brings in an interesting and novel approach to uncertainty estimation. It is strongly anchored in theoretical analysis (although with several assumptions). It is also more practical than current approaches, as it requires significantly less compute, and it performs well on different datasets compared to baselines.\nThe paper does a good job at presenting the concept development in a pedagogical and self-contained manner, which is necessary for anyone who is not familiar with the NTK literature.\nWhile some of the assumptions are strong, the discussion section in the Appendix about approximations is relevant and important.\nI really appreciated Figure 2 - it demonstrates clearly and convincingly the method’s objective and performance.\nI believe this is an impressive piece of work which could have a lasting impact on uncertainty estimation with neural networks.", "weaknesses": "I did not find important weaknesses in this work. Still, elements could be improved.\n\nWhile the experimental results on image classification are good and convincing, I believe other types of tasks (regression, for example) could strengthen the article and open the use of this approach for a wider range of tasks. \n\nRegression uncertainty estimation is what is actually showed in the RL environment, but very indirectly. The RL experiment is one of the least strong aspect of this paper. While it indirectly supports the claim about CSD generating good uncertainty estimations, it is not sufficiently rigorous to demonstrate the usefulness of the approach in RL. For a stronger case, more details about the implementation and other environments would be needed. More discussion too - for example, it is surprising that CSD works better than Bootstrap DQN where CSD is supposed to approximate ensembles with lower costs and Bootstrap DQN actually uses ensembles.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZuRmHMv3cC", "reviewer": "Reviewer_ExXT", "summary": "This paper proposes contextual similarity distillation (CSD), a method which approximates an ensemble of deep networks using a single model. The key idea behind this method is motivated by neural tangent kernels (NTK), which describes how similar the gradients of two inputs are with respect to the network parameters. For a network layer with infinite width, the NTK becomes deterministic and constant (allowing it to be expressed as a Gaussian process). CSD estimates the NTK by training with regression targets corresponding to kernel similarities between data points. Empirical results support theoretical claims against several baseline methods.", "questions": "1. Can the proposed idea of contextual regression be extended to capture aleatoric uncertainty as well, or is it primarily limited to estimating epistemic uncertainty?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Strong theoretical grounding. The proposed method builds off of well-studied prior literature (e.g., NTK and Gaussian processes),\n2. Reduced computational cost. The proposed method is able to approximate a deep ensemble using a single model, avoiding large training and inference costs. This is supported by the experimental results on tasks like OOD detection and RL navigation.\n3. Novelty. The idea to use kernel regression to analytically approximate an ensemble is novel. Furthermore, the addition of a context variable to allow the method to work on arbitrary query points.", "weaknesses": "1. Empirical results. The experiments are limited to small scale datasets like VizDoom and FashionMNIST with well-studied, but somewhat outdated, baselines. Inclusion of larger / more complex datasets and comparison against more recent methods would considerably strengthen the paper and give more context on the scalability of CSD.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I4vMcTWpmM", "reviewer": "Reviewer_ccSm", "summary": "This paper proposes contextual similarity distillation (CSD), a single-model method that tries to estimate the variance of a random initialization ensemble of DeepEnsembles by using the kernel similarities in the covariance in the Neural Tangent Model. Empirical results show CSD performs competitively and sometimes superior to ensemble-based baselines in OOD detection and sparse-reward RL environments.", "questions": "1. How does CDS perform on larger-scale datasets such as CIFAR-100 and ImageNet?\n\n2. In a large-scale dataset, where the feature dimension is higher, how about the training/inference efficiency of CDS get impacted? How is the ensemble variance estimator impacted (can the authors formally show the estimator error in this regard?)?\n\nReferences:\n\n[1] He et al., Bayesian Deep Ensembles via the Neural Tangent Kernel, NeurIPS, 2020.\n\n[2] Sergio et al., Epistemic Uncertainty and Observation Noise with the Neural Tangent Kernel, NeurIPS, 2022.\n\n[3] Seijin et al., Disentangling the Predictive Variance of Deep Ensembles through the Neural Tangent Kernel, NeurIPS, 2022.\n\n[4] Nado et al., Uncertainty Baselines: Benchmarks for Uncertainty & Robustness in Deep Learning, Bayesian Deep Learning workshop, NeurIPS 2021.\n\n[5] Yang et al., OpenOOD: Benchmarking Generalized Out-of-Distribution Detection, NeurIPS 2022 Datasets and Benchmarks Track.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is generally well-written to understand the key aspects of the proposed algorithm.\n- The direction of this paper is important because DeepEnsembles' superiority in uncertainty and robustness, but struggle with computational efficiency.\n- Empirical results are quantitatively shown in the experiments and qualitatively shown in Figures 1 and 2 to understand how the method actually works.", "weaknesses": "- Using NTK to get epistemic uncertainty is not new and has been studied in several works [1,2,3]. This limits the novelty contribution of CDS. Note that I still appreciate your contribution to the similarity distillation.\n- There is no theoretical contribution in this paper to formally explain why and how CDS works. \n- The experiments lack several computational efficiency baselines in uncertainty and robustness (e.g., BatchEnsembles, Rank1-BNN, etc. [4]) and OOD detection baselines (ReAct, DICE, etc. [5]).\n- The main contribution is improving computational efficiency, but there is no computational efficiency evaluation (training, inference speed, GPUs, etc.) in the experiments.\n- Contribution to uncertainty also needs to improve beyond OOD detection evaluation, such as Expected Calibration Error (in IID and OOD data), sharpness, etc.\n- In writing, background in Reinforcement Learning, such as MDP, Q-function, policy improvement objective, are not related to the main method in Section 3 (many mathematical notations are not reused at all). Note that I still appreciate your additional RL experiments, but I would suggest moving some RL background in Section 2 to the Appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KqJXgRBVZB", "reviewer": "Reviewer_Bv9r", "summary": "The authors propose  contextual similarity\ndistillation, an approach that explicitly estimates the variance of an ensemble\nof deep neural networks with a single model, without learning or evaluating\nsuch an ensemble in the first place. It can estimate predictive variance at inference time with a single forward pass,\nand can make use of unlabeled target-domain data or data augmentations to refine\nits uncertainty estimates.", "questions": "In line 86, shouldn't it be $\\mu \\in \\mathscr{P}(\\mathcal S)$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clear, well-written, solves a cogent problem, and presents convincing experimental evidence.", "weaknesses": "The paper only has a conceptual weakness, that is widespread in the ML literature.\n\nIf we posit that epistemic uncertainty (EU) is encoded in the parameter distribution, then a single posterior predictive is not able to gauge EU correctly, since it is the expectation of the likelihood taken wrt the posterior, $p(\\tilde{y} | \\tilde{x},D) = \\int_\\Theta p(\\tilde{y} | \\tilde{x},\\theta) p(\\theta | D) \\text{d}\\theta = \\mathbb{E}_{\\theta \\sim p(\\theta | D)}[p(\\tilde{y} | \\tilde{x},\\theta)]$, hence all the uncertainty encoded in the posterior gets ``washed away'' by taking the expectation. \n\nThis is not the case in ensembles, where we have different posterior predictives, but if we then distil all the models into only one, and estimate the variance of the latter thinking that it captures the EU, we incur the same problem as above. In symbols, we would have $p(\\tilde{y} | \\tilde{x},D) = \\sum_{k=1}^K w_k \\int_\\Theta p_k(\\tilde{y} | \\tilde{x},\\theta) p_k(\\theta | D) \\text{d}\\theta = \\mathbb{E}_{k} [p_k(\\tilde{y} | \\tilde{x},D)]$, where the $w_k$'s are the weights that the ensemble method assigns to each model $k = 1,\\ldots,K$ (which of course are nonnegative and sum up to $1$).\n\nThe way to avoid this is to keep the posterior predictives separate, e.g. by considering their induced convex hull, and use other metrics existing in the literature (e.g. in https://link.springer.com/article/10.1007/s10994-021-05946-3) to quantify EU (and also aleatoric uncertainty). This was already done in https://openreview.net/forum?id=4NHF9AC5ui, and applied to active learning (thus very close to RL) in https://dl.acm.org/doi/10.1145/3716863.3718040, where EU is used to inform the exploration of the state space. \n\nWhile the paper solves a very relevant problem, in a creative and innovative way, I urge the authors to (a) either avoid referring to epistemic uncertainty, and rather just call it uncertainty (in which case my previous point can be ignored), or (b) to add even a small discussion where they acknowledge my point above, and state that they keep the term \"epistemic\" only to conform to the existing literature. \n\nMoreover, I suggest the authors, in future work, to look into the techniques developed by the imprecise probabilistic machine learning field, that may well turn out useful for the future endeavor pointed out in the conclusion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)