# On Sequential Bayesian Inference for Continual Learning

Samuel Kessler<EMAIL_ADDRESS>
University of Oxford Adam Cobb<EMAIL_ADDRESS>
SRI International Tim G. J. Rudner<EMAIL_ADDRESS>
University of Oxford Stefan Zohren<EMAIL_ADDRESS>
University of Oxford Stephen J. Roberts<EMAIL_ADDRESS>
University of Oxford

###### Abstract

Sequential Bayesian inference can be used for _continual learning_ to prevent
catastrophic forgetting of past tasks and provide an informative prior when
learning new tasks. We revisit sequential Bayesian inference and test whether
having access to the true posterior is guaranteed to prevent catastrophic
forgetting in Bayesian neural networks. To do this we perform sequential
Bayesian inference using Hamiltonian Monte Carlo. We propagate the posterior
as a prior for new tasks by fitting a density estimator on Hamiltonian Monte
Carlo samples. We find that this approach fails to prevent catastrophic
forgetting demonstrating the difficulty in performing sequential Bayesian
inference in neural networks. From there we study simple analytical examples
of sequential Bayesian inference and CL and highlight the issue of model
misspecification which can lead to sub-optimal continual learning performance
despite exact inference. Furthermore, we discuss how task data imbalances can
cause forgetting. From these limitations, we argue that we need probabilistic
models of the continual learning generative process rather than relying on
sequential Bayesian inference over Bayesian neural network weights. In this
vein, we also propose a simple baseline called _Prototypical Bayesian
Continual Learning_ , which is competitive with state-of-the-art Bayesian
continual learning methods on class incremental continual learning vision
benchmarks.

##  1 Introduction

The goal of continual learning (CL) is to find a predictor that learns to
solve a sequence of new tasks without losing the ability to solve previously
learned tasks. One key challenge of CL with neural networks (NNs) is that
model parameters from previously learned tasks are “overwritten” during
gradient-based learning of new tasks, which leads to _catastrophic forgetting_
of previously learned abilities (French, 1999). One approach to CL hinges on
using recursive applications of Bayes’ Theorem; using the weight posterior in
a Bayesian neural network (BNN) as the prior for a new task (Kirkpatrick et
al., 2017). However, obtaining a full posterior over NN weights is
computationally demanding and we often need to resort to approximations, such
as the Laplace method (MacKay, 1992) or variational inference (Graves, 2011;
Blundell et al., 2015) to obtain a neural network weight posterior.

When performing Bayesian CL, sequential Bayesian inference is performed with
an approximate BNN posterior, not the true posterior (Schwarz et al., 2018;
Ritter et al., 2018; Nguyen et al., 2017; Ebrahimi et al., 2019; Kessler et
al., 2019; Loo et al., 2020). If we consider the performance of sequential
Bayesian inference with a variational approximation over a BNN weight
posterior then we barely observe an improvement over simply learning new tasks
with stochastic gradient descent (SGD). We will develop this statement further
in Section 2.2. So if we had access to the true BNN weight posterior, would
this be enough to prevent forgetting by sequential Bayesian inference?

Our contributions in this paper are to revisit Bayesian CL. 1) Experimentally,
we perform sequential Bayesian inference using the true Bayesian NN weight
posterior. We do this by using the gold standard of Bayesian inference
methods, Hamiltonian Monte Carlo (HMC) (Neal et al., 2011). We use density
estimation over HMC samples and use this approximate posterior density as a
prior for the next task within the HMC sampling process. Surprisingly our HMC
method for CL yields no noticeable benefits over an approximate inference
method (VCL Nguyen et al. (2017)) despite using samples from the true
posterior. 2) As a result we consider a simple analytical example and
highlight that exact inference with a misspecified model can still cause
forgetting. 3) We show mathematically that under certain assumptions task data
imbalances will cause forgetting in Bayesian NNs. 4) We propose a new
probabilistic model for CL and show that by explicitly modeling the generative
process of the data, we can achieve good performance, avoiding the need to
rely on recursive Bayesian inference over NN weights to prevent forgetting.
Our proposed model, _Prototypical Bayesian Continual Learning_ (ProtoCL), is
conceptually simple, scalable, and competitive with state of the art Bayesian
CL methods in the class-incremental learning setting.

##  2 Background

###  2.1 The Continual Learning Problem

_Continual learning_ (CL) is a learning setting whereby a model must learn to
make predictions over a set of tasks sequentially while maintaining
performance across all previously learned tasks. In CL, the model is
sequentially shown $T$ tasks, denoted $\mathcal{T}_{t}$ for $t=1,\ldots,T$.
Each task, $\mathcal{T}_{t}$, is comprised of a dataset
$\mathcal{D}_{t}=\left\\{({\bm{x}}_{i},y_{i})\right\\}_{i=1}^{N_{t}}$ which a
model needs to learn to make predictions with. More generally, tasks are
denoted by distinct tuples comprised of the conditional and marginal data
distributions, $\\{p_{t}(y|{\mathbf{x}}),p_{t}({\mathbf{x}})\\}$. After task
$\mathcal{T}_{t}$ the model will lose access to the training dataset but its
performance will be continually evaluated on all tasks $\mathcal{T}_{i}$ for
$i\leq t$. For a comprehensive review of CL scenarios see (Hsu et al., 2018;
Van de Ven & Tolias, 2019). We decompose predictors as $g=h\circ f$ such that
$\hat{y}=g({\bm{x}})$ we define $f$ as an embedding function mapping
$f:\mathcal{X}\rightarrow\mathcal{Z}$ and $h$ as a head mapping to outputs
$h:\mathcal{Z}\rightarrow\mathcal{Y}$. Some continual learning methods use a
separate head per task $\left\\{h_{i}\right\\}_{i=1}^{T}$, these methods are
called multi-headed while those that use one head are called single-headed.

###  2.2 Bayesian Continual Learning

We consider a setting in which task data arrives sequentially at time steps,
$t=1,2,\ldots,T$. At the first time step, $t=1$, the model parameterized by
${\bm{\theta}}$ receives the dataset $\mathcal{D}_{1}$ and learns the
conditional distribution $p(y_{i}|{\bm{x}}_{i},{\bm{\theta}})$ for all
$({\bm{x}}_{i},y_{i})\in\mathcal{D}_{1}$ ($i$ indexes a datapoint), the
parameters ${\bm{\theta}}$ have a prior distribution $p({\bm{\theta}})$. The
posterior predictive distribution for a test point, ${\bm{x}}_{1}^{*}$ is:

$\displaystyle p(y_{1}^{*}|{\bm{x}}_{1}^{*},\mathcal{D}_{1})=\int
p(y_{1}^{*}|{\bm{x}}_{1}^{*},{\bm{\theta}})p({\bm{\theta}}|\mathcal{D}_{1})d{\bm{\theta}}.$
(1)

Computing this posterior predictive distribution above requires
$p({\bm{\theta}}|\mathcal{D}_{1})$. For $t=2$, a CL model is required to fit
$p(y_{i}|{\bm{x}}_{i},{\bm{\theta}})$ for
$\mathcal{D}_{1}\cup\mathcal{D}_{2}$. The posterior predictive distribution
for a new test point ${\bm{x}}_{2}^{*}$ point is:

$\displaystyle
p(y_{2}^{*}|{\bm{x}}_{2}^{*},\mathcal{D}_{1},\mathcal{D}_{2})=\int
p(y_{2}^{*}|{\bm{x}}_{2}^{*},{\bm{\theta}})p({\bm{\theta}}|\mathcal{D}_{1},\mathcal{D}_{2})d{\bm{\theta}}.$
(2)

The posterior must thus be updated to reflect this new conditional
distribution. We can use repeated application of Bayes’ rule to calculate the
posterior distributions
$p({\bm{\theta}}|\mathcal{D}_{1},\ldots,\mathcal{D}_{T})$ as:

$\displaystyle
p({\bm{\theta}}|\mathcal{D}_{1},\ldots,\mathcal{D}_{T-1},\mathcal{D}_{T})$
$\displaystyle=\frac{p(\mathcal{D}_{T}|{\bm{\theta}})p({\bm{\theta}}|\mathcal{D}_{1},\ldots,\mathcal{D}_{T-1})}{p(\mathcal{D}_{T}|\mathcal{D}_{1},\ldots,\mathcal{D}_{T-1})}.$
(3)

In the CL setting we lose access to previous training datasets: however, using
repeated applications of Bayes’ rule Eq. 3, allows us to sequentially
incorporate information from past tasks in the parameters ${\bm{\theta}}$. At
$t=1$, we have access to $\mathcal{D}_{1}$ and the posterior over weights is:

$\displaystyle\log p({\bm{\theta}}|\mathcal{D}_{1})$ $\displaystyle=\log
p(\mathcal{D}_{1}|{\bm{\theta}})+\log p({\bm{\theta}})-\log
p(\mathcal{D}_{1}).$ (4)

At $t=2$, we require a posterior
$p({\bm{\theta}}|\mathcal{D}_{1},\mathcal{D}_{2})$ to calculate the posterior
predictive distribution in Eq. 2. However, we have lost access to
$\mathcal{D}_{1}$. According to Bayes’ rule, the posterior may be written as:

$\displaystyle\log p({\bm{\theta}}|\mathcal{D}_{1},\mathcal{D}_{2})$
$\displaystyle=\log p(\mathcal{D}_{2}|{\bm{\theta}})+\log
p({\bm{\theta}}|\mathcal{D}_{1})-\log p(\mathcal{D}_{2}|\mathcal{D}_{1}),$ (5)

where we used the conditional independence of $\mathcal{D}_{2}$ and
$\mathcal{D}_{1}$ given ${\bm{\theta}}$. We note that the likelihood is only
dependent upon the current task dataset, $\mathcal{D}_{2}$, and that the prior
encodes parameter knowledge from the previous task. Hence, we can use the
posterior at $t$ as a prior for learning a new task at $t+1$. For the MNIST
dataset (LeCun et al., 1998) we know that if we were to train a BNN we would
achieve a good performance by inferring $p({\bm{\theta}}|\mathcal{D})$. Hence
if we were to Split-MNIST into $5$ two-way classification tasks then we should
be able to recursively recover the multi-task posterior
$p({\bm{\theta}}|\mathcal{D})=p({\bm{\theta}}|\mathcal{D}_{1}\ldots,\mathcal{D}_{5})$.
This problem is called Split-MNIST. From Eq. 5 we require that our model with
parameters ${\bm{\theta}}$ is a sufficient statistic of $\mathcal{D}_{1}$,
making the likelihood conditionally independent of $\mathcal{D}_{1}$ given
${\bm{\theta}}$. This observation motivates the use of high-capacity
predictors, such as Bayesian neural networks, that are flexible enough to
learn $\mathcal{D}_{1}$.

