Gradient Surgery for Multi-Task Learning 논문 요약

Abstract

이 논문은 멀티태스크 학습이 왜 어려운지를 gradient 관점에서 분석한다. 저자들은 여러 태스크의 gradient가 서로 충돌하면 학습이 비효율적이거나 성능이 떨어질 수 있다고 보고, 이를 줄이기 위한 방법으로 PCGrad(Projecting Conflicting Gradients) 를 제안한다. 핵심 아이디어는 서로 충돌하는 gradient 성분만 제거하는 것이다. 이 방법은 구조에 크게 의존하지 않고, supervised learning과 reinforcement learning 모두에 적용 가능하며, 여러 실험에서 효율성과 최종 성능을 개선했다고 주장한다.

그림 1

2차원 멀티태스크 최적화 문제에서 PCGrad를 시각화한 그림.
(a) 멀티태스크 목적 함수의 지형(objective landscape)
(b)와 (c)는 (a)를 구성하는 각 개별 작업 목적 함수의 등고선도(contour plot)
(d) Adam 옵티마이저를 사용했을 때, 멀티태스크 목적 함수 위에서의 그래디언트 업데이트 경로, 경로의 끝 지점에서 두 작업의 그래디언트 벡터는 각각 파란색 화살표와 빨간색 화살표로 표시되어 있으며, 이들의 상대적 길이는 로그 스케일로 나타냈다.
(e) PCGrad가 적용된 Adam을 사용했을 때, 멀티태스크 목적 함수 위에서의 그래디언트 업데이트 경로
(d)와 (e)에서 최적화 경로는 검은색에서 노란색 방향으로 진행된다.

1. Introduction

이 섹션에서는 문제의식을 설명한다.

멀티태스크 학습은 여러 태스크가 공유 구조를 활용할 수 있어서 원래는 더 효율적이어야 한다. 하지만 실제로는 각 태스크를 따로 학습하는 것보다 오히려 더 어렵고 성능도 나빠지는 경우가 많다. 기존 연구들은 learning speed 차이, optimization plateau, architecture 문제 등을 원인으로 봤지만, 이 논문은 특히 태스크 간 gradient conflict를 핵심 원인으로 본다.

저자들은 단순히 gradient가 충돌하는 것만이 문제가 아니라, 아래 세 조건이 함께 나타날 때 특히 심각해진다고 말한다.

• 서로 반대 방향인 conflicting gradients
• 어떤 태스크 gradient가 훨씬 커서 다른 태스크를 누르는 dominating gradients
• 손실 지형의 high curvature

이 세 가지를 합쳐 논문에서는 tragic triad 라고 부른다. 그리고 이 문제를 완화하기 위해 gradient 자체를 수정하는 gradient surgery 가 필요하다고 주장한다.

그림 2

충돌하는 그래디언트와 PCGrad.
(a)에서 작업 $i$ 와 $j$ 는 서로 충돌하는 그래디언트 방향을 가지며, 이는 파괴적인 간섭(destructive interference)을 일으킬 수 있다.
(b)와 (c)는 그래디언트가 서로 충돌하는 경우의 PCGrad 알고리즘을 보여준다. PCGrad는 작업 $i$ 의 그래디언트를 작업 $j$ 의 그래디언트의 법선 벡터(normal vector) 방향으로 투영하고, 반대로 작업 $j$ 의 그래디언트도 작업 $i$ 의 그래디언트의 법선 벡터 방향으로 투영한다.
충돌하지 않는 작업 그래디언트들(d)은 PCGrad에서 변경되지 않으며, 이를 통해 건설적인 상호작용(constructive interaction)이 가능해진다.

2. Multi-Task Learning with PCGrad

2.1 Preliminaries: Problem and Notation

