Apriel-H1: Mambaハイブリッドによる効率的な推論モデルの蒸留

ServiceNow AIは、フルアテンションの教師モデルをMambaハイブリッドアーキテクチャに変換する15B推論モデルのファミリーであるApriel-H1を開発しました。このアプローチは、推論品質の損失を最小限に抑えながらスループットを最大2.1倍向上させ、ターゲットを絞った蒸留によって既存の強力なモデルに効率を後付けできることを示しています。

コアな洞察: 能力にマッチした蒸留

推論モデルを蒸留するには、一般目的のプリトレーニングデータではなく、保持したい特定の能力に合ったデータを使用する必要があります。プリトレーニングデータまたはプリトレーニングとSFTデータの混合を使用した初期の蒸留試みでは、推論品質が劇的に低下しました。

成功したアプローチでは、教師のSFTデータセットから高品質な推論トレースを使用し、以下を含みます：

  • マルチステップの数学的証明
  • 明確な論理的依存関係を持つコーディングタスク
  • 詳細な説明の連鎖を伴う科学的分析

これは、注意メカニズム（検索ヘッドや誘導ヘッドなど）を線形再帰に置き換えると、推論に使用される計算パスが乱されるため必要です。高品質な推論トレースは、ハイブリッドモデルが同じ結果に至る新しいパスを見つけるために必要な明示的な構造を提供します。さらに、チームは逆KLダイバージェンス（temperature 1）を蒸留に使用しました。これはモード探索的挙動により、学生モデルが教師の高信頼度かつ構造化された予測にコミットすることを促すからです。

Apriel-H1のパフォーマンスとベンチマーク

フラッグシップモデルであるApriel-H1-15b-Thinker-SFTは、フルアテンションの教師モデルと比較して2.1倍のスループットを達成します。推論ベンチマークへの影響は以下の通りです：

ベンチマーク 教師スコア Apriel-H1-15b-Thinker-SFTスコア
MATH500 0.90 $
ightarrow$ 0.92 改善
MTBench 8.30 $
ightarrow$ 8.58 改善
GSM8k 0.97 $
ightarrow$ 0.95 わずかな後退
GPQA 0.59 $
ightarrow$ 0.55 わずかな後退
AIME24 0.70 $
ightarrow$ 0.65 わずかな後退

Apriel-H1ファミリー全体では、チェックポイントは50層中の25〜40のMamba層の範囲です。H-30バリアントが最良のバランスを提供する一方で、H-40バリアントはレイテンシーが重要なワークロードに対してスループットを最大化します（最大3.4倍）。

段階的蒸留プロセス

フルアテンションモデルをハイブリッドに変換することは、一手順ではできません。ServiceNow AIは3段階の手順を採用しました：

ステージ1: 初期レイヤー置換

MMLUでのLeave-One-Out（LOO）分析を使用して、チームは最も重要度の低い25層を特定しました。これらは、Mamba-in-Llama（MIL）で初期化されたミキサーに置き換え、エンドツーエンドで蒸留されました。

ステージ2: プログレッシブ変換

LOO分析は25層を超えると信頼できなくなるため、チームは**MIL-Mamba-Replacement (MMR)**を実装しました。この動的ヒューリスティックは、Mambaミキサーを初期化し、100ステップのトレーニングを実行して蒸留損失を記録します。損失が低く収束するレイヤーを優先的に置換対象とします。このプロセスは増分的に行われました（25 $ ightarrow$ 27 $ ightarrow$ 30 $ ightarrow$ 34 $ ightarrow$ 37 $ ightarrow$ 40層）。

ステージ3: 最終SFT

目標のMamba層数に到達後、推論性能が安定するまで最終SFTパスを実施しました。最終的なApriel-H1-15b-Thinker-SFTモデルは、55.9Bの蒸留トークンと20.9BのSFTトークンを使用して生成されました。

実装と再現性

Apriel-H1は、Fast-LLMを使用して構築されました。これは、注意とMambaをモジュラーなミキシングインターフェースとして扱うオープンソース（Apache 2.0）トレーニングフレームワークです。これにより、構成のpatternフィールドを通じてレイヤーの順序を指定でき、ブロックタイプの自由な交換が可能になります。

本番デプロイメント

Apriel-H1はHugging Face TransformersおよびvLLMと統合されています。vLLMの統合では、連続バッチング、プレフィックスキャッシング、チャンク済みプリフィルのためにMambaキャッシュ操作を利用しています。著者らは、スループットの向上は大きいものの、ツールが成熟途中であるため、ハイブリッドモデルのデプロイには現在カスタムコードと慎重な数値検証が必要であると述べています。

技術仕様

  • Mambaハイパーパラメータ: ステートサイズ 16, DTランク 16, 内部次元 4096.
  • アーキテクチャ: Mamba-1と注意レイヤーのハイブリッド.
  • トレーニングデータ: フラッグシップモデルのための総トークン数76.8B.

Sources