💻 computer science

PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

本論文は、非決定的なブロックを破棄するのではなく、ブロック単位のテイラー展開によって近似することにより、Diffusion Transformerにおいて劣二次計算量を達成し、高い生成品質を維持しながら大幅な高速化を実現する、学習不要の区分的スパースアテンション機構であるPISAを導入するものである。

原著者： Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

公開日 2026-02-04

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大で信じられないほど詳細な壁画（高解像度のビデオや画像のようなもの）を、アーティストのチームを使って描こうとしていると想像してください。AIの世界では、この「チーム」は**拡散トランスフォーマー（Diffusion Transformer）**であり、「描画」プロセスとは、次の筆致（トークン）を決定するために、あらゆる一筆一筆（トークン）を観察することを意味します。

問題は？現在のAIチームの働き方は、ある厳格なルールに従っているようなものです。「次の動きを決めるために、壁画全体のあらゆる筆致をすべて見ろ」というルールです。壁画が大きくなる（解像度が高くなる、あるいはビデオが長くなる）につれて、これは不可能になります。アーティストたちは圧倒され、プロセスは極端に遅くなり、コンピュータはエネルギーを使い果たしてしまいます。これが、論文で言及されている「二次関数的な複雑さ（quadratic complexity）」というボトルネックです。

これを解決するために、以前の手法は**「保持するか、捨てるか（Keep or Drop）」**という戦略を試みました。彼らはこう言いました。「よし、筆致の80%を無視して、最も重要な20%だけを見よう」。

欠陥： これは、顔を描こうとしているのに、その「上位20%」のリストに入っていなかったからという理由で、目や口を無視してしまうようなものです。その結果、画像はぼやけたり、グリッチが発生したり、重要な細部が欠落したりします。

新しいアイデア：PISA（Piecewise Sparse Attention / 区分的な疎なアテンション）

この論文の著者たちは、PISAと呼ばれるよりスマートな方法を提案しています。単に「重要でない」部分を無視するのではなく、PISAは壁画を**「区分的なパズル」**として扱います。

その仕組みを、簡単な比喩を使って説明します。

1. 「正確 vs 近似」戦略

あなたが大勢の人々のために、巨大なスープを作っているシェフだと想像してください。

従来の方法（Dense Attention / 密なアテンション）： 完璧な味にするために、スープのすべてのスプーン一杯分を味わいます。正確ですが、非常に時間がかかります。
「捨てる」方法（Standard Sparse Attention / 標準的な疎なアテンション）： スープの最初の20%だけを味わい、残りは無視します。これでは、奥の方にある塩分を見逃してしまうため、スープの味が変になってしまいます。
PISAの方法：
- クリティカル・ブロック（「正確」な部分）： 最も重要な材料（メインのスパイスや肉など）を特定します。これらは正確に、注意深く味わいます。
- 非クリティカル・ブロック（「近似」の部分）： スープの残り（水、出汁、野菜など）については、一滴一滴を味わう必要はありません。代わりに、そのセクションの平均的な味に基づいた数学的なショートカット（テイラー展開）を使用して、味を推定します。

2. なぜこれが「賢い」のか

論文では、AIのアテンションにおける「重要でない」部分（非クリティカル・ブロック）には、実は非常に予測可能であるという興味深い事実を発見しました。それらは滑らかで穏やかなパターンに従っています。

予測可能であるため、それらを捨て去る必要はありません。最終的な画像の「風味」を損なうことなく、それらを非常に素早く近似することができるのです。
これは、部屋の温度を測る際、空気の全箇所を測定するのではなく、部屋の隅々の平均温度を測って推定するようなものです。

3. 結果：犠牲のないスピード

「正確な」計算と「スマートな近似」を組み合わせることで、PISAは2つのことを達成します。

スピード： 現在の最高の手法よりも2倍から2.5倍速く動作します。論文のテストでは、これまで26分かかっていたビデオ生成が、約11分に短縮されました。
品質： 画像やビデオの見た目は、低速な「完璧な」バージョンと同じくらい良好です。実際、いくつかのテストでは、情報を単に切り捨てる他の「高速な」手法よりも、PISAの方が優れた結果を出しました。

裏側にある「魔法」

論文では、再学習を必要とせずにこれを可能にするいくつかのテクニックについても触れています。

再学習不要： PISAはオリジナルの「完璧な」思考プロセスを非常に密接に模倣しているため、既存のAIモデル（Wan2.1やFLUX.1など）に組み込むだけで、そのまま動作します。AIに新しい言語を教え直す必要はありません。
「ハイブリッド」カーネル： 著者たちは、コンピュータが混乱したり速度が低下したりしないように、「正確に味わう」ことと「推定する」ことを瞬時に切り替えるカスタムプログラム（カーネル）を構築しました。

