On Sequential Bayesian Inference for Continual Learning

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 assess whether using the previous task’s posterior as a prior for a new task can prevent catastrophic forgetting in Bayesian neural networks. Our first contribution is to perform sequential Bayesian inference using Hamiltonian Monte Carlo. We propagate the posterior as a prior for new tasks by approximating the posterior via 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. Our final contribution is to propose a simple baseline called Prototypical Bayesian Continual Learning, which is competitive with the best performing Bayesian continual learning methods on class incremental continual learning computer 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 [1,2]. 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 [3]. However, obtaining a full posterior over NN weights is computationally demanding and we often need to resort to approximations, such as the Laplace method [4] or variational inference [5,6] 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 [7,8,9,10,11,12]. 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. Therefore, 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) [13]. 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. [9]) 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,\dots ,T$. Each task, ${\mathcal{T}}_t$, is comprised of a dataset ${\mathcal{D}}_t={\left\{({\mathit{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\left(y\right|\mathbf{x}),p_t\left(\mathbf{x}\right)\}$. 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\le t$. We decompose predictors as $g=h\circ f$ such that $\widehat{y}=g\left(\mathit{x}\right)$. We define f as an embedding function mapping $f:\mathcal{X}\to \mathcal{Z}$ and h as a head mapping to outputs $h:\mathcal{Z}\to \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 timesteps, $t=1,2,\dots ,T$. At the first timestep, $t=1$, that is, for task ${\mathcal{T}}_1$, the model receives the first dataset ${\mathcal{D}}_1$ and learns the conditional distribution $p\left(y_i\right|{\mathit{x}}_i,\mathit{\theta })$ for all $({\mathit{x}}_i,y_i)\in {\mathcal{D}}_1$ (i indexes a datapoint in ${\mathcal{D}}_1$). We denote the parameters $\mathit{\theta }$ as having a prior distribution $p\left(\mathit{\theta }\right)$ for ${\mathcal{T}}_1$. The posterior predictive distribution for a test point ${\mathit{x}}_1^{*}\in {\mathcal{D}}_1$ is hence:

$$\begin{array}{c}\hfill p\left(y_1^{*}\right|{\mathit{x}}_1^{*},{\mathcal{D}}_1)=\int p\left(y_1^{*}\right|{\mathit{x}}_1^{*},\mathit{\theta })p\left(\mathit{\theta }\right|{\mathcal{D}}_1)d\mathit{\theta }.\end{array}$$

(1)

We note that computing this posterior predictive distribution requires $p\left(\mathit{\theta }\right|{\mathcal{D}}_1)$. For $t=2$, a CL model is required to fit $p\left(y_i\right|{\mathit{x}}_i,\mathit{\theta })$ for $({\mathit{x}}_i,y_i)\in {\mathcal{D}}_1\cup {\mathcal{D}}_2$. The posterior predictive distribution for a new test point ${\mathit{x}}_2^{*}\in {\mathcal{D}}_1\cup {\mathcal{D}}_2$ point is:

$$\begin{array}{c}\hfill p\left(y_2^{*}\right|{\mathit{x}}_2^{*},{\mathcal{D}}_1,{\mathcal{D}}_2)=\int p\left(y_2^{*}\right|{\mathit{x}}_2^{*},\mathit{\theta })p\left(\mathit{\theta }\right|{\mathcal{D}}_1,{\mathcal{D}}_2)d\mathit{\theta }.\end{array}$$

(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\left(\mathit{\theta }\right|{\mathcal{D}}_1,\dots ,{\mathcal{D}}_T)$ as:

$$\begin{array}{cc}\hfill p\left(\mathit{\theta }\right|{\mathcal{D}}_1,\dots ,{\mathcal{D}}_{T-1},{\mathcal{D}}_T)& =\frac{p\left({\mathcal{D}}_T\right|\mathit{\theta })p\left(\mathit{\theta }\right|{\mathcal{D}}_1,\dots ,{\mathcal{D}}_{T-1})}{p\left({\mathcal{D}}_T\right|{\mathcal{D}}_1,\dots ,{\mathcal{D}}_{T-1})}.\hfill \end{array}$$

(3)

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

$$\begin{array}{cc}\hfill logp\left(\mathit{\theta }\right|{\mathcal{D}}_1)& =logp\left({\mathcal{D}}_1\right|\mathit{\theta })+logp\left(\mathit{\theta }\right)-logp\left({\mathcal{D}}_1\right).\hfill \end{array}$$

(4)

At $t=2$, we require $p\left(\mathit{\theta }\right|{\mathcal{D}}_1,{\mathcal{D}}_2)$ to calculate the posterior predictive distribution in Equation (2). However, we have lost access to ${\mathcal{D}}_1$. According to Bayes’ rule, the posterior may be written as:

$$\begin{array}{cc}\hfill logp\left(\mathit{\theta }\right|{\mathcal{D}}_1,{\mathcal{D}}_2)& =logp\left({\mathcal{D}}_2\right|\mathit{\theta })+logp\left(\mathit{\theta }\right|{\mathcal{D}}_1)-logp\left({\mathcal{D}}_2\right|{\mathcal{D}}_1),\hfill \end{array}$$

(5)

where we used the conditional independence of ${\mathcal{D}}_2$ and ${\mathcal{D}}_1$ given $\mathit{\theta }$. We note that the likelihood $p\left({\mathcal{D}}_2\right|\mathit{\theta })$ is only dependent upon the current task dataset, ${\mathcal{D}}_2$, and that the prior $p\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ encodes parameter knowledge from the previous task. Hence, we can use the posterior evaluated at t as a prior for learning a new task at $t+1$. From Equation (3), we require that our model with parameters $\mathit{\theta }$ is a sufficient statistic of ${\mathcal{D}}_1$, i.e., $p\left({\mathcal{D}}_2\right|\mathit{\theta },{\mathcal{D}}_1)=p\left({\mathcal{D}}_2\right|\mathit{\theta })$, making the likelihood conditionally independent of ${\mathcal{D}}_1$ given $\mathit{\theta }$. This observation motivates the use of high-capacity predictors, such as Bayesian neural networks, that are flexible enough to learn from ${\mathcal{D}}_1$.

Continual Learning Example: Split-MNIST

For the MNIST dataset [14], we know that if we were to train a BNN we would achieve good performance by inferring the posterior $p\left(\mathit{\theta }\right|\mathcal{D})$ and integrating out the posterior to infer the posterior predictive over a test point Equation (1). Therefore, if we were to split the dataset MNIST into 5 two-class classification tasks, then we should be able to recursively recover the multi-task posterior $p\left(\mathit{\theta }\right|\mathcal{D})=p\left(\mathit{\theta }\right|{\mathcal{D}}_1\dots ,{\mathcal{D}}_5)$ using Equation (3). This problem is called Split-MNIST [15], where the first task involves the classification of the digits $\{0,1\}$, the second task classification of the digits $\{2,3\}$, and so on.

We can define three different CL settings [16,17,18]. When we allow the CL agent to make predictions with a task identifier $\tau$ the scenario is referred to as task-incremental. The identifier $\tau$ could be used to select different heads Section 2.1, for instance. This scenario is not compatible with sequential Bayesian inference outlined in Equation (3) since no task identifier is required for making predictions. Domain-incremental learning is another scenario that does not have access to $\tau$ during evaluation and requires the CL agent to perform classification to the same output space for each task; for example, for Split-MNIST the output space is $\{0,1\}$ for all tasks, so this amounts to classifying between even and odd digits. Domain incremental learning is compatible with sequential Bayesian inference with a Bernoulli likelihood. The third scenario is class-incremental learning which also does not have access to $\tau$ but the agent needs to classify each example to its corresponding class. For Split-MNIST, for example, the output space is $\{0,\dots ,9\}$ for each task. Class-incremental learning is compatible with sequential Bayesian inference with a categorical likelihood.

2.3. Variational Continual Learning

Variational CL (VCL; Nguyen et al. [9]) simplifies the Bayesian inference problem in Equation (3) into a sequence of approximate Bayesian updates on the distribution over random neural network weights $\mathit{\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\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ that maximizes a corresponding variational objective. For the subsequent task, the prior is chosen to be $q_1\left(\mathit{\theta }\right|{\mathcal{D}}_1)$, and the goal becomes to learn a variational distribution $q_2\left(\mathit{\theta }\right|{\mathcal{D}}_2)$ that maximizes a corresponding variational objective under this prior. Denoting the recursive posterior inferred from multiple datasets by $q_t\left(\mathit{\theta }\right|{\mathcal{D}}_{1:t})$, we can express the variational CL objective for the t-th task as:

$$\begin{array}{cc}\hfill \mathcal{L}(\mathit{\theta },{\mathcal{D}}_t)& ={\mathbb{D}}_{\mathrm{KL}}\left[q_t\left(\mathit{\theta }\right)\left|\right|q_{t-1}\left(\mathit{\theta }\right|{\mathcal{D}}_{1:t-1})\right]-{\mathbb{E}}_{q_t}[logp\left({\mathcal{D}}_t\right|\mathit{\theta })].\hfill \end{array}$$

(6)

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 Equation (3). Therefore, why does single-head VCL not improve over SGD if we can recursively build up an approximate posterior using Equation (3)? 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 [13]. We then use these samples and learn a density estimator that can be used as a prior for a new task (we considered Sequential Monte Carlo, but it is unable to scale to the dimensions required for the NNs we consider [19]. HMC on the other hand has recently been successfully scaled to relatively small BNNs of the size considered in this paper [20] and ResNet models but at large computational cost [21]). HMC is considered the gold standard in approximate inference and is guaranteed to asymptotically produce samples from the true posterior (in the NeurIPS 2021 Bayesian Deep Learning Competition (https://izmailovpavel.github.io/neurips_bdl_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 $\mathit{\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 $\mathit{\theta }$ rather than a diagonal Gaussian variational posterior such as in VCL. More concretely, to propagate the posterior $p\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ we use a density estimator, defined $\widehat{p}\left(\mathit{\theta }\right|{\mathcal{D}}_1)$, to fit a probability density on HMC samples as a posterior. For the next task ${\mathcal{T}}_2$ we can use $\widehat{p}\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ as a prior for a new HMC sampling chain and so on (see Figure 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.

We use a toy dataset (Figure 3) with two classes and inputs $x\in {\mathbb{R}}^2$ [22]. 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 a hidden state size of 10. We use 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 [23]; however, these did not work robustly for HMC sampling (RealNVP 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 such as HMC.

Training a BNN with HMC on the same multi-task dataset obtains a test accuracy of $1.0$. Thus, the final posterior is suitable for continual learning under Equation (3) and we should be able to recursively arrive at the multi-task posterior with our recursive inference method with HMC. The results from Figure 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. For ${\mathcal{T}}_2$, we use $\widehat{p}\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ instead of $p\left(\mathit{\theta }\right|{\mathcal{D}}_1)$ as a prior and this will bias the HMC sampling for all subsequent tasks. In the next paragraph, we detail the measures taken to ensure that our HMC chains have converged so we are sampling from the true posterior. Moreover, we access the fidelity of the GMM density estimator with respect to the HMC samples. We also repeated these experiments with another toy dataset of five binary classification tasks where we observe similar results Appendix A.

For HMC, we ensure that we are sampling from the posterior by assessing chain convergence and effective sample sizes (Figure A5). 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 [20]. Moreover, we ensure that the GMM approximate posterior is multi-modal and has a more complex posterior in comparison to VCL, and that the GMM samples produce equivalent results to HMC samples for the current task (Figure A4). See Appendix B for details.

The 2-d benchmarks we consider in this section are from previous works and are domain-incremental continual learning problems. The domain incremental setting is also simpler [18] than the class-incremental setting and thus a good starting point when attempting to perform exact sequential Bayesian inference. Despite this, 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, more accurate, and multi-modal posteriors. Despite this, we are still not able to sequentially build up the multi-task posterior or obtain much better results than an isotropic Gaussian posterior such as 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

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

Observations $y_1,y_2,\dots ,y_t$ arrive online, and each observation is generated by a hidden variable ${\theta }_1,{\theta }_2,\dots ,{\theta }_t\sim p$ where p is a probability density function. At time t, we wish to infer the filtering distribution$p\left({\theta }_t\right|y_1,y_2,\dots ,y_t)$ [24] using sequential Bayesian inference, similarly to the Kalman filter [25]. The likelihood is $p\left(y_t\right|{\theta }_t)=\mathcal{N}(y_t;f(·;{\theta }_t),{\sigma }^2)$ such that the mean is parameterized by a linear model $y_t=f(·;{\theta }_t)+ϵ$ where $ϵ\sim \mathcal{N}(0,{\sigma }^2)$ and $f(·;{\theta }_t)={\theta }_t$. We consider a Gaussian prior over the mean parameters $\theta$ such that $p\left({\theta }_0\right)=\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};{\widehat{\theta }}_{t-1},{\widehat{\sigma }}_{t-1}^2)$ and $\mathcal{N}({\theta }_t;{\widehat{\theta }}_t,{\widehat{\sigma }}_t^2)$. By using sequential Bayesian inference we can have closed-form update equations for our posterior parameters:

$$\begin{array}{cc}\hfill {\widehat{\theta }}_t={\widehat{\sigma }}_t^2\left(\frac{y_t}{{\sigma }^2}+\frac{{\widehat{\theta }}_{t-1}}{{\widehat{\sigma }}_{t-1}^2}\right)={\widehat{\sigma }}_t^2\left(\sum _{i=1}^t\frac{y_i}{{\sigma }^2}+\frac{{\widehat{\theta }}_0}{{\widehat{\sigma }}_0^2}\right),\frac{1}{{\widehat{\sigma }}_t^2}& =\frac{1}{{\sigma }^2}+\frac{1}{{\widehat{\sigma }}_{t-1}^2}.\hfill \end{array}$$

(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. Therefore, the posterior mean is a weighted sum of the data and the final value of the posterior is not dependent on the order of the data. We consider the situation where there is a task change (this non-stationarity is referred to as a changepoint in the time-series literature, in Figure 4A at $t=110$). Concretely, for task 1 the dataset is generated according to $\mathcal{N}(-1,{\sigma }^2)$, so we want the model to regress to this task. For task 2, the data is generated according to $\mathcal{N}(1,{\sigma }^2)$ and so we want our continual learning agent to regress well to this task too. As with all continual learning benchmarks, we require our model to retain performance on past tasks and perform equally well on both tasks at the end of training at $t=220$. From Figure 4A, we can see that the linear model will regress to the first dataset well, as data is seen online and the linear model is updated online. However, as is seen in data from the second task, the linear model eventually tracks the global mean over both tasks Equation (7) rather than a mean for each task Figure 4A. This is even more pronounced when there is a task dataset imbalance Figure 4B.

The model is clearly misspecified since a linear model cannot regress to both of these tasks simultaneously. A more suitable model would be a mixture model, which is able to regress to both task datasets. Despite performing exact inference, a misspecified model can forget, Figure 4. Performance on the first task is reduced while learning the second task, this becomes even more pronounced with task dataset imbalances Figure 4B. In the case of HMC, we verified that our Bayesian neural network had a perfect performance on all tasks beforehand. In Figure 3, we had a well-specified model but struggled with exact sequential Bayesian inference, Equation (3). When learning with linear models online, 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 that 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. Therefore, in the scenario of online continual learning [26,27,28], 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 [29]. However, we struggle to fit the posterior samples from HMC to perform sequential Bayesian inference in Section 3.

We continue to use Bayesian filtering and assume a Bayesian NN where the posterior is Gaussian with a full covariance. By modeling the entire covariance, we enable modeling of how each individual weight varies with respect to all others. We do this by interpreting online learning in Bayesian NNs as filtering [30]. Our treatment is similar to Aitchison [31], who derives an optimizer by leveraging Bayesian filtering. We consider inference in the graphical model depicted in Figure 5. The aim is to infer the optimal BNN weights, ${\mathit{\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 $({\mathit{x}}_t,y_t)$ is observed at a time.

Instead of modeling the full covariance, we instead consider each parameter ${\theta }_i$ as a function of all the other parameters ${\mathit{\theta }}_{-it}$. We also assume that the values of the weights are close to those of the previous timestep [32]. To obtain the updated equations for BNN parameters given a new observation and prior, we make two simplifying assumptions as follows.

Assumption A1.

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

$$\begin{array}{c}\hfill \mathcal{L}({\mathit{x}}_t,y_t;\mathit{\theta })={\mathcal{L}}_t\left(\mathit{\theta }\right)=-\frac{1}{2}{\mathit{\theta }}^{\top }\mathit{H}\mathit{\theta }+{\mathbf{z}}_t^{\top }\mathit{\theta },\end{array}$$

(8)

the result of a second-order Taylor expansion. The Hessian is assumed to be constant with respect to $({\mathit{x}}_t,y_t)$ (but not with respect to $\mathit{\theta }$).

To construct the dynamical equation for $\mathit{\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}^{*}$:

$$\begin{array}{c}\hfill {\theta }_{it}^{*}=-\frac{1}{H_{ii}}{\mathit{H}}_{-ii}^{\top }{\mathit{\theta }}_{-it},\end{array}$$

(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 ${\mathit{\theta }}_{-it}$. The dynamics of ${\mathit{\theta }}_{-it}$ are a complex stochastic process dependent on many different variables such as the dataset, model architecture, learning rate schedule, etc.

Assumption A2.

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

$$\begin{array}{c}\hfill p\left({\mathit{\theta }}_{-i,t+1}\right|{\mathit{\theta }}_{-i,t})=\mathcal{N}((1-\vartheta ){\mathit{\theta }}_{-it},{\eta }_{-i}^2),\end{array}$$

(10)

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

$$\begin{array}{c}\hfill p\left({\theta }_{i,t+1}^{*}\right|{\theta }_{i,t}^{*})=\mathcal{N}((1-\vartheta ){\theta }_{it}^{*},{\eta }^2),\end{array}$$

(11)

where ${\eta }^2={\eta }_{-i}^2{\mathit{H}}_{-ii}^{\top }{\mathit{H}}_{-ii}$.

In simple terms, in Assumption 2, we assume a parsimonious model of the dynamics, and that the next value of ${\mathit{\theta }}_{-i,t}$ is close to their previous value according to a Gaussian, similarly to Aitchison [31].

Lemma 1.

Under Assumptions 1 and 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:

$$\begin{array}{cc}\hfill {\mu }_{t,\mathit{post}}& ={\sigma }_{t,\mathit{post}}^2\left(\frac{{\mu }_{t,\mathit{prior}}}{{\sigma }_{t,\mathit{prior}}^2\left({\eta }^2\right)}+\frac{y_t}{{\sigma }^2}g\left({\mathit{x}}_t\right)\right)\mathit{and}\frac{1}{{\sigma }_{t,\mathit{post}}^2}=\frac{g{\left({\mathit{x}}_t\right)}^2}{{\sigma }^2}+\frac{1}{{\sigma }_{t,\mathit{prior}}^2\left({\eta }^2\right)},\hfill \end{array}$$

(12)

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

See Appendix E for the derivation of Lemma 1. From Equation (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 Assumption 1 and Assumption 2, if there is a data imbalance between tasks in Equation (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\left(y\right|\mathit{x},\mathcal{D})=\int p(y|\mathit{x},\mathcal{D},\mathit{\theta })p\left(\mathit{\theta }\right|\mathcal{D})d\mathit{\theta }$ for one’s application. Additionally, in this section, we have shown that if there is a task dataset size imbalance, then we can obtain forgetting under certain assumptions.

6. Related Work

There has been a recent resurgence in the field of CL [33] given the advent of deep learning. Methods that approximate sequential Bayesian inference Equation (3) have been seminal in CL’s revival and have used a diagonal Laplace approximation [3,7]. The diagonal Laplace approximation has been enhanced by modeling covariances between neural network weights in the same layer [8]. Instead of the Laplace approximation, we can use a variational approximation for sequential Bayesian inference, named VCL [9,34]. The variational Gaussian variance of each Bayesian NN parameter can be used to pre-condition the learning rates of each weight and create a mask per task by using pruning [10]. Using richer priors has also been explored [11,35,36,37,38]. For example, one can learn a scaling of the Gaussian NN weight parameters for each task by learning a new variational adaptation parameter which can strengthen the contribution of a specific neuron [39]. The online Laplace approximation can be seen as a special case of VCL where the KL-divergence term Equation (6) is tempered and the temperature tends to 0 [12]. Gaussian processes have also been applied to CL problems leveraging inducing points to retain previous task functions [40,41].

Bayesian methods that regularize weights have not matched up to the performance of experience replay-based CL methods [42] in terms of accuracy on CL image classification benchmarks. Instead of regularizing high-dimensional weight spaces, regularizing task functions is a more direct approach to combat forgetting [43]. Bayesian NN weights can also be generated by a hypernetwork, where the hypernetwork needs only simple CL techniques to prevent forgetting [44]. In particular, one can leverage the duality between the Laplace approximation and Gaussian processes to develop a functional regularization approach to Bayesian CL [45] or using function-space variational inference [46,47].

In the next section, we propose a simple Bayesian continual learning baseline that models the data-generating continual learning process and performs exact sequential Bayesian inference in a low-dimensional embedding space. Previous work has explored modeling the data-generating process by inferring the joint distribution of inputs and targets $p(\mathbf{x},y)$ and learning a generative model to replay data to prevent forgetting [48], and by learning a generative model per class and evaluating the likelihood of the inputs given each class $p\left(\mathbf{x}\right|y)$ [49].

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 [50,51] 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 [51], which creates embedding functions for different classes that are simply class means and predictions made by nearest neighbors. ProtoCL also bears similarities to the few-shot learning model Probabilistic Clustering for Online Classification [52], developed for few-shot image classification.

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

$$\begin{array}{c}\hfill y_{i,t}\sim \mathrm{Cat}\left(p_{1:J}\right),p_{1:J}\sim \mathrm{Dir}\left({\alpha }_t\right).\end{array}$$

(13)

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

$$\begin{array}{c}\hfill {\mathit{z}}_{it}|y_{it}\sim \mathcal{N}({\overline{\mathit{z}}}_{yt},{\Sigma }_{ϵ}),{\overline{\mathit{z}}}_{yt}\sim \mathcal{N}({\mathit{\mu }}_{yt},{\Lambda }_{yt}^{-1}).\end{array}$$

(14)

See Figure 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 ${\left\{p_j\right\}}_{j=1}^J$ and class embeddings ${\left\{{\overline{z}}_{y_j}\right\}}_{j=1}^J$ is denoted in short hand as $p\left(\mathit{\theta }\right)$ with parameters ${\eta }_t=\{{\alpha }_t,{\mathit{\mu }}_{1:J,t},{\Lambda }_{1:J,t}^{-1}\}$. ProtoCL models each class prototype but does not use task-specific NN parameters or modules such as multi-head VCL. By modeling a probabilistic model over an embedding space, this allows us to use powerful embedding functions $f(·;\mathit{w})$ without having to parameterize them probabilistically, and so this approach is 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(·;\mathit{w})$ and update the parameters ${\eta }_t$ at each iteration of gradient descent as well, see Algorithm 1.

Algorithm 1 ProtoCL continual learning

1: Input: task datasets ${\mathcal{T}}_{1:T}$, initialize embedding function: $f(·;\mathit{w})$, coreset: $\mathcal{M}=\varnothing$. 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(·;\mathit{w})$ by maximizing the posterior predictive $p(\mathbf{z},y)$ Equation (18) 5:      Obtain posterior over $\mathit{\theta }$ by updating $\eta$, Equations (15)–(17). 6:    end for 7:    Add random subset from ${\mathcal{T}}_i$ to $\mathcal{M}$. 8: end for

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

$$\begin{array}{c}\hfill {\alpha }_{t+1,j}={\alpha }_{t,j}+\sum _{i=1}^{N_t}\mathbb{I}(y_t^i=j),\end{array}$$

(15)

where $N_t$ are the number of points seen during the update at timestep t. Moreover, due to Gaussian-Gaussian conjugacy, the posterior for the Gaussian prototypes is governed by:

$$\begin{array}{cc}\hfill {\Lambda }_{y_{t+1}}& ={\Lambda }_{y_t}+N_y{\Sigma }_{ϵ}^{-1}\hfill \end{array}$$

(16)

$$\begin{array}{cc}\hfill {\Lambda }_{y_{t+1}}{\mathbf{\mu }}_{y_{t+1}}& =N_y{\Sigma }_{ϵ}^{-1}{\overline{\mathbf{z}}}_{y_t}+{\Lambda }_{y_t}{\mathbf{\mu }}_{y_t},\forall y_t\in C_t,\hfill \end{array}$$

(17)

where $N_y$ are the number of samples of class y and ${\overline{\mathit{z}}}_{y_t}=(1/N_y){\sum }_{i=1}^{N_y}{z}_{yi}$, see Appendix D.2 for the detailed derivation.

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

$$\begin{array}{cc}\hfill p(\mathbf{z},y)& =\int p(\mathbf{z},y|{\mathit{\theta }}_t;{\eta }_t)p({\mathit{\theta }}_t;{\eta }_t)d{\mathit{\theta }}_t=p\left(y\right)\prod _{i=1}^{N_t}\mathcal{N}\left({\mathbf{z}}_{it}\right|y_{it};{\mathit{\mu }}_{y_t,t},{\Sigma }_{ϵ}+{\Lambda }_{y_t,t}^{-1}).\hfill \end{array}$$

(18)

where the $p\left(y\right)={\alpha }_y/{\sum }_{j=1}^J{\alpha }_j$, see Appendix D.3 for the detailed derivation. This objective can then be optimized using gradient-based optimization for learning the prototype embedding function $\mathit{z}=f(\mathit{x};\mathit{w})$.

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

$$\begin{array}{cc}\hfill p(y^{*}=j|{\mathit{x}}^{*},{\mathit{x}}_{1:t},y_{1:t})& =p(y^{*}=j|{\mathit{z}}^{*},{\mathit{\theta }}_t)=\frac{p(y^{*}=j,{\mathbf{z}}^{*}|{\mathit{\theta }}_t)}{{\sum }_{i}p(y=i,{\mathbf{z}}^{*}|{\mathit{\theta }}_t)},\hfill \end{array}$$

(19)

see Appendix D.4 for further details.

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 a 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 being evaluated on with a task identifier $\tau$. Therefore, we cannot use the task identifier to select a specific head to use for classifying a test point. Moreover, we require the CL agent to identify each class, $\{0,\dots ,9\}$ for Split-MNIST and Split-CIFAR10, and not just $\{0,1\}$ as in domain-incremental learning. Class-incremental learning is more general, realistic, and harder a problem setting, and thus important to focus on rather than other settings, despite domain-incremental learning also being compatible with sequential Bayesian inference as described in Equation (3).

Implementation. For Split-MNIST and Split-FMNIST, the baselines and ProtoCL all use two-layer NNs with a 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. [22]. For ProtoCL and all baselines that 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 [47], which is state-of-the-art among Bayesian CL methods 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. The code to reproduce all experiments can be found here https://github.com/skezle/bayes_cl_posterior. 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. Mean accuracies across all tasks over CL vision benchmarks for class incremental learning [17]. 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 [9]	✗	$33.01\pm 0.08$	$32.77\pm 1.25$
$+$ coreset	✓	$52.98\pm 18.56$	$61.12\pm 16.96$
HIBNN ${}^{*}$ [11]	✗	$85.50\pm 3.20$	$43.70\pm 20.21$
FROMP [22]	✓	$84.40\pm 0.00$	$68.54\pm 0.00$
S-FSVI [47]	✓	$\mathbf{92}.\mathbf{94}\pm \mathbf{0}.\mathbf{17}$	$80.55\pm 0.41$
ProtoCL (ours)	✓	$\mathbf{93}.\mathbf{73}\pm \mathbf{1}.\mathbf{05}$	$\mathbf{82}.\mathbf{73}\pm \mathbf{1}.\mathbf{70}$

are lower for certain baselines than in the respective papers. We use 200 data points per task, see Figure A6 for a sensitivity analysis of the performance over the Split-MNIST benchmark as a function of core size for ProtoCL. In

Table 2. Mean accuracies across all tasks over CL vision benchmarks for class incremental learning [17]. All results are averages and standard errors over 10 seeds. Training times have been benchmarked using an Nvidia RTX3090 GPU.

Method	Training Time (s) $(\downarrow )$	Split CIFAR-10 (Acc) $(\uparrow )$
FROMP [22]	$1425\pm 28$	$48.92\pm 10.86$
S-FSVI [47]	$\mathrm{44,434}\pm 91$	$50.85\pm 3.87$
ProtoCL (ours)	$\mathbf{384}\pm \mathbf{6}$	$\mathbf{55}.\mathbf{81}\pm \mathbf{2}.\mathbf{10}$
		Split CIFAR-100 (Acc) $(\uparrow )$
S-FSVI [47]	$\mathrm{37,355}\pm 1135$	$20.04\pm 2.37$
ProtoCL (ours)	$\mathbf{1425}\pm \mathbf{28}$	$\mathbf{23}.\mathbf{96}\pm \mathbf{1}.\mathbf{34}$

we show how ProtoCL is able to scale to larger and more challenging CL vision benchmarks. ProtoCL demonstrates competitive performance versus the baselines we consider while at the same time requiring a fraction of the computational cost in terms of training times benchmarked on the same GPU.

The stated aim of ProtoCL is not to provide a novel state-of-the-art method for CL, but rather to propose a simple baseline that 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. ProtoCL scales to 10 tasks of Split-CIFAR100, which, to the best of our knowledge, is the highest number of tasks and classes that have been considered by previous Bayesian continual learning methods.

8. Discussion and Conclusions

In this paper, we revisited the use of sequential Bayesian inference for CL. We can use sequential Bayes to recursively build up the multi-task posterior Equation (3). 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 are 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 does not require complex variational optimization and achieves competitive results to the 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 on 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 [53], 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.