Knowledge Distillation: LLMの圧縮

このスキルを使用する時機

以下の場合に知識蒸留を使用してください：

• モデルを圧縮 70B → 7Bで90%以上のパフォーマンスを保持
• 機能を転送 専有モデル（GPT-4）からオープンソース（LLaMA、Mistral）へ
• 推論コストを削減 小さな学生モデルを展開することで
• 専門的なモデルを作成 ドメイン固有の知識を蒸留することで
• 小さなモデルを改善 大規模な教師から合成データを使用

主要な技術：温度スケーリング、ソフトターゲット、逆KLD（MiniLLM）、ロジット蒸留、応答蒸留

論文：Hinton et al. 2015（arXiv 1503.02531）、MiniLLM（arXiv 2306.08543）、KD調査（arXiv 2402.13116）

インストール

# 標準 transformers
pip install transformers datasets accelerate

# トレーニング用
pip install torch deepspeed wandb

# オプション：MiniLLM実装
git clone https://github.com/microsoft/LMOps
cd LMOps/minillm
pip install -e .

クイックスタート

基本的な知識蒸留

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments

# 1. 教師（大規模）と学生（小規模）モデルを読み込む
teacher = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",  # 大規模な教師
    torch_dtype=torch.float16,
    device_map="auto"
)

student = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",  # 小規模な学生
    torch_dtype=torch.float16,
    device_map="cuda:0"
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf")

# 2. 蒸留損失を定義
def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    """
    ハード損失（交差エントロピー）とソフト損失（KLダイバージェンス）を組み合わせます。

    Args:
        temperature: 確率分布を軟化させます（高いほどより軟らか）
        alpha: 蒸留損失の重み（1-alphaはハード損失）
    """
    # ハード損失：正解ラベルを使用した標準的な交差エントロピー
    hard_loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), labels.view(-1))

    # ソフト損失：学生と教師間のKLダイバージェンス
    soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean') * (temperature ** 2)

    # 組み合わせられた損失
    return alpha * soft_loss + (1 - alpha) * hard_loss

# 3. トレーニングループ
for batch in dataloader:
    # 教師の順伝播（勾配なし）
    with torch.no_grad():
        teacher_outputs = teacher(**batch)
        teacher_logits = teacher_outputs.logits

    # 学生の順伝播
    student_outputs = student(**batch)
    student_logits = student_outputs.logits

    # 蒸留損失を計算
    loss = distillation_loss(
        student_logits,
        teacher_logits,
        batch['labels'],
        temperature=2.0,
        alpha=0.7  # 70%ソフト、30%ハード
    )

    # 逆伝播と最適化
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

MiniLLM（逆KLD）

出典：arXiv 2306.08543（2024）

革新：生成モデルの蒸留を改善するため、順方向KLDではなく逆KLDを使用します。

def reverse_kl_loss(student_logits, teacher_logits, temperature=1.0):
    """
    逆KLダイバージェンス：KL(Teacher || Student)
    順方向KLよりも生成モデルに適しています。
    """
    # 教師分布（目標）
    p_teacher = F.softmax(teacher_logits / temperature, dim=-1)

    # 学生分布（モデル）
    log_p_student = F.log_softmax(student_logits / temperature, dim=-1)

    # 逆KL：教師で合計し、学生は教師のすべてのモードをカバーすることを学ぶ
    reverse_kl = -(p_teacher * log_p_student).sum(dim=-1).mean()

    return reverse_kl * (temperature ** 2)

# MiniLLMでのトレーニング
for batch in dataloader:
    with torch.no_grad():
        teacher_logits = teacher(**batch).logits

    student_logits = student(**batch).logits

    # 逆KLD（生成に適している）
    loss = reverse_kl_loss(student_logits, teacher_logits, temperature=1.0)

    loss.backward()
    optimizer.step()

なぜ逆KL？

• 順方向KL（標準）：学生は教師の平均に一致することを学ぶ
• 逆KL（MiniLLM）：学生は教師のすべてのモードをカバーすることを学ぶ
• テキスト生成の多様性により適している

