Open shimopino opened 4 years ago
[arXiv:2004.02984] MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, Denny Zhou
2020-04-06
通常のBERTを圧縮し高速化を行ったMobileBERTを提案した。MobileBERTはボトルネック構造を有し、自己注意機構と順伝播をうまくバランスさせることでBERT-largeモデルを小さくしている。学習させる際には教師モデルとしてBERT-largeを学習させ、学習できた知識の転移をMobileBERTに対して行う。
近年提案されているBERTベースのモデルは非常に多くのパラメータを必要とし、レイテンシも高いためモバイル端末などのハードウェア制約がある環境に適用することが難しい。
また計算コストを落としたり圧縮させたモデルは提案されているが、どれもタスクに依存するモデルを生成する方法であり、タスクに依存しない事前学習済みモデルを作成することができていない。
本研究で提案されているMobileBERTは、BERT-Largeと同程度の深さを有しているが、自己注意機構と順伝播計算をバランスさせることで、モデルサイズを落としている。また学習を行う際には、まずは教師モデルとしてBERT-Largeを学習させ、MobileBERTに対して知識の転移を行う。
MobileBERTは既存のモデルに対して同程度の精度を発揮しながらも、約4倍サイズが小さく、約5倍の計算速度を有している。
MobileBERTモデル宇野設定は以下になる。なおこのモデルはArchitecture Searchを行った結果得られたものである。
上記の表を見てわかるようにMobileBERTで使用している隠れ層の次元数はたったの128である。MobileBERTは層数が深いが隠れ層の次元数は小さく、こういったモデルを学習させることが非常に難しい。
そこで入力と出力部分に次元数を512に調整する線形結合層を導入しているIB-BERTを教師モデルとして知識の転移を行う。次元数を512で合わせているため、MobileBERTと出力される特徴量をそのまま比較することができる。
なお入力と出力に線形結合層を導入しているボトルネックと、導入していないボトルネックを組み合わせることで、IB-BERTの精度を高く保ったままMobileBERTを圧縮させることができる。
元のBERT論文ではMulti-Head Attention(MHA)とFeed-Forward Network(FFN)のバランスは1：2に設定されている。MHAではそれぞれのヘッドが特徴量の中の異なる位置に関する異なる部分空間を処理しており、順伝播層ではモデルに非線形性を導入している。
しかしMobileBERTではMHAの特徴量のサイズのほうが大きくなってしまっている。このアンバランスさを解消するために、上記の表でわかるように順伝播部分を4つ積み重ねて計算を行っている。
今までの研究によりLayerNormalizationとGeLUがLatencyの大部分を占めることがわかっている。そこでこの2つを異なる計算に置き換える。
LayerNormalizationを削除
要素ごとの線形変換に変更。
ReLUに変更
元のBERTでモデルサイズの大部分は単語のベクトル化を行うEmbedding層が担っている。そこでMobileBERTではEmbedding層の次元数を128まで減らし、その後でカーネルサイズが3の1D畳み込み層を使用して次元数を512まで増大させている。
知識の転移に関しては層ごとの特徴量に対して行う。その際に特徴マップとAttention層の転移を行う。
BERTでは単に1つ前の層の出力を受け取り計算を行っているだけなので、出力される特徴マップがMobileBERTと教師モデルで近くなるように転移を行う。
Attention層の転移は、教師モデルとMobileBERTのMHAの各ヘッドでのAttention層の分布間距離をKL-Divergenceで計算しこれを最小化していく。
またMobileBERTの事前学習を行う際に、知識の蒸留を使用した損失関数を適用する。ここではMLM損失とNSP損失を使用し、新しいMLMに対する知識の蒸留に対する損失関数を使用する。
知識の蒸留を行う際に、知識の転移を追加のタスクとして損失関数を計算する。
知識の転移を行う場合、途中状態のIB-BERTの教師モデルがMobileBERTに対して最適とは限らないため、MobileBERTを学習させる際は最初に知識の転移の損失関数を採用し、その後で知識の蒸留の損失関数を適用する。
知識の転移を行う際、浅い層でのエラーは深い層にまで影響を与えてしまうため、各層ごとに知識の転移を行う。
上記3つの学習戦略は以下の図に表現される。JKTとPKTでは最初のEmbedding層と最後の分類層では知識の転移は行わず、単にIB-BERTモデルのパラメータをコピーしているだけである。
まずはどのようにIB-BERTのモデルを決めたのか議論していく。目的としては隠れ層の次元数を可能な限り小さくすることなので、BERT-Largeと同じサイズから初めて徐々に次元数を下げていきながら精度が大きく減少しない構造を探索した。
またMobileBERTの構造に関して、目的としてはBERT-BASEの4分の1のサイズであるため25M程度のサイズであり、MHAとFFNの最適なバランスを持つネットワークを探索した。
得られたMobileBERTに対してBookCorpusとWikipediaデータを使用して事前学習を行い、この後のタスクに関しては事前学習済みのMobileBERTに対してfine-tuningを行い性能評価を行った。
まずはGLUEベンチマークを行った。タスクに依存しないモデルでありながら、高い性能を発揮しており、また低レイテンシを実現できている。
読解力を測るSQuADデータセットに対して性能を評価した。どのモデルに対してもMobileBERTがより高い精度を発揮していることがわかる。
また同じSQuADデータセットに対して8bit量子化を行った際の性能の変化を検証した。MobileBERTのサイズを4分の1に圧縮しながらも精度の悪化が見られない。
LayerNormalizationとGeLUを削除することでどの程度レイテンシが向上するのか検証した。実応用のためにはこれらの計算を削除したほうがいいことがわかる。
提案している3つの学習戦略に対してどの戦略が精度向上に貢献しているのか検証した。AKTがほかの2つの手法に対して精度の面で劣っており、これは層ごとの転移を途中で行ってしまうと、生徒モデルに対して適したものではなく、追加の蒸留を行う必要があるためだと考えられる。
どの転移の手法がより精度に貢献しているのか検証した。特徴マップの転移が最も性能に貢献していることがわかる。
https://github.com/google-research/google-research/tree/master/mobilebert
openreview
shimopino
commented
4 years ago shimopino
commented
4 years ago
論文へのリンク
[arXiv:2004.02984] MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
著者・所属機関
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, Denny Zhou
投稿日時（YYYY-MM-DD）
2020-04-06
1. どんなもの？
通常のBERTを圧縮し高速化を行ったMobileBERTを提案した。MobileBERTはボトルネック構造を有し、自己注意機構と順伝播をうまくバランスさせることでBERT-largeモデルを小さくしている。学習させる際には教師モデルとしてBERT-largeを学習させ、学習できた知識の転移をMobileBERTに対して行う。
2. 先行研究と比べてどこがすごいの？
近年提案されているBERTベースのモデルは非常に多くのパラメータを必要とし、レイテンシも高いためモバイル端末などのハードウェア制約がある環境に適用することが難しい。
また計算コストを落としたり圧縮させたモデルは提案されているが、どれもタスクに依存するモデルを生成する方法であり、タスクに依存しない事前学習済みモデルを作成することができていない。
本研究で提案されているMobileBERTは、BERT-Largeと同程度の深さを有しているが、自己注意機構と順伝播計算をバランスさせることで、モデルサイズを落としている。また学習を行う際には、まずは教師モデルとしてBERT-Largeを学習させ、MobileBERTに対して知識の転移を行う。
MobileBERTは既存のモデルに対して同程度の精度を発揮しながらも、約4倍サイズが小さく、約5倍の計算速度を有している。
3. 技術や手法の"キモ"はどこにある？
3.1 MobileBERT
MobileBERTモデル宇野設定は以下になる。なおこのモデルはArchitecture Searchを行った結果得られたものである。
3.2 Bottleneck and Inverted-Bottleneck
上記の表を見てわかるようにMobileBERTで使用している隠れ層の次元数はたったの128である。MobileBERTは層数が深いが隠れ層の次元数は小さく、こういったモデルを学習させることが非常に難しい。
そこで入力と出力部分に次元数を512に調整する線形結合層を導入しているIB-BERTを教師モデルとして知識の転移を行う。次元数を512で合わせているため、MobileBERTと出力される特徴量をそのまま比較することができる。
なお入力と出力に線形結合層を導入しているボトルネックと、導入していないボトルネックを組み合わせることで、IB-BERTの精度を高く保ったままMobileBERTを圧縮させることができる。
3.3 Stacked Feed-Forward Networks
元のBERT論文ではMulti-Head Attention(MHA)とFeed-Forward Network(FFN)のバランスは1：2に設定されている。MHAではそれぞれのヘッドが特徴量の中の異なる位置に関する異なる部分空間を処理しており、順伝播層ではモデルに非線形性を導入している。
しかしMobileBERTではMHAの特徴量のサイズのほうが大きくなってしまっている。このアンバランスさを解消するために、上記の表でわかるように順伝播部分を4つ積み重ねて計算を行っている。
3.4 Operational Optimizations
今までの研究によりLayerNormalizationとGeLUがLatencyの大部分を占めることがわかっている。そこでこの2つを異なる計算に置き換える。
LayerNormalizationを削除
要素ごとの線形変換に変更。
ReLUに変更
3.5 Embedding Factorization
元のBERTでモデルサイズの大部分は単語のベクトル化を行うEmbedding層が担っている。そこでMobileBERTではEmbedding層の次元数を128まで減らし、その後でカーネルサイズが3の1D畳み込み層を使用して次元数を512まで増大させている。
3.6 Training Objectives
知識の転移に関しては層ごとの特徴量に対して行う。その際に特徴マップとAttention層の転移を行う。
Feature Map Transfer (FMT)
BERTでは単に1つ前の層の出力を受け取り計算を行っているだけなので、出力される特徴マップがMobileBERTと教師モデルで近くなるように転移を行う。
Attention Transfer (AT)
Attention層の転移は、教師モデルとMobileBERTのMHAの各ヘッドでのAttention層の分布間距離をKL-Divergenceで計算しこれを最小化していく。
Pre-training Distillation (PD)
またMobileBERTの事前学習を行う際に、知識の蒸留を使用した損失関数を適用する。ここではMLM損失とNSP損失を使用し、新しいMLMに対する知識の蒸留に対する損失関数を使用する。
3.7 Training Strategies
Auxiliary Knowledge Transfer
知識の蒸留を行う際に、知識の転移を追加のタスクとして損失関数を計算する。
Joint Knowledge Transfer
知識の転移を行う場合、途中状態のIB-BERTの教師モデルがMobileBERTに対して最適とは限らないため、MobileBERTを学習させる際は最初に知識の転移の損失関数を採用し、その後で知識の蒸留の損失関数を適用する。
Progressive Knowledge Transfer
知識の転移を行う際、浅い層でのエラーは深い層にまで影響を与えてしまうため、各層ごとに知識の転移を行う。
Diagram of three strategies
上記3つの学習戦略は以下の図に表現される。JKTとPKTでは最初のEmbedding層と最後の分類層では知識の転移は行わず、単にIB-BERTモデルのパラメータをコピーしているだけである。
4. どうやって有効だと検証した？
まずはどのようにIB-BERTのモデルを決めたのか議論していく。目的としては隠れ層の次元数を可能な限り小さくすることなので、BERT-Largeと同じサイズから初めて徐々に次元数を下げていきながら精度が大きく減少しない構造を探索した。
またMobileBERTの構造に関して、目的としてはBERT-BASEの4分の1のサイズであるため25M程度のサイズであり、MHAとFFNの最適なバランスを持つネットワークを探索した。
得られたMobileBERTに対してBookCorpusとWikipediaデータを使用して事前学習を行い、この後のタスクに関しては事前学習済みのMobileBERTに対してfine-tuningを行い性能評価を行った。
まずはGLUEベンチマークを行った。タスクに依存しないモデルでありながら、高い性能を発揮しており、また低レイテンシを実現できている。
読解力を測るSQuADデータセットに対して性能を評価した。どのモデルに対してもMobileBERTがより高い精度を発揮していることがわかる。
また同じSQuADデータセットに対して8bit量子化を行った際の性能の変化を検証した。MobileBERTのサイズを4分の1に圧縮しながらも精度の悪化が見られない。
LayerNormalizationとGeLUを削除することでどの程度レイテンシが向上するのか検証した。実応用のためにはこれらの計算を削除したほうがいいことがわかる。
提案している3つの学習戦略に対してどの戦略が精度向上に貢献しているのか検証した。AKTがほかの2つの手法に対して精度の面で劣っており、これは層ごとの転移を途中で行ってしまうと、生徒モデルに対して適したものではなく、追加の蒸留を行う必要があるためだと考えられる。
どの転移の手法がより精度に貢献しているのか検証した。特徴マップの転移が最も性能に貢献していることがわかる。
5. 議論はあるか？
6. 次に読むべき論文はあるか？
論文情報・リンク