Multi-Task Learning as Multi-Objective Optimization

Ozan Sener, Vladlen Koltun

Introduction

One of the most surprising results in statistics is Stein’s paradox. Stein (1956) showed that it is better to estimate the means of three or more Gaussian random variables using samples from all of them rather than estimating them separately, even when the Gaussians are independent. Stein’s paradox was an early motivation for multi-task learning (MTL) (Caruana, 1997), a learning paradigm in which data from multiple tasks is used with the hope to obtain superior performance over learning each task independently. Potential advantages of MTL go beyond the direct implications of Stein’s paradox, since even seemingly unrelated real world tasks have strong dependencies due to the shared processes that give rise to the data. For example, although autonomous driving and object manipulation are seemingly unrelated, the underlying data is governed by the same laws of optics, material properties, and dynamics. This motivates the use of multiple tasks as an inductive bias in learning systems.

A typical MTL system is given a collection of input points and sets of targets for various tasks per point. A common way to set up the inductive bias across tasks is to design a parametrized hypothesis class that shares some parameters across tasks. Typically, these parameters are learned by solving an optimization problem that minimizes a weighted sum of the empirical risk for each task. However, the linear-combination formulation is only sensible when there is a parameter set that is effective across all tasks. In other words, minimization of a weighted sum of empirical risk is only valid if tasks are not competing, which is rarely the case. MTL with conflicting objectives requires modeling of the trade-off between tasks, which is beyond what a linear combination achieves.

An alternative objective for MTL is finding solutions that are not dominated by any others. Such solutions are said to be Pareto optimal. In this paper, we cast the objective of MTL in terms of finding Pareto optimal solutions.

The problem of finding Pareto optimal solutions given multiple criteria is called multi-objective optimization. A variety of algorithms for multi-objective optimization exist. One such approach is the multiple-gradient descent algorithm (MGDA), which uses gradient-based optimization and provably converges to a point on the Pareto set (Désidéri, 2012). MGDA is well-suited for multi-task learning with deep networks. It can use the gradients of each task and solve an optimization problem to decide on an update over the shared parameters. However, there are two technical problems that hinder the applicability of MGDA on a large scale. (i) The underlying optimization problem does not scale gracefully to high-dimensional gradients, which arise naturally in deep networks. (ii) The algorithm requires explicit computation of gradients per task, which results in linear scaling of the number of backward passes and roughly multiplies the training time by the number of tasks.

In this paper, we develop a Frank-Wolfe-based optimizer that scales to high-dimensional problems. Furthermore, we provide an upper bound for the MGDA optimization objective and show that it can be computed via a single backward pass without explicit task-specific gradients, thus making the computational overhead of the method negligible. We prove that using our upper bound yields a Pareto optimal solution under realistic assumptions. The result is an exact algorithm for multi-objective optimization of deep networks with negligible computational overhead.

We empirically evaluate the presented method on three different problems. First, we perform an extensive evaluation on multi-digit classification with MultiMNIST (Sabour et al., 2017). Second, we cast multi-label classification as MTL and conduct experiments with the CelebA dataset (Liu et al., 2015b). Lastly, we apply the presented method to scene understanding; specifically, we perform joint semantic segmentation, instance segmentation, and depth estimation on the Cityscapes dataset (Cordts et al., 2016). The number of tasks in our evaluation varies from 2 to 40. Our method clearly outperforms all baselines.

Related Work

