[논문 리뷰] FAMO: Fast Adaptive Multitask Optimization

Bo Liu, Yihao Feng외 2명|arXiv (Cornell University)|2023. 06. 06.

Domain Adaptation and Few-Shot Learning인용 수 6

한 줄 요약

FAMO는 손실 이력을 활용하여 작업 간 손실 감소를 동적으로 가중하는 방식으로, 각 반복에서 O(1)의 공간과 시간 복잡도를 유지하면서도 다중 작업 최적화를 빠르고 적응적으로 수행하는 방법이다. 이는 기존의 O(k)의 공간과 시간 복잡도를 요구하는 기울기 조작 방법과는 달리, 모든 작업의 기울기를 저장하거나 계산할 필요 없이도, 감독 및 강화 학습 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하면서도 훨씬 더 효율적이다.

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는 각 단계에서 모든 작업 기울기를 계산하는 대신, 이전 손실 추세를 기반으로 동적 작업 가중치를 설정하는 메커니즘을 사용한다.
• 기울기 가중치 로그를 최적화하는 데 안정성을 높이기 위해 계수 γ를 가진 정규화 항을 도입하며, 이는 확률적 기울기 기반 방법으로 업데이트된다.
• 손실 이력과 각 반복에서 한 번의 기울기 계산만을 사용하여 작업 가중치를 적응적으로 조정함으로써, 각 작업의 손실 감소 속도가 약간씩 동일하게 유지되도록 보장한다.
• 모든 작업의 기울기를 각 업데이트에서 저장하거나 명시적으로 계산하지 않기 때문에, 이전 방법들이 겪는 O(k)의 공간 및 시간 오버헤드를 피한다.
• 손실 이력은 각 작업의 상대적 진전을 추정하는 데 사용되며, 전체 기울기 계산 없이도 균형 잡힌 최적화를 가능하게 한다.
• 표준 최적화 프레임워크와 호환되며, 기존 딥러닝 파ip라인에 원활하게 통합될 수 있다.

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단계마다 한 번만 적용된 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.