Year: 2021

Abstract:
We present a novel gradient-based multi-task learning (MTL) approach that balances training in multi-task systems by aligning the independent components of the training objective. In contrast to state-of-the-art MTL approaches, our method is stable and preserves the ratio of highly correlated tasks gradients. The method is scalable, reduces overfitting, and can seamlessly handle multi-task objectives with a large difference in gradient magnitudes. We demonstrate the effectiveness of the proposed approach on a variety of MTL problems including digit classification, multi-label image classification, camera relocalization, and scene understanding. Our approach performs favourably compared to other gradient-based adaptive balancing methods, and its performance is backed up by theoretical analysis.

heading: 1 INTRODUCTION
text: In multi-task learning (MTL, e.g., Caruana, 1993; Doersch & Zisserman, 2017), several tasks are optimized simultaneously using a single model and by leveraging shared information across tasks to improve generalization and boost performance for all tasks. MTL is a key component in real-world applications, such as natural language processing (Liu et al., 2019c), computer vision (Kendall et al., 2015; Kokkinos, 2017), and reinforcement learning (Parisotto et al., 2016; Teh et al., 2017). However, these systems are difficult to train since different tasks need to be properly balanced, which is challenging when the multi-task gradient is dominated by the gradient of one of the tasks.
MTL approaches have demonstrated rapid progress, and in general, they can be divided into two categories. The first group of methods aims at designing multi-task neural network architectures which should learn both task-shared and task-specific features by using hard and soft parameter sharing (Liu et al., 2019b; Lu et al.; Houlsby et al., 2019; Pfeiffer et al., 2021). The methods in the second group focus on multi-task optimization strategies. To learn all tasks with equal importance, some methods use gradient replacement (Yu et al., 2020; Lopez-Paz & Ranzato, 2017) or balance the individual loss functions with weights sampled from a uniform distribution, learned by using task-dependant uncertainty (Kendall et al., 2018), determined by utilising the learning speed for the task (Liu et al., 2019a;b; Lee & Son, 2020), or obtained by minimizing an auxiliary loss (Chen et al., 2018). However, combining individual loss functions into one using a weighted average can be hard and prone to errors if some task gradients are in conflict with each other. To alleviate this limitation, a variety of multi-objective optimization methods have been proposed (Désidéri, 2012; Sener & Koltun, 2018). They aim to find Pareto-stationary solutions for which the performance of any of the objectives can only be improved by deteriorating performance of some other objective. Although these methods demonstrate good performance, their main limitation is local convergence. Specifically, if the objective function is non-convex, which is typical in deep learning, the methods converge to a poor Pareto-stationary solution. Moreover, there is a gap between optimization and generalization performance as proven by Lin et al. (2019).
This work hinges on the realization that improving upon the overall performance in multi-task learning can be formulated by conservatism in the individual learning steps of uncorrelated losses. We formulate an approach that aims to guarantee consistent improvement over dynamically constructed linearly independent tasks at every step. This results in a multi-objective optimization algorithm that automatically balances training in multi-task systems by aligning orthogonal components of the objective. We address the problem of a large difference in the task gradient magnitudes through the condition number of the gradient matrix. This number defines the dominance rate in the multi-
objective task and is obtained as the ratio of maximum and minimum singular values. Fig. 1 illustrates the proposed approach and shows its performance compared to other gradient-based MTL methods when the difference in the task gradient magnitudes is large.
The contributions of this paper are summarized as follows. (i) We present a novel gradient-based MTL approach that dynamically constructs consistent losses and seamlessly handles the problem of a large difference in task gradients. (ii) We provide a computationally efficient approximation for the approach with constant time complexity and which is directly applicable to large-scale problems. (iii) We formulate a new metric, the dominance rate, as an interpretable way of establishing the maximum possible dominance in the gradient system. Additionally, we perform an extensive evaluation of the proposed approach on various MTL benchmarks and demonstrate that our method outperforms prior state-of-the-art MTL methods.

heading: 2 RELATED WORK
text: A multi-task formulation is well motivated and ubiquitous to many application domains such as computer vision (Bilen & Vedaldi, 2016; Kokkinos, 2017; Zamir et al., 2018; Nekrasov et al., 2019; Kendall et al., 2018), natural language processing (Collobert & Weston, 2008; Luong et al., 2015; Dong et al., 2015), speech processing (Seltzer & Droppo, 2013), and robotics (Wulfmeier et al., 2020; Levine et al., 2016). In this work we focus on gradient-based malti-task learning strategies and refer the interested reader to comprehensive surveys by Caruana (1993), Ruder (2017), and Crawshaw (2020) for a more detailed overview. One way to ease multi-task optimization is to balance the individual loss functions for different tasks. Prior methods to simultaneously learning multiple tasks parameterize the total objective as a weighted sum of task-specific loss functions whose weighting coefficients are manually tuned (Kendall et al., 2015; Laskar et al., 2017). However, using grid search for an optimal weighting is computationally inefficient especially when the number of tasks is large. To address this limitation, Kendall et al. (2018) utilize homoscedastic uncertainty of each task to derive a weighted multi-task loss. Chen et al. (2018) propose a heuristic to tune loss weights based on task gradient magnitudes.
One of the main challenges in MTL is the presence of conflicting gradients when tasks have gradients pointing in opposing directions, i.e. negative transfer (Lee et al., 2018). Among other approaches to mitigate the problem of negative transfer, the methods based on explicit gradient modulation have shown great potential (Lopez-Paz & Ranzato, 2017; Yu et al., 2020; Chaudhry et al., 2019). The core idea is to replace a task gradient which conflicts with some other task by a modified, non-conflicting, gradient. Yu et al. (2020) propose a ‘gradient surgery’ technique which makes the new gradient non-conflicting with the average of task gradients at the previous step. However, it requires checking multiple gradients for conflicts at each update step leading to a large computational overhead.
On the other hand, multi-objective optimization (MOO) is the process of optimizing a set of possibly contrasting objectives simultaneously (Fliege & Svaiter, 2000; Schäffler et al., 2002; Désidéri, 2012;
Peitz & Dellnitz, 2018; Poirion et al., 2017). These methods find a direction that decreases all objectives relying on multi-objective Karush–Kuhn–Tucker (KKT) conditions (Kuhn & Tucker, 1951). The work by Sener & Koltun (2018) extends the classical MGDA (Désidéri, 2012) method to a form that scales well to high-dimensional problems. However, it converges to a non-optimal Pareto-stationary solution if the objective function is non-convex. We empirically compare our approach to a number of MOO methods and demonstrate its superior performance.

