Skip to content

LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control

Conference: ECCV2026
Paper: ECCV Paper
Code: https://github.com/OptMN-Lab/LDC-MTL
Area: Optimization & Theory
Keywords: multi-task learning, loss discrepancy control, bilevel optimization, dynamic weighting, first-order optimization

TL;DR

LDC-MTL explicitly controls task-loss discrepancies through bilevel optimization and removes an empirically small gradient term to enable single-loop first-order training, achieving an aggregate relative performance-drop score of -1.31 on CelebA at 1.07 times the training time of linear scalarization.

Background & Motivation

Multi-task learning shares a model across classification, regression, or dense prediction tasks, but sharing does not mean that every task benefits equally. Different loss units, magnitudes, task difficulties, and rates of improvement can allow some tasks to dominate shared-parameter updates. Fixed-weight linear scalarization (LS) is inexpensive but can persistently favor a subset of tasks; dynamic weighting changes task importance without necessarily controlling the remaining gaps between task losses.

Another line of work addresses gradient conflict. PCGrad, CAGrad, and Nash-MTL inspect task gradients before adjusting the shared update direction, typically requiring separate computation and storage of those gradients. With 40 tasks on CelebA, these operations are much more expensive than back-propagating through one aggregate loss. The question is not whether balance is desirable, but whether inexpensive loss-level signals can produce balance without merely making every task equally poor.

The authors assign these requirements to different levels: the lower level minimizes a weighted sum of task losses, while the upper level selects weights that reduce loss discrepancies. The implementation challenge then becomes solving the bilevel problem efficiently, motivating an empirical gradient-based simplification that removes the costly inner loop. Core Idea: explicitly optimize task-loss discrepancies while preserving pressure to reduce the weighted task loss, then approximate this control with single-loop first-order training at a cost close to LS.

Method

Overall Architecture

The input remains training data with supervision for multiple tasks. A base model produces shared features and task predictions, from which the individual losses are computed. A small routing network reads the shared features and uses softmax to output task weights; model parameters are denoted by \(x\), routing parameters by \(W\), and the number of tasks by \(K\).

Training considers both the total weighted task loss and discrepancies between task losses. These are denoted by \(g\) and \(f\), respectively. The bilevel formulation ideally trains the model under given routing weights before evaluating whether its loss discrepancies are acceptable. In practice, LDC-MTL does not fully solve this lower-level problem at every step; it updates the model and routing network jointly.

This is a training-objective and solver modification, not an additional inference pipeline. The explanation below follows how discrepancy is defined, how lower-level optimality is retained, and why the inner loop can be removed, without representing individual loss terms as separate processing modules.

Key Designs

1. Explicit loss discrepancy control: a concrete balancing objective for weight learning

Given the router output, the lower-level objective is a weighted sum of task losses. Softmax makes the weights nonnegative and sum to 1, keeping the method within the familiar scalarization framework of multi-objective optimization rather than introducing a separate optimizer for every task:

\[ g(W,x)=\sum_{i=1}^{K}\sigma_i(W)l_i(x). \]

The upper-level objective \(f\) accumulates absolute differences between weighted losses adjacent in task-index order, rather than distances between every pair of tasks. It uses another set of coefficients \(\tau\), with two options in the paper: \(\tau=\sigma(W)\) or \(\tau=\mathbf{1}\). The former compares routing-weighted losses, whereas the latter compares raw losses. These encode different notions of balance: similar weighted values do not imply equal raw task errors. Equation (1) is corrupted in the extracted text, so this note describes its identifiable definition without reconstructing the full equation.

This objective gives weight learning a more specific signal than whether the total loss is small. If some tasks improve quickly while others lag behind, their loss discrepancies contribute to the update. However, the objective does not require all final losses to be equal, nor does it imply equal test accuracy across tasks. Evaluation quantities such as classification error and depth error are not the training losses themselves, and changing a loss's units changes the meaning of its discrepancy.

