[논문 리뷰] Know What You Don't Need: Single-Shot Meta-Pruning for Attention Heads

Zhengyan Zhang, Fanchao Qi외 3명|arXiv (Cornell University)|2020. 11. 07.

Topic Modeling참고 문헌 47인용 수 5

한 줄 요약

이 논문은 사전 훈련된 BERT 모델의 주의 헤드를 미세조정 이전에 최대 50%까지 단일 스텝으로 프루닝할 수 있는 Single-Shot Meta-Pruning (SMP)을 제안한다. 이 방법은 텍스트 표현의 분포를 유지하는 메타학습 기반 기준을 사용한다. SMP는 미세조정 및 추론 비용을 줄이며, 일반화 성능을 향상시켜 미세조정 기반 기준보다 유사하거나 더 좋은 성능을 달성한다. 특히 의미적 유사성 작업에서 개선된 일반화 성능을 보였다.

ABSTRACT

Deep pre-trained Transformer models have achieved state-of-the-art results over a variety of natural language processing (NLP) tasks. By learning rich language knowledge with millions of parameters, these models are usually overparameterized and significantly increase the computational overhead in applications. It is intuitive to address this issue by model compression. In this work, we propose a method, called Single-Shot Meta-Pruning, to compress deep pre-trained Transformers before fine-tuning. Specifically, we focus on pruning unnecessary attention heads adaptively for different downstream tasks. To measure the informativeness of attention heads, we train our Single-Shot Meta-Pruner (SMP) with a meta-learning paradigm aiming to maintain the distribution of text representations after pruning. Compared with existing compression methods for pre-trained models, our method can reduce the overhead of both fine-tuning and inference. Experimental results show that our pruner can selectively prune 50% of attention heads with little impact on the performance on downstream tasks and even provide better text representations. The source code will be released in the future.

연구 동기 및 목표

• 대규모 사전 훈련된 트랜스포머 모델에서 미세조정 및 추론의 계산 오버헤드를 줄이기 위해.
• 미세조정 이후가 아니라, 사전 훈련된 모델에서 불필요한 주의 헤드를 식별하고 제거하기 위해.
• 다양한 최종 작업에 걸쳐 성능을 유지하는 일반화 가능한 단일 스텝 프루닝 방법을 개발하기 위해.
• 다양한 모델 크기와 작업 간 이식 가능한 암묵적 프루닝 규칙을 메타학습을 통해 학습하기 위해.

제안 방법

• 텍스트 표현의 분포를 프루닝 후에도 유지하는 자기지도 학습 목표를 사용하여, 메타학습자(SMP)가 주의 헤드 중요도를 예측하도록 훈련한다.
• 다양한 문장 집합을 사용하여 메타학습 프레임워크를 적용하여 프루너 훈련 중 다양한 최종 작업을 시뮬레이션한다.
• 반복 최적화를 피하기 위해, 미세조정 이전에 한 번의 스텝으로 주의 헤드를 제거하는 단일 스텝 프루닝 전략을 적용한다.
• 정보량을 반영하는 점수 함수를 통해 헤드 중요도를 측정하며, 높은 점수는 유용한 헤드를 나타낸다.
• 점수 기반으로 낮은 점수를 가진 헤드를 식별하고, 비정보성으로 간주하여 제거한다.
• 최종 작업에서 프루닝된 모델의 성능 및 효율성 향상을 평가하기 위해 평가를 실시한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 트랜스포머에서 미세조정 이전에 주의 헤드를 효과적으로 프루닝하여 계산 비용을 줄일 수 있는가?
• RQ2 메타학습 기반 프루너는 다양한 최종 작업과 모델 아키텍처 간에 일반화 가능한가?
• RQ3 단일 스텝 프루닝이 반복적인 미세조정 기반 압축과 유사하거나 더 좋은 성능을 달성할 수 있는가?
• RQ4 일부 주의 헤드를 프루닝하면 특정 자연어 처리 작업에서 모델의 일반화 성능이 향상되는가?
• RQ5 프루닝된 주의 헤드에서 드러나는 암묵적 패턴은 인간이 이해하는 주의의 의미와 일치하는가?

주요 결과

• SMP는 GLUE 벤치마크 작업에서 BERT의 최대 50% 주의 헤드를 프루닝해도 성능 저하가 최소한이다.
• 헤드 50% 프루닝으로 메모리 사용량이 30% 감소하여 더 높은 배치 크기로 효율적인 훈련이 가능해졌다.
• 일부 사례에서는 성능 향상이 관찰되었으며, 특히 SentEval의 의미적 유사성 작업에서 SMP는 무작위 프루닝보다 뛰어나고 HISP-retrain와 유사한 성능을 보였다.
• 더 적은 파라미터로도 미세조정 이후 압축 기반 기준보다 유사하거나 더 좋은 결과를 달성했다.
• SMP가 학습한 암묵적 프루닝 규칙은 주의 행렬의 시각화를 통해 인간의 직관과 일치하는 것으로 확인되었다.
• 프루너는 강력한 이식성(transferability)을 보였다: BERT_BASE에서 훈련된 모델은 BERT_LARGE로도 잘 일반화되었으며, QNLI와 같은 새로운 작업에서도 우수한 성능을 보였다.