Compositional Models: Multi-Task Learning and Knowledge Transfer with Modular Networks

1 INTRODUCTION . Most modern interpretations of neural networks treat layers as transformations that convert lowerlevel features , such as pixels , into higher-level abstractions . Adopting this view , many designs that break up monolithic models into interacting modular components localize them to particular processing stages thus disallowing parameter sharing across layers . This restriction , however , can be limiting when we need to continually grow the model , be that for transferring knowledge to a new task or preventing catastrophic forgetting . It also prohibits module reuse within a model even though certain parameter-efficient architectures and complex decision making processes may involve recurrent components ( Kaiser & Sutskever , 2016 ; Randazzo et al. , 2020 ) and feedback loops ( Herzog et al. , 2020 ; Yan et al. , 2019 ; Kar et al. , 2019 ) . Here we propose a simple alternative design that represents an entire model as a mixture of modules each of which can contribute to computation at any processing stage . As opposed to many other approaches ( Zoph & Le ( 2016 ) ; Bengio ( 2016 ) ; Kirsch et al . ( 2018a ) ; Rosenbaum et al . ( 2019 ) etc . ) , we use a soft mixture of modules to obtain a more flexible model and produce a differentiable optimization objective that can be optimized end-to-end and does not involve high-variance estimators . Following our approach , the parameters of every block ( or layer ) of the network is computed as a linear combination of a set of “ template ” block parameters thus representing the entire model as : ( a ) a databank of template blocks and ( b ) vectors of “ mixture weights ” that are used to generate weights for every layer . This simple design can be utilized for a variety of applications from producing compact networks and training multi-task models capable of re-using individual modules to knowledge transfer and domain adaptation . The experimental results show that : ( a ) when used for multi-task training , our model organizes its modules where tasks share first few layers while specializing closer to the head , while ( b ) in domain adaptation problems modules instead specialize on processing the image , while sharing later processing stages . Moreover , our self-organizing model achieves promising results in multi-task learning and model personalization . The rest of the paper is organized as follows : in Section 2 we go over the related literature and discuss the existing approaches to modular networks and neural architecture search ; in Section 3 we describe the architecture of our conversational model and various aspects of learning the best module sequence for each specific task ; we introduce the experimental results in a single-task setting , multi-task learning , continual learning and domain adaptation in Section 4 and provide a final discussion in Section 5 . 2 PRIOR WORK . The idea of constructing a deep neural network as a composition of reusable computation blocks has been extensively explored in recent literature . First conditioned computation methods used evolution algorithms ( Wierstra et al . ( 2005 ) ; Floreano et al . ( 2008 ) ; Miikkulainen et al . ( 2019 ) ) to determine a suitable model architecture ; later on , reinforcement learning was used to optimize the model layout and parameters ( Zoph & Le ( 2016 ) ; Bengio ( 2016 ) ; Baker et al . ( 2016 ) ) . Recently , Kirsch et al . ( 2018a ) proposed an end-to-end modular neural network framework that learns to choose the best among several module candidates at each layer according to the input data . Conditional computation routing is even more appealing for the multitask learning application , as it allows to both learn reusable computation blocks and adjust the model to each specific task with minimal network alteration . For instance , Misra et al . ( 2016 ) use an additional set of modules for each task and enforce similarity between the corresponding task-specific modules weights . Rosenbaum et al . ( 2017 ; 2019 ) used reinforcement learning to perform task- and data-specific routing ; Sun et al . ( 2019 ) train both the model weights and the task-specific policy that determines which layers should be shared for a given task ; Ma et al . ( 2019 ) ; Maziarz et al . ( 2019 ) introduced a multitask model where each layer ’ s output is computed as a weighted sum of a set or candidate modules outputs , which is similar to our method , with two important differences : a ) in the works by Ma et al . ( 2019 ) ; Maziarz et al . ( 2019 ) , the output is computed as a linear combination of candidates outputs , not module weights , and b ) the module candidates in their methods are tied to the specific location in the network . Maninis et al . ( 2019 ) added a task-specific residual adapters to specialize the feature extractor to each task , The method proposed by Purushwalkam et al . ( 2019 ) performs zero-shot multitask learning by finding a task-specific routing using a gating mechanism . This approach is somewhat similar to ours , but is less efficient in terms of the number of parameters since a ) the method in Purushwalkam et al . ( 2019 ) uses a modular network on top of a classical ResNet ( He et al . ( 2016 ) ) while our method uses only about 100 additional parameters and b ) the modules used in Purushwalkam et al . ( 2019 ) are layer-specific and not reusable while there is no such constraint in our architecture . Wu et al . ( 2018 ) ; Newell et al . ( 2019 ) ; Guo et al . ( 2020 ) perform task-specific model compression . There also exists a number of studies on the effectiveness of modular networks in the context of visual question answering . One of the early papers Andreas et al . ( 2016 ) used a natural language parser to determine the layout of the composite network consisting of predefined modules that solve different kinds of subtasks , such as find , measure , describe etc ; further works improved this approach by switching from an external parser to a fixed ( Hu et al . ( 2017b ) ) or arbitrary query structure ( Hu et al . ( 2017a ; 2018 ) ; Pahuja et al . ( 2019 ) ) . 3 COMPOSITIONAL MODELS AND MODULE MIXTURES . 3.1 COMPOSITIONAL MODELS . In conventional convolutional deep neural networks , most layers are different from each other and do not conform to a particular fixed design . Various model blocks may have different resolutions and different numbers of filters , some blocks may have or lack a residual connection and so on . This makes most model activations at different layers incompatible with each other . Recently , it has been demonstrated ( Sandler et al. , 2019 ) that high performing convolutional networks can nevertheless be composed of identical blocks . Such networks that were called isometric networks essentially iterate on the same activation space . The architecture of an isometric network includes the following core components : ( a ) input adapter that maps model input i ∈ I into the activation space Z ; ( b ) model body , a sequence of blocks all sharing the same architecture and mapping the space Z to itself ; ( c ) output adapter mapping the output of the last block into the embedding space E , and finally ( d ) logits layer for classification models , mapping embedding spaceE into the final predicted probabilities . While isometric networks achieve near SOTA performance , the fact that they iterate on the same internal activation space Z makes them a perfect candidate for studying modular architectures . Specifically , instead of treating each block as an independently trainable mapping , we may view the entire model body as a composition of modules fk each of which may appear in a sequence more than once . In the following , we will denote a set of trainable modules as F = { f1 , f2 , . . . , fK } and their composition f = fmL ◦ · · · ◦ fm1 forming a model body1 will frequently be denoted as f ( m1 , ... , mL ) with ( m1 , . . . , mL ) being called a signature of this model . 3.2 LEARNING COMPOSITIONAL MODELS . Mixture modules Compositional models discussed in Section 3.1 can be trained on a particular task for any predefined choice of the model signature . But , while the study of hand-designed modular architectures is an interesting exploration in itself , one may want to go further and identify a sequence of modules ( m1 , m2 , . . . , mL ) together with the modules F themselves that is the most advantageous choice for solving a particular problem . In some previous explorations of modular architectures , authors used a probabilistic approach where the model body ( m1 , . . . , mL ) was sampled for every input batch or sample from the optimized probability distribution p ( m1 , m2 , . . . , mL ) . For example , in Kirsch et al . ( 2018b ) , modules2 were sampled from a distribution defined by per-layer conditional probabilities p ( m ` |x ` −1 ) with x ` −1 being the output of the previous layer ( ` − 1 ) , i.e. , x ` −1 = f ( m1 , ... , m ` −1 ) ( i ) . Since discrete nature of such a choice limits the flexibility of the model and since in practice learning such a model may be complicated due to high-variance estimator noise , we instead replace a composition of blocks from F with a composition of mixture modules , each of which is essentially a smooth superposition of blocks from F . More precisely , each mixture module hξ with ξ ∈ RK has the same architecture as blocks from F , but every parameter θ in this mixture module ( be that some convolutional kernel , or a batch normalization β or γ parameter ) is a linear combination ∑K k=1 ξkθk of the corresponding parameters θk from fk ∈ F . Notice that since each block contains multiple convolutions with different nonlinearities , a mixture module obtained by linearly combining weights from F is not a linear combination of modules from F . It is true however that if ξi = δij for some fixed3 j ( which we later refer to as “ one-hot ” vectors ) , a mixture module hξ becomes identical to fj ∈ F . Due to the fact that hξ mixes parameters from F , we will 1with L is the total number of layers in the model 2here each mk is actually a set of several chosen modules the output of which is combined 3Here δij is a Kronecker delta sometimes refer to blocks from F as templates , which are mixed together to finally form an actual unit of computation . After we choose a particular task and replace the model body with a composition of mixture modules hξL ◦ · · · ◦ hξ1 , we obtain a differentiable objective that can be optimized with respect to a model signature Ξ = ( ξ1 , . . . , ξL ) and parameters of F . In our experiments , we constrained each of ξ ` by using a softmax nonlinearity , which allowed us to guarantee that ∑ i ( ξ ` ) i = 1 for every layer ` . Of course , even with this choice of the nonlinearity , we are not guaranteed that the final signature Ξ will contain only one-hot vectors ξ ` that can be interpreted as blocks from F . Regularization of mixture weights While in our experiments , mixture weights ξ ` frequently approached a set of approximately one-hot vectors4 , one can use additional regularizers to ensure that all ξ ` are one-hot and the resulting model can be represented as a composition of individual template modules . In Appendix A , we outline two choices that we experimented with : ( a ) entropy regularization that penalizes maxi ( ξ ` ) i < 1 , but does not involve template weights and ( b ) a regularizer that penalizes Ξ and also effectively attracts individual templates towards the generated mixture weights to which they contribute the most . Weight generation interpretation Consider a single block of our isometric network and let F be a function family containing single-block tensor transformations corresponding to all possible choices of block weights . All neural networks obtained by composing L such blocks without any weighttying form a function family FL . Now , instead , consider a model composed of mixture-weight blocks with templates from F . In this case , the family of all possible single-block transformations is only controlled by a vector of mixture weights ξ and narrows down to a family further denoted as FF . A composition of mixture-weights blocks is thus FLF . In a sense , choosing templates F , we introduce a small function family that can be parameterized by just a handful of real-valued parameters Ξ . This perspective shows a connection of the proposed approach to other approaches parameterizing neural network weights like , for example , HYPERNETWORKS ( Ha et al. , 2017 ) .

Summary: This paper considers “modular multi-task learning” where parameters in each layer/task are generated as a (layer/task-specific) linear (mixture) combination of a common pool of parameters. Exploring this idea, several observations are made: (1) single task isometric model performance on ImageNet can be improved, (2) Multi-task learning is supported and parameter sharing (selecting same mixture components) emerges in early layers, with specialisation emerging in later layers. With multi-domain learning, the opposite effect is achieved with domain-wise specificity arising in earlier layers, and sharing in later layers. (3) Parameter-efficient transfer learning is supported by fine-tuning the task-specific weights for new tasks. (4) Parameter-efficient domain adaptation is supported by optimising the task-specific weights for new domains.