技術

知識蒸留とは

Q: 知識蒸留とファインチューニングの違いは？

ファインチューニングは1つのモデル自身の重みを追加学習で調整する手法。 知識蒸留は、教師モデルの出力を使って別の(通常はより小規模な)生徒モデルを新たに学習させる手法。

Knowledge Distillation / 蒸留

大規模な教師モデルの出力を模倣させることで、小規模な生徒モデルへ知識を移す学習手法

モデル圧縮軽量化

ひとことで言うと

大きくて賢いAIの知識を、小さいAIに教え込んで引き継がせる方法。

概要

知識蒸留(Knowledge Distillation)とは、学習済みの大規模な教師モデル(teacher)の出力や内部表現を模倣するように、より小規模な生徒モデル(student)を学習させる手法。生徒モデルは正解ラベルだけでなく、教師モデルが出力する確率分布(各クラスやトークンに対するソフトな予測)も学習信号として利用することで、正解ラベルだけでは伝わらない教師モデルの「自信の度合い」や誤答同士の類似性といった情報も引き継げる。 SLM(小規模言語モデル)の学習や、モバイル・エッジ向けのモデル軽量化などで広く使われる。蒸留によって生徒モデルが教師モデルと全く同じ性能に達することは少ないが、パラメータ数の割に高い性能を実現できるため、推論コストの削減が重視される実運用環境で重宝される。

背景

モデルの性能を高めるにはパラメータ数を増やすアプローチが有効だが、大規模モデルは推論コスト・メモリ使用量が大きく実運用に不向きな場合がある。知識蒸留は、大規模モデルが学習した知識をできるだけ保ちながら、より軽量なモデルへ移し替えることで性能と実用性を両立させるために考案された。

歴史

2015年: HintonらがKnowledge Distillationという名称と手法を論文で提案し、教師モデルの出力するソフトターゲットを用いた蒸留の枠組みを確立。 2019年: DistilBERT等、BERTを蒸留した軽量モデルが登場。 2023年以降: Microsoft PhiシリーズなどSLMの学習で、大規模モデルの出力を用いた蒸留的アプローチが活用される。

アーキテクチャ

教師モデルの出力層のソフトマックス出力(温度パラメータで平滑化したソフトターゲット)と、生徒モデルの出力との差を損失関数(蒸留損失)として計算し、正解ラベルに対する通常の損失と組み合わせて生徒モデルを学習する。中間層の特徴量やAttentionの分布を模倣させる、より発展的な蒸留手法も存在する。

ワークフロー

大規模な教師モデルを用意(または既存の学習済みモデルを利用) → 同じ入力に対する教師モデルの出力(ソフトターゲット)を取得 → 正解ラベルと教師の出力の両方に近づくよう生徒モデルを学習 → 評価し、蒸留の設定(温度・損失の重み等)を調整。

コード例

PyTorchによる蒸留損失の計算例

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # 教師の出力(ソフトターゲット)に近づける蒸留損失
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits / temperature, dim=-1),
        reduction="batchmean",
    ) * (temperature ** 2)
    # 正解ラベルに対する通常の損失
    hard_loss = F.cross_entropy(student_logits, labels)
    return alpha * soft_loss + (1 - alpha) * hard_loss

利点

教師モデルより少ないパラメータで、近い性能を達成できる場合がある
推論コスト・メモリ使用量を抑えられ、エッジデバイスでの実行に向く
ソフトターゲットを通じて、正解ラベルだけでは伝わらない教師の判断の傾向も引き継げる

欠点

生徒モデルの性能は教師モデルが上限となり、それを超えることは基本的にない
教師モデルの選定や損失の調整など、蒸留の設定に試行錯誤を要する
教師モデル自体の推論コストがかかるため、学習データ生成の準備に時間を要する

比較

SLM — SLMの学習では、大規模モデルからの知識蒸留がよく用いられる
ファインチューニング — ファインチューニングは既存モデル自身を追加学習するのに対し、知識蒸留は別の(通常より小さい)モデルへ知識を移す
量子化 — 知識蒸留と量子化は、どちらもモデルを軽量化する代表的な手法であり、組み合わせて使われることもある
Phi — Phiシリーズの学習データには、大規模モデルによって生成された合成データが活用されている
エッジAI — 知識蒸留は、エッジAI向けにモデルを軽量化する代表的な手法の1つ

よくある質問

知識蒸留とファインチューニングの違いは？

ファインチューニングは1つのモデル自身の重みを追加学習で調整する手法。知識蒸留は、教師モデルの出力を使って別の(通常はより小規模な)生徒モデルを新たに学習させる手法。

生徒モデルは教師モデルより高性能になれる？

基本的には教師モデルの性能が上限となるが、蒸留による正則化効果で、特定の指標では教師モデルに匹敵する例も報告されている。

参考文献

Research PaperDistilling the Knowledge in a Neural Network
Research PaperDistilBERT, a distilled version of BERT