heading: 3 PROBLEM FORMULATION AND NOTATION
text: Multi-task learning is concerned with finding the set of parameters for a parametric model that correspond to a high average performance across all tasks. More formally, we consider a MTL problem defined over the input space X and a series of targets {Y1, . . . ,YT }, where T is the number of tasks. We consider a parametric model f(x;θsh,θ1, . . . ,θT ) with shared, θsh, and task-specific parameters, θi. Each task has a corresponding differentiable loss function Li(ŷi,yi) : Yi × Yi → R. Moreover, we can include prior knowledge on user preferences p(Y) in the model. The MTL problem can be formulated as an empirical risk minimisation (ERM) problem on a data set {(xj ,y1j , . . . ,yTj )}Nj=1 of i.i.d. data points:
min θsh
θ1,...,θT
EYi∼p(Y) [ Li(θsh,θi) ] = min
θsh
θ1,...,θT
T∑ i=1 wiLi(θsh,θi) = min θsh
θ1,...,θT
T∑ i=1 wiLi(θ = θsh). (1)
The vector of preferences w defines the importance of each task: w = 1 for MTL; w 6= 1 for auxiliary task learning. We consider shared parameters, unless otherwise specified. Since the gradients of θi and θsh are linearly independent pairwise, task-specific parameters have a single gradient, and thus can be omitted without loss of generality.

heading: 4 INDEPENDENT COMPONENT ALIGNMENT
text: One of the key challenges in MTL is overfitting to one of the tasks caused by the dominance of the gradient of some linear component in the loss function. In practice, the overfitting might be defined by a dominating direction which has the fastest learning rate compared to others. This dominance can be caused by, for example, gradient interference between the individual loss functions or a significant difference in their magnitudes. Such a situation is illustrated in Fig. 2a, where the gradient for task 2 dominates task 1. Recent works on the interactions of individual tasks (Zamir et al., 2018; 2020) demonstrate that tasks may have a common component. This means that the uniqueness of a particular task is determined by another, unique, component independent to the common one. In this regard, the interference of gradients leads to a faster training rate of the joint subtask reducing the efficiency of learning unique components for each task. Therefore, we need to analyze the independent components of the gradient system, rather than the individual gradients themselves. We define dominance in the system of gradients as follows: Definition 1 (Dominance rate). The rate of dominance of the system of loss functions is the condition number of the matrix composed of the gradients of these functions over the shared parameters:
D(L1(θ), . . . ,LT (θ)) = D(Gθ) = σmax(Gθ)
σmin(Gθ) , where Gθ =
( ∇θL1(θ), . . . ,∇θLT (θ) ) . (2)
This definition is based on the singular value decomposition (SVD) and has a geometric interpretation. The dominance coefficient defines the maximum possible dominance in the gradient system and is obtained as the ratio of the maximum and minimum learning rates of the orthonormal components of the gradient system (as in Fig. 2a). Recently, Wang et al. (2020) propose a metric, OGR to quantitatively estimate the significance of overfitting of multi-modal networks for classification. Although related, the metric is obtained by using an assumption that target distribution is well approximated by a validation data set. In contrast, our measure does not require such a strong prior and solely based on gradients with respect to the training set.

heading: 4.1 MODEL-AGNOSTIC ALIGNMENT
text: For effective training over all tasks, the system of gradients of loss functions should be consistent and have a small dominance rate. We take this as a design principle, and aim for avoiding domination and
inconsistency altogether. We propose constructing a new (consistent) gradient system that best (under conditions that follow) describes the given directions, but whose dominance coefficient is equal to minimum dominance over the task gradients, i.e. unit dominance. More formally, we define a new aligned gradient as follows.
Definition 2 (θ-aligned). We define a θ-aligned gradient as ∇θaθ L(θ) := σmin(Gθ)Ĝθw, where vector w denotes a pre-selected task preference and Ĝθ is the solution to the Procrustes problem:
Ĝθ = arg min Q ‖Gθ −Q‖F , s.t. Q>Q = I. (3)
In order to calculate the θ-aligned gradient, the solution to the general Procrustes problem is needed. Schönemann (1966) shows that the problem has an analytical solution, namely:
Ĝθ = UV >, where U ,V : Gθ = UΣV >(given by a SVD). (4)
The analytical solution allows us to avoid the costly procedure for finding the minimum, which reduces the computational cost of our approach. Additionally, this leads to the dominance rate of the resulting matrix being equal to unit dominance: D(σmin(Gθ)Ĝθ) = σminσmin = 1, i.e. we have eliminated the possibility of uneven learning of different tasks. Fig. 2c shows a sketch of the geometric intuition of the aligned gradients after the alignment step.
At this stage we aim to verify that standard gradient descent with θ-aligned gradients corresponds to a sensible optimization procedure. We provide the following analysis, which states the convergence for general non-convex loss functions with standard assumptions on regularity. For a proof and more theoretical results, see App. B.
Theorem 1. Assume L1(θ), . . . ,LT (θ) are lower bounded continuously differentiable functions with Lipschitz continuous gradients with Λ > 0. Assume that the total loss function L(θ) is a weighted sum of these functions with a pre-selected vector of preferences w. Gradient descent with step size α ≤ 1Λ and θ-aligned gradient converges if the initial rate of dominance is finite, D(Gθ) < C, and task preferences vector w is consistent with the gradient system, ‖V w‖ ≥ ε.

