FAMO: Fast Adaptive Multitask Optimization
Authors: Bo Liu
Yihao Feng
Peter Stone
qiang liu
Venue: NeurIPS 2023 poster (2023)
Keywords: multitask learning
multitask optimization
conflicting gradients
knowledge transfer
TL;DR: We introduce a fast adaptive task weighting method for multitask optimization.

Abstract: One of the grand enduring goals of AI is to create generalist agents that can learn multiple different tasks from diverse data via multitask learning (MTL). However, in practice, applying gradient descent (GD) on the average loss across all tasks may yield poor multitask performance due to severe under-optimization of certain tasks. Previous approaches that manipulate task gradients for a more balanced loss decrease require storing and computing all task gradients ($\mathcal{O}(k)$ space and time where $k$ is the number of tasks), limiting their use in large-scale scenarios. In this work, we introduce Fast Adaptive Multitask Optimization (FAMO), a dynamic weighting method that decreases task losses in a balanced way using $\mathcal{O}(1)$ space and time. We conduct an extensive set of experiments covering multi-task supervised and reinforcement learning problems. Our results indicate that FAMO achieves comparable or superior performance to state-of-the-art gradient manipulation techniques while offering significant improvements in space and computational efficiency. Code is available at \url{https://github.com/Cranial-XIX/FAMO}.

FAMO: Fast Adaptive Multitask Optimization
"o Liu, & YYhho Feng, Peter Stone, "Qiang Liu "The University of Texas at Austin, *Salesforce Al Research, $Sony Al (bliu, pstone, lqiang)ecs. utexas. edu, yihaof saoof@sallsforceccm
Abstract
One of the grand enduring goals of Al is to create generalist agents that can learn multiple different tasks from diverse data via multitask learning (MTL). However, in practice, applying gradient descent (GD) on the average loss across all tasks may yield poor multitask performance due to severe under-optimization of certain tasks. Previous approaches that manipulate task gradients for more balanced loss decrease require storing and computing all task gradients (O(k) space and time where k is the number of tasks), limiting their use in large-scale scenarios. In this work, we introduce Fast Adaptive Multitask Optimization (FAMO), dy- namic weighting method that decreases task losses in a balanced way using O(1) space and time. We conduct an extensive set experiments covering multi task supervised and reinforcement learning problems. Our results indicate that FAMO achieves comparable or superior performance to state-of-the-ar gradient manipula- tion techniques while offering significant improvements in space and computational efficiency. Code is available at hhttps://// /tte::///ttta gighhu..b.... www.tttuu..ooreeeeeeeeeeeeeeeeee Coom/Craiilll -XIX/AAMO
Introduction
Large models trained on diverse data have advanced both computer vision 20 and natural language processing 4. paving the way for generalist agents capable of multitask learning (MTL) [5] Given crucial to design MTL methods that are effective in terms of the substantial size of these models, task performance and efficient in terms of space and time complexities for managing training costs and environmental impacts. This work explores such methods through the lens of optimization.
Perhaps the most intuitive way of solving an MTL problem is to optimize the average loss across all tasks. However, in practice, doing so can lead models with poor multitask performance: subset of tasks are severely under-optimized. A major reason behind such optimization failure is that subset of tasks are under-optimized because the average gradient constantly results in small (or even negative) progress on these tasks (see details in Section
To mitigate this problem, gradient manipulation methods 4K] 25
compute a new update vector in place of the gradient to the average loss, such that all task losses decrease in a more balanced way. The new update vector often determined by solving an additional optimization problem that involves all task gradients. While these approaches exhibit improved performance, they become computationally expensive when the number of tasks and the model size are large ΣMI. This is because they require computing and storing all task gradients at each iteration, thus demanding O(k) space and time complexities, not to mention the overhead introduced by solving the additional optimization problem. In contrast, the average gradient can be efficiently computed in (11) space and time per iteration because one can first average the task losses and then take the gradient of the average loss To this end, we ask the following question:
'Here, refer to the situation where single data be used compute all task losses.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Avg Loss Landscape
L contour
L contour
Time Comsumption (sec.)
10
995
33
39
-10
-10
-10
0
10
-10
10
Adam
MGDA
PCOmd
CAGrau
NI
FAMO
Parcto Front
Initial Point
Adam
MGDA
PCGrad
CAGrad
NashMTL
FAMO
-10
L
L
2
L
L
L
2.0-15-10-03000
Figure 1: Top left: The loss landscape, and individual task losses of toy 2-task learning problem represents the minimum task losses). Top right: the runtime different MTI. methods 50000 steps. Bottom: the loss trajectories of different MTL methods. ADAM fails in out of runs to reach the Pareto front due to CG. FAMO decreases task losses in balanced and is the only method matching the (1)) space/time complexity of ADAM. Experimental details and analysis provided in Section 5.1
(Q) Is it possible design multi-task learning optimizer that ensures balanced reduction in losses across all tasks while utilizing ((1) space and time per iteration?
In this work, we present Fast Adaptive Multitask Optimization (FAMO), simple yet effective adaptive task weighting method to address the above question. On the one hand, FAMO is designed to ensure that all tasks are optimized with approximately similar progress. the other hand, FAMO leverages the loss history to update the task weighting, hence bypassing the necessity of computing all task gradients. To summarize, our contributions are:
1. We introduce FAMO, an MTL optimizer that decreases task losses approximately at equal rates while using only ((1) space and time per iteration.
2. We demonstrate that FAMO performs comparably ulation methods on a wide range of standard MTL benchmarks, in terms of standard MTL metries, while being significantly computationally cheaper.
better than existing gradient manip-
2 Background
In this section, we provide the formal definition of multitask learning, then discuss its optimization challenge, and provide brief overview of the gradient manipulation methods.
Multitask Learning (MTL) MTL considers optimizing single model with parameter 0 R.' that can perform & 2 2 tasks well, where each task is associated with a loss function (,(0) TR" - Rgo Then, is common optimize the average loss across all tasks:
min Ockm
(o(0)
1 1:
(1)
Optimization Challenge Directly optimizing (T) can result in severe under-optimization a subset of tasks. A major reason behind this optimization challenge is the "generalized" conflicting gradient phenomenon, which we explain in the following. At any time step I, assume one updates the model
In this work, we assume V i, (0) 2 O, which is true for typical loss functions including mean square and cross-entropy losses. Note that one can always transform E; to be non-negative if loss lower bound is known.

---

Algorithm Fast Adaptive Multitask Optimization (FAMO)
Input: Initial parameter d. task losses [ (ensure that & 2 for instance, by l
4--6 4 + c, E = inf, (()),, learning rate a and 3, and decay (= 0.001 by default).
2: 6
initialize the task logits to all zeros
3: fort=1 1 T do
4: Compute 2, Softmax (t). e.g.
exp(Ei.t)
2 Σ exp(Ex,t)
5: Update the model parameters:
00+1 = 00 -a Σ
Zi, Vlan where as ( Zi,t (ct
lip
1=1 fie
i=1
6: Update the logits for task weighting:
V (E))
log(1, : log Cu.e - log Cx,e+1
loog l1,-11
h+1 = E - 360 7E.) where
V (E)))
7: end for
parameter using gradient descent style iterative update: 0.+1 02 -ad, where a is the step size and de is the update at time Then, we say that conflicting gradients (CG) 24 43) happens if
3i, (.(0..1) - (,(0)) x -VV 0.
In other words, certain task's loss is increasing. CG often occurs during optimization and is not inherently detrimental. However, it becomes undesirable when subset of tasks persistently undergoes under-opting optdmi-otimizziionn due to CG. In a more general sense, it is not desirable a subset of tasks has much slower learning progress compared to the rest of the tasks (even if all task losses are decreasing). This very phenomenon, which we call the "generalized" conflicting gradient, has spurred previous research mitigate each optimization stage
4
Gradient Manipulation Methods Gradient manipulation methods aim to decrease all task losses more balanced way by finding new update de at each step. de is usually a convex combination of task gradients, and therefore the name gradient manipulation (denote Vl,, (ol, for short):
VET
W1.t : Wk.t
de
Wes
where Wt
f(VeI........ Vla.)) Su.
7eT
k.t
Here, Sy {w & wT 1} is the probabilistic simplex, and w. is the task weighting across all tasks. Please refer to Appendix [ppendix[AAfor for details of five state-of-the-ee gradient manipulation methods (MGDA, PCGRAD, CAGRAD. IMTL-G, NASHMTL) and their corresponding f. Note that existing gradient manipulation methods require computing and storing & task gradients before applying f to compute de, which often involves solving an additional optimization problem. As a result, we say these methods require at least (k) space and time complexity, which makes them slow and memory inefficient when k and model size m are large.
3 Fast Adaptive Multitask Optimization (FAMO)
In this section, we introduce FAMO that addresses question Q, which involves two main ideas:
1. At each step, decrease all task losses at an equal rate as much as possible (Section 3.1
2. Amortize the computation in 1. over time (Section 3.2

---

3.1 Balanced Rate of Loss Improvement
At time 4, assume we perform the update Oc+1 0, - add, we define the rate of improvement for task
i as
l - -1,-11
r: (a de)
lift
FAMO then seeks an update de that results in the largest worst-case improvement rate across all tasks 2 dell is subtracted prevent an under-specified optimization problem where the objective can be
infinitely large):
max min (a d:) - Ide1l de [c[A]
When the step size is small, using Taylor approximation, the problem (4) can be approximated by
Vlidd 1 Ilddel (v lx.)
max min deER" [ b,2
2 d
(5)
2
Instead of solving the primal problem in (5) where d e R (m can be millions if 0 is the parameter of a neural network), we consider its dual problem:
Proposition 3.1. The dual objective of 5
v log L.
1
z arg min ze5g 2
- |i2|?
where Je
(6)
V Vlog Cx
where zi [2t,]] is the optimal combination weights of the gradients, and the optimal update direction is di Jez
Proof.
log lie) 'dd) d 1-]
max min delt id[R]
v log Ci,1) dd 1 ld]
max min de 28
Fl
b
min max
Σ
v log Ci,e) |dl
(strong duality)
2682 deRm.
Write g(d,2)= (E) vV log (.) ) d 1 ||| then by setting
0g
d' =
Z log
ad
Plugging in back, we have
1 ||d]
1
min (V log (a.)) d-
min 258
min
Σ
2
log
l
min 268
||
deR sc[k]
At the optimum, we have di Jez
The dual problem 6 can be viewed as optimizing the log objective of the multiple gradient descent algorithm (MGDA) 9 35 Similar to MGDA, (6) only involves decision variable of dimension < m. Furthermore, if the optimal combination weights 2 is an interior point of Sk, then the improvement rates (a, di) of the different tasks i equal, as we show in the following result.
'To avoid division by zero, in practice we add small constant (e.g., 1e - 8) to all losses. For the ease of notation (e.g.. ((() f.(-) le 8, omit it throughout the paper.

---

Proposition 3.2. Assume {fi are smooth and the optimal weights 2i in 66 is an interior point of Sg. then
Vi+je [k],
(()) (r)
where rf (d)) aim r d))
Proof.
Consider the Lagrangian form of
6
((2,7, 1 Σ a log e
X( Σ 21 -
12-11 1) Σ Hezi, where Vi, p4 2 0.
(7)
When z² reaches the optimum, we have OL(z, A, p)/02= 0, recall that di Je2 then
V log log [
Jee -m X, where J+
J. - X).
v log k.r
When of is an interior point of Sg, we know that μ O. Hence J7 = = This means,
Vi+ j.
lim rr(a,d) log v log " lim ry (a,d;) d;). a-0 a
a-u a
3.2 Fast Approximation by Amortizing over Time
Instead of fully solving 6 each optimization step, FAMO performs single-step gradient descent on 2, which amortizes the computation over the optimization trajectory:
x = 24 azz, where 8= V=
Σ 1-1
Z v log tip
Ji Jezc.
(8)
But then, note that
log 11,1 - log 1,,
2 Jj de = Jo
(9)
log Ee.t - log bxct+1
so we can use the change in log losses to approximate the gradient.
In practice, to ensure that always stays in Sg, we re-parameterize by & and let =; Softmax (E:). where & e IR* are the unconstrained softmax logits. Consequently, have the following approxi- mate update on & from (8)
V (1)
log E,, - llog fi.
E++ E. 38, where 8
(10)
V +x
log lu. - log f
While possible to perform gradient descent on for other gradient manipulation methods in principle, we will demonstrate in Appendix B that not all such updates can be easily approximated using the change losses.
Remark:
3.3 Practical Implementation
To facilitate practical implementation, we present two modifications to the update (1O).
Re-normalization loss, e.g..
The suggested update above is a convex combination of the gradients of the log
Σ
Zz, V...logge bit Σ (
1-1 Ree
1-1
When la, is small, the multiplicative coefficient Zr,s optimization. Therefore, we propose to multiply d' by constant ct, such that e,d" can be written as a convex combination of the task gradients just as in other gradient manipulation algorithms (see and we provide the corresponding definition of w in the following):
can quite large and result in unstable
&
Zi.t
Ziet
and d, c
W. Vitt where Wi C.
(11)
E,
5

---

Regularization As we are amortizing the computation over time and the loss objective (4((-) }s are changing dynamically, it makes sense focus more on the recent updates of 5 [46]. To this end. we put decay term such that the resulting is an exponential moving average of its gradient updates:
Ex+1 & -(0 + = 6t (1 13)6,-1 (1- 1-- 8
(12)
We provide the complete FAMO algorithm in Algorithm and its pseudocode Appendix
3.4 The Continuous Limit of FAMO
One way to characterize FAMO's behavior is to understand the stationary points of the continuous- time limit of FAMO (i.e. when step sizes (a, B) shrink to zero), From Algorithm one can derive the following non-autonomous dynamical system (assuming { are all smooth):
V (1,)
[
i [ - -4 ArJi Jen
Jez
26.
where A,
(13)
(E,)).
13 reaches its stationary points (or fixed points) when (note that C₁ > 0)
[==00 =0 = = and & =0=0 =
Σ v log =.
(14)
Therefore, the minimum points of Ei=1 log (0) are all stationary points of 13
4 Related Work
In this section, we summarize existing methods that tackle learning challenges in multitask learning (MTL). The general idea of most existing works is encourage positive knowledge transfer by sharing parameters while decreasing any potential negative knowledge transfer (a.k.a, interference) during learning. There are three major ways of doing so: task grouping, designing network architectures specifically for MTL, and designing multitask optimization methods.
Task Grouping Task grouping refers to grouping K tasks into N K clusters and learning N models for each cluster. The key is estimating the amount of positive knowledge transfer incurred by
grouping certain tasks together and then identifying which tasks should be grouped (39 45 38 36
Multitask Architecture Novel neural architectures for MTL include turdl-paaameter-shariiy methods, which decompose neural network into task-specific modules and shared feature extractor 2]. and soft- wet---aaaaaassaaeeee
using manually designed heuristics 21
methods, which learn 12 27 Recent studies extend neural architecture search for MTL
which parameters to share 33 34
by learning where to branch a network to have task-specific modules 14
The most relevant approach to our method is MTL optimization via task balancing. These methods dynamically re-weight all task losses to mitigate the confticting gradient issue 40 43 The simplest form of gradient manipulation is to re-weight the task losses based on manually designed criteria 6 13 but these methods are often heuristic and lack theoretical support. Gradient manipulation methods [35] 433 25 [6] 24 32, (6 47] propose to form new update vector at each optimization by linearly combining task gradients. The local improvements across all tasks using the new update can often be explicitly analyzed, making these methods better understood in terms of convergence. However, has been observed that gradient manipulation methods are often slow in practice, which may outweigh their performance benefits [22]. By contrast, FAMO is designed to match the performance these methods while remaining efficient in terms of memory and computation. Another recent work proposes to sample random task weights at each optimization step for MTL 23], which also computationally efficient. However, we will demonstrate empirically that FAMO performs better than this method.
Multitask Optimization

---

5 Empirical Results
We conduct experiments to answer the following question:
How does FAMO perform in terms of space/time complexities and standard MTL metrics against prior MTL optimizers standard benchmarks (e.g., supervised and reinforcement MTL problems)?
In the following, we first use toy 2-task problem demonstrate how FAMO mitigates CG while being efficient. Then we show that FAMO performs comparably or even better than state-of-ther art gradient manipulation methods on standard multitask supervised and reinforcement learning benchmarks. In addition, FAMO requires significantly lower computation time when K' is large compared to other methods. Lastly, we conduct an ablation study on how robust FAMO is to y. Each subsection first details the experimental setup and then analyzes the results.
5.1 A Toy 2-Task Example
7
L
L
25
85
5
25
05
10
10
10
-10
0g
-10
0g
-10 10
02
01
10
-10
01
10
01
Figure The average loss 7,9 and the task losses T,! and T.² for the toy example.
To better understand the optimization trajectory of FAMO, we adopt the same 2D multitask op- timization problem from NASIMTL [32] visualize how FAMO balances different loss ob- jectives. The model parameter 0 = (1..0 022 E IR2 The two tasks' objectives and their sur-
face plots are provided Appendix D and Figure 2 We compare FAMO against ADAM [I9].
MGDA [35, PCGRAD 43 CAGRAD 24. and NASHMTL 32 We then pick 5 initial points Omn {.58.5, 7.5). (-8.5,5) (0. 0), (9, 9), (10, -8)} and plot the corresponding optimization trajec- tories with different methods in Figure | Note that the toy example is constructed such that naively applying ADAM on the average loss can cause the failure of optimization for task 1.
Findings: From Figure we observe that FAMO, like all other gradient manipulation methods, mitigates the CG and reaches the Pareto front for all five runs. In the meantime, FAMO performs similarly to NASIMTL and achieves balanced loss decrease even when the two task losses are improperly scaled. Finally, as shown in the top-right of the plot, FAMO behaves similarly ADAM in terms of the training time, which is 25x faster than NASHMTL.
5.2 MTL Performance
Multitask Supervised Learning. We consider four supervised benchmarks commonly used in
prior MTL research 24 27 32 33:: NYU-v2 [1] (3 tasks), CityScapes [8] (2 tasks), QM-9 ( (11
tasks), and CelebA 28 (40 tasks). Specifically, NYU-v2 is indoor scene dataset consisting of
1449 RGBD images and dense per-pixel labeling with 13 classes. The learning objectives include image segmentation, depth prediction, and surface normal prediction based any scene image. CityScapes dataset is similar to NYU-v2 but contains 5000 street-view RGBD images with per-pixel annotations. QM-9 dataset a widely used benchmark in graph neural network learning. consists of d'>1300 130K molecules represented as graphs annotated with node and edge features. We follow the same experimental setting used in NASHMTL [32,, where the learning objective is to predict 11 properties of molecules. We use 1 10K molecules from the QM9 example in PyTorch Geometric [10] 10K molecules for validation, and the rest of 10K molecules for testing. The characteristic of this dataset is that the 11 properties are at different scales, posing a challenge for task balancing in MTL. Lastly, CelebA dataset contains 200K face images of 10K different celebrities, and each face image is provided with 40 facial binary attributes. Therefore, CelebA can be viewed as 40-task MTL problem. Different from NYU-v2, CityScapes, and QM-9, the number of tasks (K) in CelebA much larger, hence posing challenge to learning efficiency.
7

---

Segmentation
Depth
Surface Normal
Method
Angle Dist 1
Within (*
MR 1 AmX
mloU
Pix Acc
Abs Err
Rel Err
Mean
Median
11.25
22.5
30
STL
38.30
63.76
0.6754
0.2780
25.01
19.21
30.14
57.20
69.15
39.29
0.5493
0.2263
28.15
23.96
22.09
61.08
8.89
5.59
SI
38.45
64.27
0.5354
0.2201
27.60
23.37
22.53
48.57
62.32
7.89
4.39
RLW
37.17
63.77
0.5759
0.2410
28.27
24.18
22.26
47.05
60.62
11.22
7.78
DWA
39.11
65.31
0.5510
0.2285
27.61
23.18
24.17
50.18
62.39
7.67
3.57
UW
36.87
63.17
0,5446
0.2260
27.04
22.61
23.54
49.05
63.65
7.44
4.05
MGDA
30.47
59.90
0.6070
0.2555
24.88
19.45
29.18
56.88
69.36
6.00
1.38
PCGRAD
38.06
64.64
0.5550
0.2325
27.41
22.80
23.86
49.83
63.14
8.00
3.97
GRADDROP
39.39
65.12
0,5455
0.2279
27.48
22.96
23.38
49,44
62.87
7.00
3.58
CAGRAD
39.79
65.49
0,5486
0.2250
26.31
21.58
25.61
52.36
65.58
4.56
0.20
IMTL-G
39.35
65.60
0.5426
0.2256
26.02
21.19
26.
53.13
66.24
3.78
-0.76
NASHMTI.
40.13
0.5261
0.2171
25.26
28.40
68.15
.4.04
FAMO
38.88
64.90
0.5474
0.2194
25.06
19.57
29.21
56.61
68.98
3.44
-4.10
Table 1: Results NYU-v2 dataset (3 tasks). Each experiment is repeated over 3 random seeds and the
reported. The best average result marked in bold. MR and A Am% are the main metrics for MTL performance.
Method
CHOMO
ELUMO
(R*)
ZPVE
Uo
U
H
G
Gy
MR
Am
MAE 1
STL
0.07
0.18
60.6
53.9
0.50
4,53
58.8
64.2
63.8
66.2
0.07
LS
0.11
0.33
73.6
89.7
5,20
14.06
143.4
144.2
144.6
140.3
0.13
6.45
177.6
SI
0.31
0.35
149.8
135.7
1.00
4.51
55.3
55.8
55.8
55,3
0.11
3.55
77.8
RLW
0.11
0.34
76.9
92.8
5.87
15.47
156.3
1.7.11
157.6
153.0
0.14
8.00
203.8
DWA
0.11
0.33
74.1
90.6
5.09
13.99
142.3
143.0
143.4
139.3
0.13
6.27
175.3
UW
0.39
0.43
166.2
155.8
4.99
66.4
66.8
66.8
66.2
0.12
4.91
108.0
MGDA
0.22
0.37
126.8
104.6
3.23
5.69
88.4
89.4
89.3
88.0
0.12
5.91
120.5
PCGRAD
0.11
0.29
75.9
88.3
3.94
9.15
116.4
116.8
117.2
114.5
0.11
4.73
125.7
CAGRAD
0.12
0.32
83.5
94.8
3.22
6.93
114.0
114.3
114.5
112.3
0.12
5.45
112.8
IMTL-G
0.14
0.29
98.3
93.9
1.75
5.70
101.4
102.4
102.0
100.1
0.10
4.36
77.2
NASHMTL
0.10
0.25
82.9
81.9
2,43
5.38
74.5
75.0
75.1
74.2
0.09
2.09
62.0
FAMO
0.15
0.30
94.0
95.2
1.63
4.95
70.82
71.2
71.2
70.3
0.10
3.27
58.5
Table 2: Results on QM-9 dataset (11 tasks). Each experiment is repeated over 3 random seeds and the mean is reported. The best average result is marked in bold. MR and Am% are the main metrics for MTL performance.
We compare FAMO against 11 MTL optimization methods and single-task learning baseline: (1) Single task learning (STL), training an independent model (Ø for each task; (2) Linear scalarization (LS) baseline that minimizes L (3) Scale-invariant (SI) baseline that minimizes Σg log L (=)) as
SI is invariant to any scalar multiplication of task losses; (4) Dynamic Weight Average (DWA) 27
heuristic for adjusting task weights based on rates of loss changes; (5) Uncertainty Weighting
(UW) [8] uses task uncertainty as a proxy to adjust task weights; (6) Random Loss Weighting
(RLW) (23 that samples task weighting whose bog-probabilities follow the normal distribution;
(7) MGDA (35] that finds the equal descent direction for each task; (8) PCGRAD [43] proposes
to project each task gradient the normal plan of that of other tasks and combining them together in the end; (9) CAGRAD
[2 optimizes the average loss while explicitly controls the minimum decrease across tasks; (10) IMTL-G (25) finds the update direction with equal projections on task gradients; (11) GRADDROP 7 that randomly dropout certain dimensions of the task gradients based on how much they conflict; (12) NASHMTL 332 formulates MTL bargaining game and finds the solution to the game that benefits all tasks. For FAMO, we choose the best hyperparameter y é (0.0001, 0.001, 0.01) based on the validation loss. Specifically, we choose ? equals 0.01 for the CityScapes dataset and 0.001 for the rest of the datasets. See Appendix (fr results with error bars.
Evaluations: We consider two metrics 32 for MTL.: Am,, the average per-task performance drop of method relative to the STL baseline denoted b: Am% Σ (- (M Mr. (...))//. /MM x 100, where Mo.k and Mm, are the STL and m's 'ss value for metric M2. 52 1 (or 0) if the Mg is higher (or lower) the better. 2) Mean Rank (MR): the average rank of each method across tasks. For instance, if method ranks first for every task, MR will be
Findings: Results on the four benchmark datasets are provided Table 1 and 3
We observe that FAMO performs consistently well across different supervised learning MTL benchmarks compared
8

---

CityScapes
CelebA
Method
Segmentation
Depth
MR
Am%
MR
Am 4
mIoU 1
Pix Ace
Abs Err
Rel Err 4
STL
74.01
93.16
0.0125
27.77
LS
70.95
91.73
0.0161
33.83
6,50
14.11
4.15
6.28
SI
70.95
91.73
0.0161
33.83
9.25
14.11
7.20
7.83
RLW
74.57
93,41
0.0158
47.79
9.25
24.38
1.46
5.22
DWA
75.24
93.52
0.0160
44.37
6.50
21.45
3.20
6.95
UW
72.02
92.85
0.0140
30.13
6.00
3.23
5.78
MGDA
68.84
91.54
0.0309
33.50
9.75
44.14
14.85
10.93
PCGRAD
75.13
93.48
0.0154
42.07
6.75
18.29
3.17
6.65
GRADDROP
75.27
93.53
0.0157
47.54
6.00
23.73
3.29
7.80
CAGRAD
75.16
93.48
0.0141
37.60
5.75
11.64
2.48
6.20
IMTL-G
75.33
93.49
0.0135
38.41
4.00
11.10
0.84
4.67
NASHMTL
75.41
93.66
0.0129
35.02
2.00
6.82
2.84
4.97
FAMO
74.54
93.29
0.0145
32.59
6.25
8.13
1.21
4.72
Table 3: Results CityScapes tasks) and CelebA (40 tasks) datasets. Fach experiment repeated over random seeds and mean reported. The best average result marked bold. MR A Am% are the main metrics for MTL performance.
other gradient manipulation methods. particular, achieves state-of-the oftthe-rrttttttt results terms of Am on the NYU-v2 and QM-9 datasets.
Multitask Reinforcement Learning. We further apply FAMO to multitask reinforcement learning (MTRL) problems MTRL often suffers more from conflicting gradients due the stochastic nature of reinforcement learning [3]. Following CAGRAD AD 24,, we apply FA MO on the MetaWorld [44 MT10 benchmark, which consists 10 robot manipulation tasks with different reward functions. Following [37]. we use Soft Actor-Critic (SAC) [I5] as the underlying RL algorithm, and compare against baseline methods including LS (SAC with shared model) 4 Soft Modularization 42) (an MTL network that routes different modules in shared model to form different policies), PC-
GRAD 4 CAGRAD and NASHMTL [32] The experimental setting and hyperparameters all match exactly with those in CAGRAD. For NASHMTL, we report the results of applying the NASHMTL update once per {1, 50, 100} iterations
The results for all methods are provided Table 5.2
LS
NashMTL (1)
(50)
(100)
Success t (mean stderr)
FAMO(7
0.0001)
(0.001)
(0.01)
Method
1.0
Avg. Success Rate
LS (lower bound)
0.49 10.07
STL (proxy for upper bound)
0.90 20.03
0.75
PCGRAD 43
0.72 40.02
0.5
SOFT MODULARIZATION 42
0.73 20.04
CAGRAD
0.83 #0.05
0.25
NASHMTL (every
0,91 #0.03
NASHMTL
(every 50)
0.85 #0.02
1M
2M
NASHMTL
(every 100)
0.87 +0.03
Training Steps
NASHMTL (ours) (every
+0.13
11.0
NASHMTL (ours) (every 50)
0.76 +0.10
NASHMTL (ours) (every 100)
0.80 10.12
Time
UW 18
0.77 20.05
2.0
1.5
FAMO (ours)
0.83 20.05
1.0
1.0
Table 4: MTRL results (averaged over 10 runs) on the Metaworld- 10 benchmark.
Figure Training Success Rate Time.
*We could not reproduce the MTRL results of NASHMTL exactly, so we report both the results from the original paper and our reproduced results.

---

we observe that FAMO performs comparably to CAGRAD and out- performs PCGRAD and the average gradient descent baselines by large margin. FAMO also outperforms NASHMTI. based on our implementation. Moreover, FAMO is significantly faster than NASHMTL,, even when it applied once every 100 steps.
Findings: From Table 5.2
5.3
MTL Efficiency (Training Time Comparison)
Figure 4
provides the FAMO's average training time per epoch against that of the baseline methods.
LS
MGDA
PCGrud
CAGrad
IMTL-G
NushMTL
FAMO
1.91
1.33
3.50
2.57
1.48
2.15
L.17
4:51
3.00
1.07
135.40 40
3,00
1.02
53.77
1.00
241
CityScapes tasks)
NYU-v2 tasks)
QM-9 tasks)
CelebA (40 tasks)
Figure 4: Average training time per epoch for different MTL optimization methods. We report the relative training time a methed to that of linear scalarization (LSS) method (which uses the average gradient).
Findings: From the figure, we observe that FAMO introduces negligible overhead across all benchmark datasets compared to the LS method, which is, in theory, the lower bound for computation time. In contrast, methods like NASHMTL have much longer training time compared to FAMO. More importantly, the computation cost of these methods scales with the number of tasks. In addition, note that these methods also take least ((K) space to store the task gradients, which is implausible for large models in the many-task setting (i.e., when m [0] and K are large).
5.4 Ablation on
In this section, we provide the ablation study the regularization coefficient in Figure
7 =000001
7
0,001
7 = 0.01
36.08
06.16
Am1%
90
60
3.55
-3.59
CitySeapes
NYU-v2
QM-9
CelebA
Figure 5: Ablation over plot the performance of FAMO (in terms of Am using different values of from (0.0001, 0.001, 0.01) on the four supervised MTL benchmarks.
Findings: From Figure 5 we can observe that choosing the right regularization coefficient can be crucial. But except for CityScapes, FAMO performs reasonably well using all different ys. The problem with CityScapes is that one of the task losses is close to 0 at the very beginning, hence small changes task weighting can result in very different loss improvement. Therefore we conjecture that using a larger y, in this case, can help stabilize MTL.
6 Conclusion and Limitations
In this work, we introduce FAMO, fast optimization method for multitask learning (MTL) that mitigates the conflicting gradients using O(1) space and time. As multitasking large models gain more attention, we believe designing efficient but effective optimizers like FAMO for MTL is crucial. FAMO balances task losses by ensuring each task's loss decreases approximately at an equal rate. Empirically, we observe that FAMO can achieve competitive performance against the state-of-the-art MTL gradient manipulation methods. One limitation of FAMO is its dependency on the regularization parameter 7, which is introduced due to the stochastic update of the task weighting logits w. Future work can investigate more principled way of determining 7
10

---

References
[1] L. C. Blum and J.-L. Reymond. 970 million druglike small molecules for virtual screening in the chemical universe database GDB-13. J. Am. Chem. Soc., 131:8732, 2009.
[2] Felix JS Bragman, Ryutaro Tanno, Sehastien Ourselin, Daniel C Alexander, and Jorge Cardoso. Stochastic filter groups for multi-task cnns: Learning specialist and generalist convolution kernels. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1385-1394, 2019.
[3] David Bruggemann, Menelaos Kanakis, Stamatios Georgoulis, and Luc Van Gool. Automated search for resource-efficient branched multi-task networks. arXiv preprint arXiv:2008. 10292, 2020.
[4] Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott I Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv: 127/2, 2023.
[5] Rich Caruana. Multitask learning. Machine learning, 28(1))41-75, 1997.
[6] Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gra- dient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pages 794-803, PMLR, 2018.
[7] Zhao Chen, Jiquan Ngiam, Yanping Huang. Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick sign: Optimizing deep multitask models with gradient sign dropout. arXiv preprint arXiv:2010.06880 2020.
[8] 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, pages 3213-3223, 2016.
[9] Jean-Antoine Désidéri. Muiliipe-grraiinn descent algorithm (mgda) for multiobjective opti- mization. Comptes Rendus Mathematique, 350(5-6):313-313 2012.
[10] Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric.
arXiv preprint arXiv: 022282, 2019.
[11] Chris Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Efficiently identifying task groupings for multi-task learning. Advances Neural Information Processing Systems, 34:27503-27516, 2021.
[12] Yuan Gao, Haoping Bai, Zequn Jie, Jiayi Ma, Kui Jia, and Wei Liu. Mtl-nas: aank-ago neural architecture search towards general-purposs multi-task learning. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pages 11543-11552, 2020.
[13] Michelle Guo, Albert Haque, De-An Huang, Serena Yeung, and Li Fei-Fei. Dynamic task prioritization for multitask learning. In Proceedings of the European conference on computer vision (ECCV), pages 270-287, 2018.
[14] Pengsheng Guo, Chen- Yu Lee, and Daniel Ulbricht. Learning to branch for multi-tash learning. In International Conference Machine Learning, pages 3854-3863. PMLR, 2020.
[15] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning, pages 1861-1870 PMLR, 2018.
[16] Adrián Javaloy and Isabel Valera. Rotograd: Dynamic gradient homogenization for multi-task learning. arXiv preprint arXiv:2103.026311 2021.
[17] Alexandr Katrutsa, Daniil Merkulov, Nurislam Tursynbek, and Ivan Oseledets. Follow the bisector: a simple method for multi-objective optimization. arXiv preprint Niv:2007.06937 2020.

---

[18] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-tat task learning using uncertainty to weigh losses for scene geometry and semanties. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1482-7491, 2018.
[19] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[20] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, al. Segment anything. arXiv preprint arXiv:2304.02644; 2023.
[21] lasonas 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, pages
2017.
[22] 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: 15:96111-, 2022.
[23] Baijiong Lin, Feiyang Ye, and Yu Zhang. A closer look at loss weighting in multi-task learning. arXiv preprint arXiv:2/ 10,00033, 2021.
[24] Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-tast learning. Advances in Neural Information Processing Systems, 14:18878-18880, 2021.
[25] Liyang Liu, Yi Li, Zhanghui Kuang, Jing-Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Towards impartial multi-task learning. International Conference on Learning Representations, 2020.
[26] Shikun Liu, Stephen James, Andrew J Davison, and Edward Johns. Auto-lambda: Disentangling dynamic task relationships. arXiv preprint arXiv:2202.3009, 2022.
[27] 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, pages 1871-1880, 2019.
[28] Ziwei Liu, Ping Luo,, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[29] Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Philip Yu. Learning multiple tasks with multilinear relationship networks. Advances in neural information processing systems, 30, 2017.
[30] Ishan Misra, Abhinav Shrivastava, Abhinav Gupta, and Martial Hebert. Cross-stitch networks for multi-task learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3994-4003, 2016.
[31] Pushmeet Kohli Nathan Silberman, Derek Hoiem and Rob Fergus. Indoor segmentation and support inference from rgbd images. In ECCV, 2012.
[32] 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.
[33] Lucas Pascal, Pietro Michiardi, Xavier Bost, Benoit Huet, and Maria A Zuluaga. Improved optimization strategies for deep multi-task networks. arXiv preprint arXiv:2109.1 19678 2021.
[34] 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, pages 4822 4829, 2019.
[35] Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. arXiv preprint 1100001800 2018
12

---

[36] Jiayi Shen, Xiantong Zhen, Marcel Worring, and Ling Shao. Variational multi-task learning with gumbel-softmax priors. Advances Neural Information Processing Systems, 34:21031-21042, 2021.
[37] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context-based representations. arXiv preprint arXiv:2102.06177, 2021.
[38] Trevor Standley, Amir Zamir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi task learning? In International Conference Machine Learning, pages 9120-9132. PMLR, 2020.
[39] Sebastian Thrun and Joseph O' Sullivan. Discovering structure in multiple learning tasks: The algorithm. In ICML, volume 96, pages 489-497, 1996.
[40] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansheke, Mare Proesmans, Dengxin Dai, and Luc Van Gool, Multi-task learning for dense prediction tasks: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
[41] 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.
[42] Ruihan Yang. Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. arXiv preprint arXiv:2003. 13661, 2020.
[43] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. arXiv preprint arXiv:2001.0678) 2020.
[44] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, In Conference Robot Learning, pages 1094-1100. PMLR, 2020.
[45] Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. Proceedings of the IEEE conference computer vision and pattern recognition, pages 3712-3722, 2018
[46] 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.
[47] Shijie Zhu, Hui Zhao, Pengjie Wang, Hongbo Deng, Jian Xu, and Bo Zheng. Gradient deconfliction via orthogonal projections onto subspaces for multi-task learning.
13

---

A Gradient Manipulation Methods
In this section, we provide a brief overview of representative gradient manipulation methods in multiask/mulliojjetiiee optimization. Specifically, we will also discuss the connections among these methods.
Multiple Gradient Descent Algorithm (MGDA) 9 35
The MGDA algorithm is one of the MGDA, the per step update de is
earliest gradient manipulation methods for multitask learning. found by solving
- max min VL, V - Idd1. de id[R]
As a result, the solution d" of MGDA optimizes the "worst improvement" across all tasks or equivalently seeks an equal descent across all task losses as much as possible. But in practice, MGDA suffers from slow convergence since the update d' can be very small. For instance, if one task has a very small loss scale, the progress of all other tasks will be bounded by the progress on this task. Note that the original objective (6) is similar the MGDA objective in the sense that we can view optimizing (6) optimizing the log of the task losses. Hence, when we compare FAMO against MGDA, one can regard FA MO as balancing the rate loss improvement while MGDA balances the absolute improvement across task losses.
Projecting Gradient Descent (PCGRAD) 43 PCGRAD initializes dpc Vla,i, then for each task
i, PCGRAD loops over all task j (in a random order, which crucial as mentioned in [43]) and removes the "conffict"
r VVgg, Vlga if v/c Vez, O.
PCC UPC
||e
In the end, PCGRAD produces de E ufc. Due to the construction, PCGRAD will also help improve the "worst improvement" across all tasks since the "conflicts" have been removed. However, due to the stochastic iterative procedural of this algorithm, it is hard to understand PCGRAD from first principle approach.
Conflict-averse Gradient Descent (CAGRAD)
(4
dy is found by solving
max min VeT pp delt xc
s.t. |d Vla, s cl
Ei= l. CAGRAD seeks an update that optimizes the "worst improvement" as much as possible, conditioned on that the update still decreases the average loss. By controlling the hyperparameter e, CAGRAD can recover MGDA (c oo) and the vanilla averaged gradient descent (c + 0). Due to the extra constraint, CAGRAD provably converges to the stationary points of to when Osesl.
Here, lo, =
Impartial Multi-Task Learning (IMTL-G) (25 cosine similarity with any task gradients:
IMTL-G finds de such that shares the same
Vlza
Vi+j. j. df
d.
and d, Σ wa Vlr for some wy Sk.
|elell |
The constraint that de Σl Wi. is for preventing the problem from being under-determined. From the above equation, we can see that IMTL-G ignores the "size" each task gradient and only cares about the "direction". As a result, one can think of IMTL-G as variant of MGDA that applies to the normalized gradients. By doing so, IMTL-G G does not suffer from the straggler effect due to slow objectives. Furthermore, one view IMTL-G as the equal angle descent, which is also proposed in Katrutsa et al. [1],, where the objective is to find d such that
Vi
cos(d, vl.c)) cos (d, Vly.).
14

---

NASHMTL(32 NASHMTL finds de by solving a bargaining game treating the local improvement of each task loss as the utility for each task:
log (Vei,
deR s1
Note that the objective of NASHMTL implicitly assumes that there exists d such that V i, V ve 0 (otherwise we reach the Pareto front). It is easy to see that
Vha [Vl,
max 51 F=
log (Vl,,
max |d]st
log
log cos (Vli,, max [d]st F1
Σ
E1
Therefore, due to the log, NASHMTL also ignores the "size" of task gradients and only cares about Then, according the KKT condition, know:
Vh. their "directions". Moreover, denote u4 Ive.ll"
d= 1 - u
Σ u ad = O, a
Consider when &s 2. we take the equal angle descent direction: d, = (u1 + u2) (note that as uj and U2 are normalized, their bisector just their average). Then it is easy check that
d= d= (u] w2) U1 u1( + u2) U2 where a= u [11 + U2)
un (m) u22
4
As result, we can see that when h= 2, NASHMTL is equivalent to IMTL-G (or the equal angle descent). However, when k: 2, this is not in general true.
Remark Note that of these gradient manipulation methods require computing and storing K task gradients before applying to compute d, which often involves solving additional optimization problem. Hence, these methods can be slow for large K and large model sizes.
B Amortizing other Gradient Manipulation Methods
Although FAMO uses iterative update on w, it is not immediately clear whether we can apply the same amortization easily on other existing gradient manipulation methods. In this section, discuss such possibilities and point out the challenges.
Amortizing MGDA This is almost the same as FAMO, except that MGDA acts on the original task losses while FAMO acts on the log of task losses.
Amortizing PCGRAD For PCGRAD, finding the final update vector requires iteratively projecting one task gradient to the other, so there no straightforward way of bypassing the computation of task gradients.
The task weighting IMTL--G is computed by a series of matrix-matrin and matrix-vector products using task gradients (25). Hence, it is also hard to amortize its computation over time.
Amortizing IMTL--
Therefore, we focus on deriving the amortization for CAGRAD and NASIMTL..
Amortizing CAGRAD For CAGRAD, the dual objective is
min F(w) = - + clgg0 |gol| w
(15)
where go Vlu, and Iw E W. Denote
G
15

---

Now, if we take the gradient with respect to in (1 we have:
OF
Igoll
G" go
G" Jw-
(16)
dw
As result, in order to approximate this gradient, one can separately estimate:
{(0) {(0 ago)
G 90 A
((0) { ogw)
G' Iw
(17)
Igoll VITGT
|g V
Once all these are estimated, one can combine them together to perform single update on w. But note that this will require forward and backward passes through the model, making it harder to implement in practice.
Amortizing NASHMTL Per derivation from NASHMTL
32 the objective is solve for
G W.
(18)
One can therefore form objective:
min F(w) = |||''GGGG Gw 1
(19)
Taking the derivative F' with respect to w, we have
OF 2G G G'
tw G' Iw. -1 + G¹ 9n - 10-1oww - (wow).
(20)
Therefore, to approximate the gradient of w. one needs first estimate
GO Iho L(0) I(O agw)
(21)
Then we estimate
L(0) L(0 aG
GT G
(22)
Again, this results in 3 forward and backward passes through the model, let alone the overhead of resetting the model back to d (requires copy the original weights).
In short, though is possible to derive fast approximation algorithm to approximate the gradient update for some of the existing gradient manipulation methods, often involves much more complicated computation compared to that of FAMO.
C FAMO Pseudocode in PyTorch
We provide the pseudocode for FAMO in Algorithm 2 To use FAMO, one just first compute the task losses, call get__eeigtteeeeee www__weigghhdeeses.. to get the weighted loss, and do the normal backpropagation through the weighted loss. After that, one call update to update the task weighting.
D Toy Example
We provide the task objectives for the toy example in the following. The model parameter 0 (0,, 02) e 2R² and the task objectives are I and I.2:
(²))) 0.1 (() tanh (-02)l, 0.000005)) + 6,
and (²( ())) + c2))))=)f), where
(1) log
f2(0) log (max((0.5(-0) +3) tanh (-02)+ + 2), 0.000005)
(0) (-
20,
(0) (1' (-02 (2:02-8))//11 3/1 200,
c1 (0) max(tanh (0.5 0z), 0) and c2( max(tanh (-0.5 02), 0).
16

---

Algorithm 2 Implementation of FAMO in PyTorch-like Pseudocode
class FAMO:
def i__init__.. ___init__(self (self, num_tasks, min_losses, a=0.025, y=0.001):
# mii_losses (num_ tasks,) the loss lower bound for each task. self www.minnllsssee min_losses
self torch. tensor ([0.0 0] * num_tasks, - self.xi_ops torch optim (self. xi], 1r=a, weight_decayf1)
def get get_ losses):
# losses (num_tasks,
z f.softmax(self.xt -1)
D losses self min_losses + 1e-8
с 1 (z D) sum () ddtacch))
loss (c * D.log() * 2) sum()
return loss
def update prev rev__osses, curr curr_losses):
# prev_losses (num_taske tasks,)
# curr curr_losses (num_ tasks,)
delta (prev_losses self min_ 1e-8) .1gg() (curr_lossed with torch. enable_grad
self wwt.mmnnssssees + 1e-8). 1og ()
d = torch. autograd grad (F. softmax (self.u) xi, -1),
self.xi,
grad_outputs del.a delta. detach ()) [o]
self. zero, zero_grad()) self self.xi.gree grad = d self self.xi_opt.stey
E Experimental Results with Error Bars
We followed the exact experimental setup from NASHMTL 32 Therefore, the numbers for baseline methods are taken from their original paper. In the following, we provide FAMO's result with emor bars.
17

---

Segmentation
Depth
Surface Normal
Method
Angle Dist 4
Within /
A Am%
m
Pix Acc
Abs Err
Rel Etr
Mcan
Median
11.25
22.5
30
FAMO (mean)
38.88
64,90
0.5474
0.2194
25.06
19.57
29.21
56.61
68.98
-4.10
FAMO (stderr) +0.54
=0.21
#0.0016
=0.0026
#0.06
+0.09
#0.17
#0.19
=0.14
#0.39
Table 5: Results on NYU-v2 dataset (3 tasks). Each experiment is repeated over 3 random seeds and the mean is reported. The best average result is marked in bold. MR and Am% the main metrics for MTL performance.
Method
EJOMO
GL1MO
(R²)
ZPVE
U,
Н
G
Am7%
MAE
FAMO (mean)
0.15
0.30
94.0
95.2
1.63
4.95
70.82
71.2
71.2
70,3
0.10
58.5
FAMO (sIderr)
Table 6: Results QM-9 dataset (11 tasks). Each experiment is repeated over random seeds and the mean is reported. The best average result is marked in bold. MR and Am% the main metrics for MTI. performance.
CityScapes
CelebA
Method
Segmentation
Depth
Am%
Am%%
mloU
Pix Acc 1
Abs Eer
Rel Err 1
FAMO (mean)
74.54
93.29
0.0145
32.59 #1.0
8.13
1.21
FAMO (stderr)
+0.11
+0.04
=0,0009
#1.98
#0.24
Table 7: Results SctySsapes (2 tasks) and CelebA (40 tasks) datasets. Each experiment is repeated over random seeds and mean reported. The best average result marked bold. MR A Am are the main metrics for MTL performance.
18


Reviews:
[{"review_id": "Dku8jKw95r", "reviewer": "Reviewer_gjX4", "summary": "This paper proposed a novel multi-task optimization method aimed at mitigating conflict between task gradients without inducing the substantial time slowdown that typically comes with specialized multi-task optimizers.\nThe proposed method, called FAMO, aims at improving the worst-case rate of improvement across all tasks at each step.\nAt a high level, it does this by setting the weights of each task to minimize the gradient norm of the log task loss (similar to how MGDA sets weights by minimizing the gradient norm of the task loss). They show that finding these weights is the dual of finding the direction that maximizes the minimum rate of improvement across tasks. In practice however, these weights are learned over time via SGD updates after each optimization step of the main model, rather than optimized in full at each model update step.\n\nThe authors then demonstrate their method on a toy 2D loss landscape (used in prior work), showing that their method reaches the Pareto frontier of this problem, like other optimization methods (and unliked naive Adam), but has the closest overall compute time to Adam. Next, the authors compare their method to 11 other optimization methods for MTL on 4 supervised MTL and 1 RL setting finding that FAMO compares competitively (it is consistently in the top 2 or 3 performing methods) in all settings.  Finally, the authors demonstrate the performance efficiency of their method in the supervised learning settings showing that FAMO is slightly faster than comparable methods on settings with 2 or 3 tasks, and orders of magnitude faster as the number of tasks increases to 10 and 60.", "questions": "- Is there a typo in the psuedo code of algo 2? z is computed via softmax, but then the raw weights w are used for the rest of the get_weighted_loss method.", "limitations": "yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The method being proposed has clear benefits in terms of speed / efficiency, especially as the number of tasks scales, which is very important in many MTL settings.\n- The method is compared to a broad and representative set of optimization methods, and performs competitively on all benchmarks proposed even outperforming slower methods consistently.\n- Moreover, the experimental methodology seems sound; all experiments consider multiple random seeds and 2 sensible metrics to compare methods are considered.\n- The paper is well written, with good figures, clear descriptions, and the motivation is straightforward. Moreover, I find the proposed method to be explained well.", "weaknesses": "- The primary contribution of FAMO (in my view) is the decision to maximize minimum task rate of improvement, by focusing on the gradient of the log loss, rather than maximizing the minimum total improvement (i.e. MGDA). However, this is not the “fast” part. The fast part of the method comes from optimizing the task-weights slowly, via gradient updates, throughout the optimization process. The authors do well to point this out in the appendix, but the effects of these contributions should be analyzed separately. In a deep neural network it is possible that the generalization benefits come from the amortized optimization of task weights, rather than the specific weight objectives. This could easily be addressed by, for instance, comparing amortized MGDA to FAMO, which is perhaps the most comparable setting.\n- I feel as though GradNorm is also natural comparison to this method, given that it also learns task weights during training (in an amortized fashion), with the goal of balancing the average rate of improvement across tasks.\n- [21] is referenced as suggesting that current optimization methods aren’t worth the tradeoff in computation time, but iirc their conclusion was that optimization methods actually don’t help performance over the baseline, which appears to not be true in this papers experiments. This discrepancy should likely be commented on.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and comments. We address your concerns and questions in the following.\n\n---\n\n**Weakness:**\n\n**1. Compare FAMO against amortized MGDA.**\n\nPlease refer to Figure 1 above. We see that since amortized MGDA (like MGDA) is seeking equal descent for each loss, it performs poorly compared to FAMO.\n\n**2. Compare FAMO against GradNorm.**\n\nPlease refer to Figure 2 in [1] for a comparison of IMTL-G versus GradNorm on ResNet-50. It is known that GradNorm does not perform as well as gradient manipulation methods. Moreover, GradNorm actually is $\\mathcal{O}(K)$ in space and time, because in order to perform the GradNorm update, one needs to take the gradient of each task’s gradient norm. Hence it requires computing each task’s gradient first.\n\n**3. [21] claims that MTL methods are not useful.**\n- [21] only studies MGDA, PCGrad, RLW, and GradDrop. It does not consider more recent gradient manipulation methods like CAGrad, NashMTL, IMTL-G, etc. \n- For supervised learning, [21] only tries multi-MNIST (too simple) and Cityscapes (only 2-task). It hasn’t been tried on celeb or NYU-v2.\n- We understand that what [21] claims is that there always exists a set of unit scalarizations that performs on par with existing gradient manipulation methods. But it does not provide a method for how we can decide this set of task weights except for just doing a grid-search, and this scales poorly with the number of tasks. On the other hand, the main drawback of existing gradient manipulation methods is that they are too slow to be applied. Hence we find that FAMO achieves the sweet spot in that it achieves the state-of-the-art performance across a wide range of commonly used MTL benchmarks and it is very cheap to apply in practice.\n\n---\n\n**Questions:**\n\n**1. Typo in Pseudocode.**\n\nThanks for catching that! We will correct it in the final version.\n\n---\n\n**Reference:**\n\n\n[1] [Towards Impartial Multi-task Learning](https://openreview.net/pdf?id=IMPnRXEWpvr).\n"}, {"review_id": "aQ1WpQaTPj", "reviewer": "Reviewer_XmqU", "summary": "This paper proposes a dynamic convex combination multi-task learning losses reweighting so that it can decrease different task losses more balancedly while having little computational overhead than simple task loss average. Specifically, they formulate an optimization problem to find a new “gradient” direction that can achieve sufficient relative-to-magnitude decrease for the worst task loss. Instead of solving this optimization exactly every update, the authors amortize the optimization by keeping dual variables on the probability simplex across updates and update them through one-step gradient descent after every primal variable (learnable parameter) update. Experimentally, the authors show that their proposed method FAMO can achieve comparable/superior on 4 multitask supervised learning benchmarks and 1 multi-task RL benchmark, while being more time-saving compared to other gradient manipulation methods.", "questions": "**[Why is FAMO $O(1)$ in space?]** For performing backprop over the simple case of the average task loss (LS), either all the $K$ tasks’ computation graphs are computed simultaneously and stored in the memory (needed for backprop) at the same time, or the tasks’ losses are computed one after another. In the first case, the computation graph would already have space complexity of O(K), while in the second case, the time to iterate over all tasks would have time complexity of O(K). FAMO is basically calling backprop on a reweighting of the logarithm of the task losses so it should at least have the same space and time complexity as LS. So why do the authors claim FAMO is $O(1)$ in space? Besides, storing the K weights are already technically O(K) in space. I can understand that FAMO would take less time than the other multi task gradient manipulation methods which require non-amortized optimization, but I would appreciate the authors be more rigorous about their claims with more explanation.\n\n**[Why is STL so good?]** On QM-9 and MTRL results, it seems STL performs (much) better than other methods. Can the authors provide explainations on why that is the case?\n\n**[Mean Rank]** In Table 3, why is IMTL-G’s mean rank (0.84) smaller than 1? On page 7 the paper says for a method that ranks first for every task, its MR is still 1.\n\n**[Obeying the simplex constraint]** This question is relatively minor but I would suggest the authors consider alternative ways to  obey the simplex constraint other than using a softmax which couldn’t really exactly represent the corners of the simplex. Instead, one could apply $L^2$-projected gradient descent instead of gradient descent on $z$ by projecting the updated $z$ back onto the simplex. There exist efficient algorithms for simplex projection which don’t require solving the projection optimization problem iteratively: https://arxiv.org/pdf/1309.1541.pdf.", "limitations": "The authors mention their approach depending an extra weight decay hyperparameter as an limitation of their work. I don’t believe this paper needs to discuss negative societal impacts. I currently rate the paper as borderline accept because of the weaknesses described above. However, I’m open to increase my score given satisfactory author rebuttal responses.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is in general well-written, with a nice discussion about different multi-task gradient manipulation technique in Appendix A.\n2. The proposed method FAMO empirically is competitive with other methods on multiple benchmarks.", "weaknesses": "\n**[Lack of proof for significant contribution]**\n\nIn terms of methodology, as the authors mention themselves, FAMO’s motivation formulation is basically applying MGDA to the logarithm of all task losses. Thus the formulation is practical but not very novel. What distinguishes FAMO from MGDA fundamentally is FAMO’s amortized optimization of the dual variable. I will raise additional question later in this section about solidifying this contribution. In terms of experiment results, the authors claim FAMO can perform better/equal to other baseline methods in terms of (1) final performance, and (2) wall-clock time. For final performance, it is not convincing enough that FAMO can always match/outperform other MTL methods. On CityScapes and CelebA, it seems to fall behind some other methods. On QM-9, it’s also not clear to me what’s the value of multi-task learning — it seems single task learning (STL) is better (with lower MAE values) than almost all MTL methods  (including FAMO) (please do correct me if I’m misunderstanding this). I would appreciate authors to provide further clarifications about these points. In terms of wall-clock time, the authors only show the wall clock time for each epoch for different methods. However, this neglects the fact that different methods might require different number of epochs to converge. It is conceivable that a non-optimization-amortized method could take longer per epoch but converge in fewer epochs than FAMO and thus could use less total time. I think the authors should also report the total time as it should be final metric to judge the wall-clock savings.\n\n**[Understanding the impact of amortizing the optimization]** A primary contribution of the paper is to propose to amortize the optimization to find the update direction instead of solving it exactly each update. Hence I believe it is important and useful that the authors provide experiment results on the non-amortized version of FAMO, where the optimization in equation (5) is solved exactly in each iteration. With this non-amortized result, it would provide the context of whether the amortization deteriorates the performance, allowing the readers to understand the room of improvement with amortization. Besides, it would also be interesting to see the amortized vs non-amortized comparison for MGDA to further understand this point.\n\n**[Potentially incorrect proof for proposition 3.1]** In the proof for proposition 3.1, the authors use the Lagrangian of the dual by adding $cz^\\top \\mathbf{1}$ to the objective to get rid of the simplex constraint. However, the simplex constraint is not only about sum of the variables being 1, but also about having nonnegative individual values. Thus, I believe these nonnegative constraint should also be reflected in the Lagrangian. Otherwise it needs to be specifically argued that the optimal solution cannot occur on the edges of the simplex.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and comments. We address your concerns and questions in the following.\n\n---\n\n**Weakness:**\n\n**1. Lack of significant contribution. Why STL is good?**\n\nThe 2 contributions of FAMO are: 1) we propose that we should do equal-rate descent instead of equal descent, and 2) we introduce a fast amortizing trick that can be applied to drastically speed up the optimization, which so far only MGDA and FAMO could use with one extra forward pass in stochastic gradient descent setting. \nThe STL learning performance is computed using the same network but only trained on a single task. So if we have K tasks, we essentially use K more parameters. We provide the plot of the delta M over epoch for FAMO, NashMTL, amortized MGDA, and log-MGDA (the exact algorithm of FAMO) in Figure 1. From the plot, we observe that FAMO not only is fast and performs well asymptotically, but also achieves such good performance fast over epoch. \n \n**2. Non-amortized log-MGDA.**\n\nPlease refer to Figure 1 for the log-MGDA performance. It performs well but FAMO tracks its performance well.\n\n**3. Incorrect Proof**\n\nThanks for catching that. The proof is correct with the assumption that the optimal $z$ is an interior point of $\\mathcal{S}^K$, hence $z > 0$. Otherwise, we will have the extra Lagrangian multiplier $\\nu \\in \\mathbb{R}_+^K$ such that $\\nu \\odot z = 0$, i.e. $\\nu \\odot z$ is element-wise zero. Under this case, it means that it is not possible to achieve equal rate descent with any convex combination of the gradients, the best we can hope for is that we concentrate the weighting only on a single (or a few) task gradient(s). Consider an obtuse triangle with points A, B, and C, where A-B represents $g^1$ and A-C represents $g^2$ ($g^1$, $g^2$ represent $\\nabla \\log \\Delta^1$ and $\\nabla \\log \\Delta^2$), and assume the angle $\\angle ABC > 90\\degree$, then we cannot find any update along B-C that results in equal descent. The best we can hope for is update A-B, which is $g^1$. However, not that the optimal solution is still the best possible update in terms of balancing the loss decrease rate.\n\nNote that this is the same result in the original MGDA paper, except that they are working on the original objectives while FAMO works on the log objectives. We will make sure to make it clear about our assumption in the final updated version.\n\n---\n\n**Questions:**\n\n**1. Why is FAMO $\\mathcal{O}(1)$ in space?**\n\nFAMO remains  $\\mathcal{O}(1)$ in space, for the same reason that averaged gradient descent is  $\\mathcal{O}(1)$ in space. This is due to the nature of the backpropagation algorithm, as explained in Line 32-34. In general, if we weigh the loss first, say $L_w = w_1 L^1 + w_2 L^2$, since we know this weighting $w_1$ and $w_2$ beforehand, one can take advantage of the backpropagation as $$w_1 \\frac{d L^1}{d \\theta} + w_2 \\frac{d L^2}{d \\theta} = \\frac{d L_w}{d \\theta},$$ which can be computed within 1 backward pass **without** extra space. This is similar to the fact that if we compute the loss over a batch, the space to store the gradient will not scale with the batch size. This is because we average (or sum) the batch loss first before doing the backpropagation, and taking the derivative is a linear operation.\n\n**2. Why is STL so good?**\n\nIt is addressed in response to weakness 1. We mention it here again for convenience. STL reported in this paper means we use the same neural architecture but only train for a single task. So if we have $K$ tasks it means we need to train $K$ models separately. The total number of effective parameters is much larger than that of an MTL model.\n\n**3. Mean Rank < 1**\n\nThanks for catching that. It is a typo. The Mean Rank for Celeb-A should all be increased by 1. This is because in our implementation we first compute the Mean Rank (0-indexed) and then add 1 on top of it. But due to a mistake, we forgot the add 1 for the ranks in Celeb-A. It will be corrected in the final version of the paper.\n\n**4. Alternative to reparametrization (projected gradient)**\n\nThanks for the suggestion. We agree that it would be interesting to investigate whether performing projected gradient descent on w can be more efficient for FAMO. We adopt the reparameterization because then the problem becomes completely unconstrained. Usually constrained optimization is harder to solve stably in practice."}, {"review_id": "ycZ7agGApv", "reviewer": "Reviewer_DG6u", "summary": "This paper proposes a new multi-task balancing method with claimed O(1) efficiency, which is much more efficient than previous gradient manipulation methods.\n\nThe core idea is to let the multiple losses decrease at roughly the same speed.\n\n", "questions": "1. Equation 6 is transformed from Equation 5 using L'Hopital's rule, based on the step size \\alpha is reaching to 0, however, learning rate \\alpha is normally something like 0.001, which is much larger than 0. So, Equation 6 is an approximation here and should be made clear.\n\n2. Could authors elaborate a bit on why Equation 7 is the dual form for (6)?\n\n3. In Section 3.2,\"But then, assume [author's equation but hard to write for me here], one can approximate (8) by:\", may I ask why we can assume this equation? what is the intuition for this assumption?\n\n4. What is w? the final weighs losses I guess?. In the paper, authors have used v_{t} and z_{t} as task weight. And they did not clearly explain w.\n", "limitations": "No social impact", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The method is more efficient than previous gradient manipulation methods.\nAuthors conducted lots of experiments on several datasets and show that the method is performing well on them.", "weaknesses": "1. I think the contribution is little bit over-stated. When I read the title, abstract and introduction, I thought authors will present a novel method to manipulate task gradients with O(1) time and space complexity, which is awesome. Why I thought that? For example, in Figure 1, they compare FAMO with other gradient balancing methods like PCGrad and CAGrad, which balance task gradients with O(K) time and space complexity.\n\nHowever, authors actually do not manipulate gradients but weight losses. Losses weighting are not novel path with previous methods DWA and UW. If you balance losses rather gradients, of course you can have O(1) time and space complexity since you do not touch parameter gradients w.r.t to K losses. \nIn Question [highlight in the introduction], authors want to design an MTL optimizer, which should handle the gradients of parameters. I am not sure if the algorithm proposed in the paper can be called a new optimizer?\nMy suggestion is that authors can re-factory the story [title, abs, intro] as an efficient loss balancing methods, which seems more suitable to their contributions.\n\n2.  The method part is not presented very clearly, see the questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and comments. We address your concerns and questions in the following.\n\n---\n\n**Weakness:**\n\n**1. FAMO does not manipulate gradients but reweighs them. Re-weighting methods exist like UW/DWA. Should refactor the abstract, intro, method.**\n\nExisting gradient manipulation methods all produce the final manipulated gradient as a linear combination of task gradients. Hence from this perspective, they are also re-weighting methods. We think the confusion originates from the definition of “manipulation”. In our case, we mean that the final update is some function of the task gradients for achieving a certain property. From this perspective, UW or DWA approach the MTL problems from the “loss” perspective, i.e., they propose a way to manipulate losses. While gradient manipulation methods will not necessarily propose a meaningful new objective to optimize, they aim to manipulate the task gradients such that after performing the update, certain a property is achieved. In our case, it is that the task losses descend at roughly an equal rate. We will make the above more clear in the revised version of the paper.\n\nApart from the above, we point out that we compare FAMO against UW and DWA in our tables. We emphasize that re-weighting methods, so far, do not reach similar performance compared to state-of-the-art gradient manipulation methods.\n\n**2. Methods are not clear**\n\nWe address your concerns in the following.\n\n---\n\n**Questions:**\n\n**1. Eq.6 is an approximation of Eq.5, which should be made clear.**\n\nWe will make sure it is clear in the final revised version of the paper.\n\n**2. Why Eq.7 is the dual form of Eq.6?**\n\nThe proof is shown on page 16 in the appendix. In summary, the objective is convex and we can hence switch the min and max.\n\n**3. What is the intuition in Sec. 3.2 from Eq.8 to Eq.9?**\n\nThe transition from Eq. 8 to Eq. 9 follows the first-order Taylor expansion. Specifically, denote $h(\\theta_t) = \\log \\Delta_t$,\n\n$$\nh(\\theta_t) - h(\\theta_{t+1}) = h(\\theta_t) - h(\\theta_t - \\eta \\nabla \\log \\Delta_t z) \\approx \\eta \\bigg(\\nabla_\\theta h_t\\bigg)^\\top \\bigg( \\nabla \\log \\Delta_t z\\bigg) = \\eta  (\\nabla \\log \\Delta_t) ^\\top (\\nabla \\log \\Delta_t) z.\n$$\n\nWe will make it clear in the final version of the paper.\n\n**4. What is $w$, the final weight loss?**\n\n$w$ is the logits to the task weights (check Alg. 1). For instance, $z = \\textbf{softmax}(w)$."}, {"review_id": "nbrKzBW5Fy", "reviewer": "Reviewer_otYb", "summary": "The paper addresses an important issue with (most) prior MTL optimization techniques which requires the computation of all per-task gradients, during training, for obtaining the update direction. This results in a $\\mathcal{O}(K)$ requirement in space and time where $K$ is the number of tasks. The authors propose a novel approach for dynamically weighting tasks using $\\mathcal{O}(1)$ space and time.", "questions": "1. Could you provide some intuition as to why the proposed method alleviates the CG problem in MTL, only by ensuring each task’s loss decreases at an equal rate?\n2. Have you tried performing several optimization steps for approximating the solution in Eq. 7 at each training iteration?\n", "limitations": "Limitations addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper proposes a novel approach for addressing a known pain point with (most) SoTA multi-task optimization methods.\n2. The paper provides strong empirical evidence which indicates the method performs on par with SoTA MTL methods (e.g., NashMTL, IMTL-G, CAGrad) while significantly reducing the optimization time.\n3. The paper is well-written and well-structured. Connection to prior works in MTL literature is well established (Appendix A-B).\n", "weaknesses": "1. While the division of the proposed method follows a different trajectory, the proposed algorithm is fairly similar to MGDA. Hence, more discussion is needed on the connection between the proposed method and MGDA.\n2. Please provide additional details on the experimental sections (data split sizes, HPs, etc.).\n3. Hyperparameter (HP) optimization:\n    - Previous papers on MTL (e.g., CAGrad) did not include a validation split for NYU/Cityscapes. If the authors follow the same settings, how did you tune FAMO’s specific HPs? \n    - Missing details on the tuning of the lr for the task weights.\n4. From section 5.4 (and as stated by the authors), it appears the method is highly sensitive to the choice of HPs.\n5. It would be beneficial to include figures of the losses throughout the optimization process and some metric to verify that the losses decrease in an (approximately) equal rate.\n6. It would be beneficial if you could provide a figure of the task weights throughout the optimization process.\n\nMinor: \nTable 2, MR column missing indication for best-performing method.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and comments. We address your concerns and questions in the following.\n\n---\n\n**Weakness:**\n\n**1. More discussion is needed with MGDA.**\n\nThanks for the suggestion, we will make the difference and connection to MGDA more clear in the final version of the paper. Note that we have provided some comparisons in Appendix A already.\n\n**2. Add additional experimental details.**\n\nThanks for the suggestion and we will add more details in the appendix. Moreover, we will release all code for reproducing all experiments, to make it easy for others to build on our research.\n\n**3.. Hyperparameters optimization.**\n\nWe follow the same criterion as in CAGrad: we use the training loss to determine the hyperparameter.\n\n**4. Sensitivity to hyperparameters.**\n\nWe admit that to achieve the best performance of FAMO currently it requires some tuning of the $\\gamma$ parameter. But FAMO's performance remains strong across a relatively wide range of $\\gamma$s. Moreover, empirically we observe that setting it between $0.001$ to $0.01$ achieves good MTL performance. As stated in the conclusion and limitation section (Sec. 6), an interesting future direction is to make FAMO robust.\n\n**5. It will be good to include figures showing that losses are decreasing at roughly equal rates.**\n\nThanks for the suggestion, we provide it in Figure 2 above. We are showing the descent rate over the epoch for both FAMO and NashMTL. We observe that the rates for different losses are more aligned in FAMO than in NashMTL.\n\n**6. Provide a task weight throughout the optimization process.**\n\nThanks for the suggestion, we have provided a task weighting plot in Figure 3 above. From the plot, we see that both FAMO and NashMTL essentially gradually increase the task weight on task 3 (the surface normal prediction), which is the task that has the minimum loss scale.\n\n---\n\n**Questions:**\n\n**1. More intuitions on why FAMO mitigates CG.**\n\nMitigating CG is one objective, but the method should also optimize the task losses. Regarding why FAMO mitigates CG, it is because CG happens when a certain task is under-optimized due to an imbalance weighting among tasks. FAMO’s objective is to encourage an equal descent rate for all task objectives. Hence if CG happens, then definitely this objective is violated as well. Furthermore, while MGDA also mitigates CG, it encourages equal descent, which may not make sense when one of the task losses is extremely small (hence all other tasks' progress is bounded by this single task). Therefore, we believe FAMO’s objective is more reasonable in practice.\n\n**2. Have you tried multiple steps for approximating Eq.7?**\n\nThanks for the advice. We have not yet tried multiple steps for approximating Eq.7. In principle it is doable. But performing multiple steps approximation will no longer only require 1 extra forward pass. It will introduce more overhead and slow down FAMO. It is an interesting future direction to investigate the trade-off between speed and performance.\n"}]

Author rebuttal:
## Common Response to All Reviewers (with Additional Results)

---

We sincerely thank all reviewers for their comments and valuable suggestions. Per the reviewers' request, **we conduct additional experiments on NYU-v2 and summarize the results in the attached PDF. We will respond to each reviewer's questions and concerns separately via individual comments.**

---

**A. Additional Results**

**1.** MTL performance of FAMO, compared against **log-MGDA** (the exact form of FAMO), **amortized MGDA** (using FAMO's idea but applied to MGDA), over **epochs** of training (Figure 1).

**Findings:**
- Amortized MGDA performs poorly, as it pursues equal descent instead of equal rate descent.
- log-MGDA performs well and outperforms Nash-MTL at the end of training. But FAMO tracks log-MGDA's performance closely while being much faster (log-MGDA has the same computation efficiency as MGDA).
- FAMO not only is an efficient algorithm that achieves good asymptotic performance but also achieves good performance fast over epochs of training.

**2.** We provide the loss decrease rate $\frac{L^i_t - L^i_{t+1}}{L^i_{t}}$ for LS, NashMTL, and FAMO over epochs (Figure 2).

**Findings:**
- Among the three methods, we only observe that FAMO's losses decrease at a roughly equal rate. The decrease rate in NashMTL is also roughly even but not as good as in FAMO. LS has the most imbalanced loss decrease rate.
- Note that there is a bump in the middle of training caused by the learning rate scheduling (i.e., we halved the learning rate at the 100-th epoch)

**3.** We provide the task weights for the three task objectives over training steps for LS, NashMTL and FAMO (Figure 3).

**Findings:**
- NashMTL, due to its formulation, has task weights not necessarily sum to 1.
- Both NashMTL and FAMO gradually assign a higher task weighting to Task 3 (surface normal prediction), which is the task that has the minimum scale. By contrast, LS always assigns equal task weight to all three tasks.

---

**B. FAMO's contribution**

Some reviewers have concerns about FAMO's contribution. We want to emphasize 2 contributions of FAMO:
- We propose that the equal rate objective is more reasonable than the equal descent objective. And we empirically verify this claim.
- We identify a simple and efficient way of amortizing the log-MGDA. So far, only MGDA and log-MGDA, among existing MTL gradient manipulation methods, can be easily amortized by having a single extra forward pass of the network.

Combining the two contributions, the introduced FAMO algorithm achieves state-of-the-art MTL performance across popular MTL benchmarks, while having almost no overhead compared to the vanilla learning scalarization (LS) method. This will be particularly important if we go to hundreds or even a larger number of tasks, where existing gradient manipulation methods would almost certainly fail to apply due to their expensive computation.

The implementation of FAMO, together with all code to reproduce all experiment results, will be made available public upon acceptance.

Meta-review: This paper presents Fast Adaptive Multitask Optimization (FAMO), an optimizer for multi-task learning problems that aims to mitigate the conflicting gradient problem with decreased space and time requirements by iteratively updating task weightings while ensuring that all task losses decrease at roughly similar rates. The reviewers appreciated the clear empirical benefits of the proposed approach, as validated through extensive experimental evaluation. However, they also raised several substantive issues in their initial reviews, including:
* lack of assessment of the impact of the paper's two contributions separately (maximizing the minimum task rate of improvement vs. optimizing the task weights slowly via gradient updates). The proposed fix is to compare to log-MGDA in all experiments, and provide relevant discussions about the results.
* incorrectness of the proof of Proposition 3.1, notably since FAMO only achieves equal rate of task loss decrease under particular restrictive assumptions. The proposed fix is to adapt the message of Proposition 3.1 to clarify this, and remove the claim that FAMO always ensures an equal rate (rather than the best possible worst-case loss decrease rate).
* need for better discussion of the relationship between FAMO and MGDA. The proposed fix is to add more discussion of this in the paper.

The reviewers were satisfied by the authors' clarification on these concerns, under the assumption that the authors will make the above changes to the paper in the final version.
Decision: Accept (poster)