# 敵対的学習を用いた知識蒸留への中間層蒸留と対照学習の導入

鈴木 偉士 山田寛章 徳永 健伸
東京工業大学 情報理工学院
\{suzuki.t.dp@m, yamada@c, take@c\}.titech.ac.jp

## 概要

知識蒸留 (KD) とは，大規模なニューラルネット ワークを圧縮する手法の一つである。言語モデル向 け KD の中で最高性能の手法は，敵対的学習に中間層出力と対照学習を導入した CILDA と呼ばれる手法である. CILDA の学習は最大化ステップと最小化ステップに分かれているが，中間層出力と対照学習は最大化ステップでのみ活用されている. 本研究 では，最小化ステップに中間層蒸留と対照学習を 導入し，性能を向上させることを目指した.しか し, 既存手法に対して有意な差は確認できなかった ため，原因分析のために CILDA 単体の再現実験を 行ったところ, 先行研究の主張とは異なり, GLUE における複数のタスクで CILDA がそれ以前の手法 の性能を上回らないという結果を得た。

## 1 はじめに

BERT[1] やRoBERTa[2] のような事前学習済み言語モデルは, 様々な言語処理タスクにおいて, 高 い性能を発揮しているが，近年では，GPT-3[3] や， PaLM[4] に見られるように， パラメータ数が大幅に 増加し，通常のデバイスでは扱えなくなっている。

知識蒸留 (KD)[5] は，ニューラルネットワークを 圧縮する手法の一つである。 パラメータ数の大きい 教師モデルの出力と, パラメータ数の小さい生徒モ デルの出力の KL ダイバージェンスを小さくする学習をすることで，生徒モデルが教師モデルと同様の 出力が出来るようにする. 先行研究により, 知識蒸留は, 教師モデルと生徒モデルの最終的な出力間だ けでなく, 中間層の出力間の誤差の利用 $[6,7,8]$, データ抳張と敵対的学習の導入 [9], 対照学習の導入 [10] により，性能が向上することが報告されて いる. Haidar ら [11] は知識蒸留に中間層の利用, 敵対的学習, 対照学習を導入した CILDA と呼ばれる 手法を提案し，GLUE タスクにおいて SOTAを達成 した.
しかし, Haidar らの研究 [11]では, データ拡張と 敵対的学習に用いられる Generator の学習にのみ中間層の利用と対照学習の導入を行っており, 生徒モ デルの学習については通常の KD から改良されてい ない. 本研究では CILDA の生徒モデルの学習時に 中間層の利用と対照学習を導入し，さらに性能を向上させることを目指す.

本研究のソースコードは $\mathrm{GitHub}^{11}$ にて公開して いる.

## 2 関連研究

## 2.1 知識蒸留 (Vanilla KD)

知識蒸留 [5] とは，モデル圧縮の手法の 1 つであ る.この手法では， パラメータ数が小さい生徒モ デルがタスクの正解ラベルとの誤差のほかに， パラ メータ数が大きく性能の高い教師モデルの出力との 誤差を用いて学習する。損失関数は以下のとおりで ある。

$
\begin{array}{r}
L=\lambda L_{C E}+(1-\lambda) L_{K D} \\
L_{K D}=K L(\text { Teacher }(X), \operatorname{Student}(X))
\end{array}
$

ただし， $L_{C E}$ は交差エントロピー誤差， $K L($,$) は$ KL ダイバージェンス, Teacher $(X)$, Student $(X)$ は それぞれ，教師モデル，生徒モデルにデータ $X$ を 入力した際の出力する確率分布とする。学習対象は 生徒モデルのみであり，教師モデルのパラメータは 固定される.

## 2.2 RAIL-KD

RAIL-KD[8] は通常の知識蒸留に加えて中間層 の出力の情報も用いる中間層蒸留の一種である. RAIL-KD では，各エポックごとに教師モデルから 生徒モデルの中間層の数だけ前から順にランダムに 中間層を選び，生徒モデルの中間層との誤差を計算