멀티태스크 학습의 목적을 수식으로 정의한다. 전체 목적은 여러 태스크 손실의 평균 또는 합을 줄이는 것이고, 각 태스크의 gradient를 $𝑔_𝑖$ 로 둔다. 모델은 task-conditioned form으로 표현할 수 있고, 입력과 함께 task encoding을 넣는 형태를 사용한다. 즉, 이 부분은 뒤의 알고리즘 설명을 위한 기호 정리다.

2.2 The Tragic Triad: Conflicting Gradients, Dominating Gradients, High Curvature

여기서 저자들은 왜 멀티태스크 최적화가 힘든지 더 구체적으로 설명한다.

단순히 gradient 평균을 내면 될 것 같지만, 실제로는 두 태스크 gradient가 충돌하고, 한쪽 gradient 크기가 더 크고, curvature까지 크면 문제가 심해진다. 이 경우 큰 gradient를 가진 태스크 쪽으로만 업데이트가 치우치고, 다른 태스크는 오히려 손해를 보게 된다. 논문은 이를 2D 예시로 보여주며, Adam optimizer가 이런 상황에서 valley를 따라 제대로 진행하지 못하고 멈출 수 있다고 설명한다.

즉 이 섹션의 메시지는 이거다.

멀티태스크 학습이 어려운 이유는 태스크들이 서로 정보를 공유하지 못해서가 아니라, 공유하는 과정에서 gradient가 서로 방해하기 때문일 수 있다.

2.3 PCGrad: Project Conflicting Gradients

여기서 실제 방법이 나온다.

PCGrad는 두 태스크의 gradient가 충돌하면, 한 태스크 gradient를 다른 태스크 gradient의 normal plane으로 projection한다. 쉽게 말하면, 상대 태스크를 방해하는 방향 성분만 제거하는 것이다. 반대로 두 gradient가 충돌하지 않으면 그대로 둔다.

예를 들어 두 gradient $𝑔_𝑖$, $𝑔_𝑗$ 가 있고 inner product가 음수이면, $𝑔_𝑖$ 에서 $g_j$ 방향의 충돌 성분을 빼 준다. 이 과정을 현재 batch 안의 다른 태스크들에 대해 반복해서 각 태스크의 수정된 gradient를 만들고, 마지막에 합쳐 업데이트한다. 알고리즘 1이 이 절차를 정리한다.

핵심 장점은 세 가지다.

• 모델 구조에 독립적이다.
• 기존 optimizer와 함께 쓸 수 있다.
• gradient magnitude만 조절하는 게 아니라 방향까지 수정한다.

2.4 Theoretical Analysis of PCGrad

이 부분은 이론 분석이다.

먼저 convex한 2-task setting에서, 적절한 step size를 쓰면 PCGrad는 최적값으로 수렴하거나, 혹은 두 gradient가 완전히 정반대가 되는 특별한 경우에 머무를 수 있음을 보인다. 다만 실제 SGD/minibatch 환경에서는 gradient가 정확히 코사인 -1이 되는 경우가 드물기 때문에 큰 문제는 아니라고 설명한다.

그다음 더 중요한 결과는, gradient conflict가 충분히 크고, gradient magnitude 차이가 크고, curvature가 큰 상황에서는 PCGrad가 일반적인 multi-task gradient descent보다 더 낮은 loss를 만들 수 있다는 충분조건을 제시한 것이다. 즉, 논문의 tragic triad 가설을 이론적으로도 뒷받침한다.

3. PCGrad in Practice

이 섹션은 실제 구현 설명이다.

supervised learning

각 training step에서 batch를 task별로 나눈 뒤, 각 task gradient를 먼저 계산한다. 그런 다음 task gradient들 사이의 cosine similarity를 계산하고, PCGrad 규칙으로 gradient를 수정한 뒤 shared parameter에 적용한다. 중요한 점은 gradient를 다시 역전파할 필요 없이, 이미 계산한 gradient들만 가지고 처리할 수 있다는 것이다.

reinforcement learning

