要約
本論文では、空間（モデル構造）と時間（タイムステップ）のトークン疎化によりDiffusion Transformer（DiT）の計算効率を向上させる新しいフレームワークSparseDiTを提案する。
空間次元では、SparseDiTは3段構成を採用：下層ではPoolingformerにより効率的なグローバル特徴抽出、中層ではSparse-dense generation token model（SDTM）によりグローバルとローカルの特徴をバランス、上層では密トークンにより高周波詳細を精製する。
時間次元では、SparseDiTはノイズ除去段階でのトークン数を動的に調整し、タイムステップの進行に伴いトークンを徐々に増加させ、効率を保ちつつ詳細をより良く捉える。この空間・時間戦略により計算効率が向上し、生成品質を維持する。
実験により、SparseDiTは画像生成、動画生成、文生図など複数の生成タスクで優れた性能を示し、例えば512×512解像度の画像生成タスクではFLOPsを55%削減、推論速度を175%向上させつつ、ほぼ同等の生成品質を維持する。
論文タイトル：
SparseDiT: Token Sparsification for Efficient Diffusion Transformer
論文リンク：
https://arxiv.org/pdf/2412.06028
コードリンク：
https://github.com/changsn/SparseDiT
既存の問題と課題の紹介
Diffusion Transformer（DiT）は生成性能に優れるものの、自注意力とサンプリングステップの高計算複雑度が実用的応用を制限している。
既存手法の多くはサンプリングプロセスを加速して複雑度を低減するが、DiT自体の構造効率問題を無視している。
U-Net比で、DiTのトークンレベル自注意力が追加計算負荷を生むため、DiT特化の革新的手法で計算効率と生成品質のバランスを実現する必要がある。
研究価値と手法紹介
研究価値
本研究はSparseDiTによりDiTモデルの計算効率問題を効果的に解決。動的トークン疎化戦略により計算複雑度を低減し、複数生成タスクで高品質生成性能を維持する。
さらに複数実験でFLOPs削減を示し、推論速度を大幅向上、大規模応用・デプロイに重要。本研究は高品質高効率拡散モデルに拡張可能な解決策を提供。
手法紹介
SparseDiTの核心革新は空間・時間次元のトークン疎化により拡散モデルの計算効率を向上させ生成品質を維持すること。
設計は2主部：空間トークン密度管理とタイムステップ別動的トークン管理戦略。
• 空間次元：3段構成
1. 下層：Poolingformer
下層Transformerでは、Poolingformer構造で伝統的自注意力に置き換えグローバル特徴を捉える。
実験で下層自注意力の複雑計算が追加情報を生まないことが判明、全グローバル平均プーリングで効率向上。Poolingformerはkey/value計算を除去、トークンに直接グローバル平均プーリングを適用し入力に統合、負荷低減。
上図実験ではfinetuningなしで注意層をグローバルプーリングに置き換えても画像に大きな影響なし、下層注意層の有効性を示す。
2. 中層：疎密トークンモジュール
中層は疎密トークンモジュール（SDTM）で表現をグローバル構造抽出とローカル詳細強化に分ける。
疎トークンはグローバル構造捕捉で計算コスト低減、密トークンは詳細強化と訓練安定。
SDTMは相互注意層で疎密トークン変換、疎Transformerが疎トークン処理、密Transformerが疎から回復した密トークン処理、情報保持と計算節約。
3. 上層：標準Transformer
上層では標準Transformer層を密トークンで使用、高周波詳細精製に集中し品質確保。
• 時間次元：Time-wise pruning rate
動的Time-wise pruning rateはノイズ除去進行に伴うトークン密度調整の鍵。具体的に：
• 初期段階：
初期ノイズ除去は低周波構造中心のため高剪定率で資源節約、疎トークン少なく低複雑度。
• 後期段階：
進行に伴い剪定率漸減で密度増加、高周波詳細正確捕捉、トークン需要増加。
この時空二重適応戦略で詳細保持しFLOPs削減・推論加速により効率大幅向上。
実験結果
論文では条件付き画像生成、条件付き動画生成、テキストto画像生成の3タスクでSparseDiTを検証、顕著効果。
1. 条件付き画像生成
256×256解像度でSparseDiT-XLはFLOPs43%減、推論速度87%向上、FID0.11増のみ、~25%トークンで同等性能。
512×512では高剪定で優性能効率trade-off、90%超トークン剪定でFLOPs55%減・速度175%向上、FID0.09増。
これらによりSparseDiTがDiT計算負担解決、品質維持し効率向上証明。
2. 条件付き動画生成
FaceForensics、SkyTimelapse、UCF101、Taichi-HDの4公開データセット、256×256解像度。
SparseDiTは動画時間次元に高剪定、FLOPs56%減で競争力FVD維持、動画生成有効性証明。
3. テキストto画像生成実験設定：
PixArt-αベースモデル、SAMデータセットで訓練・評価。
SparseDiTは原PixArt-αより低いFID、生成速度大幅向上、テキストto画像有効性示す。