する．誤差の計算の方法には以下の 2 通りがある.

$
\begin{gathered}
L_{R A I L-K D^{l}}=\sum_{x \in X} \sum_{k=0}^{m}\left.\|\frac{\hat{h}_{a_{\kappa}, x}^{T}}{\left.\|\hat{h}_{x}^{T}\right.\|_{2}}-\frac{\hat{h}_{\kappa, x}^{S_{\theta}}}{\left.\|\hat{h}_{x}^{S_{\theta}}\right.\|_{2}}\right.\|_{2}^{2} \\
L_{R A I L-K D^{c}}=\sum_{x \in X}\left.\|\frac{\hat{h}_{x}^{T}}{\left.\|\hat{h}_{x}^{T}\right.\|_{2}}-\frac{\hat{h}_{x}^{S_{\theta}}}{\left.\|\hat{h}_{x}^{S_{\theta}}\right.\|_{2}}\right.\|_{2}^{2}
\end{gathered}
$

$\hat{h}_{*}^{*}, X, \kappa, a$ はそれぞれモデルの中間層の出力, 訓練 データ, 生徒モデルの中間層の数, 選択した中間層 のインデックスの配列である． $T$ は教師モデル， $S_{\theta}$ は生徒モデルを表す.

$L_{R A I L-K D^{l}}$ では, 1 層ずつの中間層同士の誤差の 和をとり， $L_{R A I L-K D^{c}}$ では，選択した層を連結して から誤差をとる．教師モデルと生徒モデルの中間層 の出力の次元数が異なる場合は損失関数の計算前に $h$ に対して線形変換を行い，次元数を揃える.

最終的な損失関数は以下のようになる.

$
L=\lambda_{1} L_{C E}+\lambda_{2} L_{K D}+\lambda_{3} L_{R A I L-K D}
$

## 2.3 MATE-KD

MATE-KD[9] は知識蒸留に敵対的学習を導入した 手法である. この手法では学習は最大化ステップと 最小化ステップの 2 つのステップに分かれている.

最大化ステップでは，まず訓練データに一定の確率でマスクをかける。 その後 Generator と呼ばれる 事前学習済みマスク型言語モデルを用いてマスクさ れた部分を埋めることでデータ拡張を行う. 拡張さ れたデータを用いて以下の損失関数を最大化するよ うに Generatorを学習する。

$
L_{A D V}=K L\left(\text { Teacher }\left(X^{\prime}\right), \text { Student }\left(X^{\prime}\right)\right)
$

ただし， $X^{\prime}$ は拡張されたデータである. 最大化ス テップでは生徒モデルのパラメータは固定される.

最小化ステップでは以下の損失関数を最小化する ように生徒モデルを学習する. 最小化ステップの間 は Generator のパラメータは固定される.

$
L=\lambda_{1} L_{C E}+\lambda_{2} L_{K D}+\lambda_{3} L_{A D V}
$

学習時は最大化ステップを $n_{G}$ ステップ行ってから 最小化ステップを $n_{S}$ ステップ行うことを繰り返す。

## 2.4 CILDA

CILDA[11] は知識蒸留に中間層出力の利用，敵対的学習, 対照学習を導入した手法である. MATE-KD と同様に最大化ステップと最小化ステップに分かれ ている.最大化ステップでは以下の損失関数を最大化 する。

$
\begin{array}{r}
L_{G}=\alpha_{1} L_{A D V}+\alpha_{2} L_{C R D} \\
L_{C R D}=-\log \frac{\exp \left(<\bar{h}_{k}^{T}, \bar{h}_{k}^{S_{\theta}}>/ \tau_{2}\right)}{\sum_{j=0}^{K} \exp \left(<\bar{h}_{j}^{T}, \bar{h}_{k}^{S_{\theta}}>/ \tau_{2}\right)}
\end{array}
$

ただし， $\bar{h}_{*}^{*}$ はデルの全ての中間層の出力を連結し 線形変換を施したもの， $k$ はミニバッチの中のデー タのインデックス，K はバッチサイズ，<，>はコ サイン類似度とする。

