Flattening Sharpness for Dynamic Gradient Projection Memory Benefits Continual Learning

1 Introduction . Humans have the ability to continually learn new knowledge without forgetting their previously learned ones through mediating a rich set of neurocognitive mechanisms [ 41 , 15 , 39 ] . This ability , often known as continual learning or lifelong learning [ 29 ] , is crucial for computational systems , such as deep neural networks ( DNNs ) , which are required to sequentially learn and deal with multiple tasks when implemented in the dynamically changing environment . Continual learning remains a long-standing challenge for DNNs since these networks are typically trained with stationary training batches by stochastic gradient descent methods [ 19 ] , which generally leads to an abrupt performance decrease on previously learned tasks as new tasks are learned . To address such catastrophic forgetting , we can brutally retrain an oracle network on the entire dataset containing all tasks to capture dynamic changes in the data distribution , but this methodology is obviously too inefficient to hinder the learning of novel data in real time . During the last few years , lots of research efforts have been devoted to improving the stability of DNNs on old tasks while keeping sensitive to new information . The first intuitive idea is to introduce ∗Corresponding author : Guangyong Chen . < gy.chen @ siat.ac.cn > 2The code is available at : https : //github.com/danruod/FS-DGPM 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . an independent branch for each new task while freezing the old task parameters to preserve the old knowledge [ 33 , 45 , 43 , 25 , 35 , 22 ] . However , in this way , the network will inevitably become redundant as the task number continually increases . As presented in the neurocognitive works [ 41 , 15 ] , the reactivation of neuronal activity patterns , representing old memories , plays an important role in the continual learning of humans [ 39 ] . Thus , forgetness can be effectively mitigated by training a single network for new tasks by considering diverse information stored in the memory , including the original training samples of old tasks [ 31 , 5 , 13 ] , the gradients induced from old tasks [ 9 ] and the feature subspace representing old tasks [ 34 ] . However , their continual learning performance is still limited because DNNs can easily overfit the limited information stored in the small-size memory . The overfitting problem of DNNs is often attributed to the complex loss landscape containing multiple local optima , and the sharpness of the loss landscape has been widely used to characterize the generalization gap in standard training scenarios from both theoretical and empirical perspectives [ 27 , 21 , 42 , 10 , 6 ] . While this characterization has inspired new approaches for model training with better generalization , practical algorithms that especially seek out flatter minima to effectively address the ’ sensitivity-stability ’ dilemma for continual learning have thus far been elusive . In this paper , our first contribution is to characterize the weight loss landscape for the continual learning scenario and identify that a flatter loss landscape with lower loss value often leads to better continual learning performance , as shown in Figure 1 and Figure 3 . Further , based on our characterization of the weight loss landscape , we find that the recently proposed Gradient Projection Memory ( GPM ) method [ 34 ] maintains the lowest loss value on old tasks among the previously proposed methods by taking gradient steps orthogonal to the subspace representing old tasks . However , its loss landscape on newly learned tasks is the sharpest due to the lack of sufficient subspace left for new task learning . To improve the network ’ s sensitivity , our second contribution is to predict the importance of bases spanning the subspace for old tasks , so that less important bases can be dynamically released . In particular , we introduce a soft weight to indicate the bases importance , which can be dynamically adjusted by combining the Flattening Sharpness ( FS ) to minimize the loss value and loss sharpness simultaneously . Intuitively , a basis will be regarded as important for preserving old knowledge if the gradients induced by new tasks and old ones are aligned in the opposite direction on that basis . As demonstrated through extensive experiments , our proposed method can consistently outperform the state-of-the-art methods [ 17 , 30 , 24 , 31 , 5 , 13 , 34 ] by a notable margin across a range of widely used benchmark datasets . 2 Related Work . In this section , we briefly survey the representative works of continual learning by highlighting their contributions . To simplify our presentation , this section is organized by dividing the representative works into three categories , parameter isolation , regularization-based , memory-based methods . Parameter isolation methods address forgetting by assigning a different subset of network parameters to each task . Without restrictions on network architecture , new neurons or layers or modules can be added for new tasks , while the previous task parameters can be frozen or copied to preserve old knowledge . For instance , Progressive Neural Network ( PGN ) [ 33 ] freezes the parameters trained with previous knowledge while expands the architecture by allocating new sub-networks with fixed capacity for new tasks . Dynamically Expandable Networks ( DEN ) [ 45 ] selectively retrains or expands network capacity by splitting/duplicating important units on new tasks . Reinforced Continual Learning ( RCL ) [ 43 ] uses reinforcement learning strategy to adaptively expand the network of each layer , while [ 22 ] use neural architecture search to find optimal network structures for each sequential task . Alternatively , with the architecture remaining static , a fixed part is allocated to each task . During the training of a new task , previous task parts are masked out to prevent interference . The mask sets are imposed at parameter level [ 25 ] , or unit level [ 35 ] . PackNet [ 25 ] uses iterative pruning to fully restrict gradient updates on important weights via a binary mask , whereas HAT [ 35 ] limits the update of important units recognized by the hard attention mask through gradient backpropagation . Regularization-based methods introduce an additional regularization term in the loss function to consolidate previous knowledge without using replay . This involves using knowledge distillation [ 23 , 14 ] or penalizing changes in weights deemed important for previous tasks [ 17 , 46 , 28 , 2 , 3 ] to reduce forgetting . There are many ways to measure the importance . Elastic Weight Consolidation ( EWC ) [ 17 ] identifies important weights based on the diagonal values of Fisher information matrix after training , while Synaptic Intelligence ( SI ) [ 46 ] calculates them online and over the entire learning trajectory in parameter space . Memory Aware Synapses ( MAS ) [ 2 ] estimates importance weights in an unsupervised manner , while Variational Continual Learning ( VCL ) [ 28 ] introduces a variational framework that spawned some Bayesian-based works [ 32 , 8 , 1 , 7 ] . For example , [ 32 ] recursively uses a Gaussian Laplace approximation of the Hessian to approximate the posterior after every task , [ 8 ] adjusts the learning rate according to the uncertainty defined by the probability distribution of the network weights . [ 1 ] introduces an interpretation of node-wise uncertainty on the Kullback-Leibler ( KL ) divergence term of the variational lower bound for Gaussian mean-field approximation . Our method mainly follows memory-based methods , which mitigate forgetting based on information extracted from old tasks or based on a generative model to generate pseudo samples . For example , iCaRL [ 30 ] selects and stores samples closest to the feature mean of each class . ER [ 5 , 31 ] suggests reservoir sampling under the limited and fixed budget for replay buffer . Deep Generative Replay ( DGR ) [ 36 ] trains a deep generative model in the Generative Adversarial Network ( GAN ) framework [ 11 ] to simulate past data . These previous task samples are mainly reused as model inputs for replay in the above methods . However , replay might be prone to overfitting the subset of stored samples . Alternatively , samples stored in memory can also be used to constrain the optimization of the new task loss to prevent previous task interference , thereby leaving more leeway for backward and forward transfer . Gradient Episodic Memory ( GEM ) [ 24 ] projects the estimated gradients in the feasible region , which is outlined by previous task gradients calculated from the episodic memory samples . Averaged-GEM ( A-GEM ) [ 4 ] relaxes the projection to a direction that is estimated from samples randomly selected from memory . [ 12 ] proposes a unified view of episodic memory-based continual learning methods , including GEM and A-GEM , and improves performance over these methods by using a loss-balancing update scheme . A few other works have utilized gradient information to protect previous knowledge . [ 31 , 13 ] adopt optimization-based meta-learning to enforce gradient alignment between samples from different tasks . GPM [ 34 ] minimizes interference between sequential tasks by ensuring that gradient updates only occur in directions orthogonal to the input of previous tasks . 3 The Weight Loss Landscape of Continual Learning . In this section , we first introduce our formulation of continual learning , and then characterize the weight loss landscape for the continual learning scenario from stability and sensitivity . Finally , some insights combining the weight loss landscape and continual learning are provided . 3.1 Problem Formulation . Throughout the paper , we denote scalars as a , vectors as a , matrices as A , and sets asA . We consider a supervised learning setup where T tasks are sequentially learned from their training data . Each task has an identical task descriptor , τ ∈ { 1 , 2 , . . . , T } , with its dataset Dτ = { xi , τ , yi , τ } nτi=1 containing nτ samples randomly generated from a latent distribution Dτ . At any time-step during the learning process , we minimize the empirical risk of the model on all t tasks seen so far , with just limited size of memoryM to summarize the training data of previous tasks { Dτ } t−1τ=1 . To simplify the notation , we denote LA ( w ) = 1|A| ∑ ( x , y ) ∈A [ ` ( fw ( x ) , y ) ] as the average empirical loss for the set A , where ` ( · , · ) is an arbitrary loss function ( e.g . the cross-entropy ( CE ) loss ) , |A| is the sample size of the set A , and fw is the DNN with weight w. Our final goal is to find an optimal parameter w , which minimizes the overall risk ∑T τ=1 LDτ ( w ) for all tasks . 3.2 Connection of Weight Loss Landscape and Continual Learning . After learning a new task , we visualize the weight loss landscape of each task seen so far by plotting changes in its training loss when moving the weights w in a random direction d with magnitude α following [ 21 ] : gt ( α ) = LDt ( w + αd ) = 1 |Dt| ∑ ( x , y ) ∈Dt ` ( fw+αd ( x ) , y ) , where Dt is the training set for the t-th task previously learned . To eliminate the scaling invariance of DNNs , d is sampled from a Gaussian distribution and further normalized by dl , j ← dl , j‖dl , j‖F ‖wl , j‖F , where dl , j represents the j-th filter at l-th layer of d , and ‖·‖F denotes the Frobenius norm . Compared with our visualization , [ 26 ] only consider one task and plot the loss landscape along the directions of its Hessian eigenvectors , which only reflects some of the relationship between forgetting and landscape . Considering d is randomly selected , we repeat the visualization 10 times with different d. We first study the stability of the network by plotting changes of the weight loss landscape for the first task after new task learning . In particular , We use the previously proposed ER [ 5 ] , La-MAML [ 13 ] , and GPM [ 34 ] to train a MLP network with two hidden layers on the Permuted MNIST ( PMNIST ) [ 20 ] dataset that contains 10 tasks . We also retrain the network on the entire dataset contain all passed tasks as an Oracle network . Early stopping is used to halt the training with up to 10 epochs for each task based on the validation loss as proposed in [ 35 ] . As shown in Figure 1 ( a ) , all three continual learning methods lose their stability as learning new tasks . It can be observed from Figure 1 ( b ) - ( c ) that the weight loss landscape becomes sharper and loss value increases simultaneously , when the testing accuracy of the first task continually decreases . We further evaluate the sensitivity of the network by observing the performance of the fifth task just after it has just been learned . As shown in Figure 1 ( d ) - ( e ) , ER shows the best learning capability compared with other methods with the lowest loss value and the flattest loss landscape . Thus , based on these empirical findings , we assume that lower loss value with a flatter neighbor may lead to better continual learning performance . 3.3 A Case Study of Flattening Sharpness for Vanilla ER In this part , we further validate our above assumption by Flattening Sharpness for vanilla ER ( FS-ER ) . Compared with ER that looks for a solution w that jointly minimizes the training loss of current task data and memory data , FS-ER seeks out a solution with both low loss and flat neighbor by minimizing the maximal loss in the neighbor around the parameter value . The schematic of the Flattening Sharpness ( FS ) is shown in Figure 2 . We introduce the adversarial weight perturbation ( orange dashed line ) to explicitly flatten the weight loss landscape via injecting the worst-case weight perturbation , which is calculated from the current task data and past task data sampled from the replay buffer ( Refer to Appendix C.1 and the next section for more details ) . Figure 3 respectively shows the weight loss landscapes of the fifth task after just learning the fifth task and all tasks . We find that FS-ER successfully gets a solution with a lower loss value and flatter landscape , either after the fifth task has just been learned or all ten tasks have been learned . The average testing accuracy of all tasks using FS-ER is 90.44 % , significantly higher than ER ( 86.16 % ) , which means that flattening sharpness does benefit continual learning .

Summary: The paper proposes a new continual learning method that promotes the loss shapes to be flattened. The main idea is to optimize the worst case perturbed loss function at every gradient step, and the empirical investigation of the loss values around the learned parameter corroborates their assertion. Also, experimentally, they showed promising results compared to other recent baselines.