[Paper Review] Scalarization for Multi-Task and Multi-Domain Learning at Scale

Amélie Royer, Tijmen Blankevoort+1 more|arXiv (Cornell University)|Oct 13, 2023

Domain Adaptation and Few-Shot LearningComputer Science3 citations

TL;DR

This paper investigates scalarization—training multi-task and multi-domain models by minimizing a weighted average of task losses—as a simple yet effective alternative to complex multi-task optimization (MTO) methods. Through large-scale experiments across diverse architectures and datasets, it demonstrates that optimal scalarization weights significantly boost performance, especially in imbalanced settings, and proposes population-based training to efficiently search these weights at scale, achieving SOTA performance with minimal computational overhead.

ABSTRACT

Training a single model on multiple input domains and/or output tasks allows for compressing information from multiple sources into a unified backbone hence improves model efficiency. It also enables potential positive knowledge transfer across tasks/domains, leading to improved accuracy and data-efficient training. However, optimizing such networks is a challenge, in particular due to discrepancies between the different tasks or domains: Despite several hypotheses and solutions proposed over the years, recent work has shown that uniform scalarization training, i.e., simply minimizing the average of the task losses, yields on-par performance with more costly SotA optimization methods. This raises the issue of how well we understand the training dynamics of multi-task and multi-domain networks. In this work, we first devise a large-scale unified analysis of multi-domain and multi-task learning to better understand the dynamics of scalarization across varied task/domain combinations and model sizes. Following these insights, we then propose to leverage population-based training to efficiently search for the optimal scalarization weights when dealing with a large number of tasks or domains.

Motivation & Objective

• To understand the training dynamics of scalarization in multi-task and multi-domain learning across diverse model capacities and data distributions.
• To investigate the impact of scalarization weights on gradient conflicts and model performance in imbalanced task/domain settings.
• To address the scalability challenge of hyperparameter search for large numbers of tasks or domains.
• To evaluate whether simple scalarization can match or exceed complex MTO methods in performance while reducing computational cost.

Proposed method

• Conduct a large-scale empirical analysis of scalarization across multiple vision benchmarks, model architectures (from ViT-S/4 to full-width ViT), and task/domain combinations.
• Systematically vary scalarization weights using population-based training (PBT) to efficiently explore the high-dimensional hyperparameter space.
• Compare scalarization with state-of-the-art MTO methods, including loss-based (e.g., uncertainty, IMTL-L) and gradient-based (e.g., CAGrad, PCGrad) approaches.
• Analyze gradient conflicts across layers and learning rates to assess their role in task interference.
• Use PBT to search for optimal scalarization weights in settings with many tasks, reducing search cost while maintaining performance.
• Report results across multiple metrics and standard deviations to ensure statistical robustness.

Figure 1: Table (a) summarizes the experimental setups used throughout the paper. Figure (b) reports profiling results for popular multi-task optimization (MTO) methods in a small-scale (ResNet18, batch size 16) and large-scale (ResNet50, size 128) setting, for 224px inputs: In practice, we find tha

Experimental results

Research questions

• RQ1 How does model capacity affect the performance gain of multi-task and multi-domain learning via scalarization?
• RQ2 To what extent do optimal scalarization weights improve performance in imbalanced multi-task and multi-domain settings?
• RQ3 How do gradient conflicts vary across layers and learning rates, and do they correlate with task interference?
• RQ4 Can population-based training efficiently identify optimal scalarization weights for large-scale multi-task and multi-domain learning?
• RQ5 Does uniform scalarization (equal weighting) perform comparably to complex MTO methods, or is weight tuning essential?

Key findings

• Larger model capacities amplify the benefits of multi-task and multi-domain learning, with performance gains increasing significantly as model size grows.
• Optimal scalarization weights are crucial for high-performance multi-task and multi-domain learning in imbalanced settings, but their relative ranking remains consistent across model sizes within the same architecture family.
• Gradient conflicts between tasks or domains are a natural and persistent phenomenon during training, but they are largely unaffected by model capacity or scalarization weight choices.
• Population-based training (PBT) enables efficient and effective search for optimal scalarization weights, achieving performance on par with complex MTO methods while reducing computational cost.
• Scalarization with learned weights outperforms uniform scalarization and matches or exceeds state-of-the-art MTO methods (e.g., CAGrad, PCGrad) across all evaluated benchmarks and model sizes.
• The performance of scalarization is highly sensitive to weight tuning, with even small changes in weights leading to significant performance shifts in imbalanced settings.

Figure 2: Performance of scalarization for MDL/MTL relative to SD under different model capacities; Each column corresponds to a different task/domain pair ( $T=2$ ). The first row of each plot contains a heatmap of the best performing scalarization weights $p^{\ast}$ wrt. to the average test accura