Efficient Online Bayesian Inference for Neural Bandits
This work addresses the challenge of efficient online learning for researchers and practitioners in contextual bandits, offering a scalable solution with constant memory, though it is incremental as it builds on existing Bayesian and subspace methods.
The paper tackles the problem of catastrophic forgetting in neural bandits by introducing an online Bayesian inference algorithm that combines the extended Kalman filter with a low-dimensional subspace, enabling scaling to models with ~1M parameters and constant memory usage. It demonstrates good results on benchmarks including the Deep Bayesian Bandit Showdown, MNIST, and a recommender system.
In this paper we present a new algorithm for online (sequential) inference in Bayesian neural networks, and show its suitability for tackling contextual bandit problems. The key idea is to combine the extended Kalman filter (which locally linearizes the likelihood function at each time step) with a (learned or random) low-dimensional affine subspace for the parameters; the use of a subspace enables us to scale our algorithm to models with $\sim 1M$ parameters. While most other neural bandit methods need to store the entire past dataset in order to avoid the problem of "catastrophic forgetting", our approach uses constant memory. This is possible because we represent uncertainty about all the parameters in the model, not just the final linear layer. We show good results on the "Deep Bayesian Bandit Showdown" benchmark, as well as MNIST and a recommender system.