heading: 4.2 EFFICIENT ALIGNMENT FOR ENCODER–DECODER NETWORKS
text: The main limitation of the θ-aligned approach is the requirement of multiple backward passes through the shared part of the model to calculate the gradient matrix. The backward passes are computationally demanding, which results in linear scaling of the training time with respect to the number of tasks. Therefore, it is challenging to apply the method to large-scale learning problems when the number of tasks is high. To address this limitation, we propose an efficient approximation to the θ-aligned approach, obtained under stronger assumptions on the model architecture.
Suppose that the class of hypotheses consists of factorizable functions, i.e. f(x;θsh,θ1, . . . ,θT ) = ( f(g(x,θsh),θ1), . . . , f(g(x,θsh),θT ) ) (5)
where g(·,θsh) is a representation function common to all tasks and f(·,θi) is a function that implements model functions for a specific task. This assumption corresponds to encoder–decoder type neural network architectures which are widely used in computer vision applications (e.g., Kendall et al., 2018; Sener & Koltun, 2018; Zamir et al., 2018).
We consider the original Procrustes problem (Def. 2) with new assumptions. We denote the encoded representation as Z = g(x, θsh) := g(x, θ). Suppose the Jacobian J = ∂Z ∂θ affects equally all task gradients, i.e. J
>J = λI . This assumption is strong, but it is supported by empirical evidence (as also sketched out in Fig. 3) and holds with some accuracy for a wide class of networks.
Following these assumptions, the original problem is formulated as follows:
‖Gθ−Q‖2F = ‖J>Gθ−J>Q‖2F = λ2‖GZ−P ‖2F , s.t. P>P = I. (6)
In practice, computing the matrixGZ requires only gradients over the latent representation, which can be computed
in one backward pass. Therefore, we define an efficient method of gradient alignment for encoder– decoder networks as follows. Definition 3 (Z-aligned). We define a Z-aligned gradient as ∇Zaθ L(θ) := σmin(GZ)∂Z∂θ ĜZw, where w denotes a pre-selected task preference and ĜZ is the solution to the Procrustes problem:
ĜZ = arg min P ‖GZ − P ‖F , s.t. P>P = I. (7)
The value of the rate of dominance in this case will strongly depend on the condition for the Jacobian. The value of the growth rate λ does not affect the value of the dominance coefficient:
D(σmin(GZ)JĜZ) ≈ λσmin(GZ)
λσmin(GZ) = 1. (8)
Algorithm 1 Gradient descent with alignment both for the θ-aligned and Z-aligned variants. 1: for i = 1 to T do 2: θit+1 = θ i t − η∇θiL̂i(θsh,θi) . Gradient descent on task-specific parameters
3: end for 4: if θ-alignment then 5: ComputeGθ = ( ∇θshL1(θ), . . . ,∇θshLT (θ) ) 6: Ĝθ, σ = PROCRUSTESSOLVER(Gθ) . Solve problem in Def. 2 7: ∇θL(θ) = σĜθw . Compute θ-aligned gradient 8: else if Z-alignment then 9: ComputeGZ = ( ∇ZL1(θ), . . . ,∇ZLT (θ)
) 10: ĜZ , σ = PROCRUSTESSOLVER(GZ) . Solve problem in Def. 3 11: ∇θL(θ) = σJĜZw . Compute Z-aligned gradient 12: end if 13: θsht+1 = θ sh t − η∇θL(θ) . Gradient descent on the shared parameters
14: procedure PROCRUSTESSOLVER(G) 15: FactorizeG,G = UΣV > . Compute SVD decomposition 16: Ĝ = UV > . Analytical solution of Procrustes problem 17: Compute σ = mini Σii . Find a minimal learning rate 18: return Ĝ, σ 19: end procedure
Even though our observations confirm that the Jacobian often satisfies our assumptions (see also Sec. 5), we prove (proof of the theorem below in App. B) that gradient descent with Z-aligned gradients will converge in a milder setup even if our requirements are not met. Our formulation increases the stability of the proposed approach because in worst case it behaves similarly to gradient descent without alignment.
Theorem 2. Assume L1(θ), . . . ,LT (θ) are lower bounded continuously differentiable functions with Lipschitz continuous gradients with Λ > 0. Assume that the total loss function L(θ) is a weighted sum of these functions with a pre-selected vector of preferencesw. The gradient descent with step size α ≤ 1Λ and a Z-aligned gradient converges if the Jacobian J = ∂Z ∂θ is full rank, the initial rate of dominance is finite, i.e. D(GZ) < Cσ and D(J) < Cλ, and task preferences vector w is consistent with the gradient system, ‖V w‖ ≥ ε.
Alg. 1 summarizes the approaches for the θ-aligned and the Z-aligned multi-task learning methods.

