図解されたtinyBERTモデル - BERTモデル圧縮のエッセンス

翻訳者|朱仙中

レビュー | Chonglou

導入

近年、大規模言語モデルの開発は飛躍的に進歩しました。 BERT は、さまざまな自然言語処理 ( NLP )タスクを高い精度で解決する、最も人気があり効果的なモデルの 1 つになりました。 BERTモデルに続いて、他のモデルのグループが登場し、それぞれが優れたパフォーマンスを示しました。

時間の経過とともに、大規模言語モデル (LLM) は、トレーニングの対象となるパラメータとデータの数が指数関数的に増加するため、より複雑になる傾向があることは明らかです。ディープラーニングの研究によると、この手法は多くの場合、より良い運用結果につながることが分かっています。残念ながら、機械学習の世界では大規模言語モデルに関連する多くの問題を克服してきましたが、スケーラビリティの問題が大規模言語モデルを効果的にトレーニング、保存、使用するための大きな障害となっています。

上記の問題を考慮して、大規模な言語モデルを圧縮するための多くの特別な方法が開発されてきました。この投稿では、 Transformer Distillationメソッドに焦点を当てます。このメソッドによって、 TinyBERT と呼ばれるBERTモデルのミニバージョンが生まれました。さらに、 TinyBERTモデルの学習プロセスと、 TinyBERTモデルが非常に強力であるいくつかの微妙な理由も紹介します。この記事はTinyBERT の公式論文に基づいています。

本旨

最近の投稿では、DistilBERTモデルで蒸留手法がどのように機能するかについて説明しました。つまり、蒸留の主なアイデアは、学生モデルと教師モデルの予測が同様になるように損失関数の目的を変更することです。 DistilBERTモデルでは、損失関数は、両方のモデルの出力埋め込み（類似性損失の場合）を考慮して、学生モデルと教師モデルの出力分布を比較します。

DistilBERTモデルの詳細については、次の記事「大規模言語モデル: DistilBERT — より小さく、より高速で、より安価で、より軽量」を参照してください。
「https://towardsdatascience.com/distilbert-11c8810d29fc?source=post_page-----1a928ba3082b--------------------------------」。この記事の主な内容は、教師-生徒モデルフレームワークの効率を最大化することを目標としたBERTモデル圧縮の秘密を紹介することです。

表面的には、 TinyBERTモデルの蒸留フレームワークはDistilBERTモデルからあまり変わっていません。損失関数が再度変更され、目標は学生モデルが教師モデルを模倣することです。ただし、 TinyBERTモデルの場合はさらに一歩進んで、損失関数は教師モデルと生徒モデルによって生成された結果を考慮するだけでなく、予測結果がどのように得られたかという問題も考慮します。 TinyBERTモデル論文の著者によると、TinyBERT 損失関数は3 つの部分で構成されており、教師モデルと生徒モデルのさまざまな側面をカバーしています。

埋め込み層の出力
トランスフォーマー層からエクスポートされた隠れ状態と注目行列
予測層によって出力されるロジット値

変圧器蒸留損失関数の概略図

では、教師モデルと生徒モデルの隠れ状態を比較することにはどのような意義があるのでしょうか?注意行列に隠れ状態と注意の出力を含めることで、生徒モデルは教師モデルの隠れ層の内容を学習し、教師モデルに似た層を構築できるようになります。このようにして、抽出されたモデルは元のモデルの出力だけでなく、その内部動作も模倣できます。

では、教師モデルの行動をコピーすることがなぜ重要なのでしょうか?研究者らは、 BERTモデルによって学習された注意重みは言語構造を捉えるのに役立つと主張している。したがって、別のモデルを蒸留することで、学生モデルに言語知識を習得する機会も増えます。

レイヤーマッピング

TinyBERTモデルは、エンコーダーレイヤーの数が少ない BERT の小型バージョンを単純に表します。ここで、 BERTモデルレイヤーの数をN、TinyBERTモデルレイヤーの数をM と定義します。層の数が異なるため、蒸留損失値をどのように計算するかは明確ではありません。

この目的のために、TinyBERTモデルの対応するレイヤーmに知識を抽出するためにどのBERTモデルレイヤーnが使用されるかを定義する特別な関数n=g(m)が導入されています。選択された BERT レイヤーは、トレーニング中の損失値の計算に使用されます。

