巨大で極めて詳細な壁画を壁に描こうとしている状況を想像してください。壁の前に「トークン」と呼ばれる芸術家のチームが立ち、それぞれの芸術家は、自分の特定の場所にどの色を塗るかを決定するために、他のすべての芸術家と話し合わなければなりません。これにより絵画は完璧になりますが、非常に疲弊します。1 万人の芸術家がいた場合、次の筆遣いを合意するためには 1 億回の会話が必要になります。これが現在の AI 画像生成器（拡散トランスフォーマーと呼ばれる）の仕組みです：これらは卓越していますが、非常に遅く、計算資源を大量に消費します。
この論文は、このプロセスを絵画を損なうことなく高速化するための新しい手法「CoReDiT」を紹介しています。その仕組みは、以下の 3 つの簡単なステップに分解されます。
1. 「退屈な芸術家」検出器（空間的コヒーレンス）
典型的な絵画では、顔や木のように非常に忙しく詳細な領域がある一方で、青空や滑らかな壁のように単調で退屈な背景の領域もあります。
- 問題点: 従来の AI はすべての芸術家を同様に扱い、退屈な青空を塗る芸術家にも、顔を塗る芸術家と同じ激しい会話を強要します。
- CoReDiT の解決策: CoReDiT は芸術家たちを見て、「自分と全く同じように見える人の隣に立っているのは誰か？」と問いかけます。退屈で均一な領域にいる芸術家は、その隣人が実質的にクローンです。CoReDiT はこれらの「退屈な」芸術家に高いコヒーレンス・スコアを与えます。
- 実行: 高スコア（冗長）な芸術家に対し、「今は全員と話す必要はない。休憩を取れ」と伝えます。これにより、画像の退屈な部分における高コストな会話をスキップし、膨大な時間を節約します。
2. 「穴埋め」芸術家（再構成）
もし単に「退屈な」芸術家に帰宅して作業を停止させれば、壁画には穴が開いてしまいます。絵画は壊れてパッチワークのようになってしまうでしょう。
- 問題点: トークン（芸術家）を単に削除すると、視覚的な不具合が生じます。
- CoReDiT の解決策: 穴を開けたままにするのではなく、CoReDiT は残った芸術家を使って空白を埋めます。「退屈な」芸術家は隣人と非常に似ていたため、残った芸術家は、欠落した芸術家が何を塗ったかを容易に推測できます。
- 実行: これは「伝言ゲーム」のようであり、残った芸術家が正しい色を空いた場所に囁き伝えます。これにより、実際に重労働を担っていた芸術家の数が少なくても、最終的な画像は滑らかで完全に見え、目に見える隙間が生じません。
3. 賢い管理者（段階的剪定）
絵画プロセスのすべての部分は同じではありません。
- 問題点: 画像作成の非常に初期段階では、画像は単なる静的なノイズ（テレビの雪ノイズのようなもの）であり、ほぼすべてが似通っており冗長です。しかし、最終段階で詳細が鮮明になる頃には、すべての筆遣いが重要になります。早期に芸術家を削りすぎると AI が混乱し、終盤に削りすぎると画像がぼやけてしまいます。
- CoReDiT の解決策: CoReDiT は、最適なスケジュールを学習する賢い管理者のように機能します。画像がノイズのみの段階では、AI が多くの冗長な芸術家を削除できるようにします。画像が鮮明になり詳細が増すにつれて、管理者は徐々に削除する芸術家の数を減らし、重要な詳細に十分な注意が払われるようにします。
- 実行: AI の脳（ブロック）のどの部分と、プロセスのどの段階（時間ステップ）が最も多くの削減に耐えられるかを自動的に判断し、品質を高く保つために計画をその場で調整します。
結果
この手法を使用することで、著者らは以下の結果を確認しました。
- 速度: 強力なクラウドコンピューター上では AI が1.33 倍高速に動作し、スマートフォン用チップ上では1.72 倍高速になります。
- 効率: 重い数学的作業（計算）を最大**55%**削減します。
- 品質: 画像は遅いバージョンと同じように美しく見え、明らかな亀裂や奇妙なアーティファクトは生じません。
- ボーナス: メモリ使用量が少なくなるため、通常はメモリ不足でクラッシュするスマートフォンでも、より高解像度の画像を生成できます。
要約すると、CoReDiT は、クルーがいつ休ませ、いつハードワークさせるべきかを正確に知っている、芸術プロジェクトの賢い監督者を雇うようなものであり、その結果、マスターピースがはるかに速く完成します。
技術的概要：CoReDiT
問題定義
拡散トランスフォーマ（DiT）は、アテンション機構を通じて長距離依存関係を効果的にモデル化することで、画像および動画生成において最先端のパフォーマンスを確立している。しかし、その計算コストはトークン数に対して二次的に増大するため、高解像度合成や多フレーム動画生成は、リソース制約のあるモバイルプラットフォームでは実質的に不可能なほど高価となる。潜在空間内のトークンの大部分は視覚的に冗長な領域（例：均一な背景）に対応しているが、既存のトークンプルーニング手法は、拡散推論に適用される際に以下の 3 つの主要な課題に直面している：
- 効率的な局所化: グローバルアテンションスコアの計算という二次的なオーバーヘッドを伴わずに、低サリエンシーのトークンを特定すること。
- 意味的保存: トークンをスキップする際に、視覚的連続性を維持し、アーティファクト（例：テクスチャの断絶、境界の歪み）を防止すること。これは、識別タスクとは異なり、生成タスクにおいて極めて重要である。
- 適応的スケジューリング: 異なるトランスフォーマブロックおよびノイズ除去タイムステップ間で変化する冗長性レベルに適応するプルーニング戦略を決定すること。
手法：CoReDiT
CoReDiT は、出力忠実度を維持しながら事前学習済み DiT を高速化するために設計された構造化されたトークンプルーニングフレームワークである。これは以下の 3 つの中核コンポーネントを通じて動作する：
1. 空間的整合性に基づくトークン選択
グローバルアテンションスコア（O(N2)）に依存するのではなく、CoReDiT は線形時間（O(N)）で局所的な冗長性を推定するために**空間的整合性（SC）**スコアを利用する。
- メカニズム: トークン格子は非重複のグリッドに分割される。各トークンについて、SC スコアは、そのトークンの正規化埋め込みと、そのグリッドの平均正規化埋め込みとの間の平均コサイン類似度として計算される。
- 論理: 高い SC スコアを持つトークンは、空間的近隣によってよく説明されるため、冗長であるとみなされる。
- 選択: 最も高い SC スコアを持つトップ-Kのトークンがスキップ対象（XS）として選択され、残りのトークン（XR）はマルチヘッドセルフアテンション（MHSA）レイヤーに進む。
2. 整合性ガイド型トークン再構築
トークンをスキップすることによって引き起こされる空間的不連続性を防止するため、CoReDiT は保持された近隣（YR）の出力を用いて、スキップされたトークンのアテンション出力（YS）を再構築する。
- 近似: 再構築のためのペアワイズ類似度を計算する（オーバーヘッドを追加する）代わりに、本手法は近隣の事前計算済み SC スコアを用いて類似度重みを近似する。
- 局所サブグリッド: 過度な平滑化を避けるため、再構築はより小さなサブグリッド内で行われる。
- マイクロデザイン:
- 交互グリッドサイズ: グリッド寸法を連続するブロック間で交互に変更することで、境界アーティファクトをシフトさせ、グリッド境界を越えた情報の伝播を可能にする。
- m-ストライド保持: 決定論的なストライドパターンにより、各サブグリッドに少なくとも 1 つのトークンが保持され、再構築に必要なアンカーが提供される。
3. 段階的・ブロック適応型プルーニングスケジュール
冗長性がトランスフォーマブロックおよびノイズ除去ステップ間で変化することを認識し、CoReDiT は軽量ファインチューニング中に段階的なプルーニング戦略を採用する。
- ブロック適応性: プルーニング予算は、候補トークンの SC スコアの合計によって推定される、最も高い測定された冗長性を持つブロックに割り当てられる。
- タイムステップ減衰: プルーニングは、ノイズが支配的で構造が粗い初期のノイズ除去ステップではより積極的に行われ、微細な詳細が洗練される後期のステップでは軽減される。ステップ固有の計算グラフの管理の複雑さと効率性のバランスを取るために、2 段階の減衰スケジュールが使用される。
- トレーニング: モデルは、元のモデル出力損失とプルーニング済みブロック損失を組み合わせた蒸留損失を用いてファインチューニングされ、パフォーマンスを回復させる。
主要な貢献
- 線形時間冗長性推定: グローバルアテンション分析の計算ボトルネックを回避し、O(N)時間で冗長なトークンを特定する空間的整合性スコアの導入。
- 密格子保存: 近隣からスキップされたアテンション出力を合成する再構築メカニズムにより、後続のレイヤーのためにトークン格子が密な状態を維持し、視覚的アーティファクトを防止する。
- 適応的プルーニング戦略: 深度および時間全体の実証的冗長性統計に基づいてプルーニング予算を動的に割り当てる、段階的かつブロック適応型のスケジュールにより、安定した適応を可能にする。
実験結果
CoReDiT は、PixArt-α、PixArt-Σ、およびMagicDrive-V2を含む最先端の DiT バックボーンで評価された。
- 効率性の向上:
- セルフアテンション FLOPs を最大55% 削減達成。
- 推論速度向上: クラウド GPU（Nvidia H100）で 1.33 倍、モバイル NPU（Qualcomm Snapdragon 8 Elite）で1.72 倍。
- メモリ余裕: 1600 解像度でベースラインモデルがメモリ不足（OOM）エラーに遭遇したモバイルデバイスにおいて、1920 までの解像度での生成を可能にした。
- 品質の維持:
- PixArt-α-1024で 40% プルーニングの場合、CoReDiT は FID 28.7（ベースライン 27.3 対比）および競争力のある CLIP/IS スコアを維持した。蒸留により、FID の差は 27.4 に縮小した。
- PixArt-Σ-2048では、平均 23% のプルーニング比率により、セルフアテンションレイテンシを 22% 削減し、品質指標のわずかな低下のみで済んだ。
- MagicDrive-V2（動画）では、セルフアテンション FLOPs を 39% 削減しても、フレームごとの品質（PSNR、LPIPS、SSIM）および条件整合性（mAP、mIoU）は維持されたが、時間的整合性（FVD）のわずかな低下が観察された。
- アブレーション研究:
- ランダムなトークン選択は壊滅的な品質低下（FID 644.5）をもたらしたことで、整合性に基づく選択の必要性が実証された。
- 再構築を無効化すると、品質の大幅な損失（FID 30.4 対 28.5）が生じ、再構築ステップの重要性が確認された。
- 均一プルーニング followed by ファインチューニングは、提案された段階的適応スケジュールよりも性能が劣った。
意義と主張
本論文は、CoReDiT が生成トークンプルーニングの特定の課題に対処することで、エッジデバイスへの高忠実度拡散モデルの実用的な展開を提供すると主張している。その意義は以下の点にある：
- スケーラビリティ: 計算およびメモリフットプリントの削減により、モバイルハードウェアでの高解像度合成および動画生成を可能にする。
- アーキテクチャ互換性: 構造変更や複雑なルーティングモジュールを必要とせず、既存の DiT アーキテクチャとの互換性を維持する。
- 効率性と品質のトレードオフ: 構造化された整合性ガイド型スキップが、推論コストを劇的に削減（セルフアテンション FLOPs 最大 55%）しつつ、フルモデルと同等の知覚品質を維持することを示している。
著者は、現在の手法は空間的冗長性に焦点を当てているが、将来の作業では FFN/MLP コンポーネントへの効率性の拡張と、動画の時間的整合性をさらに向上させるためのフレーム間整合性目的の組み込みを目指すと述べている。
毎週最高の computer science 論文をお届け。
スタンフォード、ケンブリッジ、フランス科学アカデミーの研究者に信頼されています。
受信トレイを確認して登録を完了してください。
問題が発生しました。もう一度お試しください。
スパムなし、いつでも解除可能。
週刊ダイジェスト — 最新の研究をわかりやすく。登録