Multi-task learning. We summarize the work most closely related to ours and refer the interested reader to reviews by Ruder (2017) and Zhou et al. (2011b) for additional background. Multi-task learning (MTL) is typically conducted via hard or soft parameter sharing. In hard parameter sharing, a subset of the parameters is shared between tasks while other parameters are task-specific. In soft parameter sharing, all parameters are task-specific but they are jointly constrained via Bayesian priors (Xue et al., 2007; Bakker and Heskes, 2003) or a joint dictionary (Argyriou et al., 2007; Long and Wang, 2015; Yang and Hospedales, 2016; Ruder, 2017). We focus on hard parameter sharing with gradient-based optimization, following the success of deep MTL in computer vision (Bilen and Vedaldi, 2016; Misra et al., 2016; Rudd et al., 2016; Yang and Hospedales, 2016; Kokkinos, 2017; Zamir et al., 2018), natural language processing (Collobert and Weston, 2008; Dong et al., 2015; Liu et al., 2015a; Luong et al., 2015; Hashimoto et al., 2017), speech processing (Huang et al., 2013; Seltzer and Droppo, 2013; Huang et al., 2015), and even seemingly unrelated domains over multiple modalities (Kaiser et al., 2017).

Baxter (2000) theoretically analyze the MTL problem as interaction between individual learners and a meta-algorithm. Each learner is responsible for one task and a meta-algorithm decides how the shared parameters are updated. All aforementioned MTL algorithms use weighted summation as the meta-algorithm. Meta-algorithms that go beyond weighted summation have also been explored. Li et al. (2014) consider the case where each individual learner is based on kernel learning and utilize multi-objective optimization. Zhang and Yeung (2010) consider the case where each learner is a linear model and use a task affinity matrix. Zhou et al. (2011a) and Bagherjeiran et al. (2005) use the assumption that tasks share a dictionary and develop an expectation-maximization-like meta-algorithm. de Miranda et al. (2012) and Zhou et al. (2017b) use swarm optimization. None of these methods apply to gradient-based learning of high-capacity models such as modern deep networks. Kendall et al. (2018) and Chen et al. (2018) propose heuristics based on uncertainty and gradient magnitudes, respectively, and apply their methods to convolutional neural networks. Another recent work uses multi-agent reinforcement learning (Rosenbaum et al., 2017).

Multi-objective optimization. Multi-objective optimization addresses the problem of optimizing a set of possibly contrasting objectives. We recommend Miettinen (1998) and Ehrgott (2005) for surveys of this field. Of particular relevance to our work is gradient-based multi-objective optimization, as developed by Fliege and Svaiter (2000), Schäffler et al. (2002), and Désidéri (2012). These methods use multi-objective Karush-Kuhn-Tucker (KKT) conditions (Kuhn and Tucker, 1951) and find a descent direction that decreases all objectives. This approach was extended to stochastic gradient descent by Peitz and Dellnitz (2018) and Poirion et al. (2017). In machine learning, these methods have been applied to multi-agent learning (Ghosh et al., 2013; Pirotta and Restelli, 2016; Parisi et al., 2014), kernel learning (Li et al., 2014), sequential decision making (Roijers et al., 2013), and Bayesian optimization (Shah and Ghahramani, 2016; Hernández-Lobato et al., 2016). Our work applies gradient-based multi-objective optimization to multi-task learning.

Multi-Task Learning as Multi-Objective Optimization

Although many hypothesis classes and loss functions have been proposed in the MTL literature, they generally yield the following empirical risk minimization formulation:

for some static or dynamically computed weights ctc^{t} per task, where L^t(θsh,θt)\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t}) is the empirical loss of the task tt, defined as \hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})\triangleq\frac{1}{N}\sum_{i}\mathcal{L}\big{(}f^{t}(\mathbf{x}_{i};{\bm{\theta}}^{sh},{\bm{\theta}}^{t}),y_{i}^{t}\big{)}.

