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: This paper proposes a method for grouping tasks to achieve stronger multi-task learning (MTL) performance. The approach is relatively simple and quite efficient in comparison to alternatives: during a single pass of learning, measure the impact that one task's gradient updates would have on other tasks. A theoretical analysis is conducted in a simplified (strongly convex) setting, and experiments show that this approach is as effective as alternatives that are much more computationally intensive. Overall this appears to be a useful technique that could be useful for many other researchers and also potentially for product teams.