Multi-Task Loss Function

• Multi-task loss functions are scalar objectives that aggregate individual task losses using weights, normalization, or advanced methods to coordinate learning across tasks.
• They employ strategies like homoscedastic uncertainty weighting, dynamic gradient-based balancing, and bilevel optimization to mitigate negative transfer and ensure fairness.
• Recent advances integrate cross-task consistency, feature distillation, and perceptual losses to enhance generalization and achieve scalable, balanced multi-task performance.

A multi-task loss function is a scalar objective for jointly training a single model to solve multiple tasks, typically by aggregating the per-task losses through explicit weighting, normalization, or more advanced multi-objective strategies. Multi-task loss functions are fundamental to multi-task learning (MTL), serving as the main mechanism for coordinating learning signals, controlling negative transfer, and achieving balanced or Pareto-efficient task trade-offs.

1. Formulations of the Multi-Task Loss

The most prevalent multi-task setup aggregates $K$ individual task losses $l_i(\theta)$ (arising from regression, classification, ranking, reconstruction, etc.) over shared parameters $\theta$:

$$L_{\textrm{MTL}}(\theta) = \sum_{i=1}^K w_i l_i(\theta)$$

where $w_i \geq 0$ are task-specific weights. This weighted sum scalarization remains the default for broad classes of deep MTL architectures (Silva et al., 2020, Verboven et al., 2020, Crawshaw et al., 2021, Kirchdorfer et al., 2024).

However, several alternative formulations have been introduced:

• Geometric Mean Aggregation: The loss can be formulated as a geometric mean for scale-invariance and implicit reweighting:

$$L_{\textrm{geo}}(\theta) = \left( \prod_{i=1}^K l_i(\theta) \right)^{1/K}$$

This is equivalent to minimizing the arithmetic mean of the log-losses and induces dynamic, scale-adaptive gradients (Chennupati et al., 2019).

• Bilevel Optimization: Recent work recasts multi-task loss balancing as a bilevel optimization, where the inner loop minimizes a weighted sum of normalized losses, and the outer loop controls the discrepancy between normalized per-task losses to ensure balanced optimization (Xiao et al., 12 Feb 2025).
• Hybrid and Task-Specific Compositions: Some applications, e.g., embedding learning or image translation, combine heterogeneous loss families (InfoNCE, MSE, perceptual losses) in batch-wise or stratified fashions, decoupling the formulation from fixed summing and delegating balancing to data sampling or explicit multi-component compositional objectives (Zhu et al., 2023, Huang et al., 2024).

2. Loss Weighting and Balancing Strategies

Weight selection ($w_i$) is critical, as naive uniform weighting or static coefficients can result in domination by numerically large or hard-to-learn tasks. Major strategies include:

• Homoscedastic Uncertainty Weighting: Introduces per-task learned parameters $\sigma_i$ capturing intrinsic output noise, yielding the loss

$$L = \sum_i \frac{1}{2\sigma_i^2} l_i(\theta) + \log \sigma_i$$

which down-weights noisy or high-variance task losses adaptively (Silva et al., 2020, Kirchdorfer et al., 2024).

• Analytical Optimal Weighting: By optimizing the uncertainty-weighted objective analytically (as in UW-SO), optimal weights are set as $w_i \propto \frac{1}{l_i(\theta)}$, followed by a softmax normalization with tunable temperature:

$$w_i = \mathrm{softmax}_i\left( \frac{1}{T l_i(\theta)} \right)$$

(Kirchdorfer et al., 2024).

• Dynamic Gradient-Based Balancing: Methods such as GradNorm, SLAW, and HydaLearn align the gradient contributions of each task, either by equalizing normalized gradient magnitudes (SLAW) or by computing per-batch task weights that maximize the main task's metric improvement (HydaLearn) (Crawshaw et al., 2021, Verboven et al., 2020). SLAW maintains an exponential moving estimate of per-task loss variance as a proxy for gradient norm, and updates $w_i$ accordingly.
• Fairness-Based and Multi-Objective Balancing: FairGrad generalizes loss aggregation to $\alpha$-fair utility maximization, where the update direction is chosen to optimize:

$\sum_{i=1}^K U_\alpha\left( g_i^\top d \right)$

with $U_\alpha$ the $\alpha$-fair utility (e.g., linear, proportional, or max-min fairness), subject to $g_i^\top d \geq 0$. The solution involves solving a generalized nonlinear system per batch (Ban et al., 2024).

• Loss Discrepancy Minimization: The LDC-MTL (BiLB4MTL) approach introduces a bilevel procedure to minimize discrepancies between normalized losses, seeking solutions where loss inequities are penalized and Pareto stationarity is ensured at $O(1)$ complexity in the task count (Xiao et al., 12 Feb 2025).

3. Advanced Multi-Task Loss Architectures

Recent MTL frameworks augment classic formulations with architectural or task-structural innovations:

• Feature-to-Feature Perceptual Losses: In medical image enhancement, multi-task customization of perceptual losses using feature-extractors trained in an MTL paradigm allows task-specific, multi-scale content and style constraints, as in CFP-loss, which aggregates learned MSE measures at different feature levels with adaptive per-task weighting (Zhu et al., 2023).
• Cross-Task Consistency and Alignment: Multi-task losses may incorporate not only per-task error terms, but also "alignment" or "consistency" losses penalizing disagreement between direct predictions and cross-task inferences (e.g., segmentation-to-depth and vice versa), promoting mutually consistent outputs (Nakano et al., 2021).
• Distillation-Augmented Losses: Knowledge distillation for MTL introduces auxiliary feature-alignment losses between the shared MTL representation (via dedicated adaptors) and frozen single-task teachers, with hyperparameter-tuned trade-offs between standard per-task losses and these distillation terms (Li et al., 2020).
• Pairwise or Structural Losses: In recommender MTL, pairwise ranking losses can be introduced to leverage the inherent sequential or relational structure among tasks (e.g., conversion must follow click), enforcing higher scores for conversion than for non-conversion clicks, augmenting the standard BCE with a margin-based pairwise objective (Durmus et al., 2024).

4. Practical Considerations in Multi-Task Loss Implementation

Robust and efficient implementation of multi-task losses raises several concerns:

• Scale Normalization: Raw task losses often differ in scale and units. Pre-normalization (division by initial loss, log-normalization) is crucial for effective joint optimization (Xiao et al., 12 Feb 2025, Chennupati et al., 2019).
• Sampling and Task Heterogeneity: In settings where tasks differ significantly in data availability or architecture, losses may be computed on a stratified, sample-wise, or alternating basis (as in hybrid losses of Piccolo2 (Huang et al., 2024)), rather than in a uniform, batch-synchronous manner.
• Hyperparameter Sensitivity: Methods with temperature parameters (softmax-based normalizations), learning rate, or router architectures (bilevel optimization) require careful tuning, as improper settings can bias task performance or delay convergence (Kirchdorfer et al., 2024, Xiao et al., 12 Feb 2025).
• Computational Overhead: Advanced gradient manipulation methods (e.g., PCGrad, MGDA, Nash-MTL, traditional GradNorm) have $O(K)$ compute/memory footprints, whereas SLAW and BiLB4MTL achieve $O(1)$ scaling by leveraging loss statistics or router-based reweighting, making them preferable in large-task settings (Xiao et al., 12 Feb 2025, Crawshaw et al., 2021).

5. Empirical Evaluation and Comparative Performance

Recent benchmarks evaluating scalarization, adaptive weighting, gradient manipulation, fairness, and distillation-based schemes present the following findings:

Method	Fundamental Mechanism	Scaling ($K$)	Noted Strengths	Noted Limitations
Scalarization	Static weights	$O(1)$	Simple, baseline	Requires tuning, imbalanced tasks
Uncertainty Wtd	Learned/task noise parameter	$O(1)$	Principled, solid on heterogeneity	Suffering from slow adaptation
GradNorm/SLAW	Gradient-based reweighting	$O(K)$, $O(1)$	Empirically strong, balanced grads	GradNorm is expensive, SLAW efficient
UW-SO	Analytic optimal weighting	$O(1)$	Matches scalarization at lower cost	Sensitive to temperature
FairGrad	$\alpha$-fairness (utility)	$O(K)$	Flexible fairness objectives	Solver overhead
BiLB4MTL	Bilevel loss discrepancy	$O(1)$	Balanced, efficient, Pareto-stationary	Requires normalization choice
Distillation	Feature-alignment	$O(1)$	Balanced sharing, improved generalization	Needs task-specific teachers
Pairwise/Rk	Pairwise ranking	$O(B^2)$	Leverages task structure	Quadratic cost, mitigated by sampling

Empirically, analytically-optimized uncertainty weighting (as in UW-SO and its softmax variant), geometric mean, and bilevel loss discrepancy minimization consistently achieve performance comparable to sophisticated multi-gradient methods but with orders-of-magnitude lower compute in large-$K$ regimes. For fair task treatment, $\alpha$-fair formulations and explicit loss-discrepancy control further reduce "task starvation" compared to traditional scalarization (Kirchdorfer et al., 2024, Xiao et al., 12 Feb 2025, Ban et al., 2024, Chennupati et al., 2019).

6. Design Trade-offs and Current Challenges

Design of effective multi-task loss functions must trade off several orthogonal objectives:

• Task Balance vs. Specialization: Simple sum-based scalarization may promote convergence on easy/high-loss tasks, starving others. Structured regularizers, Pareto optimization, and fairness-based terms can rebalance but may trade off absolute task performance for equity (Xiao et al., 12 Feb 2025, Ban et al., 2024, Chennupati et al., 2019).
• Negative Transfer Mitigation: Downweighting noisy or adversarial tasks, either via learned uncertainties or bilevel discrepancy, is essential to suppress negative transfer and enhance generalization (Silva et al., 2020, Kirchdorfer et al., 2024).
• Adaptivity Across Training: Dynamic schedules (SLAW, HydaLearn) and analytic batch-wise updates (UW-SO) adapt rapidly to training nonstationarities and per-batch anomalies, outperforming static reweighting or epoch-level updates (Crawshaw et al., 2021, Kirchdorfer et al., 2024, Verboven et al., 2020).
• Scalability and Implementation Simplicity: Increasing task count and data volume demands $O(1)$-scaling solutions and avoidance of $K$-or $B^2$-dependent operations (e.g., as in BiLB4MTL, SLAW, and analytic uncertainty weighting) (Crawshaw et al., 2021, Xiao et al., 12 Feb 2025, Kirchdorfer et al., 2024).

Ongoing research continues to address remaining open problems in robust normalization, improved fairness mechanisms, theoretical understanding of cross-task trade-offs, and efficient integration with self-supervised, contrastive, or multi-modal tasks.

---

References

(Verboven et al., 2020): HydaLearn: Highly Dynamic Task Weighting for Multi-task Learning with Auxiliary Tasks (Silva et al., 2020): Task Uncertainty Loss Reduce Negative Transfer in Asymmetric Multi-task Feature Learning (Crawshaw et al., 2021): SLAW: Scaled Loss Approximate Weighting for Efficient Multi-Task Learning (Zhu et al., 2023): Feature-oriented Deep Learning Framework for Pulmonary CBCT Enhancement with Multi-task Customized Perceptual Loss (Ban et al., 2024): Fair Resource Allocation in Multi-Task Learning (Huang et al., 2024): Piccolo2: General Text Embedding with Multi-task Hybrid Loss Training (Durmus et al., 2024): Pairwise Ranking Loss for Multi-Task Learning in Recommender Systems (Kirchdorfer et al., 2024): Analytical Uncertainty-Based Loss Weighting in Multi-Task Learning (Xiao et al., 12 Feb 2025): LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control (Chennupati et al., 2019): MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning (Nakano et al., 2021): Cross-Task Consistency Learning Framework for Multi-Task Learning (Li et al., 2020): Knowledge Distillation for Multi-task Learning