品質低下なしでサブ秒の Ideogram v4 を提供

image @fal では、Ideogram V4 を高速で提供しています。1K 解像度において、従来 2.75 秒かかっていた生成が、現在は 0.44 秒でレンダリングされ、可視的な画質の低下なしに 6 倍の速度向上を実現しました。これは、考えられるあらゆる軸で画像生成コストを削減することで達成されたものであり、そのコストは「拡散トランスフォーマー（diffusion transformer）が実行するフォワードパスの数」×「各パスのコスト」という一つの式として捉えることができます。

Contents

FP4 とシステムレベルの最適化

量子化認識蒸留（Quantization Aware Distillation）

タイムステップ蒸留（Timestep Distillation）

image 二つの戦線があります。まず、各フォワードパスを安価にすることです：トランスフォーマーを FP4 で実行し、行列乗算（matmul）の周囲にある小さな演算を融合させて、メモリの往復コストを支払わないようにします。次に、パス数を減らすことです：二つのブランチから分類器フリーガイダンス（classifier-free guidance）を一つのブランチに蒸留し、多数のノイズ除去ステップを手数少ないステップに圧縮します。これらをすべて積み重ねれば、計算資源のわずかな一部で、完全な bf16 モデルと同じ画像を得ることができます。

この投稿では、私たちが構築した順序に沿って、全体の流れを概説します。まず、高速化を実現するカーネル融合を伴う FP4 順伝播から始め、次に 4 ビットが引き起こした品質問題と、それを解決した量子化意識型蒸留について述べます。最後に、ステップ数を劇的に削減したタイムステップ蒸留について解説します。FP4 は確かに高速ですが、素朴な実装の FP4 は bf16 と比較して視覚的に明らかに劣っており、さらにその場合でも多くのステップを要し、ガイダンスブランチも 2 つ実行する必要がありました。ここでは、これらすべての課題をどのように克服したかを示します。

FP4 でトランスフォーマーを実行する

最速の推論パスは、拡散トランスフォーマーを NVFP4 で実行するものです。FP4 は 4 ビット浮動小数点形式であり、1 ビットの符号ビット、2 ビットの指数部、1 ビットの仮数部から構成されるため、各値は限られたレベルのいずれかしか取り得ません。その代わり、意味のある速度向上と大幅なメモリフットプリントの削減が得られ、これによりより多くの並行処理を可能にし、画像あたりのレイテンシを低減できます。Blackwell アーキテクチャではこの形式がハードウェアでサポートされているため、行列乗算（matmul）はネイティブに 4 ビットで実行されます。

しかし、行列乗算が実際の時間消費の大部分を占める場合にのみ、4 ビットの行列乗算は効果を発揮します。拡散トランスフォーマーでは、必ずしもそうとは限りません。大きな GEMM（汎用行列乗算）の直後には、アテンションパス上の RMSNorm や FFN 内のゲート付き SiLU といった小さな演算が常に続きます。これらを融合しない場合、その小さな演算によって GEMM の出力全体を HBM（ハイバンド幅メモリ）へ書き出し、再び読み出す必要が生じます。行列乗算のコストは下がりましたが、除去しなかったメモリの往復転送が、そのメリットを相殺してしまいます。したがって、高速な FP4 パスにとって最初の課題は、これらの演算を GEMM 自体に融合させることです。

エピローグの融合により FP4 の効果を最大化する

エプログの融合について簡潔に

GEMM は、A @ B をタイル単位で生成します。各タイル内部では、Tensor Cores が乗算加算をストリーミングし、それを SM 自身のストレージ（Hopper ではレジスタ、Blackwell では TMEM）に格納されたアキュムレータへ蓄積していきます。このメインループが実行されている間、出力に対して HBM はまだ一切アクセスされておらず、答えはオンチップ上で集約されています。

しかし、A @ B がレイヤーの最終段階となることはほとんどありません。拡散トランスフォーマーでは、FFN は通常、GEMM に続くゲート付き SiLU（SwiGLU）で構成され、アテンション経路は GEMM に続く RMSNorm で構成されます。このパターンは LLM や DiT 全体に共通しており、「大きな GEMM の後に小さな後処理演算」が繰り返されます。演算自体は軽量ですが、問題は融合されていない場合、GEMM がタイル全体を HBM に書き出し、次のカーネルがそれを再度読み込んで演算を適用し、再び書き出すという点にあります。わずかな FLOPs（浮動小数点演算）が、フルなメモリ往復と第 2 のカーネル起動に繋がってしまいます。

エピローグ融合（Epilogue fusion）とは、GEMM の直後に、結果を HBM に保存する前に適用される演算のことです。出力はグローバルメモリへ一度だけ書き出され、すでに変換済みであり、第 2 のカーネルも不要で、メモリ往復も発生しません。Hopper および Blackwell では、エピローグ処理を次のタイルのメインループと重ねて実行できるため、小規模な演算においては追加計算コストはほぼ隠蔽されます。

CUTLASS はこれを Epilogue Visitor Trees (EVT) を通じて公開しています。各アキュムレータ断片に対して実行される visit() 関数を書き、集合的なエピローグがロードとストアを処理します。このメカニズムは Colfax の EVT に関する記事で詳しく説明されています。簡潔に言えば、ポイントワイズ融合（バイアス、ReLU、キャスト）はきれいに実装されます：各出力要素は独立しており、跨るべき縮約要素はなく、CUTLASS はこれら用の事前構築済みビジターを備えています。

本稿では RMSNorm の融合に焦点を当てます。これはベクトルをその自身の RMS で正規化します：y = x · rsqrt(mean(x²) + eps) · weight。ポイントワイズ演算とは異なり、全体ヘッドにわたる rstd が必要であり、これは縮約処理です。この縮約が融合を困難にする要因ですが、なぜそうなるのかを理解するには、エピローグが実際にはデータをどのように認識しているかを検討する必要があります。

Blackwell におけるタイルと断片

GEMM は CTA タイル単位で動作します：Blackwell (SM100) では、当社の NVFP4 GEMM は 256×256 の CTA タイルを使用します。エピローグはこの全体を一度に処理するのではなく、より小さな 128×128 のエピローグタイルで順次走査し、CTA あたり 4 つのエピローグサブタイルを提供します：

各エピローグのサブタイルは128スレッドに分割されており、各スレッドは128要素の断片（FragmentSize = 128）を保持します。visit() コールバックは一度に一つの断片を受け取り、そのスレッドが責任を持つ128値となります。これがエピローグが扱う作業単位であり、何と融合できるかに影響します。

これにより、head_dim=128 の RMSNorm 融合は極めて容易になります：128要素からなるヘッドはまさに一つのスレッドの断片に相当します。visit() が呼び出された際、そのスレッドはすでに自身のレジスタ内に一つのヘッドの全128値を保持しています。RMSNorm を計算するには、これら128値に対して平方和を求め（スレッド間でのシャッフルも不要、二番目の断片も不要、追加パスも不要）、スケーリングして保存します。この集約は単一の visit() 呼び出しで完結します。

Ideogram v4 はそれほど都合よくできていません。そのアテンションでは head_dim=256 を使用するため、一つのヘッドはもはや一つの断片に収まらず、二つの断片にまたがることになります。そのため、きれいな単一 visit() での集約が崩れてしまいます。

なぜ head_dim=256 が難しいのか

これが問題となる理由を理解するには、エピローグの実行方法を見てみましょう。メインループの後、CUTLASS の集合型エピローグは CTA タイルの4つのサブタイルを固定された順序（Tile 0 → Tile 1 → Tile 2 → Tile 3）でループ処理します。各サブタイルに対して、断片ごとに visit() を一度呼び出し、結果を保存して次に進みます。これは単一の順方向パスであり、以前のサブタイルに戻ることはできません。

