　インターン生の三澤遼（みさわ (@misawann06) | Twitter）です。前回の投稿では言語モデルの事前学習方法について解説しました。

　今回は，図1の中央，言語モデルを社会実装する上で肝腎な計算量・計算コストの削減について，「モデル圧縮」に焦点を当てて解説します！
　前回の記事で紹介した，文生成に長けたGPT-3や様々な読解タスクで State-of-the-Art (SOTA) を達成したT5を筆頭に，高性能な言語モデルは非常に規模が大きく社会実装が困難とされています。例えば，前回の記事でも言及しましたが，GPT-3の事前学習には約5億円かかると言われています。また，学習済みのモデルの実運用コストも非常に高いです。こういった課題を解消することを目的とした研究や具体的なモデルを紹介していきます。

f:id:elyza_ryomisawa:20210430172204p:plain — 図1　事前学習済みモデルの分類

　前回と同様，TransformerやBERTについて知識があると分かりやすいと思います。また，前回の記事とは独立しているので，そちらをまだ読まれていない方にも一読して欲しいです！

1 量子化
- 1.1 QBERT
2 枝刈り
- 2.1 CompressingBERT
3 蒸留
- 3.1 DistilBERT
4 モジュール交換
- 4.1 BERT-of-Theseus
5 パラメーター共有
- 5.1 ALBERT
最後に
ELYZAでは仲間を募集しています！
参考文献

1 量子化

　量子化と聞いて量子力学や量子コンピュータを連想するかもしれませんが，ここでの量子化はデジタル化（離散化）と同義です。端的に言うと，浮動小数点数を整数に変換することを指します。
　ナイーブな方法として，kビットを割り当てる場合，テンソル内に現れる浮動小数点数の値の範囲を等間隔に分割し（これを「量子化の範囲」と呼ぶことにします），それぞれの区間に含まれる数を0, 1, ..., 2^k-1のうち1つで表現します。これにより情報が整数という離散的な情報になるためパラメータ数を削減できます。またビット数kが小さくなるほどパラメータ数は減少します。

1.1 QBERT

　量子化は主にコンピュータビジョンの分野で用いられていた手法で，BERTに量子化を適用したのは本論文が初めてです。
　本論文では，層ごとに適切なビット数を割り当てるmixed precision quantizationをBERTに適応するとともに，12個のアテンションヘッドを細かくグループ化し各グループで量子化の範囲を決めるgroup-wise quantizationを提案しました。

mixed precision quantization
　まず，mixed precision quantizationを導入するきっかけとなった分析を紹介します。
　各層におけるビット数の影響をヘッセ行列（各層における勾配を偏微分した関数を並べた行列程度の認識で大丈夫です）の固有値を求めることにより調べました。ヘッセ行列を調べる理由として2点挙げられます。

その固有値を見ることで局所最適解に収束しているか分かります
- ある点においてヘッセ行列が正定値対称行列であること（今回の場合は固有値が全て正であること）とその点において局所最適解（極小値）に収束していることは同値であるためです
- 負の固有値が存在すると，極小値に収束していないことが分かります
ヘッセ行列の固有値が大きい層ほどより多くのビット数が必要であるということが知られています
- 各層に割り当てるべきビット数を明らかにできます

　図2から分かるように，全ての層は同じ構造であるにもかかわらず，固有値の大きさは層ごとに違っており，層に応じたビット数を割り当てる必要があることが分かります（eigenvalue: 固有値）。
　興味深いのは，後でも言及しますが，量子化後のSQuADにおける性能がかなり低下する理由を，ヘッセ行列の分析で理解する事ができることです。図2からも分かるように，SQuADに適応したモデルにはヘッセ行列の固有値が負となる層が存在します。ここから，SQuADにおいて極小値に収束しておらず上手く学習できていないことが分かります。

f:id:elyza_ryomisawa:20210429221651j:plain — 図2　4つのデータセットでのファインチューニング後の各層におけるヘッセ行列の上位の固有値

group-wise quantization
　次にグループごとに量子化の範囲を変えるgroup-wise quantizationを図を用いて説明します。
　まず， $d$ を埋め込みベクトルの次元， $N_h$ をアテンションのヘッド数とします。図3の立方体は，アテンションの $\frac{d}{N_h}×d$ 行列を $N_h$ 個並べたテンソルを表現しています。通常の方法 (layer-wise quantanization) ではテンソル全体で計算された量子化の範囲を用います。それに対し，group-wise quantizationでは各アテンションヘッドを一つのグループとし（図3 (b) ），それぞれのグループで量子化の範囲を計算します。より一般的な場合ではグループ数を更に増やします（図3 (c) ）。
　ヘッドによって行列成分の値の範囲が異なると考えられるため，ヘッド全体で量子化の範囲を計算するのは不適切であり情報がより失われてしまいますが，グループ化によってこの問題を軽減することができます。

f:id:elyza_ryomisawa:20210429230049j:plain — 図3　layer-wise quantizationとgroup-wise quantizationの図例

