体系 › 蒸留と圧縮
·無料·論文·12分で読めます

蒸留レシピ図鑑 — logit・feature・attention・self をどう使い分けるか

蒸留のレシピは「教師のどこを学生に合わせるか」で決まります。出力(logit)・中間層(FitNet)・注意行列・自己蒸留の4種を、1枚の表と4本の式で並べ、教師がAPI越しなのか同族アーキなのかといった条件から選び方を決める地図にします。

要点

  • 蒸留レシピの違いは、教師モデルのどこを学生モデルに合わせるかの違いである。
  • 出力蒸留は最終層の確率分布だけを合わせ、教師の内部に触れず、API越しや構造が違う教師でも使える。
  • 中間層蒸留は教師と学生の隠れ層の特徴を合わせ、次元差は変換器 r で揃え、層ごとに正規化して距離を取る。
  • 注意蒸留はヘッドごとの注意行列のずれをMSEで測って平均し、ヘッド数やマスク位置の食い違いに注意が要る。
  • 自己蒸留は教師がいらないが、自分の間違いも強化されるため、効きは正則化程度に置くべきとされる。
対象textタスクtraining

Distilling the Knowledge in a Neural Network


「真似る」場所は、出口だけではない

料理の師匠から学ぶとき、真似る対象はいくつもあります。完成した皿の味だけを真似ることもできるし、下ごしらえで野菜がどう切られているかを真似ることもできる。師匠が鍋のどこを見ているか、その視線ごと真似る道もある。あるいは昨日の自分の出来映えを今日の手本にする、という独学の道もあります。

知識蒸留のレシピの違いは、突き詰めるとこれだけです。教師モデルのどこを学生モデルに合わせるか。合わせる場所が変われば、教師に要求する情報も、損失関数の形も、うまくいく条件も変わります。

蒸留の土台 — なぜ正解ラベルより教師の出力のほうが情報量が多いのか — は知識蒸留を1からで扱いました。この記事はその次の地図、「で、実際どれを使うのか」です。

4つのレシピを1枚の表で

レシピ 合わせるもの 教師に要求すること 効きやすい場面 主な落とし穴
出力蒸留(logit / response) 最終層の確率分布 出力の分布だけ(中身は不要) 教師がAPI越し、構造が全く違う 出口の情報しか流れず、深い学生に情報が足りない
中間層蒸留(feature / FitNet) 隠れ層の特徴ベクトル 中間活性の取り出し 学生が細く深い、教師と同族 次元とスケールの不一致、層の対応づけが恣意的
注意蒸留(attention) 注意行列・注意由来のマップ 各層の注意重み Transformer同士、事前学習段階での圧縮 ヘッド数・系列長・マスク位置の食い違い
自己蒸留(self) 自分(の過去/深い層)の出力 教師モデルが存在しない 教師を用意できない、正則化が欲しい 誤りを自分で強化する、伸びが小さい

この表の各行を、式・体感・コードの順に開いていきます。

レシピ1: 出力蒸留 — いちばん少ない材料で作れる

もっとも古典的で、いまも主力なのが出力蒸留です。教師と学生の最終出力の確率分布を近づけるだけで、教師の内部には一切触れません。

L=(1−α) CE(y,pS)+αT2⋅KL ⁣(pT(T) ∥ pS(T))\mathcal{L} = (1-\alpha)\,\mathrm{CE}(y, p_S) + \alpha T^2 \cdot \mathrm{KL}\!\left(p_T^{(T)} \,\|\, p_S^{(T)}\right)
(1)

言い換えると「正解ラベルに合わせる分と、温度を上げてぼかした教師の分布に合わせる分を、α\alpha の割合で混ぜる」だけです。yy は正解ラベル、CE\mathrm{CE} は交差エントロピー、pSp_S は学生の確率、pT(T)p_T^{(T)}・pS(T)p_S^{(T)} は温度 TT でならした教師・学生の分布、KL\mathrm{KL} は2つの分布のズレを測る量(KL情報量を1から)です。

温度 TT は softmax に入る値を TT で割る操作で、TT を上げるほど分布は平らになります。平らになると「1位ではないが2位の候補」の情報が浮かび上がり、そこに教師の知識が宿っている、というのが蒸留の出発点でした。式(1)の頭に付いている T2T^2 は補正項です。温度で割った分だけ勾配が 1/T21/T^2 に縮むので、掛け戻さないと温度を変えるたびに実効的な学習率が変わってしまうのです。

FIG 1温度スライダーを上げると、1位に集中していた棒が平らになり、2位・3位の候補が見えてくる。教師が渡している「ソフトな知識」はこの部分にある

このレシピの強みは要求の少なさです。教師の重みが手に入らなくても、構造がまったく違っていても、語彙さえ共有していれば実行できます。だから「まずこれを試す」が既定の順序になります。

