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

## Abstract

Abstract Multi-task learning (MTL) has been widely adopted for its ability to simultaneously learn multiple tasks. While existing gradient manipulation methods often yield more balanced solutions than simple scalarization-based approaches, they typically incur a significant computational overhead of 𝒪 ​ ( K ) \mathcal{O}(K) in both time and memory, where K K is the number of tasks. In this paper, we propose LDC-MTL, a simple and scalable loss discrepancy control approach for MTL, formulated from a bilevel optimization perspective. Our method incorporates two key components: (i) a bilevel formulation for fine-grained loss discrepancy control, and (ii) a scalable first-order bilevel algorithm that requires only 𝒪 ​ ( 1 ) \mathcal{O}(1) time and memory. Theoretically, we prove that LDC-MTL guarantees convergence not only to a stationary point of the bilevel problem with loss discrepancy control but also to an ϵ \epsilon -accurate Pareto stationary point for all K K loss functions under mild conditions. Extensive experiments on diverse multi-task datasets demonstrate the superior performance of LDC-MTL in both accuracy and efficiency.

## 1 Introduction

In recent years, Multi-Task Learning (MTL) has received increasing attention for its ability to predict multiple tasks simultaneously using a single model, thereby reducing computational overhead. This versatility has enabled a wide range of applications, including autonomous driving (Chen et al., 2018), recommendation systems (Wang et al., 2020), and natural language processing (Zhang et al., 2022).

One of the main challenges in MTL is the imbalance and discrepancy in task losses, where different tasks progress at uneven rates during training. This discrepancy stems from several sources: variations in loss magnitudes due to differing units or scales (e.g., meters vs. millimeters) (Kendall et al., 2018; Liu et al., 2019), heterogeneity in task types (e.g., regression vs. classification) (Dai et al., 2023; Lin et al., 2021), and conflicting gradient directions across tasks (Yu et al., 2020; Liu et al., 2021a). When unaddressed, such discrepancies may cause certain tasks to dominate the optimization trajectory, ultimately leading to degraded performance on others.

To mitigate this issue, MTL research generally follows two main paradigms. The first is the class of scalarization-based methods, which transform MTL into a single-objective optimization problem by aggregating task losses, typically through weighted or averaged sums. Early works adopted static weighting schemes for their simplicity and scalability (Caruana, 1997), but these often led to degraded multi-task performance relative to single-task baselines, largely due to persistent gradient conflicts under fixed weights (Xiao et al., 2024). As a remedy, more recent approaches explore dynamic loss weighting strategies that adapt during training (Kendall et al., 2018; Liu et al., 2019; Lin et al., 2021; Dai et al., 2023). However, these methods do not explicitly address loss discrepancy, allowing task interference to persist and leading to imbalanced performance across tasks.
The second line of work involves gradient manipulation techniques, which aim to promote balanced optimization by explicitly resolving gradient conflicts. These approaches seek update directions that are more equitable across tasks (Désidéri, 2012; Liu et al., 2021a; Ban & Ji, 2024; Navon et al., 2022; Yu et al., 2020; Fernando et al., 2023; Xiao et al., 2024). While often effective in reducing task interference, they typically require computing and storing gradients from all $K$ tasks at each iteration, incurring $O(K)$ time and memory costs. This scalability bottleneck poses challenges for large-scale MTL scenarios involving deep architectures and massive datasets.

In this paper, we propose a simple and scalable loss discrepancy control approach for MTL from a novel bilevel optimization perspective. Our approach comprises two key components: a bilevel formulation for fine-grained loss discrepancy control, and a scalable first-order bilevel algorithmic design. Our specific contributions are summarized as follows.

- $\bullet$
Bilevel formulation for loss discrepancy control. At the core of our bilevel formulation, the lower-level problem optimizes the model parameters by minimizing a weighted sum of individual loss functions. Meanwhile, the upper-level problem adjusts these weights to minimize the discrepancies among the loss functions, ensuring balanced learning across tasks.
- $\bullet$
Efficient algorithms with $O(1)$ time and memory cost. We develop Loss Discrepancy Control for Multi-Task Learning (LDC-MTL), a highly efficient algorithm tailored to solve the proposed bilevel problem with loss discrepancy control. Unlike traditional bilevel methods, LDC-MTL has a fully single-loop structure without second-order gradient computations, resulting in an overall $\mathcal{O}(1)$ time and memory complexity. The $2$-task toy example in [Figure˜1](https://arxiv.org/html/2502.08585v3#S0.F1) illustrates that our LDC-MTL method achieves a more balanced solution compared to other competitive approaches while maintaining superior computational efficiency.
- $\bullet$
Empirical performance. Extensive experiments show that our proposed LDC-MTL method outperforms a wide range of scalarization-based and gradient manipulation approaches across multiple supervised multi-task datasets, including QM9 (Ramakrishnan et al., 2014), CelebA (Liu et al., 2015), and Cityscapes (Cordts et al., 2016), while also demonstrating superior efficiency and scalability.
- $\bullet$
Experimental analysis. We conduct a deeper exploration showing that task losses under LDC-MTL become more concentrated and consistently lower. As a side effect of controlling loss discrepancies, gradient conflicts are also reduced. In addition, when compared to weight-swept linear scalarization (LS), LDC-MTL solutions align with the Pareto frontier (PF) and achieve consistently better results. Our experiments also suggest that the dynamic training trajectory, rather than just the final model parameters, plays a key role in the strong performance of our method. Overall, these results highlight the importance of dynamic loss discrepancy control in LDC-MTL.
- $\bullet$
Theoretical guarantees. Theoretically, we show that LDC-MTL guarantees convergence not only to a stationary point of the bilevel problem with loss discrepancy control but also to an $\epsilon$-accurate Pareto stationary point for all
$K$ individual loss functions under suitable conditions.

## 2 Related Works

Multi-task learning. MTL has recently garnered significant attention in practical applications. One line of research focuses on model architecture, specifically designing various sharing mechanisms (Kokkinos, 2017; Ruder et al., 2019). Another direction addresses the mismatch in loss magnitudes across tasks, proposing methods to balance them. For example, Kendall et al. (2018) balanced tasks by weighting loss functions based on homoscedastic uncertainties, while Liu et al. (2019) dynamically adjusted weights by considering the rate of change in loss values for each task.

Besides, one prominent approach frames MTL as a Multi-Objective Optimization (MOO) problem. Sener & Koltun (2018) introduced this perspective in deep learning, inspiring methods based on the Multi-Gradient Descent Algorithm (MGDA) (Désidéri, 2012). Subsequent work has aimed to address gradient conflicts. For instance, PCGrad (Yu et al., 2020) resolves conflicts by projecting gradients onto the normal plane, GradDrop (Chen et al., 2020) randomly drops conflicting gradients, and CAGrad (Liu et al., 2021a) constrains update directions to balance gradients. Additionally, Nash-MTL (Navon et al., 2022) formulates MTL as a bargaining game among tasks, while FairGrad (Ban & Ji, 2024) incorporates $\alpha$-fairness into gradient adjustments. Achituve et al. (2024) introduces a novel gradient aggregation approach using Bayesian inference to reduce the running time.

Prior works (Kurin et al., 2022; Xin et al., 2022) have shown that several gradient-based MTL methods fail to outperform linear scalarization (LS) with weight sweeping. In contrast, our extensive experiments demonstrate that our method yields solutions that dominate those obtained by weight-swept LS in [Figure˜5](https://arxiv.org/html/2502.08585v3#S6.F5). On the theoretical side, Zhou et al. (2022) analyzed the convergence properties of stochastic MGDA, and Fernando et al. (2023) proposed a method to reduce bias in the stochastic MGDA with theoretical guarantees. More recent advancements include a double-sampling strategy with provable guarantees introduced by Xiao et al. (2024) and Chen et al. (2024).

Bilevel optimization.
Bilevel optimization, first introduced by Bracken & McGill (1973), has been extensively studied over the past few decades. Early research primarily treated it as a constrained optimization problem (Hansen et al., 1992; Shi et al., 2005). More recently, gradient-based methods have gained prominence due to their effectiveness in machine learning applications. Many of these approaches approximate the hypergradient using either linear systems (Domke, 2012; Ji et al., 2021) or automatic differentiation techniques (Maclaurin et al., 2015; Franceschi et al., 2017). However, these methods become impractical in large-scale settings due to their significant computational cost (Xiao & Ji, 2023; Yang et al., 2024b).
The primary challenge lies in the high cost of gradient computation: approximating the Hessian-inverse vector requires multiple first- and second-order gradient evaluations, and the nested sub-loops exacerbate this inefficiency. To address these limitations, recent studies have focused on reducing the computational burden of second-order gradients. For example, some methods reformulate the lower-level problem using value-function-based constraints and solve the corresponding Lagrangian formulation (Kwon et al., 2023; Yang et al., 2024a). The work studies
convex bilevel problems and proposes a zeroth-order optimization method with finite-time
convergence to the Goldstein stationary point (Chen et al., 2023). In this work, we propose a simplified first-order bilevel method for MTL, motivated by intriguing empirical findings.

## 3 Preliminary

Scalarization-based methods. MTL aims to optimize multiple tasks (objectives) simultaneously with a single model. The straightforward approach is to optimize a weighted summation of all loss functions:
$\min_{x}L_{total}(x)=\sum_{i=1}^{K}w_{i}l_{i}(x),$
where $x\in\mathbb{R}^{d}$ denotes the model parameter, $l_{i}(x):\mathbb{R}^{d}\rightarrow\mathbb{R}_{\geq 0}$ represents the loss function of the $i$-th task and $K$ is the number of tasks.
This approach faces three key challenges: 1) loss values could differ in scale, 2) fixed weights can lead to significant gradient conflicts, potentially allowing one task to dominate the learning process (Xiao et al., 2024; Wang et al., 2024); and 3) the overall performance is highly sensitive to the weighting of different losses (Kendall et al., 2018). Consequently, such methods often struggle with performance imbalances across tasks.

Gradient manipulation methods. To mitigate gradient conflicts, gradient manipulation methods dynamically compute an update $d^{t}$ at each epoch to balance progress across tasks, where $t$ is the epoch index. The update $d^{t}$ is typically a convex combination of task gradients, expressed as:

$$ $\displaystyle d^{t}=G(x^{t})w^{t},\;\;\text{where}\;w^{t}=h(G(x^{t})),$ $$

