SALAD: 動画生成 Diffusion Transformer 向けの高疎密アテンション実現のための効率的な線形アテンション微調整

本論文「SALAD: Achieve High-Sparsity Attention via Efficient Linear Attention Tuning for Video Diffusion Transformer」は、動画生成における Diffusion Transformer (DiT) の推論コストを削減しつつ、生成品質を維持するための新しいアプローチを提案しています。以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 背景と問題定義

• 課題: 動画生成モデル（Diffusion Transformer）は、長いシーケンス長（高解像度・長時間）を扱う際、フルアテンションの二次的な計算複雑度 ($O(N^2)$) により、推論時の遅延とメモリ消費が甚大になります。
• 既存手法の限界:
  • 学習不要の疎密アテンション (Training-free): 計算コストを削減できますが、達成できる疎密度は限定的（例：30k トークンで 40-60%）であり、加速効果も限定的です。
  • 学習ベースの疎密アテンション (Training-based): 高い疎密度（80-95%）を達成できますが、大規模なデータセット（例：Koala-36M）と膨大な計算リソース（例：182 GPU 時間以上）を必要とし、実用性が低いです。
  • LoRA による微調整: 既存の疎密モデルを LoRA (Low-Rank Adaptation) で微調整しても、極端な高疎密条件下では、長距離依存性の欠如によりアーティファクト（テキストと映像の不一致、物体の重複など）が発生し、フルアテンションベースラインに匹敵する品質を回復させることが困難です。

2. 提案手法：SALAD

著者らは、SALAD (High-Sparsity Attention paralleling with Linear Attention for Diffusion Transformer) を提案しました。これは、疎密アテンションと並列に軽量な線形アテンション枝を追加し、効率的に高疎密化を実現するアーキテクチャです。

2.1 基本的なアーキテクチャ

• 並列構造: 疎密アテンション（Sparse Attention）の枝と、線形アテンション（Linear Attention）の枝を並列に配置します。
• 共有パラメータ: 両方の枝は、事前学習済みの重みから初期化された同じ Query ($Q$), Key ($K$), Value ($V$) 投影行列を共有します。線形枝は追加パラメータを約 4.99% しか増やしません。
• 線形アテンションの役割: 疎密アテンションが見過ごす「グローバルなトークン間相互作用」を補完するために機能します。線形アテンションは $O(N)$ の計算複雑度を持つため、推論時のオーバーヘッドは最小限に抑えられます。

2.2 核心技術：マルチレベル静的・動的スケーリング戦略

単純に両枝の出力を足し合わせるだけでは、線形アテンションの低ランク性が疎密アテンションの高ランク構造を抑制し、性能が低下することが発見されました。これを解決するため、以下の 2 段階の制御戦略を導入しています。

1. レイヤごとの静的スケーリング (Layer-wise Static Scaling):
   • 各 Transformer レイヤーの線形アテンション出力後に、ゼロ初期化された投影層（Projection Layer）を挿入します。
   • これにより、微調整の初期段階で線形枝の寄与を小さく保ち、安定した学習を可能にします。
2. 入力・タイムステップ認識型の動的スケーリング (Input-Timestep-Aware Dynamic Scaling):
   • 現在の隠れ状態とデノイジングのタイムステップ $t$ を入力として受け取り、線形枝の寄与を動的に調整するスケーリング係数 $s$ を生成します。
   • シグモイド関数を用いて $s \in (0, 1)$ に制限し、線形枝が疎密枝を圧倒しないように制御します。
   • これにより、入力や生成段階（タイムステップ）に応じて、必要な情報量だけ線形枝を適応的に補完できます。

2.3 推論時の枝のドロップ (Post-tuning Branch Dropping)

• 動的スケーリング係数が非常に小さい場合、その線形枝はほとんど寄与していないと判断し、推論時にその枝をスキップ（ドロップ）します。
• これにより、さらに計算コストを削減しつつ、品質を維持または向上させることが可能です。

3. 主要な貢献

1. 高効率な高疎密アテンションアーキテクチャ: 疎密アテンションの並列に軽量な線形枝を追加し、極端な高疎密度（90%）を維持しながら生成品質を回復させる新しい設計を提案しました。
2. バランス制御の重要性の解明と実装: 線形枝と疎密枝の「ランクのアンバランス」を分析し、それを制御する「マルチレベル静的・動的スケーリング戦略」を開発しました。これにより、線形枝が補完的な役割を果たすように最適化されています。
3. 極めて効率的な微調整: 大規模データセットや長時間の学習を必要とせず、2,000 枚の動画サンプル、1,600 ステップ未満の学習、30 GPU 時間以内でモデルを微調整可能にしました。

4. 実験結果

• 評価モデル: Wan2.1-1.3B (480p, 30k トークン; 720p, 100k トークン) および Wan2.1-14B (480p) に対して評価を行いました。
• 性能:
  • 疎密度: 90% の高疎密度を達成しました。
  • 速度向上: 異なるモデルとシーケンス長において、1.52 倍〜2.03 倍の推論速度向上を実現しました。
  • 品質: VBench スコア（被写体の一貫性、背景の一貫性、画質、テキスト一貫性）において、フルアテンションベースラインと同等、あるいはそれ以上の品質を維持しました。
  • 比較: 既存の学習不要手法や LoRA 微調整のみでは、90% 疎密条件下で品質が劣化しますが、SALAD はそれを克服しました。
• パラメータ効率: 学習可能パラメータは LoRA のみ（r=256）の場合よりも少なく、追加パラメータは約 5% 程度です。

5. 意義と結論

SALAD は、動画生成モデルにおける「計算効率」と「生成品質」のトレードオフを劇的に改善する画期的な手法です。

• 実用性: 大規模な事前学習や膨大なデータ収集なしに、既存のモデルを効率的に最適化できるため、実社会での導入障壁が大幅に下がります。
• 理論的洞察: 疎密アテンションと線形アテンションの組み合わせにおいて、単なる出力の加算ではなく、**「ランクのバランス」と「動的な重み付け」**が品質回復の鍵であることを示しました。
• 将来展望: 高解像度・長時間の動画生成を現実的なリソースで可能にする基盤技術として、今後の動画生成モデルの標準的な最適化手法となる可能性があります。

要約すれば、SALAD は「疎密アテンションの計算効率」と「線形アテンションのグローバル情報補完能力」を、高度に制御された微調整戦略で融合させることで、高品質かつ高速な動画生成を実現する手法です。