Branched Multi-Task Networks: Deciding What Layers To Share

1 INTRODUCTION . Deep neural networks are usually trained to tackle different tasks in isolation . Humans , in contrast , are remarkably good at solving a multitude of tasks concurrently . Biological data processing appears to follow a multi-tasking strategy too ; instead of separating tasks and solving them in isolation , different processes seem to share the same early processing layers in the brain – see e.g . V1 in macaques ( Gur & Snodderly , 2007 ) . Drawing inspiration from such observations , deep learning researchers began to develop multi-task networks with branched architectures . As a whole , multi-task networks ( Caruana , 1997 ) seek to improve generalization and processing efficiency through the joint learning of related tasks . Compared to the typical learning of separate deep neural networks for each of the individual tasks , multi-task networks come with several advantages . First , due to their inherent layer sharing ( Kokkinos , 2017 ; Lu et al. , 2017 ; Kendall et al. , 2018 ; Guo et al. , 2018 ; Liu et al. , 2019 ) , the resulting memory footprint is typically substantially lower . Second , as features in the shared layers do not need to be calculated repeatedly for the different tasks , the overall inference speed is often higher ( Neven et al. , 2017 ; Lu et al. , 2017 ) . Finally , multi-task networks may outperform their single-task counterparts ( Kendall et al. , 2018 ; Xu et al. , 2018 ; Sener & Koltun , 2018 ; Maninis et al. , 2019 ) . Evidently , there is merit in utilizing multi-task networks . When it comes to designing them , however , a significant challenge is to decide on the layers that need to be shared among tasks . Assuming a hard parameter sharing setting1 , the number of possible network configurations grows quickly with the number of tasks . As a result , a trial-and-error procedure to define the optimal architecture becomes unwieldy . Resorting to neural architecture 1In this setting , the input is first encoded through a stack of shared layers , after which tasks branch out into their own sequence of task-specific layers ( Kokkinos , 2017 ; Lu et al. , 2017 ; Kendall et al. , 2018 ; Guo et al. , 2018 ; Sener & Koltun , 2018 ) . Alternatively , a set of task-specific networks can be used in conjunction with a feature sharing mechanism ( Misra et al. , 2016 ; Liu et al. , 2019 ; Ruder et al. , 2019 ) . The latter approach is termed soft parameter sharing in the literature . search ( Elsken et al. , 2019 ) techniques is not a viable option too , as in this case , the layer sharing has to be jointly optimized with the layers types , their connectivity , etc. , rendering the problem considerably expensive . Instead , researchers have recently explored more viable alternatives , like routing ( Rosenbaum et al. , 2018 ) , stochastic filter grouping ( Bragman et al. , 2019 ) , and feature partitioning ( Newell et al. , 2019 ) , which are , however , closer to the soft parameter sharing setting . Previous works on hard parameter sharing opted for the simple strategy of sharing the initial layers in the network , after which all tasks branch out simultaneously . The point at which the branching occurs is usually determined ad hoc ( Kendall et al. , 2018 ; Guo et al. , 2018 ; Sener & Koltun , 2018 ) . This situation hurts performance , as a suboptimal grouping of tasks can lead to the sharing of information between unrelated tasks , known as negative transfer ( Zhao et al. , 2018 ) . In this paper , we go beyond the aforementioned limitations and propose a novel approach to decide on the degree of layer sharing between tasks in order to eliminate the need for manual exploration . To this end , we base the layer sharing on measurable levels of task affinity or task relatedness : two tasks are strongly related , if their single task models rely on a similar set of features . Zamir et al . ( 2018 ) quantified this property by measuring the performance when solving a task using a variable sets of layers from a model pretrained on a different task . However , their approach is considerably expensive , as it scales quadratically with the number of tasks . Recently , Dwivedi & Roig ( 2019 ) proposed a more efficient alternative that uses representation similarity analysis ( RSA ) to obtain a measure of task affinity , by computing correlations between models pretrained on different tasks . Given a dataset and a number of tasks , our approach uses RSA to assess the task affinity at arbitrary locations in a neural network . The task affinity scores are then used to construct a branched multitask network in a fully automated manner . In particular , our task clustering algorithm groups similar tasks together in common branches , and separates dissimilar tasks by assigning them to different branches , thereby reducing the negative transfer between tasks . Additionally , our method allows to trade network complexity against task similarity . We provide extensive empirical evaluation of our method , showing its superiority in terms of multi-task performance vs computational resources . 2 RELATED WORK . Multi-task learning . Multi-task learning ( MTL ) ( Caruana , 1997 ; Ruder , 2017 ) is associated with the concept of jointly learning multiple tasks under a single model . This comes with several advantages , as described above . Early work on MTL often relied on sparsity constraints ( Yuan & Lin , 2006 ; Argyriou et al. , 2007 ; Lounici et al. , 2009 ; Jalali et al. , 2010 ; Liu et al. , 2017 ) to select a small subset of features that could be shared among all tasks . However , this can lead to negative transfer when not all tasks are related to each other . A general solution to this problem is to cluster tasks based on prior knowledge about their similarity or relatedness ( Evgeniou & Pontil , 2004 ; Abernethy et al. , 2009 ; Agarwal et al. , 2010 ; Zhou et al. , 2011 ; Kumar & Daume III , 2012 ) . In the deep learning era , MTL models can typically be classified as utilizing soft or hard parameter sharing . In soft parameter sharing , each task is assigned its own set of parameters and a feature sharing mechanism handles the cross-task talk . Cross-stitch networks ( Misra et al. , 2016 ) softly share their features among tasks , by using a linear combination of the activations found in multiple single task networks . Sluice networks ( Ruder et al. , 2019 ) extend cross-stitch networks and allow to learn the selective sharing of layers , subspaces and skip connections . In a different vein , multi-task attention networks ( Liu et al. , 2019 ) use an attention mechanism to share a general feature pool amongst task-specific networks . In general , MTL networks using soft parameter sharing are limited in terms of scalability , as the size of the network tends to grow linearly with the number of tasks . In hard parameter sharing , the parameter set is divided into shared and task-specific parameters . MTL models using hard parameter sharing are often based on a generic framework with a shared off-the-shelf encoder , followed by task-specific decoder networks ( Neven et al. , 2017 ; Kendall et al. , 2018 ; Chen et al. , 2018 ; Sener & Koltun , 2018 ) . Multilinear relationship networks ( Long et al. , 2017 ) extend this framework by placing tensor normal priors on the parameter set of the fully connected layers . Guo et al . ( 2018 ) proposed the construction of a hierarchical network , which predicts increasingly difficult tasks at deeper layers . A limitation of the aforementioned approaches is that the branching points are determined ad hoc , which can easily lead to negative transfer if the predefined task groupings are suboptimal . In contrast , in our branched multi-task networks , the degree of layer sharing is automatically determined in a principled way , based on task affinities . Our work bears some similarity to fully-adaptive feature sharing ( Lu et al. , 2017 ) , which starts from a thin network where tasks initially share all layers , but the final one , and dynamically grows the model in a greedy layer-by-layer fashion . Task groupings , in this case , are decided on the probability of concurrently simple or difficult examples across tasks . Differently , ( 1 ) our method clusters tasks based on feature affinity scores , rather than example difficulty , which is arguably a better criterion ; ( 2 ) the tree structure is determined offline using the precalculated affinities for the whole network , and not online in a greedy layer-by-layer fashion , which promotes task groupings that are optimal in a global , rather than local , sense ; ( 3 ) our approach achieves significantly better results , especially on challenging datasets featuring numerous tasks , like Taskonomy . Neural architecture search . Neural architecture search ( NAS ) ( Elsken et al. , 2019 ) aims to automate the construction of the network architecture . Different algorithms can be characterized based on their search space , search strategy or performance estimation strategy . Most existing works on NAS , however , are limited to task-specific models ( Zoph & Le , 2017 ; Liu et al. , 2018b ; Pham et al. , 2018 ; Liu et al. , 2018a ; Real et al. , 2019 ) . This is to be expected as when using NAS for MTL , layer sharing has to be jointly optimized with the layers types , their connectivity , etc. , rendering the problem considerably expensive . To alleviate the heavy computation burden , a recent work ( Liang et al. , 2018 ) implemented an evolutionary architecture search for multi-task networks , while other researchers explored more viable alternatives , like routing ( Rosenbaum et al. , 2018 ) , stochastic filter grouping ( Bragman et al. , 2019 ) , and feature partitioning ( Newell et al. , 2019 ) . In contrast to traditional NAS , the proposed methods do not build the architecture from scratch , but rather start from a predefined backbone network for which a layer sharing scheme is automatically determined . Transfer learning . Transfer learning ( Pan et al. , 2010 ) makes use of the knowledge obtained when solving one task , and applies it to a different but related task . Our work is loosely related to transfer learning , as we use it to measure levels of task affinity . Zamir et al . ( 2018 ) provided a taxonomy for task transfer learning to quantify such relationships . However , their approach scales unfavorably w.r.t . the number of tasks , and we opted for a more efficient alternative proposed by Dwivedi & Roig ( 2019 ) . The latter uses RSA to obtain a measure of task affinity , by computing correlations between models pretrained on different tasks . In our method , we use the performance metric from their work to compare the usefulness of different feature sets for solving a particular task . Loss weighting . One of the known challenges of jointly learning multiple tasks is properly weighting the loss functions associated with the individual tasks . Early work ( Kendall et al. , 2018 ) used the homoscedastic uncertainty of each task to weigh the losses . Gradient normalization ( Chen et al. , 2018 ) balances the learning of tasks by dynamically adapting the gradient magnitudes in the network . Dynamic task prioritization ( Guo et al. , 2018 ) prioritizes the learning of difficult tasks . Zhao et al . ( 2018 ) observed that two competing tasks can cause the destructive interference of the gradient , and proposed a modulation module to alleviate this problem . Sener & Koltun ( 2018 ) cast multi-task learning as a multi-objective optimization problem , with the overall objective of finding a Pareto optimal solution . Note that , addressing the loss weighting issue in MTL is out of the scope of this work . In fact , all our experiments are based on a simple uniform loss weighing scheme .

Summary: This paper presents a method to infer multi-task networks architecture, more specifically to determine which part of the network should be shared among different tasks. The main idea is to first train a specific, standalone encoder/decoder network for each task. Subsequently, a task affinity factor is computed by looking at the similarity (or, more likely the dissimilarity) of an holdout set of images feature representations. Knowing these dissimilarities (based on RDM), one can cluster the tasks and make similar tasks share more of their layers. Computational budget can also be taken into consideration in the form of the number of parameters. Results on Cityscapes, Taskonomy, and CelebA datasets show, to some extent, improvements against the state of the art.