導入された関数n = g(m)には2つの推論制約があります。

g(0)=0です。これは、 BERTモデルの埋め込みレイヤーがTinyBERTモデルの埋め込みレイヤーに直接マッピングされることを意味し、これは理にかなっています。
g(M+1)=N+1です。この式は、 BERTモデルの予測レイヤーがTinyBERTモデルの予測レイヤーにマッピングされることを示しています。条件1≤m≤mを満たすTinyBERTモデル内の他のすべてのレイヤーについては、n=g(m)の対応する関数値をマッピングする必要があります。ここで、このような関数が定義されていると仮定します。
TinyBERTモデルのセットアップに関する問題については、この記事の後半で説明します。

コンバータ蒸留

1. 埋め込み層の蒸留

生の入力は最初にトークン化され、次に学習された埋め込みレイヤーにマッピングされてからモデルに渡されます。これらの埋め込みレイヤーは、モデルの最初のレイヤーとして使用されます。すべての可能な埋め込みレイヤーは、行列形式で表現できます。生徒モデルと教師モデルの埋め込み層の違いを比較するには、それぞれの埋め込み行列 E に対して標準的な回帰メトリックを使用できます。たとえば、 Transformer Distillation では、回帰メトリックとして平均二乗誤差( MSE )を使用します。

生徒モデルと教師モデルの埋め込み行列のサイズが異なるため、平均二乗誤差を使用してそれらの要素を適切に比較することはできません。これにより、生徒モデルの埋め込み行列に学習可能な重み行列 Wが乗算され、結果として教師モデルの埋め込み行列と同じ形状の行列が生成される理由が説明されます。

埋め込み層の蒸留損失関数。

生徒モデルと教師モデルの埋め込み空間は異なるため、行列Wは生徒モデルの埋め込み空間を教師モデルの埋め込み空間に線形変換する上で重要な役割を果たします。

2. コンバーター層蒸留

変圧器層蒸留損失関数の可視化

2A. 注意層の蒸留

トランスフォーマーのマルチヘッドアテンションメカニズムの中核は、豊富な言語知識を含む複数のアテンションマトリックスを生成することです。教師モデルの注意の重みをシフトすることで、生徒モデルも重要な言語概念を理解できるようになります。このアイデアを実装するために、損失関数を使用して、生徒モデルと教師モデルの注意重みの差を計算します。

TinyBERTモデルでは、すべての注意層が考慮され、各層の最終的な損失値は、すべてのヘッドの対応する学生モデルと教師モデルの注意行列間の平均二乗誤差値の合計に等しくなります。

層蒸留の損失関数は次のように計算されることに注意してください。

注意層の抽出に使用される注意行列 A は、ソフトマックス出力 softmax(A) ではなく、正規化されていないことに注意してください。研究者によると、この微妙な違いにより収束が早まり、パフォーマンスが向上するという。

2B. 隠れ層の蒸留

豊富な言語知識をキャプチャするというアイデアを実現するために、トランスフォーマー層の出力に蒸留操作も適用されます。

隠れ層蒸留損失関数の計算式。

ここで、重み行列 W は、前述の埋め込み層蒸留に使用される重み行列と同じ役割を果たします。

3. 予測層の蒸留

最後に、生徒モデルが教師モデルの出力を再現するように、予測層損失関数が使用されます。 2 つのモデルによって予測されるロジットベクトル間のクロスエントロピーを計算します。

予測層蒸留損失関数計算式

場合によっては、ロジットは出力分布の滑らかさを制御する温度パラメータ T によって除算されることに留意する価値があります。 TinyBERTモデルでは、温度パラメータT は 1 に設定されています。

損失方程式

TinyBERTモデルでは、各レイヤーにはそのタイプ特性に応じた独自の損失関数があります。いくつかのレイヤーの重要度の多寡を考慮するために、対応する損失値に定数 a を掛けます。最終的な損失関数は、すべての TinyBERTモデルレイヤーの損失値の加重合計に等しくなります。

TinyBERTモデルにおける損失関数の計算式

広範囲にわたる実験により、 3 つの損失要素のうち、変圧器層の蒸留損失がモデルのパフォーマンスに最も大きな影響を与えることが示されています。

モデルトレーニング