Although the weighted summation formulation (1) is intuitively appealing, it typically either requires an expensive grid search over various scalings or the use of a heuristic (Kendall et al., 2018; Chen et al., 2018). A basic justification for scaling is that it is not possible to define global optimality in the MTL setting. Consider two sets of solutions θ{\bm{\theta}} and θˉ\bar{{\bm{\theta}}} such that L^t1(θsh,θt1)<L^t1(θˉsh,θˉt1)\hat{\mathcal{L}}^{t_{1}}({\bm{\theta}}^{sh},{\bm{\theta}}^{t_{1}})<\hat{\mathcal{L}}^{t_{1}}(\bar{{\bm{\theta}}}^{sh},\bar{{\bm{\theta}}}^{t_{1}}) and L^t2(θsh,θt2)>L^t2(θˉsh,θˉt2)\hat{\mathcal{L}}^{t_{2}}({\bm{\theta}}^{sh},{\bm{\theta}}^{t_{2}})>\hat{\mathcal{L}}^{t_{2}}(\bar{{\bm{\theta}}}^{sh},\bar{{\bm{\theta}}}^{t_{2}}), for some tasks t1t_{1} and t2t_{2}. In other words, solution θ{\bm{\theta}} is better for task t1t_{1} whereas θˉ\bar{{\bm{\theta}}} is better for t2t_{2}. It is not possible to compare these two solutions without a pairwise importance of tasks, which is typically not available.

Alternatively, MTL can be formulated as multi-objective optimization: optimizing a collection of possibly conflicting objectives. This is the approach we take. We specify the multi-objective optimization formulation of MTL using a vector-valued loss L\mathbf{L}:

The goal of multi-objective optimization is achieving Pareto optimality.

A solution θ{\bm{\theta}} dominates a solution θˉ\bar{{\bm{\theta}}} if L^t(θsh,θt)L^t(θˉsh,θˉt)\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})\leq\hat{\mathcal{L}}^{t}(\bar{{\bm{\theta}}}^{sh},\bar{{\bm{\theta}}}^{t}) for all tasks tt and L(θsh,θ1,,θT)L(θˉsh,θˉ1,,θˉT)\mathbf{L}({\bm{\theta}}^{sh},{\bm{\theta}}^{1},\ldots,{\bm{\theta}}^{T})\neq\mathbf{L}(\bar{{\bm{\theta}}}^{sh},\bar{{\bm{\theta}}}^{1},\ldots,\bar{{\bm{\theta}}}^{T}).

A solution θ{\bm{\theta}}^{\star} is called Pareto optimal if there exists no solution θ{\bm{\theta}} that dominates θ{\bm{\theta}}^{\star}.

The set of Pareto optimal solutions is called the Pareto set (Pθ\mathcal{P}_{{\bm{\theta}}}) and its image is called the Pareto front (PL={L(θ)}θPθ\mathcal{P}_{\mathbf{L}}=\{\mathbf{L}({\bm{\theta}})\}_{{\bm{\theta}}\in\mathcal{P}_{{\bm{\theta}}}}). In this paper, we focus on gradient-based multi-objective optimization due to its direct relevance to gradient-based MTL.

In the rest of this section, we first summarize in Section 3.1 how multi-objective optimization can be performed with gradient descent. Then, we suggest in Section 3.2 a practical algorithm for performing multi-objective optimization over very large parameter spaces. Finally, in Section 3.3 we propose an efficient solution for multi-objective optimization designed directly for high-capacity deep networks. Our method scales to very large models and a high number of tasks with negligible overhead.

As in the single-objective case, multi-objective optimization can be solved to local optimality via gradient descent. In this section, we summarize one such approach, called the multiple gradient descent algorithm (MGDA) (Désidéri, 2012). MGDA leverages the Karush-Kuhn-Tucker (KKT) conditions, which are necessary for optimality (Fliege and Svaiter, 2000; Schäffler et al., 2002; Désidéri, 2012). We now state the KKT conditions for both task-specific and shared parameters:

There exist α1,,αT0\alpha^{1},\ldots,\alpha^{T}\geq 0 such that t=1Tαt=1\sum_{t=1}^{T}\alpha^{t}=1 and t=1TαtθshL^t(θsh,θt)=0\sum_{t=1}^{T}\alpha^{t}\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})=0

For all tasks tt, θtL^t(θsh,θt)=0\nabla_{{\bm{\theta}}^{t}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})=0

Any solution that satisfies these conditions is called a Pareto stationary point. Although every Pareto optimal point is Pareto stationary, the reverse may not be true. Consider the optimization problem

