[Paper Review] In Defense of the Unitary Scalarization for Deep Multi-Task Learning

Vitaly Kurin, Alessandro De Palma+3 more|arXiv (Cornell University)|Jan 11, 2022

Sparse and Compressive Sensing Techniques20 citations

TL;DR

The paper argues that unitary scalarization (minimizing the sum of task losses) is competitive with, or superior to, specialized multi-task optimizers (SMTOs) across supervised and reinforcement learning, when paired with standard regularization and stabilization techniques. It further analyzes SMTOs as regularizers and provides a unified experimental and theoretical perspective.

ABSTRACT

Recent multi-task learning research argues against unitary scalarization, where training simply minimizes the sum of the task losses. Several ad-hoc multi-task optimization algorithms have instead been proposed, inspired by various hypotheses about what makes multi-task settings difficult. The majority of these optimizers require per-task gradients, and introduce significant memory, runtime, and implementation overhead. We show that unitary scalarization, coupled with standard regularization and stabilization techniques from single-task learning, matches or improves upon the performance of complex multi-task optimizers in popular supervised and reinforcement learning settings. We then present an analysis suggesting that many specialized multi-task optimizers can be partly interpreted as forms of regularization, potentially explaining our surprising results. We believe our results call for a critical reevaluation of recent research in the area.

Motivation & Objective

• Assess whether unitary scalarization matches or outperforms complex SMTOs on popular MTL benchmarks.
• Investigate the role of regularization and stabilization techniques in improving unitary scalarization performance.
• Analyze SMTOs as regularizers and understand their convergence and generalization behavior.
• Provide a unified experimental pipeline and reproducible code for evaluating MTL optimizers.

Proposed method

• Formulate MTL as minimization of the sum of per-task losses (unitary scalarization) and contrast with SMTOs that use per-task gradients.
• Review and implement representative SMTOs (MGDA, IMTL, PCGrad, GradDrop, RLW) and analyze their computational overhead.
• Evaluate on supervised learning benchmarks (Multi-MNIST, CelebA, Cityscapes) and reinforcement learning benchmarks (Meta-World with SAC).
• Incorporate standard single-task regularization (dropout, L2, early stopping) and reward normalization in RL to stabilize learning.
• Analyze SMTOs through a regularization lens and perform ablations to compare against regularized unitary scalarization.

Experimental results

Research questions

• RQ1 Does unitary scalarization perform on par with SMTOs across diverse MTL settings and datasets?
• RQ2 Can standard regularization and stabilization techniques bridge gaps between unitary scalarization and SMTOs?
• RQ3 Do SMTOs act primarily as regularizers affecting generalization rather than strictly improving optimization?
• RQ4 What are the computational trade-offs (time, memory) between unitary scalarization and SMTOs in deep MTL?
• RQ5 How do SMTOs compare in reinforcement learning contexts with shared representations?

Key findings

• Unitary scalarization often matches or outperforms SMTOs on supervised learning benchmarks when coupled with regularization and stabilization.
• SMTOs can exhibit higher training time and memory overhead, with no consistent superiority over unitary scalarization.
• Regularization (dropout, L2) and early stopping reduce overfitting and can bridge performance gaps for unitary scalarization.
• Analyses suggest SMTOs may act as regularizers by delaying overfitting and influencing convergence properties.
• In RL settings, unitary scalarization achieves competitive performance with significantly lower compute compared to SMTOs like PCGrad.