まとめ

PISAを、AIアートスタジオの**「スマートなマネージャー」**と考えてください。

古いマネージャー： 「仕事の80%を無視しろ！」（結果：質の低いアート）。
PISAマネージャー： 「重要なディテールには100%の力を注ぎ、背景のノイズには素早くスマートな推測を使え。」（結果：高速で高品質なアート）。

この論文は、「推測」することは決して怠慢なショートカットではなく、傑作を損なうことなく時間を節約するための、数学的に裏付けられた方法であることを証明しています。

技術要約: PISA - 効率的な拡散トランスフォーマーのための区分的疎性アテンション (Piecewise Sparse Attention)

1. 問題提起

拡散トランスフォーマー (DiT) は、高忠実度の画像および動画生成における標準となっています。しかし、高解像度や長尺の動画のために入力シーケンス長 ( $L$ ) が増加すると、自己アテンション機構の二次関数的複雑性 ( $O(L^2)$ ) がスケーラビリティのボトルネックとなります。

既存の解決策には、以下のような重大な制限があります：

疎性アテンション ("Keep-or-Drop"): 非決定的なキー・バリュー (KV) ブロックを破棄する手法は、計算を加速させますが、高い疎性比においてコンテキストの喪失による性能低下を招きます。また、選択のオーバーヘッドがメリットを上回ってしまう短いシーケンスでは苦戦します。
線形/ハイブリッド・アテンション: アテンション分布を変更する手法（特徴量マッピングや、疎なブランチと線形ブランチの加法的な組み合わせなど）は、学習済みモデルの固有の正規化を乱すことがよくあります。これにより、事前学習済み重みの直接的な再利用を妨げる、高価な再学習やファインチューニングが必要になります。
ハード・トランケーション (強制的切断): ブロックを直接ドロップすることは、元のアテンション分布からの急激な逸脱を引き起こし、生成コンテンツの品質低下を招きます。

核心となる課題は、事前学習済みの重みと互換性を保ち、再学習を必要とせずに、フルアテンションと同等の精度を維持しながら、劣二次的な複雑性で推論を加速させるメカニメントを設計することです。

2. 手法: 区分的疎性アテンション (PISA)

PISAは、従来の「保持か破棄か (keep-or-drop)」という戦略を超えた、新しい**「厳密か近似か (exact-or-approximate)」**パラダイムを導入します。これは、アテンションを、劣二次的な複雑性でフルアテンションの全範囲をカバーする区分的なプロセスとして扱います。

主要な洞察

著者らは事前学習済みのDiTを分析し、2つの重要な特性を発見しました：

分布の安定性: 非決定的な（重要でない）ブロックのソフトマックス前アテンションスコアは、ゼロまたは負の値を中心とした対称的なベル型の分布を示します。これにより、これらは単純な破棄や厳密な計算ではなく、テイラー展開による正確な近似が非常に容易になります。
正規化の一貫性: 学習なしの互換性を維持するためには、近似は単に異なるブランチからの出力を加算するのではなく（これは重み付き和の規則に反します）、ソフトマックスの分子と分母の内部に統合される必要があります。

コア・アルゴリズム

PISAはKVブロックを2つの集合に分割します：

決定的なブロック ( $S_i$ ): 高いサリエンス（顕著性）情報を保持するために、厳密に計算されます。
非決定的なブロック ( $U_i$ ): ブロック単位のテイラー展開を用いて効率的に近似されます。

出力 $o_t$ は、重み付き値集計を正規化することで導出されます：
$o_t = \frac{N_t}{D_t}$

ここで、分母 ( $D_t$ ) と分子 ( $N_t$ ) は以下で構成されます：

厳密な疎項: 選択されたブロック $j \in S_i$ に対して計算されます。
ブロック単位の零次近似: 選択されなかったブロックに対して、ブロックの重心 $\bar{k}_j$ と値の総和を用いて近似されます。
グローバルな一次補正: すべての未選択ブロックに対して個別の一次項を計算する代わりに（これはメモリ制限の原因となります）、PISAはハイブリッド次数近似を採用します。全ブロックにわたるブロック単位の一次行列の平均であるグローバル統計量 $\bar{H}$ を使用して、テールセット全体の一次勾配を近似します。

分子 $N_t$ の定式化は以下の通りです：
$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$
ここで、 $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$ です。

共分散を考慮したルーティング

