Linear Mode Connectivity in Multitask and Continual Learning

1 INTRODUCTION . One major consequence of learning multiple tasks in a continual learning ( CL ) setting — where tasks are learned sequentially , and the model can only have access to one task at a time — is catastrophic forgetting ( McCloskey & Cohen , 1989 ) . This is in contrast to multitask learning ( MTL ) , where the learner has simultaneous access to all tasks , which generally learns to perform well on all tasks without suffering from catastrophic forgetting . This limitation hinders the ability of the model to learn continually and efficiently . Recently , several approaches have been proposed to tackle this problem . They have mostly tried to mitigate catastrophic forgetting by using different approximations of the multitask loss . For example , some regularization methods take a quadratic approximation of the loss of previous tasks ( e.g . Kirkpatrick et al. , 2017 ; Yin et al. , 2020 ) . As another example , rehearsal methods attempt to directly use compressed past data either by selecting a representative subset ( e.g . Chaudhry et al. , 2019 ; Titsias et al. , 2019 ) or relying on generative models ( e.g . Shin et al. , 2017 ; Robins , 1995 ) . In this work , we depart from the literature and start from the non-conventional question of understanding “ What is the relationship , potentially in terms of local geometric properties , between the multitask and the continual learning minima ? ” . Our work is inspired by recent work on mode con- ∗Equal contribution 1The code is available at : https : //github.com/imirzadeh/MC-SGD nectivity ( Draxler et al. , 2018 ; Garipov et al. , 2018 ; Frankle et al. , 2020 ) finding that different optima obtained by gradient-based optimization methods are connected by simple paths of non-increasing loss . We try to understand whether the multitask and continual solutions are also connected by a manifold of low error , and what is the simplest form this manifold can take . Surprisingly , we find that a linear manifold , as illustrated in Fig . 1 right , reliably connects the multitask solution to the continual ones , granted that the multitask shares same initialization with the continual learning as described below . This is a significant finding in terms of understanding the phenomenon of catastrophic forgetting through the lens of loss landscapes and optimization trajectory and also for designing better continual learning algorithms . To reach this conclusion , we consider a particular learning regime described in Fig . 1 left , where after learning the first task using the data D1 , we either sequentially learn a second task obtaining ŵ2 or continue by training on both tasks simultaneously ( i.e. , train on D1 + D2 ) , obtaining the multitask solution w∗2 . We investigate the relationship between the two solutions ŵ2 and w∗2 . Note that w∗2 is not the typical multitask solution , which would normally start from w0 and train on both datasets . We chose this slightly non-conventional setup to minimize the potential number of confounding factors that lead to discrepancies be- tween the two solutions ( Fort et al. , 2019 ) . We also rely on the observation from ( Frankle et al. , 2020 ) that initialization can have a big impact on the connectivity between the solutions found on the same task , and sharing the same starting point , as we do between ŵ2 and w∗2 , might warrant a linear path of low error between the two solutions . Moreover , Neyshabur et al . ( 2020 ) noted that in the context of transfer learning , there is no performance barrier between two minima that start from pre-trained weights , which suggests that the pre-trained weights guide the optimization to a flat basin of the loss landscape . In contrast , barriers clearly exist if these two minima start from randomly initialized weights . Our contributions can be summarized as follows : 1 . To the best of our knowledge , our work is the first to study the connectivity between continual learning and multitask learning solutions . 2 . We show that compared to conventional similarity measures such as Euclidean distance or Central Kernel Alignment ( Kornblith et al. , 2019 ) , which are incapable of meaningfully relating these minima , the connectivity through a manifold of low error can reliably be established . And this connecting path is linear , even when considering more than 20 tasks in a row . 3 . Motivated by this , we propose an effective CL algorithm ( Mode Connectivity SGD or MC-SGD ) that is able to outperform several established methods on standard CL benchmarks . 1.1 RELATED WORK . With the trending popularity of deep learning , continual learning has gained a critical importance because the catastrophic forgetting problem imposes key challenges to deploy deep learning models in various applications ( e.g Lange et al. , 2019 ; Kemker et al. , 2018 ) . A growing body of research has attempted to tackle this problem in recent years ( e.g Parisi et al. , 2018 ; Toneva et al. , 2018 ; Nguyen et al. , 2019 ; Farajtabar et al. , 2019 ; Hsu et al. , 2018 ; Rusu et al. , 2016 ; Li et al. , 2019 ; Kirkpatrick et al. , 2017 ; Zenke et al. , 2017 ; Shin et al. , 2017 ; Rolnick et al. , 2018 ; Lopez-Paz & Ranzato , 2017 ; Chaudhry et al. , 2018b ; Riemer et al. , 2018 ; Mirzadeh et al. , 2020 ; Wallingford et al. , 2020 ) . Among these works , our proposed MC-SGD bares most similarities to rehearsal based methods such us ( e.g . Shin et al. , 2017 ; Chaudhry et al. , 2018b ) and regularization based methods ( e.g . Kirkpatrick et al. , 2017 ; Zenke et al. , 2017 ) similar to ( Titsias et al. , 2019 ) . Following ( Lange et al. , 2019 ) , one can categorize continual learning methods into three general categories , based on how they approach dealing with catastrophic forgetting . Experience replay : Experience replay methods build and store a memory of the knowledge learned so far ( Rebuffi et al. , 2016 ; Lopez-Paz & Ranzato , 2017 ; Shin et al. , 2017 ; Riemer et al. , 2018 ; Rios & Itti , 2018 ; Zhang et al. , 2019 ) . As examples , Averaged Gradient Episodic Memory ( A-GEM ) ( Chaudhry et al. , 2018b ) builds an episodic memory of parameter gradients , while ER-Reservoir ( Chaudhry et al. , 2019 ) uses a Reservoir sampling method to maintain the episodic memory . Regularization : These methods explicitly apply regularization techniques to ensure parameters do not change too much ( Kirkpatrick et al. , 2017 ; Zenke et al. , 2017 ; Lee et al. , 2017 ; Aljundi et al. , 2018 ; Kolouri et al. , 2019 ) . They can either a Bayesian ( Nguyen et al. , 2017 ; Titsias et al. , 2019 ; Schwarz et al. , 2018 ; Ebrahimi et al. , 2019 ; Ritter et al. , 2018 ) or frequentist views ( Farajtabar et al. , 2019 ; He & Jaeger , 2018 ; Zeng et al. , 2018 ) . For instance , Orthogonal Gradient Descent ( OGD ) ( Farajtabar et al. , 2019 ) projects the prediction gradients from new tasks on the subspace of previous tasks ’ gradients to preserve the knowledge . Parameter isolation : Parameter isolation methods allocate different subsets of the parameters to each task ( Rusu et al. , 2016 ; Yoon et al. , 2018 ; Jerfel et al. , 2019 ; Rao et al. , 2019 ; Li et al. , 2019 ) . From the stability-plasticity perspective , these methods implement gating mechanisms that improves the stability and controls the plasticity by activating different gates for each task . Masse et al . ( 2018 ) proposes a bio-inspired approach for a context-dependent gating that activates non-overlapping subset of parameters for any specific task . Supermask in Superposition ( Wortsman et al. , 2020 ) is another parameter isolation method that starts with a randomly initialized , fixed base network and for each task finds a subnetwork ( supermask ) such that the model achieves good performance . Recently , Mirzadeh et al . ( 2020 ) have shown that dropout implicitly creates different pathways or gates for tasks , therefore , it reduces their mutual interference and leads to less forgetting . Continual learning as a problem expands beyond dealing with catastrophic forgetting , one of the hopes behind sequential learning is that it can potentially enable positive forward transfer , as one can build on the previously acquired knowledge . In this sense it connects to topics such as MetaLearning ( Beaulieu et al. , 2020 ; Jerfel et al. , 2019 ; He & Jaeger , 2018 ; Riemer et al. , 2018 ) , Few-Shot Learning ( Wen et al. , 2018 ; Gidaris & Komodakis , 2018 ) , Multi-task and Transfer Learning ( He & Jaeger , 2018 ; Jerfel et al. , 2019 ) . It also aims to work in scenarios where task boundaries are not well defined or provided , or when the data distribution shifts slowly or when a multi-task solution does not exist ( Rao et al. , 2019 ; Aljundi et al. , 2019 ; He et al. , 2019 ; Kaplanis et al. , 2019 ) . Mode connectivity ( Draxler et al. , 2018 ; Garipov et al. , 2018 ) , is a novel tool to understand the loss landscape of deep neural networks . It postulates that different optima obtained by gradient-based optimization methods are connected by simple paths of non-increasing loss ( i.e. , low-loss valleys ) . Recently , various works provided different theoretical explanations for mode connectivity in different scenarios ( Venturi et al. , 2019 ; Kuditipudi et al. , 2019 ) either by relying on over-parametrization or concepts such as noise-stability or dropout-stability . ( Neyshabur et al. , 2020 ) investigated the connection between minima obtained by pre-trained models versus freshly initialized ones . They note that there is no performance barrier between solutions coming from pre-trained models , but there can be a barrier between solutions of different randomly initialized models . ( Frankle et al. , 2020 ) shows that different minima that share the same initialization point are connected by a linear path , even with weight pruning . We rely on this observation when designing our setting that the multitask and continual learning minima share a common starting point . 2 THE RELATION BETWEEN MULTITASK AND CONTINUAL MINIMA . One question driving this work is understanding the relationship between the two different solutions : multitask learning vs. continual learning . In particular , we are interested in scenarios where a multitask solution for the considered tasks exists ( for a discussion when this does not hold see ( He et al. , 2019 ) ) and when both learning regimes have the same objective , finding a solution that performs well on all tasks . We posit that this difference can not be reliably explained by simple and typical distance metrics used to measure similarity between the trained models . In particular we consider Euclidean distance and Central Kernel Alignment ( CKA ) ( Kornblith et al. , 2019 ) . However , these solutions are connected by paths of low error , and , provided that learning starts from the same initial conditions , these paths can have a linear form . In Fig . 2 left column we can see the performance of Naive SGD for multitask and continual learning . Details on the experimental setup can be found in the Appendix C. The dashed line represents the multitask solution at convergence , which achieves strong performance on all tasks . It further shows the performance of all tasks during the sequential learning experience ( each point represents the performance after learning another task ) , highlighting how performance on past tasks degrades considerably . Note that as described in Fig . 1 and further detailed in Appendix D.2 Fig . 15 , in the multitask learning scenario , tasks are added sequentially to the loss , to construct a parallel with the continual learning setting . This will be the case throughout this work . Eucledian distance . It might be reasonable to expect that the less the parameter changes , the less the forgetting will be . One can motivate this heuristic on a Taylor expansion of the loss , as done in ( Mirzadeh et al. , 2020 ) , where , the forgetting is defined as : L1 ( ŵ2 ) − L1 ( ŵ1 ) ≈ 1 2 ( ŵ2 − ŵ1 ) > ∇2L1 ( ŵ1 ) ( ŵ2 − ŵ1 ) ≤ 1 2 λmax1 ‖ŵ2 − ŵ1‖2 . ( 1 ) Here , L1 ( w ) is the empirical loss for task 1 , and λmax1 is the largest eigenvalue of its Hessian at ŵ1 . Note that all terms of the Taylor expansion are multiplied with ŵ2 − ŵ1 and hence the norm of this delta will affect the amount of forgetfulness . In fact , this is frequently done when pruning neural networks ( Zhu & Gupta , 2018 ; Han et al. , 2015 ) , where weights are zeroed out based on magnitude , producing minimal Euclidean distance to unprunned model . But , as observed in Fig . 2 ( middle column ) , Eucledian distance does not correlate with not suffering from catastrophic forgetting : the CL solution on task 1 ( ŵ1 ) is closer to the CL solution of task 5 ( ŵ5 ) than the multitask solution of tasks 5 ( w∗5 ) . One explanation could be that the bound defined by Eq . ( 1 ) will not be tight if the vector ŵ5 − ŵ1 does not lie in the direction of the largest eigenvector . Appendix D.1 contains further details on this topic . Centered Kernel Alignment . Centered Kernel Alignment ( CKA ) ( Kornblith et al. , 2019 ) measures the similarity of two representations on the same set of examples . Given N examples and two activation outputs on these examples , R1 ∈ RN×d1 and R2 ∈ RN×d2 , CKA is defined by : CKA ( R1 , R2 ) = ‖R > 1 R2‖2F ‖R > 1 R1‖F ‖R > 2 R2‖F , ( 2 ) where , ‖.‖F is the Frobenius norm . Recent work by Ramasesh et al . ( 2020 ) studies catastrophic forgetting on the CIFAR dataset by measuring the CKA similarity score of different layers of ŵ1 and ŵ2 . They argue that the later layers suffer more from catastrophic forgetting by showing that the CKA similarity of initial layers decreases less after training on sequential tasks . However , the CKA score suffers from a few shortcomings . If the number of training epochs per task is small ( e.g. , in streaming case ) , the CKA does not change much , even though the accuracy for previous tasks drops drastically . For instance , in Fig . 2 right column , we show that the pairwise CKA between different layers of the first task minimum ( ŵ1 ) and CL and multitask minima of task 2 and task 5 are roughly the same . Although , the phenomenon observed in ( Ramasesh et al. , 2020 ) is still realizable by a very tiny margin . Moreover , we can see that the multitask minimum of task 2 ( w∗2 ) , and task 5 ( w∗5 ) are more similar to ŵ1 compared to CL minima ( ŵ2 and ŵ5 ) .

Summary: The paper studies the relation between the geometry of solutions of continual (CL) and multi-task learning (MTL). Towards this end, the authors empirically identify that all the solutions of CL (i.e. solutions obtained after each task) and MTL are connected by a linear region of low error. This is a very interesting finding and, to my knowledge, has not been studied previously in the CL literature. Based on this observation, the authors propose a memory and regularization-based CL method, MC-SGD, that ensures that the final CL solution is linearly connected to all the task’s solutions. The authors further demonstrate that the solution of the MTL lies in the region where the Hessian of the loss function is low and hence the regularization-based approaches that make use of curvature information (e.g.) EWC, are a promising direction for CL. Experiments are conducted on Permuted and Rotated MNIST, Split CIFAR benchmarks. MC-SGD performs strongly compared to other baselines.