2. Value-function penalty: narrowing gaps without abandoning task learning

Minimizing discrepancy alone could yield an unhelpful solution in which all losses are equally large, so lower-level optimality must be retained. Let \(x^*(W)\) be an optimal model for the weighted loss under given routing parameters. The authors penalize the gap between the current weighted loss and this optimal value:

\[ p(W,x)=g(W,x)-g(W,x^*(W)),\qquad F_\lambda(W,x)=f(W,x)+\lambda p(W,x). \]

For an exact lower-level optimum, this gap is nonnegative; requiring it to be no greater than zero is equivalent to requiring the lower level to be optimized. A finite positive \(\lambda\) relaxes the need to reach lower-level optimality at every step: a larger value emphasizes reducing the weighted loss and approaching its lower-level solution, while a smaller value emphasizes reducing task discrepancies. Balance is therefore a compromise optimized jointly with task learning, not a post-training equalization of results.

Direct optimization remains expensive because updating the router involves a gradient of the lower-level optimal value. A conventional first-order penalty method maintains an auxiliary variable \(z\) with the same dimensionality as the model and performs \(N\) lower-level gradient steps per outer iteration to approximate \(x^*(W)\). Even without Hessians, this auxiliary model and inner loop can cancel the efficiency advantage sought through loss weighting.

3. Small-gradient omission: reducing the bilevel approximation to one aggregate backward pass

The authors inspect the routing gradient at the auxiliary lower-level solution, \(\nabla_W g(W,z_N)\). In the diagnostic experiment, the inner loop uses \(N=50\); on Cityscapes, the ratio between the corresponding gradient norm at the current model and at the auxiliary solution exceeds 100 on average. This observation supports omitting the latter, but it is empirical evidence from particular experiments, not a condition guaranteed for every model or training stage.

After omitting that term, the auxiliary lower-level solution no longer needs to be computed, removing both \(z\) and the inner loop. Following the update directions in Algorithm 1, the practical first-order aggregate objective can be written compactly as:

\[ J(W,x)=f(W,x)+\lambda g(W,x). \]

Both the model and routing parameters descend along their respective gradients of \(J\). Here, \(J\) is a compact expression of the simplified update, not a claim of unconditional equivalence to the original bilevel problem. The omitted term is precisely the approximation that requires empirical and theoretical justification.

The aggregate objective requires a constant number of backward passes, without materializing a full-model gradient separately for each task or computing second-order gradients. The claimed \(O(1)\) time and memory should be understood as gradient-aggregation overhead relative to task count, not as total network training cost being independent of that count: every task loss must still be evaluated, and the task heads and \(K\)-dimensional routing output remain.

Loss & Training

During training, the supervised task losses contribute to \(g\), loss discrepancy contributes to \(f\), and a positive penalty coefficient controls their relative influence. The implementation uses the same learning rate for the model and router, although the paper notes that these can be set independently. No gradient-projection term is added; the observed reduction in conflict is a side effect of controlling loss discrepancy.

The paper adopts learning-rate settings from prior experiments and selects the penalty coefficient through grid search. The reported CelebA result uses \(\lambda=0.01\). The supplied main-text cache does not include the appendix, so it does not justify filling in every dataset's training epochs, batch size, network configuration, or specific choice of \(\tau\).

The theoretical analysis requires an upper-level Lipschitz condition, smoothness of the relevant functions, and a lower-level Polyakโ€“ลojasiewicz (PL) condition. Absolute-value discrepancies are themselves nonsmooth; the authors propose a smooth absolute-value replacement, so the theorem does not directly apply to every original nonsmooth implementation.

For the simplified Algorithm 1, the omitted routing gradient must additionally remain sufficiently small, of order \(O(\epsilon)\), throughout the analyzed iterations. Under these conditions and suitable hyperparameters, the authors establish approximate stationarity for the penalized bilevel problem and approximate Pareto stationarity for the original tasks, reporting an iteration complexity of \(O(\epsilon^{-2})\).

