紹介する論文

【筆者】Kunbo Ding, Weijie Liu, Yuejian Fang, Zhe Zhao, Qi Ju, Xuefeng Yang, Rong Tian, Zhu Tao, Haoyan Liu, Han Guo, Xingyu Bai, Weiquan Mao, Yudong Li, Weigang Guo, Taiqiang Wu, and Ningyuan Sun.
【タイトル】Multi-stage Distillation Framework for Cross-Lingual Semantic Similarity Matching.
【採択会議】NAACL2022 Findings

どんな論文？

新たなモデル蒸留手法を提案した論文です．
本論文で扱うSTSタスク（後述）はモデルサイズに大きく影響されるそうで，以下の図のようにモデルサイズが大きければ大きいほど性能が高くなる傾向があるらしいです．
しかし，モデルサイズが大きくなればなるほど計算コストや時間的コストが大きくなってしまうため，本論文では性能をなるべく維持しつつモデルサイズを小さくする手法について提案しています．

モデル蒸留とは

サイズの大きい（パラメータ数の多い）モデルを教師モデルとし，その知識を小さいモデル（生徒モデル）の学習に利用する手法です．
これにより教師モデルの性能をなるべく落とさずに生徒モデルへ知識を継承することが期待できます．

モデル蒸留には，教師モデルから生徒モデルにラベル情報を継承する手法やベクトル表現ごと継承する方法など様々な手法がありますが，今回は生徒モデルのベクトルを教師モデルのベクトルに近づけるように学習することで性能を維持したままサイズの小さなモデルを構築します．

タスクについて

モデル蒸留自体は色々なタスク（機械翻訳，読解，文章生成，文章分類）について研究されていますが，この論文では異言語間におけるSTS(Semantic Textual Similarity)タスクを扱っています．
STSタスクとは，文章Aと文章Bが意味的に近いかどうかを評価するタスクです．
今回は同一言語と異言語間でそれぞれSTSタスクを解いています．

技術的な解説

以下が今回のモデル蒸留手法の全体図です．
4段階の学習を行うことで性能を余り落とさずにモデルサイズを50%に引き下げています．

前準備

この論文では3種類の機械学習モデルを使用しているので，まずはそちらの紹介をします．

教師モデル

学習済みのSentence-BERT(通称SBERT)を教師モデルとして使用します．
この教師モデルからなるべく性能を落とさずにモデルサイズのみを小さくすることがこの研究の目標です．また，Sentence-BERTはSTSタスクで高い性能を示しているモデルのため，今回のタスクにぴったりです．
以下Sentence-BERTの情報↓
オリジナルの論文：https://arxiv.org/abs/1908.10084
解説記事：https://data-analytics.fun/2020/08/04/understanding-sentence-bert/

アシスタントモデル

続いてアシスタントモデルです．
モデルの蒸留を行う際に，モデル間のギャップが大きいと性能が低下することが分かっており，そのギャップを補うためのモデルがアシスタントモデルです．(Mirzadeh et al., 2020)
研究者が生徒に直接教えるのではなく，研究者に教わった講師の人が生徒に分かりやすく教える，というイメージです．

そのため，このアシスタントモデルでは教師モデルより少しサイズの小さいモデルを使用します．

生徒モデル

最後に生徒モデルです．
この生徒モデルに知識を付けさせることで，教師モデルに近い性能を得られるように学習させていきます．
生徒モデルはアシスタントモデルの最初から３層目までのTransformerモデルの重みで初期化したものを使用します．
つまり，概略図だと以下の部分ですね．

ただ，異言語間のタスクでは2つ以上の言語を扱うために語彙Vの埋め込み次元が非常に大きくなります．
そこで，埋め込み層Vと隠れ層Hの間に「ボトルネック層」を挿入します．
ボトルネック層では $V × H$ から $V×B+B×H$ に次元を削減することで計算的な負荷を減らします．

Stage 1

まずは教師モデルからアシスタントモデルへ蒸留を行います．
ここは非常に単純で，教師モデルとアシスタントモデルに全く同じ入力$s_i$を行い，それぞれの出力に対して MSE(Mean Squared Error, 平均二乗誤差) を計算します．
加えて$s_i$を他言語に変換したものをアシスタントモデルに入力し，この出力ともMSEを計算します．
このように2つの文章に対してMSEを計算することで，同一言語の知識はもちろん異言語の知識もアシスタントモデルに蒸留する事ができます．
計算したMSEを損失として学習を行います．