最小化ステップでは以下の損失関数に拡張された データと元データを入力して最小化する.

$
L=\lambda_{1} L_{C E}+\lambda_{2} L_{K D}
$

## 3 提案手法

提案手法を本論文では CILDA+minILD と呼ぶ. CILDA+minILDでは， CILDA の最小化ステップに中間層出力の利用と対照学習を導入する。そのために 損失関数に $L_{C R D}$ を追加する．損失関数は以下のよ うになる。

$
L=\lambda_{1} L_{C E}+\lambda_{2} L_{K D}+\lambda_{3} \frac{1}{\log K} L_{C R D}
$

$K$ はバッチサイズである. $L_{C R D}$ の値はバッチサイ ズの対数に依存し，バッチサイズが大きいほど值が 大きくなる.そのため $L_{C R D}$ の係数に $\frac{1}{\log K}$ をかけ る. 最大化ステップにおいても同様に $L_{C R D}$ の係数 に $\frac{1}{\log K}$ をいる。

## 4 実験

## 4.1 実験設定

## 4.1.1 データセット

蒸留対象のタスクは GLUE[12] の中から，学習に 要する時間の短い CoLA，MRPC，RTE，STS-B の 4 タスクを利用する。評価指標は先行研究 [11] に倣 い，それぞれマシューズ相関係数，f1 スコア，正解率，ピアソンの相関係数とする.

## 4.1.2 モデル

教師モデルには, RoBERTa-large[2] を元に, 各タ スク毎に教師モデルをファインチューニングし, dev set において各評価指標で最高性能を示したモ デルを用いた. ファインチューニングは，バッチ サイズ $128 ， 3$ エポックとして，異なる 4 つの学
習率 1e-05，2e-05，5e-05，1e-04 について各 5 回行った. 生徒モデルには DistilRoBERTa[13]を使用 した. MATE-KD, CILDA, CILDA+minILDで用いる Generator には RoBERTa-large[2]を用いた.

## 4.1.3 比較手法

提案手法と比較する手法は生徒モデルのみを用 いた学習 (w/o KD), 通常の知識蒸留 (Vanilla KD), RAIL-KD，MATE-KD，CILDAを用いた. RAIL-KD の $L_{R A I L-K D}$ には $L_{R A I L-K D^{c}}$ を用いた。

## 4.1.4 ハイパーパラメータ

学習率を 1e-05，2e-05，5e-05，1e-04 の 4 通りか らハイパーパラメータ探索で選択し，バッチサイズ はw/o KD，Vanilla KD，RAIL-KD，MATE-KDでは 128, CILDA, CILDA+minILDでは 64 とした. 損失関数の係数は, Vanilla KDでは $\lambda$ は $1 / 2$, RAIL-KD と MATE-KD の $\lambda_{1}, \lambda_{2}, \lambda_{3}$ はすべて $1 / 3$ とした. CILDA の $\lambda_{1}, \lambda_{2}$ は元データを入力する際にはともに $1 / 3$ とし，拡張されたデータに対してはそれぞれ $2 / 9$, $1 / 9$ とした. $\alpha_{1}, \alpha_{2}$ はともに $1 / 2, \tau_{2}$ は 2 とした. CILDA+minILD では元データに対しては $\lambda_{1}=2 / 9$ , $\lambda_{2}=2 / 9, \lambda_{3}=2 / 9$, 拡張されたデータに対して は $\lambda_{1}=1 / 6, \lambda_{2}=1 / 12, \lambda_{3}=1 / 12$ とした. RAIL-KD, CILDA，CILDA+minILDでの中間層の線形変換後の 次元数は 128 とした。 ハイパーパラメータ探索では 5 エポックの実験を 1 回行い，5 エポックの時点で dev set でもっとも性能の高いハイパーパラメータを 選択した. MATE-KD, CILDA, CILDA+minILD では $n_{G}=10, n_{S}=100$ とした.