RMSNorm では、ヘッド全体に対する rstd（標準化された逆平方根）が必要であり、これは rstd = rsqrt(mean(x²) + eps) という計算で求められます。その後、y = x * rstd * weight となります。これを計算するには両方の半分における二乗和が必要ですが、それらは別々の visit() コールで到着します：最初の半分は Tile 0 で、2 番目の半分は Tile 1 です。Tile 1 を見て rstd を計算できる頃には、Tile 0 はすでに訪問されて保存されています。後から戻って修正する方法はありません。このセクションの残りは、それでもどうやって融合したかを示すものであり、最初の半分をレジスタ内に保持する 2 パスの再訪スケジュールです。

再訪融合

CUTLASS のエピローグ集合体は、固定された単一パスのループ順序で動作します：各サブタイルを一度だけ訪問し、保存して次に進みます。これは私たちの問題を解決しません：Tile 0 を正規化するには Tile 1 を先に確認する必要があるため、戻る必要があります。そこで私たちは、標準的な単一の前方パスではなく、減少計算が必要とする順序（すでに訪れたタイルに戻ってくることも含む）で訪問者がサブタイルを歩くことを可能にするエピローグを作成しました。

各行グループごとのスケジュールは 3 回のタッチです：

Tile 0、パス 0：これはヘッドの最初の半分です。その二乗和を計算して次に進みますが、2 番目の半分がない限り正規化はまだできません。

タイル 1、パス 1：頭の後半が到着します。これで両方の半分が揃いましたので、この半分の二乗和を計算し、最初の半分に加算して rstd を取得します。これを正規化して、この半分を保存します。

タイル 0、パス 2（再訪）：rstd が既知になったので、タイル 0 に戻ってそれをスケーリングします。これはシングルパスのエピローグでは行えない追加の手順です。

すべての処理がオンチップ内で完結します。グローバルメモリへの往復も、第 2 のカーネル実行も不要です。カスタムエピローグが 3 回のタッチスケジュールを歩き回り、ビジターがそのタッチ間で集約状態を引き継ぎ、rstd が準備できたらスケーリングを適用します。これにより、HBM に到達する前に正規化が融合されます。

ゲート付き SiLU (SwiGLU) の融合

拡散トランスフォーマー（および LLaMA などの大規模言語モデル）でよく見られる別の一般的なパターンに、ゲート付き SiLU があります。1 つの GEMM で出力 C を生成した後、レイヤーはそれを半分に分けて乗算することで完了します。

C = A @ B

