Posterior Meta-Replay for Continual Learning

1 Introduction . In recent years , a variety of continual learning ( CL ) algorithms have been developed to overcome the need to train neural networks with an independent and identically distributed ( i.i.d . ) sample . Most CL literature focuses on the particular scenario of continually learning a sequence of T tasks with datasets D ( 1 ) , . . . , D ( T ) . Because only access to the current task is granted , successful training of a discriminative model that captures p ( Y | X ) has to occur without an i.i.d . training sample from the overall joint D ( 1 : T ) i.i.d.∼ p ( X ) p ( Y | X ) . The advantages of a Bayesian approach for solving this problem are numerous and include the ability to drop all i.i.d . assumptions across and within tasks in a mathematically sound way , the ability to revisit tasks whenever new data becomes available , and access to principled uncertainty estimates capturing both data and parameter uncertainty . Up until now , Bayesian approaches to CL essentially focused on finding a combined posterior distribution via a recursive Bayesian update p ( W | D ( 1 : T ) ) ∝ p ( W | D ( 1 : T−1 ) ) p ( D ( T ) |W ) . Because the posterior of the previous task is used as prior for the next task , these approaches are also known as prior-focused [ 17 ] . In theory , the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . above recursive update can always recover the posterior p ( W | D ( 1 : T ) ) , independently of how the data is presented . However , because proper Bayesian inference is intractable , approximations are needed in practice , which lead to errors that are recursively amplified . As a result , whether solutions that are easily found in the i.i.d . setting can be obtained via the approximate recursive update strongly depends on factors such as task ordering , task similarity and the considered family of distributions . These factors limit the effectiveness of the recursive update and have a detrimental effect on the performance of prior-focused methods , especially in task-agnostic CL settings . To overcome these limitations , we propose an alternative Bayesian approach to CL that does not rely on the recursive update to learn distinct tasks and instead aims to learn task-specific posteriors ( Fig . 1 , refer to SM F.1 for a detailed discussion of the graphical model ) . In this view , finding trade-off solutions across tasks is not required , and knowledge transfer can be explicitly controlled for each task via the prior , which is no longer prescribed by the recursive update and can thus be set freely . By introducing probabilistic extensions of task-conditioned hypernetworks [ 91 ] , we show how task-specific posteriors can be learned with a single shared meta-model , an approach we term posterior meta-replay . This approach introduces two challenges : forgetting at the level of the hypernetwork , and the need to know task identity to correctly condition the hypernetwork . We empirically show that forgetting at the meta-level can be prevented by using a simple regularizer that replays parameters of previous posteriors . In task-agnostic inference settings , often referred to as class-incremental learning in the context of classification benchmarks [ 88 ] , the main hurdle therefore becomes task inference at test time . Here we focus on this task-agnostic setting , arguably the most challenging but also the most natural CL scenario , since the obtained models can be deployed just like those obtained via i.i.d . training ( e.g. , irrespective of the sequential training , the final model will be a classifier across all classes ) . In order to explicitly infer task identity from unseen inputs without resorting to generative models , we thoroughly study the use of principled uncertainty that naturally arises in Bayesian models . We show that results obtained in this task-agnostic setting with our approach constitute a leap in performance compared to prior-focused methods . Furthermore we show that limitations in task inference via predictive uncertainty are not related to our CL solution , but depend instead on the combination of approximate inference method , architecture , uncertainty measure and prior . Finally , we investigate how task inference can be further improved through several extensions . We summarize our main contributions below : • We describe a Bayesian CL framework where task-conditioned posterior parameter distributions are continually learned and compressed in a hypernetwork . • In a series of synthetic and real-world CL benchmarks we show that our task-conditioned hypernetworks exhibit essentially no forgetting , both for explicitly parameterized and implicit posterior distributions , despite using the parameter budget of a single model . • Compared to prior-focused methods , our approach leads to a leap in performance in task-agnostic inference while maintaining the theoretical benefits of a Bayesian approach . • Our approach scales to modern architectures such as ResNets , and remaining performance limitations are linked to uncertainty-based out-of-distribution detection but not to our CL solution . • Finally , we show how prominent existing Bayesian CL methods such as elastic weight consolidation can be dramatically improved in task-agnostic settings by introducing a small set of task-specific parameters and explicitly inferring the task . 2 Related Work . Continual learning . CL algorithms attempt to mitigate catastrophic interference while facilitating transfer of skills whenever possible . They can be coarsely categorized as ( 1 ) regularization-methods that put constraints on weight updates , ( 2 ) replay-methods that mimic pseudo-i.i.d . training by rehearsing stored or generated data and ( 3 ) dynamic architectures which can grow to allocate capacity for new knowledge [ 71 ] . Most related to our work is the study from von Oswald et al . [ 91 ] that introduces task-conditioned hypernetworks for CL , and already considers task inference via predictive uncertainty in the deterministic case . Our framework can be seen as a probabilistic extension of their work , which provides task-specific point estimates via a shared meta-model ( cf . Sec . 3 ) . Follow-up work also achieves task inference via predictive uncertainty , e.g. , Wortsman et al . [ 94 ] use it to select a learned binary mask per task that modulates a random base network . Here we complement these studies by thoroughly exploring task inference via several uncertainty measures , disclosing the factors that limit task inference and highlighting the importance of parameter uncertainty . A variety of methods tackling CL have been derived from a Bayesian perspective . A prominent example are prior-focused methods [ 17 ] , which incorporate knowledge from past data via the prior and , in contrast to our work , aim to find a shared posterior for all data . Examples include ( Online ) EWC [ 38 , 78 ] and VCL [ 65 , 54 ] . Other methods like CN-DPM [ 46 ] use Bayes ’ rule for task inference on the joint p ( X , C ) , where C is a discrete condition such as task identity . An evident downside of CN-DPM is the need for a separate generative and discriminative model per condition . More generally , such an approach requires meaningful density estimation in the input space , a requirement that is challenging for modern ML problems [ 64 ] . Other Bayesian CL approaches consider instead task-specific posterior parameter distributions . Lee et al . [ 47 ] learn separate task-specific Gaussian posterior approximations which are merged into a single posterior after all tasks have been seen . CBLN [ 49 ] also learns a separate Gaussian posterior approximation per task but later tries to merge similar posteriors in the induced Gaussian mixture model . Task inference is thus required and achieved via predictive uncertainty , although for a more reliable estimation all experiments consider batches of 200 samples that are assumed to belong to the same task . Tuor et al . [ 85 ] also learn a separate approximate posterior per task and use predictive uncertainty for task-boundary detection and task inference . In contrast to these approaches , we learn all task-specific posteriors via a single shared meta-model and remain agnostic to the approximate inference method being used . A conceptually related approach is MERLIN [ 33 ] , which learns task-specific weight distributions by training an ensemble of models per task that is used as training set for a task-conditioned variational autoencoder . Importantly , MERLIN requires a fine-tuning stage at inference , such that every drawn model is fine-tuned on stored coresets , i.e. , a small set of samples withheld throughout training . By contrast , our approach learns the parameters of an approximate Bayesian posterior p ( W | D ( t ) ) per task t , and no fine-tuning of drawn models is required . Bayesian neural networks . Because neural networks are expressive enough to fit almost any data [ 98 ] and are often deployed in an overparametrized regime , it is implausible to expect that any single solution obtained from limited data generalizes to the ground truth p ( Y | w , X ) ≈ p ( Y | X ) almost everywhere on p ( x ) . By contrast , Bayesian statistics considers a distribution over models , explicitly handling uncertainty to acknowledge data insufficiencies . This distribution is called the posterior parameter distribution p ( W | D ) ∝ p ( D |W ) p ( W ) , which weights models based on their ability to fit the data ( via the likelihood p ( D |W ) ) , while considering only plausible models according to the prior p ( W ) . Predictions are made by marginalizing over models ( for an introduction see MacKay [ 57 ] ) . Bayesian neural networks ( BNN ) apply this formalism to network parameters w , whereas for practical reasons hyperparameters like architecture are chosen deterministically [ 56 ] . While a deterministic discriminative model can only capture aleatoric uncertainty ( i.e. , uncertainty intrinsic to the data p ( Y | X ) ) , a Bayesian treatment allows to also capture epistemic uncertainty by being uncertain about the model ’ s parameters ( parameter uncertainty ) . This proper treatment of uncertainty is of utmost importance for safety-critical applications , where intelligent systems are expected to know what they don ’ t know . However , due to the complexity of modelling highdimensional distributions at the scale of modern deep learning , BNNs still face severe scalability issues [ 82 ] . Here , we employ several approximations to the posterior based on variational inference [ 5 ] from prior work , ranging from simple and scalable methods with a mean-field variational family like Bayes-by-Backprop ( BbB , [ 6 ] ) to methods with complex but rich variational families like the spectral Stein gradient estimator [ 79 ] . For more details see Sec . 3 and SM C . 3 Methods . In this section we describe our posterior meta-replay framework ( Fig . 2 ) . We start by introducing task-conditioned hypernetworks as a tool to continually learn parameters of task-specific posteriors , each of which is learned using variational inference ( SM C.1 ) . We then explain how the framework can be instantiated for both simple , explicit posterior approximations , and complex ones parametrized by an auxiliary network , and describe how forgetting can be mitigated through the use of a metaregularizer . We next explain how predictive uncertainty , naturally arising from a probabilistic view of learning , can be used to infer task identity for both PosteriorReplay methods , and PriorFocused methods that use a multihead output . Finally , we outline ways to boost task inference . Task-conditioned hypernetworks . Traditionally , hypernetworks are seen as neural networks that generate the weights w of a main network M processing inputs as ŷ = fM ( x , w ) [ 22 , 77 ] . Here , we consider instead hypernetworks that learn to generate θ , the parameters of a distribution qθ ( W ) over main network weights . By taking low-dimensional task embeddings e ( t ) as inputs and computing θ ( t ) = fTC ( e ( t ) , ψ ) , task-conditional ( TC ) computation is possible . Sampling is realized by transforming a base distribution p ( Z ) via a weight generator ( WG ) fWG ( z , θ ( t ) ) , whose choice determines the family of distributions considered for the approximation ( i.e. , the variational family ) . In our framework , weights w ∼ qθ ( t ) ( W ) are directly used for inference without requiring any fine-tuning . Importantly , all learnable parameters are comprised in the TC system , which can be designed to have less parameters than the main network , i.e. , dim ( ψ ) + ∑ t dim ( e ( t ) ) < dim ( w ) . Such constraint is vital to ensure fairness when comparing different CL methods , and is enforced in all our computer vision experiments . Additional details can be found in SM C.2 . Posterior-replay with explicit distributions . Different families of distributions can be realized within our framework . In the special case of a point estimate qθ ( t ) ( W ) = δ ( W − θ ( t ) ) , the WG system can be omitted altogether as it corresponds to the identity θ ( t ) = fWG ( z , θ ( t ) ) . This reduces our solution to the deterministic CL method introduced by von Oswald et al . [ 91 ] , which we refer to as PosteriorReplay-Dirac . However , capturing parameter uncertainty is a key ingredient of Bayesian statistics that is necessary for more robust task inference ( cf . Sec . 4.2 ) . We thus turn as a first step to explicit distributions qθ ( t ) ( W ) , for which the WG system samples according to a predefined function . We refer as PosteriorReplay-Exp to finding a mean-field Gaussian approximation via the BbB algorithm ( SM C.3.1 , [ 6 ] ) . In this case , θ ( t ) corresponds to the set of means and variances that define a Gaussian for each weight , which is directly generated by the TC . In the SM , we also report results for another instance of explicit distribution ( cf . SM C.3.2 ) . Posterior-replay with implicit distributions . Since the expressivity of explicit distributions is limited , we also explore the more diverse variational family of implicit distributions [ 19 , 31 ] . These are parametrized by a WG that now takes the form of an auxiliary neural network , making the parameters θ ( t ) of the approximate posterior dependent on the chosen WG architecture . This setting , referred to as PosteriorReplay-Imp , results in a hierarchy of three networks : a TC network generates task-specific parameters θ ( t ) for the approximate posterior , which is defined through an arbitrary base distribution p ( Z ) and the WG hypernetwork , which in turn generates weights w for a main network M that processes the actual inputs of the dataset D ( t ) . Interestingly , the TC now plays the role of a hyper-hypernetwork as it generates the weights of another hypernetwork ( Fig . 2a and Fig . 2c ) . Variational inference commonly resorts to optimizing an objective consisting of a data-dependent term and a prior-matching term KL ( qθ ( W ) || p ( W ) ) . Estimating the prior-matching term when using implicit distributions is not straightforward since we do not have analytic access to the density nor the entropy of qθ ( t ) ( W ) . To overcome this challenge , we resort to the spectral Stein gradient estimator ( SSGE , SM C.4.2 , [ 79 ] ) . This method is based on the insight that direct access to the log-density is not required , but only to its gradient with respect to W. Noticing that this quantity appears in Stein ’ s identity , the authors consider a spectral decomposition of the term and use the Nyström method to approximate the eigenfunctions . We test an alternative method for dealing with implicit distributions in the SM that is based on estimating the log-density ratio ( SM C.4.1 ) . As an additional challenge introduced by the use of implicit distributions , the support of qθ ( t ) ( W ) is limited to a low-dimensional manifold when using an inflating architecture for WG , causing the priormatching term to be ill-defined . To overcome this , we investigate the use of small noise perturbations in WG outputs ( SM C.4.3 ) . Normalizing flows [ 70 ] can also be utilized as WG architectures to gain analytic access to qθ ( t ) ( W ) , albeit at the cost of architectural constraints such as invertibility . Overcoming forgetting via meta-replay . Since all learnable parameters are part of the TC system , forgetting only needs to be addressed at this meta-level . With Ltask the task-specific loss ( SM Eq . 3 ) and D ( ·||· ) a divergence measure between distributions , the loss for task t becomes : L ( t ) ( ψ , E , D ( t ) ) = Ltask ( ψ , e ( t ) , D ( t ) ) + β ∑ t′ < t D ( qθ ( t′ , ∗ ) ( W ) ||qθ ( t′ ) ( W ) ) ( 1 ) where E = { e ( t′ ) } tt′=1 is the set of task embeddings up to the current task , β is the strength of the regularizer and θ ( t ′ , ∗ ) ≡ fTC ( e ( t ′ , ∗ ) , ψ∗ ) are the parameters of the posterior approximation obtained from a checkpoint of the TC parameters before learning task t : { ψ∗ } ∪ { e ( t′ , ∗ ) } t′ < t . The checkpointed meta-model allows replaying posterior distributions from the past and retaining them via divergence minimization as detailed below , hence the name posterior meta-replay . Importantly , the loss on task t only depends on the corresponding dataset D ( t ) , but knowledge transfer across tasks is possible because task-specific models are learned through a shared meta-model.1 Since the computation required to compute this regularizer linearly scales with the number of tasks , we also explore stochastically regularizing on a subset of randomly selected tasks in each update ( cf . SM D.3 ) , and show that this does not impair performance . Notably , our PosteriorReplay method does not incur in a significant increase of runtime or memory usage ( SM F.2 ) . The evaluation of Eq . 1 requires estimates of a divergence measure to prevent changes in learned posterior approximations of past tasks . Because these are required at every loss evaluation and need to be cheap to compute , we do not consider sample-based estimates but only estimates that directly utilize posterior parameters . This goal is easy to achieve for families of posterior approximations that possess an analytic expression for a divergence measure ( e.g. , Gaussian distributions ) . More specifically , for PosteriorReplay-Exp , we consider the forward KL , backward KL and the 2-Wasserstein distance but did not observe that the specific choice of divergence measure is crucial in practice ( cf . SM C.3.1 and Table S14 ) . In all other cases , approximations are required . Specifically , we resort in our experiments to the use of an L2 regularizer at the output of the TC network : β ∑ t′ < t ‖fTC ( e ( t ′ , ∗ ) , ψ∗ ) − fTC ( e ( t ′ ) , ψ ) ‖22 ( 2 ) Perhaps surprisingly , we observe that this crude regularization , reminiscent to the one used in von Oswald et al . [ 91 ] for point estimates , does not harm performance and leads to models that exhibit 1While this type of transfer is rather implicit , explicit knowledge transfer can be realized via task-specific priors ( cf . SM F.7 ) . virtually no forgetting . However , we discuss in SM F.3 how this isotropic regularization could be improved given that the KL is locally approximated by a norm ‖·‖F induced by the Fisher information matrix F on qθ ( W ) . Task inference . A system with task-specific solutions requires access to task identity when processing unseen samples . In our framework , this amounts to selecting the correct task embedding to condition the TC . Although auxiliary systems can be used to infer task identity [ 24 , 91 ] , here we exploit predictive uncertainty , assuming task identity can be inferred from the input alone . For a task inference approach based on predictive uncertainty to work , the properties of the input data distribution p ( t ) ( x ) need to be reflected in the uncertainty measure , i.e. , uncertainty needs to be low for in-distribution data and high for out-of-distribution ( OOD ) data [ 11 ] . Uncertainty-based task inference with deterministic discriminative models only captures aleatoric uncertainty , making its overall validity debatable . Indeed , aleatoric uncertainty is only calibrated in-distribution and its OOD behavior is hard to foresee . Instead , we argue that epistemic uncertainty arising naturally in a Bayesian setting is crucial for robust uncertainty-based task inference . In our case , epistemic uncertainty stems from the fact that the posterior parameter distribution p ( W|D ( t ) ) in conjunction with the network architecture induces a distribution over functions . If this distribution captures a rich set of functions , a diversity of predictions on OOD data can be expected even if those functions agree in-distribution ( cf . Sec . 4.2 ) . Note , however , that inducing such diverse distribution over functions is not straightforward with neural networks , and that more research is required to justify uncertainty-based OOD detection [ 11 ] . We explore two different ways to quantify uncertainty for task inference . In Ent the task t leading to the lowest entropy on the predictive posterior p ( y | D ( t ) ; x̃ ) is selected , where p ( y | D ( t ) ; x̃ ) =∫ W p ( y | W ; x̃ ) p ( W|D ( t ) ) dW is approximated via Monte-Carlo with samples from qθ ( t ) ( W ) . This approach captures both aleatoric and epistemic uncertainty when used in a probabilistic setting . In Agree the task leading to the highest agreement in predictions across models drawn from qθ ( t ) ( W ) is selected . This approach exclusively measures epistemic uncertainty and can therefore only be estimated in a probabilistic setting . Although we generally consider task inference for individual samples , we also explore batch-wise ( BW ) task inference for batches of 100 samples that are assumed to belong to the same task . Such approach drastically boosts task inference simply due to a statistical accumulation effect when having above chance level task inference for single inputs . Intuitively , BW corresponds to accumulating evidence to decrease uncertainty ( e.g. , an agent looking at an object from multiple perspectives ) . Further details can be found in SM C.6 , and using uncertainty for task-boundary detection when training without explicit access to task identity is explored in SM D.8 . Facilitating task inference through coresets . A key advantage of Bayesian statistics is the ability to update models as new evidence arrives . When continually learning a sequence of tasks , posteriors may for example undergo a post-hoc fine-tuning on stored coresets to mitigate catastrophic forgetting of earlier tasks . Specifically , given a dataset split D ( t ) \ C ( t ) ∪ C ( t ) , one can perform a final update p ( W|D ( t ) ) ∝ p ( W | D ( t ) \ C ( t ) ) p ( C ( t ) | W ) using a stored coreset C ( t ) in conjunction with an already learned posterior approximation for p ( W | D ( t ) \ C ( t ) ) that now acts as a prior . Interestingly , access to coresets after training on all tasks can also be exploited to facilitate task inference via predictive uncertainty . Here , we explore this idea by encouraging task-specific models to produce uncertain predictions for OOD samples ( i.e. , coresets from other tasks ) , an approach that we denote CS ( SM C.7 ) , and in which we store 100 inputs per task . Intriguingly , training on OOD inputs makes these become in-distribution , and therefore renders model agreement ( Agree ) inapplicable for task inference , which we empirically observe . Improving prior-focused CL . We investigate ways to improve PriorFocused methods within our framework . First , we endow them with implicit posterior approximations qθ ( W ) parametrized by a WG hypernetwork , an approach we refer to as PriorFocused-Imp . Because the posterior is shared across tasks , no TC system is required and the parameters θ can be directly optimized via SSGE . Second , we enrich PriorFocused methods with a small set of task-specific parameters that enable uncertainty-based task inference for prior-focused methods too . Specifically , the learned parameters w consist of a set of shared weights φ and a set of task-specific output heads with weights { ξ ( t ) } Tt=1 . This approach is in contrast with how PriorFocused methods like Online EWC are commonly deployed in task-agnostic inference settings , where the softmax output grows as new tasks arrive ( e.g. , [ 88 ] ) . The use of a growing softmax causes the model class parametrized by w to change over time , and therefore violates the Bayesian assumption that the approximate posterior is obtained from a model class containing the ground-truth model . We show that this leads to limitations that can be overcome by a multihead approach . For Online EWC , we refer to the growing softmax and multihead scenarios as EWC-growing and EWC-multihead , respectively ( cf . SM C.5.2 ) . We also explore the prior-focused instantiation of BbB , known as VCL ( cf . SM C.5.1 ) .

Summary: The authors propose a new method for (offline) continual learning using bayesian hypernetworks to obtain a distribution over a task-based parameter generating function. This meta perspective enables replay at the level of level of the hypernet parameters (rather than at the input level), while the bayesian approach provides reliable uncertainty estimates which the authors show can be used to perform task inference at test time.   The hypernet setup, (which is really a 3-level hierarchy) comprises a 1) Task-Conditioned network, which maps a task embedding to hypernet parameters $\theta_t$, 2) a hypernet with params $\theta_t$ remaps a noise sample $z$ to a set of weights $w$, 3) which are used as parameters for the downstream classifier.   The authors provide two toy setups to illustrate the validity of their method, before benchmarking their method on Split-Cifar and Split Mnist.