TAG: Task-based Accumulated Gradients for Lifelong learning

1 INTRODUCTION . Lifelong learning ( LLL ) , also known as continual learning , is a setting where an agent continuously learns from data belonging to different tasks ( Parisi et al. , 2019 ) . Here , the goal is to maximize performance on all the tasks arriving in a stream without replaying the entire datasets from past tasks ( Riemer et al. , 2018 ) . Approaches proposed in this setting involve investigating the stability-plasticity dilemma ( Mermillod et al. , 2013 ) in different ways where stability refers to preventing the forgetting of past knowledge and plasticity refers to accumulating new knowledge by learning new tasks ( Mermillod et al. , 2013 ; De Lange et al. , 2019 ) . Unlike human beings , who can efficiently assess the correctness and applicability of the past knowledge ( Chen & Liu , 2018 ) , neural networks and other machine learning models often face various issues in this setting . Whenever data from a new task arrives , these models often tend to forget the previously obtained knowledge due to dependency on the input data distribution , limited capacity , diversity among tasks , etc . This leads to a significant drop in performance on the previous tasks - also known as catastrophic forgetting ( McCloskey & J. Cohen , 1989 ; Robins , 1993 ) . Recently there has been an ample amount of research proposed in LLL ( De Lange et al. , 2019 ) . Several methods , categorized as Parameter Isolation methods , either freeze or add a set of parameters as their task knowledge when a new task arrives . Another type of methods , known as Regularization-based methods , involve an additional regularization term to tackle the stability-plasticity dilemma . There are approaches based on approximate Bayesian inference , where parameters are sampled from a distribution , that suggest controlling the updates based on parameter uncertainty ( Blundell et al. , 2015 ; Adel et al. , 2019 ; Ahn et al. , 2019 ) . But these approaches are computationally expensive and often depend on the choice of prior ( Zenke et al. , 2017 ; Nguyen et al. , 2018 ) . Another class of methods , namely Replay-based methods , store a subset of examples from each task in a replay buffer . These methods apply gradient-based updates that facilitate a high-level transfer across different tasks through the examples from the past tasks that are simultaneously available while training . As a result , these methods tend to frame LLL as an i.i.d . setting . While replay-based methods are currently state-of-the-art in several LLL tasks , it is important that we explore various ways to tackle the original non-i.i.d . problem ( Hadsell et al. , 2020 ) . Hence the focus of this paper is to design efficient replay-free methods for LLL . While adaptive gradient descent based optimizers such as Adam ( Kingma & Ba , 2014 ) have shown superior performance in the classical machine learning setup , many existing works in LLL employ the conventional stochastic gradient descent for parameter update ( Lopez-Paz & Ranzato , 2017 ; Chaudhry et al. , 2019 ; Farajtabar et al. , 2020 ) . Adaptive gradient descent based optimizers accumulate gradients to regulate the magnitude and direction of updates but often struggle when the dataset arrives in a non-i.i.d . manner , from different tasks , etc . These optimizers tend to ‘ over-adapt ’ on the most recent batch and hence suffer from poor generalization performance ( Keskar & Socher , 2017 ; Chen et al. , 2018 ; Mirzadeh et al. , 2020 ) . While exploiting the adaptive nature of the gradient descent based optimizers , we alleviate catastrophic forgetting by introducing Task-based Accumulated Gradients ( TAG ) that is a wrapper around existing optimizers . The key contributions of our work are as follows : • We define a task-aware adaptive learning rate for the parameter update step that is also aware of relatedness among tasks in LLL . • As the knowledge base , we propose to additively accumulate the directions ( or gradients ) that the network took while learning a specific task instead of storing past examples . • We empirically show that our method prevents catastrophic forgetting and exhibits knowledge transfer if the tasks are related without introducing more parameters . Our proposed method , described in Section 3 , achieves state-of-the-art results and outperforms several Replay-free methods on complex datasets like Split-miniImageNet , Split-CUB , etc . For smaller episodic memory , our method also outperforms the Replay-based methods as shown in Section 4 . Note that we propose a new approach to optimizing based on the past gradients and as such it could potentially be applied along with the existing LLL methods including replay-based methods . We demonstrate the effectiveness of doing the same in the experiments . 2 RELATED WORK . Methods proposed in LLL are broadly categorized into three classes : Regularization-based , Parameter Isolation and Replay-based methods ( De Lange et al. , 2019 ; Masana et al. , 2020 ) . Regularization-based methods prevent a drastic change in the network parameters as the new task arrives to mitigate forgetting . These methods further are classified as data-focused ( Li & Hoiem , 2017 ; Triki et al. , 2017 ) and prior-focused methods ( Nguyen et al. , 2018 ; Ebrahimi et al. , 2019 ) . In particular , Elastic Weight Consolidation ( EWC ) ( Kirkpatrick et al. , 2017 ) , a prior-focused method , regularizes the loss function to minimize changes in the parameters important for previous tasks . Yet , when the model needs to adapt to a large number of tasks , the interference between task-based knowledge is inevitable with fixed model capacity . Parameter Isolation methods ( Rusu et al. , 2016 ; Xu & Zhu , 2018 ; Serra et al. , 2018 ) such as ( Aljundi et al. , 2017 ) assign a model copy to every new task that arrives . These methods alleviate catastrophic forgetting in general , but they rely on a strong base network and work on a small number of tasks . Another closely related methods , called Expansion-based methods , handle the LLL problem by expanding the model capacity in order to adapt to new tasks ( Sodhani et al. , 2018 ; Rao et al. , 2019 ) . Li et al . ( 2019 ) propose to learn task-specific model structures explicitly while retaining model primitives sharing , decoupling from model parameter estimation . Replay-based methods maintain an ‘ episodic memory ’ , containing a few examples from past tasks , that is revisited while learning a new task ( Riemer et al. , 2018 ; Jin et al. , 2020 ) . For instance , Averaged Gradient Episodic Memory ( A-GEM ) ( Chaudhry et al. , 2018b ) , alleviating computational inefficiency of GEM ( Lopez-Paz & Ranzato , 2017 ) , uses the episodic memory to project the gradients based on hard constraints defined on the episodic memory and the current mini-batch . Experience Replay ( ER ) ( Chaudhry et al. , 2019 ) uses both replay memory and input mini-batches in the optimization step by averaging their gradients to mitigate forgetting . Task-relatedness ( Li & Hoiem , 2017 ; Jerfel et al. , 2019 ; Shaker et al. , 2020 ) or explicitly learning task representations ( Yoon et al. , 2017 ) is also an alternative approach studied in LLL . Efficient Lifelong Learning Algorithm ( ELLA ) ( Ruvolo & Eaton , 2013 ) maintains sparsely shared basis vectors for all the tasks and refines them whenever the model sees a new task . Rao et al . ( 2019 ) perform dynamic expansion of the model while learning task-specific representation and task inference within the model . Orthogonal Gradient Descent ( OGD ) ( Farajtabar et al. , 2020 ) maintains a space based on a subset of gradients from each task . As a result , OGD often faces memory issues during run-time depending upon the size of the model and the subset ( Bennani & Sugiyama , 2020 ) . Unlike OGD , we accumulate the gradients and hence alleviate the memory requirements by orders for each task . A recent work ( Mirzadeh et al. , 2020 ) argues that tuning the hyper-parameters gives a better result than several state-of-the-art methods including A-GEM and ER . They introduce Stable SGD that involves an adjustment in the hyper-parameters like initial learning rate , learning rate decay , dropout , and batch size . They present this gain in performance on simplistic benchmarks like Permuted MNIST ( Goodfellow et al. , 2013 ) , Rotated MNIST and Split-CIFAR100 ( Mirzadeh et al. , 2020 ) . Another related work ( Gupta et al. , 2020 ) also employs an adaptive learning rate while requiring a small episodic memory . But it is based on a meta-learning setting and hence beyond the scope of our paper . 3 METHOD . 3.1 LIFELONG LEARNING SETUP . In this section , we introduce the notations and the LLL setup used in the paper . We focus on the standard task-incremental learning scenario which is adopted in the numerous state-of-the-art LLL methods . It involves solving new tasks using an artificial neural network with a multi-head output where each head is associated with a unique task and the task identity is known beforehand ( Lopez-Paz & Ranzato , 2017 ; van de Ven & Tolias , 2019 ) . We denote the current task as t and any of the previous tasks by τ . The model receives new data of the form { X ( t ) , D ( t ) , Y ( t ) } where X ( t ) are the input features , D ( t ) is the task descriptor ( that is a natural number in this work ) and Y ( t ) is the target vector specific to the task t. We consider the ‘ single-pass per task ’ setting in this work following ( Lopez-Paz & Ranzato , 2017 ; Riemer et al. , 2018 ; Chaudhry et al. , 2019 ) . It is more challenging than the multiple pass setting used in numerous research works ( Kirkpatrick et al. , 2017 ; Rebuffi et al. , 2017 ) . The goal is to learn a classification model f ( X ( t ) ; θ ) , parameterized by θ ∈ RP to minimize the loss L ( f ( X ( t ) ; θ ) , Y ( t ) ) for the current task t while preventing the loss on the past tasks from increasing . We evaluate the model on a held-out set of examples of all the tasks ( ≤ t ) seen in the stream . 3.2 TASK-BASED ACCUMULATED GRADIENTS . The specific form of our proposed method depends on the underlying adaptive optimizer . For ease of exposition , we describe it as a modification of RMSProp ( Tieleman & Hinton , 2012 ) here and call it TAG- RMSProp . The TAG versions of other methods such as the Adagrad ( Duchi et al. , 2011 ) and Adam are available in Appendix A.1 . A Naive RMSProp update , for a given learning rate η , looks like the following : Vn = βVn−1 + ( 1− β ) g2n ; 1 ≤ n ≤ N θn+1 = θn − η√ Vn + gn ( 1 ) where θn is the parameter vector at step n in the epoch , gn is the gradient of the loss , N is the total number of steps in one epoch , Vn is the moving average of the square of gradients ( or the second moment ) , and β is the decay rate . We will use TAG-optimizers as a generic terminology for the rest of the paper . We maintain the second moment V ( t ) n for each task t in the stream and store it as the knowledge base . When the model shifts from one task to another , the new loss surface may look significantly different . We argue that by using the task-based second moment to regulate the new task updates , we can reduce the interference with the previously optimized parameters in the model . We define the second moment V ( t ) n for task t for TAG-RMSProp as : V ( t ) n = β2V ( t ) n−1 + ( 1− β2 ) g2n where β2 is constant throughout the stream . We use V ( t ) n to denote a matrix that stores the second moments from all previous tasks , i.e. , V ( t ) n = { V ( 1 ) N , ... , V ( t−1 ) N , V ( t ) n } of size ( t × P ) . Hence , the memory required to store these task-specific accumulated gradients increases linearly as the number of tasks in the setting . Note that each V ( τ ) N ( where τ < t ) vector captures the gradient information when the model receives data from a task τ and does not change after the task τ is learned . It helps in regulating the magnitude of the update while learning the current task t. To alleviate the catastrophic forgetting problem occurring in the Naive RMSProp , we replace Vn ( in Eq . 1 ) to a weighted sum of V ( t ) n . We propose a way to regulate the weights corresponding to V ( t ) n for each task in the next section .

Summary: In this paper, the authors propose a new optimization method for continual learning. The authors propose a task-aware optimizer to adapt the learning rate for each task. The proposed method is evaluated on several datasets to show its effectiveness.