知識蒸留とは？大規模AIモデルを軽量化するKnowledge Distillation

Kerasで知識蒸留を実装する：Teacher-StudentモデルのDistillation Loss

知識蒸留（Knowledge Distillation）は、大きなTeacherモデルの知識を小さなStudentモデルに転移する技術です。以下ではKerasのカスタム訓練ステップを使い、KLダイバージェンスとクロスエントロピーを組み合わせたDistillation Lossを実装します。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

# --- Teacherモデル（大きめのモデル） ---
def build_teacher(input_shape, num_classes):
    model = keras.Sequential([
        layers.Input(shape=input_shape),
        layers.Conv2D(64, 3, padding="same", activation="relu"),
        layers.Conv2D(64, 3, padding="same", activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(128, 3, padding="same", activation="relu"),
        layers.Conv2D(128, 3, padding="same", activation="relu"),
        layers.MaxPooling2D(2),
        layers.GlobalAveragePooling2D(),
        layers.Dense(256, activation="relu"),
        layers.Dropout(0.3),
        layers.Dense(num_classes)  # logitsを出力（softmaxなし）
    ], name="teacher")
    return model

# --- Studentモデル（小さいモデル） ---
def build_student(input_shape, num_classes):
    model = keras.Sequential([
        layers.Input(shape=input_shape),
        layers.Conv2D(16, 3, padding="same", activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, padding="same", activation="relu"),
        layers.MaxPooling2D(2),
        layers.GlobalAveragePooling2D(),
        layers.Dense(64, activation="relu"),
        layers.Dense(num_classes)  # logitsを出力
    ], name="student")
    return model

# --- 蒸留トレーナー ---
class Distiller(keras.Model):
    def __init__(self, student, teacher):
        super().__init__()
        self.student = student
        self.teacher = teacher
        self.teacher.trainable = False  # Teacherは固定

    def compile(self, optimizer, temperature=3.0, alpha=0.7):
        super().compile(optimizer=optimizer)
        self.temperature = temperature
        self.alpha = alpha
        self.student_loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        self.distillation_loss_fn = keras.losses.KLDivergence()

    def train_step(self, data):
        x, y = data

        # Teacherの予測（勾配計算不要）
        teacher_logits = self.teacher(x, training=False)

        with tf.GradientTape() as tape:
            # Studentの予測
            student_logits = self.student(x, training=True)

            # 通常のクロスエントロピー損失（hard label）
            student_loss = self.student_loss_fn(y, student_logits)

            # 蒸留損失（soft label）
            # 温度付きsoftmaxで確率分布を軟化
            teacher_soft = tf.nn.softmax(teacher_logits / self.temperature, axis=-1)
            student_soft = tf.nn.softmax(student_logits / self.temperature, axis=-1)
            distillation_loss = self.distillation_loss_fn(teacher_soft, student_soft)
            distillation_loss *= self.temperature ** 2  # 温度のスケーリング補正

            # 合計損失: alpha * 蒸留損失 + (1-alpha) * 通常損失
            total_loss = (self.alpha * distillation_loss
                          + (1 - self.alpha) * student_loss)

        # Studentのみ更新
        grads = tape.gradient(total_loss, self.student.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.student.trainable_variables))

        return {
            "total_loss": total_loss,
            "student_loss": student_loss,
            "distillation_loss": distillation_loss
        }

# --- 実行例（MNISTで蒸留） ---
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1).astype("float32") / 255.0
x_test = x_test.reshape(-1, 28, 28, 1).astype("float32") / 255.0

teacher = build_teacher((28, 28, 1), 10)
student = build_student((28, 28, 1), 10)

# Teacherを先に学習
teacher.compile(optimizer="adam",
                loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=["accuracy"])
teacher.fit(x_train, y_train, epochs=5, batch_size=128, verbose=1)
teacher_acc = teacher.evaluate(x_test, y_test, verbose=0)[1]
print(f"Teacher accuracy: {teacher_acc:.4f}")
print(f"Teacher params: {teacher.count_params():,}")

# 蒸留でStudentを学習
distiller = Distiller(student=student, teacher=teacher)
distiller.compile(optimizer="adam", temperature=3.0, alpha=0.7)
distiller.fit(x_train, y_train, epochs=5, batch_size=128, verbose=1)

# Student単体の精度を確認
student.compile(optimizer="adam",
                loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=["accuracy"])
student_acc = student.evaluate(x_test, y_test, verbose=0)[1]
print(f"Student accuracy (distilled): {student_acc:.4f}")
print(f"Student params: {student.count_params():,}")
# Teacher: ~300K params → Student: ~30K params で精度低下を最小限に

蒸留の核心はtemperatureとalphaの2つのハイパーパラメータにあります。温度を高くするとTeacherのsoftmax出力が平坦になり、不正解クラス間の相対的な関係性（dark knowledge）がStudentに伝わりやすくなります。alphaは蒸留損失と通常損失のバランスを制御し、0.7前後で蒸留損失を重視するのが一般的です。この手法によりパラメータ数を1/10に圧縮しても精度低下を1〜2%に抑えられるケースが多く、エッジデバイスへのデプロイで威力を発揮します。