Désidéri (2012) showed that either the solution to this optimization problem is and the resulting point satisfies the KKT conditions, or the solution gives a descent direction that improves all tasks. Hence, the resulting MTL algorithm would be gradient descent on the task-specific parameters followed by solving (3) and applying the solution (t=1Tαtθsh\sum_{t=1}^{T}\alpha^{t}\nabla_{{\bm{\theta}}^{sh}}) as a gradient update to shared parameters. We discuss how to solve (3) for an arbitrary model in Section 3.2 and present an efficient solution when the underlying model is an encoder-decoder in Section 3.3.

2 Solving the Optimization Problem

The optimization problem defined in (3) is equivalent to finding a minimum-norm point in the convex hull of the set of input points. This problem arises naturally in computational geometry: it is equivalent to finding the closest point within a convex hull to a given query point. It has been studied extensively (Makimoto et al., 1994; Wolfe, 1976; Sekitani and Yamamoto, 1993). Although many algorithms have been proposed, they do not apply in our setting because the assumptions they make do not hold. Algorithms proposed in the computational geometry literature address the problem of finding minimum-norm points in the convex hull of a large number of points in a low-dimensional space (typically of dimensionality 2 or 3). In our setting, the number of points is the number of tasks and is typically low; in contrast, the dimensionality is the number of shared parameters and can be in the millions. We therefore use a different approach based on convex optimization, since (3) is a convex quadratic problem with linear constraints.

Before we tackle the general case, let’s consider the case of two tasks. The optimization problem can be defined as minααθshL^1(θsh,θ1)+(1α)θshL^2(θsh,θ2)22\min_{\alpha\in}\|\alpha\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{1}({\bm{\theta}}^{sh},{\bm{\theta}}^{1})+(1-\alpha)\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{2}({\bm{\theta}}^{sh},{\bm{\theta}}^{2})\|_{2}^{2}, which is a one-dimensional quadratic function of α\alpha with an analytical solution:

where []+,1[\cdot]_{+,{1\atop\intercal}} represents clipping to $asas[a]_{+,{1\atop\intercal}}=\max(\min(a,1),0).WefurthervisualizethissolutioninFigure1.Althoughthisisonlyapplicablewhen. We further visualize this solution in Figure 1. Although this is only applicable whenT=2$, this enables efficient application of the Frank-Wolfe algorithm (Jaggi, 2013) since the line search can be solved analytically. Hence, we use Frank-Wolfe to solve the constrained optimization problem, using (4) as a subroutine for the line search. We give all the update equations for the Frank-Wolfe solver in Algorithm 2.

3 Efficient Optimization for Encoder-Decoder Architectures

The MTL update described in Algorithm 2 is applicable to any problem that uses optimization based on gradient descent. Our experiments also suggest that the Frank-Wolfe solver is efficient and accurate as it typically converges in a modest number of iterations with negligible effect on training time. However, the algorithm we described needs to compute θshL^t(θsh,θt)\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t}) for each task tt, which requires a backward pass over the shared parameters for each task. Hence, the resulting gradient computation would be the forward pass followed by TT backward passes. Considering the fact that computation of the backward pass is typically more expensive than the forward pass, this results in linear scaling of the training time and can be prohibitive for problems with more than a few tasks.

We now propose an efficient method that optimizes an upper bound of the objective and requires only a single backward pass. We further show that optimizing this upper bound yields a Pareto optimal solution under realistic assumptions. The architectures we address conjoin a shared representation function with task-specific decision functions. This class of architectures covers most of the existing deep MTL models and can be formally defined by constraining the hypothesis class as

where gg is the representation function shared by all tasks and ftf^{t} are the task-specific functions that take this representation as input. If we denote the representations as \mathbf{Z}=\big{(}\mathbf{z}_{1},\ldots,\mathbf{z}_{N}\big{)}, where zi=g(xi;θsh)\mathbf{z}_{i}=g(\mathbf{x}_{i};{\bm{\theta}}^{sh}), we can state the following upper bound as a direct consequence of the chain rule:

where Zθsh2\left\|\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}}\right\|_{2} is the matrix norm of the Jacobian of Z\mathbf{Z} with respect to θsh{\bm{\theta}}^{sh}. Two desirable properties of this upper bound are that (i) ZL^t(θsh,θt)\nabla_{\mathbf{Z}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t}) can be computed in a single backward pass for all tasks and (ii) Zθsh22\left\|\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}}\right\|_{2}^{2} is not a function of α1,,αT\alpha^{1},\ldots,\alpha^{T}, hence it can be removed when it is used as an optimization objective. We replace the t=1TαtθshL^t(θsh,θt)22\left\|\sum_{t=1}^{T}\alpha^{t}\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})\right\|_{2}^{2} term with the upper bound we have just derived in order to obtain the approximate optimization problem and drop the Zθsh22\left\|\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}}\right\|_{2}^{2} term since it does not affect the optimization. The resulting optimization problem is

We refer to this problem as MGDA-UB (Multiple Gradient Descent Algorithm – Upper Bound). In practice, MGDA-UB corresponds to using the gradients of the task losses with respect to the representations instead of the shared parameters. We use Algorithm 2 with only this change as the final method.

Although MGDA-UB is an approximation of the original optimization problem, we now state a theorem that shows that our method produces a Pareto optimal solution under mild assumptions. The proof is given in the supplement.

Assume Zθsh\frac{\partial\mathbf{Z}}{\partial\mathbf{\theta}^{sh}} is full-rank. If α1,,T\alpha^{1,\ldots,T} is the solution of MGDA-UB, one of the following is true:

t=1TαtθshL^t(θsh,θt)=0\sum_{t=1}^{T}\alpha^{t}\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t})=0 and the current parameters are Pareto stationary.

t=1TαtθshL^t(θsh,θt)\sum_{t=1}^{T}\alpha^{t}\nabla_{{\bm{\theta}}^{sh}}\hat{\mathcal{L}}^{t}({\bm{\theta}}^{sh},{\bm{\theta}}^{t}) is a descent direction that decreases all objectives.

This result follows from the fact that as long as Zθsh\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}} is full rank, optimizing the upper bound corresponds to minimizing the norm of the convex combination of the gradients using the Mahalonobis norm defined by ZθshZθsh\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}}^{\intercal}\frac{\partial\mathbf{Z}}{\partial{\bm{\theta}}^{sh}}. The non-singularity assumption is reasonable as singularity implies that tasks are linearly related and a trade-off is not necessary. In summary, our method provably finds a Pareto stationary point with negligible computational overhead and can be applied to any deep multi-objective problem with an encoder-decoder model.

Experiments

We evaluate the presented MTL method on a number of problems. First, we use MultiMNIST (Sabour et al., 2017), an MTL adaptation of MNIST (LeCun et al., 1998). Next, we tackle multi-label classification on the CelebA dataset (Liu et al., 2015b) by considering each label as a distinct binary classification task. These problems include both classification and regression, with the number of tasks ranging from 2 to 40. Finally, we experiment with scene understanding, jointly tackling the tasks of semantic segmentation, instance segmentation, and depth estimation on the Cityscapes dataset (Cordts et al., 2016). We discuss each experiment separately in the following subsections.

The baselines we consider are (i) uniform scaling: minimizing a uniformly weighted sum of loss functions 1TtLt\frac{1}{T}\sum_{t}\mathcal{L}^{t}, (ii) single task: solving tasks independently, (iii) grid search: exhaustively trying various values from {cttct=1}\{c^{t}\in|\sum_{t}c^{t}=1\} and optimizing for 1TtctLt\frac{1}{T}\sum_{t}c^{t}\mathcal{L}^{t}, (iv) Kendall et al. (2018): using the uncertainty weighting proposed by Kendall et al. (2018), and (v) GradNorm: using the normalization proposed by Chen et al. (2018).