## 4.1.5 動作環境

本研究では PyTorch フレームワークを使用し，モ デルは Huggingface Transformers のものを使用した. また，GPUには 4 枚の Nvidia 社製 RTX A6000を使用した。節 4.1.4 に記載したバッチサイズは 4 枚の GPU の合計であり, GPU1 枚あたりのバッチサイズ は上記の 4 分の 1 である.

## 4.1.6 本実験

本実験ではハイパーパラメータ探索で選んだ学習率を用いて 20 エポックの実験を 5 回行う. 1 エポッ クごとにチェックポイントを保存し, 各実験で最も 評価指標の高いモデルを使用する. その後, 有意水準を 0.05 として並べ替え検定を片側検定で行う.

## 4.2 結果

表 1 に各実験の GLUE の dev set における平均を 示す. 括弧内の数字は標準偏差である. タスクごと に最良のものを太字としており，下線が引かれてい るものよりも統計的に有意に評価指標が高い. 教師 モデルの行は, 平均的に最も性能が高い学習率にお ける結果である.

提案手法である CILDA+minILD は使用したタス クでは MRPC 以外で MATE-KDを超える結果にな らなかった. また，先行研究では MATE-KDを超え る結果を示した CILDA も MRPC 以外で MATE-KD を超えることはなかった。

## 5 CILDA の有効性の再検証

本研究では，CILDA のスコアが MATE-KDを超え ず先行研究とは異なる結果となった. そのことが CILDA+minILD が有効でなかった理由であるという 仮説をたて，CILDA の有効性を検証した。

CILDA のスコアが低かった原因としてハイパー パラメータが適切でなかったからである可能性があ る. 先行研究では，バッチサイズは $8,16,32$ か ら, 学習率は 1e-05, 2e-05, 4e-06 から選択しており 本研究で使用したハイパーパラメータとは異なる.

## 5.1 再検証の実験設定

この仮説を検証するため，CILDAを実験を先行研究のハイパーパラメータに合わせて実験を行っ た. $n_{G}$ は CoLA， MRPC，RTE では 20，STS-B では 10 とした. $n_{S}$ はすべてのタスクにおいて 100 と した. ハイパーパラメータ探索は 5 エポックの
表 2 RoBERTa の dev set における結果

表 3 BERT の dev set における結果

実験を 5 回行い，5 エポック時点での評価指標の 平均が最も高いバッチサイズ，学習率の組を探索する. その後本実験を行う. 本実験の手順は節 4.1.6 に示したものと同じである. 実験では GPU に RTX A6000を 1 枚使用した。追加検証のため RoBERTa，BERT[1], BART[14]の 3 通りのモデルを 用いて実験を行った. BERT の実験では教師モデ ルと Generator を BERT-large (cased), 生徒モデル を DistilBERT[13], BART の実験では教師モデルと Generator を BART-large, 生徒モデルを BART-base とした. BART の実験では教師モデルのファイン チューニングを 10 エポックとし，各エポックで チェックポイントを保存し, 最良のチェックポイン トを用いた。

## 5.2 再検証実験結果

表 2，3，4 亿結果を示す. RoBERTa の実験では，節 4.2 と比較して CILDA のスコアは向上したが,傾向は変わらず MARC のみ MATE-KD を上回る結果となった. BERT の実験でも RoBERTa の実験と表 4 BART の dev set における結果

同様の傾向となった. BART の実験では RoBERTa, BERT の場合と異なり，タスク毎に最高性能の手法 が異なっているため，知識蒸留に中間層蒸留，敵対的学習，対照学習を導入することが BARTにお いては有効であるとは確認できなかった. ただし， Vanilla KD は w/o KD と比較して高いスコアが得ら れた。そのため知識蒸留自体は BARTにおいても有効である可能性がある．表 4 には記載していないが CoLA，RTE，STS-B では w/o KD と Vanilla KD の間 に有意差は確認できず，MRPC のみ Vanila KD の方 が有意に高いという結果であった.

