SmoothCache: A Universal Inference Acceleration Technique for Diffusion Transformers

拡散モデル [30, 8] は、画像、音声、動画など様々な領域で高品質なサンプルを生成できる強力な生成モデルのクラスとして台頭してきた。当初、このノイズ除去ステップにはU-Net アーキテクチャが好まれていた。これは、潜在空間における大域的および局所的な情報の両方を捉える強力な能力を持ち、データの微細な詳細を再構築する上で重要であったためである [25, 26]。

しかしながら、U-Netの固有の限界、特に高次元データや長いシーケンスに効果的にスケールすることの困難さから、代替アーキテクチャの探索が行われた。Diffusion Transformer (DiT) [24] は、トランスフォーマーアーキテクチャのスケーラビリティを活用するためのU-Netに対する解決策として提案された [32]。長距離依存関係を扱う能力と固有のスケーラビリティで知られるトランスフォーマーは、様々なシーケンスベースのタスクで非常に効果的であることが証明されている。DiTはこれらの利点を拡散フレームワークに適応させることに成功し、画像生成だけでなく、音声 [5, 12, 14]、動画 [36, 11, 21]、3D生成 [23] など多様な領域にも拡張して強力な性能を示した。

その成功にもかかわらず、DiTモデルは多くのトランスフォーマーベースのアーキテクチャと同様に、推論時の計算効率に課題を抱えている。トランスフォーマー内のアテンションとフィードフォワード層の中核的な操作は二次の時間複雑性を持つ。拡散モデルの文脈では、これらの操作が複数のノイズ除去ステップにわたって繰り返し実行されるため、この二次依存性は大きな計算オーバーヘッドにつながる。特に高解像度の出力や長いシーケンスを生成する際に顕著である。この高い推論コストは、計算リソースが限られているか、厳しい遅延要件がある実世界のアプリケーションにこれらのモデルを展開する上で大きな課題となっている。

SmoothCacheは、実装が簡単で訓練不要でありながら、特定の拡散プロセスに基づいてキャッシングの決定を行い、均一なスケジュールを上回る性能を発揮するキャッシングアプローチを導入することで、これらの制限に対処している。このモデルに依存しないフレームワークは、1回のキャリブレーション推論パスと単一のハイパーパラメータ $\alpha$ だけで、様々なモデル、サンプリング、ソルバー構成に適応し、最適な性能を達成する。

拡散プロセスは、データ$\mathbf{x}_{0}\sim q(\mathbf{x}_{0})$に$T$ステップにわたってガウシアンノイズを加えることで変換し、シーケンス$\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{T}$を生成する。これはマルコフ連鎖としてモデル化され、各ステップは以下に従う：

ここで、$\beta_{t}$はタイムステップ$t$におけるノイズの分散であり、$\mathbf{x}_{T}\sim\mathcal{N}(0,\mathbf{I})$は拡散ステップ中に加えられたノイズの累積効果によるものである。逆プロセスの目的は、ノイズの加わったサンプル$\mathbf{x}_{T}$から元のデータ$\mathbf{x}_{0}$を復元することである。これは、逆遷移をパラメータ化するニューラルネットワークによって学習される：

ここで、$\mu_{\theta}(\mathbf{x}_{t},t)$は予測された平均であり、$\Sigma_{\theta}(t)$は分散であり、両者ともモデルによって学習される。完全な生成プロセスは以下のように定義される：

生成モデルは、負の対数尤度の変分下界を最小化することによって訓練され、各タイムステップで$\mathbf{x}_{t}$のノイズを除去し、最終的に学習された分布からサンプルを生成することを学習する。従来、$p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})$はU-Netスタイルのモデルアーキテクチャで近似されてきたが、最近の研究では、Diffusion Transformer (DiT)アーキテクチャを使用し始めており、特にビデオ生成のようなより複雑なタスクにおいて、より良くスケールすることが示されている。DiTアーキテクチャは、従来のTransformerにおけるSelf-attention、Cross-attention、Feed-forwardレイヤーを含む繰り返しブロックで構成されており、これらは通常、モデルの訓練と推論の両方において計算上のボトルネックとなっている。

拡散モデルの重要な特性の一つで、キャッシング技術の開発を促進したのは、隣接するタイムステップにおけるレイヤー出力間の高いコサイン類似度である。この類似性のパターンは、画像、ビデオ、音声拡散など、さまざまなモダリティにわたる様々な生成モデルとソルバーで観察されている。この高い類似性は、拡散プロセス内に効率を改善するために活用できる計算の冗長性が存在することを示唆している。

