Feature Partitioning for Efficient Multi-Task Architectures

1 INTRODUCTION . Multi-task learning allows models to leverage similarities across tasks and avoid overfitting to the particular features of any one task ( Caruana , 1997 ; Zamir et al. , 2018 ) . This can result in better generalization and more robust feature representations . While this makes multi-task learning appealing for its potential performance improvements , there are also benefits in terms of resource efficiency . Training a multi-task model should require less data , fewer training iterations , and fewer total parameters than training an equivalent set of task-specific models . In this work we investigate how to automatically search over high performing multi-task architectures while taking such resource constraints into account . Finding architectures that offer the best accuracy possible given particular resource constraints is nontrivial . There are subtle trade-offs in performance when increasing or reducing use of parameters and operations . Furthermore , with multiple tasks , one must take into account the impact of shared operations . There is a large space of options for tweaking such architectures , in fact so large that it is difficult to tune an optimal configuration manually . Neural architecture search ( NAS ) allows researchers to automatically search for models that offer the best performance trade-offs relative to some metric of efficiency . Here we define a multi-task architecture as a single network that supports separate outputs for multiple tasks . These outputs are produced by unique execution paths through the model . In a neural network , such a path is made up of a subset of the total nodes and operations in the model . This subset may or may not overlap with those of other tasks . During inference , unused parts of the network can be ignored by either pruning out nodes or zeroing out their activations ( Figure 1 ) . Such architectures mean improved parameter efficiency because redundant operations and features can be consolidated and shared across a set of tasks . We seek to optimize for the computational efficiency of multi-task architectures by finding models that perform as well as possible while reducing average node use per task . Different tasks will require different capacities to do well , so reducing average use requires effectively identifying which tasks will ask more of the model and which tasks can perform well with less . In addition , performance is affected by how nodes are shared across tasks . It is unclear when allocating resources whether sets of tasks would benefit from sharing parameters or would instead interfere . When searching over architectures , differences in resource use can be compared at different levels of granularity . Most existing work in NAS and multi-task learning searches over the allocation and use of entire layers ( Zoph & Le , 2016 ; Fernando et al. , 2017 ; Rosenbaum et al. , 2017 ) , we instead partition out individual feature channels within a layer . This offers a greater degree of control over both the computation required by each task and the sharing that takes place between tasks . The main obstacle to address in searching for effective multi-task architectures is the vast number of possibilities for performing feature partitioning as well as the significant amount of computation required to evaluate and compare arrangements . A naive brute search over different partitioning strategies is prohibitively expensive . We leverage our knowledge of the search space to explore it more effectively . We propose a parameterization of partitioning strategies to reduce the size of the search space by eliminating unnecessary redundancies and more compactly expressing the key features that distinguish different architectures . In addition , the main source of overhead in NAS is evaluation of sampled architectures . It is common to define a surrogate operation that can be used in place of training a full model to convergence . Often a smaller model will be trained for a much shorter number of iterations with the hope that the differences in accuracy that emerge early on correlate with the final performance of the full model . We propose a strategy for evaluating multi-task architectures using feature distillation which provides much faster feedback on the effectiveness of a proposed partitioning strategy while correlating well with final validation accuracy . In this work we provide : • a parameterization that aids automatic architecture search by providing a direct and compact representation of the space of sharing strategies in multi-task architectures . • an efficient method for evaluating proposed parameterizations using feature distillation to further accelerate the search process . • results on Visual Decathlon ( Rebuffi et al. , 2017 ) to demonstrate that our search strategy allows us to effectively identify trade-offs between parameter use and performance on diverse and challenging image classification datasets . 2 RELATED WORK . Multi-Task Learning : There is a wide body of work on multi-task learning spanning vision , language , and reinforcement learning . The following discussion will center around designing multi-task architectures for deep learning in vision ( Ruder , 2017 ; Caruana , 1997 ; Zhang & Yang , 2017 ) . There are many obstacles to overcome in multi-task architecture design , but the most pressing concerns depend largely on how the problem setting has been defined . Two distinguishing factors include : • Task Ordering : Are all tasks available at all times or are they presented one after the other ? • Fixed vs Learned Strategies : Is a uniform strategy applied across tasks or is a task-specific solution learned ? The former is important as work in which tasks are presented sequentially must address catastrophic forgetting ( French , 1999 ) . This is less of a concern in our work as we train on all tasks at once . As for the latter , finding a solution tuned to a specific set of tasks requires the same sort of outer-loop optimization seen in neural architecture search ( Zoph & Le , 2016 ) which is time-consuming and expensive . The contributions presented in this work seek to make this process more manageable . Multi-Task Architectures : A strong baseline for multi-task architectures is the use of a single shared network ( Caruana , 1997 ; Kaiser et al. , 2017 ) . Deep networks are overparameterized in such a way that the same layers can be applied across different domains while producing features that are useful to different ends . Using a shared architecture is common in reinforcement learning to train a single agent to perform many tasks with a uniform observation and action space ( Espeholt et al. , 2018 ; Sharma et al. , 2017 ) . A common technique to train single shared models well in both reinforcement learning and vision is distillation of multiple models into one ( Beaulieu et al. , 2018 ; Yim et al. , 2017 ; Rusu et al. , 2015 ; He et al. , 2018 ; Chou et al. , 2018 ) . In work where tasks are presented sequentially , the focus is often to build on top of an existing network while not disrupting its ability to solve its original task ( Mallya & Lazebnik , 2018 ; Rebuffi et al. , 2018 ; Rosenfeld & Tsotsos , 2017 ) . Currently , many methods freeze the network ’ s weights so they are not changed while learning a new task . Examples include masking out specific filter weights ( Mallya & Lazebnik , 2018 ) or introducing auxiliary layers ( Rebuffi et al. , 2018 ) . Another approach is to dynamically expand a network with additional capacity for each new task ( Yoon et al. , 2018 ) . All of these methods build on top of a fixed model , meaning that new tasks must perform the computation required for the original task as well as take additional steps to be task-specific . It is also common to build multi-task architectures from sets of layers that are run in parallel ( Misra et al. , 2016 ; Rosenbaum et al. , 2017 ; Fernando et al. , 2017 ; Meyerson & Miikkulainen , 2017 ) . Cross-stitch networks compute activations for each task as a learned weighted sum across these layers ( Misra et al. , 2016 ) . This sort of soft attention over features can be seen in other multi-task architecture work as well ( Liu et al. , 2018 ; Ruder et al. , 2017 ) . There are approaches to search over paths through these layers such that each task has a unique , optimal execution path ( Rosenbaum et al. , 2017 ; Fernando et al. , 2017 ) . Similar to work in single-task NAS , the best path is found by either reinforcement learning or evolutionary algorithms ( Fernando et al. , 2017 ; Liang et al. , 2018 ) . The optimal trade-offs in parameter sharing may occur at a more fine-grained level than entire layers , so instead of working with parallel blocks of layers we divide up individual feature channels . Neural Architecture Search : There are three main areas in which contributions are made for more effective architecture search : search space , optimization , and sample evaluation . Search space : With a well-designed search space , it is possible to randomly sample and arrive at high performing solutions ( Li & Talwalkar , 2019 ; Liu et al. , 2017b ) . In general , NAS operates in a discrete space where entire layers are included or not . We instead propose a continuous search space where slight changes can be made in how resources are allocated across tasks . This allows alternatives for optimization that would not apply in other NAS work . Optimization : Leading approaches either use reinforcement learning or genetic algorithms for NAS ( Zoph & Le , 2016 ; Real et al. , 2018 ; Pham et al. , 2018 ) . This search is difficult and the tradeoffs between approaches are unclear ( Li & Talwalkar , 2019 ) . We test the effectiveness of random sampling and evolutionary strategies optimization ( Mania et al. , 2018 ; Wierstra et al. , 2008 ) . Evaluating Samples : Training a model to convergence is time-consuming and resource intensive . It is not realistic to sample thousands of architectures and train them all . Instead one must use a cheaper form of evaluation . Some options include preserving weights across samples for faster training ( Pham et al. , 2018 ) , successive halving ( Kumar et al. , 2018 ) , progressive steps to increase complexity ( Liu et al. , 2017a ) , as well as techniques to model the expected performance of sampled architectures ( Deng et al. , 2017 ; Brock et al. , 2017 ; Baker et al. , 2017 ) . It is unclear how well surrogate functions correlate with final model performance ( Zela et al. , 2018 ) . We investigate the use of distillation for performing this evaluation . 3 MULTI-TASK FEATURE PARTITIONING . Sharing in the context of multi-task architecture search is often adjusted at the level of individual layers ( Rosenbaum et al. , 2017 ; Fernando et al. , 2017 ) , but given that state-of-the-art architectures work so well across datasets and tasks we choose to preserve the layer level execution path for each task and focus on sub-architectural changes that can be made . In this case , by deciding whether or not individual feature channels within each layer are available . This way , every task experiences the exact same organization of layers , but a unique calculation of intermediate layer features . Concretely , given a feature tensor F ∈ Rc×h×w we define a binary mask mf ∈ { 0 , 1 } c for each task , and during a forward pass of the network multiply F by mf to zero out all channels not associated with a particular task . We further define a mask for the backward pass , mb ∈ { 0 , 1 } c , whose non-zero elements are a subset of the non-zero elements in mf . Gradients are calculated as usual through standard backpropagation , but any weights we wish to leave unchanged will have their gradients zeroed out according to mb . Together , these masks can capture the training dynamics seen in many existing multi-task architecture designs ( Rosenbaum et al. , 2017 ; Rebuffi et al. , 2018 ) . For example , one can devote an outsized proportion of features to a task like ImageNet classification , then make these features available during forward inference on a new smaller dataset . A backward mask , mb , can then be defined to ensure that the ImageNet weights remain untouched when finetuning on the new task . There are a number of advantages to allocating resources at the channel level . There is enough flexibility to allow fine-grained control allotting specific weights to particular subsets of tasks . And after training , it is straightforward to prune the network according to any mask configuration . Meaning for simple tasks that only require a small subset of channels we can use a fraction of compute at test time while still leveraging the advantages of joint training with other tasks . An important implementation detail is that masks are applied every other layer . Consider , for example , making a task use half the model . You might think to set half the values of mf to one and apply it to each layer . But that would mean c2 inputs and c 2 outputs at each layer which only uses one quarter of the original model . Instead , applying a mask at every other layer produces the desired behavior of allocating half the model to the task .

Summary: This paper studies multi-task learning (MTL) from the deep learning perspective where a number of layers are shared between tasks followed by specific heads for each task. One of the main challenges in this problem is to decide the best configuration among a large number of possible ones (e.g., the number of layers , number of neurons, when to stop the shared part of the network). In this paper, the authors fix the network architecture, and learn which filters (among the already learned ones) should be dedicated to (and hence fine-tuned for) a specific, and which ones should be shared between multiple tasks.