先行研究では CILDA が MATE-KD を超える性能 を示したと報告されているが，どのモデルを使った 実験でも，そのような結果を再現できなかった. そ のため, MATE-KD と CILDA の違いである $L_{C R D}$ の 利用は有効でなかったと考えられる。提案手法の CILDA+minILD は $L_{C R D}$ を CILDA の最小化ステッ プでも用いるという手法であるため， $L_{C R D}$ が有効 な損失関数でなかったことが CILDA+minILD の性能が低かった原因と考えられる。

## 6 終わりに

本研究では, CILDA の最小化ステップに中間層出力の利用と対照学習を適用した CILDA+minILDを 提案した.しかし, 従来の手法である MATE-KD よ りも低い性能となり，その原因が CILDA の性能が 低いことにあると結論づけた．今後は，本研究で扱 わなかった GLUE の他のタスクでも同様の結果が得 られるかの検証，CILDA の性能が MATE-KD よりも 低かった原因の究明，本研究で最も高い性能を発揮 した MATE-KD を改良することによる知識蒸留の性能の向上などを行っていきたい.

## 参考文献

[1] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

[2] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. arXiv, Vol. abs/1907.11692, , 2019.

[3] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv, Vol. abs/2005.14165, , 2020.

[4] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. arXiv, Vol. abs/2204.02311, , 2022.

[5] Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. arXiv, Vol. abs/1503.02531, , 2015 .

[6] Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for BERT model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4323-4332, Hong Kong, China, November 2019. Association for Com- putational Linguistics.

[7] Peyman Passban, Yimeng Wu, Mehdi Rezagholizadeh, and Qun Liu. Alp-kd: Attention-based layer projection for knowledge distillation. Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, No. 15, pp. 13657-13665, May 2021.

[8] Md Akmal Haidar, Nithin Anchuri, Mehdi Rezagholizadeh, Abbas Ghaddar, Philippe Langlais, and Pascal Poupart. RAIL-KD: RAndom intermediate layer mapping for knowledge distillation. In Findings of the Association for Computational Linguistics: NAACL 2022, pp. 1389-1400, Seattle, United States, July 2022. Association for Computational Linguistics.

[9] Ahmad Rashid, Vasileios Lioutas, and Mehdi Rezagholizadeh. MATE-KD: Masked adversarial TExt, a companion to knowledge distillation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 1062-1071, Online, August 2021. Association for Computational Linguistics.

[10] Siqi Sun, Zhe Gan, Yuwei Fang, Yu Cheng, Shuohang Wang, and Jingjing Liu. Contrastive distillation on intermediate representations for language model compression. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 498-508, Online, November 2020. Association for Computational Linguistics.

[11] Md Akmal Haidar, Mehdi Rezagholizadeh, Abbas Ghaddar, Khalil Bibi, Phillippe Langlais, and Pascal Poupart. CILDA: Contrastive data augmentation using intermediate layer knowledge distillation. In Proceedings of the 29th International Conference on Computational Linguistics, pp. 4707-4713, Gyeongju, Republic of Korea, October 2022. International Committee on Computational Linguistics.

[12] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics.

[13] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv, Vol. abs/1910.01108, , 2019.

[14] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7871-7880, Online, July 2020. Association for Computational Linguistics.

## A ハイパーパラメータ

ハイパーパラメータ探索の結果，採用したハイ パーパラメータを表 $5 ， 6 ， 7$ に示す. 数字が 1 つ の場合は学習率，2つある場合は上が学習率，下が バッチサイズである. CILDA 検証の行は 5 章の実験 のものを表す.

表 5 RoBERTa の実験

表 6 BERT の実験

表 7 BART の実験

## B 教師モデルの性能

RoBERTa，BERT，BART の実験において，実際 に使用した教師モデルの性能を表 8 に示す.

表 8 教師モデルの性能


 鈴木 偉士 山田寛章 徳永 健伸 東京工業大学 情報理工学院 \{suzuki.t.dp@m, yamada@c, take@c\}.titech.ac.jp  