応答蒸留

# 教師から合成データを生成し、学生に模倣をさせてトレーニング

# 1. 教師から合成応答を生成
prompts = ["AIを説明してください:", "MLとは何ですか?", "NLPを定義してください:"]

teacher_responses = []
for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors='pt').to(teacher.device)
    outputs = teacher.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    teacher_responses.append(response)

# 2. 教師の応答で学生をトレーニング（標準的なファインチューニング）
train_dataset = [
    {"text": f"{prompt}\n{response}"}
    for prompt, response in zip(prompts, teacher_responses)
]

# 3. 学生をファインチューニング
trainer = Trainer(
    model=student,
    args=TrainingArguments(output_dir="./student", num_train_epochs=3, learning_rate=2e-5),
    train_dataset=train_dataset,
)
trainer.train()

コアコンセプト

1. 温度スケーリング

目的：確率分布を軟化させて、教師の不確実性を露出させます。

# 低温度（T=1）：シャープな分布
logits = [3.0, 2.0, 1.0]
probs_T1 = softmax(logits / 1.0)  # [0.67, 0.24, 0.09]

# 高温度（T=4）：ソフトな分布
probs_T4 = softmax(logits / 4.0)  # [0.42, 0.34, 0.24]

# より高いTは相対的ランキングについてより多くの情報を明かします

ルール：蒸留にはT=2-5を使用してください（2が一般的なデフォルト）。

2. 損失関数コンポーネント

# 総損失 = alpha * soft_loss + (1 - alpha) * hard_loss

# ソフト損失：教師の知識から学ぶ
soft_loss = KL(student || teacher)

# ハード損失：正解ラベルから学ぶ
hard_loss = CrossEntropy(student_output, true_labels)

# 一般的な値：
alpha = 0.5  # バランス
alpha = 0.7  # 教師の知識を強調
alpha = 0.3  # ラベルを強調

3. 順方向KLD対逆KLD

# 順方向KL：KL(Student || Teacher)
# - 学生は教師の平均的な動作に一致する
# - モードシーキング：学生は教師の最高確率モードに焦点を当てる
# - 分類に適している

# 逆KL：KL(Teacher || Student)
# - 学生は教師のすべての動作をカバーする
# - モードカバリング：学生は多様な動作を学ぶ
# - 生成に適している（MiniLLM）

トレーニング戦略

戦略1：ロジット蒸留

# 学生を教師のロジットに直接一致するようトレーニング

def logit_distillation_trainer(student, teacher, dataloader, temperature=2.0):
    optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5)

    for epoch in range(3):
        for batch in dataloader:
            # ロジットを取得
            with torch.no_grad():
                teacher_logits = teacher(**batch).logits

            student_logits = student(**batch).logits

            # ロジット上のMSE（KLDの代替）
            loss = F.mse_loss(student_logits, teacher_logits)

            # またはKLDを使用
            # loss = F.kl_div(
            #     F.log_softmax(student_logits/temperature, dim=-1),
            #     F.softmax(teacher_logits/temperature, dim=-1),
            #     reduction='batchmean'
            # ) * (temperature ** 2)

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

    return student

戦略2：2段階蒸留

# ステージ1：教師から蒸留
student = distill(teacher, student, epochs=5)

# ステージ2：タスク固有のデータでファインチューニング
student = fine_tune(student, task_data, epochs=3)

# 単一段階よりも優れたタスクパフォーマンスになります

戦略3：マルチティーチャー蒸留

# 複数の専門家教師から学ぶ

def multi_teacher_distillation(student, teachers, batch):
    """教師のアンサンブルから蒸留します。"""
    teacher_logits_list = []

    # すべての教師からロジットを取得
    with torch.no_grad():
        for teacher in teachers:
            logits = teacher(**batch).logits
            teacher_logits_list.append(logits)

    # 教師予測を平均化
    avg_teacher_logits = torch.stack(teacher_logits_list).mean(dim=0)

    # 学生はアンサンブルから学ぶ
    student_logits = student(**batch).logits
    loss = F.kl_div(
        F.log_softmax(student_logits, dim=-1),
        F.softmax(avg_teacher_logits, dim=-1),
        reduction='batchmean'
    )

    return loss