l_{stage1} = \frac{1}{|N|}\sum_{i=1}^N[(h_T^{si}-h_A^{si})^2+(h_T^{si}-h_A^{ti})^2]

Stage 2

次にアシスタントモデルから生徒モデルへの蒸留を行います．
ただ，ここで蒸留に用いるのは最終的な出力ではなく，途中にある隠れ層の出力です．
アシスタントモデルと生徒モデルに全く同じ入力$s_i$，$t_i$を使用し，アシスタントモデルの埋め込み層の出力と生徒モデルの隠れ層の出力でMSEを計算します．
何故このように途中のベクトル表現を使ってMSEを計算するのでしょうか？
それは生徒モデルのボトルネック層に理由があります．
生徒モデルのボトルネック層は語彙の次元数を減らすために挿入された未学習の層です．
本来は$V × H$であった語彙の次元数を$V×B+B×H$に削減してしまっているのでパラメータ数が減り，表現力が低下してしまっています．
そこでアシスタントモデルの埋め込み層の出力に近づけることによって，少ない次元数でも余り性能を落とさずに計算が可能になります．
なお，ここでも計算したMSEを損失として学習を行います．

l_{stage2} = \frac{1}{|N|}\sum_{i=1}^N[(h_{Ae}^{si}-h_{Be}^{si})^2+(h_{Be}^{ti}-h_{Ae}^{ti})^2]

Stage 3

今度はアシスタントモデルの最終的な出力と生徒モデルの最終的な出力で蒸留を行います．
ここは非常に単純で，ある言語の文章Aと別の言語だが意味が同じ文章Bをそれぞれアシスタントモデルと生徒モデルに入力し，その出力でMSEを計算する．
そしてそのMSEを最小化するように学習を行います．

l_{stage3} = \frac{1}{|N|}\sum_{i=1}^N[(h_{A}^{si}-h_{S}^{si})^2+(h_{S}^{ti}-h_{A}^{ti})^2]

Stage 4

最後に教師モデルから生徒モデルへの蒸留です．
イメージとしては，生徒が十分育ってきたので教師に直に教えてもらうという感じですね．
まずはStage1と同様に全く同じ文章Aをそれぞれ教師モデルと生徒モデルに入力します．
続いて生徒モデルにのみ文章Aと意味は同じだが言語が違う文章Bを入力します．

次に教師モデルの出力と生徒モデルの出力のそれぞれで全通りの出力に対してコサイン類似度を計算します．
以下の図がわかりやすいと思います．

そしてそれらの全通りのコサイン類似度を生徒モデルと教師モデルのそれぞれで加算して2乗したものを損失$l_1$とします．

l_{1} = \frac{1}{N^2}\sum_{i=1}^N\sum_{j=1}^N(\phi(h_{T}^{si},h_{T}^{sj})+\phi(h_{S}^{si},h_{S}^{tj}))^2

次にStage1と全く同様に教師モデルの出力と生徒モデルの出力のそれぞれでMSEを計算います．

l_{2} = \frac{1}{|N|}\sum_{i=1}^N[(h_{T}^{si}-h_{S}^{si})^2+(h_{T}^{si}-h_{S}^{ti})^2]

Stage4の損失は以下のように$l_1$と$l_2$の和とています．

l_{stage4} = l_{1} + l_{2}

以上で教師モデルから蒸留を行った生徒モデルが完成しました．

手法の評価

最後にここまでの手法の評価です．
評価指標としてはスピアマンの相関係数を使用しています．
方法としては，文章Aに関してデータセット内で似た意味を持つ文章郡に対して，モデルが文章郡に対して出力したどれだけ似ているかのランキングと，元の文章郡のランキングとの相関を計算することで評価を行っています，

同一言語間のSTSタスク

以下は同一言語間におけるSTSタスクで検証した結果です．
「今日はいい天気ですね」と「今日は素敵な天気ですね」の類似度を当てる，みたいなタスクです．

異言語間のSTSタスク

以下は異言語間におけるSTSタスクで検証した結果です．
「こんにちは」と「Hello」の類似度を当てる，みたいなタスクです．

提案手法は通常の蒸留手法に比べて同一言語性能は0.3%，異言語性能は0.9%性能が低下した代わりにエンコーダサイズが75%，埋め込みサイズを83.2%圧縮できました．

まとめ

ここまで読んでいただきありがとうございました．
この論文では異言語間STSのためのモデルを使用していましたが，この4段階のモデル蒸留が他のタスクにも有効かが気になってきますね．