LRN: Limitless Routing Networks for Effective Multi-task Learning
Authors: Ryan Wickman
Xiaofei Zhang
Weizi Li
Venue: ICLR 2022 Submitted (2022)
Keywords: multi-task learning
MTL
reinforcement learning
machine learning
routing networks
modular networks

Abstract: Multi-task learning (MTL) is a field involved with learning multiple tasks simultaneously typically through using shared model parameters. The shared representation enables generalized parameters that are task invariant and assists in learning tasks with sparse data. However, the presence of unforeseen task interference can cause one task to improve at the detriment of another. A recent paradigm constructed to tackle these types of problems is the routing network, that builds neural network architectures from a set of modules conditioned on the input instance, task, and previous output of other modules. This approach has many constraints, so we propose the Limitless Routing Network (LRN) which removes the constraints through the usage of a transformer-based router and a reevaluation of the state and action space. We also provide a simple solution to the module collapse problem and display superior accuracy performance over several MTL benchmarks compared to the original routing network.

Under review as a conference paper at ICLR 2022
LRN: LIMITLESS ROUTING NETWORKS FOR EFFEC- TIVE MULTI-TASK LEARNING
A Anonymous authors Paper under double-blind review
ABSTRACT
Multi-task learning (MTL) aims at learning multiple tasks simultaneously typi- cally through shared model parameters. The shared representation enables gener- alized parameters that are task invariant and assists in learning tasks with sparse data. However, the presence of unforeseen task interference can cause one task to improve the detriment of another. A recent paradigm constructed to tackle these types of problems is the routing network, which builds neural network architec- tures from a set modules conditioned on the input instance, task, and previous output of other modules. This approach has many constraints, we propose the Limitless Routing Network (LRN) which removes the constraints through the us- age of transformer-based router and reevaluation of the state and action space. We also provide a simple solution to the module collapse problem and display superior accuracy performance over several MTL benchmarks compared the original routing network.
INTRODUCTION
Multi-task leamning (Caruana, 1997) is the objective of learning multiple tasks simultaneously with some shared representation, such as shared parameters a deep neural network. The shared pa- rameters can enable knowledge gained from one task to be leveraged for successfully completing another, which is called positive transfer. However, this is not always the case. As model parameters updated, not only can information be lost through catastrophic forgetting (McCloskey & Cohen, 1989; Ratcliff, 1990; Kirkpatrick et al., 2017), but can move a direction that could negatively affect the performance of another task, called negative transfer.
In (Rosenbaum et al., 2017), researchers propose the routing network which is comprised of router and set of modules, where the modules in their experiments consist of fully-connected (FC) layers. The router constructs a subset of a neural network's layers by recursively routing the output of a previously selected module through another, conditioned on the input instance, task, and output of the previous module until particular depth is reached. This allows for conditionally using differ- ent model parameters based on the input in hopes it can reduce negative transfer and catastrophic forgetting while preserving positive transfer. However, this approach has several constraints: 1) It requires the module output size to fixed as the router requires the output of a module as part of the state; 2) Each module belongs strictly to particular depth; and 3) The total number of modules at every level of depth is fixed.
As these limitations enforce constraints on the feasible region of modular architecture, we propose the Limitless Routing Network (LRN) which removes the aforementioned constraints through the usage of transformer-basel router and a reevaluation of the state and action space. We also extend routing to every layer in the constructed network and increase the diversity of modules to show the capabilities and generalizability of our approach. In addition to removing the constraints and increasing functionality, we propose a simple solution to help prevent module collapse (Kirsch et al., 2018; Rosenbaum et al., 2019), whereby the router only routes subset of modules regardless of the input, through an auxiliary reward.
Through empirical evidence given by testing on all the same MTL benchmarks discussed in (Rosen- baum et al., 2017), we show superior testing accuracy compared to the original routing network, In summary, our contributions are as follows:

---