Pareto stationarity means that some nonnegative task weights summing to 1 produce a sufficiently small weighted sum of task gradients. This is weaker than finding a globally Pareto-optimal solution and does not imply that every task outperforms STL. The extracted theorem's step-size and penalty-parameter expressions are corrupted, so their precise bounds are not reconstructed here.

Key Experimental Results

Main Results

The paper evaluates 2 tasks on Cityscapes, 3 on NYU-v2, 11 on QM9, and 40 on CelebA, using one NVIDIA A6000. The Cityscapes and NYU-v2 table captions explicitly report means over 3 random seeds; the main text refers to the unavailable appendix for standard deviations.

The aggregate metric \(\Delta m\%\) measures the average relative performance drop against single-task learning (STL): compute each evaluation quantity's relative change from STL, reverse the sign for higher-is-better quantities, then average and multiply by 100. Lower is better; negative values indicate average improvement, whereas positive values indicate an average deficit relative to STL. It is neither classification accuracy nor improvement relative to LS. MR is the mean rank across evaluation metrics, also lower-is-better.

The table below selects comparison methods appearing across the original Tables 1โ€“3. Every entry is the reported \(\Delta m\%\), without recomputing or mixing individual task metrics.

Method Cityscapes, 2 tasks NYU-v2, 3 tasks CelebA, 40 tasks QM9, 11 tasks
LS 22.60 5.59 4.15 177.6
FAMO 8.13 -4.10 1.21 58.5
Nash-MTL 6.82 -4.04 2.84 62.0
FairGrad 5.18 -4.66 0.37 57.9
LDC-MTL -0.57 -4.40 -1.31 49.5

The aggregate Cityscapes advantage is not a win on every metric: LDC-MTL has segmentation mIoU of 74.53, below FairGrad's 75.72, but its depth relative error of 26.79 is better than FairGrad's 32.25. On NYU-v2, FairGrad's aggregate score of -4.66 also improves on LDC-MTL's -4.40, precluding a claim that the proposed method is best on every benchmark.

Ablation Study

The original Table 4 analyzes the distribution of losses across 40 CelebA tasks. Its standard deviation measures dispersion between task losses, not uncertainty across random seeds. This is a mechanism-analysis table rather than a component-removal ablation.

Method Mean loss Loss standard deviation Minimum loss Maximum loss
LS 0.287 0.195 0.032 0.729
FAMO 0.231 0.141 0.027 0.621
GO4Align 0.238 0.154 0.026 0.660
LDC-MTL 0.189 0.133 0.029 0.538

The simultaneous decreases in mean and standard deviation support lower and more concentrated losses, which is more informative than a reduction in variance alone. However, LDC-MTL does not have the lowest minimum loss, reinforcing that the objective is not to beat every method on every statistic.

The original Table 5 reports penalty-coefficient sensitivity on CelebA. This parameter sweep is not presented as a nonexistent module-removal experiment.

Config CelebA aggregate score, lower is better Note
FairGrad 0.37 Comparison method
LDC-MTL, ฮป = 0.005 -1.27 Weaker weighted-loss penalty
LDC-MTL, ฮป = 0.008 -1.16 Parameter sweep
LDC-MTL, ฮป = 0.01 -1.31 Best listed configuration
LDC-MTL, ฮป = 0.02 -0.96 Stronger weighted-loss penalty

Key Findings

  • Efficiency gains are most prominent with many tasks. Figure 4 normalizes LS training time to 1.00. LDC-MTL takes 1.09, 1.06, 1.11, and 1.07 on Cityscapes, NYU-v2, QM9, and CelebA, respectively; Nash-MTL takes 12.49 on CelebA. These factors are all relative to LS, so 12.49 is not the exact speedup of LDC-MTL over Nash-MTL.
  • Training trajectories reveal more than final weights. In Figure 6, QM9 routing weights change substantially early in training and approach similar values later, yet LS performs much worse in aggregate. This supports the value of dynamic training without establishing final weights as the only controlled variable.
  • Conflict reduction is an observation, not a hard constraint. The authors randomly select 8 of CelebA's 40 tasks and inspect gradient cosine similarities after the 15th epoch. This supports reduced conflict but does not prove that every task pair is conflict-free throughout training.

