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: The paper proposes a new algorithm for addressing continual learning (CL) by leveraging modular architectures. In contrast to (the few) existing modular CL works, this one focuses on the task-agnostic setting, where the agent is not informed of the task ID. The proposed architecture assigns two components to each module: a functional component and a structural component. The latter is in charge of locally determining whether each data point (actually each mini-batch, in practice) is in-distribution for the given module. The training proceeds in stages, adding new modules whenever the new data point is out-of-distribution for all modules, then training the module such that its output is in-distribution for modules at later depths, and finally training the entire model together. Extensive experiments are carried out to demonstrate the strengths of the proposed approach.