###  2.3 Variational Continual Learning

Variational CL (VCL; Nguyen et al. (2017)) simplifies the Bayesian inference
problem in Eq. 5 into a sequence of approximate Bayesian updates on the
distribution over random neural network weights $\bm{\theta}$. To do so, VCL
uses the variational posterior from previous tasks as a prior for new tasks.
In this way, learning to solve the first task entails finding a variational
distribution $q_{1}({\bm{\theta}}|\mathcal{D}_{1})$ that maximizes a
corresponding variational objective. For the subsequent task, the prior is
chosen to be $q_{1}({\bm{\theta}}|\mathcal{D}_{1})$, and the goal becomes to
learn a variational distribution $q_{2}({\bm{\theta}}|\mathcal{D}_{2})$ that
maximizes a corresponding variational objective under this prior. Denoting the
recursive posterior inferred from multiple datasets by
$q_{t}({\bm{\theta}}|\mathcal{D}_{1:t})$, we can express the variational CL
objective for the $t$-th task as:

$\displaystyle\mathcal{L}({\bm{\theta}},\mathcal{D}_{t})$
$\displaystyle=\mathbb{D}_{\textrm{KL}}\left[q_{t}({\bm{\theta}})||q_{t-1}({\bm{\theta}}|\mathcal{D}_{1:t-1})\right]-\mathbb{E}_{q_{t}}[\log
p(\mathcal{D}_{t}|{\bm{\theta}})].$ (6)

Figure 1:  Accuracy on Split-MNIST for various CL methods with a two-layer
BNN. We compare a NN trained with SGD (single-headed) with VCL. We consider
single-headed (SH) and multi-head (MH) VCL variants.

When applying VCL to the problem of Split-MNIST Figure 1, we can see that
single-headed VCL barely performs better than SGD when remembering past tasks.
Multi-headed VCL performs better, despite not being a requirement from
sequential Bayesian inference Eq. 5. So why does single-head VCL not improve
over SGD if we can recursively build up an approximate posterior using Eq. 5?
We hypothesize that it could be due to using a variational approximation of
the posterior and so we are not actually strictly performing the Bayesian CL
process described in Section 2.2. We test this hypothesis in the next section
by propagating the true BNN posterior to verify whether we can recursively
obtain the true multi-task posterior and so improve on single-head VCL and
prevent catastrophic forgetting.

##  3 Bayesian Continual Learning with Hamiltonian Monte Carlo

To perform inference over BNN weights we use the HMC algorithm (Neal et al.,
2011). We then use these samples and learn a density estimator that can be
used as a prior for a new task111We considered Sequential Monte Carlo, but it
is unable to scale to the dimensions required for the NNs we consider (Chopin
et al., 2020). HMC on the other hand has recently been successfully scaled to
BNNs (Cobb & Jalaian, 2021; Izmailov et al., 2021).. HMC is considered the
gold standard in approximate inference and is guaranteed to asymptotically
produce samples from the true posterior222In the NeurIPS 2021 Bayesian Deep
Learning Competition, the goal was to find an approximate inference method
that is as “close” as possible to the posterior samples from HMC.. we use
posterior samples of ${\bm{\theta}}$ from HMC and then fit a density estimator
over these samples, to use as a prior for a new task. This allows us to use a
multi-modal posterior distribution over ${\bm{\theta}}$. In contrast, to a
diagonal Gaussian variational posterior like in VCL. More concretely, to
propagate the posterior $p({\bm{\theta}}|\mathcal{D}_{1})$ we use a density
estimator, defined $\hat{p}({\bm{\theta}}|\mathcal{D}_{1})$, to fit a
probability density on HMC samples as a posterior. For the next task
$\mathcal{T}_{2}$ we can use $\hat{p}({\bm{\theta}}|\mathcal{D}_{1})$ as a
prior for a new HMC sampling chain and so on (see Fig. 2). The density
estimator priors need to satisfy two key conditions for use within HMC
sampling. Firstly, that they are a probability density function. Secondly,
that they are differentiable with respect to the input samples.

Figure 2: Illustration of the posterior propagation process; priors in blue
are in the top row and posterior samples on the bottom row. This is a two step
process where we first perform HMC with an isotropic Gaussian prior for
$\mathcal{T}_{1}$ then perform density estimation on the HMC samples from the
posterior to obtain $\hat{p}_{1}(\theta|\mathcal{D}_{1})$. This posterior can
then be used as a prior for the new task $\mathcal{T}_{2}$ and so on.

We use a toy dataset (Fig. 3) with two classes and inputs
${\bm{x}}\in\mathbb{R}^{2}$ (Pan et al., 2020). Each task is a binary
classification problem where the decision boundary extends from left to right
for each new task. We train a two layer BNN, with hidden state size of $10$.
We use a Gaussian Mixture Models (GMM) as a density estimator for
approximating the posterior with HMC samples. We also tried Normalizing Flows
which should be more flexible (Dinh et al., 2016) however these did not work
robustly for HMC sampling333RealNVP was very sensitive to the choice of random
seed, the samples from the learned distribution did not give accurate
predictions for the current task and led to numerical instabilities when used
as a prior within HMC sampling.. To the best of our knowledge we are the first
to incorporate flexible priors into the sampling methods like HMC.

Training a BNN with HMC on the same multi-task dataset gets a test accuracy of
$1.0$. Thus, the final posterior is suitable for continual learning under Eq.
3 we should be able to recursively arrive at the multi-task posterior with our
recursive inference method with HMC. The results from Fig. 3 demonstrate that
using HMC with an approximate multi-modal posterior fails to prevent
forgetting and is less effective than using multi-head VCL. In fact, multi-
head VCL clearly outperforms HMC indicating that the source of the knowledge
retention is not through the propagation of the posterior but through the task
specific heads. We also repeated these experiments with another toy dataset of
five binary classification tasks where we observe similar results (Fig. 7).

For HMC we ensure that we are sampling from the posterior by assessing chain
convergence and effective sample sizes (Fig. 11). The effective sample size
measures the autocorrelation in the chain. The effective sample sizes for the
HMC chains for our BNNs are similar to the literature (Cobb & Jalaian, 2021).
Also, we ensure that our GMM approximate posteriors are multi-modal and so has
a more complex posterior in comparison to VCL, and that the GMM samples
produce equivalent results to HMC samples for the current task (Fig. 10). See
Appendix B for details.

Figure 3: On the left is the toy dataset of $5$ distinct $2$-way
classification tasks which involve classifying circles and squares (Pan et
al., 2020). Also, continual learning binary classification test accuracies
over $10$ seeds. The pink solid line is a multi-task (MT) baseline accuracy
using SGD/HMC with the same model as for the CL experiments.

We are not able to perform sequential Bayesian inference in BNNs despite using
HMC which is considered the gold standard of Bayesian deep learning. HMC and
density estimation with a GMM produces richer, accurate and multi-modal
posteriors. Despite this we are still not able to sequentially build up the
multi-task posterior or get much better results than an isotropic Gaussian
posterior like single-head VCL. The weak point of this method is the density
estimation, the GMM removes probability mass over areas of the BNN weight
space posterior which is important for the new task. This demonstrates just
how difficult a task it is to model BNN weight posteriors. In the next
section, we study a different analytical example of sequential Bayesian
inference and look at how model misspecification and task data imbalances can
cause forgetting in Bayesian CL.

##  4 Bayesian Continual Learning and Model Misspecification

Figure 4: Posterior estimate of the filtering distribution Eq. 7 for two
different scenarios with two tasks or changepoint.

We now consider a simple analytical example where we can perform the
sequential Bayesian inference Eq. 3 in closed form using conjugacy. We
consider a simple setting where data points arrive online, one after another.

Observations $y_{1},y_{2},\ldots,y_{t}$ arrive online, each observation is
generated by a hidden variable $\theta_{1},\theta_{2},\ldots,\theta_{t}\sim p$
where $p$ is a probability density function. At time $t$ we wish to infer the
_filtering distribution_ $p(\theta_{t}|y_{1},y_{2},\ldots,y_{t})$ (Doucet et
al., 2001) using sequential Bayesian inference, similarly to the _Kalman
filter_ (Kalman, 1960). The likelihood is
$p(y_{t}|\theta_{t})=\mathcal{N}(y_{t};f(\,\cdot\,;\theta_{t}),\sigma^{2})$
such that $y_{t}=f(\,\cdot\,;\theta_{t})+\epsilon$ where
$\epsilon\sim\mathcal{N}(0,\sigma^{2})$ and
$f(\,\cdot\,;\theta_{t})=\theta_{t}$. We consider a Gaussian prior over the
mean parameters $\theta$ such that
$p(\theta_{0})=\mathcal{N}(\theta_{0};0,\sigma_{0}^{2})$. Since the conjugate
prior for the mean is also Gaussian, the prior and posterior are
$\mathcal{N}(\theta_{t-1};\hat{\theta}_{t-1},\hat{\sigma}^{2}_{t-1})$ and
$\mathcal{N}(\theta_{t};\hat{\theta}_{t},\hat{\sigma}^{2}_{t})$. By using
sequential Bayesian inference we can have closed-form update equations for our
posterior parameters:

$\displaystyle\hat{\theta}_{t}=\hat{\sigma}^{2}_{t}\left(\frac{y_{t}}{\sigma^{2}}+\frac{\hat{\theta}_{t-1}}{\hat{\sigma}_{t-1}^{2}}\right)=\hat{\sigma}^{2}_{t}\left(\sum_{i=1}^{t}\frac{y_{i}}{\sigma^{2}}+\frac{\hat{\theta}_{0}}{\hat{\sigma}_{0}^{2}}\right),\quad\quad\frac{1}{\hat{\sigma}_{t}^{2}}$
$\displaystyle=\frac{1}{\sigma^{2}}+\frac{1}{\hat{\sigma}_{t-1}^{2}}.$ (7)

