FAMO: Fast Adaptive Multitask Optimization

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.

Strengths

1. The paper proposes a novel approach for addressing a known pain point with (most) SoTA multi-task optimization methods. 2. 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. 3. The paper is well-written and well-structured. Connection to prior works in MTL literature is well established (Appendix A-B).

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. 2. Please provide additional details on the experimental sections (data split sizes, HPs, etc.). 3. Hyperparameter (HP) optimization: - 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? - Missing details on the tuning of the lr for the task weights. 4. From section 5.4 (and as stated by the authors), it appears the method is highly sensitive to the choice of HPs. 5. 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. 6. It would be beneficial if you could provide a figure of the task weights throughout the optimization process. Minor: Table 2, MR column missing indication for best-performing method.

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? 2. Have you tried performing several optimization steps for approximating the solution in Eq. 7 at each training iteration?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations addressed.

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. The core idea is to let the multiple losses decrease at roughly the same speed.

Strengths

The method is more efficient than previous gradient manipulation methods. Authors 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. However, 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. In 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? My 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. 2. The method part is not presented very clearly, see the questions.

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. 2. Could authors elaborate a bit on why Equation 7 is the dual form for (6)? 3. 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? 4. 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

No social impact

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.

Strengths

1. The paper is in general well-written, with a nice discussion about different multi-task gradient manipulation technique in Appendix A. 2. The proposed method FAMO empirically is competitive with other methods on multiple benchmarks.

Weaknesses

**[Lack of proof for significant contribution]** In 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. **[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. **[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.

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. **[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? **[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. **[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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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.

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. The proposed method, called FAMO, aims at improving the worst-case rate of improvement across all tasks at each step. At 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. The 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.

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. - 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. - Moreover, the experimental methodology seems sound; all experiments consider multiple random seeds and 2 sensible metrics to compare methods are considered. - 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. - 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. - [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.

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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

yes

Follow-up to authors' rebuttal (Part I)

I want to thank the reviewer’s rebuttal response and the effort in creating new figures to compare FAMO with non-amortized log-MGDA. I have the following comments: 1. **Understanding the two contributions separately** As also pointed out by Reviewer gjX4, the two contributions of the paper **[** **1)** using logarithm of task losses in MGDA) and **2)** using amortized optimization for fast optimization **]** are two separate ideas and deserve analyses separately. Previously, the main paper hasn’t analyzed these two points separately, but directly provide the results of the combination of these two ideas through FAMO. In the rebuttal, the authors have already made an improvement in providing the result of log-MGDA on the NYU-v2 dataset in Figure 1. This allows us to understand 1) and 2) separately. To further strengthen these analyses, I believe it is important to add the result of log-MGDA on every dataset in Section 5.2, and in addition, also through the time measurements in Section 5.3. This will give the readers the complete information about the delta of improvement in each of the two ideas separately. With these changes, the authors might also need to add additional sentences in **Findings** to describe these results. 2. **Proposition 3.1 is indeed incorrect** As the authors have agreed, without an additional assumption that the optimal $z$ is an interior point of $\mathcal{S}^k$, the conclusion of proposition 3.1 does not hold. However, I don’t think the authors should fix the proof by simply adding this assumption because this assumption is too restrictive: for example, when $g^1 \coloneqq \nabla \log \Delta^1 $ and $g^2 \coloneqq \nabla \log \Delta^2$ are pointing in exactly the same direction (thus no conflicting gradient), as long as $g^2$ is longer than $g^1$, the norm minimizing solution in the 1-simplex should put all of its weight on $g^1$ and none of its weight on $g^2$. In this case, the optimal $z$ **has to be on the corner (instead of the interior)** of the 1-simplex and there is **NO equal rate** of improvement across these two task losses. Based on this simple example, the authors should not claim FAMO (or log-MGDA) to be aiming at equal rate of task loss decrease (this is mentioned in line 41 as the first contribution). The authors claim that proposition 3.1 is the same result as that in the original paper except that the authors are working with log objectives. However, the original MGDA paper never claims they are optimizing for equal **amount** of decrease across task losses. Instead, the original MGDA paper only aims to find a solution on the pareto-frontier by considering multi-task learning as a multi-objective optimization problem. Thus the authors also shouldn’t claim log-MGDA or FAMO as optimizing for equal **rate** of decrease across task losses. In terms of high-level intuition, the authors’ original formulation of the objective in equation (5) is a min-max problem, where the goal is *not to ensure a notion of equality* among the tasks but only to improve the rate of improvement of the *worst performing task*. This intuition would also imply that getting equal rate of improvement to be unlikely under this min-max formulation.

Follow-up to authors' rebuttal (Part II)

3. **Why is FAMO $O(1)$ in space?** I acknowledge the authors’ intension to claim that FAMO only needs to store $1$ model gradients but not $K$ gradients. However, what I originally want to point out is that claiming FAMO (or averaged loss gradient descent) is $O(1)$ in space is a bit ambiguous. This is because the space requirement in reverse mode (backprop) gradient computation does **not only** depend on the size of parameters whose gradient is to be computed, but **more importantly** also on the entire computation graph these parameters are involved in. (This is because we need to store all the intermediate computations in memory for the backward gradient computation). In the case of FAMO and averaged loss gradient descent, a single model’s parameters are still involved in the computation of $K$ different losses. Thus the entire computation graph’s size would still scale with $O(K)$ instead of $O(1)$. I would recommend the authors make a distinction between the space used to compute the gradients (which unavoidably is $O(K)$ for all methods) and the space needed to store the necessary gradients and use it to compute the final update direction (FAMO is $O(1)$ in this regard). 4. **Further clarifying goals of multitask learning** I believe the typical narrative of multitask learning is to improve the performance of each task by leveraging the similarity across tasks. However, many experiments in the paper (QM-9, CityScapes, CelebA, and MRTL) all show that doing multitask learning cannot outperform single task learning in terms of performance. In contrast, as the authors have explained in the rebuttal response, the benefit is instead in model storage, where we sacrifice a bit of individual task performance (compared to STL) in order to avoid storing $O(K)$ models but only $1$ model. Maybe this is an implicit assumption common to researchers working with these benchmarks, but I would recommend the authors be more explicit about this in the main paper.

Questions to Rebuttal by Authors

Thanks for Authors' reply! I feel my original rating is not high enough. Just a few more questions to confirm my understanding: (1) In equation 10, the gradient of w is basically the gradient of z multiply with the derivative of softmax(w), right? where is 1/eta in equation 9? (2) So this is the first time of using amortized computation in gradient manipulation methods? how did the previous methods like MGDA, CAGrad and IMTL-G to solve z, directly solve z at every iteration using some convex optimizer like Frank-Wolfe? (3) In the new experiments Figure1, seems amortized MGDA is the worst, so this means the trade-off between efficiency and performance? Why amortized MGDA is so bad?