Highlights & Insights

  • Separate balancing preferences from task optimization. The discrepancy objective defines balance for the router, while the weighted-loss penalty prevents a sole pursuit of numerical similarity. This perspective clarifies whether a multi-task method optimizes total loss, progress, or the distribution across tasks.
  • The main engineering simplification targets the bilevel solver. Instead of accepting an expensive bilevel procedure and optimizing its implementation afterward, the method inspects which theoretical term is empirically small. Omitting it removes both the auxiliary model and inner loop, directly connecting the computational benefit to the mechanism.
  • Assess balance through both mean and dispersion. Dispersion alone can reward collective deterioration, whereas Table 4 reports the mean, standard deviation, and extremes together. Transferring this idea to other multi-objective settings should retain both magnitude and distribution as evaluation dimensions.

Limitations & Future Work

  • Loss units still matter. Comparing raw or routing-weighted losses does not automatically provide invariance to unit changes. Heterogeneous tasks warrant a more explicit sensitivity study of normalization and the choice of \(\tau\).
  • The simplification requires an additional assumption. The small-gradient observation covers the experiments shown, while PL and smoothness conditions are not unconditional properties of general deep networks. Future work could examine when this gradient becomes large and whether a few adaptive inner steps should be restored.
  • Balance does not imply universal superiority to STL. QM9's score of 49.5 remains positive, and the NYU-v2 aggregate result does not outperform FairGrad. Applications with unequal task preferences also require explicit preference control rather than a default preference for similar task losses.
  • Reproduction is bounded by the available material. The cache ends at the references, omits the cited appendix, and contains corrupted equation extraction. This note does not infer missing hyperparameters, standard deviations, precise convergence bounds, or component-removal experiments. The authors propose recommendation systems as future work, not as an already validated application.
  • Compared with LS, UW, and DWA: These methods use fixed weights or dynamically weight tasks through uncertainty or learning progress; LDC-MTL explicitly targets task-loss gaps. It preserves the low cost of scalarization but introduces a penalty coefficient and conditions for the bilevel approximation.
  • Compared with PCGrad, CAGrad, Nash-MTL, and FairGrad: These methods address conflict or fairness at the task-gradient level; LDC-MTL works at the loss level without retaining every independent task gradient. The NYU-v2 results also show that inexpensive computation and the best metric value are different claims.
  • Compared with FAMO and GO4Align: These belong to the efficient dynamic-balancing approaches considered in the paper, while LDC-MTL organizes optimization around explicit discrepancy and a first-order bilevel approximation. The paper labels its NYU-v2 reproduction of GO4Align as GO4Align*, which should not be equated with published conclusions under every setting.
  • Compared with ConsMTL: ConsMTL assigns shared and task-specific parameters to different levels, giving the optimization variables different roles. The paper includes it in timing comparisons but not in the main performance tables, so these results do not establish an accuracy advantage over ConsMTL.

Rating

  • Novelty: 4/5. Explicit discrepancy control and single-loop bilevel approximation form a clear combination, but build on existing scalarization and first-order penalty methods.
  • Experimental Thoroughness: 4/5. The study covers 4 benchmarks, 2โ€“40 tasks, efficiency, loss distributions, and parameter sensitivity; full component ablations and reproduction details still require appendix support.
  • Writing Quality: 4/5. The progression from motivation to solver approximation is clear, but complexity accounting and additional theoretical conditions require careful distinction.
  • Value: 4/5. Useful when improving aggregate multi-task performance while retaining training costs close to LS is a practical priority.