技術的サマリー：DiffATS：整列テンソル空間における拡散

1. 問題定義

高解像度の時空間場（例：動画、気象データ、流体力学、偏微分方程式（PDE）の解）の直接拡散モデリングは、計算量的に実行不可能である。標準的な拡散モデル（DM）は、完全なデータ次元における反復的な関数評価を必要とし、膨大なトレーニングおよびサンプリングコストをもたらす。例えば、$84 \times 720 \times 1440$ のサイズのテンソル上で気象予測モデルをトレーニングするには、数十個の TPU で数日を要する。

これを緩和するため、潜在拡散モデル（LDM）は通常、データを低次元の潜在空間に圧縮するために変分オートエンコーダ（VAE）を採用する。しかし、科学データの場合、市販の事前学習済みオートエンコーダはしばしば利用できず、タスク固有のオートエンコーダをトレーニングすることは高価であり、高度な最適化を必要とする。離散コサイン変換（DCT）や固定 Tucker 分解に基づくオートエンコーダ不要のプリミティブは存在するが、これらは生データの本質的な構造を最適に近似しない可能性のある定義済みの基底に依存している。

本論文は以下の問いを提起する：事前学習済みオートエンコーダなしで、生成モデリングを改善するために、高解像度時空間データの本質的な低ランク構造を利用し、データ依存かつパラメータ効率的なプリミティブを構築することは可能か？

2. 手法

提案されるフレームワーク**整列テンソル空間における拡散（Diffusion in Aligned Tensor Space: DiffATS）**は、Tucker 分解と直交プロクラステス（OP）アライメントに基づくテンソルプリミティブの新たな構築を通じて、計算的および表現上の課題に対処する。

2.1 自然な Tucker プリミティブの退化

著者らは、標準的な Tucker 分解の因子（コアテンソルとモードごとの因子行列）を直接拡散トレーニングのプリミティブとして使用することは最適ではないと特定した。Tucker 分解は一意ではない。単一のテンソルは、直交回転によって関連するコアテンソルと因子行列の無限の組み合わせで表現され得る。この「ゲージ曖昧性」は、プリミティブ空間を分散させ、写像の連続性を破り、拡散モデルのスコア関数の学習を複雑にする。

2.2 直交プロクラステス（OP）アライメント

この退化を解決するため、DiffATS は直交プロクラステス（OP）アライメントを導入する。

1. アンカー選択: 本手法は、トレーニングデータから「メドイド」アンカー行列を選択する。これらのアンカーは、因子行列の実証分布内の最も中心的または代表的な部分空間を表す。
2. アライメント: 各データサンプルについて、因子行列を対応するアンカー行列にアライメントする。これは、フロベニウスノルム距離を最小化する最適な直交回転を見つけることにより行われる。これにより、各グラスマン同値類から一意の標準的な代表が選択される。

2.3 グラスマンプリミティブ

OP アライメントに基づき、本論文は 2 つのプリミティブのファミリーを定義する。

• 行列グラスマンプリミティブ（MGPs）: 低ランク行列（例：1 次元 PDE または画像パッチ）向け。プリミティブは、アライメントされた特異ベクトルと変換された特異値から構成される。
• テンソルグラスマンプリミティブ（TGPs）: 低多重線形ランクテンソル（例：動画または 2 次元 PDE）向け。プリミティブは、コアテンソルと選択されたモードに対するアライメントされた因子行列から構成される。

これらのプリミティブはデータ適応的（特定のデータセットの構造から導出される）かつパラメータ効率的であり、元のデータを $3.9$倍から$210$ 倍の範囲で圧縮する。

2.4 理論的保証

著者らは、特定の「アンカーオーバーラップ」条件（データの部分空間がアンカーと非自明な重なりを持つ場合）の下で、元のテンソル空間からプリミティブ空間への写像が同相写像であることを証明する。この理論的結果は以下のことを保証する。

• 写像は全単射（一対一かつ上への写像）である。
• 写像は連続であり、データ分布の局所的な位相構造を保存する。
• 表現は非退化であり、拡散モデルがプリミティブの忠実な分布を学習することを保証する。

