Approximate Thompson Sampling via Epistemic Neural Networks

TL;DR

This paper explores using epistemic neural networks to efficiently approximate Thompson sampling in complex environments, demonstrating that small networks can match large ensembles' performance at lower computational cost.

Contribution

It introduces the epinet, a small neural network for uncertainty estimation, enabling scalable approximate Thompson sampling in complex settings.

Findings

ENNs effectively approximate TS in bandit and RL environments.

The epinet matches large ensemble performance with significantly less computation.

Quality of joint predictive distributions is crucial for TS effectiveness.

Abstract

Thompson sampling (TS) is a popular heuristic for action selection, but it requires sampling from a posterior distribution. Unfortunately, this can become computationally intractable in complex environments, such as those modeled using neural networks. Approximate posterior samples can produce effective actions, but only if they reasonably approximate joint predictive distributions of outputs across inputs. Notably, accuracy of marginal predictive distributions does not suffice. Epistemic neural networks (ENNs) are designed to produce accurate joint predictive distributions. We compare a range of ENNs through computational experiments that assess their performance in approximating TS across bandit and reinforcement learning environments. The results indicate that ENNs serve this purpose well and illustrate how the quality of joint predictive distributions drives performance. Further, we demonstrate that the \textit{epinet} -- a small additive network that estimates uncertainty -- matches the performance of large ensembles at orders of magnitude lower computational cost. This enables effective application of TS with computation that scales gracefully to complex environments.