heading: 5 EXPERIMENTS
text: We empirically demonstrate the effectiveness of the proposed approach on a number of different multi-task learning benchmarks. Specifically, we tackle the problems of scene understanding, multidigit, and multi-label classification by evaluating our method on the CITYSCAPES (Cordts et al., 2016), MULTIMNIST (Sabour et al., 2017), and CELEBA (Liu et al., 2015) data sets, respectively. Furthermore, we experiment with image-based localization by jointly predicting camera orientation and translation on the 7SCENES data set (Shotton et al., 2013).
Baselines. In this work, we consider the following MTL baseline methods: (1) uniform scaling: optimizing a uniformly weighted sum of individual task objectives, i.e. 1T ∑ t Lt; (2) Uncertainty (Kendall et al., 2018): using uncertainty to find weighting coefficients for each task; (3) MGDA (Désidéri, 2012): using the gradients of each task to solve an optimization problem for an update; (4) MGDA-UB (Sener & Koltun, 2018): optimizing an upper bound for the MGDA optimization objective; (5) GradNorm (Chen et al., 2018): normalizing the gradients to balance the learning of multiple tasks, and (6) PCGrad (Yu et al., 2020): performing gradient projection to avoid
the negative interactions between tasks gradients. The proposed approach and the baseline methods are implemented using the PyTorch framework (Paszke et al., 2019).
Analysis of computational complexity. The main computational cost of all algorithms depends on the number of backward passes over the shared parameters for each task and for each iteration of the gradient descent. We denote the number of tasks as T and the number of iterations in the Frank–Wolfe solver (Jaggi, 2013; Sener & Koltun, 2018) as K. By Def. 2, the θ-aligned gradient requires only the precomputed matrix Gθ and therefore, it performs T backward passes on one step. In contrast, the Z-aligned gradient can be computed in a single pass, which is the best possible case. The computational cost for each MTL method is provided in Table 1. We consider the time complexity of the SVD decomposition constant.

heading: 5.1 SYNTHETIC MULTI-TASK PROBLEM
text: To illustrate the proposed method on a simple problem, we consider a multi-task optimization objective consisting of two tasks presented in Fig. 1 which was originally proposed by Yu et al. (2020). Specifically, the multi-task gradient of the objective is dominated by the gradient of one of the tasks leading to suboptimal training dynamics. As a result, the standard optimizers (e.g., Adam) improve only the dominating task (the magenta trajectory in Fig. 1c) and struggle to optimize the whole objective. In contrast, the proposed approach is more stable and outperforms MGDA (Désidéri, 2012) and PCGrad (Yu et al., 2020). Moreover, it demonstrates favourable convergence rate compared to the baselines: 5k and 500k iterations are required for our method and PCGrad, respectively.

heading: 5.2 MULTI-TASK CLASSIFICATION
text: We tackle the problem of multi-task classification and evaluate the proposed method on MULTIMNIST digit classification and celebrity face image data in CELEBA.
MULTIMNIST. Following the same procedure proposed by Sener & Koltun (2018), we convert digit classification into a multi-digit classification problem. Specifically, for each image of the original MNIST data set, we randomly sample a different image from a uniform distribution. The two images are then overlaid in such a way that one of the images is put at the top-left and the other one at the bottom-right. Similarly to Sener & Koltun (2018); Yu et al. (2020), we generate 60k such images and directly evaluate the proposed approach and the baselines. The resulting task in the multi-task classification problem is to classify the digit on the top left (task-L) and bottom right (task-R), respectively. The results are summarized in Table 2. The proposed MTL approach θ-aligned surpasses all baselines by a noticeable margin: up to +0.06% for Task-L and +0.16% for Task-R compared to the best performing contender. To emphasize the significance of the gradient dominance problem and its influence on the results, we report the coefficient of dominance (Def. 1) for each method, which shows large variability for the baselines.
CELEBA. We try our method on a multi-label classification problem by using the CELEBA data set (Liu et al., 2015), which consists of 200k face images annotated with 40 attributes. Each attribute can be obtained by solving a binary classification problem and thus we convert it to 40- way MTL following (Sener & Koltun, 2018; Yu et al., 2020). For all of the methods we use the same CNN architecture based on ResNet-18 as proposed by Sener & Koltun (2018). The binary classification error averaged over all 40 attributes is used for evaluating performance and is reported in Table 3 (mean ± std over 10 random seeds). Our approach achieves 91.36% classification accuracy, which outperforms other strong methods such as MGDA-UB (Sener & Koltun, 2018) and Uncertainty (Kendall et al., 2018), indicating the effectiveness of the proposed method when the number of tasks is large. The percentage of error per attribute is illustrated in Fig. 4.

heading: 5.3 SCENE UNDERSTANDING
text: We also apply the proposed MTL approach to scene understanding. Specifically, we aim to jointly solve the problem of semantic segmentation, instance segmentation, and depth estimation. We follow Sener & Koltun (2018) and use the CITYSCAPES data sets (Cordts et al., 2016) consisting of video frames captured in the streets of urban cities and labelled with instance, semantic segmentations, and depth maps. Similarly to Kendall et al. (2018), we use a CNN model with a UNet architecture comprising a shared encoder based on ResNet-50 and three independent task-specific decoders. In our experiments, we resize all training and validation images to 256×512 and use standard pixel-wise loss functions for each task: negative log likelihood loss for semantic segmentation and L1 loss for instance segmentation and depth estimation. In Table 3, we benchmark our approach in all three tasks and show that the proposed method performs favourably compared to other state-of-the-art MTL methods.

heading: 5.4 CAMERA RELOCALIZATION
text: Finally, we evaluate the performance of the proposed MTL approach on the task of image-based localization. Image-based localization, or camera relocalization refers to the problem of estimating the 6 degree-of-freedom (DoF) camera pose from visual data with respect to a known environment. It has been shown that camera relocalization can be solved by casting it as a regression problem (Kendall et al., 2018; Laskar et al., 2017; Melekhov et al., 2017; Walch et al., 2017; Noha et al., 2018), where a 7-dimensional camera pose vector p = [t, r] is directly estimated by a CNN. The camera pose vector p consists of a translation component t = [t1, t2, t3] and an orientation component r = [q1, q2, q3, q4], represented by a quaternion. Following Kendall et al. (2018), we utilize the Microsoft 7SCENES (Shotton et al., 2013) data set which is a common benchmark for indoor pose regression methods. The data set comprises RGB-D image sequences of seven different indoor scenes from a handheld Kinect camera, and each scene provides training and testing image sequences consisting of 1000–7000 frames at 640×480 resolution. Similar to Walch et al. (2017); Laskar et al. (2017), we use ResNet-34 as a backbone network in all our experiments to make the comparison between the MTL methods fair. For each method, we compute the median error of camera translation, t, and orientation, r, per scene. We also report the growth rate,R, as a measure of improvement of multi-task learning methods over the uniform scaling approach. It is defined asR = 1− PiP where Pi is the mean translation (orientation) error of a MTL method; P is the mean translation (orientation)
error of the uniform scaling method. The results are provided in Table 4. In terms of camera orientation error, the proposed Z-aligned approach is superior to other MTL methods for the majority of scenes. Regarding camera translation accuracy, Z-aligned can outperform MGDA-UB (Sener & Koltun, 2018) and improve the averaged translation error by 0.21 m showing comparable performance to PCGrad (Yu et al., 2020). See App. D.2 for details on the evaluation pipeline and additional results.

