Zyphra、ZAYA1-8B-Diffusion-Previewをリリース：自己回帰LLMから変換された初のMoE拡散モデル、最大7.7倍の高速化

Zyphra（サンフランシスコ拠点のAIラボ）は、ZAYA1-8B-Diffusion-Previewをリリースしました。これは、拡散言語モデルにおける初期の研究成果のプレビュー版です。このリリースは、既存の自己回帰言語モデルを離散拡散モデルに変換しても、評価性能に体系的な損失がなく、AMDハードウェア上で顕著な推論高速化を実現できることを示しています。

自己回帰デコードの問題 この重要性を理解するには、まずほとんどの言語モデルが現在どのようにテキストを生成しているかを理解する必要があります。標準的な大規模言語モデルは自己回帰型です。つまり、トークンを1つずつ順次デコードします。新しいトークンごとに、アテンションメカニズムは以前に生成されたすべてのトークンを振り返り、GPUメモリからそれらの保存表現（KVキャッシュ）をロードする必要があります。重要なのは、バッチ内のユーザーごとにトークン履歴が異なるため、各ユーザーのKVキャッシュは個別にロードする必要があり、リクエスト間で共有できないことです。これによりボトルネックが発生します。GPUが実際の計算よりもメモリからのデータ移動に多くの時間を費やすと、システムは計算制約ではなくメモリ帯域幅制約になります。これにより、メモリ帯域幅よりも高速に計算FLOPsを拡張している最新のGPUハードウェアを推論中に効率的に利用することが制限されます。

拡散モデルは代替手段を提供します。拡散モデルは、トークンを1つずつ生成する代わりに、N個のトークンの複数のドラフトを同時に生成し、このドラフトプロセスを複数回繰り返します。ブロック内のすべてのNトークンが同じKVキャッシュを共有するため、操作はメモリ帯域幅制約から計算制約に移行し、GPUをより効率的に利用できます。ZAYA1-8B-Diffusion-Previewでは、モデルはブロック内の各トークンに対してマスクからトークンへの単一ステップ変換（つまり、反復的ノイズ除去ではなく、マスクされていないトークンを直接予測）を実行します。

自己回帰から拡散への変換：ゼロからのトレーニング不要 拡散言語モデルをゼロからトレーニングすることは技術的に難しく、確立されたレシピはほとんどありません。Zyphraチームは、ゼロからのトレーニングよりも変換を好む理由として2つ挙げています。第一に、ゼロからのトレーニングは単純に難しく、既知のレシピがほとんどないこと。第二に、拡散モードでトレーニングする利点はないことです。なぜなら、トレーニングはすでに計算制約であり、拡散が解決するメモリ帯域幅のボトルネックは推論時にのみ現れるからです。つまり、拡散の利点はすべて推論時の利点であり、既存の事前学習スタックをそのまま再利用できます。

TiDARレシピに基づき、ZyphraはZAYA1-8B-baseチェックポイントを取得し、さらに6000億トークンの拡散変換中間トレーニング（コンテキスト長32k）を実施し、続いて5000億トークンのネイティブコンテキスト拡張（128k）、そして拡散教師ありファインチューニング（SFT）フェーズを実施しました。

ZAYA1-8B-Diffusion-Previewは、自己回帰LLMから変換された初のMoE拡散モデルであり、AMD GPUでトレーニングされた初の拡散言語モデルです。Zyphraは、ベースの自己回帰チェックポイントと比較して評価性能の低下は最小限であり、LCB-v6などの一部のベンチマークでは改善が見られたと報告しています。その理由として、改善された中間トレーニングデータセットと、因果的自己回帰と比較した拡散スタイルのブロック内非因果推論の表現力の高さを挙げています。