From Equation 7 the posterior mean follows a Gaussian distribution where the
posterior mean is a sum of the online observation and the online prior. So the
posterior mean is a weighted sum of the data. So, if the observations are non-
stationary: if there is task change (more commonly referred to as a
changepoint in the time-series literature). Then the mean parameter will
encapsulate a global mean over both tasks rather than a mean for each task
Fig. 4. The model is clearly misspecified since a single parameter cannot
model both of these tasks together. _Despite performing exact inference a
misspecified model can forget_ , Fig. 4. In the case of HMC we verified that
our Bayesian neural network had perfect performance on all tasks _beforehand_.
In Section 3 we had a well specified model but struggled with exact sequential
Bayesian inference Eq. 3. With this $1$-d online learning scenario we are
performing exact inference, however we have a misspecified model. It is
important to disentangle model misspecification and exact inference, and
highlight that model misspecification is a caveat which has not been
highlighted in the CL literature as far as we are aware. Furthermore we can
only ensure that our models are well specified if we have access to data from
all tasks a priori. So in the scenario of _online continual learning_ (De
Lange et al., 2021) we cannot know if our model will perform well on all past
and future tasks without making assumptions on the task distributions.

##  5 Sequential Bayesian Inference and Imbalanced Task Data

Neural Networks are complex models with a broad hypothesis space and hence are
a suitably well-specified model when tackling continual learning problems
(Wilson & Izmailov, 2020). However we struggle to fit the posterior samples
from HMC to perform sequential Bayesian inference Section 3.

We continue to use Bayesian filtering and assume a Bayesian NN where the
posterior is Gaussian with a full covariance. By modelling the entire
covariance we enable modelling how each individual weight varies with respect
to all others. We do this by interpreting online learning in Bayesian NNs as
filtering (Ciftcioglu & Türkcan, 1995). Our treatment is similar to Aitchison
(2018) who derives an optimizer by leveraging Bayesian filtering. We consider
inference in the graphical model depicted in Fig. 5. The aim is to infer the
optimal BNN weights, ${\bm{\theta}}^{*}_{t}$ at $t$ given a single observation
and the BNN weight prior. The previous BNN weights are used as a prior for
inferring the posterior BNN parameters. We consider the online setting where a
single data point $({\bm{x}}_{t},y_{t})$ is observed at a time.

Instead of modelling the full covariance we instead consider each parameter
$\theta_{i}$ as a function of all the other parameters ${\bm{\theta}}_{-it}$.
We also assume that the values of the weights are close to those of the
previous timestep (Jacot et al., 2018). To obtain the update equations for BNN
parameters given a new observation and prior we make two simplifying
assumptions as follows.

$\theta^{*}_{t-1}$$\theta^{*}_{t-2}$$\theta^{*}_{t}$$\theta^{*}_{t+1}$$y_{t-1}$$y_{t-2}$$y_{t}$$y_{t+1}$$x_{t-1}$$x_{t-2}$$x_{t}$$x_{t+1}$$\,\ldots\,$$\,\ldots\,$

Figure 5: Graphical model for filtering. Grey and white nodes are observed and
latent variables respectively.

######  Assumption 5.1.

_For a Bayesian neural network with output $f({\bm{x}}_{t};{\bm{\theta}})$ and
likelihood $\mathcal{L}({\bm{x}}_{t},y_{t};{\bm{\theta}})$, the derivative
evaluated at ${\bm{\theta}}_{t}$ is
${\mathbf{z}}_{t}=\partial\mathcal{L}({\bm{x}}_{t},y_{t};{\bm{\theta}})/\partial{\bm{\theta}}|_{{\bm{\theta}}={\bm{\theta}}_{t}}$
and the Hessian is ${\bm{H}}$. We assume a quadratic loss for a data point
$({\bm{x}}_{t},y_{t})$ of the form:_

$\displaystyle\SwapAboveDisplaySkip\mathcal{L}({\bm{x}}_{t},y_{t};{\bm{\theta}})=\mathcal{L}_{t}({\bm{\theta}})=-\frac{1}{2}{\bm{\theta}}^{\top}{\bm{H}}{\bm{\theta}}+{\mathbf{z}}^{\top}_{t}{\bm{\theta}},$
(8)

_the result of a second-order Taylor expansion. The Hessian is assumed to be
constant with respect to $({\bm{x}}_{t},y_{t})$ (but not with respect to
${\bm{\theta}}$)._

To construct the dynamical equation for ${\bm{\theta}}$, consider the gradient
for the $i$-th weight while all other parameters are set to their current
estimate at the optimal value for the $\theta_{it}^{*}$:

$\displaystyle\theta_{it}^{*}=-\frac{1}{H_{ii}}{\bm{H}}_{-ii}^{\top}{\bm{\theta}}_{-it},$
(9)

since $z_{it}=0$ at a mode. The equation above shows us that the dynamics of
the optimal weight $\theta^{*}_{it}$ is dependent on all the other current
values of the parameters ${\bm{\theta}}_{-it}$. The dynamics of
${\bm{\theta}}_{-it}$ will be a complex stochastic process dependent on many
different variables: such as the dataset, model architecture, learning rate
schedule, etc.

######  Assumption 5.2.

_Since reasoning about the dynamics of ${\bm{\theta}}_{-it}$ are intractable,
we assume that at the next time-step the optimal weights are close to the
previous timesteps with a discretized Ornstein-Uhlenbeck process for the
weights ${\bm{\theta}}_{-it}$ with reversion speed
$\vartheta\in\mathbb{R}_{+}$ and noise variance $\eta^{2}_{-i}$:_

$\displaystyle
p({\bm{\theta}}_{-i,t+1}|{\bm{\theta}}_{-i,t})=\mathcal{N}((1-\vartheta){\bm{\theta}}_{-it},\eta_{-i}^{2}),$
(10)

_this implies that the dynamics for the optimal weight is defined by_

$\displaystyle
p(\theta^{*}_{i,t+1}|\theta^{*}_{i,t})=\mathcal{N}((1-\vartheta)\theta^{*}_{it},\eta^{2}),$
(11)

_where $\eta^{2}=\eta^{2}_{-i}{\bm{H}}^{\top}_{-ii}{\bm{H}}_{-ii}$._

In simple terms, in 5.2 we assume a parsimonious model of the dynamics. That
the next value of ${\bm{\theta}}_{-i,t}$ is close to their previous value
according to a Gaussian, similarly to Aitchison (2018).

######  Lemma 5.3.

Under Assumptions 5.1 and 5.2 the dynamics and likelihood are Gaussian. Thus
we are able to infer the posterior distribution over the optimal weights using
Bayesian updates and by linearizing the BNN the update equations for the
posterior of the mean and variance of the BNN for a new data point are:

$\displaystyle\mu_{t,\text{post}}$
$\displaystyle=\sigma^{2}_{t,\text{post}}\left(\frac{\mu_{t,\text{prior}}}{\sigma^{2}_{t,\text{prior}}(\eta^{2})}+\frac{y_{t}}{\sigma^{2}}g({\bm{x}}_{t})\right)\qquad\text{and}\qquad\frac{1}{\sigma^{2}_{t,\textrm{post}}}=\frac{g({\bm{x}}_{t})^{2}}{\sigma^{2}}+\frac{1}{\sigma^{2}_{t,\textrm{prior}}(\eta^{2})},$
(12)

where we drop the notation for the $i$-th parameter, the posterior is
$\mathcal{N}(\theta^{*}_{t};\mu_{t,\text{post}},\sigma^{2}_{t,\text{post}})$
and $g({\bm{x}}_{t})=\frac{\partial
f({\bm{x}}_{t};\theta^{*}_{it})}{\partial\theta^{*}_{it}}$ and
$\sigma^{2}_{t,\textrm{prior}}$ is a function of $\eta^{2}$.

See Appendix E for the derivation of Lemma 5.3. From Eq. 12 we can notice that
the posterior mean depends linearly on the prior and a data dependent term and
so will behave similarly to our previous example in Section 4. Under 5.1 and
5.2 then if there is a data imbalance between tasks in Eq. 12, then the data
dependent term will dominate the prior term if there is more data for the
current task.

In Section 3 we showed that it is very difficult with current machine learning
tools to perform sequential Bayesian inference for simple CL problems with
small Bayesian NNs. When we disentangle Bayesian inference and model
misspecification we show showed that misspecified models can forget despite
exact Bayesian inference. The only way to ensure that our model is well
specified is to show that the multi-task posterior produces reasonable
posterior predictive distributions $p(y|{\bm{x}},\mathcal{D})=\int
p(y|{\bm{x}},\mathcal{D},{\bm{\theta}})p({\bm{\theta}}|\mathcal{D})d{\bm{\theta}}$
for one’s application. Additionally, in this section we have shown that if
there is a task dataset size imbalance then we can get forgetting under
certain assumptions.

##  6 Related Work

There has been a recent resurgence in the field of CL (Thrun & Mitchell, 1995)
given the advent of deep learning. Methods which approximate sequential
Bayesian inference Eq. 5 have been seminal in CL’s revival and have used a
diagonal Laplace approximation (Kirkpatrick et al., 2017; Schwarz et al.,
2018). The diagonal Laplace approximation have been enhanced by modelling
covariances of between neural network weights in the same layer (Ritter et
al., 2018). Instead of the Laplace approximation we can use a variational
approximation for sequential Bayesian inference (Nguyen et al., 2017; Zeno et
al., 2018). Using richer priors has also been explored (Ahn et al., 2019;
Farquhar et al., 2020; Kessler et al., 2019; Mehta et al., 2021; Kumar et al.,
2021; Loo et al., 2020). Gaussian processes have also been applied to CL
problems leveraging inducing points to retain previous task functions (Titsias
et al., 2019; Kapoor et al., 2021).

Bayesian methods which regularize weights have not matched up to the
performance of experience replay based CL methods (Buzzega et al., 2020) in
terms of accuracy on CL image classification benchmarks. Instead of
regularizing high dimension weight spaces, regularizing task functions is a
more direct approach to combat forgetting (Benjamin et al., 2018). In
particular, one can leverage the duality between the Laplace approximation and
Gaussian Processes to develop a functional regularization approach to Bayesian
CL (Swaroop et al., 2019) or using function-space variational inference
(Rudner et al., 2022a; b).

##  7 Prototypical Bayesian Continual Learning

We have shown that sequential Bayes over NN parameters is very difficult
(Section 3), and is only suitable for situations where the multi-task
posterior is suitable for all tasks. We now show that a more fruitful approach
is to model the full data-generating process of the CL problem and we propose
a simple and scalable approach for doing so. In particular, we represent
classes by prototypes (Snell et al., 2017; Rebuffi et al., 2017) to prevent
catastrophic forgetting. We refer to this framework as Prototypical Bayesian
Continual Learning, or ProtoCL for short. This approach can be viewed as a
probabilistic variant of iCarl (Rebuffi et al., 2017), which creates embedding
functions for different classes which are simply class means and predictions
are made by nearest neighbors. ProtoCL also bears similarities to the few-shot
learning model Probabilistic Clustering for Online Classification (Harrison et
al., 2019), developed for few-shot image classification.