heading: 6 DISCUSSIONS AND CONCLUSIONS
text: We have presented a new approach to gradient-based multi-task learning (MTL) that balances training in multi-task systems by aligning the independent components of the gradient of the training objective. In Sec. 4.1, we derived the ‘θ-aligned’ approach and the conditions for it to converge (details in App. B). To ensure practicality, we extended our approach to incorporate heuristics motivated by common encode–decoder networks, which was formulated into the ‘Z-aligned’ approach in Sec. 4.2. Additionally, we formulate a new metric, the dominance rate (Def. 1) that provides an interpretable way of establishing the maximum dominance in a gradient system and gives insights to pre-existing MTL approaches.
We also performed an extensive set of experiments for analyzing the practical benefits of our approach. To gain insights, we provided results for a toy benchmark problem and compared our approach to other recent gradient-based MTL approaches on MULTIMNIST to study the effect of the rate of dominance. As for more pure benchmarking, we provided results for MTL problems in image classification (CELEBA), scene understanding (CITYSCAPES), and camera relocalization (7SCENES), each being large-scale and well-known benchmarks in their respective domains. We show state-of-the-art performance on MULTIMNIST, CITYSCAPES, and 7SCENES, while the differences between methods remains small for CELEBA. The results are particularly interesting given the smaller computational requirements in the Z-aligned method (see Table 1).
A reference implementation of the methods presented in this paper is currently available as supplementary material to this submission.
REPRODUCIBILITY STATEMENT
This paper proposes a new method for multi-task learning, and the contributions of the paper are technical of nature. To back up the statements in the main paper, we have provided proofs for the given theorems in the appendix (included after the references). Additionally, we have also provided further theorems and proofs that help the reader assess the applicability of the proposes methodology to other optimization setups.
We also include a code package as a supplementary zip file to this double-blind submission. This reference implementation helps understand technical details, and the included scripts can be used for reproducing the numerical experiments in the paper. Upon acceptance, the codes will be available as a GitHub repository under the MIT license.

