Source-linked AI summary

Revisit the Imbalance Optimization in Multi-task Learning: An Experimental Analysis

Yihang Guo, Tianyuan Yu, Liang Bai, Yanming Guo, Yirun Ruan, William Li, Weishi Zheng

arXiv:2509.23915v1cs.CV

TL;DR

Task interference in multi-task learning can leave some tasks performing worse than single-task models, and the factors shaping this imbalance extend beyond optimizer choice. The paper systematically analyzes optimization, architectures, Vision Foundation Model initialization, and data, finding a strong correlation with task-gradient norms. Scaling task losses by those norms achieves performance comparable to exhaustive grid search.

  • Problem

    Task interference can make multi-task learning perform worse than single-task learning on some tasks, and the roles of optimization, architecture, initialization, and data in this imbalance remain unclear.

    Source: §I · pp. 1–2
  • Method

    The paper systematically evaluates multitask optimizers, architectures, Vision Foundation Model initialization, and data properties, then tests scaling task losses by gradient norms.

    Source: §I · p. 2
  • Results

    Task-gradient norms are strongly correlated with imbalance, and scaling task losses by their norms achieves performance comparable to exhaustive grid search.

    Source: §I · p. 2
  • Takeaways & Limitations

    The analysis points to controlling gradient dynamics as a direct route toward stable multi-task learning.

    Source: §I · p. 2

Abstract

from arXiv · show

Multi-task learning (MTL) aims to build general-purpose vision systems by training a single network to perform multiple tasks jointly. While promising, its potential is often hindered by "unbalanced optimization", where task interference leads to subpar performance compared to single-task models. To facilitate research in MTL, this paper presents a systematic experimental analysis to dissect the factors contributing to this persistent problem. Our investigation confirms that the performance of existing optimization methods varies inconsistently across datasets, and advanced architectures still rely on costly grid-searched loss weights. Furthermore, we show that while powerful Vision Foundation Models (VFMs) provide strong initialization, they do not inherently resolve the optimization imbalance, and merely increasing data quantity offers limited benefits. A crucial finding emerges from our analysis: a strong correlation exists between the optimization imbalance and the norm of task-specific gradients. We demonstrate that this insight is directly applicable, showing that a straightforward strategy of scaling task losses according to their gradient norms can achieve performance comparable to that of an extensive and computationally expensive grid search. Our comprehensive analysis suggests that understanding and controlling gradient dynamics is a more direct path to stable MTL than developing increasingly complex methods.

I. INTRODUCTION

Multi-task learning can suffer from task interference that leaves some tasks worse off than single-task models, while the factors shaping this imbalance remain incompletely understood. The paper systematically examines optimization methods, architectures, initialization, and data, and finds that task-gradient norms are strongly correlated with imbalance.

I. INTRODUCTION

  • Task interference during joint training can degrade performance on some tasks relative to single-task learning, motivating broader analysis of the imbalance problem.Prior work had largely focused on optimization strategies, leaving interactions with architectures, initialization, and data less clear.Source: §I · pp. 1–2
  • The study examines architectures, Vision Foundation Model initialization, and data quality and quantity alongside optimization methods.It reports that VFM choice and data characteristics have limited effects on imbalance, while VFMs provide strong initialization without preventing imbalance.Source: §I · p. 2
  • Existing multitask optimization methods perform inconsistently across datasets, while advanced architectures often still depend on costly grid-searched loss weights.Vanilla MTL often underperforms single-task learning, but grid search can mitigate the imbalance.Source: §I · p. 2
  • Optimization imbalance is strongly correlated with task-specific gradient norms, and scaling losses to balance those norms achieves performance comparable to exhaustive grid search.The analysis distinguishes this association from gradient-conflict angle.Source: §I · p. 2
  • Scaling task losses according to their gradient norms achieves performance comparable to an exhaustive grid search.The paper presents this as a straightforward, efficient alternative to costly weight search.Source: §I · p. 3

II. RELATED WORK

Related work addresses multitask imbalance through loss weighting, gradient updates, distillation, and specialized architectures. The paper situates its analysis as broader than prior optimizer-focused studies by also examining architecture, initialization, and data.

A. Multi-Task Balanced Optimization

  • Loss-weighting methods adjust task contributions using uncertainty, loss trends, task importance, hold-out performance, or gradient information.Examples include UW, DWA, FAMO, DTP, GradNorm, MGDA, and NashMTL [29].Source: A. Multi-Task Balanced Optimization · p. 3

