PCGrad - Gradient Surgery for Multi-Task Learning 논문 설명 및 구현

Multi-Task Learning은 하나의 모델로 여러가지 Task를 동시에 학습 시키는 방식입니다. 전반적인 성능이 향상되기도 하고, 성능을 끌어 올리고자 하는 주요 Task 가 있을 때 보조 Task를 잘 설계/선택하여 Multi-Task Learning 방식으로 주요 Task의 성능을 향상시키기도 합니다. 그러나 Multi Task에 사용되는 각각의 Task 관점에서 보면 개별적으로 학습 되었을 때 보다 성능이 떨어지는 경우가 제법 발생합니다. 

 

이에 이 원인을 상세히 분석하고 Multi Task Learning에서 성능을 안정적으로 향상시키고자 하는 방법들이 연구되어 왔습니다. 대표적인 기법중에 하나가 학습시 Task 간 Gradient를 조절하는 방식으로 접근하는 PCGrad 라는 기법입니다. 이 방식은 Yu et al., Gradient Surgery for Multi-Task Learning. NeurIPS. 2020. 논문에서 제안되었습니다.

 

본 포스트에서는  해당 논문을 설명하고 한국어 NLP Multi Task + RoBERTa 모델에 해당 기법을 적용하여 구현하고 실험한 내용 (https://github.com/han0ah/pcgrad-transformers) 을 설명합니다. 

---

Part1. PCGrad 논문 설명

1. Introduction

1-1. Overview

•  Multi Task Learning은 어려운 최적화 문제이며, 각 Task 를 개별적으로 학습시켰을 때보다 성능이 떨어지는 경우가 발생하기도 합니다.
• Multi task learning을 하게 되면 아무래도 성격이 다른 task를 학습하는 것이니, 학습 과정에서 task간 gradient가 충돌하는 상황(conflicting gradient)이 발생합니다. 본 논문에서는 이 때 어떤 상황에서 이 conflicting gradient 가 문제가 되는지를 분석하였습니다.
• 이 분석에 기반하여 해법으로 학습 시 gradient를 수정하는 Model-agnostic 한 방법을 제안합니다. 간략히 얘기하면Task간 gradient 벡터 끼리 서로 Projection 시키는 방식입니다. 

1-2. 문제가 되는 상황과 PCGrad 효과에 대한 Visualization 

논문에 소개된 Parameter 2개로 이루어진 모델을 2개의 Task에 대해서 학습 할 때 각 Task 의 paramter에 따른 Loss 등고선 그래프입니다. 특정 같은 위치(같은 paramter 값일 때)에서 각 task 가 가지는 gradient 는 위 그림에 점과 화살표로 표기하였습니다. 등고선에서 보이듯이,  Task1의 해당 위치는 경사가 급격하여 상대적으로 큰 gradient를 가지게 되고 Task2는 반대입니다. 또한 gradient의 방향도 다릅니다. (진한 주황색일 수록 Loss 값이 작은 구간입니다.)

 

이 때 두 Task간 gradient 를 단순히 합 또는 평균을 낸 후 optimize 하면, 왼쪽 그림과 같이 두 Task가 모두 최적화 되는 방향과 동떨어진 방향으로 이동하게 됩니다. 오른쪽 그림은 본 논문에서 사용된 기법을 활용하면 안정적으로 최적화가 될 수 있다는 것을 보여줍니다. 

2. Conflicting Gradients 및 문제가 되는 상황 정의

저자들은 위와 같이 두 Gradient Vector 사이 각의 cosine 값이 0보다 작을 때, 즉 직관적으로 두 벡터 사이 각이 90도 보다 클 때를 conflicg gradient 라고 정의합니다. 그런데 이 conflicting gradient 가 항상 문제는 아닙니다. task 간 gradient 끼리 단순히 평균을 내는 것이 Multi Task Objective를 최적화 하는데 올바른 솔루션을 제공할 때도 있습니다. 

 

특정 상황에서 이런 Conflicting Gradient 가 성능 저하를 유발하고 그 상황들을 저자들은 다음과 같이 2가지 케이스로 설명합니다.

1) 먼저 gradient magnitued similarity 가 작을 때 (=gradient magnitued 차이가 클 때)입니다. 즉 overview의 그림과 같이 특정 Task 의 gradient 값이 훨씬 클 때 입니다. 직관적으로 해당 상황에서 단순히 task 간 gradient 를 평균 내면, 훨씬 큰 gradient 값을 가지는 Task가 최적화 방향을 지배하게 됩니다.

 

2) 두번째로 multi-task curvature 가 높은 양수 값을 가질 때 입니다. curvature 는 한국어로 보통 곡률로 번역되고, 저자들이 정의한 multi-task curvature 는 loss의 이계도 함수가 활용됩니다. 이 부분은 정확히 이해 못했고, 일단은 피상적으로 아래정도로만 이해 했습니다.

loss의 이계도 함수 ( gradient 를 또 미분한 값 ) 이 크다 --> parameter를 조금만 움직여도 gradient가 급격히 커진다 --> 그래서 특정 optimal로 가속이 빠진다. ??

 

3. Solution - PCGrad!

이에 대한 Solution 으로 PCGrad를 제안합니다. 간략히는 두 벡터간 Conflicting Gradient( Cosine 값 < 0 ) 이 발생하면 한벡터를 다른 벡터의 normal plane 으로 projection 시킵니다. 논문에 나온 diagram 을 조금더 자세히 그려보았는데 다음 그림과 같습니다.

