OpenRouter蒸留によりGPT-5 Pro上でMixtralをトレーニングした方法

1. 概要と動機

Shannon AIの構築無修正AIモデルのAIレッドチーム研究には、フロンティアレベルの能力をオープンウェイトアーキテクチャに転送する必要がありました。私たちの解決策：OpenRouter APIを介してGPT-5 ProからMixtralのMixture-of-Expertsフレームワークに知識を蒸留することです。

なぜGPT-5 Proなのか？

GPT-5 Proは現在の能力の最前線を代表し、以下の点で優れています：

• 複雑な多段階推論
• コード生成と分析
• 微妙な言語理解
• 幅広い知識範囲

なぜMixtralなのか？

Mixtralのアーキテクチャは、私たちの研究に独自の利点を提供します：

• 完全な透明性を可能にするオープンウェイト
• 効率的なMoE設計（アクティブパラメータはわずか12.9B/39B）
• ファインチューニングのための強力なベースライン能力
• 研究変更を許可するApache 2.0ライセンス

2. 蒸留アーキテクチャ

OpenRouter統合

OpenRouterの統合APIを利用してGPT-5 Proにアクセスし、いくつかの利点を得ました：

• コスト効率：直接APIアクセスと比較して競争力のある価格設定
• レート制限：大規模生成のための管理されたスループット
• フォールバックルーティング：データ収集の継続性を保証する自動フェイルオーバー
• 応答キャッシング：類似プロンプトのコスト削減

import openai
from typing import Generator

class OpenRouterDistillation:
    def __init__(self):
        self.client = openai.OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.environ["OPENROUTER_API_KEY"]
        )
        self.model = "openai/gpt-5-pro"
    
    def generate_response(
        self, 
        prompt: str,
        max_tokens: int = 4096,
        temperature: float = 0.7
    ) -> str:
        """Generate GPT-5 Pro response for distillation."""
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
            temperature=temperature,
            extra_headers={
                "HTTP-Referer": "https://shannon.ai",
                "X-Title": "Shannon AI Distillation"
            }
        )
        return response.choices[0].message.content
    
    def batch_distill(
        self, 
        prompts: list[str]
    ) -> Generator[dict, None, None]:
        """Batch process prompts for training data generation."""
        for prompt in prompts:
            response = self.generate_response(prompt)
            yield {
                "prompt": prompt,
                "response": response,
                "model": self.model,
                "timestamp": datetime.utcnow().isoformat()
            }

3. データ収集パイプライン

プロンプトキュレーション戦略

包括的な能力転送を確実にするため、当社のプロンプトは複数のドメインにわたって慎重にキュレーションされました：

• 推論（35%）：数学、論理、科学分析
• コード（25%）：20以上の言語での生成、デバッグ、説明
• 知識（20%）：事実に基づくクエリ、合成、分析
• 創造性（10%）：執筆、ブレインストーミング、アイデア出し
• レッドチーム (10%):エッジケース、敵対的プロンプト、境界テスト

品質フィルタリング

すべてのGPT-5 Proの応答がトレーニングに適しているわけではありませんでした。厳格なフィルタリングを適用しました。

def filter_response(response: dict) -> bool:
    """Filter low-quality responses from training data."""
    
    # Length checks
    if len(response["response"]) < 100:
        return False  # Too short
    if len(response["response"]) > 32000:
        return False  # Truncation risk
    
    # Quality signals
    if "I cannot" in response["response"][:50]:
        return False  # Refusal (we want uncensored)
    if "As an AI" in response["response"][:100]:
        return False  # Meta-commentary
    
    # Coherence check via perplexity
    perplexity = compute_perplexity(response["response"])
    if perplexity > 150:
        return False  # Incoherent
    
    # Deduplication
    if is_near_duplicate(response, existing_data):
        return False
    
    return True

フィルタリング後、約1.8Mの高品質ペアをトレーニング用に保持しました。

4. トレーニング手法

ステージ1: 教師ありファインチューニング (SFT)

フィルタリングされたGPT-5 Proの応答に対する標準SFTによる初期能力転送:

# Shannon V1 SFT Configuration
model:
  base: mistralai/Mixtral-8x7B-v0.1  # or 8x22B for Deep
  dtype: bfloat16
  load_in_4bit: false

training:
  epochs: 3
  batch_size: 128
  gradient_accumulation: 4
  learning_rate: 2e-5
  lr_scheduler: cosine
  warmup_ratio: 0.03
  weight_decay: 0.01
  max_seq_length: 8192

data:
  train_path: /data/gpt5_distilled_train.jsonl
  eval_path: /data/gpt5_distilled_eval.jsonl
  format: sharegpt

lora:  # For efficient fine-tuning
  r: 64
  alpha: 128
  dropout: 0.05
  target_modules: 
    - q_proj
    - k_proj
    - v_proj
    - o_proj
    - gate_proj
    - up_proj
    - down_proj

ステージ2: 直接選好最適化 (DPO)

モデルの挙動をさらに調整し、拒否を減らすために、選好ペアを使用してDPOを適用しました。

• 選択済み:完全で役立つGPT-5 Proの応答
• 拒否済み:拒否、部分的な応答、または低品質な出力

インフラストラクチャ

トレーニングは、当社の専用計算クラスターで実施されました。

• ハードウェア:8× NVIDIA H100 80GBノード
• フレームワーク:PyTorch 2.1 + DeepSpeed ZeRO-3
• トレーニング時間:8×7Bで約72時間、8×22Bで約168時間
• 総計算量:約15,000 H100-時間

5. 結果とベンチマーク

トレーニング後の評価により、知識の転送が成功したことが示されました。

*GPT-5 Proは、安全トレーニングのため、ほとんどのレッドチームプロンプトを拒否します。

6. 学んだ教訓

うまくいったこと

• 多様なプロンプトが不可欠でした。狭いデータセットは能力の崩壊につながりました。
• 制約緩和のためのDPOは、モデルが一般的な拒否を回避することを効果的に教えました。
• OpenRouterの信頼性により、数か月にわたる一貫したデータ収集が可能になりました。
• 品質フィルタリングは、最終モデルの一貫性を大幅に改善しました。

克服された課題

• レート制限:複数のAPIキーにわたる分散収集が必要でした。
• 応答のばらつき:GPT-5 Proの確率的性質により、プロンプトごとに複数のサンプルが必要でした。
• コスト管理:慎重なプロンプトエンジニアリングにより、平均応答長が30%削減されました。
• MoEの不安定性:エキスパート層に特化した学習率スケジューリングが必要でした。

今後の方向性

当社の蒸留パイプラインは進化を続けています。今後の改善点には以下が含まれます。

• リアルタイム選好学習を伴うオンライン蒸留
• GPT-5 Pro + Claude + Geminiを組み合わせたマルチティーチャー蒸留
• 混合エキスパートファインチューニングによる専門ドメインエキスパート