弱みは強みと対称で、出口の1点からしか情報が流れないこと。学生が深いと下のほうの層まで信号が届きにくく、教師が持っている「途中でどんな特徴に分解しているか」は、確率分布に潰れてしまってからでは復元できません。ここから先の3つのレシピは、いずれもこの一点への対処だと思って読めます。

レシピ2: 中間層蒸留 — 途中の下ごしらえを合わせる

出口だけでは足りないなら、途中も合わせよう、というのがFitNetの発想です。教師のある層の出力(hint)と、学生のある層の出力(guided layer)を直接近づけます。

Lhint=∥r(hS(j))−hT(i)∥22\mathcal{L}_{\text{hint}} = \left\| r\big(h_S^{(j)}\big) - h_T^{(i)} \right\|_2^2
(2)

この式が言っているのは要するに、比べる場所を出口から途中へ移した、というだけのことです。教師と学生の棚を1段ずつ選び、そこに載っているものがどれだけ違うかを罰する。出口の確率分布ではなく、答えを作る途中の中間生成物どうしを突き合わせています。

「学生の第 jj 層の特徴ベクトル hS(j)h_S^{(j)} を、変換器 rr で形を合わせてから、教師の第 ii 層の特徴 hT(i)h_T^{(i)} との差の2乗を取る」という意味です。∥⋅∥22\|\cdot\|_2^2 は要素ごとの差を2乗して足した平方距離です。

rr が要るのは、細い学生と太い教師では特徴の次元がそもそも違うからです。全結合層や 1×1 畳み込みを1枚挟んで揃えます。学習が終われば捨てる、足場のようなパーツです。

勘所は2つ。1つは層の対応づけで、教師12層・学生4層なら3層ごとに1本と等間隔で結ぶのが素直な既定値です。もう1つはスケール。特徴の大きさは層ごとに桁が違うため、生のMSEを足すと値の大きい層が損失全体を支配します。層ごとに正規化してから距離を取るのが定石です(正規化層を1から)。

FIG 22本のベクトルを回すと、内積と余弦(向きの一致度)が別々に動く。「向きだけ合わせる」のか「長さまで合わせる」のか — 中間層蒸留で損失関数を選ぶとは、この2つのどちらを問題にするかを決めること

ここが中間層蒸留でいちばん判断の要るところです。MSEは向きと大きさの両方を合わせにいき、余弦距離は向きだけを問題にします。教師と学生で表現のスケールが根本的に違うなら、大きさまで一致させようとするのは無理な注文です。

レシピ3: 注意蒸留 — 「どこを見ているか」を写す

Transformer同士の蒸留では、さらに一歩進んで注意行列そのものを合わせられます。注意行列は「どのトークンがどのトークンを参照しているか」の表なので(Attention機構を1から)、これを写すのは師匠の視線を写すことに相当します。

Lattn=1h∑k=1hMSE ⁣(AS(k), AT(k))\mathcal{L}_{\text{attn}} = \frac{1}{h}\sum_{k=1}^{h} \mathrm{MSE}\!\left(A_S^{(k)},\, A_T^{(k)}\right)
(3)

つまり、答えではなく答えの出し方を採点する式です。「どのトークンがどのトークンを見ているか」の表を教師と学生で2枚重ね、マス目ごとのずれを測り、ヘッドの本数で平均する。出てくる答えが同じでも、見ている場所が違えば罰が入ります。

「hh 個のヘッドそれぞれについて、学生の注意行列 AS(k)A_S^{(k)} と教師の AT(k)A_T^{(k)} の差を測って平均する」という意味です。A(k)A^{(k)} は kk 番目のヘッドの(系列長×系列長)の表です。

画像側では、注意そのものではなく特徴マップの二乗和から作った「どこが強いか」の空間マップを合わせる手法(attention transfer)が提案されました。中身は違っても「モデルが何に反応しているかの分布を写す」という思想は共通です。

実装で気をつける点が3つ。まずヘッド数で、学生と教師で数が違うと1対1の対応づけが成立しません。MiniLMが「最終層の注意関係だけを写す」方向に寄せたのは、この問題を回避する意味もあります。次にマスク。パディング位置は両者とも重みがほぼ0なので、そこを含めてMSEを取ると「元から合っている部分」で損失が薄まります。最後にsoftmax前か後か。スコア(前)と確率(後)では勾配の性質が変わるので、論文の実装を追うときは必ず確認します。

レシピ4: 自己蒸留 — 教師がいないときの蒸留

ここまでは教師が要りました。教師がいなくても蒸留の枠組みは使えます。

生まれ変わり型(Born-Again)は、普通に学習したモデルを教師にして、同じ構造の学生を蒸留で育てます。容量が同じなのに学生のほうが良くなることがある、という直感に反する結果が報告されました。ソフトな目標が正則化として働くから、という説明がされています。