SmoothCacheの目的は、訓練不要でモデルに依存しない層出力のキャッシュと再利用の戦略を提供し、最小限のエラーを導入することである。しかし、我々は、異なるモデルアーキテクチャ、ソルバー、モダリティにわたって、タイムステップ間の層出力をキャッシュする単一の最適な静的スキームは存在しない可能性があると考えている。例えば、図2に示すように、連続するタイムステップの層出力間の平均表現誤差を調べると、DiT-XLラベル-画像モデルでは一般に後半のタイムステップで層がより大きな差異を示すのに対し、Open-Soraテキスト-動画モデルでは拡散の最初と最後のタイムステップで層がより敏感であることが観察される。この発見は、異なるモデルが誤差曲線に基づいて異なる利益を得られるように、一般化可能な技術で類似性原理を適用する必要性を強調している。

$L_{t}$をタイムステップ$t$における特定の層の出力とし、$L_{t-k}$を将来の拡散タイムステップ$t-k$における同じ層の出力とする。タイムステップ間の層類似性の観察により、$L_{t-k}\sim L_{t}$である。したがって、拡散プロセス中に$L_{t-k}$を計算する代わりに、以前に計算された$L_{t}$を使用して関数を近似できる。層が計算されるたびに、我々の手法は$L_{t}$をキャッシュに保存し、将来の層計算の代わりにアクセスして使用できるようにする。我々は、図LABEL:fig:smoothcacheに示すように、残差接続の前の出力における前述の計算のボトルネックにキャッシングを適用する。これには、DiT-XLモデルのSelf-attentionとFeed-forward層、StableAudioモデルのSelf-attention、Cross-attention、Feed-forward層、そしてOpenSoraモデルの空間ブロックと時間ブロックの両方におけるSelf-attention、Cross-attention、Feed-forward層が含まれる（図3参照）。図4は、これらのコンポーネントが生成中に発生するほぼすべての計算を占めていることを示している。また、計算の分布がモデルごとに異なることにも注目する。

以前にキャッシュされた出力を使用するかどうかを決定するために、以下の問題を定義する。$t$を現在のタイムステップ、$t+k$をキャッシュが以前に埋められたタイムステップ、$i_{j}$をタイプ$i$の$j$番目の層とする。ここで、$i\in S=\{attn,f\!f\!n,...\}$はモデルアーキテクチャに依存する。我々の手法は、計算された出力とキャッシュされた出力の間の損失$\mathcal{L}(L_{i_{j},t},L_{i_{j},t+k})$が、層依存のハイパーパラメータ$\alpha_{i_{j}}>0$によって制限される場合、キャッシングが効果的であるという主要な仮説に導かれる。

$\mathcal{L}(L_{i_{j},t},L_{i_{j},t+k})$を計算することは、$L_{i_{j},t}$を評価せずには不可能であり、これは層計算をスキップすることによる利点を失わせる。特定のDiTアーキテクチャに対して、我々は2つの異なるサンプルの層表現誤差の差が無視できるほど小さい閾値内にあることを顕著に観察した。これは図2に示す10個のキャリブレーションサンプルで生成されたプロットの95%信頼区間から分かる。この発見は、特定のモデル入力$\mathcal{L}(L_{i_{j},t},L_{i_{j},t+k})$の誤差曲線が、十分に大きなキャリブレーション入力セットの平均誤差曲線によって近似できることを示唆している。言い換えれば、$\tilde{L}_{i_{j},t}$が層$L_{i_{j},t}$のキャリブレーション出力を表す場合、$\mathcal{L}(L_{i_{j},t},L_{i_{j},t+k})\sim\mathcal{L}(\tilde{L}_{i_{j},t},\tilde{L}_{i_{j},t+k})$となる。

