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 investigates the backbone networks that are less prone to catastrophic forgetting. Two modifications to existing backbones are found to be useful: a mask attached to the gate function of the Highway Network, and layer-norm (without tuning parameters). Experiments on top of existing popular learning algorithms designed for continual learning (EWC, ER and HAT) show that these modifications work.