ICLR 2022｜スピードこそすべて！完全なバイナリ化 BiBERT による超圧縮

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

近年、事前トレーニング済みの言語モデルは自然言語処理において優れたパフォーマンスを発揮していますが、パラメータの数が多いため、実際のハードウェアデバイスへの展開が困難になっています。

最近、機械学習のトップカンファレンスであるICLR 2022の採択論文の結果が正式に発表されました。少なくとも9つの研究が、ニューラルネットワークの量子化の方向で関連する進歩を示しました。

この論文では、自然言語タスク向けの最初の完全にバイナリ量子化された BERT モデルである BiBERT を紹介します。これにより、FLOP とモデルサイズが最大 56.3 倍と 31.2 倍節約されます。

この研究は、北京航空航天大学、南洋理工大学、百度の劉翔龍教授のチームによって共同で完了しました。

事前トレーニング済みの言語モデルは自然言語処理では優れたパフォーマンスを発揮しますが、パラメータの数が多いため、実際のハードウェアデバイスへの展開が困難になります。

既存のモデル圧縮方法には、パラメータの量子化、蒸留、プルーニング、パラメータの共有などがあります。

その中で、パラメータ量子化法は、浮動小数点パラメータを固定小数点表現に効率的に変換し、モデルをコンパクトにします。

研究者らはQ-BERT[1]、Q8BERT[2]、GOBO[3]など多くの解決策を提案しているが、定量モデルは依然として表現力が限られており、最適化が難しいという深刻な問題に直面している。

幸いなことに、一般的に使用される補助的な最適化方法としての知識蒸留により、量子化モデルは完全精度の教師モデルの特徴表現を模倣することができ、それによって精度損失の問題をより適切に解決することができます。

この論文では、北京航空航天大学、NTU、百度の研究者らが、重み、アクティベーション、埋め込みを1 ビットに量子化する (アクティベーションを 4 ビット以上に維持しながら重みを 1 ビットに量子化するのではなく) BiBERT を提案しました。

これにより、モデルは推論中にビット単位の操作を使用できるようになり、モデルを実際のハードウェアに展開したときに推論が大幅に高速化されます。

我々は、2値化プロセス中のBERTモデルのパフォーマンス損失を研究しました。情報理論に基づいて、著者らは、順方向伝播における2値化後の注意機構の情報劣化問題を解決するために効率的なBi -Attention機構を導入し、逆方向伝播における蒸留の最適化方向不一致問題を解決するためにDirection-Matching Distillation法を提案しました。

△ 図1 BiBERTの全体フレームワーク

BiBERT は、BERT モデルの完全な 2 値化の実現可能性を初めて実証しました。GLUE データセットでの精度は、既存の BERT モデル 2 値化アルゴリズムを大幅に上回り、より高いビット表現のモデルをも上回りました。

モデルの計算とボリュームに関して、BiBERT は理論的には FLOP とモデルサイズをそれぞれ56.3 倍と31.2 倍削減できます。

方法

バイアテンション: バイナリアテンションメカニズム

私たちの研究によると、BERT モデルの注意メカニズムでは、ソフトマックス関数によって得られた正規化された注意重みは確率分布に従うと考えられており、それを直接 2 値化すると完全な情報損失につながり、情報エントロピーが 0 に退化することがわかりました (図 2 を参照)。

△ 図2 ソフトマックス関数に2値化を直接適用すると、完全な情報損失が発生する

この情報劣化を軽減するための一般的な対策は、符号関数を適用する前に入力テンソルの分布をシフトすることです。これは次のように表すことができます。

その中で、シフトパラメータは二値化のしきい値としても考えられ、二値化後のエントロピーを最大化することを期待しています。

ソフトマックス関数は順序を保持することに注意してください。つまり、バイナリ表現の情報エントロピーを最大化する固定しきい値があることを意味します。

Hard Attention [4]に触発されて、著者はbool関数を使用してAttentionの重みAを2値化します。

bool関数を適用することで、注目重みの中で値の低い要素が0に二値化されるため、エントロピー値が最大の注目重みでキー部分の要素をフィルタリングすることができます。

最終的なバイナリアテンションメカニズムは次のように表現できます。

このうち、BVは符号関数の2値化によって得られる値であり、BAはバイナリアテンションウェイトであり、これは慎重に設計されたビットアフィン行列乗算である。合計シフトで構成される BAMM 演算子は、トレーニングと推論の表現を揃え、効率的なビット計算を実行するために使用されます。

DMD: 方向マッチング蒸留

著者は、注意の重みは 2 つのバイナリ活性化を直接乗算することによって得られることを発見しました。

そのため、決定のエッジの値は反対側に簡単に二値化される可能性があるため、注意の重みを直接最適化すると、トレーニング中に最適化の方向の不一致が発生することがよくあります。（図3参照）

△ 図3 最適化方向の不一致

そのため、著者は、アクティベーションを蒸留するために上流のクエリ、キー、および値マトリックスの類似性マトリックスを構築する新しい蒸留スキームを設計しました。

ここで、||·|| はL2正則化を表します。これまでの研究では、このように構築された行列は、ネットワークの特定のパターンの意味的理解を反映し、スケールや数値サイズの影響を無視し、特徴間の本質的な相関関係をより安定して表現でき、バイナリネットワークと完全精度ネットワーク間の知識転送により適していると考えられていることが示されています。

したがって、蒸留損失は、隠れ層の損失、予測、および上記の活性化類似性マトリックスの合計として表すことができます。

ここでLは変圧器の層数を表し、

。