ほとんどの自然言語処理モデル (BERT を含む)の開発プロセスは、おおまかに2 つの段階に分けられることに注意することが重要です。

このモデルは、言語構造に関する一般的な知識を得るために、大規模なデータコーパスで事前トレーニングされています。
特定の下流タスクを解決するために、別のデータセットでモデルを微調整します。

同じアイデアに従って、研究者は新しいフレームワーク TinyBERTを開発しました。その学習プロセスも上記と同様の2 つの段階で構成されます。これら 2 つのトレーニング段階では、トランスフォーマー蒸留アルゴリズムを使用して、 BERTモデルの知識をTinyBERTモデルに変換します。

ステージ1：通常の蒸留。教師モデルとして、 TinyBERT は事前トレーニング済みの (微調整は不要) BERTモデルを通じて言語構造の豊富な常識を獲得します。より少ないレイヤーとパラメータを使用することにより、この段階以降、 TinyBERTモデルのパフォーマンスは BERTモデルよりも一般的にわずかに低下します。

ステージ 2:タスク固有の蒸留。今回は、微調整されたBERTモデルが教師モデルの役割を果たします。パフォーマンスをさらに向上させるために、研究者の提案に従って、トレーニングデータセットにデータ拡張手法が適用されました。実験結果によると、タスク固有の蒸留操作の後、TinyBERTモデルはBERTモデルと同等のパフォーマンスを実現します。

TinyBERT モデルのトレーニングプロセスの概略図

データ拡張

タスク固有の蒸留のために、特別なデータ拡張技術が導入されています。このデータ拡張技術では、まず特定のデータセットからシーケンスが抽出され、次に一定の割合の単語が次の2 つの方法のいずれかで置き換えられます。

単語が同じ単語としてラベル付けされている場合、それは BERT モデルによって予測され、予測された単語がシーケンス内の元の単語に置き換えられます。
単語が複数のサブワードにトークン化されている場合、それらのサブワードは最も類似した GloVe 埋め込み(単語表現のグローバルベクトル)に置き換えられます。

モデルサイズが大幅に縮小されたにもかかわらず、説明したデータ拡張メカニズムにより、TinyBERT がより多様な例から学習できるようになり、パフォーマンスに大きな影響を与えます。

データ強化技術の概略図

モデルの設定

わずか 1450 万個のパラメータを持つ TinyBERTモデルは、基本BERTモデルの約 7.5 倍小さくなります。詳細な比較は以下の図に示されています。

基本的なBERTモデルとTinyBERTモデルの比較

レイヤーマッピングについては、著者らは統一された戦略を提案した。この戦略によれば、レイヤーマッピング関数は、各 TinyBERT レイヤーを3 つのBERTレイヤーの最初のレイヤーに順番にマッピングします(g(m)=3*m)。著者らは他の戦略（すべてのボトムまたはトップの BERT レイヤーを使用するなど）も研究しましたが、統合戦略のみが最良の実験結果を示しました。この結論は、異なる抽象化レベルからの知識を転送し、転送される情報をより多様化できるため、より論理的であると思われます。

異なるレイヤーマッピング戦略に基づく: この図は、GLUEデータセットに基づくパフォーマンス比較結果を示しています。

トレーニング実装プロセスに関しては、TinyBERTモデルは英語版Wikipedia（2500M語）でトレーニングされており、そのハイパーパラメータのほとんどはBERTモデルライブラリで使用されているものと同じです。

結論は

トランスフォーマー蒸留は自然言語処理における重要な手段です。 Transformerベースのモデルは現在、機械学習で最も強力なモデルの 1 つであることを考慮すると、 Transformer蒸留を適用して効率的に圧縮することで、さらに活用することができます。その最も優れた例の 1 つは、 BERTモデルよりも7.5倍圧縮されたTinyBERTモデルです。

パラメータが大幅に削減されたにもかかわらず、実験では、TinyBERTモデルのパフォーマンスは基本的にBERT ベースモデルと同等であることが示されています。GLUEベンチマークデータセットのテスト結果では、 TinyBERTモデルが77.0% のスコアを達成しており、これは BERTモデルの79.5% のスコアとそれほど変わりません。明らかに、これは素晴らしい成果です!最後に、量子化やプルーニングなどの他の一般的な圧縮手法をTinyBERTモデル圧縮アルゴリズムに適用して、モデルを小さくすることができます。