Gradient update:

  • Gradient-update methods modify gradients with respect to shared parameters, but capturing gradient information can consume substantial computation.Examples include PCGrad, GradDrop, CAGrad, RotoGrad, and FairGrad [32].Source: Gradient update · p. 3

Distillation approaches and analysis on MTO methods:

  • Distillation approaches transfer task-specific or foundation-model representations into a unified model, while this study expands analysis beyond optimizer-only comparisons.The paper examines optimization, architecture, and data factors, unlike prior analyses focused only on multitask optimization methods –.Source: Distillation approaches and analysis… · p. 4

B. Multi-Task Architecture

  • Specialized multitask architectures seek to share relevant information while retaining task-specific elements, and are divided into two broad categories.The supplied passage introduces the categories without detailing them.Source: B. Multi-Task Architecture · p. 4

Encoder layer:

The section situates imbalance analysis within supervised MTL, which shares learning across tasks, and contrasts it with single-task models. It also surveys architectural approaches and positions imbalance analysis as an orthogonal direction within broader MTL research.

Encoder layer:

  • Encoder-layer methods use branching, cross-task attention, soft attention, gating, expert hybrids, or learnable cue vectors to tailor feature extraction and sharing.The cited examples include mechanisms for selecting useful features and controlling feature propagation between tasks.Source: Encoder layer · p. 4

Decoder layer:

  • Decoder-layer methods extract task-relevant information from shared features through cross-head distillation, pixel affinities, multi-scale interaction, bridge features, or transformer-based long-range correlations.MTI-Net conveys complementary information across tasks using multi-scale feature maps.Source: Decoder layer · p. 4
  • The paper analyzes how CNN- and ViT-based architectural choices influence task imbalance in MTL.The analysis focuses on architecture’s effects on imbalanced optimization.Source: Decoder layer · p. 4

C. Other Multi-Task Learning Research

  • MTL research spans task grouping, cross-task relations, continual learning, and semi/partially supervised learning, while this work focuses orthogonally on factors leading to task imbalance.The paper hopes its imbalance findings will be useful across these broader research areas.Source: C. Other Multi-Task Learning Research · p. 4

A. Single-task Learning (STL)

  • The supervised setting assumes images and labels for T tasks from one training distribution, while STL trains an independent encoder-decoder model for each task.Each task’s loss L_t is the objective for its independently trained model.Source: §III · p. 4, A. Single-task Learning (STL) · p. 4

B. Multi-task learning (MTL)

Vanilla MTL shares an encoder across task-specific decoders and jointly minimizes weighted task losses. The paper frames its performance gap against single-task learning as an imbalance problem and describes weighting and gradient-update approaches to address it.

B. Multi-task learning (MTL)

C. Imbalanced problem in MTL

  • The imbalanced problem is the challenge of matching or exceeding single-task results while jointly training one network across tasks.Jointly minimizing task losses can yield subpar performance relative to task-specific networks.Source: C. Imbalanced problem in MTL · p. 5

D. Existing Methods for Addressing the Imbalance Problem

  • Prior approaches dynamically weight task losses or modify shared-parameter gradients; UW estimates weights, while CAGrad uses gradient-update strategies.The paper denotes these approaches as G(·) in its shared-encoder update formulation.Source: D. Existing Methods for Addressing the… · p. 5

Multi-task achitectures:

The section describes task-specific encoder and decoder designs for sharing information, notes that imbalance remains despite other strategies, and specifies the datasets and metrics used in the analysis.

Multi-task achitectures:

  • Multi-task architectures augment shared features with task-specific encoder designs α_t and decoder modules β_t that accumulate information across tasks.Examples include MTAN attention modules and MTI-Net feature propagation and prediction-distillation modules.Source: Multi-task achitectures · p. 5

Other strategies:

  • Despite existing methods, the paper identifies imbalance and the effects of optimization, architecture, and VFMs as needing systematic analysis.VFMs are presented as a possible strategy for addressing imbalance.Source: Other strategies · p. 5

Datasets:

  • Experiments use NYUD-v2 and PASCAL-Context, plus Replica, a 3D reconstruction dataset with segmentation, depth, and surface-normal tasks.The datasets cover indoor scenes, dense prediction, and synthetic 3D reconstruction settings.Source: Datasets · p. 5

Evaluation Metrics:

  • Evaluation uses mIoU, maxF, odsF, RMSE, and mErr for task-specific metrics, plus Δm for average per-task improvement relative to a single-task baseline.For Δm, the sign is reversed when a lower metric value indicates better performance.Source: Evaluation Metrics · pp. 5–6

Model Backbone:

