# Contextual Similarity Distillation: Ensemble Uncertainties with a Single Model

arXiv: https://arxiv.org/abs/2503.11339 · HTML: https://mlsanity.com/paper/2503.11339/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 51 of its 68 references that are in the map.
- First posted on arXiv 14 Mar 2025.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 51)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [Unifying Count-Based Exploration and Intrinsic Motivation](https://mlsanity.com/paper/1606.01868/) (2016)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [Deep Exploration via Bootstrapped DQN](https://mlsanity.com/paper/1602.04621/) (2016)
- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Is Q-learning Provably Efficient?](https://mlsanity.com/paper/1807.03765/) (2018)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [Deep Learning for Classical Japanese Literature](https://mlsanity.com/paper/1812.01718/) (2018)
- [What Are Bayesian Neural Network Posteriors Really Like?](https://mlsanity.com/paper/2104.14421/) (2021)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Randomized Ensembled Double Q-Learning: Learning Fast Without a Model](https://mlsanity.com/paper/2101.05982/) (2021)
- [ViZDoom: A Doom-based AI Research Platform for Visual Reinforcement Learning](https://mlsanity.com/paper/1605.02097/) (2016)
- [Deep Exploration via Randomized Value Functions](https://mlsanity.com/paper/1703.07608/) (2017)
- [Bayesian Neural Networks: An Introduction and Survey](https://mlsanity.com/paper/2006.12024/) (2020)

## Built on by (most cited first, 3 of 3)

- [Epistemic Artificial Intelligence is Essential for Machine Learning Models to Truly 'Know When They Do Not Know'](https://mlsanity.com/paper/2505.04950/) (2025)
- [Universal Value-Function Uncertainties](https://mlsanity.com/paper/2505.21119/) (2025)
- [On the Equivalence of Random Network Distillation, Deep Ensembles, and Bayesian Inference](https://mlsanity.com/paper/2602.19964/) (2026)

## Abstract (arXiv)

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.