[論文レビュー] TernaryBERT: Distillation-aware Ultra-low Bit BERT

Wei Zhang, Lu Hou+5 more|arXiv (Cornell University)|Sep 27, 2020

Topic Modeling参考文献 37被引用数 10

ひとこと要約

TernaryBERT は、知識蒸留に配慮した、超低ビット量子化手法を提案する。この手法は、語彙埋め込みとトランスフォーマー層に対して、混合粒度で重みを {-1, 0, +1} に三値化する。近似に基づく手法と損失に配慮した三値化を併用する。全精度 BERT と同等の性能を達成しながら、モデルサイズを 14.9 倍に削減し、GLUE および SQuAD ベンチマークで既存の量子化および圧縮手法を上回る性能を発揮する。

ABSTRACT

Transformer-based pre-training models like BERT have achieved remarkable performance in many natural language processing tasks.However, these models are both computation and memory expensive, hindering their deployment to resource-constrained devices. In this work, we propose TernaryBERT, which ternarizes the weights in a fine-tuned BERT model. Specifically, we use both approximation-based and loss-aware ternarization methods and empirically investigate the ternarization granularity of different parts of BERT. Moreover, to reduce the accuracy degradation caused by the lower capacity of low bits, we leverage the knowledge distillation technique in the training process. Experiments on the GLUE benchmark and SQuAD show that our proposed TernaryBERT outperforms the other BERT quantization methods, and even achieves comparable performance as the full-precision model while being 14.9x smaller.

研究の動機と目的

• リソース制約のあるデバイス上で全精度 BERT モデルの高い計算コストとメモリ使用量に対処すること。
• 1〜2 ビットの超低ビット量子化における精度低下を、三値化と知識蒸留を組み合わせることで軽減すること。
• 語彙埋め込みやトランスフォーマー層などの BERT の異なるコンponents における最適な三値化粒度を調査すること。
• 三値重み制約によるモデル容量の低下にもかかわらず、高い性能を維持する手法を開発すること。

提案手法

• 近似に基づく手法と損失に配慮した量子化戦略を併用して、BERT 重みを {-1, 0, +1} に三値化する。
• トランスフォーマー重みにはレイヤー単位の粒度、語彙埋め込みにはヘッド単位の粒度を適用する。
• 全精度教師モデルからの知識蒸留を活用し、三値化された学生モデルの学習を支援する。
• ログティスに加えて、中間表現とアテンションマップに対しても知識蒸留損失を適用し、知識の保持を促進する。
• 交差エントロピー損失と蒸留損失を組み合わせたハイブリッド学習目的関数を採用し、一般化性能を向上させる。
• 消去実験および比較のため、損失に配慮した量子化（LAQ）を用いて 3 ビット量子化への拡張を実施する。

実験結果

リサーチクエスチョン

• RQ1 BERT 重みの三値化は、モデルサイズを著しく削減しながらも、全精度 BERT と同等の性能を達成できるか？
• RQ2 三値化の粒度選択（レイヤー単位 vs. ヘッド単位）は、BERT の異なるコンponents におけるモデル精度にどのように影響するか？
• RQ3 知識蒸留は、超低ビット量子化された BERT モデルにおける精度低下をどの程度軽減できるか？
• RQ4 異なる蒸留損失（ログティス、中間特徴、アテンションマップ）は、三値化 BERT における性能向上にどの程度寄与するか？
• RQ5 精度とモデルサイズの観点から、TernaryBERT は最先端の量子化および圧縮手法と比較してどのように差をつけるか？

主な発見

• TernaryBERT は 28 MB のモデルサイズで MNLI-m の精度を 83.5% 達成し、全精度 BERT と比較して 14.9 倍の小型化を実現した。
• 2 ビット重みの混合精度 Q-BERT や GOBO よりも高い精度を達成し、MNLI および SQuAD ベンチマークで両者を上回った。
• 知識蒸留は性能向上に顕著な寄与を示した：トランスフォーマーレイヤーとログティスからの蒸留を削除すると、MNLI-m 精度が 5% 以上低下した。
• 中間表現とアテンションマップへの蒸留は、ログティスのみに依存する場合よりも、性能向上に寄与した。
• データ拡張と微調整済み BERT 初期化の併用により、特に CoLA および RTE で性能向上が見られ、両者を適用した際には CoLA で最大 9.7% の向上を達成した。
• SQuAD v1.1 では、全精度 BERT と同等の性能（F1 87.4/87.4）を達成しながらも、モデルサイズを 14.9 倍に削減した。これは、強力な圧縮効率を示している。