すべての$\alpha_{i_{j}}$を見つけるためのハイパーパラメータ探索は、層の数に基づいて指数関数的な探索空間を持ち、非常にコストがかかる。キャッシング問題を簡略化するために、我々はすべての層のキャッシングを導くための単一のハイパーパラメータ$\alpha>0$を定義する。$\mathcal{L}$をタイプ$i$のすべての$N$層の平均L1相対誤差として定義する。これは、すべてのタイプ$i$と位置的深さ$j$のネットワークにおける層間の真の表現誤差を比較するために選択される。さらに、特定の層をキャッシュすることで、ネットワーク内の同じタイプの将来の層にエラーが導入される可能性があることを認識している。例えば、$L_{i_{j-1},t}$を$L_{i_{j-1},t+k}$で近似すると、ノイズが導入され、キャリブレーション誤差$\mathcal{L}(\tilde{L}_{i_{j},t},\tilde{L}_{i_{j},t+k})$が真の誤差$\mathcal{L}(L_{i_{j},t},L_{i_{j},t+k})$を正しく近似しなくなり、不適切なキャッシング決定につながる可能性がある。層のキャッシングのカスケード効果を軽減するために、タイプ$i$のすべての層のキャッシング決定をグループ化し、$L_{i,t+k}$内のすべての$j$層が$L_{i,t}$を近似するようにする。したがって、キャッシュされた出力$L_{i_{j},t+k}$は、以下の条件が満たされる場合に$L_{i_{j},t}$の計算の代わりに使用される：

キャッシュされた出力の使用は計算コストが低く、拡散推論プロセスの大幅な高速化を可能にする。式4は、キャッシュされる特定の拡散プロセスの特性について事前の仮定を行わないため、複数のアーキテクチャとモダリティにわたって適用できる。さらに、キャッシング決定はキャリブレーション誤差にのみ依存するため、モデルの実行時に変更されない。これにより、既存のグラフコンパイル最適化との互換性が確保される。

モデル、データセット、およびソルバー 我々の手法の有効性をDiTアーキテクチャ全体で実証し、モデルに依存しないという主張を証明するために、我々はSmoothCacheを複数の候補拡散モデルで評価し、さまざまなモダリティと異なる拡散ステップ数にわたって検証する。

テキストから画像生成 まず、オリジナルのDiT-XL-256x256ラベルから画像拡散モデルを選択し、公開されたモデルの重み²²2https://github.com/facebookresearch/DiTを使用する。DiTファミリーのモデルはImageNet1kデータセットで訓練されているため、評価用に50,000枚の256x256画像を生成する。我々はDDIMソルバーを用いてSmoothCacheをテストし、スケール1.5のクラシファイアフリーガイダンスを使用する。

テキストから動画生成 次に、OpenSORAテキストから動画モデルを選択する。我々は事前訓練されたOpenSORA v1.2モデルを使用し、VBenchプロンプトスイートからの946のプロンプトを用いて生成された2秒間の480p動画（アスペクト比9:16）に対してVBench評価プロトコルを使用してその性能を評価する。推論にはフラッシュアテンションとbfloat16データ型を使用する。我々はRectified Flow [16]ソルバーを用いてSmoothCacheをテストし、1000 Eulerステップにわたる30サンプリングステップで、デフォルトのCFGスケール7.0を使用する。

テキストから音声生成 最後に、Stable Audio Openテキストから音声モデルを選択する。我々はオリジナルの論文[5]で説明されている推論プロトコルに従う。Stable Audio Openの推論は、クラシファイアフリーガイダンススケール7.0で100ステップのDPM-Solver++ (3M) SDEを用いて実行する。

評価指標 効率性と品質のトレードオフを実証するために、各ドメインで様々な標準的な指標を使用する。我々は完全な拡散プロセスの乗算累積演算（MACs）と遅延を記録する。ベースラインのサンプリング遅延から加速率を導出する。生成品質を測定するために、対応するドメインの一般的な指標で候補モデルを評価する。 DiT-XL-256x256については、指定された設定で50,000枚の画像を生成し、FID、IS、およびsFIDを報告する。 OpenSORAについては、VBench[9]を使用してモデルのパフォーマンスを評価し、VBenchスイートのプロンプトに基づいて946の動画を生成する。最終的なスケーリングされたVBenchスコアを報告する。 最後に、Stable Audio Openについては、Stable Audio Open技術報告書[5]で説明されている評価プロトコルを同じ評価コード³³3https://github.com/Stability-AI/stable-audio-metricsを使用して厳密に適用し、AudioCaps [10]、歌唱プロンプトなしのMusicCaps [2]、および歌唱プロンプトなしのSong Describer [22]に対するCLAP、FD_OpenL3、およびKL_PaSST指標を報告する。すべての速度結果は単一のH100-80G GPUで測定され、50回の実行にわたって平均化される。