with
$G(x^{t})=[\nabla l_{1}(x^{t}),\nabla l_{2}(x^{t}),\dots,\nabla l_{K}(x^{t})]^{\top}.$
The weight vector $w^{t}$ is determined by a function $h(\cdot):\mathbb{R}^{K\times d}\rightarrow\mathbb{R}^{K}$, which varies depending on the specific method.
However, these methods often require computing and storing the gradients of all $K$ tasks during each epoch, making them less scalable and resource-intensive, particularly in large-scale scenarios.
Therefore, it is necessary to develop lightweight methods that achieve balanced performance.

Pareto concepts. Solving the MTL problem is challenging because it is difficult to identify a common $x$ that achieves the optima for all tasks. Instead, a widely accepted target is finding a Pareto stationary point. Suppose we have two points $x_{1}$ and $x_{2}$. It is claimed that $x_{1}$ dominates $x_{2}$ if $l_{i}(x_{1})\leq l_{i}(x_{2})\;\forall i\in[K]$, and $\exists j\;l_{j}(x_{1})<l_{j}(x_{2})$. A point is Pareto optimal if it is not dominated by any other point, implying that no task can be improved further without sacrificing another. Besides, a point $x$ is a Pareto stationary point if $\min_{w\in\mathcal{W}}\|G(x)w\|=0$.

## 4 Loss Discrepancy Control for Multi-Task Learning