拡散サンプラーの仕組み 推論中、ZAYA1-8B-Diffusion-Previewは16トークンのドラフトを同時に生成します。これらのトークンの一部は、投機的デコードから借用したサンプリング基準に基づいて受け入れられます。重要なのは、同じモデルが単一のフォワードパス内で提案者と検証者の両方の役割を果たすことで、従来のEAGLEやdFlashのような2つの別々のモデルを実行するオーバーヘッドを排除している点です。メモリ帯域幅が非常に制約される状況では、受け入れられたトークンのほとんどは自己回帰デコードに対する無料の高速化を意味します。GPUはすでにロードされており、追加トークンの計算コストは非常に低くなります。

Zyphraチームは、速度と品質のトレードオフが異なる2つのサンプラーを報告しています：

• ロスレス拡散サンプラー：標準的な投機的デコード受入基準min(1, p(x)/q(x))を使用。pは自己回帰モデルのロジット分布、qは拡散モデルの分布。拒否された場合、次のトークンはp(x)-q(x)の残差分布からサンプリングされます。このサンプラーは体系的な評価劣化なしに4.6倍の高速化を実現します。
• ロジットミキシングサンプラー：まず拡散提案者と自己回帰モデルのロジットを混合し、平均化された分布を検証に使用します。これにより、検証ロジットが拡散ロジットに近くなるため受入率が向上しますが、品質にはいくらか影響があります。このサンプラーは7.7倍の高速化を実現します。速度と品質のトレードオフは実行時に選択できます。

重要な注意点として、ZAYA1-8B-Diffusion-PreviewはRLトレーニングを受けていないベース中間チェックポイントであるため、Zyphraは標準的な精度ベンチマークではなくpass@評価を使用して、RLトレーニング後のモデルの最終的な可能性をより適切に表現しています。これらの数値を他のモデルの報告されたベンチマークと比較する読者は、この点に留意する必要があります。

Zyphraチームはまた、拡散による高速化はマルチトークン予測（MTP）やEAGLE3などのさまざまな投機的デコード戦略といった代替手法よりも高いと指摘しています。TiDARスタイルの拡散モデルは単一のフォワードパスのみを使用するため、dFlashと同等の受入率でも大幅な高速化が得られます。

アーキテクチャ詳細 ZAYA1-8B-Diffusion-Previewは、順序制約生成を使用する単一ステップ投機的拡散モデルです。つまり、拡散モデルはプレフィックスから始まる連続した部分列のトークンのみを生成できます。この制約により、制約なしのマスク拡散目標やセットブロックデコードと比較してトレーニングの安定性が大幅に向上し、これがZyphraがTiDARレシピを採用した主な理由です。

このモデルは、ZyphraのCCAアテンションバリアントを使用しています。CCAはアテンションにおけるプリフィルFLOPsを大幅に削減し、拡散にとって直接的利点があります。なぜなら、拡散はデコードをプリフィル類似の操作に変換するからです。つまり、CCAにより、モデルは計算限界に達する前に、より多くのトークンを並列に拡散できるようになります。

具体的には、アーキテクチャはCCGQAを使用し、クエリヘッドとキーヘッドの比率は4:1です。この設計選択の背後には、MLA（マルチヘッド潜在アテンション）を意図的に避けている点があります。MLAの高い演算強度はCCGQAと比較してミスマッチと見なされたためです。ブロック拡散は同じキャッシュにアクセスするため、演算強度はブロックサイズとフォワードパスあたりのブロック数に比例します。AMD MI300xハードウェア（bf16）では、システムは1フォワードパスあたりおよそ3つのブロックサイズ提案をサポートし、MI355xでは約5つに増加します。CCGQAは2倍圧縮でも動作するため、ZyphraはTiDAR中間トレーニングに関連する追加のトレーニングFLOPsを負担できました。AMD GPUハードウェアのより大きなVRAM容量により、拡散トレーニング全体がさらに効率的になりました。

実際には、理論的な高速化を達成することはより困難です。なぜなら、拡散には追加の運用オーバーヘッドがあり、拡散モデルの推論スタックは自己回帰推論の成熟したツールに比べて大幅に最適化されていないからです。