実装の詳細 前述のように、我々は残差接続の前のレイヤーにキャッシングを適用する。DiT-XL-256x256モデルでは、これにはDiTブロック内のSelf-attentionとFeed-forwardモジュールが含まれる。OpenSora V1.2では、時間的Self-attention、Cross-attention、Feed-forwardモジュール、およびそれらに相当する空間的変種（合計6種類のモジュール）をキャッシュする。 Stable Audio Openでは、Self-attention、Cross-attention、Feed-forwardモジュールをキャッシュする。各アーキテクチャ、ソルバー、およびプリセットの拡散ステップ数に対して、いくつかのランダムサンプルに対するレイヤー出力の差を計算することで表現誤差$L(F_{i_{j},t},F_{i_{j},t+k})$を取得する。 すべてのモダリティについて、キャリブレーションには10個のサンプルを使用する。アブレーション研究では、キャリブレーションサンプル数の選択はあまり重要ではないことが示されており、これは図2で観察できる。 また、キャリブレーション中に条件付き生成を行ったかどうかに基づいて、誤差曲線を生成するサンプルを選択する。これはOpenSoraとStable Audio Openでのみ行う。 DiT-XL-256x256とStable Audio Openでは$k\in\{1,2,3\}$を固定する。これは、4タイムステップ以上の差でレイヤー表現誤差が大きくなりすぎると判断したためである。OpenSORAでは、特にネットワークのクロスアテンションコンポーネントの誤差が比較的低いため、$k$は5まで上がる。 DiTについては、ヌルプロンプトを使用して無条件に生成されたサンプルでキャリブレーションを行う。OpenSoraについては、VidProM [33]からランダムにサンプリングされたプロンプトを用いて条件付きで生成された480p 2秒動画でキャリブレーションを行う。 Stable Audio Openについては、AudioCaps検証セットからランダムにサンプリングされたプロンプトを使用してキャリブレーションを行う。

我々は、DiT-XL-256x256、OpenSora、およびStable Audio Openに対するSmoothCacheの結果を、それぞれ表1、2、3に示す。加速と品質のトレードオフを公平に比較するため、DDIMやDPM++などのモデルで使用されるデフォルトのソルバーと、指定されたハイパーパラメータを持つSmoothCacheの実装を比較する。同一の加速率を持つ2つの構成において、SmoothCacheが様々な品質指標でより良い性能を示すことを発見した。これは、典型的な加速/品質のトレードオフに対する優れた性能を示している。すべての結果について、我々は5回の試行を実行し、各指標の平均値と標準偏差を報告する。

我々は、異なるステップサイズにおけるモデルの性能の変化を示すことで我々の調査のアブレーションを行い、DiT-XL-256x256モデルに対するSmoothCacheの堅牢性を示す。

キャッシング/サンプリングステップのパレートフロントの検証 我々は、静的キャッシングとの比較を通じて、我々の手法が異なるサンプリングステップにおいてより優れたパレートフロントを生成することを画像生成結果により示す。すべてのキャッシング技術はソルバーのサンプリングステップ数を変えると性能が変化することに注目する。我々は、SmoothCacheがサンプラーやモデルアーキテクチャに関して最小限の仮定しか設けていないにもかかわらず、少なくとも他のキャッシング戦略と同等か、多くの面で上回っていることを示す。表1は、サンプリングステップ数が少ない場合も多い場合も、SmoothCacheが異なるサンプリングステップ数における複数の推論速度/品質のポイントでFORAを上回っていることを示している。 我々は、SmoothCacheが複数のサンプリング設定において静的キャッシングを上回ることを観察した。L2CはDDIMサンプリングにおいてDiT-XLに対してSmoothCacheをわずかに上回るが、SmoothCacheはImageNet1kに対する高コストな訓練を必要とせず、異なるモデルアーキテクチャやサンプリング設定に一般化できる。

キャリブレーションサンプルサイズ 我々は、キャリブレーションサンプル数の選択が生成されるキャッシングスケジュールに悪影響を与えないことに注目する。経験的に、本稿で調査した3つのモデルすべてにおいて、同じ$\alpha$が与えられた場合に同じキャッシングスケジュールを確実に再生成するには、通常10サンプルで十分であることを観察した。しかし、図2に示されるように、異なる誤差曲線の信頼区間はモダリティによって異なることに注意する。