[論文レビュー] 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}.

研究の動機と目的

• 勾配の相反性により一部のタスクが遅く学習されるというマルチタスク学習における過小最適化の課題に対処すること。
• すべてのタスク勾配の保存や計算を必要とせず、すべてのタスクにおける損失低減をバランスよく保証するMTL最適化手法の設計。
• 各反復で O(1) の空間的・時間的計算量を達成することで、大規模モデルや多数のタスクにスケーラブルな手法の実現。
• 多様なMTLベンチマークにおいて、従来の勾配操作手法に比べて性能と計算効率の両面で優れるようにすること。

提案手法

• FAMO は、各ステップですべてのタスク勾配を計算するのではなく、損失履歴の傾向に基づいた動的タスク重み付け機構を用いる。
• タスク重みのログィットの最適化を安定化させるために、係数 γ を用いた正則化項を導入し、確率的勾配に基づくアプローチで更新を行う。
• 各タスクの損失が概ね同じ速度で減少するように、損失履歴と1回の勾配計算のみを用いて、タスク重みを動的に調整することで、バランスの取れた最適化を実現する。
• 各更新ですべてのタスクの勾配を保存・明示的に計算しないことで、従来手法の 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 多様なMTLベンチマークにおいて、FAMO は最先端の勾配操作手法と比べて、性能と計算効率の両面でどのように異なるか？
• RQ3 正則化係数 γ が、さまざまなデータセットにおける FAMO の安定性と性能に与える影響は何か？
• RQ4 FAMO は、高い計算コストを伴わずに、確率的強化学習設定における対立勾配を効果的に緩和できるか？

主な発見

• FAMO は、教師あり学習および強化学習のMTLベンチマークにおいて、CAGrad や NashMTL などの最先端手法と同等またはそれ以上の性能を達成する。
• MetaWorld-10 ベンチマークでは、FAMO は 0.83 ± 0.05 の成功率を達成し、PCGrad や Soft Modularization を上回り、CAGrad と同等の性能を示すが、はるかに高速である。
• 線形スカラライゼーション（LS）と比較して、FAMO は訓練オーバーヘッドをほとんど追加しない。全データセットで相対的訓練時間が 1.0 に近いが、タスク数の増加に伴い性能が著しく低下するような NASHMTL ような手法とは対照的である。
• アブレーションスタディの結果、FAMO は γ の値に強く依存せず、安定しているが、初期に損失がほぼゼロに近いタスクを含むデータセット（例：CityScapes）では性能に若干の感受性が見られる。
• FAMO は、100ステップに1回しか適用しない NashMTL の再現版に対しても、性能で上回っており、その優れた効率性と安定性を示している。
• FAMO は、NYU-v2 と QM-9 における Δm% といった複数の指標でも一貫した性能を発揮しており、タスク間の損失低減がバランスよく保たれていることが裏付けられている。

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