2025-11-17T17:13:12.426721

Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model

Saghatchian, Moghadam, Nickabadi

Diffusion models have emerged as a promising approach for generating high-quality, high-dimensional images. Nevertheless, these models are hindered by their high computational cost and slow inference, partly due to the quadratic computational complexity of the self-attention mechanisms with respect to input size. Various approaches have been proposed to address this drawback. One such approach focuses on reducing the number of tokens fed into the self-attention, known as token merging (ToMe). In our method, which is called cached adaptive token merging(CA-ToMe), we calculate the similarity between tokens and then merge the r proportion of the most similar tokens. However, due to the repetitive patterns observed in adjacent steps and the variation in the frequency of similarities, we aim to enhance this approach by implementing an adaptive threshold for merging tokens and adding a caching mechanism that stores similar pairs across several adjacent steps. Empirical results demonstrate that our method operates as a training-free acceleration method, achieving a speedup factor of 1.24 in the denoising process while maintaining the same FID scores compared to existing approaches.

academic

キャッシュ適応型トークンマージング：拡散モデルにおける動的トークン削減と冗長計算の排除

基本情報

論文ID: 2501.00946
タイトル: Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model
著者: Omid Saghatchian, Atiyeh Gh. Moghadam, Ahmad Nickabadi (アミルカビル工科大学)
分類: cs.CV (コンピュータビジョン)
発表日: 2025年1月1日 (arXiv プレプリント)
論文リンク: https://arxiv.org/abs/2501.00946
コードリンク: https://github.com/omidiu/ca_tome

要約

拡散モデルは高品質で高次元の画像生成の有望な手法となっています。しかし、これらのモデルは高い計算コストと遅い推論速度に阻まれており、その一因は自己注意機構の入力サイズに対する二次計算複雑度です。本論文は、キャッシュ適応型トークンマージング(CA-ToMe)手法を提案し、トークン間の相似性を計算し、相似度がしきい値パラメータtを超えるトークンをマージすることでこの問題に対処します。隣接ステップで観察される反復パターンと相似性頻度の変化のため、本手法は適応的しきい値の実装とキャッシュ機構の追加によってトークンマージング手法を強化します。実験結果は、本手法が訓練不要な加速手法として、ノイズ除去プロセスで1.24倍の加速を実現しながら、既存手法と同じFIDスコアを維持することを示しています。

研究背景と動機

問題定義

拡散モデルは画像生成タスクで優れた性能を示していますが、深刻な計算効率の問題に直面しています：

高い計算コスト：自己注意機構の二次複雑度により推論速度が遅い
直列ノイズ除去プロセス：並列化不可能で、各ノイズ除去ステップで重複計算が必要
冗長計算：隣接する時間ステップ間に大量の重複計算が存在

問題の重要性

拡散モデルの高レイテンシは、高速推論が必要なアプリケーションでの使用を制限
高い計算コストはモデル展開を困難にし、特にリソース制限環境では顕著
既存の加速手法は再訓練が必要であるか、品質に大きな損失をもたらす

既存手法の限界

サンプリングステップ削減手法は通常、再訓練または複雑な数値ソルバーが必要
トークン剪定手法は情報損失と性能低下をもたらす
**従来のトークンマージング(ToMe)**は固定マージ率を使用し、異なる時間ステップと層の相似性分布変化に適応できない

研究動機

観察された2つの重要な現象に基づく：

異なる時間ステップと層におけるトークン相似性分布に大きな変化が存在
隣接する推論ステップ間のトークンペアは高度の相似性を示す

核心貢献

適応的しきい値機構の提案：トークン相似性分布に基づいて動的にマージング戦略を調整し、固定マージ率に置き換える
キャッシュ機構の設計：隣接ステップ間の相似性を利用し、トークンペアをキャッシュして重複計算を削減
訓練不要な加速の実現：事前訓練済みモデルに直接適用可能で、再訓練不要
より良い品質-速度トレードオフの達成：ベースラインToMe手法と比較して、画像品質を維持しながらより高速な推論を実現

手法の詳細

タスク定義

