Posted inAI

LLMの知識蒸留（Knowledge Distillation）：モデルを軽量化して高速化と低コスト化を実現する方法

AI tutorial - IT technology blog
AI tutorial - IT technology blog

背景と目的：大規模LLM運用の現実

6ヶ月前、私のチームは一般的ですが非常に厄介なボトルネックに直面しました。GPT-4を使用した複雑なRAGパイプラインを運用していましたが、品質は素晴らしいものの、APIの利用料金が跳ね上がっていたのです。さらに重要なのは、単純な推論タスクに5〜10秒のレイテンシが発生し、ユーザーエクスペリエンスを損なっていたことでした。Llama-3-8Bのような小型モデルに対してプロンプトエンジニアリングや基本的なファインチューニングを試みましたが、推論能力の差が大きすぎました。小型モデルでは、大型モデルが持つニュアンスを再現できなかったのです。

そこで私は、知識蒸留（Knowledge Distillation: KD）に注目しました。単に生のデータで小型モデルを学習させるのではなく、KDでは「教師（Teacher）」モデル（GPT-4やLlama-3-70Bなどの巨大モデル）を使用して、「生徒（Student）」モデル（Phi-3やLlama-3-8Bなど）を導きます。目標は、生徒に教師の最終的な出力だけでなく、内部の論理構成（ロジック）を模倣させることです。

このアプローチを本番環境に適用したところ、結果は一貫して安定していました。特定の推論能力を8Bモデルに蒸留することで、推論コストを10分の1に削減し、トークン生成速度（tokens-per-second）を大幅に向上させつつ、特定のドメインタスクにおいて教師モデルの約94%の精度を維持できました。モデルの知能と運用予算のトレードオフに悩んでいるなら、蒸留は論理的な次のステップと言えるでしょう。

インストール：蒸留環境のセットアップ

蒸留を行うには、2つのモデルを同時に処理できる（あるいは少なくとも教師モデルの出力をデータセットとして処理できる）堅牢な環境が必要です。私は通常、このワークフローに最も成熟したツールを提供しているHugging Faceのエコシステムを使用します。

まず、十分なVRAMを搭載したマシンを用意してください。ローカルで70Bモデルから8Bモデルへ蒸留する場合、A100が2枚必要になるか、少なくとも教師モデルの量子化版を使用することになるでしょう。以下は、私が使用しているベースのセットアップです：

pip install -U torch transformers datasets accelerate bitsandbytes peft trl

フルパラメータの蒸留はコストがかかりすぎることが多いため、trl（Transformer Reinforcement Learning）とpeftを使用します。多くの場合、メモリと時間を節約するために、蒸留プロセス中に生徒モデルに対してLoRA（Low-Rank Adaptation）を適用します。

設定：教師・生徒ワークフローの実装

知識を蒸留する方法はいくつかありますが、LLMにおいては一般的に「レスポンスベースの蒸留（Response-based Distillation）」または「ロジットベースの蒸留（Logit-based Distillation）」に焦点を当てます。私の経験では、教師が生成したデータに対してSFTを行うレスポンスベースの蒸留が、プロダクションチームにとって最も実用的です。

1. 蒸留用データセットの生成

最初のステップは、学習セットに対して教師モデルに詳細な解説を生成させることです。生徒に論理的思考を身につけさせたいのであれば、教師はその思考プロセスを示す必要があります。これは「Chain-of-Thought（思考の連鎖）蒸留」と呼ばれることがよくあります。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 教師モデルのロード (例: Llama-3-70B-Instruct)
teacher_id = "meta-llama/Meta-Llama-3-70B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(teacher_id)
model = AutoModelForCausalLM.from_pretrained(
    teacher_id, 
    device_map="auto", 
    load_in_4bit=True # コンシューマー向けや中位のエンタープライズハードウェアに収めるために4-bitを使用
)

def generate_teacher_rationale(prompt):
    messages = [
        {"role": "system", "content": "あなたは熟練の教師です。次の質問に対して、ステップバイステップの詳細な推論を提供してください。"},
        {"role": "user", "content": prompt}
    ]
    # 標準的な生成ロジックをここに記述...
    return teacher_output

2. 生徒モデルの学習

{ "prompt": "...", "teacher_explanation": "...", "final_answer": "..." } という形式のデータセットが準備できたら、生徒モデルを学習させます。以下の設定は、このプロセスで最も信頼できる方法であるSFTTrainerを使用した学習ループの構築方法を示しています。

from trl import SFTTrainer
from transformers import TrainingArguments

# 生徒モデル: Llama-3-8B
student_id = "meta-llama/Meta-Llama-3-8B"

training_args = TrainingArguments(
    output_dir="./distilled-model",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=10,
    optim="paged_adamw_32bit",
    fp16=True,
)

trainer = SFTTrainer(
    model=student_id,
    train_dataset=distilled_dataset,
    dataset_text_field="teacher_explanation", # 生徒は教師のロジックを学習する
    max_seq_length=2048,
    args=training_args,
)

trainer.train()

単なる「最終的な回答（final_answer）」ではなく、「教師の解説（teacher_explanation）」に基づいて学習することで、生徒モデルは潜在的な思考パターンを学びます。これが、蒸留された8Bモデルが、生のデータで学習された標準的な8Bモデルを凌駕することが多い理由です。

検証とモニタリング：成功の測定

蒸留プロセスが完了した後、標準的な損失（loss）メトリクスだけに頼ることはできません。学習時の損失が低いからといって、本番環境で「賢い」モデルになるとは限らないからです。私は、生徒モデルがデプロイ可能な状態にあるかを確認するために、3段階の検証システムを使用しています。

1. LLM-as-a-Judge（評価者としてのLLM）

通常、500個のテストサンプルを用意し、GPT-4oに教師と生徒の出力を比較させます。ここで注目するのは「勝率」です。特定のドメインタスクにおいて、生徒が教師に勝つか、あるいは同等である割合が80%を超えれば、成功とみなします。

2. レイテンシとスループットのベンチマーク

この取り組みの最大の目的はパフォーマンスです。私はvLLMを使用して、蒸留済みモデルのスループットをテストします。最近の本番環境への移行では、以下のような変化が見られました：

  • 教師 (70B): 15 tokens/sec, 100万トークンあたり0.80ドル
  • 生徒 (蒸留済み8B): 95 tokens/sec, 100万トークンあたり0.05ドル

3. ドリフト・モニタリング

デプロイ後は、生徒モデルの「確信度スコア（Confidence Score）」を監視します。蒸留時に見られた値と比較して、生徒が低い対数確率（log-probability）で回答を生成した場合、教師モデルへのフォールバックをトリガーします。このハイブリッドなアプローチにより、蒸留モデルが未知の事態に遭遇しても、「大きな脳（教師）」がカバーできる体制を整えています。

蒸留は一度きりの作業ではありません。データが進化するにつれ、定期的に蒸留パイプラインを再実行し、教師から新たなニュアンスを取り込む必要があります。これは、ユーザーが期待する品質を犠牲にすることなく、本番環境のコストを低く抑え続けるための継続的な改善サイクルなのです。

Share:
Tags: