概要
- 知識蒸留は大規模モデルの実用化における重要なブリッジ:70Bパラメータの教師モデルを7Bの学生モデルに圧縮し、推論コストを90%以上削減、性能を85%以上保持
- 3つの蒸留戦略:ホワイトボックス蒸留(Logitマッチング)、グレーボックス蒸留(特徴量アライメント)、ブラックボックス蒸留(出力模倣)、それぞれ適用シーンが異なる
- 段階的蒸留3フェーズ:70B→30B→7B、各フェーズで90%以上の知識を保持、最終モデルの性能は直接蒸留に近い
- マルチタスク蒸留の実践:汎用能力+専門能力を同時に蒸留、破滅的忘却を回避
- 本記事ではLLaMA蒸留の全プロセスコードとDeepSeek-R1蒸留スキームの解析を提供
目次
知識蒸留:大規模モデル実用化の必須プロセス
なぜ知識蒸留が必要なのか?
| 次元 |
教師モデル(70B) |
学生モデル(7B) |
蒸留による効果 |
| 推論レイテンシ |
800ms/token |
80ms/token |
10× |
| GPUメモリ |
4×A100 |
1×RTX4090 |
4× |
| デプロイコスト |
$5/千リクエスト |
$0.5/千リクエスト |
10× |
| 汎用能力 |
92点 |
78点 |
-15% |
| 専門能力 |
88点 |
82点 |
-7% |
知識蒸留の進化ルート
| 段階 |
時期 |
手法 |
代表的な研究 |
| 古典的蒸留 |
2015 |
Soft Label |
HintonのKD |
| 特徴量蒸留 |
2018-2020 |
中間層アライメント |
FitNets, PKT |
| LLM蒸留 |
2023 |
Logit+特徴量 |
Alpaca, Vicuna |
| 体系的蒸留 |
2024-2026 |
段階的+マルチタスク |
DeepSeek-R1, Qwen2.5 |
2026年の主流蒸留スキーム
| スキーム |
教師モデル |
学生モデル |
能力保持 |
オープンソース |
| DeepSeek-R1 |
R1-671B |
R1-Distill-8B |
82% |
はい |
| Qwen2.5蒸留 |
Qwen2.5-72B |
Qwen2.5-7B |
85% |
はい |
| LLaMA蒸留 |
LLaMA-3-70B |
LLaMA-3-8B |
80% |
一部 |
| GPT-4蒸留 |
GPT-4 |
GPT-4o-mini |
N/A |
いいえ |
| Claude蒸留 |
Claude-3.5 |
Claude-3-Haiku |
N/A |
いいえ |
3つの蒸留戦略の詳細解説
戦略概要
┌─────────────────────────────────────────────────────────────┐
│ 3つの蒸留戦略 │
│ │
│ 1. ホワイトボックス蒸留 (White-Box) │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ 教師モデル内部にアクセス:Logit、隠れ層、アテンション重み │ │
│ │ 損失 = α·L_logit + β·L_feature + γ·L_attn │ │
│ │ 利点:蒸留効果が最も高い │ │
│ │ 欠点:教師モデルの重みが必要、計算コストが高い │ │
│ └─────────────────────────────────────────────────────┘ │
│ │
│ 2. グレーボックス蒸留 (Gray-Box) │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ 教師モデルに部分的アクセス:Logitのみまたは特徴量のみ │ │
│ │ 損失 = α·L_logit + β·L_CE │ │
│ │ 利点:効果とコストのバランス │ │
│ │ 欠点:ホワイトボックスより効果が劣る │ │
│ └─────────────────────────────────────────────────────┘ │
│ │
│ 3. ブラックボックス蒸留 (Black-Box) │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ 教師モデルの出力のみアクセス:生成テキスト+スコア │ │
│ │ 損失 = L_CE(学生出力, 教師生成テキスト) │ │
│ │ 利点:教師の重みが不要、APIのみで可能 │ │
│ │ 欠点:効果が最も低い、データ依存性が強い │ │
│ └─────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
戦略選択のデシジョンツリー
| 条件 |
推奨戦略 |
| 教師モデルの重みあり + 十分なGPU |
ホワイトボックス蒸留 |
| 教師モデルの重みあり + GPUが限定的 |
グレーボックス蒸留(Logit) |
| APIアクセスのみ |
ブラックボックス蒸留 |
| 教師モデルが非常に大きい(>100B) |
ブラックボックス蒸留+データ拡張 |
ホワイトボックス蒸留:Logitマッチング実践
コア原理
ホワイトボックス蒸留のコアは、学生モデルに教師モデルの出力分布(Soft Label)を学習させることであり、ハードラベル(Hard Label)ではない。
import torch
import torch.nn as nn
import torch.nn.functional as F
class LogitDistillationLoss(nn.Module):
def __init__(self, temperature=2.0, alpha=0.7):
super().__init__()
self.temperature = temperature
self.alpha = alpha
def forward(self, student_logits, teacher_logits, labels):
soft_loss = F.kl_div(
F.log_softmax(student_logits / self.temperature, dim=-1),
F.softmax(teacher_logits / self.temperature, dim=-1),
reduction="batchmean",
) * (self.temperature ** 2)
hard_loss = F.cross_entropy(student_logits, labels)
return self.alpha * soft_loss + (1 - self.alpha) * hard_loss
class LLMDistillationTrainer:
def __init__(
self,
teacher_model,
student_model,
tokenizer,
temperature=2.0,
alpha=0.7,
):
self.teacher = teacher_model
self.student = student_model
self.tokenizer = tokenizer
self.loss_fn = LogitDistillationLoss(temperature, alpha)
self.teacher.eval()
def distill_step(self, input_ids, attention_mask, labels):
with torch.no_grad():
teacher_outputs = self.teacher(
input_ids=input_ids,
attention_mask=attention_mask,
)
teacher_logits = teacher_outputs.logits
student_outputs = self.student(
input_ids=input_ids,
attention_mask=attention_mask,
)
student_logits = student_outputs.logits
shift_teacher = teacher_logits[..., :-1, :].contiguous()
shift_student = student_logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = self.loss_fn(shift_student, shift_teacher, shift_labels)
return loss
温度パラメータのチューニング
| 温度T |
ソフトラベル分布 |
蒸留効果 |
適用シーン |
| 1.0 |
尖鋭(ハードラベルに近い) |
低い |
非推奨 |
| 2.0 |
適度 |
良い |
汎用推奨 |
| 4.0 |
滑らか |
中程度 |
知識が豊富な場合 |
| 8.0 |
過度に平滑 |
低い |
非推奨 |
αパラメータのチューニング
| α値 |
ソフトラベル重み |
ハードラベル重み |
適用シーン |
| 0.3 |
30% |
70% |
データ品質が低い |
| 0.5 |
50% |
50% |
バランス型 |
| 0.7 |
70% |
30% |
データ品質が高い(推奨) |
| 0.9 |
90% |
10% |
教師が非常に強力 |
グレーボックス蒸留:特徴量アライメント実践
中間層特徴量アライメント
class FeatureDistillationLoss(nn.Module):
def __init__(self, teacher_layers, student_layers, projection_dim=2048):
super().__init__()
self.teacher_layers = teacher_layers
self.student_layers = student_layers
self.layer_mapping = self._build_layer_mapping()
self.projections = nn.ModuleDict({
str(s_layer): nn.Linear(
student_dim, projection_dim, bias=False
)
for s_layer, student_dim in student_layers
})
def _build_layer_mapping(self):
t_count = len(self.teacher_layers)
s_count = len(self.student_layers)
return {
s_idx: int(s_idx * t_count / s_count)
for s_idx in range(s_count)
}
def forward(self, teacher_hidden, student_hidden):
total_loss = 0.0
for s_idx, t_idx in self.layer_mapping.items():
s_feat = student_hidden[s_idx]
t_feat = teacher_hidden[t_idx]
s_proj = self.projections[str(s_idx)](s_feat)
s_norm = F.normalize(s_proj, dim=-1)
t_norm = F.normalize(t_feat, dim=-1)
total_loss += (2 - 2 * (s_norm * t_norm).sum(dim=-1)).mean()
return total_loss / len(self.layer_mapping)
class AttentionDistillationLoss(nn.Module):
def __init__(self, num_heads=32):
super().__init__()
self.num_heads = num_heads
def forward(self, teacher_attn, student_attn):
t_heads = teacher_attn.shape[1]
s_heads = student_attn.shape[1]
head_mapping = {
s: int(s * t_heads / s_heads)
for s in range(s_heads)
}
total_loss = 0.0
for s_h, t_h in head_mapping.items():
s_attn = student_attn[:, s_h]
t_attn = teacher_attn[:, t_h]
total_loss += F.mse_loss(s_attn, t_attn)
return total_loss / len(head_mapping)
統合蒸留損失
class CombinedDistillationLoss(nn.Module):
def __init__(self, config):
super().__init__()
self.logit_loss = LogitDistillationLoss(
temperature=config.temperature,
alpha=config.logit_alpha,
)
self.feature_loss = FeatureDistillationLoss(
teacher_layers=config.teacher_layers,
student_layers=config.student_layers,
)
self.attn_loss = AttentionDistillationLoss(
num_heads=config.num_heads,
)
self.w_logit = config.w_logit
self.w_feature = config.w_feature
self.w_attn = config.w_attn
def forward(self, student_outputs, teacher_outputs, labels):
l_logit = self.logit_loss(
student_outputs.logits, teacher_outputs.logits, labels
)
l_feature = self.feature_loss(
teacher_outputs.hidden_states,
student_outputs.hidden_states,
)
l_attn = self.attn_loss(
teacher_outputs.attentions,
student_outputs.attentions,
)
return (
self.w_logit * l_logit
+ self.w_feature * l_feature
+ self.w_attn * l_attn
)
| 損失の組み合わせ |
w_logit |
w_feature |
w_attn |
効果 |
| Logitのみ |
1.0 |
0 |
0 |
ベースライン |
| Logit+特徴量 |
0.5 |
0.5 |
0 |
+3% |
| Logit+アテンション |
0.5 |
0 |
0.5 |
+2% |
| 全組み合わせ |
0.4 |
0.35 |
0.25 |
+5% |
段階的蒸留3フェーズ
3フェーズの流れ
┌──────────────────────────────────────────────────────────────┐
│ 段階的蒸留3フェーズ │
│ │
│ フェーズ1: 70B → 30B │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ 教師: 70B全パラメータ → 学生: 30B │ │
│ │ 戦略: ホワイトボックス蒸留(Logit+特徴量+アテンション) │ │
│ │ データ: 500万件の高品質指示 │ │
│ │ 知識保持: 93% │ │
│ └──────────────────────────────────────────────────────┘ │
│ ↓ │
│ フェーズ2: 30B → 14B │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ 教師: 30B蒸留モデル → 学生: 14B │ │
│ │ 戦略: グレーボックス蒸留(Logit+特徴量) │ │
│ │ データ: 300万件+フェーズ1データの混合 │ │
│ │ 知識保持: 91% │ │
│ └──────────────────────────────────────────────────────┘ │
│ ↓ │
│ フェーズ3: 14B → 7B │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ 教師: 14B蒸留モデル → 学生: 7B │ │
│ │ 戦略: グレーボックス蒸留(Logit) + オンラインデータ拡張 │ │
│ │ データ: 200万件+前2フェーズのデータ │ │
│ │ 知識保持: 88% │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ 総合知識保持: 93% × 91% × 88% ≈ 74% │
│ 直接70B→7B蒸留: 約65% │
│ 段階的蒸留の優位性: +9% │
└──────────────────────────────────────────────────────────────┘
段階的蒸留の実装
class ProgressiveDistillation:
def __init__(self, config):
self.stages = config.stages
self.current_stage = 0
def get_stage_config(self, stage_idx):
stage_configs = [
{
"teacher_size": "70B",
"student_size": "30B",
"strategy": "white_box",
"data_size": 5_000_000,
"epochs": 3,
"lr": 5e-5,
"loss_weights": {
"logit": 0.4, "feature": 0.35, "attn": 0.25
},
},
{
"teacher_size": "30B",
"student_size": "14B",
"strategy": "gray_box",
"data_size": 3_000_000,
"epochs": 4,
"lr": 3e-5,
"loss_weights": {
"logit": 0.5, "feature": 0.5, "attn": 0.0
},
},
{
"teacher_size": "14B",
"student_size": "7B",
"strategy": "gray_box_logit",
"data_size": 2_000_000,
"epochs": 5,
"lr": 2e-5,
"loss_weights": {
"logit": 0.7, "feature": 0.3, "attn": 0.0
},
},
]
return stage_configs[stage_idx]
def run_stage(self, stage_idx, teacher_model, student_model, dataset):
config = self.get_stage_config(stage_idx)
trainer = DistillationTrainer(
teacher_model=teacher_model,
student_model=student_model,
train_dataset=dataset,
args=DistillationArguments(
num_train_epochs=config["epochs"],
learning_rate=config["lr"],
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
warmup_ratio=0.03,
bf16=True,
logging_steps=100,
save_strategy="epoch",
strategy=config["strategy"],
loss_weights=config["loss_weights"],
),
)
trainer.train()
return trainer.model
def run_all_stages(self, initial_teacher, dataset):
current_teacher = initial_teacher
for stage_idx in range(len(self.stages)):
stage_config = self.get_stage_config(stage_idx)
student = self._init_student(stage_config["student_size"])
distilled = self.run_stage(
stage_idx, current_teacher, student, dataset
)
current_teacher = distilled
self._save_checkpoint(distilled, stage_idx)
return current_teacher
段階的蒸留 vs 直接蒸留の比較
| 手法 |
70B→7B |
学習GPU時間 |
最終MMLU |
最終HumanEval |
| 直接蒸留 |
1ステップ |
2000 |
58.2 |
42.1 |
| 段階的(2フェーズ) |
70B→14B→7B |
2400 |
61.5 |
46.8 |
| 段階的(3フェーズ) |
70B→30B→14B→7B |
2800 |
63.1 |
48.3 |
| 段階的+データ拡張 |
3フェーズ+拡張 |
3200 |
65.2 |
51.7 |
マルチタスク蒸留と破滅的忘却
マルチタスク蒸留アーキテクチャ
class MultiTaskDistillation:
def __init__(self, teacher, student, task_weights):
self.teacher = teacher
self.student = student
self.task_weights = task_weights
self.task_buffers = {}
def compute_loss(self, batch, task_id):
with torch.no_grad():
teacher_out = self.teacher(**batch)
student_out = self.student(**batch)
logit_loss = F.kl_div(
F.log_softmax(student_out.logits / 2.0, dim=-1),
F.softmax(teacher_out.logits / 2.0, dim=-1),
reduction="batchmean",
) * 4.0
return logit_loss
def ewc_penalty(self, task_id):
if task_id not in self.task_buffers:
return 0.0
buffer = self.task_buffers[task_id]
penalty = 0.0
for name, param in self.student.named_parameters():
if name in buffer["params"]:
penalty += (
(param - buffer["params"][name]) ** 2
).sum() * buffer["importance"][name]
return penalty
def save_task_buffer(self, task_id, dataloader):
params = {}
importance = {}
for name, param in self.student.named_parameters():
params[name] = param.data.clone()
importance[name] = torch.zeros_like(param.data)
self.student.eval()
for batch in dataloader:
self.student.zero_grad()
loss = self.student(**batch).loss
loss.backward()
for name, param in self.student.named_parameters():
if param.grad is not None:
importance[name] += param.grad.data ** 2
for name in importance:
importance[name] /= len(dataloader)
self.task_buffers[task_id] = {
"params": params,
"importance": importance,
}
破滅的忘却の防御戦略
| 戦略 |
原理 |
計算オーバーヘッド |
効果 |
| EWC |
パラメータ重要度正則化 |
低い |
中程度 |
| Replay Buffer |
旧タスクデータのリプレイ |
中程度 |
良い |
| LoRA Adapter |
タスクごとに独立したアダプタ |
低い |
良い |
| Progressive Freezing |
段階的な層のフリーズ |
低い |
中程度 |
| ハイブリッド戦略 |
EWC+Replay+LoRA |
中程度 |
最良 |
DeepSeek-R1蒸留スキーム解析
DeepSeek-R1の蒸留スキームは2026年で最も成功したLLM蒸留事例の一つである:
| フェーズ |
教師モデル |
学生モデル |
データ量 |
キーテクノロジー |
| フェーズ1 |
R1-671B |
R1-Distill-70B |
80万件 |
ブラックボックス蒸留+CoTデータ |
| フェーズ2 |
R1-Distill-70B |
R1-Distill-32B |
60万件 |
ホワイトボックスLogit蒸留 |
| フェーズ3 |
R1-Distill-32B |
R1-Distill-8B |
40万件 |
グレーボックス蒸留+データ拡張 |
主なイノベーション:
- CoTデータ蒸留:最終的な回答だけでなく、推論過程も蒸留
- 拒否サンプリング:教師モデルが複数の推論パスを生成し、最適なパスを選択して蒸留
- カリキュラム学習:簡単なものから複雑なものへと段階的に蒸留データの難易度を上げる
まとめと関連情報
主要ポイントの振り返り
- 戦略選択:ホワイトボックスは効果が最も高いがコストも高い、ブラックボックスは最も柔軟だが効果が低い、グレーボックスが最適なバランス
- 段階的蒸留:3フェーズの段階的蒸留は直接蒸留より9%以上の知識保持率を向上
- マルチタスク蒸留:EWC+Replay+LoRAのハイブリッド戦略が破滅的忘却を効果的に防止
- DeepSeek-R1:CoTデータ蒸留+拒否サンプリングが現在最先端の蒸留スキーム
蒸留スキームの推奨
| シーン |
推奨スキーム |
期待効果 |
| 教師の重みあり+十分なGPU |
ホワイトボックス3フェーズ段階的 |
知識保持85%以上 |
| 教師の重みあり+GPUが限定的 |
グレーボックス2フェーズ |
知識保持80%以上 |
| APIのみ |
ブラックボックス+データ拡張 |
知識保持70%以上 |
| マルチタスクシーン |
ハイブリッド戦略+EWC |
各タスクのバランス |
蒸留データのJSON形式変換が必要ですか?JSONフォーマットツールとテキスト差分比較で、蒸留データセットを素早く処理できます。
関連記事