近似誤差を最小限に抑えるため、PISAは共分散を考慮したTop-k戦略を使用します。ブロックの選択スコアには、アテンションスコアだけでなく、ブロックの共分散行列のノルム ( $M_j = \|H_j - \bar{H}\|_2$ ) も組み込まれます。セマンティックな関連性が高い、あるいは不均一性が高いブロックは厳密な計算のために選択され、より滑らかなブロックはテイラー展開に委ねられます。

ハードウェア実装

PISAは、オンライン・ソフトマックスのプロセス内で厳密な計算と近似スキャニングを交互に行う融合カーネル (fused kernel) (FlashAttentionスタイル) として実装されています。パイプラインは以下を含みます：

事前計算: ブロックの重心、値の総和、およびグローバル統計量 $\bar{H}$ の計算。
フェーズ1 (厳密): 選択されたブロックをSRAMにロードして厳密なアテンションを実行。
フェーズ2 (近似): 未選択ブロックに対して、集約されたブロックベクトルを零次展開を用いてスキャン。
フェーズ3 (補正): グローバルな一次補正項を注入。

この設計により、メモリ効率が確保され、各ブロックに対して個別の行列をロードすることによるボトルネックが回避されます。

3. 主な貢献

新しい区分的疎性アテンション: 厳密な疎計算とブロック単位の近似を統合することで、劣二次的な複雑性でフルアテンションの全範囲のカバーを可能にする、学習不要のメカニズム。
ハイブリッド次数近似スキーム: ブロック単位の零次展開とグローバルな一次補正を組み合わせた戦略。これにより、標準的な疎性アテンションと比較して、計算コストをほとんど増やすことなく精度を大幅に向上させます。
共分散を考慮したルーティング: アテンションスコアとブロック共分散ノルムの両方に基づいてブロックを選択することで、近似の乖離を最小限に抑える誤差解析からの導出。
最先端のパフォーマンス: 多様なタスク（動画および画像生成）において、再学習を必要とせずに優れた効率性と品質を実証。

4. 実験結果

著者らは、動画生成モデル (Wan2.1-1.3B/14B, Hunyuan-Video) および画像生成モデル (SD 3.5, FLUX.1) でPISAを評価しました。

動画生成:
- Wan2.1-14B: 87.5%の疎性において、高密度アテンションに対し1.91倍の高速化 (1468s $\to$ 818s) を達成しつつ、高密度アテンションと同等またはわずかに優れた品質指標 (VBench, PSNR, SSIM) を維持しました。
- Hunyuan-Video: 一貫した品質保持とともに、2.57倍の高速化 (1651s $\to$ 641s) を達成しました。
- 比較: 速度と品質の両面において、特に微細な詳細や時間的一貫性の保持において、SpargeAttnやSVG2といった既存の疎性手法を上回りました。
画像生成:
- FLUX.1-dev: 85%の疎性において、1.2倍の高速化 (8.32s $\to$ 6.87s) を達成。特筆すべきは、PISAが高密度アテンションやSpargeAttnの両方と比較して、人間による好みの指標 (ImageReward, HPSv2) および類似性指標 (PSNR, SSIM) で高いスコアを獲得したことです。
- SD 3.5: 高密度アテンションと同等以上のFIDスコアを維持しながら、大幅な高速化 (例: SD 3.5-Mで1.4倍) を実証しました。
カーネル効率:
- PISAは、様々なシーケンス長 (4Kから32K) および密度において、FlashAttention-2/3およびSpargeAttnよりも一貫して低いレイテンシを示しました。
- 短いシーケンスで性能が低下するSpargeAttnとは異に対し、PISAは4Kトークンにおいても速度上の優位性を維持しました。

5. 重要性と主張

本論文は、PISAが「保持か破棄か」から「厳密か近似か」へのパラダイムシフトを象徴していると主張しています。その主な重要性は以下の通りです：

速度と品質のギャップを埋める: 推論速度と生成品質の間のギャップを効果的に埋め、無視できるほどのオーバーヘッドで、ほぼ損失のない近似を実現します。
学習不要の互換性: 近似をソフトマックス正規化フレームワークに直接統合することで、PISAは事前学習済みの重みの直接的な再利用を可能にし、線形またはハイブリッド・アテンション手法で必要とされる高価な再学習やファインチューニングを排除します。
スケーラビリティ: 生成コンテンツの忠実度を犠牲にすることなく、高解像度画像および長尺動画生成の両方に対する統一的な解決策を提供し、二次関数的な複雑性のボトルネックに対処します。

著者らは、彼らの手法が、事前学習済みモデルの固有の分布を保持する統一されたフレームワークの中に、ブロック単位のテイラー展開と疎性アテンションをネイティブに統合した最初の手法であることを強調しています。