RL에서도 마찬가지로 task별 policy gradient나 actor-critic gradient를 계산한 뒤 PCGrad를 적용한다. actor와 critic 모두에 쓸 수 있다. 즉, supervised/RL 모두에 통하는 일반적인 gradient post-processing 기법으로 볼 수 있다.

4. Related Work

이 섹션은 기존 연구와의 관계를 정리한다.

기존 멀티태스크 연구는 크게 두 갈래가 있다.

첫째, architecture를 바꾸는 방법이다. 예를 들어 multiple modules, routing, attention-based sharing 같은 방식이다.

둘째, 태스크별로 따로 학습한 뒤 distillation으로 합치는 방식이다.

이 논문은 그와 달리 구조를 바꾸지 않고 optimization 자체를 바꾸는 방법을 제안한다. 또 GradNorm처럼 gradient magnitude만 재조정하는 방법과도 다르다. PCGrad는 gradient 방향 자체를 수정하기 때문에 더 직접적으로 interference를 줄인다. Continual learning 쪽 gradient projection 방법들과도 닮아 있지만, 이 논문은 동시에 여러 태스크를 학습하는 multitask setting을 다룬다는 점이 다르다.

5. Experiments

이 섹션은 세 가지 질문을 검증한다.

1. PCGrad가 실제로 최적화를 쉽게 만드는가
2. 기존 멀티태스크 방법과 결합해서 더 좋아지는가
3. tragic triad 가 실제로 중요한 원인인가

5.1 Multi-Task Supervised Learning

여기서는 CIFAR-100, CelebA, NYUv2 같은 supervised multitask benchmark를 실험한다. 결과적으로 PCGrad는 단독으로도 좋은 성능을 내고, routing networks나 MTAN 같은 기존 구조 기반 방법과 함께 썼을 때도 추가 성능 향상을 보인다. 특히 NYUv2에서는 MTAN + PCGrad가 9개 지표 중 8개에서 가장 좋은 결과를 냈다고 보고한다.

즉 supervised setting에서는 PCGrad가 단순 baseline 개선뿐 아니라, 기존 state-of-the-art multi-task architecture와도 잘 결합되는 보완적 방법 이라는 점을 보여준다.

5.2 Multi-Task Reinforcement Learning

여기서는 Meta-World의 MT10, MT50 benchmark에서 SAC와 결합해 실험한다. 결과적으로 PCGrad를 붙인 SAC가 일반 single policy나 multi-head policy보다 훨씬 좋은 성능을 보였고, independent training과 비교해도 더 적은 샘플로 비슷하거나 더 좋은 결과를 냈다. 특히 멀티태스크 RL에서 성능 차이가 크게 나타난다.

또 ablation으로,

• 방향만 바꾼 경우
• 크기만 바꾼 경우
• GradNorm 같은 magnitude-only 방법

과 비교했을 때, 방향과 크기를 함께 수정하는 원래 PCGrad가 가장 좋았다고 보고한다. 이건 이 논문의 중요한 포인트다. 단순히 loss balancing이 아니라 gradient interference 방향을 직접 제거하는 것이 핵심이라는 뜻이다.

5.3 Empirical Analysis of the Tragic Triad

이 부분은 논문의 가설 검증에 가깝다.

저자들은 RL 두 태스크를 예시로 실제 학습 중의 multi-task curvature, conflicting gradients 비율, 이론 조건이 성립하는 빈도를 측정한다. 그 결과, 학습 초반이나 어려운 구간에서 tragic triad 조건이 자주 나타났고, 바로 그 구간에서 PCGrad가 특히 이점을 보였다고 설명한다.

즉 이 섹션은 “PCGrad가 잘 된다”를 넘어서, 왜 잘 되는지에 대한 실험적 근거를 제시하는 부분이다.

표 1

NYUv2 데이터셋에서의 3개 작업 학습 결과: 13개 클래스 의미론적 분할(semantic segmentation), 깊이 추정(depth estimation), 그리고 표면 법선(surface normal) 예측 결과를 나타낸다.