深さ方向の自己蒸留は、1つのネットワークの中で深い層が浅い層の教師になる方式です。浅い層に補助的な出口を付け、最終層の出力に合わせて学習させます。終わったら捨てれば元の構造に戻せますし、残せば「浅い出口で早く答える」早期終了の仕組みにもなります。

落とし穴は明快です。教師が自分なので、間違いも一緒に強化されること。外部の正しさが入らないため、系統的な誤りは修正されないどころかソフトラベル経由で引き継がれます。効きは「大きなモデルを小さくする」種類ではなく正則化に近い、と期待値を置くのが正しい構えです。

実装: 全部を1つのループに足す

4つのレシピは排他ではなく、重みを付けて足し合わせるのが普通です。

teacher.eval()
with torch.no_grad():
    t = teacher(x, output_hidden_states=True, output_attentions=True)
s = student(x, output_hidden_states=True, output_attentions=True)

kd = F.kl_div(F.log_softmax(s.logits / T, -1),      # 入力は必ず log 確率
              F.softmax(t.logits / T, -1),
              reduction="batchmean") * (T ** 2)      # T^2 補正を忘れない
hint = sum(mse(proj[j](s.hidden_states[j]), t.hidden_states[i])
           for j, i in layer_map)                    # 例: {1:3, 2:6, 3:9}
attn = sum(mse(a_s[mask], a_t[mask])                 # パディングは除く
           for a_s, a_t in zip(s.attentions, t.attentions))

loss = (1 - alpha) * ce + alpha * kd + beta * hint + gamma * attn

kl_div の第1引数が log 確率、第2引数が確率という非対称は、蒸留コードで最も多い取り違えです。log_softmax を忘れても例外は出ず、損失は下がるのに精度が伸びない、という気づきにくい壊れ方をします。reduction も "mean" ではなく "batchmean" にしないと、語彙サイズで割られてKD項だけが極端に小さくなります。

どれをいつ使うか — 条件から引く

現場ではこう使う

このレシピ選びが仕事になるのは、推論コストを下げる担当(MLエンジニア、推論基盤の担当者)が「精度をなるべく落とさずにモデルを小さくしろ」と言われたときです。最初に手が伸びるのはたいてい量子化で(量子化を1から)、それでも足りないときに蒸留が候補に上がります。

実際に触るのは、温度 T、正解ラベルとKD項の混合比 alpha、中間層・注意の損失重み beta / gamma、そして層対応表です。HuggingFace transformers なら output_hidden_states=True / output_attentions=True で中間層と注意が取り出せ、自前実装なら PyTorch の register_forward_hook を使います。

事故になる落とし穴を、起きやすい順に並べます。

  1. 教師を eval() にし忘れる。Dropout や BatchNorm が動いたままだと教師の出力が毎回揺れ、学生は動く的を撃つことになります。torch.no_grad() を忘れると教師側の勾配を保持してメモリも大きく増えます
  2. T² 補正の欠落。温度を変えた実験の比較が、実質的には学習率を変えた比較になり、結論がひっくり返ります
  3. 損失のスケールを見ていない。beta・gamma を初期値のまま走らせると、中間層のMSEがCEを何桁も上回り、学生がタスクを学ばないことがあります。開始時に各項の値をログに出すだけで防げます
  4. 教師の入力前処理が学生と違う。画像なら正規化の平均・分散、言語ならトークナイザ。ずれると教師の出力は「別の入力に対する答え」になり、いくら合わせても精度は出ません

設計レビューでよく問われるのは「なぜ温度を上げるのか、T2T^2 は何のためか」です。上げる理由は2位以下の情報を取り出すため、掛ける理由は温度で縮んだ勾配を戻して温度と学習率を独立に扱えるようにするため。2つを分けて説明できると、式の意味まで理解していることが伝わります。

まとめ

参考文献

  1. Distilling the Knowledge in a Neural Network. arXiv:1503.02531論文ページ·PDF
  2. FitNets: Hints for Thin Deep Nets. arXiv:1412.6550論文ページ·PDF
  3. Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer. arXiv:1612.03928論文ページ·PDF
  4. TinyBERT: Distilling BERT for Natural Language Understanding. arXiv:1909.10351論文ページ·PDF
  5. MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers. arXiv:2002.10957論文ページ·PDF
  6. Born-Again Neural Networks. arXiv:1805.04770論文ページ·PDF
  7. Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation. arXiv:1905.08094論文ページ·PDF

本記事は上記論文の本文にもとづいて執筆しています。数値・主張は原典を優先してください。

運営: TechAthletes · 編集方針 · 出典の付け方 · 訂正の記録

コメント

コメントにはログインが必要です