2025-11-23T17:13:17.428108

Token Pruning for Caching Better: 9 Times Acceleration on Stable Diffusion for Free

Zhang, Xiao, Tang et al.

Stable Diffusion has achieved remarkable success in the field of text-to-image generation, with its powerful generative capabilities and diverse generation results making a lasting impact. However, its iterative denoising introduces high computational costs and slows generation speed, limiting broader adoption. The community has made numerous efforts to reduce this computational burden, with methods like feature caching attracting attention due to their effectiveness and simplicity. Nonetheless, simply reusing features computed at previous timesteps causes the features across adjacent timesteps to become similar, reducing the dynamics of features over time and ultimately compromising the quality of generated images. In this paper, we introduce a dynamics-aware token pruning (DaTo) approach that addresses the limitations of feature caching. DaTo selectively prunes tokens with lower dynamics, allowing only high-dynamic tokens to participate in self-attention layers, thereby extending feature dynamics across timesteps. DaTo combines feature caching with token pruning in a training-free manner, achieving both temporal and token-wise information reuse. Applied to Stable Diffusion on the ImageNet, our approach delivered a 9$\times$ speedup while reducing FID by 0.33, indicating enhanced image quality. On the COCO-30k, we observed a 7$\times$ acceleration coupled with a notable FID reduction of 2.17.

academic

トークン剪枝によるキャッシング改善: Stable Diffusionで9倍の加速を無料で実現

基本情報

論文ID: 2501.00375
タイトル: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
著者: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
分類: cs.CV (コンピュータビジョン), cs.LG (機械学習)
発表日: 2024年12月31日
論文リンク: https://arxiv.org/abs/2501.00375
コードリンク: github.com/EvelynZhang-epiclab/DaTo

要約

Stable Diffusionはテキスト画像生成分野で顕著な成功を収めていますが、その反復的なノイズ除去メカニズムは高い計算コストと遅い生成速度をもたらします。特徴キャッシング等の方法はその有効性と単純性から注目を集めていますが、先前のタイムステップで計算された特徴を単純に再利用すると、隣接するタイムステップ間の特徴が類似するようになり、時間経過に伴う特徴の動的性が低下し、最終的に生成画像の品質に悪影響を及ぼします。本論文は、特徴キャッシングの限界に対処するため、動的認識トークン剪枝(DaTo)手法を提案しています。DaToは動的性の低いトークンを選択的に剪枝し、高動的性トークンのみが自己注意層に参加することを許可することで、タイムステップ間の特徴動的性を拡張します。ImageNetでStable Diffusionに適用した場合、この手法は9倍の加速を実現し、同時にFIDが0.33低下しました。COCO-30kでは7倍の加速が観察され、FIDが大幅に2.17低下しました。

研究背景と動機

問題背景

拡散モデルは生成モデリング分野で大きな進展を遂げており、テキスト画像生成やビデオ生成などのタスクに広く応用されています。しかし、拡散モデルの反復的なノイズ除去メカニズムは膨大な計算コストと遅い生成速度をもたらし、より広範な応用を制限しています。

既存手法の限界

拡散モデルを加速するための現在の手法は主に以下を含みます:

サンプリングステップ数の削減: DDIMなどの高速サンプラー
各ステップの計算コスト削減: 知識蒸留、構造剪枝、量子化、トークン剪枝、特徴キャッシング

このうち、特徴キャッシングはその有効性と単純性から広く支持されており、先前のタイムステップで計算された特徴を保存し、後続のタイムステップで再利用します。しかし、特徴の再利用は異なるタイムステップの特徴が類似した値を持つことを強制し、タイムステップに沿った特徴の動的性を低下させ、元の拡散プロセスを損傷させ、生成品質を低下させます。

研究動機

論文は実験的観察を通じて、元のStable Diffusionと比較して、特徴キャッシングを使用するモデルでは隣接するタイムステップ間の特徴差が大幅に減少することを発見しました。これは重要な問題を引き起こします: 特徴キャッシングを実行しながら、正しい特徴動的性を維持することは可能でしょうか?

核心的貢献

動的認識トークン剪枝(DaTo)手法の提案: 異なるタイムステップ間で特徴キャッシングによって動的性が低下したトークンを剪枝し、高動的性トークンでそれらを復元することにより、特徴キャッシングによる生成品質低下を回避します。
進化探索戦略の設計: 進化的手法を通じて最適な特徴キャッシングとトークン剪枝戦略を探索し、DaToの可能性を十分に引き出す方法を提案します。
顕著なパフォーマンス向上の実現: Stable DiffusionおよびSDXLでの広範な実験により、訓練とデータなしで、Stable Diffusionで最大9倍の加速と生成品質の損失なしを達成できることが示されています。

方法の詳細

タスク定義

本論文のタスクは、画像生成品質を維持しながら、Stable Diffusionモデルの推論プロセスを大幅に加速することです。入力はテキストプロンプト、出力は対応する高品質画像であり、制約条件はモデルの再訓練が不要なことです。

モデルアーキテクチャ

1. 動的認識トークン剪枝(DaTo)

基本トークン選択:

時間ノイズ差分スコア: t番目のタイムステップについて、前の2つの隣接するタイムステップの出力の絶対差を計算します:
```
DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
```
パッチベースのトークン選択: 画像をs×sの非重複パッチに分割し、各パッチ内で最高DiffScoreを持つトークンを基本トークンとして選択します。

CFG整列: 分類器自由ガイダンス(CFG)を処理するため、条件付き生成の基本トークン位置を無条件生成にコピーします:

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

剪枝トークン選択: コサイン類似度に基づいて基本トークンと最も類似したK個のトークンを剪枝用に選択します:

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

剪枝トークン復元: 剪枝されたトークンを、最も類似した基本トークンを直接コピーして復元します。

2. タイムステップ認識特徴キャッシング

探索空間の剪枝:

キャッシング深度dは{0, 1, 1/2}に制限
剪枝比率rは{0.3, 0.4, 0.5, 0.6, 0.7}に制限

進化探索アルゴリズム: NSGA-II多目的最適化アルゴリズムを使用し、最適化目標は以下を含みます:

推論レイテンシ
生成品質(FID)

探索プロセスには選択、交叉、変異などの標準進化操作が含まれ、最終的に最適なステップ認識戦略F(t)が得られます。

技術的革新点

動的性復元メカニズム: 低動的性トークンを選択的に剪枝し、高動的性トークンで復元することにより、特徴キャッシングによって破壊された特徴動的性分布を成功裏に復元します。
統一されたキャッシング剪枝フレームワーク: 特徴キャッシングとトークン剪枝を訓練不要なフレームワークに結合し、時間およびトークンレベルの情報再利用を実現します。
適応的戦略探索: 異なるタイムステップの異なる冗長特性に対して、最適なキャッシング深度と剪枝比率を自動的に探索する方法を提案します。