Designing Less Forgetful Networks for Continual Learning

1 INTRODUCTION . Continual learning ( McCloskey & Cohen , 1989 ) is a difficult experimental setup for most neural networks . It places a single learning agent in a dynamic environment where it must learn to adapt from a stream of tasks to make new predictions . This setup means that the underlying data are not independent and identically distributed ( non-iid ) . Thus when we use conventional gradient descent methods ( Rumelhart et al. , 1986 ) to fine-tune a pre-trained network for a new task , its weights will be modified to satisfy the learning objectives of the new task . As a result , the re-parameterised network will experience catastrophic forgetting ( McCloskey & Cohen , 1989 ; French , 1999 ) and abruptly lose the ability to solve a previously learnt task . Recently , continual learning has been of increasing interest . Most of the latest studies can be roughly categorised as the three archetypes of replay , regularisation , and dynamic architecture . Replay techniques such as gradient episodic memory ( Lopez-Paz & Ranzato , 2017 ) and experience replay ( Chaudhry et al. , 2019 ) externally store a handful of data per task . Then upon sequentially learning a new task , they re-expose copies of the stored data to the backbone network along those data of the new task . Regularisation methods like learning without forgetting ( Li & Hoiem , 2017 ) and elastic weight consolidation ( Kirkpatrick et al. , 2017 ) , on the other hand , impose the backbone network with secondary learning objectives to learn robust mappings that are less prone to parametric corruptions . Furthermore , the dynamic architectural approaches of dynamically expandable networks ( Yoon et al. , 2018 ) and hard attention to the task ( Serra et al. , 2018 ) either introduce more parameters to host new knowledge or explicitly control the backbone synaptic connections to perform inference on a subset of pre-selected weights . Though the three archetypes of continual learning methods may appear to be ideologically different , they all seek external means to retain learnt knowledge . In contrast , this paper will introduce internal modifications to mitigate forgetting . That is , we aim to create an extremely plastic yet stable network architecture that can withstand catastrophic forgetting even when it is naı̈vely implemented ( simply fine-tuned ) for sequential learning over multiple tasks . Our study was inspired by the recent work of Kuo et al . ( 2021b ) . In that 1The codes will be made publicly available once the paper is published . paper , the authors compartmentalised their backbone network . They employed a fixed feature extractor and only fine-tuned their classifiers . Unlike the conventional multiple layer perceptron ( MLP ) classifier of Figure 1 ( a ) used in most studies ( such as Mai et al . ( 2020 ) ) , their classifier which they called highway connection networks ( HCNs ) had the peculiar design shown in Figure 1 ( b ) . An HCN included a two-layer perceptron with highway connection ( Srivastava et al. , 2015 ) followed by a softmax layer . We used red to denote fully connected layers , yellow for softmax , and blue for the HCN gated unit . They demonstrated that HCNs were more robust against forgetting than MLPs ; and even showed that naı̈vely sequential learning with HCNs could often outperform MLP classifiers with continual learning techniques applied . Moreover , they found that highway connections were specifically required and could not be replaced with residual connections ( He et al. , 2016 ) . In this paper , we introduce two novel modifications to HCNs . Our new ideas draw inspiration from the stability-plasticity dilemma ( Mermillod et al. , 2013 ) . Stability is the ability to retain learnt knowledge upon parametric corruption ( Sun et al. , 2021 ) ; whereas plasticity refers to the ease of assimilating new information . Our study focuses on designing new architectures in the forward pass which indirectly but positively affect the updates in the backward pass . This kind of approach was previously investigated in weight normalisation ( Salimans & Kingma , 2016 ) . In that paper , the authors found that by reparameterising the weight vectors , they could influence the gradient covariance matrix by making it closer to an identity matrix thereby hastening network configuration . We propose the Masked Highway Connection ( MHC ) and Layer-Wise Normalisation ( LWN ) shown in Figures 1 ( c ) and 1 ( d ) respectively . While similar techniques exist for general machine learning ; their applications for continual learning are not well studied . After discussing how HCN mitigates forgetting ( Section 3 ) , we will introduce our MHC and show that HCNs can be made even more robust via self-identifying important weights that are worth fine-tuning through a masking mechanism ( Section 4 ) . Then , we will alternatively characterise forgetting as shifts that occur in previous distributions ( Ebrahimi et al. , 2020 ) ; and that the LWN can reduce such shifts ( Section 5 ) . Our experiments ( Section 6 ) indicated that both MHC and LWN were robust against forgetting and greatly outperformed their HCN and MLP baselines . In addition , both designs were compatible with the continual learning archetypes of replay , regularisation , and dynamic architecture ; and they were able to help those respective techniques to achieve new state of the art . Moreover , MHC and LWN could also be jointly implemented . Together , they often produced even more robust results and this was because they each addressed a different aspect of catastrophic forgetting . In order to suggest the best practice for MHC and LWN for continual learning , we also tested our novel classifiers over several different settings . Importantly , we found that LWN should not be incautiously substituted with batch normalisation ( Ioffe & Szegedy , 2015 ) . 2 THE CONTINUAL LEARNING FRAMEWORK . As stated in Lopez-Paz & Ranzato ( 2017 ) , the continual learning setup exposes a base learner F to a continuum of data with n locally , but not globally , iid tasks ω where ( x11 , ω1 , y 1 1 ) , . . . , ( x i φ , ωi , y i φ ) , . . . , ( x n Φ , ωn , y n Φ ) . ( 1 ) Each task corresponds to a dataset ( xiφ , y i φ ) ∈ Dωi with Φ samples φ = 1 , . . . , Φ . Once the learner F observes data from the j-th task , full access to an earlier dataset ωi where i < j is prohibited to consolidate learnt knowledge . The performance is tracked using a matrix A ∈ R ( n+1 ) ×n where each column represents a distinct task . When F is initialised , its accuracy over all tasks is recorded in row A1 , ( 1 : n ) . After learning the i-th task , the new performances are updated in row A ( i+1 ) , ( 1 : n ) . 3 HIGHWAY CONNECTIONS IMPLICITLY RE-SCALE GRADIENTS . An intermediate layer Lγ in an MLP can be represented as aLγ = H ( aL ( γ−1 ) , WLγ ) ( 2 ) where the layer output aLγ is transformed from the layer input aLγ−1 over the non-linear function H with weight WLγ . In addition , say there are Γ layers then the network output is aLΓ = ŷ . To fine-tune WLγ , we use gradient descent with the derivative of a costC with respect to the parameters ∂C ∂WLγ = ∂C ∂ŷ Γ−1∏ ζ=γ ∂aL ( ζ+1 ) ∂aLζ ︸ ︷︷ ︸ Rel : Srivastava et al . ( 2015 ) Rel : Kuo et al . ( 2021b ) ︷ ︸︸ ︷ { ∂aLγ ∂WLγ } . ( 3 ) Below , we discuss how gradient vanishing and catastrophic forgetting are closely related ( Rel : ) . Deep MLPs are known to suffer gradient vanishing ( Bengio et al. , 1994 ) . When the layers increase but the magnitudes of the individual Jacobian terms ∂aL ( ζ+1 ) ∂aLζ are small ( i.e. , ∈ [ 0 , 1 ] ) , the term∏Γ−1 ζ=γ ∂aL ( ζ+1 ) ∂aLζ diminishes to 0 ( and so does ∂C∂WLγ ) and no effective update is made to WLγ . Highway connection ( Srivastava et al. , 2015 ) is well-known for its ability to withstand gradient vanishing . It reformulates an intermediate layer as kLγ = ( 1−TLγ ) kLγ−1 + TLγ H ( kL ( γ−1 ) , WLγ ) with ( 4 ) gate TLγ = T ( kL ( γ−1 ) , QLγ ) ( 5 ) and the element-wise dot product . The colour-coded gated unit TLγ has its own weight QLγ and its values lie in the range of [ 0 , 1 ] . This new formulation has the Jacobian term ∂kLγ ∂kL ( γ−1 ) = { I , if TLγ = 0 H ′ ( kL ( γ−1 ) , WLγ ) , if TLγ = 1 . ( 6 ) By introducing the identity matrix I to the Jacobian ( and thus the [ . ] term of Equation ( 3 ) ) , highway connections prevent the vanishing of gradients and thus facilitate the training of deep networks . In Kuo et al . ( 2021b ) , the authors noted a second property unique to highway connections . While the { . } term of Equation ( 3 ) of an MLP was ∂aLγ∂WLγ = ∂H ( aL ( γ−1 ) , WLγ ) ∂WLγ , a highway connection had the alternative form of ∂kLγ ∂WLγ = TLγ ∂H ( kL ( γ−1 ) , WLγ ) ∂WLγ which included the gated unit . Since each entry in TLγ lied in [ 0 , 1 ] , highway connections thus implicitly scaled down the magnitudes of the gradients used in gradient descent . This was desirable for continual learning because less changes were made to a pre-trained weight WLγ ; and HCNs were hence less susceptible to forgetting . 4 IMPROVING THE HCN ARCHITECTURE . Though highway connection is effective against catastrophic forgetting , its best practices and shortcomings remain unclear . In this section , we review the different roles that the gated unit TLγ plays in the forward pass ( inference ) and in the backward pass ( update ) . Furthermore , we will show that there exists an interesting structural connection between the HCN gated unit TLγ and the derivative ∂H ( kL ( γ−1 ) , WLγ ) ∂WLγ which can be exploited to make the network even more robust to forgetting . Overleaf in the top half of Figure 2 , we show that HCN performs a point-to-point gradient reduction . On the left , we superimpose a dotted box on top of HCN . The inference and update that occur within this box are illustrated on the right . Without loss of generality , the scenario in the figure is for 3 hidden dimensions and that each TLγ entry scales down its corresponding feature of the fully connected layer output kLγ during inference . Note that the entries of kLγ are shown in different shades of red ; and they correspond to specific rows in the weight WLγ . Likewise for update , we highlight the rows of ∂H ( kL ( γ−1 ) , WLγ ) ∂WLγ in different shades of green ; and we show how each row of this derivative is scaled down , in a point-to-point fashion , by the entries of the gated unit TLγ . But does the entirety of weight WLγ need to be updated at all times ? Two prior studies , the lottery ticket hypothesis ( LTH ) ( Frankle & Carbin , 2019 ) and the uncertainty-regularised continual learning ( UCL ) ( Ahn et al. , 2019 ) , suggested not . In network pruning , LTH found that only a small subset of the dense network needed to be updated to achieve competitive performances . Whereas UCL showed that continual learning could be achieved by regularising specific nodes in the weight that deviated from a pre-determined distribution . They both considered the structural relationship among nodes in the base learner . A similar approach can be applied to HCN . The lower-half of Figure 2 presents the inner workings of our Masked Highway Connecton ( MHC ) classifier . We add a binary Mask to HCN and change the formulation from Equations ( 4 ) and ( 5 ) to kLγ = ( 1−GLγ ) kLγ−1 + GLγ H ( kL ( γ−1 ) , WLγ ) with ( 7 ) a masked gate GLγ = Mask TLγ = Mask T ( kL ( γ−1 ) , QLγ ) . ( 8 ) The Mask keeps the largest neural values in TLγ with the top-K function . As shown in Figure 2 , when TLγ ∈ R3 and that K = 1 , the MHC masked gate GLγ nullifies the two rows with weaker activation values . The mask also nullifies specific updates during optimisation . In our experiments in Section 6 , we will demonstrate that this minimalistic update can reduce excessive weight modification thereby improving network stability for the sequential learning setup .

Summary: The paper proposes modifications to the highway connection network (HCN) architecture, which exhibits some inherent resistance to catastrophic forgetting.  These modification introduce fairly simple masking rule for selection of portion of weights to update and normalisation that (presumably) further adjust the internal representation to be more selective in which weights to update for a given task.  Empirical evaluation shows that these modification are highly effective in alleviating catastrophic forgetting.