Model. ProtoCL models the generative CL process. We consider classes
$j\in\\{1,\ldots,J\\}$, generated from a categorical distribution with a
Dirichlet prior:

$\displaystyle y_{i,t}\sim\text{Cat}(p_{1:J}),\quad
p_{1:J}\sim\textrm{Dir}(\alpha_{t}).$ (13)

Images are embedded into a embedding space by an encoder,
${\bm{z}}=f({\bm{x}};{\bm{w}})$ with parameters ${\bm{w}}$. The per class
embeddings are Gaussian whose mean has a prior which is also Gaussian:

$\displaystyle{\bm{z}}_{it}|y_{it}\sim\mathcal{N}(\bar{{\bm{z}}}_{yt},\Sigma_{\epsilon}),\quad\bar{{\bm{z}}}_{yt}\sim\mathcal{N}({\bm{\mu}}_{yt},\Lambda^{-1}_{yt}).$
(14)

Figure 6: Overview of ProtoCL.

See Fig. 6 for an overview of the model. To alleviate forgetting in CL,
ProtoCL uses a coreset of past task data to continue to embed past classes
distinctly as prototypes. The posterior distribution over class probabilities
$\\{p_{j}\\}_{j=1}^{J}$ and class embeddings $\\{\bar{z}_{y_{j}}\\}_{j=1}^{J}$
is denoted in short hand as $p({\bm{\theta}})$ with parameters
$\eta_{t}=\\{\alpha_{t},{\bm{\mu}}_{1:J,t},\Lambda^{-1}_{1:J,t}\\}$. ProtoCL
models each class prototype but does not use task specific NN parameters or
modules like multi-head VCL. By modeling a probabilistic model over an
embedding space this allows us to use powerful embedding functions
$f(\,\cdot\,;{\bm{w}})$ without having to parameterize them probabilistically
and so this approach will be more scalable than VCL, for instance.

Inference. As the Dirichlet prior is conjugate with the Categorical
distribution and likewise the Gaussian over prototypes with a Gaussian prior
over the prototype mean, we can calculate posteriors in closed form and update
the parameters $\eta_{t}$ as new data is observed without using gradient based
updates. We optimize the model by maximizing the posterior predictive
distribution and use a softmax over class probabilities to perform
predictions. We perform gradient-based learning of the NN embedding function
$f(\,\cdot\,;{\bm{w}})$ and update the parameters, $\eta_{t}$ at each
iteration of gradient descent as well, see Algorithm 1.

Sequential updates. We can obtain our parameter updates for the Dirichlet
posterior by Categorical-Dirichlet conjugacy:

$\displaystyle\SwapAboveDisplaySkip\alpha_{t+1,j}=\alpha_{t,j}+\sum^{N_{t}}_{i=1}\mathbb{I}(y^{i}_{t}=j),$
(15)

where $N_{t}$ are the number of points seen during the update at time step
$t$. Also, due to Gaussian-Gaussian conjugacy the posterior for the Gaussian
prototypes is governed by:

$\displaystyle\Lambda_{y_{t+1}}$
$\displaystyle=\Lambda_{y_{t}}+N_{y}\Sigma^{-1}_{\epsilon}$ (16)
$\displaystyle\Lambda_{y_{t+1}}{\bm{\mu}}_{y_{t+1}}$
$\displaystyle=N_{y}\Sigma^{-1}_{\epsilon}\bar{{\bm{z}}}_{y_{t}}+\Lambda_{y_{t}}{\bm{\mu}}_{y_{t}},\,\forall
y_{t}\in C_{t},$ (17)

where $N_{y}$ are the number of samples of class $y$ and
$\bar{{\bm{z}}}_{y_{t}}=(1/N_{y})\sum^{N_{y}}_{i=1}z_{yi}$, see Section D.2
for the detailed derivation.

Objective. We optimize the posterior predictive distribution of the prototypes
and classes:

$\displaystyle p({\mathbf{z}},y)$ $\displaystyle=\int
p({\mathbf{z}},y|{\bm{\theta}}_{t};\eta_{t})p({\bm{\theta}}_{t};\eta_{t})d{\bm{\theta}}_{t}=p(y)\prod^{N_{t}}_{i=1}\mathcal{N}({\mathbf{z}}_{it}|y_{it};{\bm{\mu}}_{y_{t},t},\Sigma_{\epsilon}+\Lambda^{-1}_{y_{t},t}).$
(18)

Where the $p(y)=\alpha_{y}/\sum^{J}_{j=1}\alpha_{j}$, see Section D.3 for the
detailed derivation. This objective can then be optimized using gradient based
optimization for learning the prototype embedding function
${\bm{z}}=f({\bm{x}};{\bm{w}})$.

Predictions. To make a prediction for a test point ${\bm{x}}^{*}$ the class
with the maximum (log)-posterior predictive is chosen, where the posterior
predictive is:

$\displaystyle p(y^{*}=j|{\bm{x}}^{*},{\bm{x}}_{1:t},y_{1:t})$
$\displaystyle=p(y^{*}=j|{\bm{z}}^{*},{\bm{\theta}}_{t})=\frac{p(y^{*}=j,{\bm{z}}^{*}|{\bm{\theta}}_{t})}{\sum_{i}p(y=i,{\bm{z}}^{*}|{\bm{\theta}}_{t})},$
(19)

see Section D.4 for further details.

Algorithm 1 ProtoCL continual learning

1: Input: task datasets $\mathcal{T}_{1:T}$ , initialize embedding function:
$f(\,\cdot\,;{\bm{w}})$, coreset: $\mathcal{M}=\emptyset$.

2: for $\mathcal{T}_{1}$ to $\mathcal{T}_{T}$ do

3: for each batch in $\mathcal{T}_{i}\cup\mathcal{M}$ do

4: Optimize $f(\cdot;{\bm{w}})$ by maximizing the posterior predictive
$p({\mathbf{z}},y)$ Eq. 18

5: Obtain posterior over ${\bm{\theta}}$ by updating $\eta$, Eqs. 15, 16 and
17.

6: end for

7: Add random subset from $\mathcal{T}_{i}$ to $\mathcal{M}$.

8: end for

Preventing forgetting. As we wish to retain the class prototypes. We make use
of coresets: experience from previous tasks. At the end of learning a task
$\mathcal{T}_{t}$, we retain subset $\mathcal{M}_{t}\subset\mathcal{D}_{t}$
and augment each new task dataset to ensure that posterior parameters
$\eta_{t}$ and prototypes are able to retain previous task information.

Class incremental learning. In this CL setting we do not tell the CL agent
which task it is currently training and evaluating with a task identifier
$\tau$. So we cannot use the task identifier to select a specific head to use
for classifying a test point, or use the task identifier to condition the
model in another way (such that $p(y|{\bm{x}},\tau)$ where $\tau$ is the task
identifier). Also we require the CL agents to classify each class in the CL
benchmarks, for example $\\{0,\ldots,9\\}$ for Split-MNIST and Split CIFAR10
and not just $\\{0,1\\}$ for each task as commonly performed (these are called
_task-incremental_ and _domain-incremental_ learning). We believe that these
are more realistic and the most important settings to evaluate on.

Table 1: Mean accuracies across all tasks over CL vision benchmarks for _class incremental learning_ (Van de Ven & Tolias, 2019). All results are averages and standard errors over $10$ seeds. ∗Uses the predictive entropy to make a decision about which head for class incremental learning. Method |  Coreset |  Split-MNIST |  Split-FMNIST
---|---|---|---
VCL (Nguyen et al., 2017) |  ✗ |  $33.01\pm 0.08$ |  $32.77\pm 1.25$
$\quad\quad+$ coreset |  ✓ |  $52.98\pm 18.56$ |  $61.12\pm 16.96$
HIBNN∗ (Kessler et al., 2019) |  ✗ |  $85.50\pm 3.20$ |  $43.70\pm 20.21$
FROMP (Pan et al., 2020) |  ✓ |  $84.40\pm 0.00$ |  $68.54\pm 0.00$
S-FSVI (Rudner et al., 2022b) |  ✓ |  $\bm{92.94\pm 0.17}$ |  $80.55\pm 0.41$
ProtoCL (ours) |  ✓ |  $\bm{93.73\pm 1.05}$ |  $\bm{82.73\pm 1.70}$
Table 2: Mean accuracies across all tasks over CL vision benchmarks for _class incremental learning_ (Van de Ven & Tolias, 2019). All results are averages and standard errors over $10$ seeds. ∗Uses the predictive entropy to make a decision about which head for class incremental learning. Training times have been benchmarked using an Nvidia RTX3090 GPU. Method |  Training time (sec) $(\downarrow)$ |  Split CIFAR-10 (acc) $(\uparrow)$
---|---|---
FROMP (Pan et al., 2020) |  $1425\pm 28$ |  $48.92\pm 10.86$
S-FSVI (Rudner et al., 2022b) |  $44434\pm 91$ |  $50.85\pm 3.87$
ProtoCL (ours) |  $\bm{384\pm 6}$ |  $55.81\pm 2.10$
|  |  Split CIFAR-100 (acc)
S-FSVI (Rudner et al., 2022b) |  $37355\pm 1135$ |  $20.04\pm 2.37$
ProtoCL (ours) |  $\bm{1425\pm 28}$ |  $\bm{23.96\pm 1.34}$

Implementation. For Split-MNIST and Split-FMNIST the baselines and ProtoCL all
use two layer NNs with hidden state size of 200. For Split-CIFAR10 and Split-
CIFAR100, the baselines and ProtoCL use a four layer convolution neural
network with two fully connected layers of size $512$ similarly to Pan et al.
(2020). For ProtoCL and all baselines which rely on replay we fix the size of
the coreset to $200$ points per task. For all ProtoCL models we allow the
prior Dirichlet parameters to be learned and set their initial value to $0.7$
found by a random search over MNIST with ProtoCL. An important hyperparameter
for ProtoCL is the embedding dimension of the Gaussian prototypes for Split-
MNIST and Split-FMNIST this was set to $128$ while for the larger vision
datasets this was set to $32$ found using grid-search.

