技術サマリー: Kaleido

問題提起

ビデオ拡散トランスフォーマー（vDiT）は、高忠実度なビデオ生成のための最先端アーキテクチャとして台頭していますが、深刻な計算上のボトルネックに直面しています。推論プロセスは多くのタイムステップにわたる反復的なデノイジングを伴い、そこでは自己注意（self-attention）メカニズムが実行時間の約68%を占める支配的な要因となります。ビデオの解像度や持続時間が増加するにつれ、トークンシーケンスの長さが増大し、自己注意の複雑さは二次関数的（$O(N^2 \times d)$）にスケールします。

既存の加速戦略の多くは、意味的な疎性（semantic sparsity）に基づいて重要度の低いアテンションスコアを削減する、大規模言語モデル（LLM）からのスパースアテンション技術を適応させていますが、これらはビデオデータに対して根本的に不適合です。テキストとは異なり、ビデオトークンは、すべてのトークンが空間的な一貫性と時間的な一貫性に寄与する、高密度で連続的な視覚的意味論を持っています。アテンションの相関を不用意に削減すると、動きの不連続性やフリッカー（ちらつき）アーティファクトが発生します。さらに、蒸留（distillation）によってディフュージョン・タイムステップが削減されると、自己注意が唯一の支配的なボトルネックとなるため、タイムステップ削減技術だけでは不十分となります。

手法

Kaleidoは、vDiTの潜在空間におけるトークンレベルの疎性ではなく、チャネル単位の時空間相関を利用するアルゴリズムとハードウェアの協調設計（algorithm-hardware co-design）を通じて、これらの課題に対処します。

1. アルゴリズム: チャネル単位の再利用（Channel-Wise Reuse）

核心となる洞察は、vDiTモデルが回転位置エンコーディング（RoPE）を利用しており、これがトークンのチャネルを時間（$t$）、水平（$x$）、垂直（$y$）の3つの異なるグループに分割しているという点にあります。これらのグループは異なる周波数で位置情報をエンコードしており、特定のアテンションスコアのパターンを誘発します。

• 相関の発見: 著者らは、特定のチャネルグループ内の隣接するトークンが、そのエンコーディング方向に対して直交する次元に沿って高い類似性を示すことを証明しました。例えば、$t$チャネルは空間次元（$x$および$y$）にわたって高い相関を持ち、$x$チャネルは時間次元にわたって相関を持ちます。
• 再利用メカニズム: 計算をスキップする（スパースアテンションのように）のではなく、Kaleidoは部分的な結果を再利用します。
  • 類似度チェック: 計算の前に、システムは定義された軸に沿った隣接トークン間の類似度（$\Delta$）を計算します。
  • 閾値処理: 2つの閾値 $\theta_{th1}$ と $\theta_{th2}$ が再利用戦略を決定します：
    • $\Delta < \theta_{th1}$ の場合：完全再利用（Full Reuse）。そのトークンは、前のトークンの部分的なアテンションスコアを完全に再利用します。
    • $\theta_{th1} \le \Delta < \theta_{th2}$ の場合：部分的再利用（Partial Reuse）。そのトークンは、前の結果の最も重要なビット（例：上位3ビット）を再利用し、残りの下位ビットのみを計算します。
    • $\Delta \ge \theta_{th2}$ の場合：完全計算（Full Computation）。
  • 適用範囲: この戦略は、自己注意、クロスアテンション、およびMLP層に適用されます。MLPの場合、再利用は重みではなく活性化関数に対して適用されます。

2. ハードウェア: 再構成可能なシストリック・アクセラレータ

再利用アルゴリズムによって導入される不規則なデータアクセスパターンは、標準的なシストリック・アレイとは互換性がありません。Kaleidoはこれをサポートするためにカスタムハードウェアアーキテクチャを導入しています。

• 再構成可能なプロセッシングエレメント（PE）: PEアレイは、4つの実行モードをサポートするために追加のレジスタとマルチプレクサで拡張されています。これらのモードにより、PEは単一のタイル内で動的にフル計算、フル再利用、および部分的再利用（ハーフビット計算）を処理できます。PEは、オンチップメモリからロードし直すことなく、前のフレームからの部分的な結果を蓄積できます。
• データディスパッチャ: 低いPE利用率を引き起こす不規則な疎性を緩和するために、専用のデータディスパッチャが導入されています。これは以下で構成されます：
  • スケジューラ: 同様の再利用パターン（バイナリ文字列として表現される）を持つトークンをクラスターにグループ化します。
  • マッチャー: クラスター内の互換性のある入力チャネルをペアリングしてデータストリームをマージし、PEアレイに供給される隣接行が同時に処理されるようにすることで、利用率を最大化します。

主な貢献

1. vDiTの体系的な特性評価: 本論文は、vDiTの潜在空間における時空間相関を体系的に説明し、RoPEに基づくチャネル分割が特有のアテンションパターンの根源であることを初めて明らかにしました。
2. チャネル単位の再利用アルゴリズム: チャネルレベルで部分的な結果を再利用することで冗長な計算をスキップする軽量なアルゴリズムです。トークンレベルのマージとは異なり、これはピクセルレベルの精度を維持するため、大幅に高い生成品質（従来の協調設計に対して17 dB以上のPSNR向上）を実現します。
3. 協調設計されたアクセラレータ・アーキテクチャ: 不規則な疎性を処理する、カスタマイズされたPEとデータディスパッチャを備えた再構成可能なシストリック・アレイであり、PE利用率を約69%から約85%に向上させます。
4. 包括的な評価: 4つの主要なvDiTモデル（HunyuanVideo, Wan, CogVideoX, TurboDiffusion）を用い、最先端のGPUおよびアクセラレータと比較した広範な評価を行いました。

結果

16nmプロセス技術、128x128 PEアレイ構成で評価：

• パフォーマンス: Kaleidoは、最先端のアクセラレータ（例：AdapTiV, Exion）に対して最大5.9倍、Nvidia A100 GPUに対して6.6倍の高速化を達成しました。
• エネルギー効率: システムは、A100ベースラインと比較して最大16.0倍のエネルギー節約を実現します。
• 生成品質: Kaleidoは優れたビデオ品質を維持し、HunyuanVideoで29.9 dB、Wanで24.5 dBのPSNRスコアを達成しました。これは、Cambricon-Dのような近似ベースのベースラインよりも17 dB以上高い値です。また、ソフトウェア最適化されたスパースアテンション手法（MInference, SVGなど）よりも、画像品質指標（PSNR, SSIM, LPIPS）において優れています。
• スケーラビリティ: ビデオ解像度が480pから720pに増加するにつれて、パフォーマンスの差は拡大します。これは、高解像度ほどアルゴリズムが利用する時空間相関が強まるためです。

重要性

本論文は、Kaleidoをビデオ生成加速における必然的な進化として位置付けています。LLMベースのスパースアテンション技術をビデオに単純に転用することは、視覚データの連続性のために効果的ではないと主張しています。ビデオ生成の「プルーニング（削減）」から、vDiT独自のチャネル単位の相関に基づく「再利用」へと焦点を移すことで、Kaleidoは出力品質を損なうことなく、計算上のボトルネックと高忠実度の要求を同時に解決する原理的なアプローチを提供します。この研究は、次世代のビデオ拡散モデルの効率を解き放つためには、アルゴリズムとハードウェアの協調設計が不可欠であることを示しています。