　結果は図4の通りです。w-bitsは重み (weight) に，e-bitsは埋め込み (embedding) に用いるビット数を表しています。ここでは，Baseline（通常のBERT），DirectQ（group-wise quantizationとmixed precision quantizationを用いないモデル），QBERT（group-wise quantizationのみを用いるモデル），QBERT MP（group-wise quantizationとmixed precision quantizationを用いるモデル）を比較しています。図4から分かるように，QBERTは全てのデータセットにおいてDirectQを大幅に上回っていて，ビット数が少なくなるほどその差は大きくなっています。
　更に，mixed precision quantizationも加えることで性能は向上します。例えば，SST-2という感情分析のデータセットでは，Q-BERTのw-bitsが2であるとき，QBERT MPはモデルサイズがたった5MBしか増えないにも関わらず精度は7ポイント以上増加しています。
　また，先ほど言及したように，SQuADでは量子化によって性能が大幅に低下しています。　

f:id:elyza_ryomisawa:20210429230751j:plain — 図4　4つのデータセットにおける性能の比較

2 枝刈り

　枝刈りとは，モデルの不要な部分を取り除く方法のことです。パラメーター数を減らしモデル圧縮するだけでなく，モデルを正則化する効果もあります。先行研究では，枝刈りをdropoutやL0正則化と見なすことができるとされています。

2.1 CompressingBERT

　CompressingBERTは，magnitude weight pruningという枝刈り手法を採用しました。これは，値が0に近い重みを不要と見なし取り除く方法です。具体的には，以下の5ステップを行います。

　1. 取り除く重みの割合を決めます
　2. 決めた割合の重みが取り除かれるように，重みの閾値を決めます
　3. 閾値以下の重みの値を0にします（情報消去: information deletion）
　4. 性能が枝刈り前と同等になるまで学習します
　5. 必要であれば，ステップ1に戻り枝刈りをする重みの割合を増加させます

　なお，枝刈りで値が0になった重みの値を固定しその後更新されないようにします（正則化: regularization）。情報消去と正則化の寄与を調べるため，正則化を行わない，つまり値が0になった重みを下流タスク時に更新できるようにする方法（単なる情報消去）とも比較した結果が図5です。
　図5はGLUEというベンチマークの検証データセットにおける枝刈りとその他の手法の比較結果を示しています。縦軸は，左のグラフでは検証データにおける精度，右では訓練損失を表していて，横軸はともに枝刈りする重みの割合を示しています。この図から，効果的な枝刈りの割合は30-40%である事が分かります。この程度の枝刈り率の場合，枝刈りをしていないモデルや情報消去を使ったモデルと性能ははほとんど変わりません。しかし，これ以上の重みを取り除くと性能は減少していきます。

f:id:elyza_ryomisawa:20210410162248j:plain — 図5　各種枝刈りの効果

　ちなみに，magnitude weight pruningは，Tensorflowのパッケージを使って実装できます。

www.tensorflow.org 　

3 蒸留

　蒸留とは，サイズの小さい生徒モデルが，サイズが大きく十分に事前学習された教師モデルと同じ出力をするように訓練し，教師モデルと同程度の性能の生徒モデルを作ることを目指す方法です。具体的には，以下のような蒸留損失関数 (distillation loss) を最適化することで，蒸留を行います。

$\displaystyle L_{\rm ce} = \sum_i t_i \log s_i$

　ここで， $L_\rm{ce}$ はクロスエントロピーロス (Cross Entropy) であることを明記していて， $t_i$ , $s_i$ はそれぞれ教師 (Teacher) モデル，生徒 (Student) モデルが $i$ 番目のデータに割り振る確率を表しています。
　今回は詳しく解説しませんが，蒸留はモデル圧縮だけでなくマルチタスク用モデルの学習にも使われます。それぞれのタスクに特化した複数の教師モデルを単一の生徒モデルが模倣することで，生徒モデルが複数のタスクを解けるようになります。マルチタスク学習という用途では，蒸留よりも強力なteacher annealingという手法も提案されています。こちらに関しては，弊社のTwitterで紹介しているので是非チェックしてください。

【生徒が先生を超える蒸留手法】
特定タスクに特化した複数モデルの蒸留から通常の教師あり学習へと徐々に移行するteacher annealingを提案。通常の蒸留とは異なり，先生（特化モデル）を生徒（マルチタスクモデル）が超えることができる。https://t.co/0kzOWBbq4n pic.twitter.com/HWXyAqRcIJ
— ELYZA, Inc. (@ELYZA_inc) 2021年4月7日

3.1 DistilBERT

　蒸留がBERTへ適応されたのはDistilBERTが初めてでした。
　図6, 7から分かるように，ほぼ全てのタスクでELMoを上回りました。更に，BERTと比べてサイズを40%削減し60%早く訓練ができつつも，BERTの言語理解能力の97%に匹敵する能力を獲得しました。

f:id:elyza_ryomisawa:20210410162427j:plain — 図6　DistilBERTの性能

f:id:elyza_ryomisawa:20210410164050j:plain — 図7　DistilBERTの軽量性

　ちなみに，HuggingFaceで日本語を含む多言語DistilBERTを使うことができます。図8では4/28のYahooニュースの記事の一文を入力していて[MASK]の部分を「雨」と予測できていますね。

