[논문 리뷰] KDLSQ-BERT: A Quantized Bert Combining Knowledge Distillation with Learned Step Size Quantization

Jing Jin, Liang Cai외 3명|arXiv (Cornell University)|2021. 01. 15.

Topic Modeling참고 문헌 40인용 수 13

한 줄 요약

KDLSQ-BERT는 지식 증류(KD)와 학습된 스텝 사이즈 양자화(Least Squares Quantization, LSQ)를 결합한 새로운 BERT 양자화 방법을 제안한다. 이 방법은 고정밀도 교사 BERT에서 양자화된 학생 모델로 지식을 전달함으로써 고정밀도, 초저비트 추론을 가능하게 한다. LSQ 학습 중에 지식을 전달함으로써 14.9배의 모델 압축을 달성하면서도 정확도 손실를 최소화하였으며, GLUE 및 SQuAD 벤치마크에서 2비트 정밀도에서도 최신 기술을 능가한다.

ABSTRACT

Recently, transformer-based language models such as BERT have shown tremendous performance improvement for a range of natural language processing tasks. However, these language models usually are computation expensive and memory intensive during inference. As a result, it is difficult to deploy them on resource-restricted devices. To improve the inference performance, as well as reduce the model size while maintaining the model accuracy, we propose a novel quantization method named KDLSQ-BERT that combines knowledge distillation (KD) with learned step size quantization (LSQ) for language model quantization. The main idea of our method is that the KD technique is leveraged to transfer the knowledge from a "teacher" model to a "student" model when exploiting LSQ to quantize that "student" model during the quantization training process. Extensive experiment results on GLUE benchmark and SQuAD demonstrate that our proposed KDLSQ-BERT not only performs effectively when doing different bit (e.g. 2-bit $\sim$ 8-bit) quantization, but also outperforms the existing BERT quantization methods, and even achieves comparable performance as the full-precision base-line model while obtaining 14.9x compression ratio. Our code will be public available.

연구 동기 및 목표

• 자원 제약이 있는 장치에 대규모이고 계산 비용이 큰 BERT 모델을 구현하는 데 도전하는 것.
• BERT의 초저비트(예: 2비트) 양자화에서 일반적으로 관찰되는 정확도 저하 문제를 해결하는 것.
• 지식 증류와 학습된 스텝 사이즈 양자화(LSQ)를 결합하여 양자화 성능을 향상시키는 것.
• 모델 아키텍처를 변경하지 않고도 효율적이고 하드웨어 우수한 양자화를 가능하게 하는 것.
• 고압축 비율을 달성하면서도 전체 정밀도 BERT와 경쟁 가능한 성능을 유지하는 것.

제안 방법

• 학습 중에 전체 정밀도 교사 BERT에서 양자화된 학생 BERT로 지식을 전달하기 위해 지식 증류를 활용한다.
• 학생 모델의 가중치와 활성화에 대해 학습된 스텝 사이즈 양자화(LSQ)를 적용하여 학습 가능한 스케일 요소를 갖춘 미분 가능한 양자화를 가능하게 한다.
• LSQ 스케일 요소에 대한 새로운 초기화 방법을 도입하여 학습 속도를 가속화하고 수렴성을 향상시킨다.
• 역전파 중에 증류 손실과 진리값 손실을 결합하여 학습 안정성과 성능 향상을 높인다.
• 증류 신호와 감독 신호를 동시에 사용하여 학생 모델을 종단 간(end-to-end)으로 학습시켜, 특히 저비트 폭에서 효과적인 지식 전달을 가능하게 한다.
• 삼성 또는 혼합 정밀도 접근 방식의 제한을 피하기 위해 설계상 임의의 비트 폭 양자화(예: 2비트, 4비트)를 지원한다.

실험 결과

연구 질문

• RQ1 지식 증류를 학습된 스텝 사이즈 양자화(LSQ)와 결합할 경우, 양자화된 BERT 모델의 성능 향상에 기여할 수 있는가?
• RQ2 기존의 양자화 방법과 비교해 볼 때, KDLSQ-BERT는 초저비트 폭(예: 2비트)에서도 높은 정확도를 유지하는 데 얼마나 효과적인가?
• RQ3 증류 손실과 진리값 손실의 조합은 단독으로 사용할 경우보다 더 나은 수렴성과 정확도를 제공하는가?
• RQ4 특히 파rameter 수와 추론 속도 측면에서 성능을 희생시키지 않고 얼마나 높은 수준의 모델 압축을 달성할 수 있는가?
• RQ5 제안된 방법은 GLUE 및 SQuAD 벤치마크를 포함한 다양한 NLP 작업에 일반화 가능한가?

주요 결과

• KDLSQ-BERT는 2비트 양자화에서 418MB에서 28MB로 14.9배의 모델 압축률을 달성하면서도 평균 GLUE 점수 83.987을 유지한다.
• 4비트 양자화에서 KDLSQ-BERT는 평균 GLUE 점수 84.414를 기록하며 54MB(7.7배 압축)를 달성했고, 전체 정밀도 TinyBERT(258MB, 1.6배 압축)를 능가한다.
• 가중치, 활성화, 임bedding에 대해 2-2-4 비트 양자화를 적용했을 때, KDLSQ-BERT는 GLUE에서 MNLI 점수 85.013과 RTE 점수 89.989를 기록하여 전체 정밀도 BERT 성능에 가까워졌다.
• SQuAD 1.1에서 2-2-4 양자화를 적용한 KDLSQ-BERT는 EM 점수 88.996을 기록했으며, 전체 정밀도 BERT(88.696)보다 3.699점 낮게 기록했지만, 강력한 제로샷 일반화 능력을 보였다.
• 4-4-8 양자화에서 KDLSQ-BERT는 평균 SQuAD 2.0 F1 점수 84.086를 기록했으며, 전체 정밀도 TinyBERT를 1.361점 초월했다.
• 증류 손실과 진리값 손실의 조합은 특히 저비트 폭에서 정확도 향상에 뚜렷한 기여를 하여 双중 감독의 효과를 확인했다.