SlowGuess/NeurIPS_2023_OCR huggingface dataset:
†Bo Liu, ‡Yihao Feng, †,§ Peter Stone, †Qiang Liu
†The University of Texas at Austin, ‡Salesforce AI Research, §Sony AI {bliu, pstone, lqiang}@cs.utexas.edu, yihaof@salesforce.com
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 https://github.com/Cranial-XIX/FAMO.
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 the substantial size of these models, it is crucial to design MTL methods that are effective in terms of 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 to models with poor multitask performance: a subset of tasks are severely under-optimized. A major reason behind such optimization failure is that a 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 2).
To mitigate this problem, gradient manipulation methods [43, 25, 7, 24] 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 is 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 [41]. This is because they require computing and storing all task gradients at each iteration, thus demanding  $\mathcal{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  $\mathcal{O}(1)$  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, we refer to the situation where a single data  $x$  can be used to compute all task losses.
Avg Loss Landscape
Figure 1: Top left: The loss landscape, and individual task losses of a toy 2-task learning problem ( $\star$  represents the minimum of task losses). Top right: the runtime of different MTL methods for 50000 steps. Bottom: the loss trajectories of different MTL methods. ADAM fails in 1 out of 5 runs to reach the Pareto front due to CG. FAMO decreases task losses in a balanced way and is the only method matching the  $\mathcal{O}(1)$  space/time complexity of ADAM. Experimental details and analysis are provided in Section 5.1.
(Q) Is it possible to design a multi-task learning optimizer that ensures a balanced reduction in losses across all tasks while utilizing  $\mathcal{O}(1)$  space and time per iteration?
In this work, we present Fast Adaptive Multitask Optimization (FAMO), a 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. On 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  $\mathcal{O}(1)$  space and time per iteration.
2. We demonstrate that FAMO performs comparably to or better than existing gradient manipulation methods on a wide range of standard MTL benchmarks, in terms of standard MTL metrics, while being significantly computationally cheaper.
In this section, we provide the formal definition of multitask learning, then discuss its optimization challenge, and provide a brief overview of the gradient manipulation methods.
Multitask Learning (MTL) MTL considers optimizing a single model with parameter  $\theta \in \mathbb{R}^m$  that can perform  $k \geq 2$  tasks well, where each task is associated with a loss function  $\ell_i(\theta): \mathbb{R}^m \to \mathbb{R}_{\geq 0}$ . Then, it is common to optimize the average loss across all tasks:
$$
\min  _ {\theta \in \mathbb {R} ^ {m}} \left\{\ell_ {0} (\theta) := \frac {1}{k} \sum_ {i = 1} ^ {k} \ell_ {i} (\theta) \right\}. \tag {1}
$$
Optimization Challenge Directly optimizing (1) can result in severe under-optimization of 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  $t$ , assume one updates the model
2In this work, we assume  $\forall i,\ell_i(\theta)\geq 0$  , which is true for typical loss functions including mean square and cross-entropy losses. Note that one can always transform  $\ell_{i}$  to be non-negative if a loss lower bound is known.
Algorithm 1 Fast Adaptive Multitask Optimization (FAMO)
1: Input: Initial parameter  $\theta_0$ , task losses  $\{\ell_i\}_{i=1}^k$  (ensure that  $\ell_i \geq \epsilon > 0$ , for instance, by  $\ell_i \gets \ell_i - \ell_i^* + \epsilon$ ,  $\ell_i^* = \inf_{\theta} \ell_i(\theta)$ ), learning rate  $\alpha$  and  $\beta$ , and decay  $\gamma$  ( $= 0.001$  by default).
2:  $\xi_{1} \gets 0$ . // initialize the task logits to all zeros
3: for  $t = 1:T$  do
4: Compute  $z_{t} = \operatorname{Softmax}(\xi_{t})$ , e.g.,
$$
z _ {i, t} = \frac {\exp (\xi_ {i , t})}{\sum_ {i = 1} ^ {k} \exp (\xi_ {i , t})}.
$$
5: Update the model parameters:
$$
\theta_ {t + 1} = \theta_ {t} - \alpha \sum_ {i = 1} ^ {k} \left(c _ {t} \frac {\tilde {z} _ {i , t}}{\ell_ {i , t}}\right) \nabla \ell_ {i, t}, \quad \text {w h e r e} c _ {t} = \left(\sum_ {i = 1} ^ {k} \frac {\tilde {z} _ {i , t}}{\ell_ {i , t}}\right) ^ {- 1}.
$$
6: Update the logits for task weighting:
$$
\xi_ {t + 1} = \xi_ {t} - \beta \Big (\delta_ {t} + \gamma \xi_ {t} \Big) \text {w h e r e} \delta_ {t} = \left[ \begin{array}{c} \nabla^ {\top} z _ {1, t} (\xi_ {t}) \\ \vdots \\ \nabla^ {\top} z _ {k, t} (\xi_ {t}) \end{array} \right] ^ {\top} \left[ \begin{array}{c} \log \ell_ {1, t} - \log \ell_ {1, t + 1} \\ \vdots \\ \log \ell_ {k, t} - \log \ell_ {k, t + 1}. \end{array} \right].
$$
parameter using a gradient descent style iterative update:  $\theta_{t + 1} = \theta_t - \alpha d_t$  where  $\alpha$  is the step size and  $d_{t}$  is the update at time  $t$ . Then, we say that conflicting gradients (CG) [24, 43] happens if
$$
\exists i, \quad \ell_ {i} (\theta_ {t + 1}) - \ell_ {i} (\theta_ {t}) \approx - \alpha \nabla \ell_ {i} (\theta_ {t}) ^ {\top} d _ {t} > 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 a subset of tasks persistently undergoes under-optimization due to CG. In a more general sense, it is not desirable if 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 to mitigate it at each optimization stage [43].
Gradient Manipulation Methods Gradient manipulation methods aim to decrease all task losses in a more balanced way by finding a new update  $d_{t}$  at each step.  $d_{t}$  is usually a convex combination of task gradients, and therefore the name gradient manipulation (denote  $\nabla \ell_{i,t} = \nabla_{\theta}\ell_{i}(\theta_{t})$  for short):
$$
d _ {t} = \left[ \begin{array}{c} \nabla \ell_ {1, t} ^ {\top} \\ \vdots \\ \nabla \ell_ {k, t} ^ {\top} \end{array} \right] ^ {\top} w _ {t}, \quad \text {w h e r e} w _ {t} = \left[ \begin{array}{c} w _ {1, t} \\ \vdots \\ w _ {k, t} \end{array} \right] = f \left(\nabla \ell_ {1, t}, \dots , \nabla \ell_ {k, t}\right) \in \mathbb {S} _ {k}. \tag {2}
$$
Here,  $\mathbb{S}_k = \{w\in \mathbb{R}_{\geq 0}^k\mid w^\top \mathbf{1} = 1\}$  is the probabilistic simplex, and  $w_{t}$  is the task weighting across all tasks. Please refer to Appendix A for details of five state-of-the-art gradient manipulation methods (MGDA, PCGRAD, CAGRAD, IMTL-G, NASHMTL) and their corresponding  $f$ . Note that existing 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. As a result, we say these methods require at least  $\mathcal{O}(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  $t$ , assume we perform the update  $\theta_{t + 1} = \theta_t - \alpha d_t$ , we define the rate of improvement for task  $i$  as
$$
r _ {i} (\alpha , d _ {t}) = \frac {\ell_ {i , t} - \ell_ {i , t + 1}}{\ell_ {i , t}}. ^ {3} \tag {3}
$$
FAMO then seeks an update  $d_t$  that results in the largest worst-case improvement rate across all tasks  $\left(\frac{1}{2} \|d_t\|\right)$  is subtracted to prevent an under-specified optimization problem where the objective can be infinitely large):
$$
\max  _ {d _ {t} \in \mathbb {R} ^ {m}} \min  _ {i \in [ k ]} \frac {1}{\alpha} r _ {i} (\alpha , d _ {t}) - \frac {1}{2} \| d _ {t} \| ^ {2}. \tag {4}
$$
When the step size  $\alpha$  is small, using Taylor approximation, the problem (4) can be approximated by
$$
\max  _ {d _ {t} \in \mathbb {R} ^ {m}} \min  _ {i \in [ K ]} \frac {\nabla \ell_ {i , t} ^ {\top} d _ {t}}{\ell_ {i , t}} - \frac {1}{2} \| d _ {t} \| ^ {2} = (\nabla \log \ell_ {i, t}) ^ {\top} d _ {t} - \frac {1}{2} \| d _ {t} \| ^ {2}. \tag {5}
$$
Instead of solving the primal problem in (5) where  $d \in \mathbb{R}^m$  ( $m$  can be millions if  $\theta$  is the parameter of a neural network), we consider its dual problem:
Proposition 3.1. The dual objective of (5) is
$$
z _ {t} ^ {*} \in \underset {z \in \mathbb {S} _ {k}} {\arg \min } \frac {1}{2} \| J _ {t} z \| ^ {2}, \quad \text {w h e r e} \quad J _ {t} = \left[ \begin{array}{c} \nabla \log \ell_ {1, t} ^ {\top} \\ \vdots \\ \nabla \log \ell_ {k, t} ^ {\top} \end{array} \right], \tag {6}
$$
where  $z_{t}^{*} = [z_{t,i}^{*}]$  is the optimal combination weights of the gradients, and the optimal update direction is  $d_{t}^{*} = J_{t}z_{t}^{*}$ .
$$
\begin{array}{l} \max  _ {d \in \mathbb {R} ^ {m}} \min  _ {i \in [ k ]} \left(\nabla \log \ell_ {i, t}\right) ^ {\top} d - \frac {1}{2} \| d \| ^ {2} \\ = \max  _ {d \in \mathbb {R} ^ {m}} \min  _ {z \in \mathbb {S} _ {k}} \left(\sum_ {i = 1} ^ {k} z _ {i} \nabla \log \ell_ {i, t}\right) ^ {\top} d - \frac {1}{2} \| d \| ^ {2} \\ = \min  _ {z \in \mathbb {S} _ {k}} \max  _ {d \in \mathbb {R} ^ {m}} \left(\sum_ {i = 1} ^ {k} z _ {i} \nabla \log \ell_ {i, t}\right) ^ {\top} d - \frac {1}{2} \| d \| ^ {2} \quad (\text {s t r o n g d u a l i t y}) \\ \end{array}
$$
Write  $g(d,z) = \left(\sum_{i = 1}^{k}z_{i}\nabla \log \ell_{i,t}\right)^{\top}d - \frac{1}{2}\| d\|^{2}$ , then by setting
$$
\frac {\partial g}{\partial d} = 0 \quad \Longrightarrow \quad d ^ {*} = \sum_ {i = 1} ^ {k} z _ {i} \nabla \log \ell_ {i, t}.
$$
Plugging in  $d^{*}$  back, we have
$$
\max  _ {d \in \mathbb {R} ^ {m}} \min  _ {i \in [ k ]} \left(\nabla \log \ell_ {i, t}\right) ^ {\top} d - \frac {1}{2} \| d \| ^ {2} = \min  _ {z \in \mathbb {S} _ {k}} \frac {1}{2} \left\| \sum_ {i = 1} ^ {k} z _ {i} \nabla \log \ell_ {i, t} \right\| ^ {2} = \min  _ {z \in \mathbb {S} _ {k}} \frac {1}{2} \| J _ {t} z \| ^ {2}.
$$
At the optimum, we have  $d_t^* = J_t z_t^*$ .
The dual problem in (6) can be viewed as optimizing the log objective of the multiple gradient descent algorithm (MGDA) [9, 35]. Similar to MGDA, (6) only involves a decision variable of dimension  $k \ll m$ . Furthermore, if the optimal combination weights  $z_{t}^{*}$  is an interior point of  $\mathbb{S}_k$ , then the improvement rates  $r_i(\alpha, d_t^*)$  of the different tasks  $i$  equal, as we show in the following result.
3To avoid division by zero, in practice we add a small constant (e.g.,  $1e - 8$ ) to all losses. For the ease of notation (e.g.,  $\ell_i(\cdot)\gets \ell_i(\cdot) + 1e - 8$ , we omit it throughout the paper.
Proposition 3.2. Assume  $\{\ell_i\}_{i = 1}^k$  are smooth and the optimal weights  $z_{t}^{*}$  in (6) is an interior point of  $\mathbb{S}_k$ , then
$$
\forall i \neq j \in [ k ], \quad r _ {i} ^ {*} \left(d _ {t} ^ {*}\right) = r _ {j} ^ {*} \left(d _ {t} ^ {*}\right),
$$
where  $r_i^* (d_t^*) = \lim_{\alpha \to 0}\frac{1}{\alpha} r_i(\alpha ,d_t^*)$
Proof. Consider the Lagrangian form of (6)
$$
\mathcal {L} (z, \lambda , \mu) = \frac {1}{2} \left\| \sum_ {i = 1} ^ {k} z _ {i} \nabla \log \ell_ {i, t} \right\| ^ {2} + \lambda \left(\sum_ {i = 1} ^ {k} z _ {i} - 1\right) - \sum_ {i = 1} ^ {k} \mu_ {i} z _ {i}, \text {w h e r e} \forall i, \mu_ {i} \geq 0. \tag {7}
$$
When  $z^{*}$  reaches the optimum, we have  $\partial \mathcal{L}(z,\lambda ,\mu) / \partial z = 0$  , recall that  $d_t^* = J_tz_t^*$  , then
$$
J _ {t} ^ {\top} J _ {t} z ^ {*} = - \mu - \lambda , \quad \text {w h e r e} \quad J _ {t} = \left[ \begin{array}{c} \nabla \log \ell_ {1, t} ^ {\top} \\ \vdots \\ \nabla \log \ell_ {k, t} ^ {\top} \end{array} \right] \quad \Longrightarrow \quad J _ {t} ^ {\top} d _ {t} ^ {*} = - (\mu + \lambda).
$$
When  $z_{t}^{*}$  is an interior point of  $\mathbb{S}_k$ , we know that  $\mu = 0$ . Hence  $J_{t}^{\top}d_{t}^{*} = -\lambda$ . This means,
$$
\forall i \neq j, \quad \lim  _ {\alpha \to 0} \frac {1}{\alpha} r _ {i} (\alpha , d _ {t} ^ {*}) = \nabla \log \ell_ {i, t} ^ {\top} d _ {t} ^ {*} = \nabla \log \ell_ {j, t} ^ {\top} d _ {t} ^ {*} = \lim  _ {\alpha \to 0} \frac {1}{\alpha} r _ {j} (\alpha , d _ {t} ^ {*}).
$$
3.2 Fast Approximation by Amortizing over Time
Instead of fully solving (6) at each optimization step, FAMO performs a single-step gradient descent on  $z$ , which amortizes the computation over the optimization trajectory:
$$
z _ {t + 1} = z _ {t} - \alpha_ {z} \tilde {\delta}, \quad \text {w h e r e} \quad \tilde {\delta} = \nabla_ {z} \frac {1}{2} \left\| \sum_ {i = 1} ^ {k} z _ {i, t} \nabla \log \ell_ {i, t} \right\| ^ {2} = J _ {t} ^ {\top} J _ {t} z _ {t}. \tag {8}
$$
But then, note that
$$
\frac {1}{\alpha} \left[ \begin{array}{c} \log \ell_ {1, t} - \log \ell_ {1, t + 1} \\ \vdots \\ \log \ell_ {k, t} - \log \ell_ {k, t + 1} \end{array} \right] \approx J _ {t} ^ {\top} d _ {t} = J _ {t} ^ {\top} J _ {t} z _ {t}, \tag {9}
$$
so we can use the change in log losses to approximate the gradient.
In practice, to ensure that  $z$  always stays in  $\mathbb{S}_k$ , we re-parameterize  $z$  by  $\xi$  and let  $z_{t} = \operatorname{Softmax}(\xi_{t})$ , where  $\xi_{t}\in \mathbb{R}^{K}$  are the unconstrained softmax logits. Consequently, we have the following approximate update on  $\xi$  from (8):
$$
\xi_ {t + 1} = \xi_ {t} - \beta \delta , \quad \text {w h e r e} \delta = \left[ \begin{array}{c} \nabla^ {\top} z _ {1, t} (\xi) \\ \vdots \\ \nabla^ {\top} z _ {k, t} (\xi) \end{array} \right] ^ {\top} \left[ \begin{array}{c} \log \ell_ {1, t} - \log \ell_ {1, t + 1} \\ \vdots \\ \log \ell_ {k, t} - \log \ell_ {k, t + 1} \end{array} \right]. \tag {10}
$$
Remark: While it is possible to perform gradient descent on  $z$  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 in losses.
3.3 Practical Implementation
To facilitate practical implementation, we present two modifications to the update in (10).
Re-normalization The suggested update above is a convex combination of the gradients of the log loss, e.g.,
$$
d ^ {*} = \sum_ {i = 1} ^ {k} z _ {i, t} \nabla \log \ell_ {i, t} = \sum_ {i = 1} ^ {k} \left(\frac {z _ {i , t}}{\ell_ {i , t}}\right) \nabla \ell_ {i, t}.
$$
When  $\ell_{i,t}$  is small, the multiplicative coefficient  $\frac{z_{i,t}}{\ell_{i,t}}$  can be quite large and result in unstable optimization. Therefore, we propose to multiply  $d^{*}$  by a constant  $c_{t}$ , such that  $c_{t}d^{*}$  can be written as a convex combination of the task gradients just as in other gradient manipulation algorithms (see (2) and we provide the corresponding definition of  $w$  in the following):
$$
c _ {t} = \left(\sum_ {i = 1} ^ {k} \frac {z _ {i , t}}{\ell_ {i , t}}\right) ^ {- 1} \quad \text {a n d} \quad d _ {t} = c _ {t} d ^ {*} = \sum_ {i = 1} ^ {k} w _ {i} \nabla \ell_ {i, t}, \quad \text {w h e r e} w _ {i} = c _ {t} \frac {z _ {i , t}}{\ell_ {i , t}}. \tag {11}
$$
Regularization As we are amortizing the computation over time and the loss objective  $\{\ell_i(\cdot)\} s$  are changing dynamically, it makes sense to focus more on the recent updates of  $\xi$  [46]. To this end, we put a decay term on  $w$  such that the resulting  $\xi_{t}$  is an exponential moving average of its gradient updates:
$$
\xi_ {t + 1} = \xi_ {t} - \beta (\delta_ {t} + \gamma \xi_ {t}) = - \beta (\delta_ {t} + (1 - \beta \gamma) \delta_ {t - 1} + (1 - \beta \gamma) ^ {2} \delta_ {t - 2} + \dots). \tag {12}
$$
We provide the complete FAMO algorithm in Algorithm 1 and its pseudocode in Appendix C.
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  $(\alpha, \beta)$  shrink to zero). From Algorithm 1, one can derive the following non-autonomous dynamical system (assuming  $\{\ell_t\}$  are all smooth):
$$
\left[ \begin{array}{c} \dot {\theta} \\ \dot {\xi} \end{array} \right] = - c _ {t} \left[ \begin{array}{c} J _ {t} z _ {t} \\ A _ {t} J _ {t} ^ {\top} J _ {t} z _ {t} + \frac {\gamma}{c _ {t}} \xi_ {t} \end{array} \right], \text {w h e r e} A _ {t} = \left[ \begin{array}{c} \nabla^ {\top} z _ {1, t} (\xi_ {t}) \\ \vdots \\ \nabla^ {\top} z _ {k, t} (\xi_ {t}) \end{array} \right]. \tag {13}
$$
(13) reaches its stationary points (or fixed points) when (note that  $c_{t} > 0$ )
$$
\left[ \begin{array}{l} \dot {\theta} \\ \dot {\xi} \end{array} \right] = 0 \quad \Longrightarrow J _ {t} z _ {t} = 0 \text {a n d} \xi_ {t} = 0 \quad \Longrightarrow \sum_ {i = 1} ^ {k} \nabla \log \ell_ {i, t} = 0. \tag {14}
$$
Therefore, the minimum points of  $\sum_{i=1}^{k} \log \ell_i(\theta)$  are all stationary points of (13).
In this section, we summarize existing methods that tackle learning challenges in multitask learning (MTL). The general idea of most existing works is to 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, 11].
Multitask Architecture Novel neural architectures for MTL include hard-parameter-sharing methods, which decompose a neural network into task-specific modules and a shared feature extractor using manually designed heuristics [21, 29, 2], and soft-parameter-sharing methods, which learn which parameters to share [30, 34, 12, 27]. Recent studies extend neural architecture search for MTL by learning where to branch a network to have task-specific modules [14, 3].
Multitask Optimization The most relevant approach to our method is MTL optimization via task balancing. These methods dynamically re-weight all task losses to mitigate the conflicting gradient issue [40, 43]. The simplest form of gradient manipulation is to re-weight the task losses based on manually designed criteria [6, 13, 18], but these methods are often heuristic and lack theoretical support. Gradient manipulation methods [35, 43, 25, 7, 16, 24, 32, 26, 47] propose to form a 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, it 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 of 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 is also computationally efficient. However, we will demonstrate empirically that FAMO performs better than this method.
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 on standard benchmarks (e.g., supervised and reinforcement MTL problems)?
In the following, we first use a toy 2-task problem to demonstrate how FAMO mitigates CG while being efficient. Then we show that FAMO performs comparably or even better than state-of-the-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  $\gamma$ . Each subsection first details the experimental setup and then analyzes the results.
5.1 A Toy 2-Task Example
Figure 2: The average loss  $L^0$  and the two task losses  $L^1$  and  $L^2$  for the toy example.
To better understand the optimization trajectory of FAMO, we adopt the same 2D multitask optimization problem from NASHMTL [32] to visualize how FAMO balances different loss objectives. The model parameter  $\theta = (\theta_{1},\theta_{2})\in \mathbb{R}^{2}$ . The two tasks' objectives and their surface plots are provided in Appendix D and Figure 2. We compare FAMO against ADAM [19], MGDA [35], PCGRAD [43], CAGRAD [24], and NASHMTL [32]. We then pick 5 initial points  $\theta_{\mathrm{init}}\in \{(-8.5,7.5),(-8.5,5),(0,0),(9,9),(10,-8)\}$  and plot the corresponding optimization trajectories with different methods in Figure 1. 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.