[论文解读] FAMO: Fast Adaptive Multitask Optimization

Bo Liu, Yihao Feng+2 more|arXiv (Cornell University)|Jun 6, 2023

Domain Adaptation and Few-Shot Learning被引用 6

一句话总结

FAMO 是一种快速、自适应的多任务优化方法，通过利用损失历史动态加权任务，仅需每次迭代 O(1) 的空间和时间复杂度，即可在不同任务间平衡损失减少。它在监督学习和强化学习基准上实现了最先进性能，且显著优于现有需要存储和计算所有任务梯度的梯度调控方法，后者需 O(k) 的空间和时间复杂度。

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}.

研究动机与目标

• 解决多任务学习中因梯度冲突导致部分任务学习缓慢的优化不足问题。
• 设计一种多任务学习优化器，确保所有任务的损失减少保持平衡，且无需存储或计算所有任务的梯度。
• 在每次迭代中实现 O(1) 的空间和时间复杂度，使其可扩展至大规模模型和大量任务。
• 在多样化多任务学习基准上，性能和计算效率均优于现有梯度调控方法。

提出的方法

• FAMO 使用基于历史损失趋势的动态任务加权机制，而非在每一步计算所有任务的梯度。
• 引入一个系数为 γ 的正则化项，以稳定任务加权对数的优化，后者通过基于随机梯度的策略进行更新。
• 通过仅使用损失历史和每次迭代的单次梯度计算，自适应调整任务权重，确保每个任务的损失以相近速率减少。
• 通过不在每次更新时存储或显式计算所有任务的梯度，避免了先前方法所需的 O(k) 空间和时间开销。
• 利用损失历史估计每个任务的相对进展，实现在无需完整梯度计算的前提下实现平衡优化。
• 该方法与标准优化框架兼容，并可无缝集成至现有深度学习流水线。

Figure 1: Top left: The loss landscape, and individual task losses of a toy 2-task learning problem ( $\filledstar$ 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

实验结果

研究问题

• RQ1 我们能否设计一种多任务学习优化器，确保所有任务的损失减少保持平衡，同时在每次迭代中保持 O(1) 的空间和时间复杂度？
• RQ2 在多样化多任务学习基准上，FAMO 与最先进梯度调控方法相比，在性能和计算效率方面表现如何？
• RQ3 正则化系数 γ 对 FAMO 在不同数据集上的稳定性和性能有何影响？
• RQ4 FAMO 是否能在不产生高计算成本的前提下，有效缓解随机强化学习设置中的梯度冲突？

主要发现

• FAMO 在监督学习和强化学习多任务学习基准上，性能与 CAGrad 和 NashMTL 等最先进方法相当或更优。
• 在 MetaWorld-10 基准上，FAMO 的成功率达到 0.83 ± 0.05，优于 PCGrad 和 Soft Modularization，且与 CAGrad 相当，同时显著更快。
• 与线性标量法（LS）相比，FAMO 引入的训练开销可忽略不计，所有数据集上的相对训练时间均接近 1.0，而像 NashMTL 这类方法在任务数量增加时扩展性差。
• 消融研究显示，FAMO 对 γ 的不同取值具有鲁棒性，但在如 CityScapes 这类初始时某项任务损失接近零的数据集上，性能对 γ 更为敏感。
• 即使每 100 步仅应用一次，FAMO 的性能仍优于我们复现的 NashMTL，表明其具有更高的效率和稳定性。
• FAMO 保持了任务间损失减少的平衡性，表现为在 NYU-v2 和 QM-9 数据集上的多个指标（包括 Δm% 和 QM-9）上表现一致。

Figure 2: The average loss $L^{0}$ and the two task losses $L^{1}$ and $L^{2}$ for the toy example.