Epistemic Neural Networks

Ian Osband; Zheng Wen; Seyed Mohammad Asghari; Vikranth Dwaracherla; Morteza Ibrahimi; Xiuyuan Lu; Benjamin Van Roy

Source-linked AI summary

Epistemic Neural Networks

Ian Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, Benjamin Van Roy

arXiv:2107.08924v8Jul 19, 2021cs.LGcs.AIstat.ML

TL;DR

The paper addresses how neural networks can distinguish what they know from what they do not know without the prohibitive cost of large ensembles. It introduces epinets as lightweight uncertainty supplements and ENNs as a broader interface, showing that epinets can outperform ensembles of hundreds of particles at only slightly more than one particle’s computational cost.

Problem
Conventional marginal predictions do not distinguish genuine ambiguity from uncertainty that additional data could resolve, while large ensembles make effective joint predictions computationally expensive.
Method
The paper introduces an epinet that supplements a conventional neural network with uncertainty estimates and defines ENNs as an interface encompassing BNNs and non-BNN approaches.
Results
Epinets produce joint predictions that outperform ensembles of hundreds of particles at a computational cost only slightly more than one particle.
Takeaways & Limitations
Joint-prediction quality provides a basis for evaluating uncertainty methods beyond BNNs, including additions to large pretrained models with modest incremental computation.
Takeaways & Limitations
The theoretical analysis considers a linear-Gaussian epinet with Gaussian noise and context vectors sampled uniformly from a unit sphere.

Abstract

from arXiv · show

Intelligence relies on an agent's knowledge of what it does not know. This capability can be assessed based on the quality of joint predictions of labels across multiple inputs. In principle, ensemble-based approaches produce effective joint predictions, but the computational costs of training large ensembles can become prohibitive. We introduce the epinet: an architecture that can supplement any conventional neural network, including large pretrained models, and can be trained with modest incremental computation to estimate uncertainty. With an epinet, conventional neural networks outperform very large ensembles, consisting of hundreds or more particles, with orders of magnitude less computation. The epinet does not fit the traditional framework of Bayesian neural networks. To accommodate development of approaches beyond BNNs, such as the epinet, we introduce the epistemic neural network (ENN) as an interface for models that produce joint predictions.

1 Introduction

Conventional neural networks cannot tell whether uncertain predictions reflect genuine ambiguity or insufficient data, while joint predictions can make that distinction. The paper introduces epinets and ENNs to improve joint uncertainty predictions beyond computationally expensive ensembles and outside the traditional BNN framework.

Motivation: Joint predictions distinguish genuine ambiguity from uncertainty that additional data could resolve, unlike conventional marginal predictions.For identical inputs, conditioning on one label can either leave the other uncertain or make it predictable, revealing whether uncertainty is resolvable.
Motivation: BNNs provide principled joint predictions, but practical ensemble implementations are often limited to ten or fewer particles by computational cost.Larger ensembles improve joint-prediction quality, but training many statistically plausible networks becomes expensive for large models.
Contribution: The epinet supplements conventional neural networks to estimate uncertainty, improving joint predictions without materially changing marginal prediction quality.It uses privileged base-network features and an epistemic index, then adds its output to the base network with a stop gradient.
Contribution: On ImageNet, the epinet-enhanced ResNet outperforms a 100-particle ensemble on joint log-loss with total parameters below two particles.The cited comparison concerns joint prediction quality; marginal classification error and marginal log-loss change little after adding the epinet.
Contribution: The paper introduces ENNs as an interface for uncertainty models, including useful approaches such as epinets that are not Bayesian neural networks.All BNNs are ENNs, but ENNs need not represent distributions over base-network parameters.

2 Related work

Prior Bayesian deep-learning work emphasizes posterior inference and diverse evaluation targets, but computational cost and the lack of unified evaluation remain challenges. The paper motivates joint-prediction quality as a common basis for evaluating uncertainty methods beyond BNNs.

