[논문 리뷰] KroneckerBERT: Learning Kronecker Decomposition for Pre-trained Language Models via Knowledge Distillation

Marzieh S. Tahaei, Ella Charlaix외 3명|arXiv (Cornell University)|2021. 09. 13.

Topic Modeling참고 문헌 35인용 수 10

한 줄 요약

KroneckerBERT는 트랜스포머 블록 내의 임bedding 레이어와 모든 선형 변환에 대해 크로네커 분해를 적용하고 중간 레이어 지식 증류를 결합하여 BERT 기반 모델의 극한 압축을 위한 새로운 프레임워크를 제안한다. 19배의 압축 비율(기존 BERT_BASE 크기의 5%)을 달성했을 때, GLUE 및 SQuAD에서 최신 기술을 초월하는 성능을 기록하며 높은 압축률임에도 불구하고 성능 저하 없이 우수한 성능을 확보한다.

ABSTRACT

The development of over-parameterized pre-trained language models has made a significant contribution toward the success of natural language processing. While over-parameterization of these models is the key to their generalization power, it makes them unsuitable for deployment on low-capacity devices. We push the limits of state-of-the-art Transformer-based pre-trained language model compression using Kronecker decomposition. We use this decomposition for compression of the embedding layer, all linear mappings in the multi-head attention, and the feed-forward network modules in the Transformer layer. We perform intermediate-layer knowledge distillation using the uncompressed model as the teacher to improve the performance of the compressed model. We present our KroneckerBERT, a compressed version of the BERT_BASE model obtained using this framework. We evaluate the performance of KroneckerBERT on well-known NLP benchmarks and show that for a high compression factor of 19 (5% of the size of the BERT_BASE model), our KroneckerBERT outperforms state-of-the-art compression methods on the GLUE. Our experiments indicate that the proposed model has promising out-of-distribution robustness and is superior to the state-of-the-art compression methods on SQuAD.

연구 동기 및 목표

• 낮은 용량의 엣지 디바이스에 대규모 과도하게 파ram터화된 사전 학습된 언어 모델을 구현하는 데 도전하는 것.
• 상당한 성능 저하 없이 10배 이상의 극한 압축(압축 인자 >10)을 가능하게 하는 압축 프레임워크를 개발하는 것.
• 표준 행렬 분해를 넘어서 트랜스포머 블록 내의 임베딩 레이어와 모든 선형 레이어에 크로네커 분해를 적용하는 것.
• 교수 모델의 중간 레이어 출력을 활용한 이중 단계 지식 증류 전략을 통해 압축 모델의 성능을 향상시키는 것.
• 기본 NLP 벤치마크에서 모델 크기, 추론 속도, 정확도 간의 상호 상충 관계를 평가하는 것.

제안 방법

• 임베딩 레이어, 멀티헤드 어텐션, 피드포워드 네트워크 내의 가중치 행렬에 크로네커 분해를 적용하여 더 작은 행렬의 크로네커 곱으로 표현함으로써 파라미터 수를 감소시킴.
• 각 가중치 행렬 W를 W = A ⊗ B 방식으로 분해함. 여기서 A와 B는 더 작은 행렬이며, 이는 모델 크기의 극적 감소를 가능하게 함.
• 사전 학습 및 미세조정 단계 모두에 지식 증류를 적용하며, 원본 BERT_BASE를 교수 모델로 삼아 학습자 모델의 중간 레이어 표현을 지도함.
• 압축 프레임워크를 적용하여 BERT_BASE를 KroneckerBERT로 변환하고, 지식 증류의 영향을 검증하기 위해 추론 분석을 실시함.
• 스마트폰과 고성능 CPU에서 추론 속도를 평가하여 압축에 따른 실제 러닝 타임 향상 여부를 분석함.
• 압축 모델을 엣지 디바이스에 배포하고 지연 시간을 측정하기 위해 TorchScript를 사용해 직렬화함.

실험 결과

연구 질문

• RQ1 크로네커 분해가 트랜스포머 기반 모델의 임베딩 레이어와 모든 선형 레이어에 효과적으로 적용되어 극한 압축을 가능하게 할 수 있는가?
• RQ2 중간 레이어 지식 증류는 크로네커 압축 모델에서 성능 저하를 어떻게 완화하는가?
• RQ3 고압축 비율(예: 19배)에서 KroneckerBERT의 표준 NLP 벤치마크 성능은 어떠한가?
• RQ4 크로네커 분해로 감소된 FLOPs가 실제로 엣지 디바이스에서 추론 속도 향상으로 이어지는가?
• RQ5 외부 분포 데이터에 대해 KroneckerBERT는 최신 기술 수준의 압축 모델과 비교해 정확도 및 강인성 측면에서 어떻게 성능을 내는가?

주요 결과

• KroneckerBERT는 19배의 압축 비율(기존 BERT_BASE 크기의 5%)을 달성했으며, 이 압축 수준에서 기존 방법들을 크게 능가하는 성능을 기록함.
• GLUE 벤치마크에서 19배 압축 시 KroneckerBERT는 최신 기술 수준의 압축 BERT 모델을 능가하며, 우수한 압축-성능 균형을 확보함.
• SQuAD에서 KroneckerBERT는 최신 기술 수준의 압축 방법들보다 뛰어난 성능을 보이며, 강력한 강인성과 일반화 능력을 입증함.
• 추론 분석 결과, 사전 학습 및 미세조정 단계 모두에서 지식 증류가 필수적임을 확인하였으며, 이중 단계 KD가 가장 높은 성능을 기록함.
• 스마트폰과 같은 엣지 디바이스에서 KroneckerBERT는 19배 압축 시 최대 2.4배의 속도 향상을 기록함으로써 실질적인 추론 이점이 있음을 확인함.
• 고성능 CPU에서는 크로네커 행렬 곱셈의 순차적 성격으로 인해 속도 향상이 관찰되지 않음. 이는 하드웨어 인식 최적화의 중요성을 강조함.