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 presents a method of composing stacked neural network blocks by linearly combining module "template" weights.  The work extends "isometric networks", in which each block in the stack has the same operational structure, by parameterizing each block using a mixture of the weights from a bank of K blocks ("templates").  In multitask learning experiments, the mixtures naturally learn to share common components between tasks, while learning task-specific components where needed.  Further experiments describe behavior of the system applied to transfer learning and domain adaptation scenarios.