Monotonic Transformation Invariant Multi-task Learning

TL;DR

This work addresses the instability caused by arbitrary monotonic scaling of task losses in multi-task learning. It introduces DiBS-MTL, a monotonic-transformation-invariant adaptation of the Direction-based Bargaining Solution that operates on normalized task gradients to achieve Pareto-stationary updates even in nonconvex settings. The authors prove subsequential convergence to a Pareto stationary point under standard assumptions and demonstrate through extensive experiments that DiBS-MTL outperforms state-of-the-art baselines when task losses are poorly scaled, while remaining competitive on standard benchmarks. The approach offers a principled, efficient alternative for robust, multi-task optimization in heterogeneous loss landscapes.

Abstract

Multi-task learning (MTL) algorithms typically rely on schemes that combine different task losses or their gradients through weighted averaging. These methods aim to find Pareto stationary points by using heuristics that require access to task loss values, gradients, or both. In doing so, a central challenge arises because task losses can be arbitrarily scaled relative to one another, causing certain tasks to dominate training and degrade overall performance. A recent advance in cooperative bargaining theory, the Direction-based Bargaining Solution (DiBS), yields Pareto stationary solutions immune to task domination because of its invariance to monotonic nonaffine task loss transformations. However, the convergence behavior of DiBS in nonconvex MTL settings is currently not understood. To this end, we prove that under standard assumptions, a subsequence of DiBS iterates converges to a Pareto stationary point when task losses are nonconvex, and propose DiBS-MTL, an adaptation of DiBS to the MTL setting which is more computationally efficient that prior bargaining-inspired MTL approaches. Finally, we empirically show that DiBS-MTL is competitive with leading MTL methods on standard benchmarks, and it drastically outperforms state-of-the-art baselines in multiple examples with poorly-scaled task losses, highlighting the importance of invariance to nonaffine monotonic transformations of the loss landscape. Code available at https://github.com/suryakmurthy/dibs-mtl

Figures (5)

• Figure 1: Loss functions in the nonconvex example (\ref{['sec:toy']}). We employ the nonaffine transformation $h(\mathcal{L}_1) = \mathop{\mathrm{sign}}\nolimits(\mathcal{L}_1)\cdot \mathcal{L}_1^4$.
• Figure 2: DiBS-MTL successfully avoids task domination due to its invariance to monotonic nonaffine transformations, while baseline MTL methods visibly favor task 1 when $\mathcal{L}_1$ is transformed, and some baseline methods even fail to reach the Pareto front. $\bullet$ and denote initializations and the Pareto front respectively. To better visualize results for the transformed case, we retain the original $\mathcal{L}_1$ axis. Results for multi-step DiBS-MTL and additional baselines are included in \ref{['fig:extra_trajectories']}.
• Figure 3: Average task success rate on the MT10 benchmark, averaged over all tasks and 10 random seeds. Transparent curves show per-episode success rates, while solid curves denote rolling averages computed over a window of 51 episodes. 1-step DiBS-MTL demonstrates strong robustness to diverse reward transformations and consistently and significantly outperforms all baselines across all transformed settings.
• Figure 4: Runtime comparisons across nonconvex demonstrative example and multi-task RL settings.
• Figure 5: Additional experiments performed in the demonstrative two-objective example. We observe that, \ref{['eq: dibs mtl']} exhibits the same invariance to monotone non-affine transforms as T-stepDiBS-MTL from \ref{['fig:trajectories']}, and also demonstrates similar behaviour, tracing very similar trajectories as T-stepDiBS-MTL ($T=5$) from the same initializations.

Theorems & Definitions (4)

• Definition 1: Pareto stationarity
• Remark 1: Invariance of DiBS
• Theorem 1
• proof