In this section, we present our bilevel loss discrepancy control framework for multi-task learning. As illustrated in [Figure˜2](https://arxiv.org/html/2502.08585v3#S4.F2), this framework contains a fine-grained bilevel loss discrepancy control procedure and a simplified first-order optimization pipeline. We first introduce the motivation and high-level idea that guide the overall design before moving to the technical details.

Figure: Figure 2: Our bilevel loss discrepancy control pipeline for multi-task learning. First, different task losses will be computed. Then, the lower-level problem optimizes the model parameter $x^{t}$ by minimizing the weighted sum of task losses, and the upper-level problem optimizes the router model parameter $W^{t}$ for fine-grained loss discrepancy control.
Refer to caption: x2.png

### 4.1 Motivation and High-Level Idea

Loss discrepancy commonly arises in MTL for several reasons.
First, loss functions may operate on different scales due to task heterogeneity
(e.g., classification vs. regression, as in QM9 (Ramakrishnan et al., 2014)) or because
they are measured in varying units (e.g., meters, centimeters, or millimeters;
(Kendall et al., 2018)).
Second, losses can evolve at different rates depending on task difficulty (Si et al., 2025).
As a result, although all task losses are expected to converge toward zero, their convergence
speeds are often imbalanced when treated with equal importance.

To mitigate this issue, some prior approaches, e.g., (Si et al., 2025), introduce
a ranking-based strategy that optimizes only a subset of tasks with the largest current losses.
While effective in reducing discrepancy, such methods require solving an additional
non-smooth optimization problem, which limits scalability. Moreover, Roh et al. (2020) proposes to reduce
loss gaps across different groups to enhance model fairness. In preference-based MOO, the goal is to identify a Pareto-optimal
point such that
$p_{1}l_{1}=p_{2}l_{2}=\dots=p_{K}l_{K},$
where $p_{1},\dots,p_{K}$ denote user-specified preferences.
In balanced MTL, these preferences are naturally equal, reducing the goal to aligning
task losses so that they remain close to one another.

Motivated by this perspective, we explore a natural question:
can loss discrepancy in MTL be addressed by explicitly reducing the mutual gaps among task losses,
while simultaneously ensuring that the solutions remain sufficiently close to the
Pareto-stationary set? In the next section, we provide a positive solution through a bilevel optimization formulation.

### 4.2 Bilevel formulation for loss discrepancy control

Building on this motivation, we propose a bilevel optimization-based approach, where the upper-level problem aims to mitigate the loss discrepancy, while the lower-level problem ensures that the solution
remains within the Pareto-stationary set.
Specifically, we consider the following formulation:

$$ $\displaystyle\min_{W}\sum_{i=1}^{K-1}|\tau_{i}{l}_{i}(x^{*})-\tau_{i+1}{l}_{i+1}(x^{*})|:=f(W,x^{*})$ $\displaystyle\text{s.t. }x^{*}\in\arg\min_{x}\sum_{i=1}^{K}\sigma_{i}(W){l}_{i}(x):=g(W,x),$ (1) $$

where we denote $x^{*}=x^{*}(W)$ for notational convenience. Note that we define a routing function $\sigma(W)\in\mathbb{R}^{K}$, which is parameterized by a small neural network with a softmax output layer. It takes the shared feature as the input and the output weights $K$ for different tasks. For the upper-level weight vector $\tau=(\tau_{1},...,\tau_{K})$, it controls the loss discrepancy, and we provide two effective options: (i) $\tau=\sigma(W)$ and (ii) $\tau={\mathbf{1}}$ that work well in experiments. It can be seen from [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2) that the lower-level problem minimizes the weighted sum of losses w.r.t. the model parameters $x$, while the upper-level problem minimizes the accumulated weighted loss gaps w.r.t. the parameters $W$, controlling the loss discrepancy among tasks. Notably, the optimal solution does not require all losses to be equal, and there will be a trade-off between loss minimization and loss discrepancy control, as shown in [Remark˜1](https://arxiv.org/html/2502.08585v3#Thmremark1).

### 4.3 Scalable first-order algorithm design

To enable large-scale applications, we adopt an efficient first-order method to solve the problem in [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2). Inspired by recent advances in first-order bilevel optimization (Kwon et al., 2023; Yang et al., 2023), we reformulate it into an equivalent constrained optimization problem as follows.

$$ $\displaystyle\min_{W,x}f(W,x)\;\;\text{s.t. }\underbrace{\sum_{i=1}^{K}\sigma_{i}(W){l}_{i}(x)-\sum_{i=1}^{K}\sigma_{i}(W){l}_{i}(x^{*})}_{\text{penalty function }p(W,x)}\leq 0.$ $$

Then, given a penalty constant $\lambda>0$, penalizing $p(W,x)$ into the upper-level loss function yields

$$ $\displaystyle\min_{W,x}f(W,x)+\lambda\sum_{i=1}^{K}\Big(\sigma_{i}(W){l}_{i}(x)-\sigma_{i}(W){l}_{i}(x^{*})\Big).$ (2) $$

Intuitively, a larger $\lambda$ allows more precise training on model parameters $x$ such that $x$ converges closer to $x^{*}$. Conversely, a smaller $\lambda$ prioritizes upper-level loss discrepancy control during training. The main challenge of solving the penalized problem above lies in the updates of $W$, as shown below:

$$ $\displaystyle W^{t+1}=W^{t}-\alpha\Big(\nabla_{W}f(W^{t},x^{t})+\lambda\big(\nabla_{W}g(W^{t},x^{t})-\nabla_{W}g(W^{t},z_{N}^{t})\big)\Big),$ (3) $$

where $t$ is the epoch index, $\alpha$ is the step size, and $z_{N}^{t}$ is an approximation of $x_{t}^{*}\in\arg\min_{x}g(W^{t},x)$ through the following loop of $N$ iterations each epoch.

$$ $\displaystyle z_{n+1}^{t}=z_{n}^{t}-\beta\nabla_{z}g(W^{t},z_{n}^{t}),n=0,1,...N-1,$ (4) $$

where $N$ is typically chosen to be sufficiently large, ensuring that $z_{N}^{t}$ closely approximates $x_{t}^{*}$ (the full algorithm is provided in Algorithm [B](https://arxiv.org/html/2502.08585v3#A2) in the appendix).
Consequently, this sub-loop of iterations incurs significant computational overhead, driven by the high dimensionality of $z$ (matching that of the model parameters) and the large value of $N$.

###### Remark 1 .

This formulation encourages task losses to be close but not necessarily equal. The penalty constant $\lambda$ controls the trade-off between minimizing the sum of weighted losses and reducing the loss discrepancy among tasks.

Figure: Figure 3: Gradient norm values during the training process on the Cityscapes dataset. Similar phenomena have also been observed in other datasets.
Refer to caption: x3.png

###### Remark 1 .

#### Scalable Algorithm.

Our experiments show that the gradient norm $\|\nabla_{W}g(W^{t},z_{N}^{t})\|$ remains small, typically orders of magnitude smaller than the gradient norm $\|\nabla_{W}g(W^{t},x^{t})\|$, which is used to update the outer parameters $W$. This behavior is illustrated in [Figure˜3](https://arxiv.org/html/2502.08585v3#S4.F3) in the Appendix. Specifically, we set $N=50$ during training. On average, the ratio $\|\nabla_{W}g(W^{t},x^{t})\|/\|\nabla_{W}g(W^{t},z_{N}^{t})\|$ exceeds 100, despite some fluctuations. Under these conditions, the term $\nabla_{W}g(W,z_{N}^{t})$ can be safely neglected, thereby eliminating the need for the expensive loop in [eq.˜4](https://arxiv.org/html/2502.08585v3#S4.E4).
This approximation has been effectively utilized in large-scale applications, such as fine-tuning large language models, to reduce memory and computational costs (Shen et al., 2024a). It also serves as a foundation for our proposed algorithm, Loss Discrepancy Control for Multi-Task Learning (LDC-MTL), described in Algorithm [4.3](https://arxiv.org/html/2502.08585v3#S4.SS3). LDC-MTL employs a fully single-loop structure, which requires only a single gradient computation for both variables per epoch, resulting in a $\mathcal{O}(1)$ time and memory cost. In [Section˜6](https://arxiv.org/html/2502.08585v3#S6), we show that our LDC-MTL method attains both an $\epsilon$-accurate stationary point for the bilevel problem in [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2) and an $\epsilon$-accurate Pareto stationary point for the original loss functions under mild conditions.

## 5 Empirical Results

In this section, we conduct extensive practical experiments under multi-task classification, regression, and mixed settings to demonstrate the effectiveness of our method. Full experimental details can be found in [Appendix˜A](https://arxiv.org/html/2502.08585v3#A1). All experiments are conducted on one NVIDIA A6000.

Baselines and evaluation. To demonstrate the effectiveness of our proposed method, we evaluate its performance against a broad range of baseline approaches.
The compared methods include scalarization-based algorithms, such as Linear Scalarization (LS), Scale-Invariant (SI), Random Loss Weighting (RLW) (Lin et al., 2021), Dynamic Weight Average (DWA) (Liu et al., 2019), Uncertainty Weighting (UW) (Kendall et al., 2018), FAMO (Liu et al., 2024), and GO4Align (Shen et al., 2024b). We also benchmark against gradient manipulation methods, including Multi-Gradient Descent Algorithm (MGDA) (Désidéri, 2012), PCGrad (Yu et al., 2020), GradDrop (Chen et al., 2020), CAGrad (Liu et al., 2021a), IMTL-G (Liu et al., 2021b), MoCo (Fernando et al., 2023), Nash-MTL (Navon et al., 2022), and FairGrad (Ban & Ji, 2024). To provide a comprehensive evaluation, we report the performance of each individual task and employ one additional metric: $\bf\Delta m\%$ to quantify overall performance (Maninis et al., 2019). The $\bf\Delta m\%$ metric measures the average relative performance drop of a multi-task model compared to its corresponding single-task learning (STL). Formally, it is defined as:
$\Delta m\%=\frac{1}{K}\sum_{i=1}^{K}(-1)^{\delta_{k}}(M_{m,k}-M_{b,k})/M_{b,k}\times 100,$
where $M_{m,k}$ and $M_{b,k}$ represent the performance of the $k$-th task for the multi-task model $m$ and single-task model $b$, respectively. The indicator $\delta_{k}=1$ if higher values indicate better performance and 0 otherwise.

Table: Table 1: Results on Cityscapes (2-task) dataset. Each experiment is repeated 3 times with different random seeds, and the average is reported. The best results are highlighted in bold, while the second-best results are indicated with underlines. Following prior works Liu et al. (2021a); Fernando et al. (2023); Xiao et al. (2024); Ban & Ji (2024), we report the mean values of $\Delta m\%$ for all results in the main text, with standard deviations provided in [Appendix˜A](https://arxiv.org/html/2502.08585v3#A1).

### 5.1 Experimental results

Results on the four benchmark datasets are provided in [Table˜1](https://arxiv.org/html/2502.08585v3#S5.T1), [Table˜2](https://arxiv.org/html/2502.08585v3#S5.T2), [Table˜6](https://arxiv.org/html/2502.08585v3#A1.T6), and [Table˜7](https://arxiv.org/html/2502.08585v3#A1.T7) in the appendix. We observe that LDC-MTL outperforms existing methods on both the CelebA and QM9 datasets, achieving the lowest performance drops of $\Delta m\%=-1.31$ and $\Delta m\%=49.5$, respectively. Detailed results for the QM9 dataset illustrate that it achieves a balanced performance across all tasks.
These results highlight the effectiveness of our method in handling a large number of tasks in both classification and regression settings. Meanwhile, it achieves the lowest performance drop, with $\Delta m\%=-0.57$ on the Cityscapes dataset, while delivering comparable results on the NYU-v2 dataset, where the detailed results are shown in [Table˜7](https://arxiv.org/html/2502.08585v3#A1.T7) in the appendix.
These findings highlight the capability of LDC-MTL to effectively handle mixed MTL scenarios.

Table: Table 2: Results on CelebA (40-task), QM9 (11-task), and NYU-v2 (3-task) datasets. The best results are highlighted in bold, while the second-best results are indicated with underlines.

Efficiency comparison. We compare the running time of well-performing approaches in [Figure˜4](https://arxiv.org/html/2502.08585v3#S5.F4). In particular, our method introduces negligible overhead compared to LS with at most a $1.11\times$ increase, aligning with other $\mathcal{O}(1)$ methods such as GO4Align and FAMO.
In contrast, gradient manipulation methods, which take the computational cost $\mathcal{O}(K)$, become significantly slower in many-task scenarios. For example, Nash-MTL requires approximately $12\times$ more training time than LDC-MTL on the CelebA dataset.

### 5.2 Experimental analysis

Loss discrepancy and gradient conflict.
To demonstrate the effectiveness of our bilevel formulation for loss discrepancy control, we conduct a detailed analysis of the loss distribution on the CelebA dataset, comparing linear scalarization (LS) with our proposed method.
As shown in [Figure˜7](https://arxiv.org/html/2502.08585v3#A1.F7) (left) and statistics in [Table˜3](https://arxiv.org/html/2502.08585v3#S5.T3), the distribution of all 40 task-specific losses reveals that our approach yields more concentrated and consistently lower values. Moreover, we randomly select 8 out of 40 tasks and check the gradient cosine similarity among them. [Figure˜7](https://arxiv.org/html/2502.08585v3#A1.F7) (right) illustrates the cosine similarities of task gradients after the 15th epoch, which shows that the gradient conflict is mitigated as a side-effect.

Table: Table 3: Statistics of loss values for LDC-MTL and LS on the CelebA dataset. Lower mean and std indicate better and more stable performance.

Comparison with weight-swept LS.
Prior work has noted that certain multi-task learning methods, such as RLW and PCGrad, do not outperform LS, which can approximate the Pareto front through weight sweeping (Kurin et al., 2022; Xin et al., 2022). To this end, we conduct a careful comparison between our LDC-MTL and LS with weight sweeping on the Cityscapes dataset. From [Table˜9](https://arxiv.org/html/2502.08585v3#A1.T9) in the appendix, even after careful weight sweeping, LS does not outperform our approach. Furthermore, the scatter plot in [Figure˜5](https://arxiv.org/html/2502.08585v3#S6.F5) (a-b) illustrates that LDC-MTL solutions form the Pareto frontier.

Figure: Figure 4: Time scale comparison among well-performing approaches, with LS considered the reference method for standard time.
Refer to caption: x4.png

Impact of task weight. To further investigate the importance of dynamic weighting, we provide the evolution of task weights for the 11 tasks on the QM9 dataset in [Figure˜5](https://arxiv.org/html/2502.08585v3#S6.F5) (right). It is evident that the task weights adapt meaningfully during the
early stages of training and gradually converge later on. Besides, the converged weights are nearly equal, recovering to LS. However, LS does not perform well as shown in [Table˜6](https://arxiv.org/html/2502.08585v3#A1.T6). These findings suggest that, for dynamic-weight methods, the *training trajectory*, not just the final weights, plays a critical role in achieving strong performance.

Table: Table 4: Parameter tuning results on the CelebA dataset.

Hyperparameter sensitivity. In our method, hyperparameters include the step size $\alpha$ and the penalty constant $\lambda$. For the step size, we adopt the settings from prior experiments without extensive tuning. While we use the same step size for updates to both $W$ and $x$ in our implementation, these can be adjusted independently in practice. For the penalty constant $\lambda$, we determine optimal values through a grid search and provide additional experimental results in [Table˜4](https://arxiv.org/html/2502.08585v3#S5.T4) and [Table˜8](https://arxiv.org/html/2502.08585v3#A1.T8) in the appendix.

## 6 Theoretical Analysis

###### Definition 1 .

Given $L>0$, a function $\ell$ is said to be $L$-Lipschitz-continuous on $\mathcal{X}$ if it holds for any $x,x^{\prime}\in\mathcal{X}$ that $\|\ell(x)-\ell(x^{\prime})\|\leq L\|x-x^{\prime}\|$. A function $\ell$ is said to be $L$-Lipschitz-smooth
if its gradient is $L$-Lipschitz-continuous.

###### Definition 2 (Pareto stationarity) .

We say $x$ is an $\epsilon$-accurate Pareto stationary point for loss functions $\{l_{i}(x)\}$ if $\min_{w\in\mathcal{W}}\|G(x)w\|^{2}=\mathcal{O}(\epsilon)$, where $G(x)=[\nabla l_{1}(x),\nabla l_{2}(x),...,\nabla l_{K}(x)]^{\top}$.

Inspired by Shen & Chen (2023), we also define the following two surrogates of the original bilevel problem in [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2).

###### Definition 3 .

Define two surrogate bilevel problems as

$$ $\displaystyle\mathcal{BP}_{\lambda}:$ $\displaystyle\min_{W,x}f(W,x)+\lambda(g(W,x)-g(W,x^{*})),\mathcal{BP}_{\epsilon}:\min_{W,x}f(W,x)\;\;\text{s.t. }g(W,x)-g(W,x^{*})\leq\epsilon,$ $$

where $\mathcal{BP}_{\lambda}$ is the penalized bilevel problem, and $\mathcal{BP}_{\epsilon}$ recovers to the original problem if $\epsilon=0$.

###### Assumption 1 (Lipschitz and smoothness) .

There exists a constant $L$ such that the upper-level function $f(W,\cdot)$ is $L$-Lipschitz continuous. There exists constants $L_{f}$ and $L_{g}$ such that functions $f(W,x)$ and $g(W,x)$ are $L_{f}$- and $L_{g}$-Lipschitz-smooth.

###### Assumption 2 (Polyak-Lojasiewicz (PL) condition) .

The lower-level function $g(W,\cdot)$ satisfies the $\frac{1}{\mu}$-PL condition if there exists a $\mu>0$ such that given any $W$, it holds for any feasible $x$ that $\|\nabla_{x}g(W,x)\|^{2}\geq\frac{1}{\mu}(g(W,x)-g(W,x^{*})).$

Lipschitz continuity and smoothness are standard assumptions in the study of bilevel optimization (Ghadimi & Wang, 2018; Ji et al., 2021). While the absolute values in the upper-level function in [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2) are non-smooth, they can be easily modified to ensure smoothness, such as by using a soft absolute value function of the form $y=\sqrt{x^{2}+\gamma}$ where $\gamma$ is a small positive constant.
Moreover, the PL condition can be satisfied in over-parameterized neural network settings (Mei et al., 2020; Frei & Gu, 2021). The following theorem presents the convergence analysis of our algorithms.

Figure: Figure 5: (a–b) Comparison between weight-swept LS and LDC-MTL on the Cityscapes dataset. In both cases, LDC-MTL solutions lie on the Pareto frontier. (c) Weight change of 11 tasks using LDC-MTL during the training on the QM9 dataset.
Refer to caption: x5.png

###### Theorem 1 .

Suppose Assumptions [1](https://arxiv.org/html/2502.08585v3#Thmassumption1)- [2](https://arxiv.org/html/2502.08585v3#Thmassumption2) are satisfied. Select hyperparameters

$$ $\displaystyle\alpha\in\Big(0,\frac{1}{L_{f}+\lambda(2L_{g}+L_{g}^{2}\mu)}\Big],\;\beta\in(0,\frac{1}{L_{g}}],\lambda=L\sqrt{3\mu{\epsilon}^{-1}},\text{ and }N=\Omega(\log(\alpha t)).$ $$

(i) Our method with the updates [eq.˜3](https://arxiv.org/html/2502.08585v3#S4.E3) and [eq.˜4](https://arxiv.org/html/2502.08585v3#S4.E4) (i.e., Algorithm [B](https://arxiv.org/html/2502.08585v3#A2) in the appendix) finds an $\epsilon$-accurate stationary point of the problem $\mathcal{BP}_{\lambda}$. If this stationary point is a local/global solution to $\mathcal{BP}_{\lambda}$, it is also a local/global solution to $\mathcal{BP}_{\epsilon}$. Furthermore, it is also an $\epsilon$-accurate Pareto
stationary point for loss functions $l_{i}(x),i=1,...,K.$

(ii) Moreover, if $\|\nabla_{W}g(W^{t},z_{N}^{t})\|=\mathcal{O}(\epsilon)$ for $t=1,...,T$. The simplified method in Algorithm [4.3](https://arxiv.org/html/2502.08585v3#S4.SS3) also achieves the same convergence guarantee as that in $(i)$.

The complete proof is provided in [Theorem˜2](https://arxiv.org/html/2502.08585v3#Thmtheorem2).
In the first part of [Theorem˜1](https://arxiv.org/html/2502.08585v3#Thmtheorem1), we establish a connection between the stationarity of $\mathcal{BP}_{\lambda}$ and Pareto stationarity. Secondly, it introduces an additional gradient vanishing assumption, which has been validated in our experiments. It demonstrates that our simplified LDC-MTL method can also attain an $\epsilon$-accurate stationary point for the problem $\mathcal{BP}_{\lambda}$ and an $\epsilon$-accurate Pareto stationary point for the original loss functions.

###### Definition 1 .

###### Definition 2 (Pareto stationarity) .

###### Definition 3 .

###### Assumption 1 (Lipschitz and smoothness) .

###### Assumption 2 (Polyak-Lojasiewicz (PL) condition) .

###### Theorem 1 .

## 7 Conclusion

We introduced LDC-MTL, a scalable loss discrepancy control approach for multi-task learning based on bilevel optimization. Our method achieves efficient loss discrepancy control with only $\mathcal{O}(1)$ time and memory complexity while guaranteeing convergence to both a stationary point of the bilevel problem and an $\epsilon$-accurate Pareto stationary point for all task loss functions. Extensive experiments demonstrate that LDC-MTL outperforms existing methods in both accuracy and efficiency, highlighting its effectiveness for large-scale MTL. For future work, we plan to explore the application of our method to broader multi-task learning problems, including recommendation systems.

## References

- Achituve et al. (2024)
Idan Achituve, Idit Diamant, Arnon Netzer, Gal Chechik, and Ethan Fetaya.
Bayesian uncertainty for gradient aggregation in multi-task learning.
*arXiv preprint arXiv:2402.04005*, 2024.
- Badrinarayanan et al. (2017)
Vijay Badrinarayanan, Alex Kendall, and Roberto Cipolla.
Segnet: A deep convolutional encoder-decoder architecture for image segmentation.
*IEEE Transactions on Pattern Analysis and Machine Intelligence*, 39(12):2481–2495, 2017.
- Ban & Ji (2024)
Hao Ban and Kaiyi Ji.
Fair resource allocation in multi-task learning.
*arXiv preprint arXiv:2402.15638*, 2024.
- Bracken & McGill (1973)
Jerome Bracken and James T McGill.
Mathematical programs with optimization problems in the constraints.
*Operations Research*, 21(1):37–44, 1973.
- Caruana (1997)
Rich Caruana.
Multitask learning.
*Machine Learning*, 28:41–75, 1997.
- Chen et al. (2023)
Lesi Chen, Jing Xu, and Jingzhao Zhang.
Bilevel optimization without lower-level strong convexity from the hyper-objective perspective.
*arXiv preprint arXiv:2301.00712*, 2023.
- Chen et al. (2024)
Lisha Chen, Heshan Fernando, Yiming Ying, and Tianyi Chen.
Three-way trade-off in multi-objective learning: Optimization, generalization and conflict-avoidance.
*Advances in Neural Information Processing Systems*, 36, 2024.
- Chen et al. (2018)
Yaran Chen, Dongbin Zhao, Le Lv, and Qichao Zhang.
Multi-task learning for dangerous object detection in autonomous driving.
*Information Sciences*, 432:559–571, 2018.
- Chen et al. (2020)
Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov.
Just pick a sign: Optimizing deep multitask models with gradient sign dropout.
*Advances in Neural Information Processing Systems*, 33:2039–2050, 2020.
- Cordts et al. (2016)
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele.
The cityscapes dataset for semantic urban scene understanding.
In *Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, pp. 3213–3223, 2016.
- Dai et al. (2023)
Yanqi Dai, Nanyi Fei, and Zhiwu Lu.
Improvable gap balancing for multi-task learning.
In *Uncertainty in Artificial Intelligence*, pp. 496–506. PMLR, 2023.
- Désidéri (2012)
Jean-Antoine Désidéri.
Multiple-gradient descent algorithm (mgda) for multiobjective optimization.
*Comptes Rendus Mathematique*, 350(5-6):313–318, 2012.
- Domke (2012)
Justin Domke.
Generic methods for optimization-based modeling.
In *Artificial Intelligence and Statistics*, pp. 318–326. PMLR, 2012.
- Fernando et al. (2023)
Heshan Fernando, Han Shen, Miao Liu, Subhajit Chaudhury, Keerthiram Murugesan, and Tianyi Chen.
Mitigating gradient bias in multi-objective learning: A provably convergent approach.
In *International Conference on Learning Representations*, 2023.
- Franceschi et al. (2017)
Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil.
Forward and reverse gradient-based hyperparameter optimization.
In *International Conference on Machine Learning*, pp. 1165–1173. PMLR, 2017.
- Frei & Gu (2021)
Spencer Frei and Quanquan Gu.
Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent.
*Advances in Neural Information Processing Systems*, 34:7937–7949, 2021.
- Ghadimi & Wang (2018)
Saeed Ghadimi and Mengdi Wang.
Approximation methods for bilevel programming.
*arXiv preprint arXiv:1802.02246*, 2018.
- Hansen et al. (1992)
Pierre Hansen, Brigitte Jaumard, and Gilles Savard.
New branch-and-bound rules for linear bilevel programming.
*SIAM Journal on Scientific and Statistical Computing*, 13(5):1194–1217, 1992.
- Ji et al. (2021)
Kaiyi Ji, Junjie Yang, and Yingbin Liang.
Bilevel optimization: Convergence analysis and enhanced design.
In *International Conference on Machine Learning*, pp. 4882–4892. PMLR, 2021.
- Kendall et al. (2018)
Alex Kendall, Yarin Gal, and Roberto Cipolla.
Multi-task learning using uncertainty to weigh losses for scene geometry and semantics.
In *Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, pp. 7482–7491, 2018.
- Kokkinos (2017)
Iasonas Kokkinos.
Ubernet: Training a universal convolutional neural network for low-, mid-, and high-level vision using diverse datasets and limited memory.
In *Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, pp. 6129–6138, 2017.
- Kurin et al. (2022)
Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and Pawan K Mudigonda.
In defense of the unitary scalarization for deep multi-task learning.
*Advances in Neural Information Processing Systems*, 35:12169–12183, 2022.
- Kwon et al. (2023)
Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert D Nowak.
A fully first-order method for stochastic bilevel optimization.
In *International Conference on Machine Learning*, pp. 18083–18113. PMLR, 2023.
- Lin et al. (2021)
Baijiong Lin, Feiyang Ye, Yu Zhang, and Ivor W Tsang.
Reasonable effectiveness of random weighting: A litmus test for multi-task learning.
*arXiv preprint arXiv:2111.10603*, 2021.
- Liu et al. (2021a)
Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu.
Conflict-averse gradient descent for multi-task learning.
*Advances in Neural Information Processing Systems*, 34:18878–18890, 2021a.
- Liu et al. (2024)
Bo Liu, Yihao Feng, Peter Stone, and Qiang Liu.
Famo: Fast adaptive multitask optimization.
*Advances in Neural Information Processing Systems*, 36, 2024.
- Liu et al. (2021b)
Liyang Liu, Yi Li, Zhanghui Kuang, J Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang.
Towards impartial multi-task learning.
In *Proceedings of the International Conference on Learning Representations (ICLR) 2021*, 2021b.
- Liu et al. (2019)
Shikun Liu, Edward Johns, and Andrew J Davison.
End-to-end multi-task learning with attention.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 1871–1880, 2019.
- Liu et al. (2015)
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.
Deep learning face attributes in the wild.
In *Proceedings of the IEEE International Conference on Computer Vision*, pp. 3730–3738, 2015.
- Maclaurin et al. (2015)
Dougal Maclaurin, David Duvenaud, and Ryan Adams.
Gradient-based hyperparameter optimization through reversible learning.
In *International Conference on Machine Learning*, pp. 2113–2122. PMLR, 2015.
- Maninis et al. (2019)
Kevis-Kokitsi Maninis, Ilija Radosavovic, and Iasonas Kokkinos.
Attentive single-tasking of multiple tasks.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 1851–1860, 2019.
- Mei et al. (2020)
Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans.
On the global convergence rates of softmax policy gradient methods.
In *International Conference on Machine Learning*, pp. 6820–6829. PMLR, 2020.
- Navon et al. (2022)
Aviv Navon, Aviv Shamsian, Idan Achituve, Haggai Maron, Kenji Kawaguchi, Gal Chechik, and Ethan Fetaya.
Multi-task learning as a bargaining game.
*arXiv preprint arXiv:2202.01017*, 2022.
- Ramakrishnan et al. (2014)
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld.
Quantum chemistry structures and properties of 134 kilo molecules.
*Scientific Data*, 1(1):1–7, 2014.
- Roh et al. (2020)
Yuji Roh, Kangwook Lee, Steven Euijong Whang, and Changho Suh.
Fairbatch: Batch selection for model fairness.
*arXiv preprint arXiv:2012.01696*, 2020.
- Ruder et al. (2019)
Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard.
Latent multi-task architecture learning.
In *Proceedings of the AAAI Conference on Artificial Intelligence*, volume 33, pp. 4822–4829, 2019.
- Sener & Koltun (2018)
Ozan Sener and Vladlen Koltun.
Multi-task learning as multi-objective optimization.
*Advances in Neural Information Processing Systems*, 31, 2018.
- Shen & Chen (2023)
Han Shen and Tianyi Chen.
On penalty-based bilevel gradient descent method.
In *International Conference on Machine Learning*, pp. 30992–31015. PMLR, 2023.
- Shen et al. (2024a)
Han Shen, Pin-Yu Chen, Payel Das, and Tianyi Chen.
Seal: Safety-enhanced aligned llm fine-tuning via bilevel data selection.
*arXiv preprint arXiv:2410.07471*, 2024a.
- Shen et al. (2024b)
Jiayi Shen, Cheems Wang, Zehao Xiao, Nanne Van Noord, and Marcel Worring.
Go4align: Group optimization for multi-task alignment.
*arXiv preprint arXiv:2404.06486*, 2024b.
- Shi et al. (2005)
Chenggen Shi, Jie Lu, and Guangquan Zhang.
An extended kuhn–tucker approach for linear bilevel programming.
*Applied Mathematics and Computation*, 162(1):51–63, 2005.
- Si et al. (2025)
Zhaofeng Si, Shu Hu, Kaiyi Ji, and Siwei Lyu.
Meta-learning with heterogeneous tasks.
In *International Joint Conference on Artificial Intelligence*, pp. 74–94. Springer, 2025.
- Silberman et al. (2012)
Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus.
Indoor segmentation and support inference from rgbd images.
In *Computer Vision–ECCV 2012: 12th European Conference on Computer Vision, Florence, Italy, October 7-13, 2012, Proceedings, Part V 12*, pp. 746–760. Springer, 2012.
- Wang et al. (2020)
Menghan Wang, Yujie Lin, Guli Lin, Keping Yang, and Xiao-ming Wu.
M2grl: A multi-task multi-view graph representation learning framework for web-scale recommender systems.
In *Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining*, pp. 2349–2358, 2020.
- Wang et al. (2024)
Yudan Wang, Peiyao Xiao, Hao Ban, Kaiyi Ji, and Shaofeng Zou.
Finite-time analysis for conflict-avoidant multi-task reinforcement learning.
*arXiv preprint arXiv:2405.16077*, 2024.
- Xiao & Ji (2023)
Peiyao Xiao and Kaiyi Ji.
Communication-efficient federated hypergradient computation via aggregated iterative differentiation.
In *International Conference on Machine Learning*, pp. 38059–38086. PMLR, 2023.
- Xiao et al. (2024)
Peiyao Xiao, Hao Ban, and Kaiyi Ji.
Direction-oriented multi-objective learning: Simple and provable stochastic algorithms.
*Advances in Neural Information Processing Systems*, 36, 2024.
- Xin et al. (2022)
Derrick Xin, Behrooz Ghorbani, Justin Gilmer, Ankush Garg, and Orhan Firat.
Do current multi-task optimization methods in deep learning even help?
*Advances in neural information processing systems*, 35:13597–13609, 2022.
- Yang et al. (2023)
Yifan Yang, Peiyao Xiao, and Kaiyi Ji.
Achieving $\mathcal{O}(\epsilon^{-1.5})$ complexity in hessian/jacobian-free stochastic bilevel optimization.
*arXiv preprint arXiv:2312.03807*, 2023.
- Yang et al. (2024a)
Yifan Yang, Hao Ban, Minhui Huang, Shiqian Ma, and Kaiyi Ji.
Tuning-free bilevel optimization: New algorithms and convergence analysis.
*arXiv preprint arXiv:2410.05140*, 2024a.
- Yang et al. (2024b)
Yifan Yang, Peiyao Xiao, and Kaiyi Ji.
Simfbo: Towards simple, flexible and communication-efficient federated bilevel learning.
*Advances in Neural Information Processing Systems*, 36, 2024b.
- Yu et al. (2020)
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn.
Gradient surgery for multi-task learning.
*Advances in Neural Information Processing Systems*, 33:5824–5836, 2020.
- Zhang et al. (2022)
Zhihan Zhang, Wenhao Yu, Mengxia Yu, Zhichun Guo, and Meng Jiang.
A survey of multi-task learning in natural language processing: Regarding task relatedness and training methods.
*arXiv preprint arXiv:2204.03508*, 2022.
- Zhou et al. (2022)
Shiji Zhou, Wenpeng Zhang, Jiyan Jiang, Wenliang Zhong, Jinjie Gu, and Wenwu Zhu.
On the convergence of stochastic multi-objective gradient manipulation and beyond.
*Advances in Neural Information Processing Systems*, 35:38103–38115, 2022.

## Appendix A Experiment details

### A.1 Experimental setup

Image-Level Classification. CelebA (Liu et al., 2015), one of the most widely used datasets, is a large-scale facial attribute dataset containing over 200K celebrity images. Each image is annotated with 40 attributes, such as the presence of eyeglasses and smiling. Following the experimental setup in Ban & Ji (2024), we treat CelebA as a 40-task multi-task learning (MTL) classification problem, where each task predicts the presence of a specific attribute. Since all tasks involve binary classification with the same binary cross-entropy loss function, we do not apply any normalization for both options of $\tau=\bf 1$ and $\tau=\sigma$.
The network architecture consists of a 9-layer convolutional neural network (CNN) as the shared model, with multiple linear layers serving as task-specific heads. We train the model for 15 epochs using the Adam optimizer with a batch size of 256.

Regression. QM9 (Ramakrishnan et al., 2014) dataset is another widely used benchmark for multi-task regression problems in quantum chemistry. It contains 130K molecules represented as graphs, and 11 properties to be predicted. Though all tasks share the same loss function, mean squared error, they exhibit significantly varying scales: a phenomenon commonly observed in regression tasks but less prevalent in classification tasks, as shown in [Figure˜6](https://arxiv.org/html/2502.08585v3#A1.F6). To mitigate this scale discrepancy, we first adopt the logarithmic normalization such that $\tilde{l}_{i}=\log\left(\frac{l_{i}}{l_{i,0}}\right)$, where $l_{i,0}$ represents the initial loss value for the $i$-th task at each epoch, motivated by Liu et al. (2024), for both options of $\tau=\bf 1$ and $\tau=\sigma$. Our experiments demonstrate that this initialization approach stabilizes training by reducing large fluctuations caused by significant scale variations.
Following the experimental setup in Liu et al. (2024); Navon et al. (2022), we use the same model and data split, 110K molecules for training, 10k for validation, and the rest 10k for testing. The model is trained for 300 epochs with a batch size of 120. The learning rate starts at 1e-3 and is reduced whenever the validation performance stagnates for 5 consecutive epochs.

Dense Prediction.
The Cityscapes dataset (Cordts et al., 2016) consists of 5000 street-scene images designed for two tasks: 7-class semantic segmentation (a classification task) and depth estimation (a regression task). Similarly, the NYU-v2 dataset (Silberman et al., 2012) is widely used for indoor scene understanding and contains 1449 densely annotated images. It includes one pixel-level classification task, semantic segmentation, and two pixel-level regression tasks, 13-class depth estimation, plus surface normal prediction. These datasets provide benchmarks for evaluating the performance of our method in mixed multi-task settings. Since the number of tasks is small and the loss values exhibit minimal variation, we applied rescaled normalization when selecting $\tau=\sigma$ and no normalization when selecting $\tau=\bf 1$. The rescaled normalization normalizes loss values by rescaling each task’s loss using its initial loss value $l_{i}^{\prime}$, such that $\tilde{l}_{i}=\frac{l_{i}}{l_{i}^{\prime}}$. The resulting normalized loss reflects the training progress and ensures comparability across tasks. We follow the same experimental setup described in Liu et al. (2021a); Navon et al. (2022) and adopt MTAN (Liu et al., 2019) as the backbone, which incorporates task-specific attention modules into SegNet (Badrinarayanan et al., 2017). Both models are trained for 200 epochs, with batch sizes of 8 for Cityscapes and 8 for NYU-v2. The learning rates are initialized at 3e-4 and 1e-4 for the first 100 epochs and reduced by half for the remaining epochs, respectively.

In a word, the hyperparameter choices are summarized in [Table˜5](https://arxiv.org/html/2502.08585v3#A1.T5).

**Table 5: Training hyperparameters combination and the best results per dataset.**
| Dataset | Stepsize | Penalty constant | $\bf\tau$ | $\Delta m\%$ |
| --- | --- | --- | --- | --- |
| CelebA | 1e-03 | 0.01 | $\sigma(W)$ | -1.31$\pm$0.26 |
| QM9 | 1e-03 | 0.05 | $\sigma(W)$ | 49.5$\pm$3.64 |
| Cityscapes | 3e-04 | 0.1 | $\bf 1$ | -0.57$\pm$1.17 |
| NYU-v2 | 8e-05 | 0.05 | $\bf 1$ | -4.40$\pm$0.74 |

### A.2 Loss and gradient scales

We show that loss can vary significantly in scale. For example, as shown in [Figure˜6](https://arxiv.org/html/2502.08585v3#A1.F6), we observed that the loss values during training across 11 regression tasks in the QM9 dataset exhibit substantial differences in magnitude, with loss ratios exceeding 1000 in certain instances. Moreover, we have also drawn the loss progression over time of LS. Tasks 1–4 are easier to optimize and dominate the training when using linear scalarization, leading other tasks to converge suboptimally. In contrast, LDC-MTL accounts for loss scale and adaptively reweights tasks using bilevel optimization, enabling better balance. For example, tasks 6–10 improve their final loss scale from $10^{-3}$ to $10^{-5}$.

Figure: Figure 6: Curves of loss values during the training process for all 11 tasks on the QM9 dataset using different methods. The loss values vary significantly across different tasks. For LS, some tasks converge suboptimally due to the loss discrepancy.
Refer to caption: x6.png

Moreover, our experiments reveal that the gradient norm $\|\nabla_{W}g(W^{t},z_{N}^{t})\|$ remains sufficiently small, typically orders of magnitude smaller than the gradient norm $\|\nabla_{W}g(W^{t},x^{t})\|$, which is used to update outer parameters $W$. This behavior is illustrated in [Figure˜3](https://arxiv.org/html/2502.08585v3#S4.F3). Specifically, we set $N=50$ during training. In average, the ratio $\|\nabla_{W}g(W^{t},x^{t})\|/\|\nabla_{W}g(W^{t},z_{N}^{t})\|$ exceeds 100, despite some fluctuations.

### A.3 Toy example

To better understand the benefits of our method, we illustrate the training trajectory along with the training time in a toy example of 2-task learning following the same setting in FAMO (Liu et al., 2024). The loss functions $L_{1}(x),L_{2}(x)$, where $x$ is the model parameter, of two tasks are listed below.

$$ $\displaystyle L_{1}(x)$ $\displaystyle=0.1\times(c_{1}(x)f_{1}(x)+c_{2}(x)g_{1}(x)),\;\;L_{2}(x)=c_{1}(x)f_{2}(x)+c_{2}(x)g_{2}(x)\;\;\text{where}$ $\displaystyle f_{1}(x)$ $\displaystyle=\log{\big(\max(|0.5(-x_{1}-7)-\tanh{(-x_{2})}|,~~0.000005)\big)}+6,$ $\displaystyle f_{2}(x)$ $\displaystyle=\log{\big(\max(|0.5(-x_{1}+3)-\tanh{(-x_{2})}+2|,~~0.000005)\big)}+6,$ $\displaystyle g_{1}(x)$ $\displaystyle=\big((-x_{1}+7)^{2}+0.1*(-x_{2}-8)^{2}\big)/10-20,$ $\displaystyle g_{2}(x)$ $\displaystyle=\big((-x_{1}-7)^{2}+0.1*(-x_{2}-8)^{2})\big/10-20,$ $\displaystyle c_{1}(x)$ $\displaystyle=\max(\tanh{(0.5*x_{2})},~0)~~\text{and}~~c_{2}(x)=\max(\tanh{(-0.5*x_{2})},~0).$ (5) $$

In [Figure˜1](https://arxiv.org/html/2502.08585v3#S0.F1), the black dots represent 5 chosen initial points $\{(-8.5,7.5),(-8.5,5),(0,0),(9,9),(10,-8)\}$ while the black stars represent the converging points on the Pareto front. We use the Adam optimizer and train each method for 50k steps.
Our method can always converge to balanced results efficiently. We use Adam optimizer with a learning rate of 1e-3. The training time is recalculated according to real-time ratios in our machine. We find that LS and MGDA do not converge to balanced points, while FAMO converges to balanced results to some extent. Meanwhile, our method with rescale normalization can always converge to balanced results efficiently.

### A.4 Detailed results

Here we provide detailed results of NYU-v2 in [Table˜7](https://arxiv.org/html/2502.08585v3#A1.T7) and QM9 in [Table˜6](https://arxiv.org/html/2502.08585v3#A1.T6).

Table: Table 6: Detailed results of on QM9 (11-task) dataset. Each experiment is repeated 3 times, and the average is reported. The best results are highlighted in bold, while the second-best results are indicated with underlines.

Table: Table 7: Results on NYU-v2 (3-task) dataset. Each experiment is repeated 3 times with different random seeds, and the average is reported.

### A.5 Parameter tuning

In our experiments, the penalty constant $\lambda$ requires some tuning effort, whereas the choice of step size had relatively less impact and did not require extensive tuning. We have included additional results on the CelebA and Cityscapes datasets in [Table˜4](https://arxiv.org/html/2502.08585v3#S5.T4) and [Table˜8](https://arxiv.org/html/2502.08585v3#A1.T8) that explore the effect of varying $\lambda$, and found that values in the range $\lambda\in[0.02,0.1]$ consistently yield strong performance. Besides, we use $\Delta m\%$ as the evaluation metric, which aggregates performance across all tasks based on their definitions. As a result, it may exhibit slightly higher variance. This behavior is consistent with other experiments; for example, similar variance patterns can be observed in Table 5 in Xiao et al. (2024).

Table: Table 8: Additional results on Cityscapes (2-task) dataset with different $\lambda$ values.

### A.6 Loss discrepancy and gradient conflict

To demonstrate the effectiveness of our bilevel formulation for loss discrepancy control, we conduct a detailed analysis of the loss distribution on the CelebA dataset, comparing linear scalarization (LS) with our proposed method. As shown in [Figure˜7](https://arxiv.org/html/2502.08585v3#A1.F7) (left) and statistics in [Table˜3](https://arxiv.org/html/2502.08585v3#S5.T3), the distribution of all 40 task-specific losses reveals that our approach yields more concentrated and consistently lower values.

Except for the loss discrepancy, we randomly select 8 out of 40 tasks and have checked the gradient cosine similarity among tasks on the CelebA dataset. [Figure˜7](https://arxiv.org/html/2502.08585v3#A1.F7) (right) illustrates the cosine similarities of task gradients after the 15th epoch, which shows that the gradient conflict is mitigated.

Figure: Figure 7: Left: comparison of loss distributions for LDC-MTL and LS on the CelebA dataset. Right: cosine similarities of task gradients for LS and LDC-MTL, respectively; LDC-MTL exhibits much higher gradient similarity among tasks, suggesting reduced gradient conflict.
Refer to caption: x7.png

### A.7 Comparison with weight-swept LS

We have operated a weight sweep on linear scalarization, and the result is shown in [Table˜9](https://arxiv.org/html/2502.08585v3#A1.T9). From this table, we find that even after a careful weight sweep, LS does not perform better
than our method. Moreover, we have also provided a scatter plot in [Figure˜5](https://arxiv.org/html/2502.08585v3#S6.F5) (left) using results in [Table˜8](https://arxiv.org/html/2502.08585v3#A1.T8) and [Table˜9](https://arxiv.org/html/2502.08585v3#A1.T9) following Figure 8 in Xin et al. (2022). In both cases, our method forms the Pareto frontier.

Table: Table 9: Additional results on Cityscapes (2-task) dataset with different weights of LS.

### A.8 Ablation study

Task orders. To investigate the impact of loss ordering in the upper-level function $f(W,x^{*})$ in [section˜4.2](https://arxiv.org/html/2502.08585v3#S4.Ex2), we randomly shuffled the order of loss values before computing the (weighted) loss gaps on the CelebA dataset. The results are reported in [Table˜10](https://arxiv.org/html/2502.08585v3#A1.T10) where CelebA(reorder) represents random reordering on task losses before computing $f(W,x)$. Our findings indicate that reordering the losses does not lead to significant performance differences. This suggests that the effect of task ordering is minimal, likely because the loss values are already on a comparable scale.

**Table 10: Additional result on the loss orders with CelebA.**
| Dataset | $\Delta m\%$ |
| --- | --- |
| CelebA | -1.31 |
| CelebA(reorder), seed=0 | -1.30 |
| CelebA(reorder), seed=1 | -1.00 |
| CelebA(reorder), seed=2 | -1.62 |

## Appendix B Additional information

Here, we present the complete version of the double-loop algorithm for solving the penalized bilevel problem, $\mathcal{BP}_{\lambda}$ in [eq.˜2](https://arxiv.org/html/2502.08585v3#S4.E2), as detailed in Algorithm [B](https://arxiv.org/html/2502.08585v3#A2). Notably, the local or global solution of $\mathcal{BP}_{\lambda}$ obtained by Algorithm [B](https://arxiv.org/html/2502.08585v3#A2) also serves as a local or global solution to $\mathcal{BP}_{\epsilon}$, as established by Proposition 2 in Shen & Chen (2023).

## Appendix C Analysis

In the analysis, we need the following definitions.

$$ $\displaystyle x_{t}^{*}=$ $\displaystyle\arg\min_{x}g(W^{t},x),$ $\displaystyle G(x)=$ $\displaystyle[\nabla l_{1}(x),\nabla l_{2}(x),...,\nabla l_{K}(x)]$ $\displaystyle F(\theta^{t})=$ $\displaystyle f(\theta^{t})+\lambda p(\theta^{t}),\Phi(\theta^{t})=f(\theta^{t})+\lambda g(\theta^{t}),$ $\displaystyle\text{where}\;\theta^{t}=$ $\displaystyle(W^{t},x^{t}),p(\theta^{t})=g(W^{t},x^{t})-g(W^{t},x_{t}^{*})$ $\displaystyle\nabla f(W,x)=$ $\displaystyle(\nabla_{W}f(W,x),\nabla_{x}f(W,x)),\nabla g(W,x)=(\nabla_{W}g(W,x),\nabla_{x}g(W,x)).$ (6) $$

###### Lemma 1 .

Let $(W,x)$ be a solution to the $\mathcal{BP}_{\epsilon}$. This point is also an $\epsilon$-accurate Pareto stationarity point for $\{l_{i}(x)\}$ satisfying

$$ $\min_{w\in\mathcal{W}}\|G(x)w\|^{2}=\mathcal{O}(\epsilon).$ $$

###### Proof.

According to the definition of $\mathcal{BP}_{\epsilon}$, its solution $(W,x)$ satisfies that

$$ $\displaystyle g(W,x)-g(W,x^{*})\leq\epsilon.$ (7) $$

Further, according to Assumption [1](https://arxiv.org/html/2502.08585v3#Thmassumption1), we can obtain

$$ $\displaystyle g(W,x)\geq g(W,x^{*})+\nabla_{x}g(W,x^{*})(x-x^{*})+\frac{1}{2L_{g}}\|\nabla_{x}g(W,x)-\nabla_{x}g(W,x^{*})\|^{2}.$ $$

Since $x^{*}\in\arg\min_{x}g(W,x)$ and $g(W,x)=\sum_{i=1}^{K}\sigma_{i}(W){l}_{i}(x)$, we have $\nabla_{x}g(W,x^{*})=0$ and $\nabla_{x}g(W,x)=\sum_{i=1}^{K}\sigma_{i}(W)\nabla_{x}{l}_{i}(x)={G}(x)\sigma(W)$ . We can obtain,

$$ $\displaystyle\|{G}(x)\sigma(W)\|^{2}\leq 2L_{g}(g(W,x)-g(W,x^{*}))=\mathcal{O}(\epsilon),$ (8) $$

where the last inequality follows from [eq.˜7](https://arxiv.org/html/2502.08585v3#A3.E7). Furthermore, since we have used softmax at the last layer of our neural network, $\sigma(W)$ belongs to the probability simplex $\mathcal{W}$. Thus, we can derive

$$ $\min_{w\in\mathcal{W}}\|{G}(x)w\|^{2}\leq\|{G}(x)\sigma(W)\|^{2}=\mathcal{O}(\epsilon).$ $$

Thus, the solution $(W,x)$ to the $\mathcal{BP}_{\epsilon}$ also satisfies Pareto stationarity of the
loss functions $\{l_{i}\}$

∎

###### Theorem 2 (Restatement of Theorem ˜ 1 ) .

Suppose Assumptions [1](https://arxiv.org/html/2502.08585v3#Thmassumption1)-[2](https://arxiv.org/html/2502.08585v3#Thmassumption2) are satisfied. Select hyperparameters

$$ $\displaystyle\alpha\in(0,\frac{1}{L_{f}+\lambda(2L_{g}+L_{g}^{2}\mu)}],\;\beta\in(0,\frac{1}{L_{g}}],\lambda=L\sqrt{3\mu{\epsilon}^{-1}},\text{ and }N=\Omega(\log(\alpha t)).$ $$

(i) Our method with the updates [eq.˜3](https://arxiv.org/html/2502.08585v3#S4.E3) and [eq.˜4](https://arxiv.org/html/2502.08585v3#S4.E4) (i.e., Algorithm [B](https://arxiv.org/html/2502.08585v3#A2) in the appendix) finds an $\epsilon$-accurate stationary point of the problem $\mathcal{BP}_{\lambda}$. If this stationary point is a local/global solution to $\mathcal{BP}_{\lambda}$, it is also a local/global solution to $\mathcal{BP}_{\epsilon}$. Furthermore, it is also an $\epsilon$-accurate Pareto
stationary point for loss functions $l_{i}(x),i=1,...,K.$

(ii) Moreover, if $\|\nabla_{W}g(W^{t},z_{N}^{t})\|=\mathcal{O}(\epsilon)$ for $t=1,...,T$. The simplified method in Algorithm [4.3](https://arxiv.org/html/2502.08585v3#S4.SS3) also achieves the same convergence guarantee as that in $(i)$.

###### Proof.

We start with the first half of our theorem. Directly from Theorem 3 in Shen & Chen (2023), Algorithm [B](https://arxiv.org/html/2502.08585v3#A2) achieves an $\epsilon$-accurate stationary point of $\mathcal{BP}_{\lambda}$ with $\widetilde{\mathcal{O}}(\epsilon^{-1.5})$ iterations such that

$$ $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\|\nabla f(W^{t},x^{t})+\lambda(\nabla g(W^{t},x^{t})-\nabla g(W^{t},x_{t}^{*}))\|^{2}\leq\frac{F(W^{0},x^{0})}{\alpha T}+\frac{10L^{2}L_{g}^{2}}{T}=\mathcal{O}(\epsilon).$ $$

Recall that $F(W^{0},x^{0})=f(W^{0},x^{0})+\lambda(g(W^{0},x^{0})-g(W^{0},x_{0}^{*}))$. According to the Proposition 2 in Shen & Chen (2023) by setting $\delta=\epsilon$ therein, we can have $g(W^{T},x^{T})-g(W^{T},x_{T}^{*})\leq\epsilon$ if this stationary point is local/global solution to $\mathcal{BP}_{\lambda}$. Then by using [Lemma˜1](https://arxiv.org/html/2502.08585v3#Thmlemma1), we know that this $\epsilon$-accurate stationary point is also an $\epsilon$-accurate Pareto stationary point of loss functions $\{{l}_{i}(x)\}$ satisfying

$$ $\min_{w\in\mathcal{W}}\|G(x^{T})w\|^{2}=\mathcal{O}(\epsilon).$ $$

The proof of the first half of our theorem is complete.

Then, for the second half, since we have built the connection between the stationarity of $\mathcal{BP}_{\lambda}$ and Pareto stationarity, we prove that the single-loop Algorithm [4.3](https://arxiv.org/html/2502.08585v3#S4.SS3) achieves an $\epsilon$-accurate stationary point of $\mathcal{BP}_{\lambda}$. Recall that

$$ $\displaystyle\|\nabla f(W^{t},$ $\displaystyle x^{t})+\lambda(\nabla g(W^{t},x^{t})-\nabla g(W^{t},x_{t}^{*}))\|^{2}$ $\displaystyle\overset{(i)}{\leq}$ $\displaystyle 2\|\nabla f(W^{t},x^{t})+\lambda\nabla g(W^{t},x^{t})\|^{2}+2\lambda^{2}\|\nabla g(W^{t},x_{t}^{*})\|^{2}$ $\displaystyle\overset{(ii)}{=}$ $\displaystyle 2\|\nabla f(W^{t},x^{t})+\lambda\nabla g(W^{t},x^{t})\|^{2}+2\lambda^{2}\|\nabla_{W}g(W^{t},x_{t}^{*})\|^{2}$ $\displaystyle\overset{(iii)}{\leq}$ $\displaystyle 2\|\nabla f(W^{t},x^{t})+\lambda\nabla g(W^{t},x^{t})\|^{2}+4\lambda^{2}\|\nabla_{W}g(W^{t},x^{*}_{t})-\nabla_{W}g(W^{t},z^{t}_{N})\|^{2}$ $\displaystyle+4\lambda^{2}\|\nabla_{W}g(W^{t},z^{t}_{N})\|^{2},$ (9) $$

where $(i)$ and $(iii)$ both follow from Young’s inequality, and $(ii)$ follows from $\nabla_{x}g(W^{t},x_{t}^{*})=0$. Besides, recall that $z_{N}^{t}$ is the intermediate output of the subloop in Algorithm [B](https://arxiv.org/html/2502.08585v3#A2). We next provide the upper bounds of the above three terms on the right-hand side (RHS). For the first term, we utilize the smoothness of $\Phi(\theta^{t})=\nabla f(\theta^{t})+\lambda\nabla g(\theta^{t})$ where $L_{\Phi}=L_{f}+\lambda L_{g}$ and $\theta^{t}=(W^{t},x^{t})$.

$$ $\displaystyle\Phi(\theta^{t+1})\leq$ $\displaystyle\Phi(\theta^{t})+\langle\nabla\Phi(\theta^{t}),\theta^{t+1}-\theta^{t}\rangle+\frac{L_{\Phi}}{2}\|\theta^{t+1}-\theta^{t}\|^{2}$ $\displaystyle\overset{(i)}{\leq}$ $\displaystyle\Phi(\theta^{t})-\frac{\alpha}{2}\|\nabla\Phi(\theta^{t})\|^{2},$ $$

where $(i)$ follows from $\alpha\leq\frac{1}{L_{\Phi}}=\mathcal{O}(\lambda^{-1})$. Thus, we can obtain

$$ $\displaystyle\|\nabla\Phi(\theta^{t})\|^{2}\leq\frac{2}{\alpha}(\Phi(\theta^{t})-\Phi(\theta^{t+1})).$ (10) $$

Then for the second term on the RHS in [appendix˜C](https://arxiv.org/html/2502.08585v3#A3.Ex21), we follow the same step in the proof of Theorem 3 in Shen & Chen (2023) and obtain

$$ $\displaystyle 4\lambda^{2}\|\nabla_{W}$ $\displaystyle g(W^{t},x^{*}_{t})-\nabla_{W}g(W^{t},z_{N}^{t})\|^{2}$ $\displaystyle\leq$ $\displaystyle 4\lambda^{2}L_{g}^{2}\mu\Big(1-\frac{\beta}{2\mu}\Big)^{N}(g(W^{t},x^{t})-g(W^{t},x^{*}_{t}))$ $\displaystyle\overset{(i)}{\leq}$ $\displaystyle 4\lambda^{2}L_{g}^{2}\Big(1-\frac{\beta}{2\mu}\Big)^{N}\|\nabla_{x}g(W^{t},x^{t})\|^{2}$ $\displaystyle=$ $\displaystyle 4\lambda^{2}L_{g}^{2}\Big(1-\frac{\beta}{2\mu}\Big)^{N}\Big\|\frac{x^{t+1}-x^{t}+\alpha\nabla_{x}f(W^{t},x^{t})}{\alpha\lambda}\Big\|^{2}$ $\displaystyle\overset{(ii)}{\leq}$ $\displaystyle 8\lambda^{2}L_{g}^{2}\Big(1-\frac{\beta}{2\mu}\Big)^{N}\Big(\frac{\|\theta^{t+1}-\theta^{t}\|^{2}}{\alpha^{2}\lambda^{2}}+\frac{L^{2}}{\lambda^{2}}\Big)$ $\displaystyle\overset{(iii)}{\leq}$ $\displaystyle\frac{1}{2\alpha^{2}}\|\theta^{t+1}-\theta^{t}\|^{2}+\frac{2L^{2}L_{g}^{2}}{\alpha^{2}t^{2}}$ $\displaystyle=$ $\displaystyle\frac{1}{2}\|\nabla\Phi(\theta^{t})\|^{2}+\frac{2L^{2}L_{g}^{2}}{\alpha^{2}t^{2}},$ (11) $$

where $(i)$ follows from the PL condition, $(ii)$ follows from Young’s inequality and Assumption [1](https://arxiv.org/html/2502.08585v3#Thmassumption1), and $(iii)$ follows from the selection on $N\geq\max\{-\log_{c_{\beta}}(16L_{g}^{2}),-2\log_{c_{\beta}}(2\alpha t)\}$ with $c_{\beta}=1-\frac{\beta}{2\mu}$. Lastly, for the last term at the RHS in [appendix˜C](https://arxiv.org/html/2502.08585v3#A3.Ex21), we have,

$$ $\displaystyle 4\lambda^{2}\|\nabla_{W}g(W^{t},z^{t}_{N})\|^{2}=\mathcal{O}(\lambda^{2}\epsilon^{2}),$ (12) $$

where this inequality follows from our experimental observation. Furthermore, substituting [eq.˜10](https://arxiv.org/html/2502.08585v3#A3.E10), and [appendix˜C](https://arxiv.org/html/2502.08585v3#A3.Ex27) into [appendix˜C](https://arxiv.org/html/2502.08585v3#A3.Ex21) yields

$$ $\displaystyle\|\nabla f(W^{t},x^{t})$ $\displaystyle+\lambda(\nabla g(W^{t},x^{t})-\nabla g(W^{t},x^{*}_{t}))\|^{2}$ $\displaystyle\leq$ $\displaystyle\frac{5}{2}\|\nabla\Phi(\theta^{t})\|^{2}+\frac{2L^{2}L_{g}^{2}}{\alpha^{2}t^{2}}+4\lambda^{2}\|\nabla_{W}g(W^{t},z^{t}_{N})\|^{2}$ $\displaystyle\leq$ $\displaystyle\frac{5}{\alpha}(\Phi(\theta^{t})-\Phi(\theta^{t+1}))+\frac{2L^{2}L_{g}^{2}}{\alpha^{2}t^{2}}+4\lambda^{2}\|\nabla_{W}g(W^{t},z^{t}_{N})\|^{2}.$ (13) $$

Therefore, telescoping the above inequality yields,

$$ $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}$ $\displaystyle\|\nabla f(W^{t},x^{t})+\lambda(\nabla g(W^{t},x^{t})-\nabla g(W^{t},x^{*}_{t}))\|^{2}$ $\displaystyle=$ $\displaystyle\mathcal{O}(\frac{\lambda}{\alpha T}+\frac{1}{\alpha^{2}T}+\lambda^{2}\epsilon^{2}).$ $$

According to the parameter selection that $\lambda=\mathcal{O}(\epsilon^{-\frac{1}{2}})$, $\alpha=\mathcal{O}(\epsilon^{\frac{1}{2}})$, and $T=\mathcal{O}(\epsilon^{-2})$, we can obtain

$$ $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\|\nabla f(W^{t},x^{t})+\lambda(\nabla g(W^{t},x^{t})-\nabla g(W^{t},x^{*}_{t}))\|^{2}=\mathcal{O}(\epsilon).$ $$

Therefore, Algorithm [4.3](https://arxiv.org/html/2502.08585v3#S4.SS3) can achieve a stationary point of $\mathcal{BP}_{\lambda}$ with $\mathcal{O}(\epsilon^{-2})$ iterations. If this stationary point is a local/global solution to $\mathcal{BP}_{\lambda}$, it is also a solution to $\mathcal{BP}_{\epsilon}$ according to Proposition 2 in Shen & Chen (2023). Then, by using [Lemma˜1](https://arxiv.org/html/2502.08585v3#Thmlemma1), we know this stationary point is also an $\epsilon$-accurate Pareto stationary point of the original loss functions. The proof is complete.
∎

###### Lemma 1 .

###### Proof.

###### Theorem 2 (Restatement of Theorem ˜ 1 ) .

###### Proof.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Could you provide a performance comparison between Algorithm $1$ and Algorithm $2$? Additionally, what do you think of a hybrid strategy that runs Algorithm $2$ until $||\nabla_W g(W^t, z^t_N)||$ drops below a threshold relative to $||\nabla_W g(W^t, x^t)||$, and then switches to Algorithm $1$ for the remainder of training?

2. Could you include a comparison with ConsMTL $[1]$, a recent state-of-the-art gradient-manipulation method that formulates MTL as a bi-level optimization over shared and task-specific parameters? This would strengthen the claim that LDC-MTL consistently outperforms gradient-based baselines.

3. In Appendix A-$7$, it is stated that the scatter plots in Figures $5 (a)$ and $5 (b)$ follow the same setting as Figure $8$ in $[2]$. Does this imply that the axis labels in Figures $5 (a–b)$ might have been mislabeled? Additionally, could you include FAMO and GO4Align (each trained with distinct random seeds) in this Pareto front analysis for a more comprehensive comparison?

$[1]$ Qin, Xiaohan, Xiaoxing Wang, and Junchi Yan. "Towards Consistent Multi-Task Learning: Unlocking the Potential of Task-Specific Parameters." Proceedings of the Computer Vision and Pattern Recognition Conference. 2025.

$[2]$ Xin, Derrick, et al. "Do current multi-task optimization methods in deep learning even help?." Advances in neural information processing systems 35 (2022): 13597-13609.
soundness: 3
strengths: The paper's main strengths lie in its clear bilevel formulation for controlling task-wise loss discrepancies, which provides a principled alternative to existing loss balancing methods. The proposed approach is both scalable and lightweight, featuring a single-loop update that avoids the typical $\mathcal{O}(K)$ gradient storage overhead found in many multi-task algorithms. Theoretical analysis further reinforces the work by providing convergence guarantees that ensure Pareto stationarity under mild assumptions. Empirically, the paper presents comprehensive experiments across four benchmark datasets and compares against competitive baselines, demonstrating consistent improvements. Moreover, the additional analyses on gradient conflict reduction and loss variance offer valuable insights into the behavior and interpretability of the proposed method.
summary: The paper proposes LDC-MTL, a scalable loss discrepancy control method for multi-task learning (MTL). It formulates MTL as a bilevel optimization problem, where the lower level minimizes a weighted sum of task losses and the upper level minimizes the discrepancies among those losses. A first-order, single-loop approximation is introduced, leading to an $\mathcal{O}(1)$ time and memory algorithm that scales with the number of tasks. Theoretical analysis establishes convergence both to a stationary point of the bilevel problem and to an $\epsilon$-accurate Pareto-stationary point. Experiments on CelebA ($40$ tasks), QM9 ($11$ tasks), NYU-v$2$ ($3$ tasks), and Cityscapes ($2$ tasks) show competitive or superior $\Delta m\\%$ compared to scalarization and gradient-manipulation baselines, with runtime close to linear scalarization (LS).
weaknesses: 1. The claimed $O(1)$ efficiency critically depends on the empirical observation that $||\nabla_W g(W^t,z^t_N)||$ remains small. While this is illustrated for the Cityscapes dataset $(K = 2)$, providing similar empirical evidence for other datasets, particularly CelebA $(K = 40)$, would strengthen the generality of this claim.

2. The experimental analyses on loss discrepancy and gradient conflict (Table $3$, Figures $6$ and $7$) are conducted only against linear scalarization. For a fairer assessment, comparisons should also include the strongest loss balancing baselines, such as GO4Align, which achieves top performance in the main experiments.

---

confidence: 3
contribution: 3
initial_rating: 4
presentation: 3
questions: See the weaknesses part.
soundness: 3
strengths: 1.  The proposed method avoids solving the complex bi-level structure in BLO. It also achieves good performance in MTL. 
2. The paper structure is clear and easy to follow.
3. A theoretical proof to support the effectiveness of the proposed method.
summary: This work proposed a first-order bi-level optimization approach to solve the multi-task learning problem. The proposed method mainly contains two components: (i) a bilevel formulation for fine-grained loss discrepancy control, and (ii) a scalable first-order bilevel algorithm that requires only constant order time and memory.
weaknesses: 1. Lacking $\mathcal{O}(1)$ baselines such as "Smooth Tchebycheff Scalarization for Multi-Objective Optimization, ICML 2024."

2. Using bi-level optimization to solve the MTL problem is not new. The "related work" part does not mention those works, such as "Multi-objective meta-learning, AIJ" and "A first-order multi-gradient algorithm for multi-objective bi-level optimization, ECAI 2024".

3. The proposed approach is quite like a penalty-based bi-level optimization approach, such as "On Penalty-based Bilevel Gradient Descent Method". Though the author cites some papers  (Kwon et al., 2023; Yang et al., 2023), but does not discuss the differences and how similar these approaches are. This is important for evaluating the theoretical contributions of this paper. In my opinion, the theoretical result does not show an advantage.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: Refer to weaknesses.
soundness: 2
strengths: 1. The bilevel formulation provides a clear and principled avenue for directly controlling loss discrepancies while aiming for balanced performance across tasks (Section 4).

2. The authors conducted extensive experiments to verify the effectiveness of the proposed LDC-MTL.

3. The authors provide a thorough theoretical analysis, showing that the algorithm achieves $\epsilon$-Pareto stationarity under standard Lipschitz and PL conditions.
summary: This paper introduces LDC-MTL, a scalable multi-task learning method that tackles the persistent issue of loss discrepancy across task objectives. Motivated by the difficulty of balancing tasks whose losses operate on different scales or progress at different speeds, LDC-MTL formulates the learning challenge as a bilevel optimization problem: the lower-level optimizes a weighted sum of task losses, and the upper-level adjusts the weights to minimize their discrepancies. The authors propose a single-loop, first-order algorithm that claims $\mathcal{O}(1)$ memory and time overhead per iteration, compared to $\mathcal{O}(K)$ for previous gradient manipulation techniques. Experiments across standard MTL benchmarks demonstrate the efficiency and competitive accuracy of LDC-MTL.
weaknesses: 1. The authors claim that LDC-MTL achieves $\mathcal{O}(1)$ memory and time overhead per iteration. However, this property has already been stated in prior work such as FAMO. In fact, all loss-based MTL methods (e.g., UW, DWA, and more recently GO4Align) inherently possess this characteristic, as they aggregate task losses directly rather than computing gradients for each task separately—thus requiring only a single backward pass per iteration. Therefore, this advantage cannot be regarded as a distinctive contribution of the proposed method.

2. The idea of employing bilevel optimization for multi-task learning is not novel. Recent works, including GO4Align and ConsMTL [2], have already adopted similar formulations, making the conceptual contribution of this paper relatively incremental.

3. In Figure 1, the authors present the performance of LDC-MTL and other methods on a toy example. However, since there is no dominating solution along the Pareto front, convergence to the same Pareto stationary point cannot be considered a meaningful advantage in the absence of any prior preference among tasks.

4. The proposed method involves several hyperparameters, such as the learning rate $\alpha$ and the penalty coefficient $\lambda$. The authors determine these through grid search, which contradicts the claimed advantage of computational efficiency, as it implies multiple runs are needed to find suitable values. More thorough ablation studies and analyses are necessary to demonstrate the robustness of the method and to provide practical guidance for hyperparameter selection.

5. Confusing supplementary material. The authors appear to provide code in the supplementary material, yet upon inspection, there seems to be no actual implementation corresponding to the proposed LDC-MTL method. This raises concerns about reproducibility and transparency; the authors should ensure that the provided materials are complete and directly related to the presented work.

6. Missing references. The paper overlooks several key related works [1,2,3], some of which report superior empirical performance (though this does not affect my evaluation of the paper). In the MTL field, there are no universally accepted quantitative metrics for superiority—measures such as $\Delta m\\%$ and MR are only proxy indicators, and it is well known that $\Delta m\\%$ tends to favor baselines with lower initial performance. Therefore, I do not insist that the authors include these methods in their experiments, but they should at least acknowledge and differentiate them properly in the Related Work section.

[1] Independent component alignment for multi-task learning. CVPR 2023

[2] Towards Consistent Multi-Task Learning: Unlocking the Potential of Task-Specific Parameters. CVPR 2025

[5] Rep-MTL: Unleashing the Power of Representation-level Task Saliency for Multi-Task Learning. ICCV 2025

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to the Weaknesses.
soundness: 3
strengths: 1. Employing a bi-level paradigm for MTL becomes popular recently, thus this paper is timely.
2. Developing efficient MTL algorithms is crucial in the era of large foundation models, and this paper addresses a timely and important issue.
3. Extensive evaluation demonstrates excellent performance across mainstream MTL datasets.
summary: This paper investigates MTL from the perspective of loss balancing. Specifically, it proposes a bi-level loss balancing paradigm to dynamically adjust the task weights and promote the balanced optimization. Besides, it also proposes a first-order alternative for efficiency purposes. Extensive experiments have demonstrated the proposed method's ability to achieve competitive performance across multiple mainstream MTL datasets.
weaknesses: 1. It remains unclear how Eqn. (1) facilitates balanced task learning speeds, especially given that $\tau$ is fixed at 1 in most experimental settings.
2. The paper lacks sufficient details regarding the router model, which plays a critical role in the proposed framework.
3. The proposed method adopts a different training schedule (e.g., learning rate, batch size) compared to prior MTL approaches [1][2]; please provide a discussion on GPU memory usage under this setting.
4. It is recommended to report another widely adopted MTL metric, Mean Rank (MR), for a more comprehensive evaluation [1][2].
5. Please specify which version of LDC-MTL was used for evaluation.

Reference:

[1] Fair Resource Allocation in Multi-Task Learning. ICML 2024.

[2] FAMO: fast adaptive multitask optimization. NeurIPS 2023.


Decision: Reject