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: The authors introduce an extension to Bayesian continual learning whereby updates in successive tasks are constrained to be close to a “trust region” of the parameters found in preceding tasks. Exploiting the Bayesian nature of the model, the trust region in question can be defined in terms of the prior precision matrix. The authors show that the method, when applied to RNNs, is effective against baselines in a) a suite of neuroscience related tasks, and b) stroke MNIST datasets.