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, 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 (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 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.

翻译：人工智能的长期核心目标之一是创建通用型智能体，使其能够通过多任务学习（MTL）从多样化的数据中学习多个不同任务。然而，基于所有任务平均损失进行梯度下降（GD）可能导致某些任务严重欠优化，从而降低多任务性能。以往通过调控任务梯度以实现更均衡损失下降的方法，需要存储并计算所有任务的梯度（空间和时间复杂度均为O(K)，其中K为任务数量），这限制了其在大规模场景中的应用。本研究提出快速自适应多任务优化（FAMO）——一种采用O(1)空间和时间复杂度实现任务损失均衡下降的动态加权方法。我们开展了涵盖多任务监督学习与强化学习问题的广泛实验。结果表明，FAMO在达到或超越最先进梯度操控技术性能的同时，显著提升了空间与计算效率。代码已开源：https://github.com/Cranial-XIX/FAMO。