#P 는 전체 네트워크 파라미터 수를 의미한다. 우리는 멀티태스크 아키텍처와 가중치 부여 방식의 조합 중 가장 좋은 성능을 낸 경우를 굵게 표시하였다. 또한 각 작업별로 가장 높은 검증 성능(validation score)은 박스로 표시하였다. 기호들은 기존 방법들을 나타낸다: ∗ : [28], † : [33], ‡ : [40]

다른 방법들의 성능은 Liu 등 [33]에서 보고된 값을 사용하였다.

표 2

CIFAR-100 멀티태스크 결과. PCGrad를 routing network와 결합했을 때, 성능이 크게 향상된다.

표 3

CelebA 결과. CelebA의 전체 40개 작업에 대한 평균 분류 오류(classification error)를 제시한다. 이 데이터셋에서 PCGrad는 기존 방법인 Sener와 Koltun [53]보다 더 우수한 성능을 보인다.

그림 3

왼쪽의 두 그래프는 각각 MT10과 MT50에서의 학습 곡선(learning curves)을 보여준다.

PCGrad는 성공률(success rate) 과 데이터 효율성(data efficiency) 측면 모두에서 다른 방법들보다 훨씬 뛰어난 성능을 보인다. 가장 오른쪽 그래프에서는, PCGrad로 수정된 그래디언트의 크기(magnitude)만 사용할 때와 방향(direction)만 사용할 때에 대한 ablation study 결과, 그리고 GradNorm [8] 과의 비교를 제시한다.

PCGrad는 이 두 ablation 버전과 GradNorm보다 모두 더 좋은 성능을 보이는데, 이는 멀티태스크 학습에서 그래디언트의 방향과 크기를 모두 수정하는 것이 중요하다는 점을 보여준다.

그림 4

정리 2에서 논의한 이론적 조건들에 대한 실험적 분석으로, 두 개의 강화학습 작업인 reach 와 press button top 에 대해 학습 초기 100 iteration을 보여준다.

왼쪽: 멀티태스크 곡률(multi-task curvature)의 추정값이다. 학습 전반에 걸쳐 높은 멀티태스크 곡률이 존재함을 관찰할 수 있으며, 이는 정리 2의 조건 (b)를 뒷받침하는 증거를 제공한다.

가운데: 실선은 두 작업 그래디언트 사이의 cosine similarity가 양수인 그래디언트의 비율을 나타낸다. 점선과 파선은 각각 cosine similarity가 음수인 iteration들 중에서, 정리 2의 조건 (a)와 조건 (b)의 함의인 $ξ(g_1,g_2)≤1$ 이 성립하는 iteration의 비율을 보여준다.

오른쪽: SAC와 SAC+PCGrad가 각 작업에서 달성한 평균 return을 나타낸다.

가운데와 오른쪽 그래프를 통해, Adam 과 Adam+PCGrad 모두 Task 2를 아직 해결하지 못한 동안에는 조건 (a)가 대부분의 시간 동안 성립한다는 것을 알 수 있다. 그리고 Adam+PCGrad가 Task 2를 학습하기 시작하자마자, 조건 (a)가 성립하는 비율은 감소하기 시작한다. 이 관찰은 조건 (a)가 PCGrad가 멀티태스크 학습에서 뛰어난 성능을 내는 핵심 요인임을 시사한다.

6. Conclusion

결론에서는 다시 한 번 핵심을 정리한다.

이 논문은 멀티태스크 학습의 주요 어려움으로 gradient conflict, high curvature, large gradient difference를 지목했고, 이를 완화하는 간단한 방법으로 PCGrad를 제안했다. PCGrad는 구현이 쉽고 모델에 독립적이며, supervised learning과 RL 모두에서 효율성과 성능을 개선했다. 저자들은 이 아이디어가 멀티태스크 학습뿐 아니라 meta-learning, continual learning, imitation learning 등 다른 분야에도 확장될 수 있다고 본다.