각 Task gradient는 는 다른 모든 Task 의 gradient에 대해서 순차적으로 이 projection 과정을 반복합니다. 그래서 총 N개의 Task가 있을 때, 이 projection 과정은 아래 알고리즘과 같이, N*(N-1) 번 반복됩니다.

4. 논문 실험 결과

Vision/강화학습 쪽 Task 들에 적용하여 실험을 하였고 PCGrad 가 다른 Multi-Task 방법론 보다 좋다, 또는 PCGrad는 Model-agnostic 하니 다른 Multi-Task 방법론과 결합하였을 때 좋다라는 내용입니다. 

---

Part2. 한국어 NLP Multi-Task 에 적용 실험

https://github.com/han0ah/pcgrad-transformers

Task(Dataset) 설정

2개의 Task 를 선택하여 실험하였고 한국어 benchmark 에서 많이 사용되며, 똑같이 Text-Pair 분류 형태의 task인 KLUE-NLI 와 PAWS-KR 를 선택하였습니다. 데이터셋은 다음 링크에서 train 과 dev 셋들을 사용했습니다.

• KlUE-NLI : https://klue-benchmark.com/tasks/68/data/description
• PAWS-KR : https://github.com/monologg/KoELECTRA/tree/master/finetune/data/paws

실험의 편리성을 위해 학습셋은 똑같이 20,000 만 선택하여 진행하였습니다.

모델

...
class RobertaForMTL(RobertaPreTrainedModel):
    # Task 0 : PAWS 2-class
    # Task 1 : Klue-NLI 3-class
    _TASK_NUM_CLASS = [2,3]

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config

        self.roberta = RobertaModel(config, add_pooling_layer=True)
        self.classifier_list = nn.ModuleList(
            [nn.Linear(config.hidden_size, self._TASK_NUM_CLASS[i]) for i in range(2)]
        )
 ...

공유하는 부분으로 klue-roberta-base RoBERTa 모델을 사용하였고, task 별 Layer로는 단일 Fully Connected Layer 를 사용하였습니다.

 

학습 구현

PCGrad 저자들이 쉽게 사용할 수 있는 PCGrad 라이브러리를 제공하고 있습니다. https://github.com/WeiChengTseng/Pytorch-PCGrad 

GitHub - WeiChengTseng/Pytorch-PCGrad: Pytorch reimplementation for "Gradient Surgery for Multi-Task Learning"

해당 저장소에서 pcgrad.py 를 가져온 후 PCGrad 학습시는 위 저장소 README 에 적혀있는 것처럼 몇줄만 수정해 주면 됩니다. 상세 차이는  https://github.com/han0ah/pcgrad-transformers 에서 train.py 와 train_pcgrad.py 를 보시면 됩니다.

실험 결과

	baseline		+pcgrad
	Loss	Accuracy	Loss	Accuracy
PAWS-KR	0.4793	0.8030	0.4071	0.8325
KLUE-NLI	0.4432	0.8486	0.4365	0.8520

Validation(Dev) 셋에 대한 결과입니다. 단일 hyperparmeter로 한번만 실험했고, validation loss 가 가장 작을때를 선택했습니다. 일단 이 실험 결과에서는 pcgrad 가 확연히 나은 성능을 보입니다. 학습 시행이나 hyperparmetr search에 따라 결과가 달라질 수도 있을 것 같습니다. 또 아시다시피 이런 분류 task 들은 validation loss가 조금 올라가도 후반 epoch 에 accuracy가 더 좋을 때가 자주 발생해서, accuracy 기준으로 모델을 선택한다면 pcgrad 가 확연히 좋지는 않고 애매합니다... 

---

Part3. 마치며

일단 PCGrad는 논문에서 보여주듯이 상단한 이론적 백그라운드가 있고, 알고리즘도 직관적인데, 저자들이 라이브러리까지 제공해줘서 구현이 훨씬 더 심플하다는 장점이 있습니다. 단 단점은 학습 시 많은 메모리를 필요로 합니다. 

 

핵심 알고리즘인 gradient 벡터간 projection을 구하는 부분을 보면, 결국 이 task gradient 벡터라는 것이 task 간 공유하는 모델의 모든 파라미터 입니다. 실험한 모델 같은 경우 RoBERTA Base size 를 공유 부분으로 썼는데 그럼 벡터하나의 크기가 약 1억개 입니다. 이런 거대한 vector 끼리 projection을 구하는 과정에서 추가적인 학습 시간도 필요하고 많은 메모리를 필요로 합니다. 그래서 모델 사이즈가 조금 더 커지거나 여러개 TASK를 다룰 때는 일반적인 학습 환경에서 어려운 부분이 있습니다.

 

또한 추가적으로,최근 이러한 multi-task learning optimization 기법들이 추가적인 학습 시간/메모리/복잡도를 필요로하는 것에 비해 단순히 task 간 gradient 를 더하고 평균 내는 방법과 비교했을 때 성능 향상에 미치는 영향이 미미하다는 연구들이 나오고 있습니다. 

In Defense of the Unitary Scalarization for Deep Multi-Task Learning. NeurIPS. 2022

Do Current Multi-Task Optimization Methods in Deep Learning Even Help?. NeruIPS. 2022

 

일단 논문을 읽어보고 몇가지 실험을 돌려본 후 제 의견은 PLM 을 활용하는 거대한 모델이나 Task 가 많은 복잡한 상황에서는 PCGrad 적용에 어려움이 있고, 그 성능향상은 아직 의문이나, 비교적 작은 모델 상황에서는 구현도 비교적 간단하니 성능을 끌어올려야 할 때 한번쯤 시도해볼만한 방식인 것 같습니다.