output = SiLU(C[:, :N//2]) * C[:, N//2:]

これにより出力幅が半分に削減され、2 つの半分が 1 つになります。RMSNorm の場合と同じく、計算自体は軽量ですが、融合されない第 2 のカーネルでは HBM への往復を完全に支払う必要があります。CUTLASS にはこの「ペア + N/2 の集約」パターン用の組み込みエピローグがないため、これを融合するにはカスタムビジターが必要です。

課題はレイアウトにあります。出力列 n に対して、対応するペアは gate = C[:, n] と up = C[:, n + N/2] であり、これらは N の半分の距離だけ離れています。これらの 2 つの値は、異なるフラグメントやサブタイル、さらには異なる CTAs に配置される可能性が十分にあります。タイル間での同期を行うと、エピローグのパラレルモデルが破綻してしまいます。

代わりに、重みパッキング時に B の列を一度並べ替えます。GEMM はその後、2 つの分離した半分ではなく、インターリーブされたペア [gate0, up0, gate1, up1, ...] を生成します。隣接性が保証されるため、カスタム EVT ビジターが同じフラグメント内の隣接要素に対して SiLU と乗算を適用し、縮小された結果を保存します。これは 1 つのカーネルで完結し、タイル間の通信は不要です。

重みパッキング、ビジターのフック、EVT の配線に関する完全なウォークスルーは、「Crafting Efficient Kernels with Epilogue Fusion」に記載されています。

imageOG 重みは FP4 量子化では高品質な画像を提供できない

落とし穴：FP4 は画質にコストを強いる

エピローグの融合により、高速な FP4 フォワードパスが実現します。しかし、それ単独では画質が悪化する結果にもなります。標準的な FP4 推論は bf16 に比べて視覚的に劣っており、特に色再現において顕著でした。赤色が褪せ、彩度が低下し、ロゴや製品、植物など色彩に依存するプロンプトでは、bf16 の参照画像と比較して無彩色のように見えました。

これらの値ごとの数レベルは、各トークンが大きな平均値への小さな寄与に過ぎない大規模言語モデルにとっては問題ありません。しかし、拡散トランスフォーマーはより寛容ではありません。なぜなら、それは潜在空間に対して多数のステップで反復的にノイズ除去を行うためであり、各ステップの出力が次のステップの入力となるからです。あるステップでの速度がわずかに間違っていると、次のステップでは誤った潜在値となり、さらにその次のステップでは誤った入力としてフィードバックされます。このエラーは雪だるま式に膨らんでいきます。FP4 では、この累積効果が彩度の低下として現れました：モデルは正しい構成を生成していましたが、bf16 モデルと比較して色が平坦で鮮やかさに欠ける結果となりました。

ポストプロセッシングのトリック（輝度補正、色の再バランス調整など）は本質的な解決にはなりませんでした。なぜなら、エラーはデコード前の潜在値に組み込まれてしまっており、単なる表示空間の問題ではないからです。私たちはモデルの外側ではなく、内部でこれを修正する必要がありました。

何が機能しなかったか

最終的なレシピにたどり着く前に、いくつかの試行を行いました。これらは、何が重要で何がそうでないかを教えてくれました：

出力に対するポストプロセッシング。最初の直感は、生成後に色を修正することでした。デコードされた画像に対して彩度補正、輝度調整、色の再バランス調整を試みました。しかし、どれも機能しませんでした。彩度の低下は表示の問題ではなく、潜在空間の軌道における問題です。修正すべきピクセルが得られる頃には、ノイズ除去器は何回かの反復でわずかに間違ったステップを踏んでおり、色情報はすでに失われています。最終画像からそれを元に戻すことはできません。

素朴な蒸留。最も単純な方法を試しました：学生モデルを FP4 に量子化し、同じ入力に対して bf16 の教師モデルを実行して、両者の予測間の MSE を最小化します。しかし、量子化プロセスを通じて勾配は流れず、学生モデルは量子化された状態で凍結され、単に教師モデルと比較されるだけでした。損失は低下しましたが、画像の品質は向上しませんでした。これにより、我々は苦い経験から、拡散モデルの学習損失が画像の品質を予測できないことを学びました。損失値が数千ステップにわたって低下しても、出力はほとんど変化しないのです。

これらの結果から、2 つの重要な教訓を得ました：量子化誤差はトレーニング終了後ではなく、トレーニング中に解決する必要があること、そして最終的な評価は数値的ではなく視覚的に行うべきことです。

量子化意識型蒸留 (Quantization-aware distillation)

我々がこれに苦闘していた頃、LLM や VLM 向けに「量子化意識型蒸留 (QAD)」という手法が論文として発表されました(paper)。後から見れば単純なアイデアです：モデルを単に量子化して運命に任せるのではなく、完全精度の教師モデルから量子化された学生モデルへと知識を移すのです。学生モデルは、自身の量子化ノイズという制約下で、教師モデルが予測したであろう内容を予測することを学びます。教師モデルは凍結された bf16 状態のまま、学生モデルは量子化され、学生モデルの出力が教師モデルの出力に一致するようにトレーニングを行います。

その論文では、LLM に関するものであったため、ロジットに対して KL 発散を用いました。私たちはトークンを予測しているのではなく、潜在空間における速度場を予測しています。そのため損失関数を交換しました：学生モデルの速度予測と教師モデルの速度予測との間の MSE です。精神は同じですが、目的は異なります。

以前試した単純な知識蒸留との決定的な違いは、QAD では学生モデルが単に FP4 で実行されて教師と比較されるだけではない点です。学生モデルは自身の量子化を通じて学習します。勾配は直線推定器（straight-through estimator）を介して FP4 の丸め処理を通って逆伝播するため、重みは自身が生成する量子化ノイズに対して積極的に適応します。これが本質的なポイントです — 重みは、4 ビットに丸められても出力が損なわれないパラメータ空間の領域へと移動します。単純な知識蒸留では、固定された量子化状態の学生モデルを教師と比較して期待するだけですが、QAD では量子化自体を最適化の一部として組み込みます。

順伝播では、学生モデルのbf16重みと活性化値は行列積のためにFP4に量子化され、その後再び非量子化されます。この往復処理によって出力には実際の量子化ノイズが含まれます。凍結されたbf16教師モデルが目標速度を生成します。損失関数は、学生モデルのノイズ付き予測と教師モデルのクリーンな目標との間のMSE（平均二乗誤差）です。

逆伝播では、FP4の丸め処理は自動微分グラフから切り離されます。丸め関数の勾配はほぼ全域でゼロとなるため直接微分できず、Straight-Through Estimator (STE) がこれを迂回して、勾配を完全精度のまま流すようにします。順伝播では量子化ノイズが伝搬しますが、逆伝播ではこれが無視されます。学習が進むにつれて重みは、4ビットへの丸めでも出力が劣化しないパラメータ空間の領域に収束していきます。

推論時にはこれらの処理は一切不要です。学習済みの重みをエクスポートし、本番環境のパイプラインに読み込み、実際のカーネルを用いて静的なFP4量子化を適用して画像を生成します。学生モデルは、ガイダンス係数 = 1 の条件付き分岐としてのみ動作します。教師モデルもCFG（Classifier-Free Guidance）もSTEも使用しません。

imageCFGを一つの順伝播に統合

Classifier-free guidance (CFG) は、拡散モデルがプロンプトに向かって誘導する仕組みです。推論時には、各ステップごとにトランスフォーマーを 2 回実行します：1 回はプロンプトに条件付けされたもの（条件付き）、もう 1 回はテキストなしのもの（非条件付き）。これは、すべてのノイズ除去ステップで 2 倍のコストがかかります。私たちはこれを蒸留に統合して、無料で実現できることに気づきました。

トレーニング中、凍結された bf16 の教師モデルは両方のブランチを実行し、ガイド付きの CFG ベロシティを生成します。学生モデルは条件付きブランチのみであり、そのガイド付きベロシティを直接予測するように訓練されます。推論時には、 trained conditional branch だけを guidance = 1 で実行し、2 回のフォワードパスではなく 1 回で済ませます。つまり、FP4 の高速化に加えて、CFG をサービングパスから完全に削除することでトランスフォーマーのコストを半減させました。

これは QAD 論文を超える部分です。その論文は量子化精度の回復にとどまりますが、私たちは量子化精度を回復し、さらに CFG ブランチを蒸留で除去します。その結果、最終的な推論グラフはより小さく（fp4）、かつ短くなります（1 つのブランチ）。

QAD が機能した理由

QAD を実際に視覚品質を回復させるためには、コアとなるレシピBeyond いくつかの判断が必要でした。特に重要だったのは以下の点です：

損失の組み合わせ。リアルデータの品質を活用するために、flow-matching loss と蒸留損失を組み合わせて使用しています。純粋な蒸留だけでは、望む完全な視覚品質を回復するには不十分でした。

ガイダンススケジュール。教師ターゲットに対して固定された単一のガイダンスは使用しませんでした。ノイジーな中間タイムステップでは高いガイダンスを、クリーンな終端付近では低いガイダンスを使用しました。これは、モデルが初期段階で強い誘導を受け、収束するにつれてより穏やかな誘導を受けることで利益を得るためです。

QAD による品質回復と CFG の単一ブランチへの統合により、コスト乗数のうち 2 つは処理されました。残りの 1 つが最大の課題となります。

imageQAD は量子化誤差を回復し、同様の出力を提供します

タイムステップの蒸留

これで各フォワードパスは可能な限り安価になりました：FP4 での 4 ビット行列乗算、エピローグの融合によるノルムと活性化値の往復コストの排除、そして 2 つのガイダンスブランチに代わる単一のブランチです。まだ手をつけていないのは、フォワードパスの数、すなわちディノイジングステップ数であり、これが最大のレバー（調整要素）となります。

拡散トランスフォーマーは、多くの反復を通じて潜在をディノイズすることで品質を獲得し、各反復が完全なフォワードパスです。もしモデルに少数のステップで同じ作業を行わせることを教えられれば、レイテンシとスループットにおける最大の要因を縮小できます。これがタイムステップの蒸留です。また、これはパイプラインの中で最も多くの要素を持ち、文献も最も多く、画像品質を静かに損なう方法が最も多い部分でもあります。

基盤：分布マッチング蒸留

自然な最初のアイデアは回帰です：教師モデルを実行し、生徒モデルを実行して、両者の間の MSE（平均二乗誤差）を最小化します。上記の QAD 作業がまさにこれを行っており、生徒と教師が同じ経路をたどる場合には機能します。しかし、数ステップで完了する生徒は、多数のステップを要する教師と同じ経路をたどりません。生徒に中間的な潜在変数を一致させるよう求めることはできません。なぜなら、生徒はそれらのほとんどをスキップしているからです。回帰という手法は、生徒が自由に離れるべき軌道に生徒を固定してしまいます。

分布マッチング蒸留（DMD）は目標を再定義します：サンプルごとに教師と一致させるのではなく、分布同士を一致させます。特定の経路上で生徒が何を行うかは問題ではなく、生徒が生み出す画像の雲が、教師が生み出す画像の雲のように見えるかを重視します。

このメカニズムでは 2 つのスコアネットワークを保持します：1 つはリアルなスコア（凍結された教師）であり、もう 1 つは偽物のスコア（生徒の現在の分布を追跡する第 2 のネットワーク）です。生徒を更新するために、DMD は生徒サンプルを取得し、それを再ノイズ化して、両方のスコアネットワークに進行方向を問います。2 つのスコアの差は、生徒の分布を教師の分布へと移動させる勾配方向となります。このスコアの差分は、定数倍を除けば、そのノイズレベルにおける 2 つの分布間の KL 発散（Kullback-Leibler divergence）の勾配です。ペアされたターゲットも、特定の潜在変数に対する MSE も不要です。必要なのは「これらの 2 つの分布を一致させる」という指示だけです。

imageDMD2: 実際に出荷されているバージョン

DMD2 は、バニラの DMD を整理し、現代の少ステップ画像モデルが派生するレシピへと昇華させました。具体的には、回帰項を削除してデータフリーな手法とし、二時間スケール更新を採用して学生モデルの変化に伴い偽スコアを正確に保ち、GAN 損失を二次的なシャープニング項として追加し、さらに多ステップサンプリングをサポートして学生モデルが単一のステップ数に縛られないようにしています。

それぞれ異なるレバーを追う一連のファミリー

DMD2 は一連の後継作の始まりとなりました。最適な選択は、何を最適化するかによって異なります。DMDR、Flash-DMD、AdvDMD、GNDM は強化学習と報酬信号を蒸留に組み込みます。DP-DMD は多様性を維持します。AMD は安定性を向上させます。