Results. ProtoCL produces good results on CL benchmarks on par or better than
S-FSVI (Rudner et al., 2022b) which is state-of-the-art on the smaller CL
benchmarks while being a lot more efficient to train and without requiring
expensive variational inference. ProtoCL can flexibly scale to larger CL
vision benchmarks producing better results than S-FSVI. Code to reproduce all
experiments can be found here anonymous.4open.science/r/bayes_cl_exploration.
All our experiments are in the more realistic class incremental learning
setting, which is a harder setting than those reported in most CL papers, so
the results in Table 1 are lower for certain baselines than in the respective
papers. We use $200$ data points per task, see Figure 12 for a sensitivity
analysis of the performance over the Split-MNIST benchmark as as function of
core size for ProtoCL.

The stated aim of ProtoCL is not provide a novel state-of-the-art method for
CL, but rather to propose a simple baseline which takes an alternative route
than weight-space sequential Bayesian inference. We can achieve strong results
that mitigate forgetting, namely by modeling the generative CL process and
using sequential Bayesian inference over a few parameters in the class
prototype embedding space. We argue that modeling the generative CL process is
a fruitful direction for further research rather than attempting sequential
Bayesian inference over the weights of a BNN.

##  8 Discussion & Conclusion

In this paper we have revisited the use of sequential Bayesian inference for
CL. We can use sequential Bayes to recursively build up the multi-task
posterior Eq. 5. Previous methods have relied on approximate inference and see
little benefit over SGD. We test the hypothesis of whether this poor
performance is due to the approximate inference scheme by using HMC in two
simple CL problems. HMC asymptotically samples from the true posterior and we
use a density estimator over HMC samples to use as a prior for a new task
within the HMC sampling process. This density is multi-modal and accurate with
respect to the current task but is not able to improve over using an
approximate posterior. This demonstrates just how challenging it is to work
with BNN weight posteriors. The source of error comes from the density
estimation step. We then look at an analytical example of sequential Bayesian
inference where we perform exact inference however due to model
misspecification, we observe forgetting. The only way to ensure a well
specified model is to assess the multi-task performance over all tasks a
priori. This might not be possible in online CL settings. We then model an
analytical example over Bayesian NNs and under certain assumptions show that
if there is task data imbalances then this will cause forgetting. Because of
these results, we argue against performing weight space sequential Bayesian
inference and instead model the generative CL problem. We introduce a simple
baseline called ProtoCL. ProtoCL doesn’t require complex variational
optimization and achieves competitive results to state-of-the-art in the
realistic setting of class incremental learning.

This conclusion should not be a surprise since the latest Bayesian CL papers
have all relied multi-head architectures or inducing points/coresets to
prevent forgetting, rather than better weight-space inference schemes. Our
observations are in line with recent theory from (Knoblauch et al., 2020)
which states that optimal CL requires perfect memory. Although the results
were shown with deterministic NNs the same results follow for BNN with a
single set of parameters. Future research directions include enabling coresets
of task data to efficiently and accurately approximate the posterior of a BNN
to remember previous tasks.

##  9 Acknowledgments

We would like to thank Sebastian Farquhar, Laurence Aitchison, Jeremias
Knoblauch and Chris Holmes for discussions. Thanks to Phil Ball for help with
writing the paper. SK acknowledges funding from the Oxford-Man Institute of
Quantitative Finance. TGJR acknowledges funding from the Rhodes Trust,
Qualcomm, and the Engineering and Physical Sciences Research Council (EPSRC).
This material is based upon work supported by the United States Air Force and
DARPA under Contract No. FA8750-20-C-0002. Any opinions, findings and
conclusions or recommendations expressed in this material are those of the
author(s) and do not necessarily reflect the views of the United States Air
Force and DARPA.

## References

  * Ahn et al. (2019) Hongjoon Ahn, Sungmin Cha, Donggyu Lee, and Taesup Moon.  Uncertainty-based continual learning with adaptive regularization.  _Advances in neural information processing systems_ , 32, 2019.
  * Aitchison (2018) Laurence Aitchison.  Bayesian filtering unifies adaptive and non-adaptive neural network optimization methods.  _arXiv preprint arXiv:1807.07540_ , 2018.
  * Benjamin et al. (2018) Ari S Benjamin, David Rolnick, and Konrad Kording.  Measuring and regularizing networks in function space.  _arXiv preprint arXiv:1805.08289_ , 2018.
  * Blundell et al. (2015) Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In _International Conference on Machine Learning_ , pp. 1613–1622. PMLR, 2015.
  * Buzzega et al. (2020) Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, and Simone Calderara.  Dark experience for general continual learning: a strong, simple baseline.  _Advances in neural information processing systems_ , 33:15920–15930, 2020.
  * Chopin et al. (2020) Nicolas Chopin, Omiros Papaspiliopoulos, et al.  _An introduction to sequential Monte Carlo_ , volume 4.  Springer, 2020.
  * Ciftcioglu & Türkcan (1995) Ö Ciftcioglu and E Türkcan.  Adaptive training of feedforward neural networks by Kalman filtering.  1995\.
  * Cobb & Jalaian (2021) Adam D Cobb and Brian Jalaian.  Scaling Hamiltonian Monte Carlo Inference for Bayesian Neural Networks with Symmetric Splitting.  _Uncertainty in Artificial Intelligence_ , 2021.
  * De Lange et al. (2021) Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars.  A continual learning survey: Defying forgetting in classification tasks.  _IEEE transactions on pattern analysis and machine intelligence_ , 44(7):3366–3385, 2021.
  * Dinh et al. (2016) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real NVP.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Doucet et al. (2001) Arnaud Doucet, Nando De Freitas, and Neil Gordon.  An introduction to sequential Monte Carlo methods.  In _Sequential Monte Carlo methods in practice_ , pp. 3–14. Springer, 2001.
  * Ebrahimi et al. (2019) Sayna Ebrahimi, Mohamed Elhoseiny, Trevor Darrell, and Marcus Rohrbach.  Uncertainty-guided continual learning with Bayesian neural networks.  _arXiv preprint arXiv:1906.02425_ , 2019.
  * Farquhar et al. (2020) Sebastian Farquhar, Michael A Osborne, and Yarin Gal.  Radial bayesian neural networks: Beyond discrete support in large-scale bayesian deep learning.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1352–1362. PMLR, 2020.
  * French (1999) Robert M French.  Catastrophic forgetting in connectionist networks.  _Trends in cognitive sciences_ , 3(4):128–135, 1999.
  * Graves (2011) Alex Graves.  Practical variational inference for neural networks.  _Advances in neural information processing systems_ , 24, 2011.
  * Harrison et al. (2019) James Harrison, Apoorva Sharma, Chelsea Finn, and Marco Pavone.  Continuous meta-learning without tasks.  _arXiv preprint arXiv:1912.08866_ , 2019.
  * Henning et al. (2021) Christian Henning, Maria Cervera, Francesco D’Angelo, Johannes Von Oswald, Regina Traber, Benjamin Ehret, Seijin Kobayashi, Benjamin F Grewe, and João Sacramento.  Posterior meta-replay for continual learning.  _Advances in Neural Information Processing Systems_ , 34:14135–14149, 2021.
  * Hsu et al. (2018) Yen-Chang Hsu, Yen-Cheng Liu, Anita Ramasamy, and Zsolt Kira.  Re-evaluating continual learning scenarios: A categorization and case for strong baselines.  _arXiv preprint arXiv:1810.12488_ , 2018.
  * Izmailov et al. (2021) Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Wilson.  What Are Bayesian Neural Network Posteriors Really Like?  _arXiv preprint arXiv:2104.14421_ , 2021.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _arXiv preprint arXiv:1806.07572_ , 2018.
  * Kalman (1960) Rudolph Emil Kalman.  A new approach to linear filtering and prediction problems.  1960\.
  * Kapoor et al. (2021) Sanyam Kapoor, Theofanis Karaletsos, and Thang D Bui.  Variational auto-regressive Gaussian processes for continual learning.  In _International Conference on Machine Learning_ , pp. 5290–5300. PMLR, 2021.
  * Kessler et al. (2019) Samuel Kessler, Vu Nguyen, Stefan Zohren, and Stephen Roberts.  Hierarchical Indian buffet neural networks for Bayesian continual learning.  _arXiv preprint arXiv:1912.02290_ , 2019.
  * Kirkpatrick et al. (2017) James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell.  Overcoming catastrophic forgetting in neural networks.  _Proceedings of the National Academy of Sciences_ , 114(13):3521–3526, 2017.  ISSN 0027-8424.  doi: 10.1073/pnas.1611835114.
  * Knoblauch et al. (2020) Jeremias Knoblauch, Hisham Husain, and Tom Diethe.  Optimal continual learning has perfect memory and is NP-hard.  In _International Conference on Machine Learning_ , pp. 5327–5337. PMLR, 2020.
  * Kumar et al. (2021) Abhishek Kumar, Sunabha Chatterjee, and Piyush Rai.  Bayesian structural adaptation for continual learning.  In _International Conference on Machine Learning_ , pp. 5850–5860. PMLR, 2021.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Loo et al. (2020) Noel Loo, Siddharth Swaroop, and Richard E Turner.  Generalized variational continual learning.  _arXiv preprint arXiv:2011.12328_ , 2020.
  * MacKay (1992) David JC MacKay.  A practical Bayesian framework for backpropagation networks.  _Neural computation_ , 4(3):448–472, 1992.
  * Mehta et al. (2021) Nikhil Mehta, Kevin Liang, Vinay Kumar Verma, and Lawrence Carin.  Continual learning using a bayesian nonparametric dictionary of weight factors.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 100–108. PMLR, 2021.
  * Neal et al. (2011) Radford M Neal et al.  MCMC using Hamiltonian dynamics.  _Handbook of Markov chain Monte Carlo_ , 2(11):2, 2011.
  * Nguyen et al. (2017) Cuong V Nguyen, Yingzhen Li, Thang D Bui, and Richard E Turner.  Variational continual learning.  _arXiv preprint arXiv:1710.10628_ , 2017.
  * Pan et al. (2020) Pingbo Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard E Turner, and Mohammad Emtiyaz Khan.  Continual deep learning by functional regularisation of memorable past.  _arXiv preprint arXiv:2004.14070_ , 2020.
  * Petersen et al. (2008) Kaare Brandt Petersen, Michael Syskind Pedersen, et al.  The matrix cookbook.  _Technical University of Denmark_ , 7(15):510, 2008.
  * Rebuffi et al. (2017) Sylvestre-Alvise Rebuffi, Alexander Kolesnikov, Georg Sperl, and Christoph H Lampert.  ICARL: Incremental classifier and representation learning.  In _Proceedings of the IEEE conference on Computer Vision and Pattern Recognition_ , pp. 2001–2010, 2017.
  * Ritter et al. (2018) Hippolyt Ritter, Aleksandar Botev, and David Barber.  Online structured Laplace approximations for overcoming catastrophic forgetting.  _arXiv preprint arXiv:1805.07810_ , 2018.
  * Rudner et al. (2022a) Tim G. J. Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal.  Tractabe Function-Space Variational Inference in Bayesian Neural Networks.  2022a.
  * Rudner et al. (2022b) Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal.  Continual Learning via Sequential Function-Space Variational Inference.  In _Proceedings of the 38th International Conference on Machine Learning_ , Proceedings of Machine Learning Research. PMLR, 2022b.
  * Schwarz et al. (2018) Jonathan Schwarz, Wojciech Czarnecki, Jelena Luketina, Agnieszka Grabska-Barwinska, Yee Whye Teh, Razvan Pascanu, and Raia Hadsell.  Progress & compress: A scalable framework for continual learning.  In _International Conference on Machine Learning_ , pp. 4528–4537. PMLR, 2018.
  * Snell et al. (2017) Jake Snell, Kevin Swersky, and Richard S Zemel.  Prototypical networks for few-shot learning.  _arXiv preprint arXiv:1703.05175_ , 2017.
  * Swaroop et al. (2019) Siddharth Swaroop, Cuong V Nguyen, Thang D Bui, and Richard E Turner.  Improving and understanding variational continual learning.  _arXiv preprint arXiv:1905.02099_ , 2019.
  * Thrun & Mitchell (1995) Sebastian Thrun and Tom M Mitchell.  Lifelong robot learning.  _Robotics and autonomous systems_ , 15(1-2):25–46, 1995.
  * Titsias et al. (2019) Michalis K Titsias, Jonathan Schwarz, Alexander G de G Matthews, Razvan Pascanu, and Yee Whye Teh.  Functional regularisation for continual learning with Gaussian processes.  _arXiv preprint arXiv:1901.11356_ , 2019.
  * Van de Ven & Tolias (2019) Gido M Van de Ven and Andreas S Tolias.  Three scenarios for continual learning.  _arXiv preprint arXiv:1904.07734_ , 2019.
  * Wilson & Izmailov (2020) Andrew G Wilson and Pavel Izmailov.  Bayesian deep learning and a probabilistic perspective of generalization.  _Advances in neural information processing systems_ , 33:4697–4708, 2020.
  * Zeno et al. (2018) Chen Zeno, Itay Golan, Elad Hoffer, and Daniel Soudry.  Task agnostic continual learning using online variational bayes.  _arXiv preprint arXiv:1803.10123_ , 2018.