Our initial experiments are on MultiMNIST, an MTL version of the MNIST dataset (Sabour et al., 2017). In order to convert digit classification into a multi-task problem, Sabour et al. (2017) overlaid multiple images together. We use a similar construction. For each image, a different one is chosen uniformly in random. Then one of these images is put at the top-left and the other one is at the bottom-right. The resulting tasks are: classifying the digit on the top-left (task-L) and classifying the digit on the bottom-right (task-R). We use 60K examples and directly apply existing single-task MNIST models. The MultiMNIST dataset is illustrated in the supplement.

We use the LeNet architecture (LeCun et al., 1998). We treat all layers except the last as the representation function gg and put two fully-connected layers as task-specific functions (see the supplement for details). We visualize the performance profile as a scatter plot of accuracies on task-L and task-R in Figure 3, and list the results in Table 5.

In this setup, any static scaling results in lower accuracy than solving each task separately (the single-task baseline). The two tasks appear to compete for model capacity, since increase in the accuracy of one task results in decrease in the accuracy of the other. Uncertainty weighting (Kendall et al., 2018) and GradNorm (Chen et al., 2018) find solutions that are slightly better than grid search but distinctly worse than the single-task baseline. In contrast, our method finds a solution that efficiently utilizes the model capacity and yields accuracies that are as good as the single-task solutions. This experiment demonstrates the effectiveness of our method as well as the necessity of treating MTL as multi-objective optimization. Even after a large hyper-parameter search, any scaling of tasks does not approach the effectiveness of our method.

2 Multi-Label Classification

Next, we tackle multi-label classification. Given a set of attributes, multi-label classification calls for deciding whether each attribute holds for the input. We use the CelebA dataset (Liu et al., 2015b), which includes 200K face images annotated with 40 attributes. Each attribute gives rise to a binary classification task and we cast this as a 40-way MTL problem. We use ResNet-18 (He et al., 2016) without the final layer as a shared representation function, and attach a linear layer for each attribute (see the supplement for further details).

We plot the resulting error for each binary classification task as a radar chart in Figure 2. The average over them is listed in Table 4.2. We skip grid search since it is not feasible over 40 tasks. Although uniform scaling is the norm in the multi-label classification literature, single-task performance is significantly better. Our method outperforms baselines for significant majority of tasks and achieves comparable performance in rest. This experiment also shows that our method remains effective when the number of tasks is high.

3 Scene Understanding

To evaluate our method in a more realistic setting, we use scene understanding. Given an RGB image, we solve three tasks: semantic segmentation (assigning pixel-level class labels), instance segmentation (assigning pixel-level instance labels), and monocular depth estimation (estimating continuous disparity per pixel). We follow the experimental procedure of Kendall et al. (2018) and use an encoder-decoder architecture. The encoder is based on ResNet-50 (He et al., 2016) and is shared by all three tasks. The decoders are task-specific and are based on the pyramid pooling module (Zhao et al., 2017) (see the supplement for further implementation details).

Since the output space of instance segmentation is unconstrained (the number of instances is not known in advance), we use a proxy problem as in Kendall et al. (2018). For each pixel, we estimate the location of the center of mass of the instance that encompasses the pixel. These center votes can then be clustered to extract the instances. In our experiments, we directly report the MSE in the proxy task. Figure 4 shows the performance profile for each pair of tasks, although we perform all experiments on all three tasks jointly. The pairwise performance profiles shown in Figure 4 are simply 2D projections of the three-dimensional profile, presented this way for legibility. The results are also listed in Table 5.

MTL outperforms single-task accuracy, indicating that the tasks cooperate and help each other. Our method outperforms all baselines on all tasks.

4 Role of the Approximation