入力：拡散モデルのノイズ除去プロセスにおけるトークンシーケンス出力：適応的マージングとキャッシュ最適化による加速推論プロセス制約：生成画像品質の大幅な低下を回避

モデルアーキテクチャ

1. 適応的トークンマージング機構

従来のToMe手法は固定比率rでトークンをマージしますが、CA-ToMeは相似性しきい値tを導入します：

核心思想：

画像をsx × syのストライド領域に分割
各ストライド領域の左上角トークンをターゲットトークンとして選択
ソーストークンとターゲットトークン間のコサイン相似度を計算
しきい値tを超える相似度のトークンペアのみをマージ

利点分析：

シナリオA：ほとんどのトークン相似度が低い場合、固定マージ率は相似度の低いトークンの強制マージを引き起こし、情報損失をもたらします。適応的しきい値は高相似度トークンのみのマージを保証
シナリオB：ほとんどのトークンが高度に相似している場合（ノイズ除去初期など）、固定マージ率はマージ数を制限します。適応的しきい値はより多くのトークンのマージを許可し、効率を向上

2. キャッシュ機構の設計

ジャッカード距離分析に基づき、隣接ステップ間のトークンペアの高相似性を発見：

$JaccardDistance(A_n, A_{n+1}) = 1 - \frac{|A_n \cap A_{n+1}|}{|A_n \cup A_{n+1}|}$

ここでAnはn番目のステップのすべてのソース-ターゲットトークンペアのセットを表します。

実装戦略：

チェックポイント(checkpoints)を設定し、特定の時間ステップでのみ相似性行列を計算
非チェックポイントステップで以前計算されたトークンペアを再利用
相似性行列の重複計算オーバーヘッドを大幅に削減

技術的革新点

動的適応性：相似性分布に基づいて自動的にマージング戦略を調整し、固定パラメータの限界を回避
時間次元の最適化：時間ステップ間の冗長性を利用し、キャッシュにより計算量を削減
層レベルの選択的適用：計算集約的なU-Netトップ層（D1およびU1）に特に最適化を適用
再訓練不要：即座に使用可能な加速手法として、既存モデルに直接適用可能

実験設定

データセット

ImageNet-1kデータセット：512×512解像度の2000枚の画像を生成（クラスあたり2枚、合計1000クラス）
検証セット：5000枚のImageNet-1k検証画像を使用してFIDスコアを計算
プロンプトテンプレート："A high-quality photograph of a classname."

評価指標

FID (Fréchet Inception Distance)：生成画像品質を測定する主要指標
推論時間：2000枚の画像生成の平均時間
PSNR：ピーク信号対雑音比、ピクセルレベルの再構成品質を測定
SSIM：構造類似性指数、空間的および構造的一貫性を評価

比較手法

ベースライン：元のStable Diffusion v1.5
ToMe：従来のトークンマージング手法(r=50%)

実装詳細

ハードウェア：Tesla V100S GPU
拡散ステップ数：50ステップのPLMSサンプリング
CFGスケール：7.5
ストライドサイズ：固定2×2
適用層：U-NetのD1およびU1層のみに適用

実験結果

主要結果

モデル	FID	平均時間(s)	加速比
ベースライン	33.66	7.61±0.001	1.0×
ToMe	34.16	6.39±0.006	1.19×
CA-ToMe	34.05	6.09±0.001	1.24×

主要な発見：

CA-ToMeは最速の推論速度(6.09s)を実現
FIDスコア(34.05)はToMe(34.16)を上回り、ベースライン(33.66)に近い
速度と品質の最適なバランスを達成

しきい値パラメータ分析

しきい値t	FID	平均時間(s)	PSNR	SSIM
0.4	35.28	6.07±0.007	27.90	0.191
0.5	35.46	6.07±0.004	27.909	0.208
0.6	35.56	6.10±0.005	27.908	0.218
0.7	34.30	6.23±0.002	27.910	0.234
0.8	33.80	6.58±0.004	27.904	0.239
0.9	33.42	6.92±0.003	27.907	0.238

観察結果：

しきい値0.4～0.6範囲内の変化は小さい。ほとんどのトークン相似度≥0.6のため
しきい値0.7は品質-速度の最適なバランスを提供
より高いしきい値は品質を向上させるが速度を低下

