[논문 리뷰] Self-Distillation for Further Pre-training of Transformers

Seanie Lee, Minki Kang외 3명|arXiv (Cornell University)|2022. 09. 30.

Multimodal Machine Learning Applications인용 수 4

한 줄 요약

이 논문은 도메인 특화된 레이블이 없는 데이터에 대해 시각 및 언어 트랜스포머의 추가 사전 훈련을 위한 정규화 방법으로 자기-디스틸리케이션(self-distillation)을 제안한다. 마스킹된 오토에코딩 동안, 추가 사전 훈련된 모델을 교사로 사용하여 초기 사전 훈련된 모델의 복제본(학생)을 표현 디스틸리케이션을 통해 안내함으로써, 다운스트림 분류 작업에서 일반화 성능을 향상시키고, 비전 트랜스포머에서의 과적합을 감소시킨다.

ABSTRACT

Pre-training a large transformer model on a massive amount of unlabeled data and fine-tuning it on labeled datasets for diverse downstream tasks has proven to be a successful strategy, for a variety of vision and natural language processing tasks. However, direct fine-tuning of the pre-trained model may be suboptimal if there exist large discrepancies across data domains for pre-training and fine-tuning. To tackle this issue, several previous studies have proposed further pre-training strategies, where we continue to pre-train the model on the target unlabeled dataset before fine-tuning. However, all of them solely focus on language models and we empirically find that a Vision Transformer is vulnerable to overfitting as we continue to pretrain the model on target unlabeled data. In order to tackle this limitation, we propose self-distillation as a regularization for a further pre-training stage. Specifically, we first further pre-train the initial pre-trained model on the target unlabeled data and then consider it as a teacher for self-distillation. Then we take the same initial pre-trained model as a student and enforce its hidden representations to be close to those of the teacher while optimizing the student with a masked auto-encoding objective. We empirically validate the efficacy of self-distillation on a variety of benchmark datasets for image and text classification tasks. Experimentally, we show that our proposed method outperforms all the relevant baselines. Theoretically, we analyze the proposed method with a simplified model to understand how self-distillation for further pre-training can potentially help improve the performance of the downstream tasks.

연구 동기 및 목표

• 사전 훈련 데이터와 파인튜닝 데이터 사이의 도메인 스위치 문제를 해결하되, 특히 표준 추가 사전 훈련이 과적합을 유도하는 비전 트랜스포머에서의 문제를 다루기 위해.
• 기존의 정규화 방법이 사전 훈련 초기화에서의 가중치 업데이트를 제약함으로써 적응을 방해하는 한계를 극복하기 위해.
• 데이터 증강 기반에 의존하지 않고도 추가 사전 훈련 동안 일반화 성능을 향상시키는 방법을 개발하기 위해.
• 제한된 레이블 데이터를 가진 저자원 다운스트림 작업에 대해 사전 훈련된 트랜스포머의 효과적인 적응을 가능하게 하기 위해.
• 자기-디스틸리케이션의 일반화 오차 감소 및 파인튜닝 성능 향상 기여에 대한 이론적 및 실증적 검증을 수행하기 위해.

제안 방법

• 목표 도메인의 레이블이 없는 데이터에 대해 초기 사전 훈련된 모델(예: ViT 또는 RoBERTa)을 마스킹된 오토에코딩(MAE)을 사용하여 추가 사전 훈련하여 교사 모델을 생성한다.
• 추가 사전 훈련된 모델의 인코더를 교사로 사용하여 지식 디스틸리케이션을 위한 목표 표현을 제공한다.
• 원래의 사전 훈련된 모델의 복제본을 학생으로 취하고, MAE 사전 훈련 동안 교사의 은닉 표현을 따라하도록 훈련한다.
• 학생 모델을 마스킹된 오토에코딩 복원 손실과 학생-교사 표현 간 일치를 강제하는 디스틸리케이션 손실을 조합한 목적함수로 최적화한다.
• 자기-디스틸리케이션 이후, 학생 모델을 레이블이 있는 다운스트림 데이터셋에 대해 파인튜닝한다.
• 이 방법은 시각 및 언어 모델 모두에 적용되며, 디스틸리케이션 라운드 수와 초기화로부터의 가중치 거리에 미치는 영향에 대한 분석도 포함된다.

실험 결과

연구 질문

• RQ1 표준 추가 사전 훈련이 과적합을 유도하는 비전 트랜스포머에서 자기-디스틸리케이션이 추가 사전 훈련을 효과적으로 정규화할 수 있는가?
• RQ2 다운스트림 일반화 성능 측면에서 자기-디스틸리케이션은 표준 파인튜닝 및 추가 사전 훈련과 비교해 어떻게 성능을 내는가?
• RQ3 첫 번째 라운드 자기-디스틸리케이션은 초기 사전 훈련된 가중치와 최종 파인튜닝된 가중치 간 L2 거리 감소에 어떤 역할을 하는가?
• RQ4 이론적 분석이 제안하는 바와 같이, 자기-디스틸리케이션은 다운스트림 파인튜닝 단계에서 일반화 갭을 줄이는가?
• RQ5 제한된 레이블 데이터를 가진 다양한 이미지 및 텍스트 분류 벤치마크에서 이 방법의 성능은 어떠한가?

주요 결과

• 표준 파인튜닝 및 추가 사전 훈련 베이스라인에 비해 자기-디스틸리케이션은 이미지 및 텍스트 분류 벤치마크에서 일반화 성능을 크게 향상시킨다.
• CIFAR-100 및 ImageNet-1K 데이터셋에서, 특히 저자원 환경에서의 성능에서 베이스라인보다 높은 정확도를 기록한다.
• 첫 번째 라운드 자기-디스틸리케이션은 가장 두드러진 정규화 효과를 가지며, 이후 라운드보다 초기 및 최종 모델 가중치 간 L2 거리를 더 크게 감소시킨다.
• 실증 결과는 초기화로부터의 가중치 거리 감소가 일반화 성능 향상과 상관관계가 있음을 확인하여, 가중치 근접도와 일반화 경계 사이의 이론적 연결을 지지한다.
• 단순화된 모델을 사용한 이론적 분석은 자기-디스틸리케이션이 파인튜닝 단계에서 감독 손실의 일반화 경계를 엄격히 감소시킴을 보여준다.
• 이 방법은 비전 및 언어 모델 모두에서 효과적이며, ViT와 RoBERTa를 포함한 다양한 모델에 대해 넓은 적용 가능성을 입증한다.