heading: B PROOFS OF THEOREMS
text: Synopsis. In this theorem we prove that the worst case performance of the Z-aligned gradient is no worse than in standard gradient descent. The constraints mentioned in the theorem always hold in practice.
Lemma 1. Assume L(θ) to be continuously differentiable and ∇L(θ) to be Lipschitz continuous with Λ > 0. Then, the following restriction is valid for any descent with the step αr:
L(θt)− L(θt+1) ≥ α〈∇L(θt), r 〉 − α2Λ
2 ‖r‖2. (9)
Proof. First, we substitute δ = −αr. Then, according to the fundamental theorem of calculus we receive the equality:
L(θt+1)− L(θt) = L(θt + δ)− L(θt) = ∫ 1
0
〈∇L(θt + sδ), δ 〉ds. (10)
Adding and substracting the value 〈∇L(θt), δ 〉 = ∫ 1 0 〈∇L(θt), δ 〉ds leads to:
L(θt+1)− L(θt) = 〈∇L(θt), δ 〉+ ∫ 1
0
〈∇L(θt + sδ)−∇L(θt), δ 〉ds. (11)
Since the gradient satisfies the Lipschitz condition: ‖∇L(θt + sδ)−∇L(θt)‖ ≤ Λ‖θt + sδ − θt‖ and due to inequality 〈x, y 〉 ≤ ‖x‖‖y‖, we can transform the integral, make the reverse substitution and get the final result:
L(θt+1)− L(θt) ≤ −α〈∇L(θt), r 〉+ α2Λ‖r‖2 ∫ 1
0
sds. (12)
Theorem 1. Assume L1(θ), . . . ,LT (θ) are lower bounded continuously differentiable functions with Lipschitz continuous gradients with Λ > 0. Assume that the total loss function L(θ) is a weighted sum of these functions with a pre-selected vector of preferences w. Gradient descent with step size α ≤ 1Λ and θ-aligned gradient converges if:
1. Gθ = UΣV > (given by a SVD) and Σ = diag{σ1, . . . , σR}, where σ1σR < C.
2. ‖ projHV w‖ ≥ ε, whereHV = span(V ).
Proof. Following the assumptions we find that the cumulative loss function is continuously differentiable and has a Lipschitz continuous gradient (denote this constant as Λ) and therefore it satisfies Lemma 1:
L(θt)− L(θt+1) ≥ α〈∇θL(θt),∇θaθ L(θt) 〉 − α2Λ
2 ∥∥∥∇θaθ L(θt)∥∥∥2. (13) By definition of θ-aligned gradient we get:
〈∇θL(θt),∇θaθ L(θt) 〉 = R∑
r,r′=1
σrσRw >(vru>r )(ur′v>r′)w = σR R∑
r=1
σr(w >vr) 2, (14)
〈∇θaθ L(θi),∇θaθ L(θi) 〉 = R∑
r,r′=1
σ2Rw >(vru>r )(ur′v>r′)w = σ2R R∑
r=1
(w>vr) 2. (15)
Since α ≤ 1Λ and ∀r : σr σR > 1, we can rewrite Eq. (13):
L(θt)− L(θt+1) ≥ σ2R α
2 R∑ r=1 ( 2 σr σR − 1 )
︸ ︷︷ ︸ >1
( w>vr )2 ︸ ︷︷ ︸
=‖w‖2 cos2(w,vr)
> ασ2R
2
ε2 σ21 σ21 . (16)
Following the assumption σRσ1 > C. Moreover σ1 = maxx 6=0 ‖Gθx‖ ‖x‖ , therefore σ1 ≥ ‖∇θL(θt)‖ ‖w‖ . Then:
L(θt)− L(θt+1) > αε2C2
2‖w‖2 ‖∇θL(θt)‖2. (17)
The sequence of L(θt) is monotonically decreasing and bounded (under assumption) and therefore converges. Then L(θt) − L(θt+1) → 0 if t → ∞. Therefore, we have a local convergence of gradient descent:
‖∇θL(θt)‖2 < 2‖w‖2
αC2 2
( L(θt)− L(θt+1) ) → 0 as t→∞. (18)
The same estimate appears in case of gradient descent. Therefore, the convergence of the proposed method is similar to that of gradient descent.
Theorem 2. Assume L1(θ), . . . ,LT (θ) are lower bounded continuously differentiable functions with Lipschitz continuous gradients with Λ > 0. Assume that the total loss function L(θ) is a weighted sum of these functions with a pre-selected vector of preferencesw. The gradient descent with step size α ≤ 1Λ and a Z-aligned gradient converges if:
1. GZ = UΣV > (given by a SVD) and Σ = diag{σ1, . . . , σR}, where σ1σR < Cσ . 2. ‖ projHV w‖ ≥ ε, whereHV = span(V ).
3. J = ∂Z∂θ is full rank.
4. J = PSQ> (given by a SVD) and S = diag{λ1, . . . , λK}, where λ1λK < Cλ.
Proof. Following the assumptions we find that the cumulative loss function is continuously differentiable and has a Lipschitz continuous gradient (with constant Λ > 0) and therefore it satisfies Lemma 1:
L(θt)− L(θt+1) ≥ α〈∇θL(θt),∇Zaθ L(θ) 〉 − α2Λ
2 ‖∇Zaθ L(θ)‖2. (19)
Since α ≤ 1Λ , we can rewrite this inequality:
L(θt)− L(θt+1) ≥ α
2 〈2∇θL(θt)−∇Zaθ L(θt),∇Zaθ L(θt) 〉. (20)
Consider the left vector of a dot product in the right part of Eq. (20). By definition of Z-aligned gradient we get:
∇Zaθ L(θt) = R∑ r=1 σR(w >vr)Jur, (21)
2∇θL(θt)−∇Zaθ L(θt) = R∑ r=1 (2σr − σR)(w>vr)Jur. (22)
Following the assumption J>J is positive definite. Any positive definite matrix is congruent to a diagonal with positive and ordered eigenvalues on the main diagonal. Thus replacing all eigenvalues λ2i with the smallest one λ 2 K does not increase the right side of Eq. (19) resulting in the inequality:
〈2∇θL(θt)−∇Zaθ L(θt),∇Zaθ L(θt) 〉 ≥ σ2Rλ2K R∑ r=1 ( 2 σr σR − 1 )
︸ ︷︷ ︸ >1
( w>vr )2 ︸ ︷︷ ︸
=‖w‖2 cos2(w,vr)
. (23)
Therefore:
L(θt)− L(θt+1) ≥ αε2σ2Rλ 2 K
2σ21λ 2 1
σ21λ 2 1. (24)
Following the assumption σRσ1 > Cσ and λK λ1 > Cλ. Moreover σ1 = maxx 6=0 ‖GZx‖ ‖x‖ ≥ ‖∇ZL(θ)‖ ‖w‖ and λ1 = ‖J‖. Therefore, we get the final bound:
L(θt)− L(θt+1) ≥ αε2C2σC 2 λ
2‖w‖2 ‖GZw‖2‖J‖2 ≥
αε2C2σC 2 λ
2‖w‖2 ‖∇θL(θt)‖2. (25)
The sequence of L(θt) is monotonically decreasing and bounded (under assumption) and therefore converges. Then L(θt)− L(θt+1)→ 0 if t→∞. Therefore, we have local convergence of gradient descent:
‖∇θL(θt)‖2 < 2‖w‖2
αC2σC 2 λε 2
( L(θt)− L(θt+1) ) → 0 as t→∞. (26)