Bayesian deep learning: Bayesian neural networks represent epistemic uncertainty through approximate posterior inference over base-network parameters, whose computational cost becomes prohibitive at scale.The cited work notes that even approximate inference methods can become impractical for large models.
Efficient uncertainty methods: Ensemble-based BNNs approximate posterior distributions with independently trained particles, so computational cost scales with the number of particles.Related approaches reduce cost by ensembling only parts of networks or by using architectures designed to incorporate uncertainty estimates.
Evaluation: Bayesian deep-learning research has emphasized varied targets such as accuracy, calibration, out-of-distribution performance, and active learning rather than unified evaluation.These improvements can also be obtained by methods that do not approximate posterior distributions, motivating more direct evaluation criteria.
ENN perspective: The ENN perspective evaluates methods through joint predictions because they are linked to effective decision making, exploration, and adaptation.This criterion permits approaches beyond the traditional BNN framework and provides a basis for comparing their predictive quality.

3 Epistemic neural networks

ENN architectures augment conventional neural networks with an epistemic index to produce expressive joint predictions that distinguish reducible from irreducible uncertainty. The framework evaluates these predictions through joint log loss and generalizes beyond Bayesian neural networks.

3 Epistemic neural networks: An ENN extends a conventional network with an epistemic index z sampled from a reference distribution PZ, allowing predictions to vary across epistemic states.Typical reference distributions include finite uniform distributions and standard Gaussians.
3 Epistemic neural networks: Joint predictions assign probabilities to combinations of labels across multiple inputs, while multiplying marginals imposes independence and cannot distinguish ambiguity from insufficient data.The distinction matters because good downstream decisions can require accurate joint predictions.
3 Epistemic neural networks: Integrating predictions over z creates joint distributions that need not equal products of marginal predictions, enabling ENNs to represent dependencies across inputs.This integration distinguishes inevitable ambiguity from uncertainty that additional data could resolve.
3.1 Evaluating ENN performance: Small marginal log loss can still yield decisions no better than random guessing, whereas small joint log loss guarantees near-optimal expected reward.Theorems 1 and 2 motivate evaluating ENNs with joint predictions rather than marginal predictions alone.
3.2 ENNs versus BNNs: For a base network, every BNN can be expressed as an ENN, but some ENNs cannot be expressed as BNNs with that same base network.The distinction concerns whether the architectures can represent the same distributions over functions and joint predictive distributions.
3.2 ENNs versus BNNs: BNNs approximate uncertainty through distributions over base-network parameters, whereas ENNs focus on representing uncertainty over functions and can provide different computation–prediction trade-offs.The epinet is presented as an ENN that lies outside the BNN framework for the same base network.

4 The epinet

The epinet supplements a conventional neural network with an indexed uncertainty component, producing an ENN that can express epistemic uncertainty through output variation across indices. It uses standard components and losses, while theoretical analysis shows convergence toward optimal joint predictions in a linear-Gaussian setting.

4.1 Architecture: The epinet adds an index-dependent network to a conventional neural network, allowing output variation across epistemic indices to represent uncertainty.The epinet receives selected base-network features and an epistemic index, then adds its output to the base prediction with a stop gradient.
4.1 Architecture: Multiple indexed predictions can share one base-network forward pass, yielding computational savings when the epinet is much smaller than the base network.The additive architecture permits multiple ENN samples without recomputing the base network.
4.1 Architecture: A fixed prior network supplies prior uncertainty, while a learnable epinet is trained to produce statistically plausible predictions across likely index values.Variation across indices after combining the learnable and prior components represents predictive epistemic uncertainty.
4.2 Training loss function: The epinet uses standard regularized marginal losses, with prior functions and perturbations providing a route to effective joint predictions.The paper analyzes this mechanism through perturbed losses and relates it to prior-function approaches.
4.3 How can this work?: In a linear-Gaussian model, an epinet trained with the specified procedure converges in distribution to the target function as the index dimension grows.The paper presents this theorem as a sanity check for approaching optimal joint predictions with standard loss training.