本番環境へのデプロイ

完全なトレーニングスクリプト

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

def train_distilled_model(
    teacher_name="meta-llama/Llama-2-70b-hf",
    student_name="meta-llama/Llama-2-7b-hf",
    output_dir="./distilled-llama-7b",
    temperature=2.0,
    alpha=0.7,
):
    # モデルを読み込む
    teacher = AutoModelForCausalLM.from_pretrained(teacher_name, torch_dtype=torch.float16, device_map="auto")
    student = AutoModelForCausalLM.from_pretrained(student_name, torch_dtype=torch.float16)
    tokenizer = AutoTokenizer.from_pretrained(teacher_name)

    # 蒸留を使用したカスタムトレーナー
    class DistillationTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False):
            # 学生の順伝播
            outputs_student = model(**inputs)
            student_logits = outputs_student.logits

            # 教師の順伝播（勾配なし）
            with torch.no_grad():
                outputs_teacher = teacher(**inputs)
                teacher_logits = outputs_teacher.logits

            # 蒸留損失
            soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
            soft_student = F.log_softmax(student_logits / temperature, dim=-1)
            soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean') * (temperature ** 2)

            # ハード損失
            hard_loss = outputs_student.loss

            # 組み合わせ
            loss = alpha * soft_loss + (1 - alpha) * hard_loss

            return (loss, outputs_student) if return_outputs else loss

    # トレーニング引数
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        learning_rate=2e-5,
        warmup_steps=500,
        logging_steps=100,
        save_steps=1000,
        bf16=True,
        gradient_checkpointing=True,
    )

    # トレーニング
    trainer = DistillationTrainer(
        model=student,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )

    trainer.train()
    student.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)

# 使用方法
train_distilled_model(
    teacher_name="meta-llama/Llama-2-70b-hf",
    student_name="meta-llama/Llama-2-7b-hf",
    temperature=2.0,
    alpha=0.7
)

ベストプラクティス

1. ハイパーパラメータの選択

# 温度
T = 1.0  # シャープ（知識転移が少ない）
T = 2.0  # 標準（バランスが良い）
T = 5.0  # ソフト（知識転移が多い）

# アルファ（重み）
alpha = 0.5  # バランス
alpha = 0.7  # 教師の知識を強調
alpha = 0.9  # 強い蒸留

# ルール：より高いT+より高いアルファ=より強い蒸留

2. モデルサイズ比

# 良い比率（教師/学生）
70B / 7B = 10×    # 優秀
13B / 1B = 13×    # 良好
7B / 1B = 7×      # 許容可能

# 大きすぎるギャップは避ける
70B / 1B = 70×    # 大きすぎる、効果がない

3. データの品質

# 最良：教師生成データ+実データ
train_data = {
    "teacher_generated": 70%,  # 多様で高品質
    "real_data": 30%            # 正解
}

# 避ける：実データのみ（教師を十分に活用しない）

評価

from transformers import pipeline

# 学生と教師を比較
teacher_pipe = pipeline("text-generation", model=teacher)
student_pipe = pipeline("text-generation", model=student)

prompts = ["量子コンピューティングを説明してください:", "AIとは何ですか?"]

for prompt in prompts:
    teacher_out = teacher_pipe(prompt, max_new_tokens=100)
    student_out = student_pipe(prompt, max_new_tokens=100)

    print(f"プロンプト: {prompt}")
    print(f"教師: {teacher_out[0]['generated_text']}")
    print(f"学生: {student_out[0]['generated_text']}")
    print(f"一致品質: {calculate_similarity(teacher_out, student_out):.2f}")

リソース

• Hinton et al. 2015（基礎）：https://arxiv.org/abs/1503.02531
• MiniLLM（逆KLD）：https://arxiv.org/abs/2306.08543
• LLMのKD調査（2024）：https://arxiv.org/abs/2402.13116
• MiniLLM GitHub：https://github.com/microsoft/LMOps/tree/main/minillm