大規模モデルの知識蒸留実践:教師モデルから学生モデルへの圧縮の道

AI与大数据

概要

  • 知識蒸留は大規模モデルの実用化における重要なブリッジ: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
デプロイコスト $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万件 グレーボックス蒸留+データ拡張

主なイノベーション:

  1. CoTデータ蒸留:最終的な回答だけでなく、推論過程も蒸留
  2. 拒否サンプリング:教師モデルが複数の推論パスを生成し、最適なパスを選択して蒸留
  3. カリキュラム学習:簡単なものから複雑なものへと段階的に蒸留データの難易度を上げる

まとめと関連情報

主要ポイントの振り返り

  1. 戦略選択:ホワイトボックスは効果が最も高いがコストも高い、ブラックボックスは最も柔軟だが効果が低い、グレーボックスが最適なバランス
  2. 段階的蒸留:3フェーズの段階的蒸留は直接蒸留より9%以上の知識保持率を向上
  3. マルチタスク蒸留:EWC+Replay+LoRAのハイブリッド戦略が破滅的忘却を効果的に防止
  4. DeepSeek-R1:CoTデータ蒸留+拒否サンプリングが現在最先端の蒸留スキーム

蒸留スキームの推奨

シーン 推奨スキーム 期待効果
教師の重みあり+十分なGPU ホワイトボックス3フェーズ段階的 知識保持85%以上
教師の重みあり+GPUが限定的 グレーボックス2フェーズ 知識保持80%以上
APIのみ ブラックボックス+データ拡張 知識保持70%以上
マルチタスクシーン ハイブリッド戦略+EWC 各タスクのバランス

蒸留データのJSON形式変換が必要ですか?JSONフォーマットツールテキスト差分比較で、蒸留データセットを素早く処理できます。

関連記事

ブラウザローカルツールを無料で試す →

#大模型知识蒸馏#LLM蒸馏#模型压缩#教师学生模型#大模型推理优化#2026