## Supplementary Material

##  Appendix A The Toy Gaussians Dataset

See Fig. 7 for a visualization of the toy Gaussians dataset which we use as a
simple CL problem. This is used for evaluating our method for propagating the
true posterior by using HMC for posterior inference and then using a density
estimator on HMC samples as a prior for a new task. We construct $5$, $2$-way
classification problems for CL. Each $2$-way task involves classifying
adjacent circles and squares Fig. 7. With a $2$ layer network with $10$
neurons we get a test accuracy of $1.0$ for the multi-task learning of all $5$
tasks together. Hence according to Eq. 3 a BNN with the same size should be
able to learn all $5$ binary classification tasks continually by sequentially
building up the posterior.

##  Appendix B HMC implementation details

We set the prior for $\mathcal{T}_{1}$, to
$p_{1}({\bm{\theta}})=\mathcal{N}(0,\tau^{-1}\mathbb{I})$ with $\tau=10$. We
burn-in the HMC chain for $1000$ steps and sample for $10000$ more steps and
run $20$ different chains to obtain samples from our posterior, which we then
pass to our density estimator. We use a step size of $0.001$ and trajectory
length of $L=20$, see Appendix C for further implementation details of the
density estimation procedure. For the GMM we optimize for the number of
components by using a holdout set of HMC samples.

##  Appendix C Density Estimation Diagnostics

We provide plots to show that the HMC chains indeed sample from the posterior
have converged in Figure 9 and Figure 11. We run $20$ HMC sampling chains and
randomly select one chain to plot for each seed (of $10$). We run HMC over
$10$ seeds and aggregate the results Figure 3 and Figure 7. The posteriors
$p(\theta|\mathcal{D}_{1}),\ldots$ are approximated with a GMM and used as a
prior for the second task and so forth.

We provide empirical evidence to show that the density estimators have fit to
HMC samples of the posterior in Figure 8 and Figure 10. Where we show the
number of components of the GMM density estimator which we use as a prior for
a new task are all multi-modal posteriors. We show the BNN accuracy when
sampling BNN weights from our GMM all recover the accuracy of the converged
HMC samples. The effective sample size (ESS) of the $20$ chains which is a
measure of how correlated the samples are (higher is better). The reported ESS
values for our experiments are in line with previous work which uses HMC for
BNN inference (Cobb & Jalaian, 2021).

Figure 7: Continual learning binary classification accuracies from the toy
Gaussian dataset similar to (Henning et al., 2021) using $10$ random seeds.
The pink solid line is a multi-task (MT) baseline test accuracy using SGD/HMC.
Figure 8: Diagnostics from using a GMM prior fit to samples of the posterior
generated from HMC, all results are for $10$ random seeds. Left, effective
sample sizes (ESS) of the resulting HMC chains of the posterior, all are
greater than those reported in other works using HMC for BNNs (Cobb & Jalaian,
2021). Middle, the accuracy of the BNN when using samples from the GMM density
estimator instead from the samples from HMC. Right, The optimal number of
components of each GMM posterior fitted with a holdout set of HMC samples by
maximizing the likelihood. Figure 9: Convergence plots from a one randomly
sampled HMC chain (of $20$) for each task over $10$ different runs (seeds) for
$5$ tasks from the toy Gaussians dataset similar to Henning et al. (2021)
(visualized in Fig. 7). We use a GMM density estimator as the prior
conditioned on previous task data. Figure 10: Diagnostics from using a GMM to
fit samples of the posterior HMC samples, all results are for $10$ random
seeds on the toy dataset from Pan et al. (2020) (and visualized in Fig. 3).
Left, effective sample sizes (ESS) of the resulting HMC chains of the
posterior, all are greater than those reported in other works using HMC for
BNNs (Cobb & Jalaian, 2021). Middle left. the current task accuracy from HMC
sampling. Middle right, the accuracy of the BNN when using samples from the
GMM density estimator instead of the converged HMC samples. Right, The optimal
number of components of each GMM posterior fitted with a holdout set of HMC
samples by maximizing the likelihood. Figure 11: Convergence plots from a
randomly sampled HMC chain (of $20$) for each task over $10$ different seeds
for $5$ tasks from the toy dataset from (Pan et al., 2020) (see Fig. 3 for a
visualization of the data). We use a GMM density estimator as a prior.

##  Appendix D Prototypical Bayesian Continual Learning

ProtoCL models the generative process of CL where new tasks are comprised of
new classes $j\in\\{1,\ldots,J\\}$ of a total of $J$ and can be modeled by
using a categorical distribution with a Dirichlet prior:

$\displaystyle y_{i,t}\sim\text{Cat}(p_{1:J}),\quad
p_{1:J}\sim\textrm{Dir}(\alpha_{t}).$ (D.1)

We learn a joint embedding space for our data with a NN,
${\bm{z}}=f({\bm{x}};{\bm{w}})$ with parameters ${\bm{w}}$. The embedding
space for each class is Gaussian whose mean has a prior which is also
Gaussian:

$\displaystyle{\bm{z}}_{it}|y_{it}\sim\mathcal{N}(\bar{{\mathbf{z}}}_{yt},\Sigma_{\epsilon}),\quad\bar{{\mathbf{z}}}_{yt}\sim\mathcal{N}({\bm{\mu}}_{yt},\Lambda^{-1}_{yt}).$
(D.2)

By ensuring that we have an embedding per class and using a memory of past
data, we ensure that the embedding does not drift. The posterior parameters
are $\eta_{t}=\\{\alpha_{t},{\bm{\mu}}_{1:J,t},\Lambda^{-1}_{1:J,t}\\}$.

###  D.1 Inference

As the Dirichlet prior is conjugate with the Categorical distribution and so
is the Gaussian distribution with a Gaussian prior over the mean of the
embedding, then we can calculate posteriors in closed form and update our
parameters as we see new data online without using gradient based updates. We
perform gradient based learning of the NN embedding function
$f(\,\cdot\,;{\bm{w}})$ with parameters ${\bm{w}}$. We optimize the model by
maximizing the log-predictive posterior of the data and use the softmax over
class probabilities to perform predictions. The posterior over class
probabilities $\\{p_{j}\\}_{j=1}^{J}$ and class embeddings
$\\{\bar{{\bm{z}}}_{y_{j}}\\}_{j=1}^{J}$ is denoted as $p(\theta)$ for short
hand and has parameters are
$\eta_{t}=\\{\alpha_{t},{\bm{\mu}}_{1:J,t},\Lambda^{-1}_{1:J,t}\\}$ are
updated in closed form at each iteration of gradient descent.

###  D.2 Sequential updates

We can obtain our posterior:

$\displaystyle p({\bm{\theta}}_{t}|\mathcal{D}_{t})$ $\displaystyle\propto
p(\mathcal{D}_{t}|{\bm{\theta}}_{t})p({\bm{\theta}}_{t})$ (D.3)
$\displaystyle=\prod^{N_{t}}_{i=1}p({\bm{z}}^{i}_{t}|y^{i}_{t};\bar{{\bm{z}}}_{y_{t}},\Sigma_{\epsilon,y_{t}})p(y^{i}_{t}|p_{1:J})p(p_{i:J};\alpha_{t})p(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{y_{t},t},\Lambda^{-1}_{y_{t},t})$
(D.4)
$\displaystyle=\mathcal{N}(\mu_{t+1},\Sigma_{t+1})\textrm{Dir}(\alpha_{t+1}),$
(D.5)

where $N_{t}$ is the number of data points seen during update $t$.
Concentrating on the Categorical-Dirichlet conjugacy:

$\displaystyle\textrm{Dir}(\alpha_{t+1})$ $\displaystyle\propto
p(p_{1:J};\alpha_{t})\prod^{N_{t}}_{i=1}p(y^{i}_{t};p_{i:J})$ (D.6)
$\displaystyle\propto\prod^{J}_{j=1}p_{j}^{\alpha_{j}-1}\prod^{N_{t}}_{i=1}\prod^{J}_{j=1}p_{j}^{\mathbb{I}(y_{t}^{i}=j)}$
(D.7)
$\displaystyle=\prod^{J}_{j=1}p_{j}^{\alpha_{j}-1+\sum_{i=1}^{N_{t}}\mathbb{I}(y_{t}^{i}=j)}.$
(D.8)

