title: HydaLearn: Highly Dynamic Task Weighting for Multi-task Learning with Auxiliary Tasks
authors: Sam Verboven Muhammad Hafeez Chaudhary Jeroen Berrevoets Wouter Verbeke
text: 
  
  
  
  
  
  
  
  
  
  theoremTheorem
  
  
  
  
  
  
  
  
  
  
  Each Part Matters: Local Patterns Facilitate 
   Cross-view Geo-localization
      Tingyu Wang,
          Zhedong Zheng,
          Chenggang Yan,
          Jiyong Zhang, Member, IEEE,
          Yaoqi Sun,
          Bolun Zheng,
          and Yi Yang, Senior Member, IEEE
  Copyright  2021 IEEE. Personal use of this material is permitted. However, permission to use this material for any other purposes must be obtained from the IEEE by sending an email to pubs-permissions@ieee.org.
  Tingyu Wang, Chenggang Yan, Jiyong Zhang, Yaoqi Sun and Bolun Zheng are with the Intelligent Information Processing Lab, Hangzhou Dianzi University, Hangzhou 310018, China (e-mail: wongtyu@hdu.edu.cn; cgyan@hdu.edu.cn; jzhang@hdu.edu.cn; syq@hdu.edu.cn; blzheng@hdu.edu.cn). Chenggang Yan is the Corresponding Author.
  Zhedong Zheng, Yi Yang are with the Australian Artificial Intelligence Institute, University of Technology Sydney, NSW 2007, Australia (e-mail: Zhedong.Zheng@student.uts.edu.au; Yi.Yang@uts.edu.au).
  This work was partially sponsored by Zhejiang Lab’s International Talent Fund for Young Professionals (No.ZJ2020GZ021).
  
      
  ====================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================
  
  
  
  
  Multi-task learning (MTL) can improve performance on a task by sharing representations with one or more related auxiliary-tasks. Usually, MTL-networks are trained on a composite loss function formed by a constant weighted combination of the separate task losses. In practice, constant loss weights lead to poor results for two reasons: (i) the relevance of the auxiliary tasks can gradually drift throughout the learning process; (ii) for mini-batch based optimisation, the optimal task weights vary significantly from one update to the next depending on mini-batch sample composition. We introduce HydaLearn, an intelligent weighting algorithm that connects main-task gain to the individual task gradients, in order to inform dynamic loss weighting at the mini-batch level, addressing i and ii. Using HydaLearn, we report performance increases on synthetic data, as well as on two supervised learning domains.
  
  
  
  
  
  § INTRODUCTION.
  
  Through joint training of shared representations with one or more auxiliary tasks, Multi-task learning <cit.> can increase performance of neural networks on a task of interest - i.e. the main task.  
  
  
  The Problem Setting.   How much a given auxiliary task should influence the training process at each step is an open research question. The gain of learning from auxiliary tasks depends on their back propagated gradients that contribute in learning the main task. While training on a composite multi-task loss has the potential to provide a richer, better regularized training signal, MTL networks are not straightforward to train. Auxiliary tasks in fact do not always contribute to better predictions for the main task. In particular, when the signal is not sufficiently relevant, performance on the main task can deteriorate. Moreover, the usefulness of the auxiliary task gradients is subject to changes over the course of training process <cit.>.  Automated adaptive MTL is as such a high-impact problem, for which a solution would mitigate a significant part of the difficulty of training powerful MTL models. Consequently, this could lead to increased adoption of such models among practitioners. More wide-spread diffusion could occur specifically outside of high-dimensional, high-compute supervised learning problems.
  
  Slow Weight Adjustment is Not Enough.   Contemporary research on adaptive task weighting either does not allow explicit prioritization of main task performance <cit.>, or requires some implicit assumptions which are frequently violated. For instance: (i) the main task gradient direction is consistently desirable <cit.> or; (ii) task usefulness only changes gradually over the learning process <cit.>. However, the composition of the mini-batch ultimately determines the training signal, and is expected to be highly variable. This variability is desirable for stochastic optimization of non-convex loss surfaces <cit.>, but causes the optimal loss weighting in multi-task learning to drastically change from mini-batch to mini-batch. Accordingly, the usefulness of a single back propagated task gradient is not only dependent on: (i) the parametrization of the model; but, we argue, (ii) the sample composition of the mini-batch.
  Contrary to the slower changing parametrization of the model, mini-batch statistics can differ significantly from one update to the next. Hence, slow weight adjustment over the course of training does not suffice. 
  
  
  
  
  
  
  Our Contribution.   We introduce and evaluate HydaLearn; an adaptive task weighting algorithm designed to handle high variance estimates of the exact gradients. On the level of the individual mini-batch, we estimate the usefulness of individual task gradients based on a metric of choice. Subsequently, these usefulness measures are used to compute task weights which are in turn used to construct a composite gradient. This composite gradient then updates the model parameters. 
  
  We extensively evaluate HydaLearn on a synthetic toy example as well as two real-world datasets, and demonstrate the effectiveness of the algorithm compared to a both logical and state-of-the-art baselines. The latter datasets comprise two separate supervised learning tasks. First, we perform in-hospital mortality prediction, aided by length-of-stay prediction using the MIMIC-III dataset <cit.>. Second, we cast default prediction as MTL through inclusion of prepayment prediction as an auxiliary task, using the Fannie Mae Single Family Loan Performance dataset. 
  
  
  
  
  
  
  § RELATED LITERATURE
  
  
  Multi-task Learning.
  The benefit from MTL is dependent on sample size, number of tasks, and the intrinsic data dimensionality <cit.>. The dominant paradigm to reap this benefit is hard parameter sharing <cit.>. In hard parameter sharing <cit.>, all parameters in a set of hidden layers are fully shared between the tasks. This fully shared hidden layer stack (or encoder) learns a shared representation, and is generally followed by one or more layers which do not share parameters, i.e., task-specific layers. These task-specific heads of the network are often called decoders and learn a the task-specific representation. 
  
  In MTL, a multitask loss is usually used, composed of a linearly weighted combination of individual task losses. These weights are constants and generally chosen through expensive search procedures <cit.> or Bayesian optimization <cit.>. 
  
  The objective in MTL can be either to simultaneously maximise performance of all tasks <cit.>, or leverage auxiliary tasks insofar they improve the main task <cit.>, i.e., not care about the performance on the auxiliary tasks. The latter objective is the focus of this work.
  
  
  Auxiliary Task Learning.   It is generally assumed that related tasks will improve and unrelated tasks can hamper performance <cit.>, although even unrelated tasks can be exploited <cit.>. However, there is still no theoretically grounded definitive definition of task relatedness. Empirically, adding auxiliary tasks has shown to improve performance for <cit.>. The concept of adding auxiliary tasks is most common in complex, high dimensional domains such as autonomous vehicle control <cit.>, reinforcement learning <cit.>,  and natural language processing <cit.>. 
  
  However, MTL can also be a valid approach on common low-compute problems.
  For example, the implicit data augmentation inherent to MTL can help to learn from imbalanced data <cit.>, which is commonly encountered in fields such as fraud detection <cit.> and default prediction <cit.>. Different strands of work have proposed strategies for auxiliary task selection or training. Our work contributes to the latter, more specifically by dynamically learning the task loss weights.
  
  
  Adaptive Task Weighting in Multi-task Learning.   The optimal task loss distribution changes as parametrization of the model gradually changes. To handle this, task weighting should be dynamically adapted throughout the training cycle. 
  Most of the existing work on adaptive task learning considers the setting where one tries to optimize all the included tasks together <cit.>. Different approaches have been proposed including
  explicit prioritization of difficult tasks <cit.> , minimization of negative conflict between gradients <cit.>, and balancing of the task gradient norms <cit.>. Furthermore, homeostatic uncertainty can be used to inform task weighting <cit.>. Finally, <cit.> focus on adapting the learning rates of the tasks, which under vanilla stochastic gradient descent (SGD) is equivalent to scaling the gradients.
  
  Adaptive Task Weighting for Auxiliary Tasks.   The algorithms most relevant to ours are proposed in <cit.> and <cit.> and aim to explicitly maximize performance on the main task <cit.>. <cit.> only train on the auxiliary task if its gradient aligns with the main task gradient, as defined by the so-called cosine similarity measure. In <cit.>, the relative task weights of auxiliary tasks are updated every N steps using a variant of online cross-validation <cit.>, where N is a hyper-parameter. 
  
  Small Mini-Batch Learning.   Neural networks are usually trained with the mini-batch SGD method <cit.> or one of its variants <cit.>. A common belief among researchers dictates that smaller batch sizes yield better out-of-sample generalization <cit.>. This belief has seen theoretical and experimental validation <cit.>, although it has also been challenged <cit.>. Nonetheless, standard practice in most domains is still to train with small mini-batches. 
  
  Variable Usefulness of Mini-Batches.   When training a neural network, some data points are more valuable than others <cit.>. By extension, through the stochastic nature of mini-batch SGD, the number of valuable or harmful examples varies over mini-batches, i.e., some mini-batches are more valuable than others. Harder examples tend to give larger gradients, and are in some cases more useful <cit.>. Nonetheless, it has been shown that harder does not always imply more valuable <cit.>. Furthermore, the usefulness of individual examples can differ between tasks. Ultimately, the goal is to optimize performance on the main task. This implies that the gradient norm distribution of the auxiliary task is not informative for the usefulness of individual examples, and by extension batches. Accordingly, strategies which try to exploit such information <cit.> are not applicable in this context.
  
  
  
  
  § HYDALEARN
  
  
  
  
  Problem Formulation.
  We focus on a problem comprising a main task T_m and an auxiliary task T_a. As shown in Fig. <ref>, we have a number of shared layers between the task and then we have task-specific layers.   Let L_m and L_a be the loss function associated with the main and auxiliary task respectively. The two losses are combined to form a total loss as follows: 
   
      L(θ_s,t, θ_m,t, θ_a,t) =    w_m, t L_m(θ_s,t, θ_m,t) 
      
         + w_a, t L_a(θ_s,t, θ_a,t),
  
  In (<ref>), t is the training step; w_m, t, w_a, t∈ℝ_+ denote the task weights, where ℝ_+ is a set of all positive real numbers; and θ_s,t, θ_m,t, θ_a,t denote model parameters of the shared and task-specific layers. We aim to find a solution for training the model parameters and develop a method to update task weights intelligently. 
  
  The model parameters are updated using mini-batch SGD with the objective of minimizing total loss, for given task weights. Let θ_t ≜ [θ_s,t, θ_m,t, θ_a,t]. The gradient descent update at training step t can be written as: 
  
      θ_t+1 = θ_t  - α∇_θ_t L(θ_t),
  
  where α denotes the gradient descent step-size, also called the learning rate. 
  
  Proposed Solution.
  The key innovation of HydaLearn lies in the way we determine the optimal loss weights for each mini-batch separately. For optimizing the task weights w_m, t and w_a, t, we focus on maximizing gain on a given metric for the main task. Let  M_m denote the metric function for the main task. 
  This metric is calculated over a much larger set of examples than the batch size, and is thus a more stable guide of progression. At the same time we do not fully block 'bad' gradient steps, ensuring sufficient stochasticity in the learning path to generalize well.
  Furthermore, the choice of metric function is not constrained to any particular function type. This implies that non-differentiable metrics can be chosen, but also the corresponding task loss itself. Recent dynamic weight adaption methods (e.g., <cit.>) require the underlying function to be strictly differentiable. Finding an approximation that is representative of the dynamics of the underlying function over the support region of interest is not straightforward. 
  
  In the following theorem we first establish the relationship between the task weights and the gain on the metric M_m. Next, in the ensuing paragraphs, we use this expression to develop the HydaLearn algorithm.
  
  
  
  
  Let δ_m, m, t and δ_m, a, t denote the gain computed on the main task metric at training step t, with gradient-descent steps executed separately on the main and auxiliary task-loss objectives. More concrete mathematical definition of the δ's is given in Appendix <ref>.  The following relationship holds between the gain values and the task weights:
  
      w_m, tw_a, t≈δ_m, m, tδ_m, a, t.
  
  
  
  
  See Appendix <ref>. 
  
   
    With known values of δ's, and constraining the sum of the task weights, we can compute the weights from (<ref>). We impose a constraint on the combined weights, given as w_m + w_a = W. This constraint serves two purposes. First, it makes the solution to the weight optimization problem easy to find—two linear equations to solve for two unknowns. Second, it helps to restrict the total learning rate to W in the gradient descent algorithm in (<ref>). The latter constraint is also used in other work, e.g., <cit.>. It is interesting to note that the result in (<ref>) does not depend on any derivative of the loss and the metric functions and the computations of δ's in our proposed algorithm are based on metric function M_m values from two updates, further explained in the ensuing sections.
   
   The overall procedure to train the model parameters and the task weights is described in Algorithm <ref>. We call this algorithm HydaLearn, or Highly Dynamic Learning. The weight updating procedure in the algorithm aims to optimize the metric function M_m at each training step. The weights can adapt to the signal coming from the individual task losses, to realize maximal gain on the metric. Extension to cases where the metric function needs to be minimized is straightforward. Furthermore, one can choose to optimize the metric on (a subset of) either the training or validation set. 
  
  Analysis. 
  The current HydaLearn solution is applicable to two-task settings—a main and an auxiliary task. Extension to cases that entail more than one auxiliary task is planned in upcoming work. 
  
  
  
  
  The inner working of the HydaLearn algorithm can be better understood with the help of Fig. <ref>. The forward pass is a standard step, whereas backward pass is taken in three steps, in the following order: 
  
      
    * Task specific model parameters are updated.
      
    * Then for the computation of δ's we perform two, what we call, fake-updates for the shared layer model parameters—one based on main-task loss and the second based on the auxiliary-task loss. The update is called fake because the resulting model parameters are only used in the computation of δ's. After each fake update, the corresponding δ is computed. Based on the resulting values of δ's, task weights are computed. 
      
    * Finally, the new task weights are then used for the actual update of the shared layers model parameters. 
  
  
  
  
  
  Comparison with Related Approaches.   Existing methods are not geared towards handling the high-variance estimators of the gradient, i.e., mini-batches.
  While the gradient cosine similarity (Gcosim) based method proposed in <cit.> uses information about the gradient of the main task to ground the contribution of the auxiliary task, <cit.> uses feedback from past batches to take gradient-descent steps on the task weights. When the main task gradient is noisy, Gcosim can exacerbate bad updates, or block helpful auxiliary task gradients, cf., the bottom left panel in Fig. <ref>. When the gradient cosine similarity is positive, there is no normalization of the gradient size, and one loss may dominate the other, cf., top panels in Fig. <ref>. Our algorithm handles such cases differently, as illustrated in the left panels. Even though T_a gradients are dominant, more weight is given to T_m for this particular update, due to the superior direction of its gradients. Olaux <cit.>, on the other hand, is not geared towards highly varying weights since: (i) it only features a single gradient step on the task weights at a time and; (ii) this gradient step towards the task weights occurs after network parameters are updated for the current batch, i.e., the weighting is not related to the contribution of the current batch.
  
  
  
  
  
  
  In the ensuing sections we evaluate our algorithm on synthetic toy examples as well as on real world datasets. We compare results with a list of prominent algorithms, from the current state of art, that tackle similar MTL problem. There we will show that our algorithm enables redistribution of weight towards the task for which a particular mini-batch yields the most valuable signal, i.e., causes the most performance gain on the main task. 
  
  
  
  
  § EXPERIMENTS
  
  
  
  
   §.§ Baselines
  
  For performance comparison we have selected following algorithms as a baseline, where the first two can be viewed as standard baselines used in such comparative studies whereas the latter three are current state of the art. 
  
      
    * Single Task Model (STL):  A model trained only on the main task
      
    * Static Loss weights (Static): A baseline with static weights throughout the full training process
      
    * GradNorm: GradNorm <cit.> balances the training rates of all tasks by normalizing the gradient magnitudes through tuning of the multitask loss 
      
    * Gcosim: This algorithm only uses the auxiliary task gradient when it aligns with the main task gradient as defined by cosine similarity <cit.>
      
    * OLaux: Online cross-validation method for adaptive task weighting. The weights of the auxiliary task are updated based on past batches with gradient descent, every few steps <cit.>
  
  For each of these baselines, we use the same encoder-decoder backbone architecture. We assign an equal budget for optimization,  the learning rate and algorithm-specific parameters to ensure a fair comparison.
  
  
  
  
   §.§ Simulation Study
  
  Using a toy example, we first show that our model is able to efficiently learn when the usefulness of both task gradients is varying. Second, we demonstrate that our model can ignore harmful auxiliary tasks.
  
  Setup.
  We generate two toy examples:
  
      
    Exp1– Related Auxiliary Task: The main and auxiliary tasks are sampled from the tanh function class according to (<ref>). As such, the tasks are related through their common basis 𝐁, and the auxiliary task can help learning the main task. Furthermore, some Gaussian noise is added to both tasks. 
      
      
    Exp2– Unrelated Auxiliary Task: To ensure the unrelatedness of the auxiliary task, the output values for T_a are uniformly sampled, I.I.D., from the output range of T_m over the dataset, i.e., there is no systematic relatedness between the tasks. 
  
  All models are trained using the same backbone. For Exp1, a 4-layer encoder, and two 1-layer task-specific decoders with 64 and 32 neurons, respectively. For Exp2, a 2-layer encoder and two 1-layer task-specific decoders with 40 and 20 neurons were employed. A small grid-search is used to optimize algorithm-specific hyper-parameters. All other hyper-parameters are kept the same for every method.
  The exact implementation details for the toy examples are the corresponding models can be found in Appendix <ref>. The models are both trained and evaluated with standard mean squared error (MSE) as loss function L_m and L_a,  and mean absolute error (MAE) as metric function M_m. 
  
  
  
  As an implementation detail for HydaLearn, we allow down-scaling of so-called total learning rate W when both δ's are negative, as follows: 
  
      W' = W(1 + exp(-w_aw_m))^-1,
  
  where the training-step index t is omitted for brevity. Both negative δ's value means neither of the task gradients for the given mini-batch is supporting improvement on the metric M_m. One possibility was just skip updating the weights for that batch. But by doing we will not exploit information during that training step.  
  
  In another implementation detail, fitting an exponential function to the δ ratio helps to amplify the difference in gain values on the metric M_m from gradients of the  corresponding tasks. More concretely, we do this as follows:
  
      w_mw_a = (δ_mδ_a)^β,
  
  where β∈ℝ_+ can be assumed as a hyper-parameter of the HydaLearn algorithm. The β value can be used to control dynamics or stochasticity of weight changes from batch-to-batch.  
  
  Results.
  Results are shown in Fig. <ref>. From the figure we can see that for Exp1, the task weights are highly variable for HydaLearn. These drastic weight adjustments positively influences training, as HydaLearn outperforms the baseline methods. On the other hand, when T_a is harmful (Exp2), its weights are consistently low. As such, interference from harmful gradients coming from L_a is suppressed. 
  
  The possible benefits of highly adaptive weighting methods are further evidenced through comparison with the baseline methods. Fig. <ref> shows that both the static baseline and GradNorm weightings consistently lie in between the trend weight values of HydaLearn. Yet, the MSE of HydaLearn is about 50% higher. Olaux performs relatively poorly for both tasks. Its ever increasing weight allocation to the harmful auxiliary task is most surprising. Since Gcosim only considers the auxiliary task gradient when its cosine similarity with the main task gradient is positive, it is intuitively well suited for (Exp2). However, it does not completely rule out interference, as indicated by the superior test time performance of HydaLearn. The implementation details can be found in  Appendix <ref>.
  
  
  
  
  
  § EXPERIMENTS ON REAL WORLD DATASETS
  
  
  
   §.§ Datasets and Tasks
  
  We apply HydaLearn to two pairs of supervised learning tasks from the following datasets:
  
  MIMIC-III. 
  The MIMIC-III database <cit.> is comprised of de-identified data of over 60000 intensive care unit (ICU) stays. 
  MIMIC-III is a popular resource for machine learning research for a variety of tasks, such as mortality prediction <cit.>, length of stay prediction <cit.>, and sepsis prediction <cit.>. Since such tasks are often related, the database is also commonly used as a benchmark for multi-task learning algorithms <cit.>.
  Clinical data is often very noisy. Similarly, in the MIMIC dataset, the base features get recorded only sparsely, and at irregular intervals, requiring heavy imputation. 
  
  We predict in-hospital mortality (classification) as our main task using features collected in the first 48h of stay. For the auxiliary task, the length of stay (regression) is used, which ends with either death or discharge from the hospital. As such, this experiment features a combination of a classification (area under the curve (AUC)) and a regression (MSE) loss. Both losses thus operate on a different scale, which can cause imbalances during learning. Furthermore, the dataset is high dimensional relative to its sample size. This can also impede learning, but can be mitigated through the regularizing effect of MTL. Further dataset details can be found in Appendix <ref>.
  
  
  
  Fannie Mae Loan Performance.
  Data on mortgage default typically has an extreme class imbalance; the large majority of mortgage holders never default. Multitask learning is one way to amplify the signal of the minority class <cit.>. As an auxiliary task, we propose to use prepayment prediction. By jointly learning both tasks, we can incorporate signal from future prepayments in the default model - a trick known as 'using the future to predict the present' <cit.>. For preprocessing and implementation details, see Appendix <ref>.
  
  
  Setup.
  For mortality prediction and default prediction, we use 4-shared layers with 48-neuron each, and 2-layer with 24-neuron encoders, respectively. The decoders are composed of 2 task-specific layers for each task, with 12 neurons for mortality prediction, and 24 for default prediction. These backbones were found to be suitable for these problems through a random search on the 'static baseline. The networks are optimized with vanilla mini-batch SGD with batch-size 16. This choice is motivated by the fact that alternative optimizers such as Adam or RMSprop scale the learning rate. In our experience, HydaLearn performs well with Adam too. The training time is chosen through early stopping on the validation set.
  
  Since many of the baseline methods, and HydaLearn work on the gradient level, the learning rate is an important parameter. We perform a grid search over a range of sensible values of the learning rate for each method. Furthermore, we optimize the algorithm-specific hyper-parameters. The grids that were used, as well as the final picks for testing are reported in Table <ref> given in Appendix <ref>.
  
  
  
  
  
   §.§ Results
  
  
  Table <ref> compares performance of different algorithms on the MIMIC dataset. The reported results are evaluated on the hold-out or test data-set. From the table we can see that HydaLearn gives better performance on the main task-metric compared to the STL baseline. Furthermore, HydaLearn surpasses the best performing algorithm from the recent state of art in multi-task learning with dynamic weight adaption for the MIMIC dataset. The same trend extends to the Fannie Mae dataset, as reported in Table <ref>. These performance gains on the real datasets, compared to the baslines, validate the effectiveness of the HydaLearn algorithm in dynamically extracting good values of the task weights.
  
  Furthermore, from Fig. <ref> we can derive that Olaux, HydaLearn and the static weight method take very different approaches to learning default. On average, Olaux learns a very balanced weighting, while the optimal Static combination weight has a very high weight for the main task, i.e. default prediction. In contrast, HydaLearn, gives a higher weight to the prepayment task on average. For batches with no default observations, it predominantly learns its shared representation through prepayment. When a batch is sampled for which the gradient of the default task is highly informative, i.e. one containing examples of the minority class, HydaLearn can adapt and immediately allocate high weight. Such dynamics would not be possible without the highly adaptive weights that characterize our algorithm.
  
  
  
  
  
  
  
  
  
  
  
  
  
  
   §.§ Parameter Impact Analysis
   
  
  As introduced in Section <ref>, our model relies on so-called fake updates to gauge the usefulness of the task gradients for gain on the main task metric M_m. Now in this section, we perform an impact study and study how model performance changes when we disable certain component of our algorithm. Concretely, we consider following four experiments, the results for which are presented in Table <ref> and Table <ref>;
  
      
    * ExpImp-0: Same configuration as in the preceding section. Key components of this configuration include: (c-i) normalizing gradients with Euclidean norm  during fake updates, (c-ii) computation of metric M_m on the validation dataset, and (c-iii) down-scaling of total learning rate as specified in (<ref>).  
      
    * ExpImp-1: Compared to  ExpImp-0, we disable normalizing gradients by removing the normalization step (c-i) completely. rest of the configuraiton is same as in ExpImp-0. This implies that, in absence of normalization, if the gradients from the two fake update point in an equally beneficial direction for the main task, the gradient with a larger magnitude will dominate the final update. From  results in Table <ref> and Table <ref>, we observe a small degradation, which is a little bit more on the MIMIC dataset. As such the degradation for both datasets are within the measurement uncertainty limit and thus negligible. However, the relatively larger change for the MIMIC dataset could partly be because the two tasks are of different type and essentially have completely different dynamic range.   
      
    * ExpImp-2: All configuration parameters same as in ExpImp-0, except that the  δ's are computed over the training dataset. When calculating the δ's based on the validation set, signal from the validation set is incorporated in training. This has the undesirable side effect of increasing bias in validation metrics, complicating model selection and early stopping. If there is no performance trade-off, it makes sense to calculate δ's using the training set. However, this could encourage over-fitting on the training set. Based on the reported results, we see no evidence of over-fitting. But this because we used the entire training data in the computation of δ's. Although note reported in the tables, we observed strong evidences of over-fitting when relatively smaller subset of training dataset was used, which is inline with expectation.    
      
    * ExpImp-3: In this experiment, we disabled the down-scaling of the total learning rate. All other configuration parameters were kept same as in ExpImp-0. From the results, we can see that the impact of this down-scaling feature is really negligible for the two datasets. However, in some of the toy examples on synthetic data we observed non-negligible gains. Further investigation is needed to ascertain the usefulness of this particular feature for other real world datasets.      
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  § CONCLUDING REMARKS
  
  In this work we presented a novel approach to dynamic task weighting in multi-task networks. We have shown that informing gradient weighting through an external metric can match and outperform the current state-of-the-art for two supervised learning tasks. The broader impact of these findings is twofold. First, implementing powerful multi-task models is made easier for practitioners, relieving them from the burden to find suitable static task loss weights, which may still perform worse than dynamic weighting. Second, current work offers a different approach to the dynamic weighting problem compared to previous work, motivated by mathematical grounding. Potential for future work is ample. For example, derivation of an approximation of the δ from the gradients, without computing extra backward passes would significantly reduce computational overhead. Furthermore, this framework can also be applied for high-dimensional problems such as object detection and NLP.
  
  
  
  
  
  
  § APPENDIX
  
  
  
  
   §.§ Function for Task Weight Optimization
  
  Assume M_m is a differentiable function. The gain on the main task metric after taking a gradient descent step on weighted main task-loss objective w_m, t+1 L_m(θ_t) can be written as follows:  
  
      δ_m, m, t+1 =    M_m(θ_t+1) -   M_m(θ_t)
      
       =    M_m(θ_t  - α w_m, t+1∇_θ_t L_m(θ_t)) -   M_m(θ_t)
      ≈    M_m(θ_t) - α w_m, t+1∇_θ_t M_m(θ_t)^T ∇_θ_t L_m(θ_t)
         -  M_m(θ_t) 
         = -α w_m, t+1∇_θ_t M_m(θ_t)^T ∇_θ_t L_m(θ_t),
  
  where the approximation follows from the first-order Taylor series expansion. Similarly the gain on the main task metric after taking a gradient descent step on weighted auxiliary task-loss objective w_a, t+1 L_a(θ_t) can be written as follows:  
  
      δ_m, a, t+1 =    M_m(θ_t+1) -   M_m(θ_t)
      
       =    M_m(θ_t  - α w_a, t+1∇_θ_t L_a(θ_t)) -   M_m(θ_t)
      ≈    M_m(θ_t) - α w_a, t+1∇_θ_t M_m(θ_t)^T ∇_θ_t L_a(θ_t)
         -  M_m(θ_t) 
      
       =    -α w_a, t+1∇_θ_t M_m(θ_t)^T ∇_θ_t L_a(θ_t),
  
  where again the approximation is obtained by applying the first-order Taylor series expansion. Based on (<ref>) and (<ref>), we can write
  
      δ_m, m, t+1δ_m, a, t+1 = (w_m, t+1w_a, t+1)(∇_θ_t M_m(θ_t)^T ∇_θ_t L_m(θ_t)∇_θ_t M_m(θ_t)^T ∇_θ_t L_a(θ_t))
  
  
  
      w_m, t+1w_a, t+1 = (δ_m, m, t+1δ_m, a, t+1) (∇_θ_t M_m(θ_t)^T ∇_θ_t L_a(θ_t)∇_θ_t M_m(θ_t)^T ∇_θ_t L_m(θ_t))
  
  
  Assuming  
  
      ∇_θ_t M_m(θ_t)^T ∇_θ_t L_a(θ_t) ≈∇_θ_t M_m(θ_t)^T ∇_θ_t L_m(θ_t),
  
  we get
  
      w_m, t+1w_a, t+1≈δ_m, m, t+1δ_m, a, t+1.
  
  The underlying assumption is expected to hold when parameter vector dimension reasonably large, and loss functions are on the same scale. 
  
  
  
  
  
   §.§ Toy Example Details
  
  Following <cit.>, we sample two regression tasks from the following functions:
  
      f_i(x) = σ_i tanh((𝐁 + ϵ_𝐢)𝐱),
  
  where 𝐱 is the input vector. ϵ_𝐢 and 𝐁 are constant matrices representing a task-dependent and a shared component, respectively. The σ_i linearly affect the scale of the tasks.
  
  
  
  For both experiments 𝐁 was sampled I.I.D. from a Gaussian with mean 0 and variance 10 , and ϵ_i from a Gaussian with mean 0 and variance 3.5. To represent common scale differences in tasks. The scaling parameters σ_m and σ_a were set to 1 and 10, respectively. 
  
  
  
  
  
  
  
  
  
  
  
   For GradNorm and Olaux we take the recommended values <cit.>. Note that for the experiments with real data we perform hyperparameter optimisation. 
   
  
  
  
  
  
   §.§ Preprocessing and Implementation: MIMIC
  
  
  Only episodes that last longer than 48 hours are considered. For in-hospital-mortality preprocessing, we follow the same approach as for the logistic regression baseline in <cit.> [https://github.com/YerevaNN/mimic3-benchmarks] to enrich 17 base-feature dataset. First, a given sequence is divided into 7 sub-sequences. Next, features are extracted for each subsequence, based on statistical characteristics of the original timeseries variables. Specifically; mean, standard deviation, minimum, maximum, skewness and number of measurements <cit.>. This procedure yields 714 features (7 subsequences X 6 statistic features X 17 base features).
  
  For all the models, we used the same encoder/decoder set-up. We used a random sweep over a range of possible configurations of the 'static' baseline to determine the backbone and batchsize parameters. The shared layers and both task-specific heads of the network consist of 4 layers with 48 neurons, and 2 layers with 24 neurons, respectively. The learning rate, early stopping point, and algorithm-specific hyperparameters were decided using gridsearch.
  
  
  
  
   §.§ Preprocessing and implementation: Fannie Mae
  
  We use a slice of the Fannie and Mae dataset[https://www.fanniemae.com/portal/funding-the-market/data/loan-performance-data.html]. It includes data on over one million mortages. Improved sample efficiency through MTL has diminishing returns. As such, it makes sense to subsample the dataset to a reasonable size. Consequently, we take a uniformly sampled, I.I.D. slice of 10000 data points, containing mortgages that were accepted between 2000 and 2009. For prediction, we use the status at the start of 2010 to predict occurrence of default and prepayment over the next twelve months. 
  
  The continuous and categorical features were standardized and onehot encoded, respectively. The resulting 138 features are used in our experiments. The basic backbone architecture is the same for all models used in the experiments with the Fannie Mae dataset. This backbone consists of 2 24-neuron shared layers, and two 2-layer 12-neuron task-specific heads. Again, learning rate, early stopping, and algorithm-specific hyperparameters are determined for each baseline separately by grid-search.
  
  
  
  
  
  
  
  
  49
      
      
      
      urlstyle[Baesens et al.(2005)Baesens, Van Gestel, Stepanova, Van den Poel, and
        Vanthienen]baesens2005neural
      Baesens, B.; Van Gestel, T.; Stepanova, M.; Van den Poel, D.; and Vanthienen,
        J. 2005.
      Neural network survival analysis for personal loan data.
      Journal of the Operational Research Society 56(9): 1089–1098.
      
      [Baesens, Van Vlasselaer, and Verbeke(2015)]baesens2015fraud
      Baesens, B.; Van Vlasselaer, V.; and Verbeke, W. 2015.
      Fraud analytics using descriptive, predictive, and social
        network techniques: a guide to data science for fraud detection.
      John Wiley & Sons.
      
      [Bingel and Søgaard(2017)]bingel2017identifying
      Bingel, J.; and Søgaard, A. 2017.
      Identifying beneficial task relations for multi-task learning in deep
        neural networks.
      arXiv preprint arXiv:1702.08303 .
      
      [Bottou(2010)]bottou2010large
      Bottou, L. 2010.
      Large-scale machine learning with stochastic gradient descent.
      In Proceedings of COMPSTAT'2010, 177–186. Springer.
      
      [Caruana(1993)]Caruana1993MultitaskLA
      Caruana, R. 1993.
      Multitask Learning: A Knowledge-Based Source of Inductive Bias.
      In ICML.
      
      [Caruana(1997)]caruana1997multitask
      Caruana, R. 1997.
      Multitask learning.
      Machine learning 28(1): 41–75.
      
      [Caruana(2000)]caruana2000learning
      Caruana, R. 2000.
      Learning from imbalanced data: Rank metrics and extra tasks.
      In Proc. Am. Assoc. for Artificial Intelligence (AAAI) Conf,
        51–57.
      
      [Caruana, Baluja, and Mitchell(1996)]caruana1996using
      Caruana, R.; Baluja, S.; and Mitchell, T. 1996.
      Using the future to" sort out" the present: Rankprop and multitask
        learning for medical risk evaluation.
      In Advances in neural information processing systems,
        959–965.
      
      [Chen et al.(2017)Chen, Badrinarayanan, Lee, and
        Rabinovich]chen2017gradnorm
      Chen, Z.; Badrinarayanan, V.; Lee, C.-Y.; and Rabinovich, A. 2017.
      Gradnorm: Gradient normalization for adaptive loss balancing in deep
        multitask networks.
      arXiv preprint arXiv:1711.02257 .
      
      [Collobert and Weston(2008)]collobert2008unified
      Collobert, R.; and Weston, J. 2008.
      A unified architecture for natural language processing: Deep neural
        networks with multitask learning.
      In Proceedings of the 25th international conference on Machine
        learning, 160–167.
      
      [Du et al.(2018)Du, Czarnecki, Jayakumar, Pascanu, and
        Lakshminarayanan]du2018adapting
      Du, Y.; Czarnecki, W. M.; Jayakumar, S. M.; Pascanu, R.; and Lakshminarayanan,
        B. 2018.
      Adapting auxiliary losses using gradient similarity.
      arXiv preprint arXiv:1812.02224 .
      
      [Gentimis et al.(2017)Gentimis, Ala'J, Durante, Cook, and
        Steele]gentimis2017predicting
      Gentimis, T.; Ala'J, A.; Durante, A.; Cook, K.; and Steele, R. 2017.
      Predicting hospital length of stay using neural networks on mimic iii
        data.
      In 2017 IEEE 15th Intl Conf on Dependable, Autonomic and Secure
        Computing, 15th Intl Conf on Pervasive Intelligence and Computing, 3rd Intl
        Conf on Big Data Intelligence and Computing and Cyber Science and Technology
        Congress (DASC/PiCom/DataCom/CyberSciTech), 1194–1201. IEEE.
      
      [Golmant et al.(2018)Golmant, Vemuri, Yao, Feinberg, Gholami,
        Rothauge, Mahoney, and Gonzalez]golmant2018computational
      Golmant, N.; Vemuri, N.; Yao, Z.; Feinberg, V.; Gholami, A.; Rothauge, K.;
        Mahoney, M. W.; and Gonzalez, J. 2018.
      On the computational inefficiency of large batch sizes for stochastic
        gradient descent.
      arXiv preprint arXiv:1811.12941 .
      
      [Goyal et al.(2017)Goyal, Dollár, Girshick, Noordhuis, Wesolowski,
        Kyrola, Tulloch, Jia, and He]goyal2017accurate
      Goyal, P.; Dollár, P.; Girshick, R.; Noordhuis, P.; Wesolowski, L.; Kyrola,
        A.; Tulloch, A.; Jia, Y.; and He, K. 2017.
      Accurate, large minibatch sgd: Training imagenet in 1 hour.
      arXiv preprint arXiv:1706.02677 .
      
      [Guo, Pasunuru, and Bansal(2019)]guo2019autosem
      Guo, H.; Pasunuru, R.; and Bansal, M. 2019.
      Autosem: Automatic task selection and mixing in multi-task learning.
      arXiv preprint arXiv:1904.04153 .
      
      [Guo et al.(2018)Guo, Haque, Huang, Yeung, and
        Fei-Fei]guo2018dynamic
      Guo, M.; Haque, A.; Huang, D.-A.; Yeung, S.; and Fei-Fei, L. 2018.
      Dynamic task prioritization for multitask learning.
      In Proceedings of the European Conference on Computer Vision
        (ECCV), 270–287.
      
      [Harutyunyan et al.(2017)Harutyunyan, Khachatrian, Kale, Steeg, and
        Galstyan]harutyunyan2017multitask
      Harutyunyan, H.; Khachatrian, H.; Kale, D. C.; Steeg, G. V.; and Galstyan, A.
        2017.
      Multitask learning and benchmarking with clinical time series data.
      arXiv preprint arXiv:1703.07771 .
      
      [Harutyunyan et al.(2019)Harutyunyan, Khachatrian, Kale, Ver Steeg,
        and Galstyan]harutyunyan2019multitask
      Harutyunyan, H.; Khachatrian, H.; Kale, D. C.; Ver Steeg, G.; and Galstyan, A.
        2019.
      Multitask learning and benchmarking with clinical time series data.
      Scientific data 6(1): 1–18.
      
      [Hochreiter and Schmidhuber(1997)]hochreiter1997flat
      Hochreiter, S.; and Schmidhuber, J. 1997.
      Flat minima.
      Neural Computation 9(1): 1–42.
      
      [Hoffer, Hubara, and Soudry(2017)]hoffer2017train
      Hoffer, E.; Hubara, I.; and Soudry, D. 2017.
      Train longer, generalize better: closing the generalization gap in
        large batch training of neural networks.
      In Advances in Neural Information Processing Systems,
        1731–1741.
      
      [Jean, Firat, and Johnson(2019)]jean2019adaptive
      Jean, S.; Firat, O.; and Johnson, M. 2019.
      Adaptive Scheduling for Multi-Task Learning.
      arXiv preprint arXiv:1909.06434 .
      
      [Johnson, Pollard, and Mark(2017)]johnson2017reproducibility
      Johnson, A. E.; Pollard, T. J.; and Mark, R. G. 2017.
      Reproducibility in critical care: a mortality prediction case study.
      In Machine Learning for Healthcare Conference, 361–376.
      
      [Johnson et al.(2016)Johnson, Pollard, Shen, Li-wei, Feng, Ghassemi,
        Moody, Szolovits, Celi, and Mark]johnson2016mimic
      Johnson, A. E.; Pollard, T. J.; Shen, L.; Li-wei, H. L.; Feng, M.; Ghassemi,
        M.; Moody, B.; Szolovits, P.; Celi, L. A.; and Mark, R. G. 2016.
      MIMIC-III, a freely accessible critical care database.
      Scientific data 3: 160035.
      
      [Kendall, Gal, and Cipolla(2018)]kendall2018multi
      Kendall, A.; Gal, Y.; and Cipolla, R. 2018.
      Multi-task learning using uncertainty to weigh losses for scene
        geometry and semantics.
      In Proceedings of the IEEE conference on computer vision and
        pattern recognition, 7482–7491.
      
      [Keskar et al.(2016)Keskar, Mudigere, Nocedal, Smelyanskiy, and
        Tang]keskar2016large
      Keskar, N. S.; Mudigere, D.; Nocedal, J.; Smelyanskiy, M.; and Tang, P. T. P.
        2016.
      On large-batch training for deep learning: Generalization gap and
        sharp minima.
      arXiv preprint arXiv:1609.04836 .
      
      [Kingma and Ba(2014)]kingma2014adam
      Kingma, D. P.; and Ba, J. 2014.
      Adam: A method for stochastic optimization.
      arXiv preprint arXiv:1412.6980 .
      
      [Li, Liu, and Wang(2019)]li2019gradient
      Li, B.; Liu, Y.; and Wang, X. 2019.
      Gradient harmonized single-stage detector.
      In Proceedings of the AAAI Conference on Artificial
        Intelligence, volume 33, 8577–8584.
      
      [Lin et al.(2017)Lin, Goyal, Girshick, He, and
        Dollár]lin2017focal
      Lin, T.-Y.; Goyal, P.; Girshick, R.; He, K.; and Dollár, P. 2017.
      Focal loss for dense object detection.
      In Proceedings of the IEEE international conference on computer
        vision, 2980–2988.
      
      [Lin et al.(2019)Lin, Baweja, Kantor, and Held]lin2019adaptive
      Lin, X.; Baweja, H.; Kantor, G.; and Held, D. 2019.
      Adaptive Auxiliary Task Weighting for Reinforcement Learning.
      In Advances in Neural Information Processing Systems,
        4773–4784.
      
      [Maurer, Pontil, and Romera-Paredes(2016)]maurer2016benefit
      Maurer, A.; Pontil, M.; and Romera-Paredes, B. 2016.
      The benefit of multitask representation learning.
      The Journal of Machine Learning Research 17(1): 2853–2884.
      
      [Mayaud et al.(2013)Mayaud, Lai, Clifford, Tarassenko, Celi, and
        Annane]mayaud2013dynamic
      Mayaud, L.; Lai, P. S.; Clifford, G. D.; Tarassenko, L.; Celi, L. A. G.; and
        Annane, D. 2013.
      Dynamic data during hypotensive episode improves mortality
        predictions among patients with sepsis and hypotension.
      Critical care medicine 41(4): 954.
      
      [McCandlish et al.(2018)McCandlish, Kaplan, Amodei, and
        Team]mccandlish2018empirical
      McCandlish, S.; Kaplan, J.; Amodei, D.; and Team, O. D. 2018.
      An empirical model of large-batch training.
      arXiv preprint arXiv:1812.06162 .
      
      [Nemati et al.(2018)Nemati, Holder, Razmi, Stanley, Clifford, and
        Buchman]nemati2018interpretable
      Nemati, S.; Holder, A.; Razmi, F.; Stanley, M. D.; Clifford, G. D.; and
        Buchman, T. G. 2018.
      An interpretable machine learning model for accurate prediction of
        sepsis in the ICU.
      Critical care medicine 46(4): 547.
      
      [Oksuz et al.(2020)Oksuz, Cam, Kalkan, and Akbas]oksuz2020imbalance
      Oksuz, K.; Cam, B. C.; Kalkan, S.; and Akbas, E. 2020.
      Imbalance problems in object detection: A review.
      IEEE Transactions on Pattern Analysis and Machine Intelligence
        .
      
      [Pang et al.(2019)Pang, Chen, Shi, Feng, Ouyang, and
        Lin]pang2019libra
      Pang, J.; Chen, K.; Shi, J.; Feng, H.; Ouyang, W.; and Lin, D. 2019.
      Libra r-cnn: Towards balanced learning for object detection.
      In Proceedings of the IEEE Conference on Computer Vision and
        Pattern Recognition, 821–830.
      
      [Purushotham et al.(2017)Purushotham, Meng, Che, and
        Liu]purushotham2017benchmark
      Purushotham, S.; Meng, C.; Che, Z.; and Liu, Y. 2017.
      Benchmark of deep learning models on large healthcare mimic datasets.
      arXiv preprint arXiv:1710.08531 .
      
      [Rai and Daumé III(2010)]rai2010infinite
      Rai, P.; and Daumé III, H. 2010.
      Infinite predictor subspace models for multitask learning.
      In Proceedings of the Thirteenth International Conference on
        Artificial Intelligence and Statistics, 613–620.
      
      [Romera-Paredes et al.(2012)Romera-Paredes, Argyriou, Berthouze, and
        Pontil]romera2012exploiting
      Romera-Paredes, B.; Argyriou, A.; Berthouze, N.; and Pontil, M. 2012.
      Exploiting unrelated tasks in multi-task learning.
      In International conference on artificial intelligence and
        statistics, 951–959.
      
      [Ruder(2016)]ruder2016overview
      Ruder, S. 2016.
      An overview of gradient descent optimization algorithms.
      arXiv preprint arXiv:1609.04747 .
      
      [Ruder et al.(2019)Ruder, Bingel, Augenstein, and
        Søgaard]ruder2019latent
      Ruder, S.; Bingel, J.; Augenstein, I.; and Søgaard, A. 2019.
      Latent multi-task architecture learning.
      In Proceedings of the AAAI Conference on Artificial
        Intelligence, volume 33, 4822–4829.
      
      [Saeed et al.(2011)Saeed, Villarroel, Reisner, Clifford, Lehman,
        Moody, Heldt, Kyaw, Moody, and Mark]saeed2011multiparameter
      Saeed, M.; Villarroel, M.; Reisner, A. T.; Clifford, G.; Lehman, L.-W.; Moody,
        G.; Heldt, T.; Kyaw, T. H.; Moody, B.; and Mark, R. G. 2011.
      Multiparameter Intelligent Monitoring in Intensive Care II
        (MIMIC-II): a public-access intensive care unit database.
      Critical care medicine 39(5): 952.
      
      [Sener and Koltun(2018)]sener2018multi
      Sener, O.; and Koltun, V. 2018.
      Multi-task learning as multi-objective optimization.
      In Advances in Neural Information Processing Systems,
        527–538.
      
      [Smith and Le(2017)]smith2017bayesian
      Smith, S. L.; and Le, Q. V. 2017.
      A bayesian perspective on generalization and stochastic gradient
        descent.
      arXiv preprint arXiv:1710.06451 .
      
      [Suresh, Gong, and Guttag(2018)]suresh2018learning
      Suresh, H.; Gong, J. J.; and Guttag, J. V. 2018.
      Learning tasks for multitask learning: Heterogenous patient
        populations in the icu.
      In Proceedings of the 24th ACM SIGKDD International Conference
        on Knowledge Discovery & Data Mining, 802–810.
      
      [Sutton(1992)]sutton1992adapting
      Sutton, R. S. 1992.
      Adapting bias by gradient descent: An incremental version of
        delta-bar-delta.
      In AAAI, 171–176.
      
      [Yang et al.(2018)Yang, Zhang, Yu, Cai, and Luo]yang2018end
      Yang, Z.; Zhang, Y.; Yu, J.; Cai, J.; and Luo, J. 2018.
      End-to-end multi-modal multi-task vehicle control for self-driving
        cars with visual perceptions.
      In 2018 24th International Conference on Pattern Recognition
        (ICPR), 2289–2294. IEEE.
      
      [Yu et al.(2020)Yu, Kumar, Gupta, Levine, Hausman, and
        Finn]yu2020gradient
      Yu, T.; Kumar, S.; Gupta, A.; Levine, S.; Hausman, K.; and Finn, C. 2020.
      Gradient surgery for multi-task learning.
      arXiv preprint arXiv:2001.06782 .
      
      [Zhang et al.(2018)Zhang, Liao, Rakhlin, Miranda, Golowich, and
        Poggio]zhang2018theory
      Zhang, C.; Liao, Q.; Rakhlin, A.; Miranda, B.; Golowich, N.; and Poggio, T.
        2018.
      Theory of deep learning IIb: Optimization properties of SGD.
      arXiv preprint arXiv:1801.02254 .
      
      [Zhang et al.(2014)Zhang, Luo, Loy, and Tang]zhang2014facial
      Zhang, Z.; Luo, P.; Loy, C. C.; and Tang, X. 2014.
      Facial landmark detection by deep multi-task learning.
      In European conference on computer vision, 94–108. Springer.
      
      
  
  
  