2.5 DiffATS パイプライン

1. 前処理: 高次元テンソルは HOOI（高次直交反復法）または SVD を介して分解される。
2. アライメント: 因子行列は OP を介してメドイドアンカーにアライメントされる。
3. トレーニング: 拡散モデル（DiT バックボーンを使用）は、アライメントされたプリミティブ（コアテンソルとアライメントされた因子）上で直接トレーニングされる。
4. 推論: モデルはプリミティブ空間でサンプルを生成し、その後、明示的な多重線形再構成（コアにアライメントされた因子を乗算）を介して元のデータ空間に再構成される。

3. 主要な貢献

• 新規プリミティブ: 回転曖昧性をアンカー付き OP アライメントによって解決しつつ、Tucker 分解の表現力を保持する、データ適応的かつパラメータ効率的な表現としての MGPs と TGPs の導入。
• 理論的基盤: 重なり条件下で提案されたプリミティブ写像が同相写像であることを証明し、生成モデリングに使用される低ランクプリミティブファミリーに対して初めて形式的な同相写像保証を提供すること。
• オートエンコーダ不要フレームワーク: 事前学習済み深層圧縮オートエンコーダに依存せず、整列プリミティブ空間で直接トレーニングおよびサンプリングを行うエンドツーエンドフレームワーク DiffATS の提案。
• 実証的パフォーマンス: 画像、動画、PDE 解において、既存のオートエンコーダ不要ベースライン（DCTDiff、SDIFT、AvgPool）および競争力のあるニューラル演算子（FNO）を上回る DiffATS の性能を実証。同等または高い圧縮率で優れた生成品質を達成。

4. 実験結果

著者らは DiffATS を 6 つのベンチマークデータセットで評価した。

• 画像: CelebA-HQ（$1024 \times 1024$）。
• 動画: Moving MNIST。
• PDE: 1 次元 Burgers、1 次元反応拡散、2 次元 Burgers、および 2 次元カルマン渦列。

主要な知見:

• 生成品質: DiffATS は CelebA-HQ において最良の Frechet Inception Distance（FID）を達成し（23.57、AvgPool の 32.89 対）、Moving MNIST において最良の Frechet Video Distance（FVD）を達成した（415.42、AvgPool の 487.82 対）。
• PDE 予測: 条件付き生成タスク（初期条件を与えて将来の軌道を予測）において、DiffATS はすべてのベースラインを大幅に上回った。例えば、1 次元反応拡散データセットでは、最良のベースライン（SDIFT）と比較して相対 $\ell_2$ 誤差を約 $9$ 倍削減した。
• 圧縮: 本手法は、高い忠実度を維持しながら $3.9$倍から$210$ 倍の圧縮率を達成した。
• アブレーション研究: 実験により、OP アライメントを用いたデータ依存基底の使用が重要であることが確認された。アライメントを除去するか、共有基底を使用すると、パフォーマンスが大幅に低下した。

5. 意義と主張

本論文は、事前学習済みオートエンコーダが利用できない、またはトレーニングが非現実的な科学および高解像度時空間データにおいて、オートエンコーダベースの潜在拡散に対する堅牢な代替手段として DiffATS を提供すると主張している。このようなデータの本質的な低ランク構造を活用し、幾何学的アライメントを通じてテンソル分解の非一意性を解決することで、DiffATS は効率的かつ高品質な生成モデリングを可能にする。

著者らは、自らの研究を、生成タスクにおける標準的なテンソル分解の「退化」問題に対する解決策として位置づけており、安定した低次元表現を作成するための理論的に裏付けられた手法を提供している。彼らは、そのアプローチが Stable Diffusion のような大規模な事前学習済みオートエンコーダに依存する最先端の画像合成モデルを凌駕することを意図しているのではなく、そのようなリソースが不足しており、低ランク構造が支配的な特徴である時空間物理場の特定の領域を対象としていることを強調している。

著者らが示唆する将来の方向性には、同相写像保証に必要な仮定を緩和するために、マルチアンカー OP アライメントと動的ランク選択の探求が含まれる。