[논문 리뷰] Less is More: Task-aware Layer-wise Distillation for Language Model Compression

Liang Chen, Simiao Zuo외 4명|arXiv (Cornell University)|2022. 10. 04.

Topic Modeling인용 수 14

한 줄 요약

이 논문은 태스크별 필터를 사용하여 교사 모델의 은닉 표현에서 유의미한 지식만 학생 모델로 희석시키는 Task-aware Layer-wise Distillation (TED)을 제안한다. 이는 부족한 지식을 줄이고 희석 효율성을 향상시킨다. TED는 지속적 사전학습 및 미세조정 환경에서 모두 최신 기술 성능을 달성하며, GLUE 및 SQuAD 벤치마크에서 기존 희석 방법들을 능가한다.

ABSTRACT

Layer-wise distillation is a powerful tool to compress large models (i.e. teacher models) into small ones (i.e., student models). The student distills knowledge from the teacher by mimicking the hidden representations of the teacher at every intermediate layer. However, layer-wise distillation is difficult. Since the student has a smaller model capacity than the teacher, it is often under-fitted. Furthermore, the hidden representations of the teacher contain redundant information that the student does not necessarily need for the target task's learning. To address these challenges, we propose a novel Task-aware layEr-wise Distillation (TED). TED designs task-aware filters to align the hidden representations of the student and the teacher at each layer. The filters select the knowledge that is useful for the target task from the hidden representations. As such, TED reduces the knowledge gap between the two models and helps the student to fit better on the target task. We evaluate TED in two scenarios: continual pre-training and fine-tuning. TED demonstrates significant and consistent improvements over existing distillation methods in both scenarios. Code is available at https://github.com/cliang1453/task-aware-distillation.

연구 동기 및 목표

• 계층별 희석에서 지식의 중복성 문제를 해결하기 위해, 학생 모델이 교사 은닉 상태의 관련 없는 정보로 인해 유용한 표현을 학습하기 어려운 문제를 해결한다.
• 대규모 교사 모델과 소규모 학생 모델 간의 지식 격차를 줄이기 위해, 태스크 관련 특징에 집중하여 희석을 수행한다.
• 중간 계층의 가장 예측 가능한 지식만 정렬하여, 자원이 제한된 환경에서 학생 모델의 일반화 능력과 성능을 향상시킨다.
• 먼저 태스크 관련 특징을 식별하고 나서 효율적으로 희석하는 두 단계 학습 프레임워크를 개발한다.
• 지속적 사전학습 및 미세조정 설정에서 기존 희석 방법들에 비해 일관된 성능 향상을 입증한다.

제안 방법

• TED는 학생 및 교사 모델의 각 계층에 대해 태스크 인식 필터를 도입하며, 이는 태스크 전용 헤드(예: 분류 작업의 경우 소프트맥스)를 갖는 신경망으로 구현된다.
• 제1단계에서는 교사 및 학생 모델의 가중치를 동결한 상태에서, 타겟 태스크의 손실을 예측하기 위해 필터를 훈련시켜 은닉 표현으로부터 태스크 관련 특징을 추출하도록 학습한다.
• 제2단계에서는 필터를 동결(태스크 헤드 제외)하고, 각 계층에서 교사 및 학생 모델의 필터링된 표현 간의 평균 제곱오차(MSE)를 최소화하도록 학생 모델과 함께 공동으로 미세조정한다.
• 희석 손실은 대응하는 교사 및 학생 계층의 필터링된 출력 간의 평균 제곱오차로 정의되며, 이는 태스크 특화 지식에 대한 일치를 장려한다.
• 전체 학습 목표는 태스크 특화 손실, 예측 희석 손실, 그리고 학습 가능한 가중치를 가진 태스크 인식 계층별 희석 손실을 조합한다.
• 필터는 선형 투영, MLP 또는 트랜스포머 레이어로도 구현 가능하며, 추론 실험 결과 성능와 비용 간의 상충 관계를 보여준다.

실험 결과

연구 질문

• RQ1 중간 은닉 표현의 태스크 인식 필터링이 소규모 학생 모델에서 지식 희석 효율성을 향상시킬 수 있는가?
• RQ2 교사 표현에서 불필요한 정보를 걸러내면 표준 계층별 희석 대비 학생 모델의 성능 향상이 이루어지는가?
• RQ3 지속적 사전학습 및 미세조정과 같은 다양한 학습 철학에서 TED의 성능은 어떠한가?
• RQ4 필터 사전학습 후 희석을 수행하는 두 단계 학습 프레임워크가 종단 간 희석보다 더 나은 일치 및 일반화를 이끌 수 있는가?
• RQ5 다양한 필터 아키텍처(예: 선형 대비 MLP 또는 트랜스포머)가 희석 성능에 미치는 영향은 어떠한가?

주요 결과

• TED는 GLUE 벤치마크에서 표준 지식 희석 및 기존 계층별 희석 방법보다 뚜렷하게 승리하며, 특히 데이터가 적은 환경에서 두드러진 성능 향상을 보였다.
• SQuAD v1.1 및 v2.0 질문-답변 데이터셋에서, DeBERTaV3-xsmall 학생 모델을 DeBERTaV3-base 교사 모델에서 희석함으로써 최신 기술 성능을 달성했다.
• 지속적 사전학습 환경에서, TED는 Lambada 및 WikiText-103에서 제로샷 및 패기샷 전이 성능을 향상시켜 강력한 일반화 능력을 입증했다.
• 추론 실험 결과, 태스크 인식 필터가 중복 지식의 악영향을 줄여 희석 효율성과 학생 모델의 적합도를 향상시킴을 확인했다.
• 두 단계 학습 프로세스는 처음부터 필터와 학생 모델을 함께 훈련하는 것보다 더 좋은 수렴과 성능을 이끌었다.
• 더 복잡한 필터(예: MLP 또는 트랜스포머)를 사용할 경우 간단한 선형 투영 대비 성능 향상을 보였지만, 계산 비용이 증가했다.