Thus:

$\displaystyle\alpha_{t+1,j}=\alpha_{t,j}+\sum^{N_{t}}_{i=1}\mathbb{I}(y^{i}_{t}=j).$
(D.9)

Also, due to Gaussian-Gaussian conjugacy, then the posterior for the Gaussian
prototype of the embedding for each class is:

$\displaystyle\mathcal{N}({\bm{\mu}}_{t+1},\Lambda_{t+1})$
$\displaystyle\propto\prod^{N_{t}}_{i=1}\mathcal{N}({\bm{z}}_{t}^{i}|y^{i}_{t};\bar{{\bm{z}}}_{y_{t}},\Sigma_{\epsilon})\mathcal{N}(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{y_{t},t},\Lambda^{-1}_{y_{t}})$
(D.10)
$\displaystyle=\prod_{y_{t}\in\\{1,\ldots,J\\}}\mathcal{N}({\bm{z}}_{y_{t}}|y_{t};\bar{{\bm{z}}}_{y_{t}},\frac{1}{N_{y_{t}}}\Sigma_{\epsilon})\mathcal{N}(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{y_{t+1}},\Lambda^{-1}_{y_{t}})$
(D.11)
$\displaystyle=\prod_{y_{t}\in\\{1,\ldots,J\\}}\mathcal{N}(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{t+1},\Lambda^{-1}_{y_{t+1}}),$
(D.12)

where $N_{y_{t}}$ is the number of points of class $y_{t}$ from the set of all
classes $C=\\{1,\ldots,J\\}$. The update equations for the mean and variance
of the posterior are:

$\displaystyle\Lambda_{y_{t+1}}$
$\displaystyle=\Lambda_{y_{t}}+N_{y_{t}}\Sigma^{-1}_{\epsilon},\quad\forall
y_{t}\in C_{t}$ (D.13) $\displaystyle\Lambda_{y_{t+1}}{\bm{\mu}}_{y_{t+1}}$
$\displaystyle=N_{y_{t}}\Sigma^{-1}_{\epsilon}\bar{{\bm{z}}}_{y_{t}}+\Lambda_{y_{t}}{\bm{\mu}}_{y_{t}},\quad\forall
y_{t}\in C_{t}.$ (D.14)

###  D.3 ProtoCL Objective

The posterior predictive distribution we want to optimize is:

$\displaystyle p({\bm{z}},y)=\int
p({\bm{z}},y|{\bm{\theta}};\eta)p({\bm{\theta}};\eta)d{\bm{\theta}},$ (D.15)

where $p({\bm{\theta}})$ denotes the distributions over class probabilities
$\\{p_{j}\\}_{j=1}^{J}$ and mean embeddings
$\\{\bar{{\bm{z}}}_{y_{j}}\\}_{j=1}^{J}$,

$\displaystyle p({\bm{z}},y)$
$\displaystyle=\int\prod^{N_{t}}_{i=1}p({\bm{z}}_{it}|y_{it};\bar{{\bm{z}}}_{y_{t}},\Sigma_{\epsilon})p(y_{it}|p_{1:J})p(p_{1:J};\alpha_{t})p(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{y_{t},t},\Lambda^{-1}_{y_{t},t})dp_{1:J}d\bar{{\bm{z}}}_{y_{t}}$
(D.16)
$\displaystyle=\int\prod^{N_{t}}_{i=1}p({\bm{z}}_{it}|y_{it};{\bm{z}}_{y_{t}},\Sigma_{\epsilon})p(\bar{{\bm{z}}}_{y_{t}};{\bm{\mu}}_{y_{t},t},\Lambda^{-1}_{y_{t},t})d\bar{{\bm{z}}}_{y_{t}}\underbrace{\int\prod^{N_{t}}_{i=1}p(y_{it}|p_{1:J})p(p_{1:J};\alpha_{t})dp_{1:J}}_{\prod_{i}p(y_{i})=p(y)}$
(D.17)
$\displaystyle=p(y)\prod^{N_{t}}_{i=1}Z_{i}^{-1}\int\mathcal{N}(\bar{{\bm{z}}}_{y_{it}};{\bm{c}},C)d\bar{{\bm{z}}}_{y_{t}}$
(D.18)
$\displaystyle=p(y)\prod^{N_{t}}_{i=1}\mathcal{N}({\bm{z}}_{it}|y_{it};{\bm{\mu}}_{y_{t},t},\Sigma_{\epsilon}+\Lambda^{-1}_{y_{t},t}).$
(D.19)

Where in Eq. (D.18) we use §8.1.8 in (Petersen et al., 2008). The term $p(y)$
is:

$\displaystyle p(y)$ $\displaystyle=\int
p(y|p_{1:J})p(p_{1:J};\alpha_{t})dp_{1:J}$ (D.20) $\displaystyle=\int
p_{y}\frac{\Gamma(\sum^{J}_{j=1}\alpha_{j})}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\prod^{J}_{j=1}p_{j}^{\alpha_{j}-1}dp_{1:J}$
(D.21)
$\displaystyle=\frac{\Gamma(\sum^{J}_{j=1}\alpha_{j})}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\int\prod^{J}_{j=1}p_{j}^{\mathbb{I}(y=j)+\alpha_{j}-1}dp_{1:J}$
(D.22)
$\displaystyle=\frac{\Gamma(\sum^{J}_{j=1}\alpha_{j})}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\frac{\prod^{J}_{j=1}\Gamma(\mathbb{I}(y=j)+\alpha_{j})}{\Gamma(1+\sum^{J}_{j=1}\alpha_{j})}$
(D.23)
$\displaystyle=\frac{\cancel{\Gamma(\sum^{J}_{j=1}\alpha_{j})}}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\frac{\prod^{J}_{j=1}\Gamma(\mathbb{I}(y=j)+\alpha_{j})}{\sum^{J}_{j=1}\alpha_{j}\cancel{\Gamma(\sum^{J}_{j=1}\alpha_{j})}}$
(D.24) $\displaystyle=\frac{\prod^{J}_{j=1,j\neq
y}\Gamma(\alpha_{j})}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\frac{\Gamma(1+\alpha_{y})}{\sum^{J}_{j=1}\alpha_{j}}$
(D.25) $\displaystyle=\frac{\prod^{J}_{j=1,j\neq
y}\Gamma(\alpha_{j})}{\prod^{J}_{j=1}\Gamma(\alpha_{j})}\frac{\alpha_{y}\Gamma(\alpha_{y})}{\sum^{J}_{j=1}\alpha_{j}}$
(D.26)

where we use the identity $\Gamma(n+1)=n\Gamma(n)$.

Figure 12: Split-MNIST average test accuracy over all $5$ tasks for different
memory sizes. Accuracies are over $10$ seeds.

Figure 13: The evolution of the Dirichlet parameters $\alpha_{t}$ for each
class in Split-MNIST tasks for ProtoCL. All $\alpha_{j}$ are shown over $10$
seeds with $\pm 1$ standard error. By the end of training all classes are
roughly equally likely, as we have trained on equal amounts of all classes.

###  D.4 Predictions

To make a prediction for a test point ${\bm{x}}^{*}$:

$\displaystyle p(y^{*}=j|{\bm{x}}^{*},{\bm{x}}_{1:t},y_{1:t})$
$\displaystyle=p(y^{*}=j|{\bm{z}}^{*},{\bm{\theta}}_{t})$ (D.28)
$\displaystyle=\frac{p({\bm{z}}^{*}|y^{*}=j,{\bm{\theta}}_{t})p(y^{*}=j|{\bm{\theta}}_{t})}{\sum_{i}p({\bm{z}}^{*}|y^{*}=i,\eta_{t})p(y^{*}=i|{\bm{\theta}}_{t})}$
(D.29)
$\displaystyle=\frac{p(y^{*}=j,{\bm{z}}^{*}|{\bm{\theta}}_{t})}{\sum_{i}p(y=i,{\bm{z}}^{*}|{\bm{\theta}}_{t})},$
(D.30)

where ${\bm{\theta}}_{t}$ are sufficient statistics for
$({\bm{x}}_{1:t},y_{1:t})$.

Preventing forgetting. As we wish to retain the task specific prototypes, at
the end of learning a task $\mathcal{T}_{t}$ we take a small subset of the
data as a memory to ensure that posterior parameters and prototypes do not
drift, see Algorithm 1.

###  D.5 Experimental Setup

The prototype variance, $\Sigma_{\epsilon}$ is set to a diagonal matrix with
the variances of each prototype set to $0.05$. The prototype prior precisions,
$\Lambda_{yt}$ are also diagonals and initialized randomly and exponentiated
to ensure a positive semi-definite covariance for the sequential updates. The
parameters $\alpha_{j}\,\forall j$ are set to $0.78$ which was found by random
search over the validation set on MNIST. We also allow $\alpha_{j}$ to be
learned in the gradient update step in addition to the sequential update step
(lines 4 and 5 Algorithm 1), see Fig. 13 to see the evolution of the
$\alpha_{j}$ or all classes $j$ over the course of learning Split-MNIST.

For the Split-MNIST and Split-FMNIST benchmarks we use a NN with $2$ layers of
size $200$ and trained for $50$ epochs with an Adam optimizer. We perform a
grid-search over learning rates, dropout rates and weight decay coefficients.
The embedding dimension is set to $128$. For the Split-CIFAR10 and Split-
CIFAR100 benchmarks, we use the same network as Pan et al. (2020) which
consists of $4$ convolution layers and $2$ linear layers. We train the
networks for $80$ epochs for each task with the Adam optimizer with a learning
rate of $1e-3$. The embedding dimension is set to $32$. All experiments are
run a single GPU NVIDIA RTX 3090.

##  Appendix E Sequential Bayesian Estimation as Bayesian Neural Network
optimization

${\bm{\theta}}^{*}_{t-1}$${\bm{\theta}}^{*}_{t-2}$${\bm{\theta}}^{*}_{t}$${\bm{\theta}}^{*}_{t+1}$$y_{t-1}$$y_{t-2}$$y_{t}$$y_{t+1}$$x_{t-1}$$x_{t-2}$$x_{t}$$x_{t+1}$$\,\ldots\,$$\,\ldots\,$
Figure 14: Graphical model of under which we perform inference in Section 5.
Grey nodes are observed and white are latent variables.