5 The neural testbed

The Neural Testbed evaluates joint-prediction quality in controlled synthetic classification problems. Across benchmark agents, epinets retain similar marginal performance while achieving stronger joint predictions at lower computational cost.

5 The neural testbed: The Neural Testbed uses synthetic data from neural-network-based generative models to unit-test joint-prediction algorithms in a controlled environment.Its evaluation targets quality of joint predictions in classification problems.
5 The neural testbed: The epinet provides better joint predictions than benchmark agents at lower computational cost, while tuned agents perform similarly on marginal prediction.The comparison covers ensembles, ensemble+, bbb, dropout, hypermodel, and sgmcmc agents.
5 The neural testbed: The epinet outperforms an ensemble with 100 particles using a model less than twice the size of a single particle.The result is reported after tuning on the Neural Testbed.

6 ImageNet

On ImageNet, epinets improve joint predictions over large ensembles and uncertainty baselines while preserving comparable marginal scaling. These gains become more substantial on the more complex dataset.

6 ImageNet: The ImageNet experiments compare epinets with ensemble approaches and uncertainty baselines using tuned ResNet-based agents.The epinet supplements a pretrained ResNet, while the ensemble comparison uses independently trained ResNet-50 particles.
6 ImageNet: None of the tuned ensemble or uncertainty-baseline agents match epinet performance on ImageNet joint log-loss.The comparison includes sngp, dropout, mimo, and het, with marginal log-loss constrained not to degrade relative to published numbers.
6 ImageNet: Epinets outperform both ResNet and ensemble alternatives on joint log-loss by an enormous margin as model size increases.The figure compares marginal classification error, marginal log-loss, and joint log-loss across model size.
6 ImageNet: Relative to large ensembles, epinets improve joint predictions at orders of magnitude lower compute cost.The ImageNet analysis presents this as the paper’s key result while showing similar marginal scaling for ResNets with and without epinets.

7 Conclusion

The paper introduces ENNs as an interface for uncertainty modeling and presents epinets as a computationally efficient architecture for joint predictions.

ENNs provide an interface for designing and evaluating models that produce joint predictions.
Epinets supplement conventional neural networks, including large pretrained models, to estimate prediction-relevant uncertainty.
Epinets achieve joint predictions outperforming ensembles of hundreds of particles at computational cost only slightly above one particle.
The paper distinguishes ENNs from BNNs by focusing on uncertainty that matters in predictions rather than unknown network parameters.

B From predictions to decisions

Marginal predictions can support poor decisions because they do not distinguish unresolved label ambiguity from uncertainty caused by limited data. Joint predictions address this gap by bounding decision shortfall through joint prediction error.

Minimizing marginal log loss can yield decisions no better than uniformly random actions, despite optimal marginal predictions.
Joint log loss formalizes prediction quality over multiple labels and supports effective decisions where marginal log loss does not necessarily suffice.
Expected KL-divergence minimization is equivalent to minimizing expected joint log loss.
Decision shortfall is bounded by the KL-divergence between an approximate and optimal joint posterior predictive, so small joint prediction error supports good decisions for any reward function.

C ENNs versus BNNs

The paper frames ENNs as a broader uncertainty interface than BNNs and evaluates joint predictions using structured input batches designed to expose dependencies among labels.

Every BNN defined over a base network can be represented as an ENN, but some ENNs cannot be represented as BNNs over that same base network.
Joint-prediction evaluation samples input batches and measures log loss against the corresponding label batches.
Dyadic sampling creates batches by repeatedly sampling from two independent anchor points, producing correlated labels with a moderate batch size of τ = 10.
On ImageNet, evaluation re-samples τ = 10 inputs from each evaluation pair and averages the resulting joint log-losses.
Repeated images in basic dyadic batches can permit a cheating agent, motivating the more diverse augmented scheme.

E.2 Augmented dyadic sampling

Augmented dyadic sampling tests whether agents recognize shared origins among perturbed inputs while making batches less vulnerable to exploiting exact repetition. Under both sampling schemes, epinets substantially outperform the evaluated baselines.