In order to understand the role of the approximation proposed in Section 3.3, we compare the final performance and training time of our algorithm with and without the presented approximation in Table 2 (runtime measured on a single Titan Xp GPU). For a small number of tasks (3 for scene understanding), training time is reduced by 40%. For the multi-label classification experiment (40 tasks), the presented approximation accelerates learning by a factor of 25.

On the accuracy side, we expect both methods to perform similarly as long as the full-rank assumption is satisfied. As expected, the accuracy of both methods is very similar. Somewhat surprisingly, our approximation results in slightly improved accuracy in all experiments. While counter-intuitive at first, we hypothesize that this is related to the use of SGD in the learning algorithm. Stability analysis in convex optimization suggests that if gradients are computed with an error ^θLt=θLt+et\hat{\nabla}_{\bm{\theta}}\mathcal{L}^{t}=\nabla_{\bm{\theta}}\mathcal{L}^{t}+\mathbf{e}^{t} (θ{\bm{\theta}} corresponds to θsh{\bm{\theta}}^{sh} in (3)), as opposed to Z\mathbf{Z} in the approximate problem in (MGDA-UB), the error in the solution is bounded as α^α2O(maxtet2)\|\hat{\mathbf{\alpha}}-\mathbf{\alpha}\|_{2}\leq\mathcal{O}(\max_{t}\|\mathbf{e}^{t}\|_{2}). Considering the fact that the gradients are computed over the full parameter set (millions of dimensions) for the original problem and over a smaller space for the approximation (batch size times representation which is in the thousands), the dimension of the error vector is significantly higher in the original problem. We expect the l2l_{2} norm of such a random vector to depend on the dimension.

In summary, our quantitative analysis of the approximation suggests that (i) the approximation does not cause an accuracy drop and (ii) by solving an equivalent problem in a lower-dimensional space, our method achieves both better computational efficiency and higher stability.

Conclusion

We described an approach to multi-task learning. Our approach is based on multi-objective optimization. In order to apply multi-objective optimization to MTL, we described an efficient algorithm as well as specific approximations that yielded a deep MTL algorithm with almost no computational overhead. Our experiments indicate that the resulting algorithm is effective for a wide range of multi-task scenarios.

References

Appendix A Proof of Theorem 1

We begin by showing that if the optimum value of (MGDA-UB) is 0, so is the optimum value of (3). This shows the first case of the theorem. Then, we will show the second part.

Hence α1,,αT\alpha^{1},\ldots,\alpha^{T} is the solution of (3) and the optimal value of (3) is . This proves the first case of the theorem. Before we move to the second case, we state a straightforward corollary. Since Zθsh\frac{\partial\mathbf{Z}}{\partial\mathbf{\theta}^{sh}} is full rank, this equivalence is bi-directional. In other words, if α1,,αT\alpha^{1},\ldots,\alpha^{T} is the solution of (3), it is the solution of (MGDA-UB) as well. Hence, both formulations completely agree on Pareto stationarity.

In order to prove the second case, we need to show that the resulting descent direction computed by solving (MGDA-UB) does not increase any of the loss functions. Formally, we need to show that

where \mathbf{M}=\big{(}\frac{\partial\mathbf{Z}}{\partial\mathbf{\theta}^{sh}}\big{)}^{\intercal}\big{(}\frac{\partial\mathbf{Z}}{\partial\mathbf{\theta}^{sh}}\big{)}. Since M\mathbf{M} is positive definite (following the assumption), this is further equivalent to

We show that this follows from the optimality conditions for (MGDA-UB). The Lagrangian of (MGDA-UB) is

The KKT condition for this Lagrangian yields the desired result as

Appendix B Additional Results on Multi-label Classification

In this section, we present the experimental results we did not include in the main text.

In the main text, we plotted a radar chart of the binary attribute classification errors. However, we did not include the tabulated results due to the space limitations. Here we list the binary classification error of each attribute for each algorithm in Table 5.

Appendix C Implementation Details

We use the MultiMNIST dataset, which overlays multiple images together (Sabour et al., 2017). For each image, a different one is chosen uniformly in random. One of these images is placed at the top-left and the other at the bottom-right. We show sample MultiMNIST images in Figure 6.

