Gumbel-Matrix Routing for Flexible Multi-task Learning

This paper proposes a novel per-task routing method for multi-task applications . Multi-task neural networks can learn to transfer knowledge across different tasks by using parameter sharing . However , sharing parameters between unrelated tasks can hurt performance . To address this issue , routing networks can be applied to learn to share each group of parameters with a different subset of tasks to better leverage tasks relatedness . However , this use of routing methods requires to address the challenge of learning the routing jointly with the parameters of a modular multi-task neural network . We propose the Gumbel-Matrix routing , a novel multi-task routing method based on the Gumbel-Softmax , that is designed to learn fine-grained parameter sharing . When applied to the Omniglot benchmark , the proposed method improves the state-of-the-art error rate by 17 % . 1 INTRODUCTION . Multi-task learning ( Caruana , 1998 ; 1993 ) based on neural networks has attracted lots of research interest in the past years and has been successfully applied to several application domains , such as recommender systems ( Bansal et al. , 2016 ) and real-time object detection ( Girshick , 2015 ) . For instance , a movie recommendation system may optimize not only the likelihood of the user clicking on a suggested movie , but also the likelihood that the user is going to watch it . The most common architecture used in practice for multi-task learning is the so-called shared bottom , where the tasks share parameters in the early layers of the model , which are followed by task-specific heads . However , as our experiments on synthetic data show , when the tasks are unrelated , parameter sharing may actually hurt individual tasks performance . Therefore , resorting to flexible parameter sharing becomes very important . This can be achieved by manually trying several different static sharing patterns . However , this option is not scalable , since it requires significant effort . Instead , an approach where the sharing pattern is learned and adapted to the task relatedness , is preferable . At the same time , routing networks ( Rosenbaum et al. , 2018 ) have been introduced as powerful models , which route each input sample through its own path , selectively activating only parts of the network . They have shown strong performance in various settings thanks to their high flexibility . Routing networks lend themselves as a natural choice for learning sharing patterns in multi-task modeling . However , they are typically rather hard to train in practice ( see e.g. , ( Rosenbaum et al. , 2019 ) for the challenges of routing networks and references there in ) . In this work , we introduce a novel method , that learns the sharing pattern jointly with the model parameters using standard back-propagation . Assuming that the network consists of several layers , where each layer consists of several components , the core idea of the proposed method is to learn , for each component , a set of binary allocation variables indicating which tasks use this component . We rely on the Gumbel-softmax reparameterization method ( Jang et al. , 2017 ) in order to train these binary variables jointly with the parameters of the components . We provide experiments on synthetic data as well as real-world data showing that the proposed method can adapt the sharing pattern to the task relatedness , outperforming strong baselines and previous state-of-the-art methods . In summary , the paper contributions are the following : • We analyze positive and negative transfer effects on synthetic scenarios . This analysis motivates the need for task dependent parameter sharing in multi-task networks . • We propose a novel routing method based on Gumbel binary variables , that allows for learning flexible parameter sharing which adapts to the task relatedness , and can be optimized with standard back-propagation . The source code implementing the proposed method and the benchmarks described in this paper is publicly available at http : //github.com/anonymous . 2 POSITIVE AND NEGATIVE TRANSFER BETWEEN TASKS . We start with a practical example showing that besides positive transfer , negative transfer may as well occur in practice . That is , when the tasks are unrelated , allowing them to interact in a bigger model instead of training them separately harms the model performance . To show that both positive and negative transfer occurs , we generate two synthetic tasks , where the task relatedness ρ can be explicitly controlled . Our synthetic data generation process is based on that of ( Ma et al. , 2018 ) , and we describe it in detail in Appendix A.1 . We consider two edge cases : two unrelated tasks ( ρ = 0 ) , and two tasks that are the same up to noise ( ρ = 1 ) . We create a simple multi-task network . This network architecture consists of 4 parallel components , and each component contains a stack of fully connected layers . Each input example can be provided as input to any subset of the 4 parallel components ; the outputs of the components are averaged before being passed to a task-specific linear head . We chose this network to have low enough capacity , so that there is visible competition between tasks . For more information about the architecture , refer to Appendix A.2 . For this experiment , we use two hard-coded sharing patterns . The ‘ full sharing ’ pattern means that both tasks use all components , while ‘ no sharing ’ means that tasks use disjoint halves of all components . Therefore , in the ‘ no sharing ’ pattern , the tasks are completely independent . Note that regardless of the pattern , the total amount of parameters in the model remains the same ; the only difference is in which parameters get used by which tasks . In other words , ‘ no sharing ‘ corresponds to imposing a constraint that the network should be evenly divided between the tasks , while ‘ full sharing ’ leaves that to the optimization algorithm to decide . We run four experiments : one for every combination of sharing pattern ( ‘ full sharing ’ and ‘ no sharing ’ ) , and task relatedness ( ρ ∈ { 0 , 1 } ) . For each experiment , we report the L2 loss over time , averaged over the two tasks . The results are shown in Figure 1 . Since for ‘ no sharing ’ there is no interaction between the tasks , the average loss behaves in the same way irrespective of the task relatedness . For ‘ full sharing ’ , both tasks are allowed to update all parameters , and we see that while it improves performance if the tasks are related , it actually hurts for two completely unrelated tasks . Motivated by this example , we argue that general multi-task models should be able to learn flexible sharing patterns that are able to adapt to the task relatedness . In Section 4 we introduce a framework which can flexibly learn parameter sharing patterns , including but not limited to ‘ no sharing ’ and ‘ full sharing ’ . Since our method is based on routing networks , we provide some background about them first . 3 ROUTING NETWORKS . Standard neural networks process every input example in the same way . Routing networks provide a more flexible alternative , where every input is passed through a subgraph of the entire model ( also referred to as the supernetwork ) . Typically , the model is divided into layers , and some layers contain multiple parallel components ( also called modules or experts ) . Inputs to the large network are routed through some sets of components in subsequent layers . This general framework was used by multiple recent works ( Fernando et al. , 2017 ; Rosenbaum et al. , 2018 ; Ramachandran & Le , 2019 ) , although with some small differences or additional constraints depending on the routing method . In this paper , we use a relatively general view depicted in Figure 2 , which is relevant to both singletask and multi-task setups . Note that we do not impose any constraints on the set of components used for a specific input ; in particular , the number of components used can vary between layers and between individual inputs . After passing a sample through a single routed layer , the outputs of all activated components are aggregated to form an input to the next layer . In all experiments in this paper we simply use the average as the aggregation . 4 GUMBEL-MATRIX ROUTING FRAMEWORK . We design our framework so that it is able to learn to select task-conditioned subgraphs of a larger computational graph . In Figure 3 , we show an example of a routing model , where the routing is conditioned on the task only . In that example , there are two tasks , and two routed layers . Each task selects some subset of components within every layer . For every routed layer , we encode the component assignment as a binary matrix . Concretely , if there are T tasks , and C components in a given layer , the matrix has shape T x C , where the entry in the i-th row and j-th column is 1 if the i-th task uses the j-th component . We show these matrices at the bottom of Figure 3 . Our goal is to learn the routing matrices in the way that maximizes the average per-task performance of the model . In order to learn the routing matrices we condition on the task id , which implies that all samples from the same task will go through the same path in the network . We refer to our framework as the Gumbel-matrix routing framework . While at training time our method samples many different routing matrices , at the end a single matrix is selected per layer . At inference time , the matrices are fixed , and thus our routing does not add any overhead over the underlying non-routed model . 4.1 TRAINING . For each layer , we want to maintain a probability distribution over all possible binary routing matrices . To that end , we assume this distribution to be factorized , and we explicitly maintain a matrix of per-connection probabilities . Each probability is represented as a pair of two complementary logits . To perform a forward pass , we sample all binary connections independently according to the probability matrix . In principle , it is possible that for a given task , all connections are sampled to 0 . In that case , the output of the routed layer would be a zero vector , independently of the input . In practice , we found that this happens very rarely , and mostly at the beginning of training , since usually one of the connection probabilities quickly becomes close to 1 . Therefore , we did not try to devise a method which would artificially prevent an all-zeros connection pattern from being sampled . To initialize our method , we have to set the connection probabilities to some initial values . In principle , it is possible to introduce prior knowledge , and set these probabilities in a way that encourages or discourages certain patterns . However , here we consider the most general approach , where all connection probabitlies are initialized to the same constant value pinit . Setting pinit = 0.5 gives the highest routing entropy , and corresponds to the weakest prior , therefore we considered it as a default choice . However , in our experiments , we also found that for routing in large and deep networks , it is beneficial to set pinit closer to 1 , in order to enhance the trainability of the components and to stabilize the initial learning phases . In the backwards pass , only the components that are activated will get gradients , as the inactive components do not contribute to the final output of the network . However , in order to get a gradient for the connection probabilities , we would have to backpropagate through sampling . In the next section , we describe a method which we use to accomplish that .

Summary: The paper proposes to learn the routing matrix in routing networks for multi-task learning (MTL) using the gumbel softmax trick for binary random variables. It makes the model amenable for training the network and the routing matrix simultaneously, which is a relatively easier and unified training procedure compared to the original routing networks. The gumbel softmax trick technique is pretty standard. The proposed method is evaluated on two MTL datasets with comparisons to baselines on one of them.