Natural continual learning: success is a journey, not (just) a destination

1 Introduction . Catastrophic forgetting is a common feature of machine learning algorithms where training on a new task often leads to poor performance on previously learned tasks . This is in contrast to biological agents which are capable of learning many different behaviors over the course of their lives with little to no interference across tasks . The study of continual learning in biological networks may therefore help inspire novel approaches in machine learning , while the development and study of continual learning algorithms in artificial agents can help us better understand how this challenge is overcome in the biological domain . This is particularly true for more challenging continual learning settings where task identity is not provided at test time , and for continual learning in recurrent neural 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . networks ( RNNs ) , which is important due to the practical and biological relevance of RNNs . However , continual learning in these settings has recently proven challenging for many existing algorithms , particularly those that rely on parameter regularization to mitigate forgetting [ 12 , 13 , 47 ] . In this work , we address these shortcomings by developing a continual learning algorithm that not only encourages good performance across tasks at convergence but also regularizes the optimization path itself using trust region optimization . This leads to improved performance compared to existing methods . Previous work has addressed the challenge of continual learning in artificial agents using weight regularization , where parameters important for previous tasks are regularized to stay close to their previous values [ 1 , 17 , 24 , 32 , 37 , 54 ] . This approach can be motivated by findings in the neuroscience literature of increased stability for a subset of synapses after learning [ 49 , 50 ] . More recently , approaches based on projecting gradients into subspaces orthogonal to those that are important for previous tasks have been developed in both feedforward [ 40 , 53 ] and recurrent [ 12 ] neural networks . This is consistent with experimental findings that neural dynamics often occupy orthogonal subspaces across contexts in biological circuits [ 3 , 14 , 20 , 22 ] . While these methods have been found to perform well in many continual learning settings , they also suffer from several shortcomings . In particular , while Bayesian weight regularization provides a natural way to weigh previous and current task information , this approach can fail in practice due to its approximate nature and often requires additional tuning of the importance of the prior beyond what would be expected in a rigorous Bayesian treatment [ 46 ] . In contrast , while projection-based methods have been found empirically to mitigate catastrophic forgetting , it is unclear how the ‘ important subspaces ’ should be selected and how such methods behave when task demands begin to saturate the network capacity . In this work , we develop natural continual learning ( NCL ) , a new method that combines ( i ) Bayesian continual learning using weight regularization with ( ii ) an optimization procedure that relies on a trust region constructed from an approximate posterior distribution over the parameters given previous tasks . This encourages parameter updates predominantly in the null-space of previously acquired tasks while maintaining convergence to a maximum of the Bayesian approximate posterior . We show that NCL outperforms previous continual learning algorithms in both feedforward and recurrent networks . We also show that the projection-based methods introduced by Duncker et al . [ 12 ] and Zeng et al . [ 53 ] can be viewed as approximations to such trust region optimization using the posterior from previous tasks . Finally , we use tools from the neuroscience literature to investigate how the learned networks overcome the challenge of continual learning . Here , we find that the networks learn latent task representations that are stable over time after initial task learning , consistent with results from biological circuits . 2 Method . Notations We use X > , X−1 , Tr ( X ) and vec ( X ) to denote the transpose , inverse , trace , and column-wise vectorization of a matrix X . We use X ⊗ Y to represent the Kronecker product between matrices X ∈ Rn×n and Y ∈ Rm×m such that ( X ⊗ Y ) mi+k , mj+l = XijYkl . We use bold lower-case letters x to denote column vectors . Dk refers to a ‘ dataset ’ corresponding to task k , which in this work generally consists of a set of input-output pairs { x ( i ) k , y ( i ) k } such that ` k ( θ ) : = log p ( Dk|θ ) = ∑ i log pθ ( y ( i ) k |x ( i ) k ) is the task-related performance on task k for a model with parameters θ . Finally , we use D̂k to refer to a dataset generated by inputs from the kth task where { ŷ ( i ) k ∼ pθ ( y|x ( i ) k ) } are drawn from the model distributionM . 2.1 Bayesian continual learning . Problem statement In continual learning , we train a model on a set of K tasks { D1 , . . . , DK } that arrive sequentially , where the data distribution Dk for task k in general differs from D6=k . The aim is to learn a probabilistic model p ( D|θ ) that performs well on all tasks . The challenge in the continual learning setting stems from the sequential nature of learning , and in particular from the common assumption that the learner does not have access to “ past ” tasks ( i.e. , Dj for j < k ) when learning task k. While we enforce this stringent condition in this paper , our approach may be easily combined with memory-based techniques such as coresets or generative replay [ 8 , 10 , 13 , 32 , 34 , 36 , 38 , 42 , 43 , 45 , 48 ] . Bayesian approach The continual learning problem is naturally formalized in a Bayesian framework whereby the posterior after k−1 tasks is used as a prior for task k. More specifically , we choose a prior p ( θ ) on the model parameters and compute the posterior after observing k tasks according to Bayes ’ rule : p ( θ|D1 : k ) ∝ p ( θ ) k∏ k′=1 p ( Dk′ |θ ) ∝ p ( θ|D1 : k−1 ) p ( Dk|θ ) , ( 1 ) where D1 : k is a concatenation of the first k tasks ( D1 , . . . , Dk ) . In theory , it is thus possible to compute the exact posterior p ( θ|D1 : k ) after k tasks , while only observing Dk , by using the posterior p ( θ|D1 : k−1 ) after k − 1 tasks as a prior . However , as is often the case in Bayesian inference , the difficulty here is that the posterior is typically intractable . To address this challenge , it is common to perform approximate online Bayesian inference . That is , the posterior p ( θ|D1 : k−1 ) is approximated by a parametric distribution with parameters φk−1 . The approximate posterior q ( θ ; φk−1 ) is then used as a prior for task k. Online Laplace approximation A common approach is to use the Laplace approximation whereby the posterior p ( θ|D1 : k−1 ) is approximated as a multivariate Gaussian q using local gradient information [ 17 , 24 , 37 ] . This involves ( i ) finding a mode µk of the posterior during task k , and ( ii ) performing a second-order Taylor expansion around µk to construct an approximate Gaussian posterior q ( θ ; φk ) = N ( θ ; µk , Λ−1k ) , where Λk is the precision matrix and φk = ( µk , Λk ) . In this case , gradient-based optimization is used to find the posterior mode on task k ( c.f . Equation 1 ) : µk = arg max θ log p ( θ|Dk , φk−1 ) ( 2 ) = arg max θ log p ( Dk|θ ) + log q ( θ ; φk−1 ) ( 3 ) = arg max θ ` k ( θ ) − 1 2 ( θ − µk−1 ) > Λk−1 ( θ − µk−1 ) ︸ ︷︷ ︸ : = Lk ( θ ) ( 4 ) The precision matrix Λk is given by the Hessian of the negative log posterior at µk : Λk = − ∇2θ log p ( θ|Dk , φk−1 ) ∣∣ θ=µk = H ( Dk , µk ) + Λk−1 , ( 5 ) where H ( Dk , µk ) = − ∇2θ log p ( Dk|θ ) ∣∣ θ=µk is the Hessian of the negative log likelihood of Dk . Continual learning with the online Laplace approximation thus involves two steps for each new task Dk . First , given Dk and the previous posterior q ( θ ; µk−1 , Λ−1k−1 ) ( i.e . the new prior ) , µk is found using gradient-based optimization ( Equation 4 ) . This step can be interpreted as optimizing the likelihood of Dk while penalizing changes in the parameters θ according to their importance for previous tasks , as determined by the prior precision matrix Λk−1 . Second , the new posterior precision matrix Λk is computed according to Equation 5 . Approximating the Hessian In practice , computing Λk presents two major difficulties . First , because q ( θ ; φk ) is a Gaussian distribution , Λk has to be positive semi-definite ( PSD ) , which is not guaranteed for the Hessian H ( Dk , µk ) . Second , if the number of model parameters nθ is large , it may be prohibitive to compute a full ( nθ × nθ ) matrix . To address the first issue , it is common to approximate the Hessian with the Fisher information matrix ( FIM ; 17 , 30 , 37 ) : Fk = Ep ( D̂k|θ ) [ ∇θ log p ( D̂k|θ ) ∇θ log p ( D̂k|θ ) > ] ∣∣∣ θ=µk ≈ H ( Dk , µk ) ( 6 ) The FIM is PSD , which ensures that Λk = ∑k k′=1 Fk′ is also PSD . Computing Fk may still be impractical if there are many model parameters , and it is therefore common to further approximate the FIM using structured approximations with fewer parameters . In particular , a diagonal approximation to Fk recovers Elastic Weight Consolidation ( EWC ; 24 ) , while a Kronecker-factored approximation [ 31 ] recovers the method proposed by Ritter et al . [ 37 ] . We denote this method ‘ KFAC ’ and use it in Section 3 as a comparison for our own Kronecker-factored method . 2.2 Natural continual learning . While the online Laplace approximation has been applied successfully in several continual learning settings [ 24 , 37 ] , it has also been found to perform sub-optimally on a range of problems [ 12 , 46 ] . Additionally , its Bayesian interpretation in theory prescribes a unique way of weighting the contributions of previous and current tasks to the loss . However , to perform well in practice , weight regularization approaches have been found to require ad-hoc re-weighting of the prior term by several orders of magnitude [ 24 , 37 , 46 ] . These shortcomings could be due to an inadequacy of the approximations used to construct the posterior ( Section 2.1 ) . However , we show in Figure 1 that standard gradient descent on the Laplace posterior has important drawbacks even in the exact case . First , we show that exact Bayesian inference on a simple continual regression problem can produce indirect optimization paths along which previous tasks are transiently forgotten as a new task is being learned ( Figure 1A ; green ) . Second , when the loss is non-convex , we show that exact Bayesian inference can still lead to catastrophic forgetting ( Figure 1B ; green ) . An alternative approach that has found recent success in a continual learning setting involves projection based methods which restrict parameter updates to a subspace that does not interfere with previous tasks [ 12 , 53 ] . However , it is not immediately obvious how this projected subspace should be selected in a way that appropriately balances learning on previous and current tasks . Additionally , such projection-based algorithms have fixed points that are minima of the current task , but not necessarily minima of the ( negative ) Bayesian posterior . This can lead to catastrophic forgetting in the limit of long training times ( Figure 1 ; pink ) , unless the learning rate is exactly zero in directions that interfere with previous tasks . To combine the desirable features of both classes of methods , we introduce “ Natural Continual Learning ” ( NCL ) – an extension of the online Laplace approximation that also restricts parameter updates to directions which do not interfere strongly with previous tasks . In a Bayesian setting , we can conveniently express what is meant by such directions in terms of the prior precision matrix Λ . In particular , ‘ flat ’ directions of the prior ( low precision ) correspond to directions that will not significantly affect the performance on previous tasks . Formally , we derive NCL as the solution of a trust region optimization problem . This involves minimizing the posterior loss Lk ( θ ) within a region of radius r centered around θ with a distance metric of the form d ( θ , θ + δ ) = √ δ > Λk−1δ/2 that takes into account the curvature of the prior via its precision matrix Λk−1 : δ = arg min δ Lk ( θ ) +∇θLk ( θ ) > δ subject to 1 2 δ > Λk−1δ ≤ r2 , ( 7 ) where Lk ( θ + δ ) ≈ Lk ( θ ) + ∇θLk ( θ ) > δ is a first-order approximation to the updated Laplace objective . The solution to this subproblem is given by δ ∝ Λ−1k−1∇θ ` k ( θ ) − ( θ − µk−1 ) ( see Appendix A for a derivation ) , which gives rise to the NCL update rule θ ← θ + γ [ Λ−1k−1∇θ ` k ( θ ) − ( θ − µk−1 ) ] ( 8 ) for a learning rate parameter γ ( which is implicitly a function of r in Equation 7 ) . To get some intuition for this learning rule , we note that Λ−1k−1 acts as a preconditioner for the first ( likelihood ) term , which drives learning on the current task while encouraging parameter changes predominantly in directions that do not interfere with previous tasks . Meanwhile , the second term encourages θ to stay close to µk−1 , the optimal parameters for the previous task . As we illustrate in Figure 1 , this combines the desirable features of both Bayesian weight regularization and projection-based methods . In particular , NCL shares the fixed points of the Bayesian posterior while also mitigating intermediate or complete forgetting of previous tasks by preconditioning with the prior covariance . Notably , if the loss landscape is non-convex ( as it generally will be ) , NCL can converge to a different local optimum from standard weight regularization despite having the same fixed points ( Figure 1B ) . Implementation The general NCL framework can be applied with different approximations to the Fisher matrix Fk in Equation 6 ( see Section 2.1 ) . In this work , we use a Kronecker-factored approximation [ 31 , 37 ] . However , even after making a Kronecker-factored approximation to Fk for each task k , it remains difficult to compute the inverse of a sum of k Kronecker products ( c.f . Equation 5 ) . To address this challenge , we derived an efficient algorithm for making a Kroneckerfactored approximation to Λk = Fk + Λk−1 ≈ Ak ⊗Gk when Λk−1 = Ak−1 ⊗Gk−1 and Fk are also Kronecker products . This approximation minimizes the KL-divergence between N ( µk , ( Ak ⊗ Gk ) −1 ) and N ( µk , ( Λk−1 + Fk ) −1 ) ( see Appendix G for details ) . Before training on the first task , we assume a spherical Gaussian prior θ ∼ N ( 0 , p−2w I ) . The scale parameter pw can either be set to a fixed value ( e.g . 1 ) or treated as a hyperparameter , and we optimize pw explicitly for our experiments in feedforward networks . NCL also has a parameter α which is used to stabilize the matrix inversion Λ−1k−1 ≈ ( Ak−1 ⊗Gk−1 + α2I ) −1 ( Appendix E ) . This is equivalent to a hyperparameter used for such matrix inversions in OWM [ 53 ] and DOWM [ 12 ] , and it is important for good performance with these methods . The pw and α are largely redundant for NCL , and we generally prefer to fix α to a small value ( 10−10 ) and optimize the pw only . However , for our experiments in RNNs , we instead fix pw = 1 and perform a hyperparameter optimization over α for a more direct comparison with OWM and DOWM . The NCL algorithm is described in pseudocode in Appendix E together with additional implementation and computational details . Finally , while we have derived NCL with a Laplace approximation in this section for simplicity , it can similarly be applied in the variational continual learning framework of Nguyen et al . [ 32 ] ( Appendix J ) . Our code is available online1 .

Summary: This paper proposes a model for continual learning. Specifically, it intends to merge parameter regularization methods and gradient projection methods. The regularization is achieved by taking a bayesian perspective over the distribution over parameters, that is approximated by a Laplace distribution in a variational fashion, with online updates. The gradient projection is addressed by trust-region optimization, resulting in an update rule that uses the precision matrix to precondition the gradient of the current task objective, plus a term that keeps the parameters close to the previous mean. Experiments are carried out for sequential problems on a stimulus-response dataset and on stroke MNIST.