Continual Learning via Local Module Composition

1 Introduction . The goal of continual learning ( CL ) is to learn efficiently from a non-stationary stream of tasks without ( catastrophically ) forgetting previous tasks [ 62 ] . CL is often modeled as a trade-off between knowledge retention ( stability ) and knowledge expansion ( plasticity ) [ 26 , 64 ] . Parameter sharing can provide control over this trade-off . For example , learning a single model shared across tasks results in better knowledge transfer and faster learning at the expense of forgetting [ 46 , 57 ] . Conversely , learning a separate model per task eliminates forgetting but minimizes transfer and data efficiency [ 2 , 41 ] . Modular learning aims at balancing transfer and forgetting by learning a set of specialized modules that can be recomposed to solve ( new ) tasks while only updating a subset of relevant modules or adding new modules [ 6 , 47 , 27 ] . In principle , a modular learner capable of composing modules in meaningful structures can provide additional benefits including ( i ) computational gains due to only executing modules that are relevant to a task [ 47 , 4 ] ; ( ii ) memory gains due to instantiating a sub-linear number of modules w.r.t . the number of tasks ; ( iii ) systematic [ 8 ] and out-of-distribution ( OOD ) generalization [ 18 ] through knowledge recombination ; and ( iv ) biological plausibility [ 91 , 90 , 96 ] . corresponding author : oleksiy.ostapenko @ t-online.de 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Designing modular methods for CL comes with two main challenges . The first is how and when to add new modules to ensure sufficient plasticity to learn new tasks . Existing modular methods use greedy search variants , expanding the model when it improves validation performance [ 92 , 63 ] . The second challenge is how to compose that is , retrieve task-specific structural knowledge given a new task ( previously seen or not ) . Existing methods rely on a task ’ s identifier ( ID ) to retrieve task-specific structural knowledge , which comes either in the form of an optimal module layout [ 92 ] or as a model- and task-specific controller network that generates modular layouts [ 63 ] . Unfortunately , in many realistic CL scenarios task identities are unavailable at test time [ 23 , 35 , 14 ] . Lifting this limitation is challenging since standard mechanisms for task inference , for example , leveraging a task-inference model , could be subject to forgetting themselves . To address both challenges , we equip each module with a local structural component that predicts a score indicating how relevant the module is for a given input . In-distribution inputs result in high scores , while out-of-distribution inputs result in low scores . In other words , modules self-determine their relevance given an input . This local component is used for composing modules : for each datum , modules are combined at each layer according to their normalized scores without requiring a task ’ s ID ( §3 ) . The local component is also used for module expansion : a new module is instantiated if all the current modules flag their input as being locally out-of-distribution ( §3.1 ) . Further , new shallow modules ( i.e . closer to the input ) are first trained in a projection phase to maximize the relatedness scores of subsequent , deeper , modules ( §3.2 ) . This process projects the output of new modules into the representation space expected by the subsequent modules and ensures the compatibility between low- and high-level modules . In a set of studies , we explore the performance and versatility of our local structural approach , which we call Local Module Composer ( LMC ) . First , we show that LMC reaches superior or comparable performance to existing modular and non-modular methods without requiring task IDs at test time using the Continual Transfer Learning ( CTrL ) benchmark , designed to evaluate transfer and forgetting in CL [ 92 ] ( §4.1 ) . Then , we demonstrate how LMC , relying on its projection phase , can solve out-ofdistribution ( OOD ) tasks not seen during the continual training ( §4.2 ) . We also show it is possible to combine modules from independently trained models into a new model to solve tasks seen by each of the independent models without any finetuning ( §4.3 ) . Finally , an analysis of longer task sequences ( 30 and 100 tasks ) reveals that LMC tends to spawn much fewer modules to reach good performance than the fully task-aware MNTDP [ 92 ] counterpart . However , LMC reaches slightly lower accuracy on longer sequences than MNTDP , which highlights the difficulty of automatic task-ID agnostic module selection in the presence of a large number of candidate modules . In Appendix F we demonstrate the applicability of LMC in the meta-continual learning ( meta-CL ) setting , a task-agnostic setting by nature . We highlight that by relying on a local ( per-module ) structural component , LMC offers a modular CL approach that i ) does not require task IDs during test in the standard task incremental settings ; ii ) balances parameter sharing to yield strong CL performances compared to baselines that require access to the task ID ; iii ) in our experiments instantiates a sub-linear number of modules ; iv ) permits recombination of modules at test time enabling OOD generalization as well as ( v ) the ability to combine independently trained models in a third model without fine-tuning . Notably , the OOD generalization is only possible if the agent is task-agnostic in the module selection process , since OOD tasks were not observed at training , the learner has to interpolate between the learned tasks , and a ( categorial ) task ID is of no use . 2 Background : Modular Continual Learning . Let F ( x ; θ ) : X → Y be a learner parametrized with a set of parameters θ . In task-incremental CL , the learner is exposed to a sequence of tasks . Each task is composed of a training set Dt of ( x , y ) pairs and a task identifier ( ID ) t [ 92 , 46 ] . The goal is to learn an optimal θ∗ that minimizes the loss L for all observed tasks : θ∗ = arg min θ T∑ t=1 E ( x , y ) ∼Dt [ L ( F ( x ; θ ) , y ) ] . ( 1 ) The parameter sharing trade-off between tasks can be addressed through different architectural design choices for F . For example , F can be a monolithic network that shares parameters θ across all tasks . Most existing task incremental CL methods use a task-specific output head , requiring the task ID to select the output head corresponding to the task at hand [ 46 , 87 , 1 ] . At the other end of the spectrum are the expert based solutions that learn an independent model , a.k.a . expert , for each task [ 2 , 83 ] . In this case , each expert trains task-specific parameters θ = { θ ( t ) } Tt=1 . To balance parameter sharing and transfer , modular methods organize their parameters in a series of modules M = { m ( l ) k } with parameters θ = { θ ( l ) k } , where θ ( l ) k denotes the parameters of module k at layer l in F . In general , a module can be any parametric function . In our experiments , unless otherwise stated , a module consists of a single convolutional layer followed by batch-norm , ReLU activation , and a max-pooling operation . Modules can be composed conditioned on a sample , a batch of samples , or a task . Let ψ denote a specific composition of modules that gives rise to a distinct prediction function ; we make this dependence explicit : F ( x ; θ , ψ ) . Importantly , sharing modules across tasks should lead to desirable transfer properties . Veniat et al . [ 92 ] frames modular CL as finding an optimal layout ψ ( t ) for each task , where each layout selects a single module per layer per task ( hard selection ) : θ∗ , Ψ∗ = arg min θ , Ψ T∑ t=1 E ( x , y ) ∼Dt [ L ( F ( x ; θ , ψ ( t ) ) , y ) ] . ( 2 ) In this case the set of layouts Ψ = { ψ ( t ) } grows with the number of tasks , while modules can be reused across different task-specific layouts resulting in sub-linear growth pattern . They design a method called MNTDP to search the exponentially large space of modular layouts by only considering layouts resulting from adding a new module per layer to the best prior path ( past task ’ s path with the highest nearest neighbor accuracy on a new task ) starting at the top layer . This solution relies on task IDs to retrieve ψ ( t ) at test time . Another way of composing modules uses dynamic routing [ 63 , 82 , 47 , 65 ] . The module layout is generated by a structural function ψ = s ( x ) , hence different inputs take different routes through F . It is standard to approximate the structural function using a neural network ψ = s ( x ; φ ) with structural parameters φ . This framework has been applied to CL in [ 63 ] by learning a separate structural function per task ψ ( t ) = s ( x ; φ ( t ) ) . The task IDs are used to retrieve the correct structural function : θ∗ , Φ∗ = arg min θ , Φ T∑ t=1 E ( x , y ) ∼Dt [ L ( F ( x ; θ , s ( x ; φ ( t ) ) ) , y ) ] , ( 3 ) where Φ = { φ ( t ) } is the set of structural parameters for all tasks . The above methods require task IDs at both training and testing time . Next we introduce our modular CL approach that only relies on task IDs during training . 3 Local Module Composer ( LMC ) . We propose LMC , a modular approach where each module consists of a functional component f ( x ; θ ( l ) m ) and a structural component s ( x ; φ ( l ) m ) , see Figure 1 . The functional components are responsible for learning to solve the prediction task and are trained via the usual task loss L ( e.g . cross-entropy loss for classification ) . The structural components receive the corresponding functional output as their input ( see Figure 1 ) and are responsible for dynamic routing through F . Structural parameters φ are trained using a structural loss L ( st. ) computed locally at each module . Intuitively , the structural component of a module should serve as a density estimator of the outputs of the functional component . The module ’ s contribution to the layer ’ s output is proportional to the likelihood of the input sample under the estimated density . In our instantiation , the structural component produces a relatedness score : a lower score for inputs that are more likely to belong to the distribution on which a given module was trained , and a higher score for inputs that are out-of-distribution for the given module . Hence , the likelihood of the input sample is approximated by the negative relatedness score . Given an input data sample x ( 0 ) = x , the output x ( l ) of a layer l is defined as the weighted sum of the functional outputs of all |M ( l ) | local modules and used as input to the subsequent layer l + 1 : x ( l ) = |M ( l ) |∑ m=1 w ( l ) m · f ( x ( l−1 ) ; θ ( l ) m ) . ( 4 ) The functional output of the network is equal to the output of the final layer : F ( x ; θ , φ ) = x ( L ) . In the last layer F implements a single output-head per task . At training time , the task ID is available and we update only the output-head corresponding to the currently learned task . At test time , the task ID is not available and we select the output head with the highest activation weight w ( L ) m , i.e. , the last layer performs hard module selection . The module activation weights w ( l ) m are computed by normalizing the vector of local relatedness scores γ ( l ) ∈ R|M ( l ) | . Each element of γ ( l ) is obtained from the negative structural loss which approximates the likelihood of each module : γ ( l ) m = −L ( st. ) ( s [ f ( x ( l−1 ) ; θ ( l ) m ) ; φ ( l ) m ] ) , ( 5 ) w ( l ) m = softmax ( γ ( l ) ) m. ( 6 ) Modules with lower structural loss get higher activation weights . Note that in practice , it can be useful to bias the module selection towards the expected module selection in a batch , assuming that samples within a batch are likely to belong to the same task . We discuss this point further in § A.1 . Instead of using the softmax function , it is possible perform hard selection taking the module with the highest score [ 82 ] , or alternatively selecting top-k modules [ 88 ] . In both cases , LMC ’ s structural parameters stay differentiable due to the local nature of structural learning . Note that in the case of global structural objective , hard module selection would require applying tools for non-differentiable learning such as Expectation Maximization [ 47 ] or reinforcement-learning based methods [ 82 , 5 ] . The overall LMC objective consists of optimizing both functional and structural losses : θ∗ , φ∗ = arg min θ , φ T∑ t=1 E ( x , y ) ∼Dt [ L ( F ( x ; θ , φ ) , y ) + L∑ l=0 |M ( l ) |∑ m=0 L ( st. ) m ( s [ f ( x ( l−1 ) ; θ ( l ) m ) ; φ ( l ) m ] ) ] . ( 7 ) As in [ 92 ] , learning is performed w.r.t . only newly introduced modules to prevent forgetting . Structural component . We test two instantiations of the structural component s and loss L ( st. ) m . In the first one , s is an invertible neural network [ 80 ] . Here we use the invertible architecture proposed by Dinh et al . [ 21 ] . As shown by Hocquet et al . [ 37 ] , for this invertible architecture the structural objective can be defined as L ( st. ) m ( x ) = ||x||2 . Intuitively , an invertible architecture prevents L ( st. ) m from collapsing to an all-zeros solution . In the second instantiation , s and f form an autoencoder and L ( st. ) m ( x ) = ||x ( l−1 ) − x||2 is the reconstruction error with respect to the module ’ s input x ( l−1 ) . Aljundi et al . [ 2 ] used a similar idea was for selecting the most relevant expert network conditioned on a task . Unless stated otherwise , modules in the feature extractor use the autoencoder as their structural component , while output heads use invertible s — these combinations worked well in practice .

Summary: This paper tries to tackle continual learning with a modular approach and proposes the local module composition (LMC) approach. LMC composes modules dynamically by computing local relevance scores for each layer, computed with the input. This architecture makes LMC perform well for OOD tasks as well as outperforming existing approaches on a continual learning benchmark.