SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Findings: From Figure 1, 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 NASHMTL and achieves a 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 to ADAM in terms of the training time, which is  $25 \times$  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 [31] (3 tasks), CityScapes [8] (2 tasks), QM-9 [1] (11 tasks), and CelebA [28] (40 tasks). Specifically, NYU-v2 is an 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 on any scene image. CityScapes dataset is similar to NYU-v2 but contains 5000 street-view RGBD images with per-pixel annotations. QM-9 dataset is a widely used benchmark in graph neural network learning. It consists of  $>130\mathrm{K}$  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 110K 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 a 40-task MTL problem. Different from NYU-v2, CityScapes, and QM-9, the number of tasks  $(K)$  in CelebA is much larger, hence posing a challenge to learning efficiency.
Method	Segmentation	Depth	Surface Normal	MR↓	Δm%↓
mIoU ↑	Pix Acc ↑	Abs Err ↓	Rel Err ↓	Angle Dist ↓	Within t° ↑
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
LS	39.29	65.33	0.5493	0.2263	28.15	23.96	22.09	47.50	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.20	53.13	66.24	3.78	-0.76
NASHMTL	40.13	65.93	0.5261	0.2171	25.26	20.08	28.40	55.47	68.15	2.11	-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 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  $\Delta m\%$  are the main metrics for MTL performance.
Method	μ	α	εHOMO	εLUMO	(R2)	ZPVE	U0	U	H	G	cv	MR↓	Δm%↓
MAE ↓
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	157.1	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	1.07	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  $\Delta m\%$  are the main metrics for MTL performance.
We compare FAMO against 11 MTL optimization methods and a single-task learning baseline: (1) Single task learning (STL), training an independent model ( $\theta$  for each task; (2) Linear scalarization (LS) baseline that minimizes  $L^0$ ; (3) Scale-invariant (SI) baseline that minimizes  $\sum_{k} \log L^{k}(\theta)$ , as SI is invariant to any scalar multiplication of task losses; (4) Dynamic Weight Average (DWA) [27], a heuristic for adjusting task weights based on rates of loss changes; (5) Uncertainty Weighting (UW) [18] uses task uncertainty as a proxy to adjust task weights; (6) Random Loss Weighting (RLW) [23] that samples task weighting whose log-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 to the normal plan of that of other tasks and combining them together in the end; (9) CAGRAD [24] 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 [32] formulates MTL as a bargaining game and finds the solution to the game that benefits all tasks. For FAMO, we choose the best hyperparameter  $\gamma \in \{0.0001, 0.001, 0.01\}$  based on the validation loss. Specifically, we choose  $\gamma$  equals 0.01 for the CityScapes dataset and 0.001 for the rest of the datasets. See Appendix E for results with error bars.
**Evaluations:** We consider two metrics [32] for MTL: 1)  $\Delta m\%$ , the average per-task performance drop of a method  $m$  relative to the STL baseline denoted as  $b$ :  $\Delta m\% = \frac{1}{K}\sum_{k=1}^{K}(-1)^{\delta_k}(M_{m,k} - M_{b,k}) / M_{b,k} \times 100$ , where  $M_{b,k}$  and  $M_{m,k}$  are the STL and  $m$ 's value for metric  $M_k$ .  $\delta_k = 1$  (or 0) if the  $M_k$  is higher (or lower) the better. 2) **Mean Rank (MR)**: the average rank of each method across tasks. For instance, if a method ranks first for every task, MR will be 1.
Findings: Results on the four benchmark datasets are provided in Table 1, 2 and 3. We observe that FAMO performs consistently well across different supervised learning MTL benchmarks compared
Method	CityScapes	CelebA
Segmentation	Depth	MR ↓	Δm% ↓	MR ↓	Δm% ↓
mIoU ↑	Pix Acc ↑	Abs Err ↓	Rel Err ↓
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	5.89	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
to other gradient manipulation methods. In particular, it achieves state-of-the-art results in terms of  $\Delta m\%$  on the NYU-v2 and QM-9 datasets.
Multitask Reinforcement Learning. We further apply FAMO to multitask reinforcement learning (MTRL) problems as MTRL often suffers more from conflicting gradients due to the stochastic nature of reinforcement learning [43]. Following CAGRAD [24], we apply FAMO on the MetaWorld [44] MT10 benchmark, which consists of 10 robot manipulation tasks with different reward functions. Following [37], we use Soft Actor-Critic (SAC) [15] as the underlying RL algorithm, and compare against baseline methods including LS (SAC with a shared model) [44], Soft Modularization [42] (an MTL network that routes different modules in a shared model to form different policies), PCGRAD [43], 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 in Table 5.2.
Figure 3: Training Success Rate and Time.
Table 3: Results on CityScapes (2 tasks) and CelebA (40 tasks) datasets. Each experiment is repeated over 3 random seeds and the mean is reported. The best average result is marked in bold. MR and  $\Delta m\%$  are the main metrics for MTL performance.
Method	Success ↑ (mean ± stderr)
LS (lower bound)	0.49 ±0.07
STL (proxy for upper bound)	0.90 ±0.03
PCGRAD [43]	0.72 ±0.02
SOFT MODULARIZATION [42]	0.73 ±0.04
CAGRAD	0.83 ±0.05
NASHMTL [32] (every 1)	0.91 ±0.03
NASHMTL [32] (every 50)	0.85 ±0.02
NASHMTL [32] (every 100)	0.87 ±0.03
NASHMTL (ours) (every 1)	0.80 ±0.13
NASHMTL (ours) (every 50)	0.76 ±0.10
NASHMTL (ours) (every 100)	0.80 ±0.12
UW [18]	0.77 ±0.05
FAMO (ours)	0.83 ±0.05
Table 4: MTRL results (averaged over 10 runs) on the Metaworld-10 benchmark.
<sup>4</sup>We could not reproduce the MTRL results of NASHMTL exactly, so we report both the results from the original paper and our reproduced results.
Findings: From Table 5.2, we observe that FAMO performs comparably to CAGRAD and outperforms PCGRAD and the average gradient descent baselines by a large margin. FAMO also outperforms NASHMTL based on our implementation. Moreover, FAMO is significantly faster than NASHMTL, even when it is applied once every 100 steps.
5.3 MTL Efficiency (Training Time Comparison)
Figure 4 provides the FAMO's average training time per epoch against that of the baseline methods.
Figure 4: Average training time per epoch for different MTL optimization methods. We report the relative training time of a method to that of the linear scalarization (LS) 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 at least  $\mathcal{O}(K)$  space to store the task gradients, which is implausible for large models in the many-task setting (i.e., when  $m = |\theta|$  and  $K$  are large).
5.4 Ablation on  $\gamma$
In this section, we provide the ablation study on the regularization coefficient  $\gamma$  in Figure 5.
Figure 5: Ablation over  $\gamma$ : we plot the performance of FAMO (in terms of  $\Delta m\%$  using different values of  $\gamma$  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  $\gamma$ s. The problem with CityScapes is that one of the task losses is close to 0 at the very beginning, hence small changes in task weighting can result in very different loss improvement. Therefore we conjecture that using a larger  $\gamma$ , in this case, can help stabilize MTL.
In this work, we introduce FAMO, a fast optimization method for multitask learning (MTL) that mitigates the conflicting gradients using  $\mathcal{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  $\gamma$ , which is introduced due to the stochastic update of the task weighting logits  $w$ . Future work can investigate a more principled way of determining  $\gamma$ .
[36] Jiayi Shen, Xiantong Zhen, Marcel Worring, and Ling Shao. Variational multi-task learning with gumbel-softmax priors. Advances in 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 on Machine Learning, pages 9120-9132. PMLR, 2020.
[39] Sebastian Thrun and Joseph O'Sullivan. Discovering structure in multiple learning tasks: The tc algorithm. In ICML, volume 96, pages 489-497, 1996.
[40] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc 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 First. 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.06782, 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 on 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. In Proceedings of the IEEE conference on 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.
A Gradient Manipulation Methods
In this section, we provide a brief overview of representative gradient manipulation methods in multitask/multiobjective 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 earliest gradient manipulation methods for multitask learning. In MGDA, the per step update  $d_t$  is found by solving
$$
\max _ {d \in \mathbb {R} ^ {m}} \min _ {i \in [ k ]} \nabla \ell_ {i, t} ^ {\top} d - \frac {1}{2} \| d \| ^ {2}.
$$
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 in (6) is similar to the MGDA objective in the sense that we can view optimizing (6) as optimizing the log of the task losses. Hence, when we compare FAMO against MGDA, one can regard FAMO as balancing the rate of loss improvement while MGDA balances the absolute improvement across task losses.
Projecting Gradient Descent (PCGRAD) [43] PCGRAD initializes  $v_{\mathrm{PC}}^i = \nabla \ell_{i,t}$ , then for each task  $i$ , PCGRAD loops over all task  $j \neq i$  (in a random order, which is crucial as mentioned in [43]) and removes the "conflict"
$$
v _ {\mathrm {P C}} ^ {i} \gets v _ {\mathrm {P C}} ^ {i} - \frac {v _ {\mathrm {P C}} ^ {i} {} ^ {\top} \nabla \ell_ {j , t}}{\| \ell_ {j , t} \| ^ {2}} \nabla \ell_ {j, t} \quad \mathrm {i f} \quad v _ {\mathrm {P C}} ^ {i} {} ^ {\top} \nabla \ell_ {j, t} <   0.
$$
In the end, PCGRAD produces  $d_{t} = \frac{1}{k}\sum_{i=1}^{k}v_{\mathrm{PC}}^{i}$ . 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 a first principle approach.
Conflict-averse Gradient Descent (CAGRAD) [24]  $d_{t}$  is found by solving
$$
\max  _ {d \in \mathbb {R} ^ {m}} \min  _ {i \in [ k ]} \nabla \ell_ {i, t} ^ {\top} d \quad \text {s . t .} \quad \| d - \nabla \ell_ {0, t} \| \leq c \| \nabla \ell_ {0, t} \|.
$$
Here,  $\ell_{0,t} = \frac{1}{k}\sum_{i=1}^{k}\ell_{i,t}$ . CAGRAD seeks an update  $d_t$  that optimizes the "worst improvement" as much as possible, conditioned on that the update still decreases the average loss. By controlling the hyperparameter  $c$ , CAGRAD can recover MGDA ( $c \to \infty$ ) and the vanilla averaged gradient descent ( $c \to 0$ ). Due to the extra constraint, CAGRAD provably converges to the stationary points of  $\ell_0$  when  $0 \leq c < 1$ .
Impartial Multi-Task Learning (IMTL-G) [25] IMTL-G finds  $d_{t}$  such that it shares the same cosine similarity with any task gradients:
$$
\forall i \neq j, d _ {t} ^ {\top} \frac {\nabla \ell_ {i , t}}{\| \nabla \ell_ {i , t} \|} = d _ {t} ^ {\top} \frac {\nabla \ell_ {j , t}}{\| \nabla \ell_ {j , t} \|}, \text {a n d} d _ {t} = \sum_ {i = 1} ^ {k} w _ {i, t} \nabla \ell_ {i, t}, \text {f o r s o m e} w _ {t} \in \mathbb {S} _ {k}.
$$
The constraint that  $d_{t} = \sum_{i=1}^{k} w_{i,t} \nabla \ell_{i,t}$  is for preventing the problem from being under-determined. From the above equation, we can see that IMTL-G ignores the "size" of each task gradient and only cares about the "direction". As a result, one can think of IMTL-G as a variant of MGDA that applies to the normalized gradients. By doing so, IMTL-G does not suffer from the straggler effect due to slow objectives. Furthermore, one can view IMTL-G as the equal angle descent, which is also proposed in Katrutsa et al. [17], where the objective is to find  $d$  such that
$$
\forall i \neq j, \quad \cos (d, \nabla \ell_ {i, t}) = \cos (d, \nabla \ell_ {j, t}).
$$
NASHMTL[32] NASHMTL finds  $d_{t}$  by solving a bargaining game treating the local improvement of each task loss as the utility for each task:
$$
\max  _ {d \in \mathbb {R} ^ {m}, \| d \| \leq 1} \sum_ {i = 1} ^ {k} \log \left(\nabla \ell_ {i, t} ^ {\top} d\right).
$$
Note that the objective of NASHMTL implicitly assumes that there exists  $d$  such that  $\forall i$ ,  $\nabla \ell_{i,t}^{\top}d > 0$  (otherwise we reach the Pareto front). It is easy to see that
$$
\max _ {\| d \| \leq 1} \sum_ {i = 1} ^ {k} \log \left(\nabla \ell_ {i, t} ^ {\top} d\right) = \max _ {\| d \| \leq 1} \sum_ {i = 1} ^ {k} \log \left\langle \frac {\nabla \ell_ {i , t}}{\| \nabla \ell_ {i , t} \|}, d \right\rangle = \max _ {\| d \| \leq 1} \sum_ {i = 1} ^ {k} \log \cos \left(\nabla \ell_ {i, t}, d\right).
$$
Therefore, due to the log, NASHMTL also ignores the "size" of task gradients and only cares about their "directions". Moreover, denote  $u_{i} = \frac{\nabla\ell_{i,t}}{\|\nabla\ell_{i,t}\|}$ . Then, according to the KKT condition, we know:
$$
\sum_ {i} \frac {u _ {i}}{u _ {i} ^ {\top} d} - \alpha d = 0, \quad \alpha \geq 0 \quad \Longrightarrow \quad d = \frac {1}{\alpha} \sum_ {i} \frac {1}{u _ {i} ^ {\top} d} u _ {i}.
$$
Consider when  $k = 2$ , if we take the equal angle descent direction:  $d_{\angle} = (u_1 + u_2) / 2$  (note that as  $u_1$  and  $u_2$  are normalized, their bisector is just their average). Then it is easy to check that
$$
d _ {\angle} = \frac {1}{\alpha} \left(\frac {2}{u _ {1} ^ {\top} (u _ {1} + u _ {2})} u _ {1} + \frac {2}{u _ {2} ^ {\top} (u _ {1} + u _ {2})} u _ {2}\right), \text {w h e r e} \alpha = \frac {u _ {1} ^ {\top} (u _ {1} + u _ {2})}{4} = \frac {u _ {2} ^ {\top} (u _ {1} + u _ {2})}{4}.
$$
As a result, we can see that when  $k = 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 all of these gradient manipulation methods require computing and storing  $K$  task gradients before applying  $f$  to compute  $d_{t}$ , which often involves solving an 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, we discuss such possibilities and point out the challenges.
Amortizing MGDA This is almost the same as in 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 is no straightforward way of bypassing the computation of task gradients.
Amortizing IMTL-G The task weighting in IMTL-G is computed by a series of matrix-matrix and matrix-vector products using task gradients [25]. Hence, it is also hard to amortize its computation over time.
Therefore, we focus on deriving the amortization for CAGRAD and NASHMTL.
Amortizing CAGRAD For CAGRAD, the dual objective is
$$
\min  _ {w \in \mathbb {S} _ {k}} F (w) = g _ {w} ^ {\top} g _ {0} + c \| g _ {w} \| \| g _ {0} \|, \tag {15}
$$
where  $g_0 = \nabla \ell_{0,t}$  and  $g_w = \sum_{i=1}^k w_i \nabla \ell_i$ . Denote
$$
G = \left[ \begin{array}{c} \nabla \ell_ {1, t} ^ {\top} \\ \vdots \\ \nabla \ell_ {k, t} ^ {\top} \end{array} \right].
$$
Now, if we take the gradient with respect to  $w$  in (15), we have:
$$
\frac {\partial F}{\partial w} = G ^ {\top} g _ {0} + c \frac {\left\| g _ {0} \right\|}{\left\| g _ {w} \right\|} G ^ {\top} g _ {w}. \tag {16}
$$
As a result, in order to approximate this gradient, one can separately estimate:
$$
G ^ {\top} g _ {0} \approx \frac {\ell (\theta) - \ell (\theta - \alpha g _ {0})}{\alpha}
$$
$$
G ^ {\top} g _ {w} \approx \frac {\ell (\theta) - \ell (\theta - \alpha g _ {w})}{\alpha}. \tag {17}
$$
$$
\left\| g _ {0} \right\| \approx \sqrt {1 ^ {\top} G ^ {\top} g _ {0}}
$$
$$
\left\| g _ {w} \right\| \approx \sqrt {w ^ {\top} G ^ {\top} g _ {w}}
$$
Once all these are estimated, one can combine them together to perform a single update on  $w$ . But note that this will require 3 forward and backward passes through the model, making it harder to implement in practice.
Amortizing NASHMTL Per derivation from NASHMTL [32], the objective is to solve for  $w$ :
$$
G ^ {\top} G w = 1 \oslash w. \tag {18}
$$
One can therefore form an objective:
$$
\min  _ {w} F (w) = \left\| G ^ {\top} G w - 1 \varnothing w \right\| _ {2} ^ {2}. \tag {19}
$$
Taking the derivative of  $F$  with respect to  $w$ , we have
$$
\frac {\partial F}{\partial w} = 2 G ^ {\top} G \left(G ^ {\top} g _ {w} - 1 \oslash w\right) + 2 \left(G ^ {\top} g _ {w} - 1 \oslash w\right) \oslash (w \odot w). \tag {20}
$$
Therefore, to approximate the gradient of  $w$ , one needs to first estimate
$$
G ^ {\top} g _ {w} \approx \frac {L (\theta) - L (\theta - \alpha g _ {w})}{\alpha} = \eta . \tag {21}
$$
Then we estimate
$$
G ^ {\top} G (\eta - 1 \oslash w) \approx \frac {L (\theta) - L (\theta - \alpha G (\eta - 1 \oslash w))}{\alpha}. \tag {22}
$$
Again, this results in 3 forward and backward passes through the model, let alone the overhead of resetting the model back to  $\theta$  (requires a copy of the original weights).