The experiments compare backbones and MTL architectures, vanilla MTL with searched loss weights, and existing optimization methods. MTO methods can improve on uniform weighting, but results vary and few consistently outperform grid search.

Model Backbone:

  • The comparison includes ResNet, HRNet, and ViT backbones, with MTAN, MTI-Net, and InvPT architectures; models generally use pretrained ImageNet weights.Training settings vary by backbone and dataset; Replica experiments use InvPT with ViT.Source: Model Backbone · p. 6

A. How Pervasive is the Imbalance Problem in Vanilla MTL?

  • Uniform vanilla MTL produces uneven task performance relative to STL, while Grid Search can bring vanilla MTL close to or above STL on NYUD-v2.Uniform weighting performs well on some tasks, such as depth, but worse on others, such as surface normals and edges.Source: A. How Pervasive is the Imbalance… · p. 6

B. How Effective are Existing Optimization (MTO) Methods at Mitigating Imbalance?

  • The reviewed MTO strategies include dynamic loss weighting, gradient updates, and distillation to align MTL with task-specific models.Weighting criteria include uncertainty, gradient norms, loss-decrease speed, random weights, Pareto solutions, and Nash bargaining.Source: B. How Effective are Existing… · p. 6, 1) loss weighting · p. 6
  • FairGrad gains +0.59% on NYUD-v2 over Uniform, but few MTO methods consistently outperform Grid Search across scenarios.Grid Search remains competitive, and selecting optimal task-specific loss weights remains challenging.Source: 1) loss weighting · p. 6

1) loss weighting

  • MGDA reaches -7.19% in PASCAL-Context, while MGDA and RLW perform worse in most scenarios.The paper suggests MGDA’s strict Pareto constraint may be limiting, while random RLW weights may over- or under-amplify task losses.Source: 1) loss weighting · p. 6

multi-task optimization.

MTL optimization methods and specialized architectures can reduce imbalance, but performance remains dataset-dependent and may still trail single-task models. Stronger backbones help MTO methods, while both method families continue to rely on loss-weight tuning.

multi-task optimization.

  • MTO methods perform better with a stronger backbone or MTAN, but remain inconsistent across datasets; grid-searched loss weights are the most stable and obtain comparable results.Stronger backbones may provide shared features that better capture inter-task correlations and give MTO methods a better basis.Source: multi-task optimization · pp. 6–7

C. Do Specialized Architectures Fundamentally Solve the Imbalance Problem?

  • MTL architectures with uniform weights outperform single-task learning on NYUD-v2 but fall short on PASCAL-Context; fully shared ViT without cross-task design worsens imbalance.MTAN Uniform reports +0.69% on NYUD-v2 and -1.67% on PASCAL-Context.Source: C. Do Specialized Architectures… · p. 8, Figure 4 · p. 8
  • MTI-Net Uniform (+3.66%) exceeds MTI-Net Single-task (+3.45%) on NYUD-v2, but on PASCAL-Context it scores -1.75% versus +1.71%.Architecture-based methods can narrow the gap to their single-task counterparts, yet still rely on grid-searched loss weights.Source: C. Do Specialized Architectures… · p. 8
  • Architecture-based methods mitigate imbalance relative to vanilla MTL, but their advantage over corresponding single-task models is smaller than vanilla MTL’s gap to STL.The authors conclude that architecture and optimization can be complementary, while current architecture-based methods still rely on costly, unstable tuning.Source: C. Do Specialized Architectures… · p. 8

D. Do Vision Foundation Models (VFMs) Prevent Optimization Imbalance?

The analysis finds that powerful Vision Foundation Models do not eliminate task imbalance, and cleaner data improves performance without clearly worsening or easing imbalance. The data-quality tests used Replica annotations with added noise across tasks.

D. Do Vision Foundation Models (VFMs) Prevent Optimization Imbalance?

  • ImageNet-, CLIP-, SAM-, DINO-, and DINOv2-pretrained ViT-B models still show task imbalance, and none surpasses its single-task counterpart.Grid search can nevertheless produce more balanced performance than uniform weighting; the tested VFMs also show different task representation biases.Source: D. Do Vision Foundation Models (VFMs)… · p. 9, Table IV · p. 9

Data Quality:

  • Noisy annotations reduce performance for both MTL and STL, but the experiments find no clear evidence that noise exacerbates task imbalance.The authors note that independently trained single-task models also lose performance under noise.Source: Data Quality · pp. 9–10

Data Quantity:

The paper reports that fewer training samples worsen performance and can intensify imbalance, while more data alone does not necessarily address it. Its further analysis links imbalance to task-gradient norms and motivates AvgNorm, which performs comparably to grid search.

Data Quantity:

  • Training with fewer Replica samples worsens STL and MTL performance and makes MTL more imbalanced, but using more samples does not necessarily address imbalance.The comparison includes models trained with different numbers of scans per scene or different scenes, against models trained on all samples.Source: Data Quantity · p. 10

F. Further Analysis in Foundational Factors

  • Grid Search, CAGrad, and NashMTL show feature differences closer to their corresponding single-task models than Uniform, which has the worst ∆m.The authors state that single-task models can guide MTL feature representations.Source: F. Further Analysis in Foundational… · p. 10, Figure 6 · p. 10
  • Inter-task gradient similarity alone shows no clear cue for whether reducing gradient conflict mitigates imbalance, while effective methods’ task gradients converge to similar cosine similarity with the aggregate.The reported value is around 0.5 for Grid Search, CAGrad, and NashMTL; Uniform’s curves diverge.Source: F. Further Analysis in Foundational… · p. 11, Figure 7 · p. 11
  • AvgNorm dynamically scales each task’s gradient norm relative to the aggregate norm, avoiding unstable weights when an individual task’s gradient norm is extreme.The authors report a correlation between task-gradient norms and MTL performance, with better-performing methods converging to similar norms.Source: F. Further Analysis in Foundational… · pp. 11–12, Figure 8 · p. 11
  • AvgNorm obtains performance comparable to grid search and does so automatically and more efficiently; with MTAN it reports the best PASCAL-Context performance (-0.43%).Experiments found that gradients from the shared backbone’s last layer can suffice.Source: F. Further Analysis in Foundational… · p. 12, Table VI · p. 12, Table VII · p. 12

VI. CONCLUSIONS

The paper’s systematic study of task imbalance in multi-task dense prediction finds a strong gradient-norm correlation and demonstrates that scaling losses by gradient norms can match grid-search performance.

VI. CONCLUSIONS

  • Across multi-task dense prediction, task-specific gradient norms correlate strongly with imbalance, and scaling losses by those norms achieves performance comparable to exhaustive grid search.The study also finds that existing optimization methods are inconsistent, advanced architectures retain costly tuning, and VFMs and data properties offer limited mitigation.Source: §VI · p. 12

Datasets and data augmentation

Experiments span three multi-task datasets, with dataset-specific preprocessing, task losses, training schedules, and optimizer settings. The analysis also reports averaged gradient measurements across training iterations.

Datasets and data augmentation

  • Experiments use PASCAL-Context, NYUD-v2, and Replica, with dataset-specific image resizing and shared scaling, flipping, color-jitter, and normalization augmentations.DINO-v2 inputs use different crop sizes from the other models.Source: Datasets and data augmentation · p. 15

Loss functions and task weight

  • Task-specific losses are combined as a weighted sum, and Grid Search uses dataset-specific task weights.The reported weights differ between NYUD-v2 and PASCAL-Context.Source: Loss functions and task weight · p. 15

Training details

  • Training duration and checkpoint selection vary by backbone and dataset, with Replica evaluated on its test set after validation-based model selection.ResNet uses epochs, while HRNet and ViT use iterations.Source: Training details · p. 15

Others settings

MORE RESULTS

  • Gradient measurements are counted each iteration and visualized as averages over 2k NYUD-v2 or 10k PASCAL-Context iterations.The reported averages are used for presenting the detailed results.Source: MORE RESULTS · p. 16

A. Gradient conflict

The analysis finds that task-pair orthogonality alone does not track optimization imbalance; gradient norms and each task’s relation to the aggregate gradient offer more relevant signals. Gradient-norm balancing can alleviate imbalance, with the paper’s method reported as less computationally expensive while achieving comparable performance to existing methods.

Gradient similarity between tasks:

B. Gradient norm

  • On PASCAL-Context, Grid Search mostly brings task gradient norms to similar values except for Norm, while CAGrad and NashMTL converge norms across tasks.The analysis suggests Grid Search may not reach the Pareto frontier and reports no gradient-norm correlation between single-task and multi-task results.Source: B. Gradient norm · p. 16, Figure 15 · p. 18, Figure 16 · p. 18
  • Limiting task gradient norms to a consistent range can alleviate imbalance, and the paper’s method is less computationally expensive while achieving comparable performance to existing methods.CAGrad and NashMTL also balance gradient sizes well.Source: B. Gradient norm · p. 16, Figure 15 · p. 18
Loading 2509.23915v1…