heading: C METHODS WITH MOMENTUM
text: Definition 4. A function L(x) : X → R is µ-strongly convex if for any x, y ∈ X
L(y) ≥ L(x) + 〈∇L(x), x− y 〉+ µ 2 ‖x− y‖2. (27)
Definition 5. A continuously differentiable function L(x) : X → R is Λ-smooth if for any x, y ∈ X :
‖∇L(x)−∇L(y)‖ ≤ Λ‖x− y‖. (28)
Lemma 2. A θ-aligned and Z-aligned gradient can be expressed as:
∇∗aθ L(θt) = T∑ i=1 ( T∑ k=1 wk [ R∑ r=1 σR σr virv k r ]) ∇θLi(θt). (29)
Proof. We consider only the Z-aligned case. The θ-aligned case may be proved using analogous logic. According to SVD theorem, U = GZV Σ−1 in matrix form that in vector form turns into ur =
1 σr GZvr. Then, by definition of a Z-aligned gradient:
σRĜZt = σR R∑ r=1 urv > r = GZt R∑ r=1 σR σr vrv > r . (30)
If we rewrite this dot product, then we get:
σRĜZt = ( T∑ i=1 [ R∑ r=1 σR σr virv 1 r ] ︸ ︷︷ ︸
ai1
∇ZLi(θt), . . . , T∑ i=1 [ R∑ r=1 σR σr virv T r ] ︸ ︷︷ ︸
aiT
∇ZLi(θt) ) . (31)
Since σRJĜZt = ( . . . , ∑ aikJ∇ZLi(θt), . . . ) = ( . . . , ∑ aik∇θLi(θt), . . . ) , then the final statement follows from the definition of the Z-aligned gradient:
∇Zaθ L(θt) = σRJĜZtw = T∑ k=1 T∑ i=1 wkaik∇θLi(θt) = T∑ i=1 ( T∑ k=1 wkaik ) ∇θLi(θt). (32)
Similarly, this equation can be derived for θ-aligned case:
∇θaθ L(θt) = σRĜθtw = T∑ k=1 T∑ i=1 wkâik∇θLi(θt) = T∑ i=1 ( T∑ k=1 wkâik ) ∇θLi(θt). (33)
Theorem 3. Assume L1(θ), . . . ,LT (θ) are strongly convex with constant µi > 0 and smooth with constant Λi > 0. Assume µ̂t = ∑T i=1 ( ∑T k=1 wkaik ) µi and Λ̂t = ∑T i=1 ( ∑T k=1 wkaik ) Λi. The heavy ball method with αt = 4 ( √ Λ̂t+ √ µ̂t)2 and βt = max{|1 − √ αtµ̂t|, |1 − √ αtΛ̂t|} and the θ-aligned or Z-aligned gradient will converge linearly to an optimal value L(θ∗) if:
1. (θ-aligned)Gθ = UΣV > (given by a SVD) and Σ = diag{σ1, . . . , σR}, where σRσ1 > Cσ .
2. (Z-aligned) GZ = UΣV > (given by a SVD) and Σ = diag{σ1, . . . , σR}, where σRσ1 > Cσ .
3. (Z-aligned) J is full rank.
4. (both) ‖projHV w‖ ≥ ε, whereHV = span(V ).
Proof. The conditions of the theorem ensure the existence of finite and non-zero θ and Z-aligned gradients. Due to existence Lemma 2 is valid. We will not separate θ-alignment and Z-alignment in this theorem due to similarity of the proving procedure. Instead we denote ∇∗aθ L(θt) as a general aligned gradient.
Consider the convergence of a heavy ball method using matrix notation following Sun (2015). The next point is computed by definition, i.e. θt+1 = θt − αt∇∗aθ L(θt) + βt(θt − θt−1). Then we can estimate the squared discrepancy between the current solution and the optimal point using the explicit expression for aligned gradient from Lemma 2:∥∥∥∥θt+1 − θ∗θt − θ∗ ∥∥∥∥ 2 = ∥∥∥∥[θt + βt(θt − θt−1)− θ∗θt − θ∗ ] − αt [ ∇∗aθ L(θt) 0 ]∥∥∥∥ 2
= ∥∥∥∥[θt + βt(θt − θt−1)− θ∗θt − θ∗ ] − αt [∑T k=1 w ′ k∇θLk(θt) 0 ]∥∥∥∥ 2 , (34)
Note that for all gradients for some point Θkt on the line segment between θt and θ ∗ the following expression is true: ∇θLk(θt) = ∇2θLk(Θkt )(θt − θ∗), where ∇2θLk(Θkt ) is a Hessian of a function Lk. DenoteHt = ∑T k=1 w ′ k∇2θLk(Θkt ). Hence,Ht is a Hessian of a total loss function. Then we
get: ∥∥∥∥θt+1 − θ∗θt − θ∗ ∥∥∥∥
2
= ∥∥∥∥[θt + βt(θt − θt−1)− θ∗θt − θ∗ ] − αt [ Ht(θt − θ∗) 0 ]∥∥∥∥ 2
= ∥∥∥∥[(1 + βt)I − αtHt −βtII 0 ] [ θt − θ∗ θt−1 − θ∗ ]∥∥∥∥ 2
≤ ∥∥∥∥[(1 + βt)I − αtHt −βtII 0 ]∥∥∥∥ 2 ∥∥∥∥[ θt − θ∗θt−1 − θ∗ ]∥∥∥∥ 2 . (35)
By strong convexity and smoothness of Lk(θ) the eigenvalues of ∇2θLk(Θkt ) are in the interval [µk,Λk]. Hence eigenvalues of Ht are between µ̂t and Λ̂t. Therefore, following Lemma 3.1 introduced in Sun (2015), the L2 norm of a matrix in the right side of Eq. (35) is bounded:∥∥∥∥[(1 + βt)I − αtHt −βtII 0 ]∥∥∥∥ 2 ≤ max { |1− √ αtµ̂t|, |1− √ αtΛ̂t| } . (36)
Now, we can plug in step size αt = 4 ( √ Λ̂t+ √ µ̂t)2 and Eq. (36) into Eq. (35):∥∥∥∥θt+1 − θ∗θt − θ∗ ∥∥∥∥
2
≤ √ γt − 1√ γt + 1 ∥∥∥∥[ θt − θ∗θt−1 − θ∗ ]∥∥∥∥
2
≤ ∥∥∥∥[ θt − θ∗θt−1 − θ∗ ]∥∥∥∥ 2 , (37)
where γt = Λ̂tµ̂t . Hence the heavy ball method with this gradient will converge linearly to an optimal value L(θ∗).
D IMPLEMENTATION DETAILS
In this section, we provide additional experimental results and more details on the training pipelines. The proposed approach and all baseline methods are implemented in PyTorch (Paszke et al., 2019).
D.1 MULTI-LABEL CLASSIFICATION
In addition to the radar chart presented in the main text, we present more detailed results on the CELEBA data set here. Specifically, we report the binary classification accuracy of each attribute for the proposed methods and the baselines on the validation split. The results are summarized in Table 5. Following Sener & Koltun (2018) we use ResNet-18 with 40 separate fully-connected (FC) layers as task-specific functions. Table 7 shows the architecture used in the experiments in more detail. The images of the data set are resized to 64×64×3 as proposed in Sener & Koltun (2018). We empirically found that learning rate 10−3 gives better results. The learning rate is gradually decreased during training. All models are optimized using Adam solver and trained for 20 epochs with batch size 256.
D.2 CAMERA RELOCALIZATION
For image-based localization, we follow Kendall et al. (2015) and adapt ResNet-34 architecture by removing the last FC layer but keep the convolutional and pooling layers intact. As a decoder, we use two separate FC layers with 3 and 4 outputs to regress camera orientation and translation, respectively. The network architecture is presented in Table 7. The 7SCENES (Shotton et al., 2013) data set is used for training and evaluation. The data set consists of 640× 480 RGB-D images captured by a Kinect device and cover 7 different indoor scenes. As a preprocessing step, the images are re-scaled in such a way that the smaller image side is 256 pixels. All models were trained on random crops of 224×224 pixels and evaluated on central crops. We use batch size 128 and train for 120 epochs. The learning rate is initialized to 10−3 and decreased by 10 every 40 epochs. The Adam solver is used for optimization. Each method is trained 10 times with different random seeds for every scene and average performance is provided in Table 6.
Easy attributes
Attr. 0 92.39 93.05 93.09 94.02 93.96 93.51 93.69 Attr. 4 98.79 98.74 98.69 98.80 98.77 98.77 98.84 Attr. 5 94.51 95.47 95.30 95.72 95.42 95.52 95.36 Attr. 9 93.77 94.93 95.22 95.45 95.16 95.39 95.38 Attr. 10 96.12 96.12 96.20 96.40 96.07 95.77 96.04 Attr. 13 95.14 95.05 95.33 95.36 95.17 95.36 95.28 Attr. 14 96.03 96.22 96.09 96.43 96.24 96.42 96.23 Attr. 16 95.88 96.18 96.18 96.40 96.39 96.13 96.41 Attr. 17 97.27 97.67 97.19 97.91 97.76 97.70 97.78 Attr. 20 97.25 98.01 98.13 98.53 98.52 98.02 98.41 Attr. 21 89.83 92.60 93.01 93.72 93.57 93.54 93.32 Attr. 22 96.13 95.91 95.60 96.20 95.99 96.11 96.05 Attr. 24 94.62 95.00 95.15 95.75 95.64 95.65 95.46 Attr. 26 96.54 96.67 96.44 96.79 96.71 96.74 96.71 Attr. 28 94.18 94.42 94.40 94.60 94.47 94.50 94.43 Attr. 29 94.60 94.73 94.63 95.14 95.11 95.17 95.02 Attr. 30 96.00 96.56 95.91 96.79 96.86 97.02 96.94 Attr. 35 98.70 98.92 98.69 98.90 98.73 98.84 98.87 Attr. 38 96.19 96.13 95.71 96.51 96.33 96.00 96.34
Hard attributes
Attr. 1 82.09 82.67 82.08 84.91 84.27 84.91 84.79 Attr. 2 79.94 80.14 81.26 81.36 80.50 80.61 81.30 Attr. 3 82.89 82.86 83.87 84.06 83.45 83.74 83.75 Attr. 6 85.73 85.66 86.09 83.43 85.92 82.22 84.06 Attr. 7 81.76 81.84 82.68 82.70 82.48 82.84 82.56 Attr. 8 89.48 90.67 91.35 91.42 90.45 89.45 91.28 Attr. 11 81.95 84.32 84.89 85.58 85.11 83.73 85.31 Attr. 12 91.83 92.19 92.06 92.54 92.32 92.19 92.46 Attr. 18 90.64 91.02 91.77 92.14 90.82 91.21 91.67 Attr. 19 86.60 87.48 87.94 88.58 87.65 88.28 87.96 Attr. 23 93.39 92.87 92.29 92.75 92.89 93.41 92.36 Attr. 25 75.60 75.14 75.57 76.14 75.87 76.04 75.84 Attr. 27 75.82 75.72 77.26 77.62 76.43 77.25 77.24 Attr. 31 91.17 92.15 92.65 93.04 92.99 92.85 92.68 Attr. 32 80.66 81.54 82.99 83.44 83.11 82.97 83.34 Attr. 33 81.42 83.27 85.95 85.53 85.66 84.56 85.63 Attr. 34 86.37 87.29 89.55 90.48 89.17 90.31 90.32 Attr. 36 91.30 91.56 92.71 92.61 92.69 92.14 92.86 Attr. 37 88.35 88.15 88.25 88.90 87.68 88.82 88.38 Attr. 39 85.67 86.03 86.93 88.13 86.94 87.07 87.58