キャッシュ構成の比較

構成	チェックポイント設定	時間(s)	FID
CONFIG 1	0,1,2,3,5,10,15,25,35	6.18±0.02	36.14
CONFIG 2	0,10,11,12,15,20,25,30,35,45	6.13±0.001	34.33
CONFIG 3	0,8,11,13,20,25,30,35,45,46,47,48,49	6.09±0.001	34.05

CONFIG 3が最高の性能を発揮し、ジャッカード距離分析と一致し、第8、11、13ステップおよび最終ステップでより多くのチェックポイントを設定。

アブレーション実験

異なるコンポーネントの貢献を比較することで：

適応的しきい値のみ：固定マージ率と比較して画像品質を向上
キャッシュ機構のみ：計算時間を大幅に削減
完全なCA-ToMe：両技術の組み合わせで最高性能を実現

結論と考察

主要な結論

適応的しきい値は固定マージ率の限界を効果的に解決し、相似性分布に基づいてマージング戦略を動的に調整
キャッシュ機構は時間ステップ間の冗長性を利用し、重複計算を大幅に削減
CA-ToMe手法は1.24倍の加速を実現しながら、画像品質を維持またはわずかに向上
訓練不要の特性により、手法は優れた実用性と拡張性を備える

限界

しきい値パラメータの調整：異なるモデルとタスクに対して最適しきい値の調整が必要
適用範囲の制限：主にU-Netアーキテクチャの拡散モデルを対象
キャッシュオーバーヘッド：キャッシュされたトークンペア情報を保存するための追加メモリが必要
層レベルの制限：トップ層のみに適用、他の層の最適化機会を見落とす可能性

今後の方向

自動しきい値学習：最適しきい値を自動決定する手法の開発
他のアーキテクチャへの拡張：DiTなどの新型拡散モデルアーキテクチャへの適応
より精密なキャッシュ戦略：コンテンツ適応型キャッシュ機構
ハードウェア最適化：特定ハードウェア向けの最適化実装

深度評価

利点

革新性が高い：適応的思想をトークンマージングに導入し、キャッシュ機構と組み合わせて完全なソリューションを形成
実用価値が高い：訓練不要で即座に使用可能な特性により、展開が容易
実験が充分：包括的なアブレーション実験とパラメータ分析が手法の有効性を支持
理論基礎が堅実：ジャッカード距離に基づく相似性分析がキャッシュ機構に理論的根拠を提供

不足

理論分析が不十分：適応的しきい値選択に対する理論的指導が欠如
実験範囲が限定的：ImageNetのみで検証、他のデータセットとタスクの評価が不足
比較手法が少ない：主にToMeとの比較で、他の加速手法との比較が不足
品質評価が単一：主にFID指標に依存、人工評価と他の品質指標が不足

影響力

学術的貢献：拡散モデル加速に新しい思想と手法を提供
実用価値：既存の拡散モデルに直接適用可能で、広範な応用前景を持つ
再現性：完全なコード実装を提供し、再現と拡張を容易に
啓発性：適応的およびキャッシュの思想は関連研究をさらに啓発可能

適用シーン

リソース制限環境：モバイルデバイス、エッジコンピューティングなどのシーン
リアルタイムアプリケーション：高速画像生成が必要なインタラクティブアプリケーション
大規模展開：サーバー計算コストとレイテンシを低減
研究プロトタイプ：他の加速技術の基礎コンポーネントを提供

参考文献

本論文は54篇の関連文献を引用し、主に以下を含む：

拡散モデルの基礎理論1,2,3
画像生成アプリケーション4,5,18,19,20
加速技術24,25,26,27,28
トークン処理手法21,40,41,43,44
キャッシュ技術24,37,38,39

総合評価：これは拡散モデル加速分野における実用価値を持つ研究です。適応的しきい値とキャッシュ機構の巧妙な組み合わせにより、画像品質を維持しながら顕著な速度向上を実現しています。理論分析と実験範囲にはまだ改善の余地がありますが、訓練不要の特性と良好な実験結果により、高い実用価値と影響力を備えています。