Abstract

Uncertainty quantification is a critical aspect of reinforcement learning anddeep learning, with numerous applications ranging from efficient explorationand stable offline reinforcement learning to outlier detection in medicaldiagnostics. The scale of modern neural networks, however, complicates the useof many theoretically well-motivated approaches such as full Bayesianinference. Approximate methods like deep ensembles can provide reliableuncertainty estimates but still remain computationally expensive. In this work,we propose contextual similarity distillation, a novel approach that explicitlyestimates the variance of an ensemble of deep neural networks with a singlemodel, 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 ofthe predictive variance of an infinite ensemble. Specifically, we reinterpretthe computation of ensemble variance as a supervised regression problem withkernel similarities as regression targets. The resulting model can estimatepredictive variance at inference time with a single forward pass, and can makeuse of unlabeled target-domain data or data augmentations to refine itsuncertainty estimates. We empirically validate our method across a variety ofout-of-distribution detection benchmarks and sparse-reward reinforcementlearning environments. We find that our single-model method performscompetitively and sometimes superior to ensemble-based baselines and serves asa reliable signal for efficient exploration. These results, we believe,position contextual similarity distillation as a principled and scalablealternative for uncertainty quantification in reinforcement learning andgeneral deep learning.