Efficiently Identifying Task Groupings for Multi-Task Learning

1 Introduction . Many of the forefront challenges in applied machine learning demand that a single model performs well on multiple tasks , or optimizes multiple objectives while simultaneously adhering to unmovable inference-time constraints . For instance , autonomous vehicles necessitate low inference time latency to make multiple predictions on a real-time video feed to precipitate a driving action [ 27 ] . Robotic arms are asked to concurrently learn how to pick , place , cover , align , and rearrange various objects to improve learning efficiency [ 25 ] , and online movie recommendation systems model multiple engagement metrics to facilitate low-latency personalized recommendations [ 13 ] . Each of the above applications depends on multi-task learning , and advances which improve multi-task learning performance have the potential to make an outsized impact on these and many other domains . Multi-task learning can improve modeling performance by introducing an inductive bias to prefer hypothesis classes that explain multiple objectives and by focusing attention on relevant features [ 43 ] . However , it may also lead to severely degraded performance when tasks compete for model capacity or are unable to build a shared representation that can generalize to all objectives . Accordingly , finding groups of tasks that derive benefit from the positives of training together while mitigating the negatives often improves the modeling performance of multi-task learning systems . While recent work has developed new multi-task learning optimization schemes [ 28 , 10 , 45 , 53 , 11 , 50 ] , the problem of deciding which tasks should be trained together in the first place is an understudied and complex issue that is often left to human experts [ 56 ] . However , a human ’ s understanding of similarity is motivated by their intuition and experience rather than a prescient knowledge of the underlying structures learned by a neural network . To further complicate matters , the benefit or detriment induced from multi-task learning relies on many non-trivial decisions including , but not limited to , dataset characteristics , model architecture , hyperparameters , capacity , and convergence [ 51 , 49 , 46 , 47 ] . As a result , a systematic technique to determine which tasks 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . should train together in a multi-task neural network would be valuable to practitioners and researchers alike [ 5 , 6 ] . One approach to select task groupings is to exhaustively search over the 2|T |−1 multi-task networks1 for a set of tasks T . However , the cost associated with this search can be prohibitive , especially when there is a large number of tasks . It is further complicated by the fact that the set of tasks to which a model is applied may change throughout its lifetime . As tasks are added to or dropped from the set of all tasks , this costly analysis would need to be repeated to determine new groupings . Moreover , as model scale and complexity continues to increase , even approximate task grouping algorithms which evaluate only a subset of combinations may become prohibitively costly and time-consuming to evaluate . In this paper , we aim to develop an efficient framework to select task groupings without sacrificing performance . We propose to measure inter-task affinity by training all tasks together in a single multi-task network and quantifying the effect to which one task ’ s gradient update would affect another task ’ s loss . This per-step quantity is averaged across training , and tasks are then grouped together to maximize the affinity onto each task . A visual depiction of the method is shown in Figure 1 . Our suggested approach makes no assumptions regarding model architecture and is applicable to any paradigm in which shared parameters are updated with respect to multiple losses . In summary , our primary contribution is to suggest a measure of inter-task affinity that can be used to systematically and efficiently determine task groupings for multi-task learning . Our theoretical analysis shows that grouping tasks by maximizing inter-task affinity will outperform any other task grouping in the convex setting under mild conditions . Further on two challenging multi-task image benchmarks , our empirical analysis finds this approach outperforms training all tasks independently , training all tasks together ( with and without training augmentations ) , and is competitive with a state-of-the-art task grouping method while decreasing runtime by more than an order of magnitude . 2 Related Work . Task Groupings . Prevailing wisdom suggests tasks which are similar or share a similar underlying structure may benefit from training together in a multi-task system [ 9 , 8 , 4 ] . Early work in this domain pertaining to the convex setting assume all tasks share a common latent feature representation , and find that model performance can be significantly improved by clustering tasks based on the basis vectors they share in this latent space [ 26 , 30 ] . However , early convex methods to determine task groupings often make prohibitive assumptions that do not scale to deep neural networks . Deciding which tasks should train together in multi-task neural networks has traditionally been addressed with costly cross-validation techniques or high variance human intuition . An altogether 1There are ( |T | 1 ) single-task networks , ( |T | 2 ) 2-task multi-task networks , ... which yields ∑|T | i=1 ( |T | i ) = 2|T |−1 . different approach may leverage recent advances in transfer learning focused on understanding task relationships [ 54 , 3 , 15 , 58 , 2 ] ; however , [ 46 ] show transfer learning algorithms which determine task similarity do not carry over to the multi-task learning domain and instead propose a multi-task specific framework which trains between ( |T | 2 ) + |T | and 2|T | − 1 models to approximate exhaustive search performance . Our approach differs from [ 46 ] in that it computes task groupings from only a single training run . Architectures and Training Dynamics . A plethora of multi-task methods addressing what parameters to share among tasks in a model have been developed , such as Neural Architecture Search [ 19 , 47 , 49 , 35 , 22 , 39 ] , Soft-Parameter Sharing [ 40 , 14 , 52 ] , and asymmetric information transfer [ 31 , 44 , 32 ] to improve multi-task performance . Although this direction is promising , we direct our focus towards when to share tasks in a multi-task network rather than architecture modifications to maximize the benefits from training a fixed set of tasks together . Nevertheless , both approaches are complementary , and architecture augmentations seem to perform best when trained with related tasks [ 43 ] . Significant effort has also been invested to improve the optimization dynamics of MTL systems . In particular , dynamic loss reweighing has achieved performance superior to using fixed loss weights found with extensive hyperparameter search [ 28 , 18 , 37 , 10 , 45 , 33 ] . Another set of methods seek to mitigate inter-task conflict by manipulating the direction of task gradients rather than simply their magnitude [ 48 , 57 , 53 , 11 , 50 , 36 ] . In our experiments , we compare against Uncertainty Weights [ 28 ] , GradNorm [ 10 ] , and PCGrad [ 53 ] to contextualize the relative change in performance from splitting tasks into groups . We find that task grouping methods outperform all three training augmentations ; nonetheless , they are naturally complementary . In Section 5 , our results indicate enhancing the networks found by our method with PCGrad can lead to additional improvements in performance . Looking into the Future . “ Lookahead ” methods in deep learning can often be characterized by saving the current state of the model , applying one or more gradient updates to a subset of the parameters , reloading the saved state , and then leveraging the information learned from the future state to modify the current set of parameters . This approach has been used extensively in the metalearning [ 16 , 42 , 7 , 17 , 29 ] , optimization [ 41 , 21 , 55 , 23 , 24 ] , and recently auxiliary task learning domains [ 34 ] . Unlike the above mentioned methods which look into the future to modify optimization processes , our work adapts this central concept to the multi-task learning domain to characterize task interactions and assign tasks to groups of networks . 3 Task Grouping Problem Definition . We draw a distinction between inference-time latency constraints and inference-time memory budget . The former characterizes the speed at which predictions can be computed , with similarly sized models running in parallel having latency roughly equivalent to a single model running by itself . The latter relates to the number of parameters used by all models during inference , with a n-times parameter model having a similar budget to an n-group of normal-sized models . We configure our analysis to span both dimensions , but also provide analysis into only the latter in the Appendix . Given a set of tasks T , a fixed inference-time memory budget b , and latency constraint c , our aim is to assign tasks to networks such that average task performance is maximized . Additionally , each network must have a parameter count less than c ; the networks must span our set of tasks T ; and the total number of networks is less than or equal to our memory budget b . Moreover , tasks can be trained in a network without being served from this network during inference . In this case , they are used to assist the learning of the other tasks , and during inference , are served from one of the other multi-task networks . Formulating our task grouping framework in this manner aligns our work with industry trends where inference-time parameter budgets are limited and latency unmovable . More formally , for a set of n tasks T = { τ1 , τ2 , .. , τn } , we would like to construct a group of k multi-task neural networks M = { m1 , m2 , ... , mk } such that ∀τi ∈ T , ∃ exactly one multi-task network mj ∈M , parameter count of mj < c , such that mj makes an inference-time prediction for ti subject to k ≤ b where b is our memory budget . mj is the jth multi-task network which takes input X , and concurrently trains a set of tasks { τa , τc , ... , τf } , but only serves a subset of those tasks at inference . For a given performance measure P , we can then define the aggregate performance of our task grouping as ∑n i=1 P ( τi|M ) where P ( τi|M ) computes the performance of task τi from the set of models M using the model mj which the task grouping algorithm predicts the performance of τi will be highest . 4 Grouping Tasks by Measuring Inter-Task Affinity . We propose a method to group tasks by examining the effect to which one task ’ s gradient would increase or decrease another task ’ s loss . We formally define the method in Section 4.1 , describe a systematic procedure to go from inter-task affinity scores to a grouping of tasks in Section 4.2 , and provide theoretical analysis in Section 4.3 . 4.1 Inter-Task Affinity . Within the context of a hard-parameter sharing paradigm , tasks collaborate to build a shared feature representation which is then specialized by individual task-specific heads to output a prediction . Specifically , through the process of successive gradient updates to the shared parameters , tasks implicitly transfer information to each other . As a consequence , we propose to view the extent to which a task ’ s successive gradient updates on the shared parameters affect the objective of other tasks in the network as a proxy measurement of inter-task affinity . Consider a multitask loss function parameterized by { θs } ∪ { θi| i ∈ T } where θs represents the shared parameters and θi represents the task i ∈ T specific parameters . Given a batch of examples X , let Ltotal ( X , θs , { θi } ) = ∑ i∈T Li ( X , θs , θi ) , denote the total loss where Li represents the non-negative loss of task i . For simplicity of notation , we set the loss weight of each task to be equal to 1 , though our construction generalizes to arbitrary weightings . For a given training batch X t at time-step t , define the quantity θt+1s|i to represent the updated shared parameters after a gradient step with respect to the task i . Assuming stochastic gradient descent for simplicity , we have θt+1s|i : = θ t s − η∇θtsLi ( X t , θts , θ t i ) . We can now calculate a lookahead loss for each task by using the updated shared parameters while keeping the task-specific parameters as well as the input batch unchanged . That is , in order to assess the effect of the gradient update of task i on a given task j , we can compare the loss of task j before and after applying the gradient update from task i onto the shared parameters . To eliminate the scale discrepancy among different task losses , we consider the ratio of a task ’ s loss before and after the gradient step on the shared parameters as a scale invariant measure of relative progress . We can then define an asymmetric measure for calculating the affinity of task i at a given time-step t on task j as Zti j = 1− Lj ( X t , θt+1s|i , θtj ) Lj ( X t , θts , θtj ) . ( 1 ) Notice that a positive value of Zti j indicates that the update on the shared parameters results in a lower loss on task j than the original parameter values , while a negative value of Zti j indicates that the shared parameter update is antagonistic for this task ’ s performance . Our suggested measure of inter-task affinity is computed at a per-step level of granularity , but can be averaged across all steps , every n steps , or a contiguous subset of steps to derive a “ training-level ” score : Ẑi j = 1 T T∑ t=1 Zti j . In Section 5 , we find Ẑi j is empirically effective in selecting high performance task groupings and provide an ablation study along this dimension in Section 5.2 .

Summary: Given a large set of tasks, the goal of this work is to identify subsets of these tasks that can be trained together to improve performance and avoid harmful interference. In contrast to multi-task research that focuses on building a model that handles _all_ tasks as well as possible, this work looks at the cues acquired from training such a model to quickly assess which tasks to group together. This is done by seeing whether a task's gradient update helps or hurts the losses of other tasks.