For the MultiMNIST experiments, we use an architecture based on LeNet (LeCun et al., 1998). We use all layers except the final one as a shared encoder. We use the fully-connected layer as a task-specific function for the left and right tasks by simply adding two independent fully-connected layers, each taking the output of the shared encoder as input. As a task-specific loss function, we use the cross-entropy loss with a softmax for both tasks. The architecture is visualized in Figure 5.

The implementation uses PyTorch (Paszke et al., 2017). For all baselines, we searched over the set LR={LR=\{110-4,,510-4,,110-3,,510-3,,110-2,,510-2}\} of learning rates and chose the model with the highest validation accuracy. We used SGD with momentum, halving the learning rate every 30 epochs. We use batch size 256256 and train for 100100 epochs. We report test accuracy.

C.2 Multi-label classification

For multi-label classification experiments, we use ResNet-18 (He et al., 2016) without the final layer as a shared representation function. Since there are 40 attributes, we add 40 separate 2048×22048\times 2 dimensional fully-connected layers as task-specific functions. The final two-dimensional output is passed through a 2-class softmax to get binary attribute classification probabilities. We use cross-entropy as a task-specific loss. The architecture is visualized in Figure 7.

The implementation uses PyTorch (Paszke et al., 2017). We resize each CelebA image (Liu et al., 2015b) to 64×64×364\times 64\times 3. For all experiments, we searched over the set LR={LR=\{110-4,,510-4,,110-3,,510-3,,110-2,,510-2}\} of learning rates and chose the model with the highest validation accuracy. We used SGD with momentum, halving the learning rate every 30 epochs. We use batch size 256256 and train for 100100 epochs. We report attribute-wise binary accuracies on the test set as well as the average accuracy.

C.3 Scene understanding

For scene understanding experiments, we use the Cityscapes dataset (Cordts et al., 2016). We resize all images to resolution 256×512256\times 512 for computational efficiency. As a shared representation function (encoder), we use the ResNet-50 architecture (He et al., 2016) in fully-convolutional fashion. We take the ResNet-50 architecture and only use layers prior to average pooling that are fully convolutional. As a decoder, we use the pyramid pooling module (Zhao et al., 2017) and set the output sizes to 256×512×19256\times 512\times 19 for semantic segmentation (1919 classes), 256×512×2256\times 512\times 2 for instance segmentation (one output channel for the x-offset of the center location and another channel for the y-offset), and 256×512×1256\times 512\times 1 for monocular depth estimation. For instance segmentation, we use the proxy task of estimating the offset for the center location of the instance that encompasses the pixel. We directly estimate disparity instead of depth and later convert it to depth using the provided camera intrinsics. As a loss function, we use cross-entropy with a softmax for semantic segmentation, and MSE for depth and instance segmentation. We visualize the architecture in Figure 8.

We initialize the encoder with a model pretrained on ImageNet (Deng et al., 2009). We use the implementation of the pyramidal pooling network with bilinear interpolation shared by Zhou et al. (2017a). Ground-truth results for the Cityscapes test set are not publicly available. Therefore, we report numbers on the validation set. As a validation set for hyperparameter search, we randomly choose 275275 images from the training set. After the best hyperparameters are chosen, we retrain with the full training set and report the metrics on the Cityscapes validation set, which our algorithm never sees during training or hyperparameter search. As metrics, we use mean intersection over union (mIoU) for semantic segmentation, MSE for instance segmentation, and MSE for disparities (depth estimation). We directly report the metric in the proxy task for instance segmentation instead of performing a further clustering operation. For all experiments, we searched over the set LR={LR=\{110-4,,510-4,,110-3,,510-3,,110-2,,510-2}\} of learning rates and chose the model with the highest validation accuracy. We used SGD with momentum, halving the learning rate every 30 epochs. We use batch size 88 and train for 250250 epochs.