技術要約:効率的なビデオ生成のためのトークン・ラジアス・アテンション(Token Radius Attention)
1. 問題提起
ビデオ拡散トランスフォーマー(VDiT)は、複雑な空間的および時間的依存関係を捉えるための高密度な3D自己注意(self-attention)に依存しており、高忠実度なビデオ生成における支配的なアーキテクチャとなっています。しかし、高密度なアテンションは、ビデオトークン数 N に対して二次的な計算コスト (O(N2)) を招き、スケーリングとデプロイメントにおける重大な推論のボトルネックを生み出します。
このコストを軽減するための既存のアプローチは、主に**ヘッドレベル(head-level)およびブロックレベル(block-level)**の手法に分類される疎なアテンション・メカニズムを採用しています。
- ヘッドレベルの手法は、個々のアテンション・ヘッドに対して構造化された空間的または時間的なパターンを割り当てます。
- ブロックレベルの手法は、重要なクエリ–キーのブロック・ペアを選択します。
核心的な限界: 両方のパラダイムは通常、ヘッドまたはブロック内のすべてのクエリ・トークンに対して、単一の計算予算(例:固定された保持密度または半径)を共有するという「共有割り当て」の仮定を共有しています。これは、Softmax正規化が各クエリに対して独立して行われるという自己注意の本質的な性質と矛盾します。その結果、集中したアテンション分布(低エントロピー)を持つクエリは、多くのキーに注目することを強制されると計算を無駄にし、逆に拡散した分布(高エントロピー)を持つクエリは、共有予算が小さすぎると重要な相互作用を失う可能性があります。このミスマッチは、生成品質を低下させるか、最適な効率化を妨げます。
2. 手法:トークン・ラジアス・アテンション(TRA)
著者らは、明示的なクエリごとのキー・ランキングを行わずに、トークン固有の構造化された疎性を実現するトレーニングフリーのフレームワークである**トークン・ラジアス・アテンション(TRA)**を提案しています。TRAは、VDiTのアテンション・パターンを分析することから得られた2つの主要な経験的知見に基づいて動作します。
知見 I:トークン固有のアテンションの疎性
著者らは、同一レイヤーおよび同一ヘッド内のクエリ間において、「保持密度」(目標とするアテンション質量を維持するために必要な上位ランクのキーの割合)が数桁のオーダーで変化することを観察しています。決定的なことに、この保持密度とアテンション・エントロピーの間には強い対数線形関係があることを見出しました。
- メカニズム: 高エントロピーは、より大きなトークン予算を必要とする拡散したアテンション分布を示し、低エントロピーは、より少ないトークンを必要とする集中した分布を示します。
- 実装: TRAは、解析的な近似を用いて、クエリ・エントロピーをトークン固有の予算(B^≈τexp(Hi)/N)に直接マッピングすることで、高価なクエリごとのソートやランキングを排除します。
知見 II:トークン・ラジアス・アテンション・パターン
著者らは、あるクエリに対する支配的な相互作用が、空間ドメインにおいて**クエリ中心の円形近傍(circular neighborhoods)**を形成することを観察しています。正規化されたアテンション確率は、2D空間距離に対して近似的に指数関数的に減衰します。
- メカニズム: TRAは、任意のトップk個のキーを選択する代わりに、スカラー値のトークン予算を空間半径へと変換します。
- 時間的減衰: ビデオシーケンスを扱うために、空間半径は時間的距離による減衰ルール(ϕ(δ)=exp(−γδ))によって調整され、規則的な時空間サポート構造を作成します。
TRAパイプライン
- エントロピーから予算へ: 初期の高密度な「ウォームアップ」フェーズ中に、TRAは各クエリのエントロピーを計算します。このエントロピーは、特定のトークン予算へとマッピングされます。
- 予算から半径へ: 予算はクエリ固有のベース半径へと変換されます。この半径は、その後、時間的距離に基づいた各フレームの調整が行われ、2Dディスクマスクを形成します。
- 効率的な実行:
- タイル・メジャー・レイアウト(Tile-Major Layout): ビデオ・トークンをタイル・メジャー順序に再配置することで、空間的に隣接するトークン(半径内に収まるもの)がシーケンス内で連続するようにします。
- 融合CUDAカーネル(Fused CUDA Kernels): カスタムカーネルを用いて、距離計算、半径比較、ブロック・プルーニング、およびトークン投票を融合させ、不規則なトークンレベルの半径マスクを、FlashInferと互換性のある規則的なブロック疎なマスクへと変換します。
- ステップ間再利用: エントロピーはウォームアップ・フェーズでのみ計算され、デノイジング・ステップ間のエントロピー・パターンの安定性を利用して、その後の疎なステップで再利用されます。
3. 主な貢献
- トークン固有の疎性の発見: 本論文は、保持密度がクエリ間で劇的に変化するものの、アテンション・エントロピーとの予測可能な対数線形関係に従っていること、および支配的な相互作用がクエリ中心の半径パターンに従うことを明らかにしました。
- トークン・ラジアス・アテンション(TRA): 明示的なキー・ランキングを回避しながら、エントロピー・ツー・バジェット・ツー・ラジアス(entropy-to-budget-to-radius)パイプラインを通じて、クエリ固有の計算需要を規則的な時空間サポートへと変換する、トレーニングフリーのフレームワークです。
- システム・コデザイン: 著者らは、半径マスク・カーネルと融合エントロピー・カーネルを共同設計し、オーバーヘッドを最小限に抑え、既存のブロック疎なバックエンドでの効率的な実行を可能にしました。
- 包括的な検証: TRAは、Wan2.1、Wan2.2、およびHunyuanVideo(1.3Bから14Bパラメータに及ぶ)の7つの構成において、テキスト・トゥ・ビデオ(T2V)およびイメージ・トゥ・ビデオ(I2V)の両方のタスクに対して検証されています。
4. 実験結果
TRAは、VBenchベンチマーク・スイートにおいて、高密度アテンションおよび3つのトレーニングフリーの疎なベースライン(SVG、SVG2、およびRadial)と比較して評価されました。
- 効率性: TRAはアテンション相互作用のわずか9%~19%のみを保持します。テストされたすべての構成において、高密度モデルに対して1.56倍から2.05倍のスピードアップを達成しました。例えば、HunyuanVideo-13Bでは、T2Vで2.05倍、I2Vで2.02倍のスピードアップを実現しています。
- 品質: TRAは競争力のある生成品質を維持しており、多くの場合、疎な手法の中で最高のVBench総合スコアを獲得しています。Wan2.1-14B T2Vにおいて、高密度アテンションのスコアにほぼ匹敵しながら、1.75倍のスピードアップを提供しています。
- 忠実度: 定性的な結果は、TRAが他の疎な手法と比較して、被写体のアイデンティティ、構造的完全性、および時間的一貫性をより良く保持していることを示しており、ちらつきや歪みといったアーティファクトも少なくなっています。
- アブレーション研究: エントロピーによる予算マッピング、または半径マスキング(1D距離への置き換え)を削除すると、性能が著しく低下し、トークン固有の適応性と構造化された空間サポートの両方の必要性が確認されました。
5. 意義と主張
本論文は、現在の疎なアテンション・パラダイムにおける根本的な非効率性、すなわち共有割り当て予算とトークン固有のアテンション需要との間のミスマッチに対処していると主張しています。エントロピーとアテンションの疎性の間の本質的な関係を利用することで、TRAは追加のトレーニングや複雑なオンライン・ランキング・メカニズムを必要とせずに、良好な品質と効率のトレードオフを実現します。
著者らは、TRAを、キャッシュや量子化のような他の加速手法を補完する技術として位置付けており、実行されるアテンション層内での二次的なクエリ–キー相互作用コストの削減に特化しています。この研究は、トークンレベルの適応性が、構造化された半径ベースのマスキングを通じて効率的に実現できることを示しており、高忠実度なビデオ生成モデルをデプロイするためのスケーラブルな道筋を提示しています。