Under review as a conference paper at ICLR 2022
Removing the constraints from the routing network to allow for more complex model con- struction;
novel approach to prevent module collapse and increase model diversity through a sim- ple, auxiliary reward;
Using function approximation for the router to enable conditioning the constructed network on the instance level information; and
Allowing routing every layer of the modular network.
The source code of LRN can be found at httos:/ /ttt:s:://nooymmuuss hwtp::/mo.uuuu........mmmom science/ TaskRouter-A624.
RELATED WORK
Traditionally work multi-task learning (Caruana, 1997) involves manually creating the shared representation and having fixed for all tasks. For example, sharing the first few layers of a neural network, but then having different heads for each task. This is similar to what is typically done in Reinforcement Learning and Actor-Critic architectures (Konda & Tsitsiklis, 2000), where the actor and critic have different output heads, one for the policy and the other for the value function.
Our work is based on the routing network (Rosenbaum et al., 2017), which uses a multi-agent rein- forcement learning (MARL) approach and the weighted policy learner (WPL) algorithm (Abdallah & Lesser, 2006), to jointly train the router and set modules. All the benchmarks in (Rosenbaum al., 2017) on are classification tasks, where the reward is set to be +1 if the class is correctly predicted and -1 otherwise. The state space is defined by the tuple (v, f, i), where is output of the previous module, timestep. Their action space is defined over all the modules in the current timestep. which is kept to the same fixed size. addition, in (Rosenbaum et al., 2017), tabular representation adopted rep- resent the policy and found to outperform the function approximation representation. In this work, we show that function approximation does work in our modified approach. more in-depth analysis of recent work multi-lask learning with deep neural networks can be found in (Crawshaw, 2020).
one-hot encoded vector representing the current task, and is the current
similar line of work is that of neural architecture search (NAS) (Zoph & Le, 2017: Baker al., 2017: Miikkulaieen et al., 2017; Liu al., 2018; Pham et al. 2018), where the goal is to search for neural network architectures that are best fit for a task. This differs from the conditional architecture construction employed in work it primarily focuses finding singular architecture for all instances rather than having conditioned on each input instance. An in-depth survey of recent work is discussed in (Ren al., 2020).
PRELIMINARIES
For 7 Ttst episodic task, an RL agent applies its policy no (a(ass)) to choose action in state s; each timestep E 11. Then, the environment responds with reward rt and next state $t41- This sequential decision-makin process is formulated the tuple (S, A, P. R), where S is the state space, A is the action space, P(s'|s a) is the transition probability function, R is the reward function, and 7 E [0, 1] is the discount factor. The objective of the RL agent is to find Va that can maximize the sum of discounted rewards R. Σi= -
Markov Decision Process (MDP) defined by
One way to obtain mo is via policy gradient methods, which learn by updating 0 using gradient ascent. Fy is typically termed actor. The policy gradient has the general form (Schulman et al., 2018):
x V log To(a
E
where V be chosen the advantage function (s, a) (F Vs) (a, is the action-value function and (s) is the state-value function. (*) can be estimated by a parameterized model called critic (Konda & Tsitsiklis. 2000).
(a,

---

Under review as a conference paper at ICLR 2022
4 LIMITLESS ROUTING
In this section, we will describe the way we removed the constraints on the routing network and increase the diversity the modules, Then, we will describe the MDP formulation and our actor- critic model architecture. Lastly, we will discuss our solution to preventing module collapse.
4.1 REMOVING THE CONSTRAINTS OF THE ORIGINAL ROUTING NETWORK
In the original routing network (Rosenbaum et al., 2017), the output size of modules is fixed and each module is required to belong to a particular depth, and thus the total number of modules is fixed at every level of depth of the network. The rationale of a fixed output size of the modules is due to the router takes in the output of a module as part of the imput at each timestep. This prevents routing over modules with different output sizes if this constraint is not kept true and thus disallows complex model construction beyond typically a few FC layers with the same fixed output size. The only solution proposed in (Rosenbaum al., 2017) circumvent fixed output size of the modules is by increasing the number of routers per set of modules with the same output size, which is inefficient.
To increase elliciency and allow for more general model construction, we remove the abovemen- tioned constraint by creating an embedding for each module and use the embedding as the part of the input rather than the output of the previous module. More specifically, the input is updated to include the embeddings for the current modules in the model being constructed, as well as the mod- ules that are available to be used in the current timestep in the state. So, the new input to the router the embedding for the image instance, the task embedding, prodel ,...... ....,, are the embeddings of the modules currently in the model I...Iealid
(v, 2, prodel I 1g goaled ..... pealed) ), where
getting constructed, and Igalid .... realid are the modules that can be used in the current timestep. The modules that can be used in a timestep based on the current depth of the model getting con- structed, the output size of the previous modules (if it is not already used), and the input size of the module.
Due to updating the input of the router, removing the constraint of requiring each module to belong to particular depth has an easy solution. We simply allow the modules to be used any depth as long the output size of the previous module, or the size of the instance if it is the first timestep, matches the input size and it has not already been included the constructed model itself to be of variable depth size, we include "PASS" action as it was used in (Rosenbaum et al., 2017), to allow the router to skip adding module at a given timestep. Note that this design is restricted to be only used on hidden layer modules and it cannot be used on the last timestep if the last module has not been added yet. This is to enforce a valid minimum size modular network, because if we always enable the PASS action, there is chance that no modules will be selected or incorrect input and/or output sizes that do not match the current task will be picked.
the model. Furthermore, to allow
To remove the constraint of requiring the total number of modules to be fixed every level of depth of the network, we use a conditional transformer model (Keskar et al., 2019; 9; Vaswani et al., 2017) for the router network. This models the conditional probability over the valid modules particular depth:
P(Ipatid U, z, mmodol
mmdell.
where control code is given by the task and instance embedding. We sample from this dis- tribution over unused, valid modules while training to encourage exploration of various modular architectures.
4.2 MDP FORMULATION
We formulate problem Markov Decision Process (MDP) where the states defined by the tuple (v,, x, pmodet ,.... pmodel), the actions are the potential next modules (Zgald..... p the rewards the negative loss of the constructed model, and the transition probability function is deterministic as it is only based on the routers decisions. The routing network was trained using Proximal Policy Optimization (PPO) (Schulman al., 2017) using an actor-critic approach with
pralid)

---

Under review as a conference paper at ICLR 2022
Image
Image Encoder
Image Embedding
Trinformer Encoder
Allemtion
PC
Thok ID
Tisk Embealding
Provicundy
Midble Embholiings
Conditional
Unased Valid Mislales
Figure 1: The LRN architecture using actor-critic approach. The state is represented by the image instance, task, and previously selected modules and the actions are over the valid, unused modules.
generalized advantage estimation (GAE) (Schulman et al., 2018) for the advantage estimate. Since critic has variable input size that depends on timestep, we make use of a simple attention mechanism (Bahdanau et al., 2016) without the context to combine the embeddings representing the state to acquire fixed sized state embedding.
Specifically, we first use transformer encoder (Vaswani et al., 2017) to perform self-attention over the state:
& Encoder(ss).
(2)
Then, scores computed over the elements of the state:
ej We
where Wa a vector weights used to map the embeddings representing the state a single scalar weight, ej, and &t is of the embeddings for the state (i.e., the instance embedding, task embedding. and the embeddings modules currently the network).
Next, the scores normalized into weights using softmax distribution:
exp(e.)
0j - exple,)
(4)
Finally, these weights are used to combine the embeddings and output the state-value for the current timestep t;
VTA (T) We Σ و
where Wo is vector of weights used to map the weighted combination of state embeddings to state-value prediction. Both W, and W. may be represented by FC layer with a single unit.
The full LRN actor-critic router architecture is given in Figure 1. We used an image encoder that comprised of a couple convolutional layers followed by a FC layer to convert instance image into embedding that has the same size the task and module embeddings.
4.3 PREVENTING MODULE COLLAPSE
Model collapse (Kirsch al., 2018; Rosenbaum et al., 2019) occurs when the routing network lacks diversity in its selection of modules and collapses to only choosing a fixed subset of modules. There have been many proposed solutions, but all assume the original routing method. Thus, formulate a novel and simple approach as reward bonus to prevent module collapse-the reward is given to each module selection proportional to normalized count over batch:

---

Under review as a conference paper at ICLR 2022
bonus
Cj Cyain
Cmaz
Cmên
where poonss is the reward for module 4i is the number of times that 4 has been selected over the entire batch of episodes, and Gmin and Cmaz are the min and max count over all the modules, respectively. Everytime the module is used in batch of instances, the corresponding timestep's reward is augmented with the same ri boonns reward.
The auxiliary reward bonus gives higher reward to modules that not used frequently over batch, and low or near near-zero reward to those that are used the most. This encourages the usage of modules that were not used much in batch and helps prevent over constructing the current best model for each instance. Thus, this reward shaping helps favor diversity of module selection, and therefore helps prevent module collapse. In summary, the final reward for each timestep and example in batch is as follows:
Vyk: Tk -T bonus
where Ttk is the reward for timestep in example k in the current batch, re is the negative loss of the h constructed model, and r.ooius the reward bonus for module (,
EXPERIMENTS AND RESULTS
In this section, we will first explain the benchmarks and experiment setup details we used to evaluate our approach. Next, we will describe the modules we used in all the experiments. Then, we will discuss our hyperparameter tuning results for the maximum depth of the constructed model and the number of duplicate modules and results of our solution to preventing module collapse. Lastly, will compare our results to the original routing network against all the benchmarks.
5.1 EXPERIMENT SETUP
We test over the three benehmarks discussed in (Rosenbaum et al., 2017), namely MNIST-MTL. MIN-MTL, and CIFAR-MTL.
MNIST-MTL uses the MNIST dataset (Lecun et al., 1998), which consists of images of hand wwritwriteen digits. The task is usually converted to a set of one-vs-all binary classification task for each digit (10 tasks in total). The training set consists of 1 000 positive examples for every task and 1 000 negative examples from every other classes, resulting in a total number of 9 000 negative examples. The test set is created by randomly sampling 200 examples per task from a different set of examples that were not present in the training set.
MIN-MTL is derived from Mini-ImageN ImageNeet (Ravi & Larochelle, 2016; Vinyals et al., 2017), which is a subset ImageNet (Deng et al., 2009). We use the first 50 labels of ImageNet and group them into 10 tasks with labels each. We use 800 examples per task for the training set and 50 per task for the test set.
CIFAR-MTL uses the 100 image classification dataset (Krizhevsky et al., 2009). We create one task for each of the 20 coarse labels and include 500 examples for each of the 5 corresponding fine labels. So, in total, 2 500 examples are created per task for the training set and 500 per task for the test set.
In all the experiments, train for total of 100 epochs the training set with batch size of 128, a mini-bated size of 64, and mini-batch epoch for PPO. We compute the accuracy over the test set after every epoch. The transformer uses singular encoder and decoder with 4 heads, and we set dmmadet = 256 and daff = 512 as in the transformer model (Vaswani et al., 2017). All the of size 256. We 7 0.99 and AçAE 0.95 for
embeddings and hidden units of FC layers the return and advantage computation. We set the initial learning rate for both training the modules and router 0.0001 and anneal them by a factor of 0.97 after every epoch.
5

---

Under review as a conference paper at ICLR 2022
5.2 MODULE DESIGN
As LRN removes the constraints of the original routing network, our router can perform routing on every layer the neural network. Thus, we came up with set of modules that can be used to construct a fully routed modular network. There are mainly three different lypes of modules; those that can be used in the first layer, hidden layers, and the output layer. We assume every module made up of single FC convolutional layer. Furthermore, we come up with two hyperparameters that control the maximum depth of the constructed network and the number of duplicates allowed for each module described below.
The input module is a convolutional layer with input channel size matched to the image size. Every convolutional layer, at every depth, contains 32 3x3 filters with stride and padding set to 1. This is to keep the input size equal to the output size to constrain the number of possible models. However, this is not required and can be of any variable size, as long as the design is able to be placed in the constructed network (i.e., its input size matches the output of another module and vice versa). If not, the module could never be used.
The hidden layer modules consist of a convolutional layer and 2 FC layers with 256 hidden units. The first FC layer has input size equal to the oulput size of the convolutional layers, and the second with input size 256 to allow being connected to the other hidden FC layers.
The output layer modules, which are used as the head of the network, are FC layers. The are iden- tical to the 2 FC layers used as hidden layers, as to allow being connected to either convolutional or FC layer, but with output size equal to the number of classes.
We apply layer normalization (Ba et al., 2016) after every FC module, besides the output module, and instance normalization (Ulyanov et al., 2017) for convolutional modules. Furthermore, they are both followed by ReLUU activation (Nair & Hinton, 2010).
5.3 MAXIMUM DEPTH AND DUPLICATES ANALYSIS
As removed the constraint on the number of modules that can be used and what layers routing can be performed, we decide to explore what the best maximum depth of the constructed network and the number of duplicates to include for each module discussed previously. We experiment over various combinations of module depth and duplicates for all the benchmarks. The results are shown in Figures 2, and 4, where we use the naming scheme "dup_x ddepthh"" denote duplicates and y max depth.
The figures on the left give the accuracy over the test after each epoch, and the right the number of unused modules. If we jointly increase the maximum depth and the number of modules, the performance will improve, with the best appearing at 10 duplicates per module and the maximum depth of 6. However, increasing the number of duplicates generally slows down the convergence speed. This is because with more modules, the state and action space increases, causing the task complexity to increase as well. Also, if module is not used in a batch of example, its parameters will not be updated after that batch.
5.4 PREVENTING MODULE COLLAPSE ANALYSIS
For our method to prevent module collapse, we compare to the case when the reward is not used, We test on the "'up___.depth_"" model as to keep the space of possible modular networks small, but still show the effectiveness of the auxiliary reward bonus. The results are given in Figures 5, 6, and 7 where append " "_no_bonus" to the name of the ones that are not using module collapse prevention.
In the MNIST-MTL and CIFAR-MTL benchmarks, the results show clear evidence of module col- lapse as the number of unused modules never decreases when the bonus is not used. This also results in the test accuracy increasing by 16% and 28% for MIN-MTL and CIFAR-MTL, respectivll Also, in all the benchmarks, when the bonus reward is used, LRN uses more modules over the test set, thus succeeding in usage of modular diversity.
Although the results are close, MNIST-MTL. performs slighter better (about 1%) when the reward bonus for preventing module collapse is not used. We believe this is because MNIST-MTL repre-
6

---

Under review as a conference paper at ICLR 2022
dup_5 5_depth_4
0.98
dup, S. dapaphhh 6
Unused Medules
dup, _d_pthhh4
dup 10 depth 6 W
Asturacy
20
0.94
0.92
0.90
60
100
100
Figure 2: MNIST results varying duplicates and maximum depth.
50
dup. 5. dedeephhh4 44
depth
Unused Modules
to
dup, 10 depth
dup. de.eepphh.
Accuracy
0.25
0.20
20
40
60
80
100
20
40
80
80
100
Epoch
Epoch
Figure 3: MIN-MTL results on varying duplicates and maximum depth.
50
Test Accuracy
Num unused FBs
40
0.5
30
dup depth. 4
20
dup_5_ d_peepth_@
dup 10, ded.epthhh_
10
dup. deppeehhh
20
40
60
80
100
20
100
Epoch
Epoch
Figure 4: CIFAR results on varying duplicates and maximum depth.
sents easier problem and all sub-tasks of MNIST-MTL are very similar, the positive transfer between tasks is high, while catastrophic forgetting and negative transfer is low. Therefore, training the routing network tends be easier. Regardless, the results clearly show that the auxiliary reward bonus help prevent module collapse.
5.5 COMPARISON TO THE ORIGINAL ROUTING NETWORK
We follow the same procedures as discussed in (Roscnbaum et al., 2017) to recreate the experiment results the routing network. We use LRN with 10 duplicate modules and a maximum depth of 6 as we found that resulted in the best performance on all the benchmarks. The results against the original routing network given in Table 1. Our approach outperforms their approach measured over all the benchmarks. Note that fully route every layer in the modular network, while (Rosen- baum al., 2017) only routes the last layers while keeping all the initial convolutional layers fixed, So, this shows support for fully routing the network. Furthermore, we use function approximation as opposed to a tabular representation for the router policy. In (Rosenbaum et al., 2017). function results show that rectify this issue by
approximation has inferior performance. However, utilizing instance level information which can only be used by function approximation (e.g.. is trivial solution build tabular representation for pixel-based data). As function approxi- mation much more versatile than tabular representation, we expect our approach to have broad impact on multi-task learning.

---

Under review as a conference paper at ICLR 2022
dup, depphhhh__
0.98
5_o___h 4_ bonus
Unused Modules
20
Accuracy
0.96
15
10
0.92
100
100
Epoch
Epoch
Figure 5: MNIST-MTL results for preventing module collapse.
depth dup dopth no bonus
0.45
Unused Modules
Accuracy
0.40
15
0.30
0.20
60
100
80
100
Figure 6: MIN-MTL results for preventing module collapse.
5.p____eehhh__
dup. 5, depth 4. no. bonus
Unused Modules
20
Accuraty
0.5
15
mmra
Epoch
Epoch
Figure 7: CIFAR-MTL results for preventing module collapse.
Table 1: Performance comparison to the original routing network. As our experiments suggest, LRN outperforms the original routing network on all the benchmarks.
Test Accuracy
Method
MNIST-MTL
CIFAR-MTL
MIN-MTL
Routing Network
99.0%
60.0% 70.0%
35.9% 53.0 %
Limitless Routing Network (ours)
99.5%
6 CONCLUSION
In this work, we propose Limitless Routing Network (LRN) for multi-task learning. LRN removes the constraints required by the original routing network (Rosenbaum et al., 2017) in terms of fixed output size of modules, and the requirement each module belonging to a particular depth in the network. As result, LRN allows more complex model construction and outperforms the original routing network on all the benchmarks. In addition, we provide simple reward bonus to help circumvent module collapse. We also allow the of function approximation (instead of tabular representation) in the routing network, which greatly increases the potential of the routing network on complex tasks. We have conducted extensive experiments and analysis of LRN hyperparameters

---

Under review as a conference paper at ICLR 2022
demonstrate the effectiveness of our approach. In the future, would like provide more in- depth experimentation on various multi-task learning domains to further show the generality of our approach. Furthermore, that can each be composed of multiple layer provide unique forms of perception (c.g., using complex modules made up of convolutional neural networks for visual-related tasks or transformers for language-rehee relaaae tasks). Lastly, we would also like to apply LRN to multi-modal multi-learning tasks where LRN can have an array modules for each modality.
inlend to extend the Iramework use more complex modules
REFERENCES
Sherief Abdallah and Victor Lesser. Learning the task allocation game. In Proceedings of the fifth international joint conference on Autonomous agents and multiagent systems, pp. 850-857, 2006.
Jimmy Lei Ba, Jamic Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2016.
Bowen Baker, Otkrist Gupta, Nikhil Naik, and Ramesh Raskar. Designing neural network architec- tures using reinforcement learning, 2017.
Rich Caruana. Multitask learning. Machine learning, 28(1):41-75, 1997.
Michael Crawshaw. Multi-task learning with deep neural networks: A survey. abs/2009.09796, 2020. URL :/appiii////vvi.////////a/a//mm/
CoRR.
/ org/ abs 200//009.009990 09796.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEFE conference on computer vision and pattern recognition, pp. 248-255. Iece, 2009.
Nitish Shirish Keskar, Bryan McCann, Lav R. Varshney, Caiming Xiong, and Richard Socher. CTRL: A conditional transformer language model for controllable generation. CoRR, 9099 05858.
0bs/1999.05858, 2019. URL http://r/arr
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Baabsk-BBrrwinsk. Demis Hass- abis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks, 2017.
Louis Kirsch, Julius Kunze, and David Barber. Modular networks: Learning to decompose neural computation. CoRR, 0tts/8111002244 2018. URL http:: /ttp:///avvvvvv/vv org/ abs/
Vijay Konda and John N Tsitsiklis. Actor-critic algorithms. In Advances in neural information processing systems, pp. 1008-1014, Citeseer, 2000.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers features from tiny images. 2009.
Lecun, L. Bottou, Y. Bengio, and Haffner. Gradient-based learning applied to document recog- nition. Proceedings of the IEFF, 86(11) 96(11):2278-2224 1998. doi: 10. 1091//7766991
Chenxi Liu, Barret Zoph, Maxim Neumann. Jonathon Shlens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan Yuille, Jonathan Huang, and Kevin Murphy. Progressive neural architecture search, 2018.
Michael McCloskey and Neal Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165. Elsevier, 1989.
Risto Miikkulainen, Jason Zhi Liang, Elliot Meyerson, Aditya Rawal, Daniel Fink, Olivier Francon, Bala Raju, Hormoz Shahrzad, Arshak Navruzyan, Nigel Duffy, and Babak Hodjat. Evolving deep neural networks. CoRR, abs/1 abs/1103300048 2017. URL http: / rxiv ort. 1 1703. 00548.

---

Under review as a conference paper at ICLR 2022
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on International Conference on Machine Learning, ICML, pp. 807-814. Omnipress, 2010.
Hieu Pham, Melody Y. Guan, Barret Zoph, Quoc V. Le, and Jeff Dean, Efficient neural architecture search via parameter sharing, 2018.
Roger Ratcliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97(2))285, 1990.
Sachin Ravi and Hugo I Larochelle. Optimization as a model for few-shot learning. 2016.
Pengzhen Ren, Yun Xiao, Xiaojun Chang, Po-Yao Huang, Zhihui Li, Xiaojiang Chen, and Xin Wang. comprehensive survey of neural architecture search: Challenges and solutions. CoRR. abs/2006.02903, 2020. URL ht.pss: aarivv org/ abs 903.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning, 2017.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation, 2018.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Koray Kavukcuoglu, and Daan Wierstra. Match- ing networks for one shot learning, 2017.
Barret Zoph and Quoe Le. Neural architecture search with reinforcement learning, 2017.
10


Reviews:
[{"review_id": "t1bCeXBqcni", "reviewer": "Reviewer_CZQn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Although the paper identifies 3 constraints with the original routing network, it is not clear why these constraints are limiting. For example, one of the constraints is that modules could only be used at a specific depth. However, this seems to be a reasonable assumption. Why would a network module that is trained to process low-level features be used deeper in the inference path (ie to process higher-level features)? The proposed approach does not enforce this explicitly, but the authors could add an analysis to show what is the preferential position (in the inference graph) for each block.\n\nThe results for routing networks for MIN-MTL in table 1 (35.9%) seems to be inconsistent with the original paper (which presents an accuracy of ~57% - original paper figure 7.) What is the justification for this difference?\n\nThis paper does not compare with any other models, other than routing networks. Multi-task learning is a widely studied field, and thus many different approaches have been proposed to address this problem. A more extensive comparison to prior work is warranted.", "summary_of_the_paper": "The paper proposes an improvement to routing networks (Rosenbaum 2017): a multi-task learning model that builds the neural network architecture from set of modules, which can be adaptively chosen conditioned on the input sample and target task. This paper identifies several constraints with the original routing networks, namely 1) each module output dimension was constant, 2) modules were depth specific, and 3) the total number of modules available per level was fixed. The paper then proposes a model that lifts these constraints. It also introduces a bonus reward to encourage all modules to be equally used.", "summary_of_the_review": "The paper presents an improvement to a previous multi-task learning model called routing networks. The proposed improvements are not well motivated. The experimental validation is also limited given the lack of comparisons to prior work. Thus due to the limited novelty and poor experimental validation, I leaning towards rejection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "g5rBKyVAeN6", "reviewer": "Reviewer_ewVv", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work improves the existing routing network with several constraints by proposing a limitless routing network that removes the constraints using an actor-critic approach. However, it is incremental because the main structure is based on the original routing network, and on top of it, they add more embeddings not to use the previous module with a fixed output size. From this, I think the contributions of this work are rather limited. The additional suggestion to prevent module collapse is straightforward. Rigorous analyses with other works sharing a similar idea/motivation are missing as well.\n\nThe writing in Section 4 is rather hard to follow as it does not describe the meaning of variables and indices when they are first introduced, for example, l_0^model, 0_l^valid, t, and n. I could roughly understand what they are by looking at Figure 1. \nWhat does \\hat{s_{tj}} in Eq. (3) mean and how it is different from \\hat{s_t} in Eq. (2)?\nI feel there is a gap between the description and the figure for the proposed method. Descriptions on how the transformer encoder and attention are used and when we use the conditional encoder are missing in the main text.\nBefore Eq. (7), what does ‘example in a batch’ mean in the context? In Eq. (7), why do you put only the reward bonus for the module l_i to the k-th model?\nSince there are many unclear expressions, the authors should improve the paper.\n\nThe presented experiments are not effective because this paper contains two analyses for the proposed method itself and one comparison to the baseline without other approaches. In addition, the experiments are all on small-scale datasets and do not contain many tasks. From those, it is hard to judge whether the proposal is promising or not. More thorough experiments should be conducted to make the paper stronger.\n", "summary_of_the_paper": "This paper presents an improved routing network for multi-task learning by removing constraints that are applied in the existing routing network. This is achieved by introducing an actor-critic approach equipped with two different encoders. It also presents a simple approach to avoid module collapse based on a reward for each module. Experimental results show that the proposed approach performs better than its baseline approach on several MTL benchmark datasets.", "summary_of_the_review": "I have some concerns on this paper; the novelty of the proposed method and a limited set of experiments. The authors should address the concerns clearly and compare the proposal with other similar works.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JCFMRzbN8i", "reviewer": "Reviewer_v8R3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and easy to follow. They propose to route the instance among modules. And modules are not limited to a certain depth and can be connected in different orders, which is called as limitless routing network. The idea is novel and interesting. Their use of actor-critic RL algorithm to optimize the router is also a solid choice. They also introduce a reward to prevent the mode collapse.\n\nHowever, the experiments are not adequate to substantiate their claim. First, the datasets used in the paper are pretty simple, e.g. CIFAR and MNIST. Second, they do not compare to many SOTA MTL methods, such as residual adapter [A], cross-stitch [B], ASTMT [C], MTAN [D] and etc to show their performance. Also, they miss the discussion of computation cost and memory efficiency in the experiment, which are important metrics to evaluate the MTL algorithms.\n\n[A] Learning multiple visual domains with residual adapters\n[B] Cross-stitch Networks for Multi-task Learning\n[C] Attentive Single-Tasking of Multiple Tasks\n[D] End-to-End Multi-Task Learning with Attention\n", "summary_of_the_paper": "The paper proposes a limitless routing network to route the instance in a certain task through a couple of modules in the multi-task learning scheme. They adopt a policy network to generate the probability of all available modules for the next step based on instance embedding. task id and previous modules. They optimize the policy network with the actor-critic RL algorithm. They evaluate their approach on MNIST-MTL, Min-MTL and CIFAR-MTL.", "summary_of_the_review": "The proposed method looks novel and reasonable but the experiments cannot well support the effectiveness of the proposed method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "3kYgbOPtyJ_", "reviewer": "Reviewer_Xh12", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n1. This paper introduced a novel idea of using a transformer in the router RL for MTL which solves the input size limitation of the traditional routing network.\n2. The routing module is agnostic to the input type so it can be applied to different types of networks. This can enlighten the application of transformers in RL. \n\nWeakness\n1. The contributions of this paper is sort of limited since the routing network is not the SOTA architecture for MTL. The routing network is designed for limited classification problems and is hard to be scaled to the advanced MTL problems including detection, segmentation, depth estimation, etc.\n2. The paper is not well written. The paper does not introduce the challenges of the routing network clearly. Given this paper is to improve the baseline routing network, it is necessary to revisit the design and challenges of the routing network in detail.\n3. Some parts of the paper are not clearly presented. For example, the conditional encoder in Figure 1 is not introduced. How to apply the routing module to all the layers is not introduced in the implementation.\n4. Although the proposed approach outperforms the routing network baseline, the paper lacks the performance comparison with the SOTA MTL work. It is hard to verify the value of using the proposed algorithm.\n5. The paper missed the model’s effectiveness evaluation of the proposed approach. Applying the transformer based routing module can make the model size very big. The model can suffer from overfitting issues which need further ablation studies.", "summary_of_the_paper": "This paper proposed a transformer based routing module to improve the routing network for multitask learning. The main contributions of this paper can be summarized as follows.\n1. The paper introduced transformer based routing to solve the input length limitation in the previous routing network.\n2. The transformer based routing network can be generally applied to different types of layers. \n3. This paper resolved the module collapse by only route to a subset of modules regardless of the input.", "summary_of_the_review": "This paper proposed a novel idea for using the transformer in the routing network. However, the paper missed many important evaluations including comparing it with SOTA MTL approaches, model effectiveness, etc. Besides, the paper was not well presented which makes it hard to follow. Overall, the paper needs a major improvement.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed a transformer based routing network which removes the constraints in the original routing network such as the depth of a network. Multi-Task learning (MTL) based on routing has been an interesting topic in the deep learning research community.  Our reviewers have serious concerns on the experiments. The presented empirical results do not seem to be able to sufficiently support the claims in this paper. Comparing with SOTA MTL methods is needed to make the proposed method convincing.
Decision: Reject