Review summaries:
The paper presents an approach for better multi task learning in the presence of objectives with a large difference in gradient magnitudes. They study the problem and define a "The rate of dominance" metric that is based on an eigenvalue ratio of the gradient tensors. they then use this dominance metric as a coefficient to better "balance" gradients across tasks and propose an elegant solution based on the  an approximate version that is only computed on the shared parameters. They report results on common multi-task settings, and either perform on par or outperform other recent methods.  ==== After rebuttal ====  I want to thank the authors for a complete and detailed response. The answers make sense to me, and I am still in favour of acceptance - I raise my score to Accept.  

The authors aim to perform gradient task balancing on MTL setting by aligning independent components of the training objective. They claim to be able to handle large differences in gradient magnitude. Not choosing the faster learning rate dominating direction prevent task overfitting. The main contribution is based on the analysis of the individual components of the gradients instead of the individual gradients using the dominance rate. Tasks should have small dominance coefficients. Through a series of experiments the authors show the proposed technique is competitive with other multi-objective optimization approaches.    

This paper proposes a gradient-based multi-task learning method to balance multi-task training by aligning the independent components of the training objective. Experiments conducted on several MTL problems are given. The method is claimed to be scalable, robust to overfitting, and able to seamlessly handle multi-task objectives with a large difference in gradient magnitudes.