Gumbel-Matrix Routing for Flexible Multi-task Learning

Abstract: 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%.

Reviews and author responses:

We thank the reviewer for valuable comments and suggestions. Our responses to specific points are provided below.

1) Extensiveness of experiments

While our method is compared with the SotA only on Omniglot, we also included several other experiments (MNIST, synthetic data), which were aimed at better understanding the behavior of our approach. We believe that these three lines of experiments together paint a relatively broad and convincing picture.

2) Sparsity of different routing methods

We can control the sparsity level learned by our method by using the budget penalty, similarly to how it can be controlled in the sparse Mixture-of-Experts paper (P. Ramachandran et al, ICLR 2019) by varying the value of 'k'.

For Omniglot, we do not use the budget penalty, so the learned solution is indeed not very sparse. The previous SotA based on a Mixture-of-Experts imposed a sparsity level of activating approximately 60% of connections, which is again not very sparse, although on average a little more than the solutions found by our method. However, note that this prior work did not mean to trade-off sparsity for accuracy, and neither did we; the other results reported for Omniglot are not sparse and not even based on routing. Therefore, results that we list in our paper (Table 2) include a variety of methods, none of which tried to trade-off accuracy for anything. We believe this constitutes a fair comparison.

3) Definition of different static sharing patterns

“Full sharing” essentially means the shared bottom pattern i.e., all tasks share the same bottom layers and the latter are followed by task-specific heads. In our allocation matrix view, this corresponds to setting the matrix to be all ones (i.e. every task uses every component).

“No sharing” means that the network is divided, with each task getting to train a separate set of parameters. In particular, the reviewer asked about the single-task training for MNIST: the non-sharing pattern is essentially that, since each task gets to independently train 1/4 of the network.

We thank the reviewer for the valuable comments. Our responses to specific comments are provided below.

1) Novelty of our method

We agree that the Gumbel trick and the Gumbel-Softmax routing method is not new. In this work, we propose a new method for multi-task learning and not a new routing method.

While Gumbel-based routing has been already applied to multi-task learning, we claim that our formulation (in its full form) is novel for the following reasons:
- We learn flexible parameter sharing among tasks by learning binary allocation matrices indicating how each component is allocated to each task. This is in contrast to previous works, which typically consider routing with a sequence of decisions “where to route”. We argue that our formulation is more natural for multi-task learning, as it provides an explicit way to control parameter sharing between tasks (depending on their relatedness). Right now we condition on the task id, but in the future we envision conditioning on task embeddings, which can better capture the relatedness of the tasks.
- Moreover, we also introduced ways to regularize our method such as the budget penalty (see Section 4.4) that promotes sparsity of the allocation solution.

Since the proposed method is a new method for multi-task learning (and not a new routing method), we argue that evaluating different routing solvers (such as REINFORCE or RELAX) goes beyond the scope of this work. However, it is a very interesting direction and it will definitely be the focus of our future work. As per reviewer’s suggestion, this may further improve the results.

2) Hard vs soft routing decisions

Please note that our method does use hard decisions, since we use the Straight-Through variant of the Gumbel-Softmax trick (the original Gumbel-Softmax paper introduces both the soft variant, and the Straight-Through variant). If a connection is sampled to be inactive, the corresponding component will not contribute to the output and therefore will not get gradients. It will only be used to compute the gradient for the per-connection routing probability.

3) Comparing apples to apples

We believe that the Omniglot experiment is an apples-to-apples comparison, since we re-used the same architecture that achieved the previous SotA (“Diversity and Depth in Per-Example Routing Models”, ICLR 2019). We made sure that we reproduced all the details by contacting the authors of that prior work; we also used the same regularization strategies and the same optimizer. The only difference is the routing method.

4) Scalability

It is indeed the case that due to the use of Gumbel-Softmax, the backward pass needs to activate all of the components of the model. Hence, the training phase of our method is more expensive than for other sparse baselines (such as the sparsely-gated mixture-of-experts). 
However, it is important to note that inference phase of our method is pretty scalable, since it uses hard decisions (and the budget penalty promotes even sparser solutions). Therefore, we argue that our method is practical for many multi-task learning applications.

We thank the reviewer for the valuable comments. We are open to any further suggestions for improving our paper. Our responses to specific comments are provided below.

1) Routing patterns for Omniglot

For the Omniglot experiment, it was indeed the case that discarding unwanted pooling layers was one of the clearest trends learned by our method. However, as pointed out in the paper, there were still important differences in allocation patterns corresponding to different tasks. Specifically, we grouped the tasks based on the pattern (i.e. the concatenation of all binary routing matrices), and we found around 10 groups on average, while the number of tasks was 20. Notice that differences in patterns may result in arbitrarily large differences in outputs.

2) Routing patterns for MNIST

In the case of no budget penalty, the routing commonly converged to the pattern of the following form: one pair of MNIST tasks would use all components (12 components, since there were 3 layers of 4 components each), while the other pair would use all but one component (11 components). Since MNIST and MNIST-rot are still highly related, this shows that the model preferred almost full sharing, except for dropping a single connection to allow for processing the first pair of tasks differently than the second.

With budget penalty enabled, each pair would usually use three out of four components in each layer, exactly matching the budget of 75% active connections. Note that the resulting accuracy was the same with and without the budget penalty.

3) Magnitude of gains on Omniglot over the full-sharing baseline

