[Paper Review] Multi-Task Learning as Multi-Objective Optimization

Ozan Şener|arXiv (Cornell University)|Oct 10, 2018

Domain Adaptation and Few-Shot Learning340 citations

TL;DR

The paper reframes multi-task learning as a Pareto-optimal multi-objective problem, introduces a scalable gradient-based optimizer (MGDA-UB with Frank-Wolfe) for deep networks, and shows superior performance on MultiMNIST, CelebA, and Cityscapes compared to baselines.

ABSTRACT

In multi-task learning, multiple tasks are solved jointly, sharing inductive bias between them. Multi-task learning is inherently a multi-objective problem because different tasks may conflict, necessitating a trade-off. A common compromise is to optimize a proxy objective that minimizes a weighted linear combination of per-task losses. However, this workaround is only valid when the tasks do not compete, which is rarely the case. In this paper, we explicitly cast multi-task learning as multi-objective optimization, with the overall objective of finding a Pareto optimal solution. To this end, we use algorithms developed in the gradient-based multi-objective optimization literature. These algorithms are not directly applicable to large-scale learning problems since they scale poorly with the dimensionality of the gradients and the number of tasks. We therefore propose an upper bound for the multi-objective loss and show that it can be optimized efficiently. We further prove that optimizing this upper bound yields a Pareto optimal solution under realistic assumptions. We apply our method to a variety of multi-task deep learning problems including digit classification, scene understanding (joint semantic segmentation, instance segmentation, and depth estimation), and multi-label classification. Our method produces higher-performing models than recent multi-task learning formulations or per-task training.

Motivation & Objective

• Motivate MTL as a multi-objective problem due to task conflicts and competing objectives.
• Formulate MTL to seek Pareto-optimal solutions rather than a single weighted sum.
• Develop a scalable optimizer that handles high-dimensional gradients and many tasks in deep networks.
• Prove that optimizing an upper bound yields Pareto optimality under reasonable assumptions.
• Demonstrate effectiveness across diverse datasets and task sets (2–40 tasks).

Proposed method

• Formulate per-task losses into a vector-valued objective L(θsh, θ1,..., θT).
• Apply gradient-based multi-objective optimization using the MGDA/KKT framework to find descent directions.
• Solve a min-norm problem in the convex hull of task gradients to obtain task-combination weights α1,...,αT (Eq. 3).
• Use a Frank-Wolfe-based solver to compute α in scalable fashion; derive an analytical line search for the two-task case (Eq. 4).
• Introduce MGDA-UB: replace shared-gradient norm with an upper bound that uses gradients w.r.t. representations Z to require only a single backward pass (Section 3.3).
• Provide a theorem (Theorem 1) showing MGDA-UB yields a Pareto stationary point under a full-rank assumption for ∂Z/∂θsh.
• Adapt the approach to encoder-decoder architectures so updates can be computed with one backward pass and shared representation g(·; θsh).

Experimental results

Research questions

• RQ1 How to cast MTL as a multi-objective optimization problem and what is the appropriate optimality notion (Pareto optimality)?
• RQ2 Can gradient-based MGDA be scaled to high-dimensional deep networks and many tasks without prohibitive overhead?
• RQ3 Does optimizing an upper-bound (MGDA-UB) preserve Pareto optimality under realistic conditions?
• RQ4 How does the proposed method perform on tasks with varying numbers of objectives (2–40) across different domains (classification, multi-label, scene understanding)?

Key findings

• The MGDA-UB approach yields Pareto-optimal or Pareto-stationary solutions with negligible overhead.
• On MultiMNIST (2 tasks) the method matches single-task performance and surpasses other MTL baselines, demonstrating effective capacity sharing.
• On CelebA (40 tasks) the method achieves lower mean error than uniform scaling, Kendall 2018, and GradNorm.
• On Cityscapes (3 tasks) the method achieves the highest mIoU and lowest pixel disparity error among baselines.
• MGDA-UB significantly speeds training (40% reduction for 3-task scene understanding; 25× speed-up for 40-task CelebA) with comparable accuracy to the full MGDA.
• Across tasks and datasets, the method consistently outperforms baselines and supports scalable MTL with many tasks.