[論文レビュー] KDLSQ-BERT: A Quantized Bert Combining Knowledge Distillation with Learned Step Size Quantization

Jing Jin, Liang Cai+3 more|arXiv (Cornell University)|Jan 15, 2021

Topic Modeling参考文献 40被引用数 13

ひとこと要約

KDLSQ-BERTは、知識蒸留（KD）と学習可能ステップサイズ量子化（LSQ）を組み合わせた新しいBERT量子化手法を提案する。LSQ学習中に、高精度な教師BERTから量子化された学生モデルへ知識を転送することで、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スケール因子のための新規初期化手法を導入し、訓練の高速化と収束の向上を図る。
• バックプロパゲーション中に蒸留損失と真値損失を組み合わせることで、訓練の安定化と性能向上を図る。
• 蒸留信号と監督信号の両方を用いて、エンドツーエンドで学生モデルを訓練することで、低ビット幅でも効果的な知識転送を実現する。
• 設計上、任意のビット幅（例：2ビット、4ビット）の量子化をサポートする。これにより、三値や混合精度アプローチに見られる制限を回避する。

実験結果

リサーチクエスチョン

• RQ1 知識蒸留を学習可能ステップサイズ量子化（LSQ）と組み合わせることで、量子化されたBERTモデルの性能を顕著に向上させることができるか？
• RQ2 既存の量子化手法と比較して、KDLSQ-BERTは超低ビット幅（例：2ビット）でも高い精度を維持できるか？
• RQ3 蒸留損失と真値損失の組み合わせは、単独で使用する場合と比較して、収束性と精度の両面で優れた効果をもたらすか？
• RQ4 特にパラメータ数と推論速度の観点から、性能を犠牲にせずにどの程度のモデル圧縮が達成できるか？
• RQ5 提案手法は、GLUEおよびSQuADベンチマークを含む多様な自然言語処理タスクに一般化可能か？

主な発見

• KDLSQ-BERTは、2ビット量子化でも14.9倍のモデル圧縮比（418 MB から 28 MB）を達成し、平均GLUEスコアは83.987を維持する。
• 4ビット量子化では、平均GLUEスコアが84.414（54 MB、7.7倍圧縮）を達成し、フルプレシジョンのTinyBERT（258 MB、1.6倍圧縮）を上回る性能を発揮する。
• 重み・活性化・埋め込みをそれぞれ2-2-4ビットで量子化した場合、KDLSQ-BERTはGLUEでMNLIスコア85.013、RTEスコア89.989を達成し、フルプレシジョンBERTの性能に近づく。
• SQuAD 1.1では、2-2-4量子化のKDLSQ-BERTが88.996 EMを達成し、フルプレシジョンBERT（88.696）からわずか3.699ポイントの差にとどまり、優れたゼロショット一般化性能を示す。
• SQuAD 2.0では、4-4-8量子化で平均F1スコア84.086を達成し、フルプレシジョンのTinyBERTを1.361ポイント上回る。
• 蒸留損失と真値損失の組み合わせにより、特に低ビット幅で顕著な精度向上が見られ、二重監督の有効性が裏付けられる。