Augmented dyadic sampling independently perturbs inputs from each anchor, using random cropping and flipping on ImageNet while retaining the original label.
The augmented scheme makes each input differ from the others, reducing the concern that an agent could exploit exact repeated images in basic dyadic batches.
The epinet agent outperforms both the ResNet and ensemble baselines by a huge margin under basic and augmented dyadic sampling.
All agents benefit from increasing model size, while ensemble joint log-loss improves more than ResNet joint log-loss.
Figure 8 compares joint log-loss under dyadic and augmented dyadic sampling, with the augmented scheme producing higher overall loss because of added perturbations.

F.4 Ablation studies

The ablations show that epinet quality depends most on index dimension, prior scale, and weight decay, while performance is comparatively insensitive to several architectural choices. Across uncertainty baselines, methods have similar marginal quality but differ substantially in joint-prediction quality, where epinet performs best.

Ablation studies: Larger index dimensions improve both joint and marginal KL estimates, while epinet performance is relatively insensitive to hidden-layer count and sufficiently large hidden-layer width.The hidden-layer result is attributed to similar parameter counts across epinets.
Ablation studies: Prior scale and L2 weight decay materially affect epinet quality: extreme prior scales or weight decay values degrade performance.Smaller prior scale improves marginal KL, but overly small or large prior scale degrades joint KL.
Uncertainty baselines: Epinet performs much better than alternative uncertainty methods on joint predictions despite similar marginal-quality tradeoffs.After tuning, approaches are roughly similar in marginal quality but widely separated in joint-prediction quality.
Uncertainty baselines: The uncertainty-baseline comparison includes mimo, dropout, and SNGP, but their joint-prediction quality does not match a size-10 ensemble after tuning.Dropout requires 10 forward passes, while SNGP remains below a size-10 ensemble on joint log-loss.

G.3 CIFAR-10 and CIFAR-100

On CIFAR-10 and CIFAR-100, the results mirror ImageNet: larger models generally improve statistical losses, while epinets offer little marginal advantage but substantially improve joint predictions. The computational-cost analysis shows that this advantage persists under inference-FLOP accounting.

CIFAR-10 and CIFAR-100: Epinets greatly improve joint predictions over very large ensembles on CIFAR-10 and CIFAR-100 at orders-of-magnitude lower computational cost.The section reports that the ImageNet pattern reproduces across both datasets.
CIFAR-10 and CIFAR-100: Epinets provide little advantage over baseline ResNets for classification and marginal log-loss, but deliver huge improvements in joint log-loss, even against very large ensembles.Across statistical losses, larger models generally perform better.
Computational cost: With inference FLOPs as the cost measure, each epinet remains below 50% of the total network cost even when counting 1,000 epinet forwards.The additional operations can often run in parallel on modern TPUs.
Computational cost: Varying the number of independent index samples exposes a tradeoff between statistical loss and computational cost.The reported sweep uses M ∈ {10, 30, 100, 300, 1000, 3000}, with the paper’s main results using M = 1000.

G.5 Epinet ablations

ResNet-50 epinet ablations identify index dimension, prior scales, and post-training temperature as important design choices, while some training settings have little effect. A cold temperature most strongly improves joint log-loss during evaluation.

Epinet ablations: Cold-temperature rescaling below 1 most dramatically improves joint log-loss during evaluation.The temperature is applied to ENN output logits after training.
Epinet ablations: Larger index dimensions improve joint log-loss but do not necessarily improve marginal log-loss or classification error.This separates the effect on joint predictions from effects on marginal performance.
Epinet ablations: Epinet performance is sensitive to prior scales, degrading quickly when the scales become too large.The ablation varies the matched-epinet and convolutional-prior scales.
Epinet ablations: Adding a hidden layer worsens marginal and joint loss in this ablation, whereas index-sample count and label smoothing show little apparent effect.The authors suggest that the hidden-layer result could improve with more tuning.