f:id:elyza_ryomisawa:20210428220447j:plain — 図8　HuggingFaceのModel HubでのDistilBERTの実行例

huggingface.co

4 モジュール交換

　モジュール交換とは，学習済み言語モデルの各モジュールに対して，よりサイズの小さいモジュールを準備し小さいモジュールに置換しながら元々のモデルと同程度の性能を目指して学習する方法です。
　サイズの大きいモデルを模倣するという点において蒸留と似ていますが，蒸留とは違い，新たな損失関数が不要です。通常，蒸留損失関数に加え他の損失関数（生徒モデルの出力とラベルに関するクロスエントロピーなど）をそれぞれ重みづけしたものを最適化しますが，モジュール交換ではその重みを探索する負担が軽減されます。

4.1 BERT-of-Theseus

　モジュール交換はBERT-of-Theseusの論文で提案されました。
　図9のように，元々のモデル（今回は6層）の層を，今回はそれぞれ2層からなる複数のモジュールpred 1, 2, 3 (pred: predecessor) に分割し，各ステップにおいて確率pでそれぞれ1層からなるscc 1, 2, 3 (scc: sucsessor) と交換します。

f:id:elyza_ryomisawa:20210428170924j:plain — 図9　モジュール交換の参考図

　BERT-of-Theseusでは，図10 (b) のようにモジュールを交換する確率を訓練が進むとともに増加させます。つまり訓練序盤では元のモデルの補助を受けながら訓練が進み，終盤になるにつれ元のモデルからは独立して学習します。こうすることによって，難易度の低いステージ1（元のモデルのモジュールが残っている状態での学習）から難易度の高いステージ2（モデルのモジュールが全て交換された状態での学習）への移行が連続的になります。

f:id:elyza_ryomisawa:20210428234256j:plain — 図10　モジュールを交換する確率を徐々に増加させる

　実際，図11から分かるように，この方法は効果的であることが実験的に示されました。ここでは，モジュール交換確率が一定である場合と増加させる場合 (carriculum) に加え，減少させる場合 (anti-carriculum) も比較しています。この比較によりcarriculumにより性能が上がった要因が，モジュール交換確率を増加させたことなのか，それとも不定にしたことなのかを明らかにすることができます。

f:id:elyza_ryomisawa:20210429102837j:plain — 図11　モジュール交換確率を一定，減少，増加させる場合を比較

　また，図12は他モデルとの比較結果を示しています。PD-BERT, Vanilla KD, BERT-PKDは全て蒸留を用いたモデル (PD: Pre-trained Distillation, KD: Knowledge Distillation, PKD: Patient Knowledge Distillation) で，BERT-of-Theseusはこれらを上回りました。

f:id:elyza_ryomisawa:20210410162505j:plain — 図12　BERT-of-Theseusの性能

　ちなみに，BERT-of-Theseusの名前はテセウスの船に由来しています（図13）。これは物体の構成要素を全て入れ替えても本質的に同じかどうかというアイデンティティーに関する思考実験のことです。

en.wikipedia.org

f:id:elyza_ryomisawa:20210430130252p:plain — 図13　BERT-of-Theseusはテセウスの船という思考実験に由来している

5 パラメーター共有

　パラメーター共有とは，名前の通り，複数の層にわたって同じパラメーターを使うことでモデルのパラメーター数を削減する方法です。

5.1 ALBERT

　ALBERTでは，パラメーターが全層で共有されています。更に，factorized embedding parameterizationという埋め込み行列 (embedding matrix) のパラメーター数を削減する方法が取られました。
　また，前回の記事で紹介したSentence Order Predictionを導入し文章間の関係性を理解する能力が強化されました。
　図14はALBERTとBERTの比較結果を示しています。最もサイズの大きいALBERT-xxlargeは，BERT-largeよりパラメーター数が30%少ないにも関わらず，その性能を上回りました。更に，単位時間あたりの処理能力が上がり，ALBERT-largeは，BERT-largeの1.7倍の計算速度となりました。

f:id:elyza_ryomisawa:20210410162539j:plain — 図14　ALBERTの性能

最後に

　今回は計算量・計算コストの削減についてモデル圧縮の観点で解説しました。
　計算量・計算コストの削減が進む一方で，HuggingFaceを筆頭に大規模な事前学習モデルを「民主化」する動きもあります。最近の話ですと，今年の3月にEleutherAIという団体がGPT-3を再現したGPT-Neoという言語生成モデルを一般公開しました。こちらも再学習したりHuggingFaceでデモを試すこともできます。

huggingface.co

ELYZAでは仲間を募集しています！

　ELYZAでは、「未踏の領域で、あたりまえを創る」をミッションに掲げ、自然言語処理（NLP）・リテールテックの2領域にフォーカスして、「研究開発」・「共同研究」・「サービス開発」の3つの活動を行っております。
　研究開発に興味がある方も、機械学習を活用したサービス開発に興味がある方も、是非お気軽にご連絡ください。

open.talentio.com

brain.elyza.ai

note.com