We shall consider inference in the graphical model depicted in Fig. 14. The
aim is to infer the optimal BNN weights, ${\bm{\theta}}^{*}_{t}$ at time $t$
given observations and the previous BNN weights. We assume a Gaussian
posterior over weights with full covariance hence we model interactions
between all weights. We shall consider the online setting where we see one
data point $({\bm{x}}_{t},y_{t})$ at a time and we will make no assumption as
to whether the data comes from the same task or different tasks over the
course of learning.

We set up the problem of sequential Bayesian inference as a filtering problem
and we leverage the work of Aitchison (2018) which casts NN optimization as
Bayesian sequential inference. We make the reasonable assumption that the
distribution over weights is a Gaussian with full covariance. Since reasoning
about the full covariance matrix of a BNN is intractable we instead consider
the $i$-th parameter and reason about the dynamics of the optimal estimates
$\theta^{*}_{it}$ as a function of all the other parameters
${\bm{\theta}}_{-it}$. Each weight is functionally dependent on all others. If
we had access to the full covariance of the parameters then we could reason
about the unknown optimal weights $\theta^{*}_{it}$ given the values of all
the other weights ${\bm{\theta}}_{-it}$. However, since we do not have access
to the full covariance another approach is to reason about the dynamics of
$\theta^{*}_{it}$ given the dynamics of ${\bm{\theta}}_{-it}$ and assume that
the values of the weights are close to those of the previous time-step (Jacot
et al., 2018) and so we cast the problem as a dynamical system.

Consider a quadratic loss of the form:

$\displaystyle\mathcal{L}({\bm{x}}_{t},y_{t};{\bm{\theta}})=\mathcal{L}_{t}({\bm{\theta}})=-\frac{1}{2}{\bm{\theta}}^{\top}{\bm{H}}{\bm{\theta}}+{\mathbf{z}}^{\top}_{t}{\bm{\theta}},$
(E.31)

which we can arrive by simple Taylor expansion where ${\bm{H}}$ is the Hessian
which is assumed to constant across data points but not across the parameters
${\bm{\theta}}$. If the BNN output takes the form
$f({\bm{x}}_{t};{\bm{\theta}})$, then the derivative evaluated at
${\bm{\theta}}_{t}$ is
${\bm{z}}_{t}=\frac{\partial\mathcal{L}({\bm{x}}_{t},y_{t};{\bm{\theta}})}{\partial{\bm{\theta}}}|_{{\bm{\theta}}={\bm{\theta}}_{t}}$.
To construct the dynamical equations for our weights, consider the gradient
for a single datapoint:

$\displaystyle\frac{\partial\mathcal{L}_{t}({\bm{\theta}})}{\partial{\bm{\theta}}}=-{\bm{H}}{\bm{\theta}}+{\mathbf{z}}_{t}.$
(E.32)

If we consider the gradient for the $i$-th weight while all other parameters
are set to their current estimate:

$\displaystyle\frac{\partial\mathcal{L}(\theta_{i},{\bm{\theta}}_{-i})}{\partial\theta_{i}}=-H_{ii}\theta_{it}-{\bm{H}}_{-ii}^{\top}{\bm{\theta}}_{-it}+z_{ti}.$
(E.33)

When the gradient is set to zero we recover the optimal value for
$\theta_{it}$, denoted as $\theta_{it}^{*}$:

$\displaystyle\theta_{it}^{*}=-\frac{1}{H_{ii}}{\bm{H}}_{-ii}^{\top}{\bm{\theta}}_{-it}.$
(E.34)

since $z_{ti}=0$ at the modes. The equation above shows us that the dynamics
of the optimal weight $\theta^{*}_{it}$ is dependent on all the other current
values of the parameters ${\bm{\theta}}_{-it}$. That is, the dynamics of
$\theta^{*}_{it}$ will be governed by the dynamics of the weights
${\bm{\theta}}_{-it}$. The dynamics of ${\bm{\theta}}_{-it}$ will be a complex
stochastic process dependent on many different variables. Since reasoning
about the dynamics is intractable we instead assume a discretized Ornstein-
Uhlenbeck process for the weights ${\bm{\theta}}_{-it}$ with reversion speed
$\vartheta\in\mathbb{R}_{+}$ and noise variance $\eta^{2}_{-i}$:

$\displaystyle
p({\bm{\theta}}_{-i,t+1}|{\bm{\theta}}_{-i,t})=\mathcal{N}((1-\vartheta){\bm{\theta}}_{-it},\eta_{-i}^{2}),$
(E.35)

this implies that the dynamics for the optimal weight is defined by

$\displaystyle
p(\theta^{*}_{i,t+1}|\theta^{*}_{i,t})=\mathcal{N}((1-\vartheta)\theta^{*}_{it},\eta^{2}),$
(E.36)

where $\eta^{2}=\eta^{2}_{-i}{\bm{H}}^{\top}_{-ii}{\bm{H}}_{-ii}$. This same
assumption is made in Aitchison (2018). This assumes a parsimonious model of
the dynamics. Together with our likelihood:

$\displaystyle
p(y_{t}|{\bm{x}}_{t};{\bm{\theta}}^{*}_{t})=\mathcal{N}(y_{t};f({\bm{x}}_{t};{\bm{\theta}}^{*}_{t}),\sigma^{2})$
(E.37)

where $f(\,\cdot\,,{\bm{\theta}})$ is a neural network prediction with weights
${\bm{\theta}}$, we can now define a linear dynamical system for the optimal
weight $\theta^{*}_{i}$ by linearizing the the Bayesian NN (Jacot et al.,
2018) and by using the transition dynamics in Eq. E.36. Thus we are able to
infer the posterior distribution over the optimal weights using Kalman filter
like updates (Kalman, 1960). As the dynamics and likelihood are Gaussian, then
the prior and posterior are also Gaussian, for ease of notation we drop the
index $i$ such that $\theta^{*}_{it}=\theta^{*}_{t}$:

$\displaystyle p(\theta^{*}_{t}|({\bm{x}},y)_{t-1},\ldots,({\bm{x}},y)_{1})$
$\displaystyle=\mathcal{N}(\mu_{t,\text{prior}},\sigma^{2}_{t,\text{prior}})$
(E.38) $\displaystyle
p(\theta^{*}_{t}|({\bm{x}},y)_{t},\ldots,({\bm{x}},y)_{1})$
$\displaystyle=\mathcal{N}(\mu_{t,\text{post}},\sigma^{2}_{t,\text{post}})$
(E.39)

By using the transition dynamics and the prior we can obtain closed form
updates:

$\displaystyle p(\theta^{*}_{t}|({\bm{x}},y)_{t-1},\ldots,({\bm{x}},y)_{1})$
$\displaystyle=\int
p(\theta^{*}_{t}|\theta^{*}_{t-1})p(\theta^{*}_{t-1}|({\bm{x}},y)_{t-1},\ldots,({\bm{x}},y)_{1})d\theta^{*}_{t-1}$
(E.40)
$\displaystyle\mathcal{N}(\theta^{*}_{t};\mu_{t,\text{prior}},\sigma^{2}_{t,\text{prior}})$
$\displaystyle=\int\mathcal{N}(\theta^{*}_{t};(1-\vartheta)\theta^{*}_{t-1},\eta^{2})\mathcal{N}(\theta^{*}_{t-1};\mu_{t-1,\text{post}},\sigma^{2}_{t-1,\text{post}})d\theta^{*}_{t-1}.$
(E.41)

Integrating out $\theta^{*}_{t-1}$ we can get updates for the prior for the
next timestep as follows:

$\displaystyle\mu_{t,\text{prior}}$
$\displaystyle=(1-\vartheta)\mu_{t-1,post}$ (E.42)
$\displaystyle\sigma^{2}_{t,\text{prior}}$
$\displaystyle=\eta^{2}+(1-\vartheta)^{-2}\sigma^{2}_{t-1,\text{post}}.$
(E.43)

The updates for obtaining our posterior parameters: $\mu_{t,\text{post}}$ and
$\sigma^{2}_{t,\text{post}}$, comes from applying Bayes’ theorem:

$\displaystyle\log\mathcal{N}(\theta^{*}_{t};\mu_{t,\text{post}},\sigma^{2}_{t,\text{post}})$
$\displaystyle\propto\log\mathcal{N}(y_{t};f({\bm{x}}_{t};\theta^{*}_{t}),\sigma^{2})+\log\mathcal{N}(\theta^{*}_{t};\mu_{t,\text{prior}},\sigma^{2}_{t,\text{prior}}),$
(E.44)

by linearizing our Bayesian NN such that $f({\bm{x}}_{t},\theta_{0})\approx
f({\bm{x}}_{t},\theta_{0})+\frac{\partial
f({\bm{x}}_{t};\theta^{*}_{t})}{\partial\theta^{*}_{t}}(\theta^{*}_{t}-\theta_{0})$
and by substituting into Eq. E.44 we obtain our update equation for the
posterior of the mean of our BNN parameters:

$\displaystyle-\frac{1}{2\sigma^{2}_{t,\text{post}}}(\theta^{*}_{t}-\mu_{t,\text{post}})^{2}$
$\displaystyle=-\frac{1}{2\sigma^{2}}(y-g({\bm{x}}_{t})\theta^{*}_{t})^{2}-\frac{1}{2\sigma^{2}_{t,\text{prior}}}(\theta^{*}_{t}-\mu_{t,\text{prior}})^{2}$
(E.45) $\displaystyle\mu_{t,\text{post}}$
$\displaystyle=\sigma^{2}_{t,\text{post}}\left(\frac{\mu_{t,\text{prior}}}{\sigma^{2}_{t,\text{prior}}}+\frac{y}{\sigma^{2}}g({\bm{x}}_{t})\right),$
(E.46)

where $g({\bm{x}}_{t})=\frac{\partial
f({\bm{x}}_{t};\theta^{*}_{t})}{\partial\theta^{*}_{t}}$, and the update
equation for the variance of the Gaussian posterior is:

$\displaystyle\frac{1}{\sigma^{2}_{t,\text{post}}}$
$\displaystyle=\frac{g({\bm{x}}_{t})^{2}}{\sigma^{2}}+\frac{1}{\sigma^{2}_{t,\text{prior}}}.$
(E.47)

From our update equations Eq. E.46 and Eq. E.47 we can notice that the
posterior mean depends linearly on the prior and an additional data dependent
term. These equations are similar to the filtering example in Section 4. So
under certain assumptions above, a BNN should behave similarly. If there
exists a task data imbalance then the data term will dominate the prior term
in Eq. E.46 and could lead to forgetting of previous tasks.