Even though the improvement on top of full sharing for Omniglot is not very large, full sharing is actually a pretty strong baseline; even stronger than previous SotA based on a sparse Mixture-of-Experts (P. Ramachandran et al, ICLR 2019). Our interpretation of this result is that in the case of limited data (Omniglot has very few samples per class), it is hard to learn task-specific routing without incurring an accuracy drop due to optimization difficulties. Since our routing method managed to learn task-conditioned routing and improve the accuracy, while the methods from previous works did not, we consider our Omniglot result to be a strong one.

4) Other comments

We are happy to move the result from Appendix A.3 to Section 2, if that helps the paper.

Also, the reviewer proposed considering the case of limited data and generalization. However, note that Omniglot might already be seen as such a case, and our experiments show that Gumbel-Matrix routing does produce solutions that generalize well.

In many ways this work is well presented. However, I have major concerns regarding the novelty of the proposed method and the theoretical rationale for the key design choices. Although the authors do cite and discuss (Rosenbaum et al., 2019), what is very much not clear to me is how the Gumbel-Matrix Routing proposed in this work differs from past work using the Gumbel Softmax within routing networks. It seems like past work even focused on using only the task for routing, so it is not clear to me how the approach here is really novel in comparison. Even if there is some distinction I am missing, the high level idea is clearly not that new. Additionally, there is not much theoretical discussion about what the Gumbel Softmax adds to routing networks. 

The bias/variance tradeoff of Gumbel Softmax / RELAX / REINFORCE was already highlighted in (Rosenbaum et al., 2019). Can the performance of the model on the settings tested be attributed to this tradeoff? If so, would a RELAX model perform even better? Moreover, there is not much discussion of important implications of using the Gumbel Softmax trick in the context of routing. First, as the authors acknowledge, but don't really elaborate on, using the Gumbel Softmax means we must backprop through every possible routing choice in each layer. As a result, the Gumbel approach results in a large scaling of computation with the number of modules, limiting the applicability to more ambitious settings. Moreover, while a clear motivation of this work is eliminating interference between tasks, it is not really explained how Gumbel Softmax does this and how it compares to hard routing decisions in this respect. During backprop, the computation it very similar to mixtures of experts models, and should contain more interference than hard routing. Can you explicitly show that the shape of the Gumbel distribution results in less interference between modules during learning than the standard mixtures of experts softmax approach? 

Furthermore, (Rosenbaum et al., 2019) found that a number of RL based models outperform Gumbel Softmax when routing on multi-task settings of CIFAR-100 and the Stanford Corpus of Implicatives. The authors do not provide any explanation for why this approach did not succeed in their settings. This also leads me to doubt how impressive the results presented here are as there is really not any apples to apples comparison with the same architecture and different routing decisions. In Tables 1 and 2 the best baseline is full sharing. This indicates to me that the performance difference with other cited baselines has to do with different architecture choices and not changes in the routing policy itself. The experiments can be much improved by discussing why past approaches to Gumbel based routing have failed and by thoroughly comparing to other methods for just the routing decisions with the same base architecture as done in prior work.  Unfortunately, in its current form, there is not enough context provided for the community to understand the implications of the proposed approach in the submitted draft even though it achieves good performance. 

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. 

In terms of methodology, using gumbel trick for learning routing matrix seems new to my knowledge. Although the trick has been applied to other problems and is used in a standard way. I like the idea of using this trick to make the learning of routing network unified under optimization compared to the learning in the original routing network. 

However, the experiments seem not extensive enough to demonstrate its superiority and efficiency. The method is only compared with other state of the art methods on one dataset. More experiments on various datasets and neural network architectures will be more convincing to me. I am also interested in how does the sparsity of the different routing models compare to each other? It would be unfair if some models trade performance for sparsity compared to the method proposed in this paper. Also it would be interesting to see how the learned routing matrix pattern could say something about the relatedness of different tasks.
Regarding "full sharing", is it different tasks trained together with the same network? 
And another minor question for the experiments on MNIST, what are the accuracies for single task learning using same architecture?

Overall, I find the idea of using gumbel trick for learning routing networks interesting. However, I feel the experiments are not sufficient and I would encourage the authors to conduct more experiments and comparisons.


This paper applies the Gumbel-softmax to optimizing task-specific routing in deep multi-task learning. Experiments demonstrate improvements of the method over no sharing or full sharing, and it is used to achieve s-o-t-a results in the Omniglot MTL benchmark.

Although the end results are good, and the approach is well-motivated, I am leaning to reject, because the experiments have not made clear when the method works and how it behaves. The improvements over the full-sharing baselines appear fairly small, and in the analysis it appears the model is mainly discarding unneeded pooling layers. Is there some real task-specific routing that the method is able to take advantage of? Maybe an experiment where full-sharing is detrimental, i.e., because there are some highly unrelated tasks, would help to highlight how the approach selects appropriate module subsets for each task. E.g., what are the routing patterns in Section 6.1 that are the same within each pair of MNIST tasks, but different across task pairs? Is there a way to visualize differences between routing of different Omniglot tasks?

Similarly, the experiment in Section 2 is interesting, but the conclusion that negative transfer exists is not novel. Is there a way to include the Gumbel approach in these synthetic experiments to show that it addresses this issue? E.g., something like the result in A.3 could be promoted to Section 2. More compelling synthetic datasets could be generated by the method in A.1. for the case where tasks are somewhat related, in which case we can actually see if how the sharing occurs. Could Gumbel see a bigger boost in these synthetic experiments if training data were limited and generalization was tested instead of training loss? 


Meta-review: This paper proposed to use Gumbel softmax to optimize the routing matrix in routing network for multitask learning.  All reviewers have a consensus on rejecting this paper.  The paper did not clearly explain how and why this method works, and the experiments are not sufficient.

Decision: reject