Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling

技術要約：Tensor Train Diffusion

問題提起
本論文は、複雑で高次元の確率密度 $p_{\text{target}} = \rho_{\text{target}}/Z$ からサンプリングするという課題に取り組んでいる。ここで、非正規化密度 $\rho_{\text{target}}$ は点ごとに評価可能であるが、正規化定数 $Z$ は計算不能であり、また $p_{\text{target}}$ からのサンプルも利用できない。これは、分子動力学やベイズ統計を含む計算科学における中心的な問題である。

拡散モデル（Diffusion Models）は、ノイズ付加の確率微分方程式（SDE）を逆転させることで、このようなタスクに対して強力なフレームワークを提供するが、通常はデノイジング・スコアマッチングを通じてサンプルからスコア関数を学習する必要がある。この特定の状況下では、スコアはサンプルにアクセスすることなく、非正規化密度から直接導出されなければならない。時間の逆転ダイナミクスは、ハミルトン・ヤコビ・ベルマン（HJB）偏微分方程式（PDE）によって支配される。PINNs（Physics-Informed Neural Networks）や軌道ベースの手法といった既存の高次元PDE解法は、長い学習時間、ハイパーパラメータへの高い敏感さ、および $\rho_{\text{target}}$ とその高次導関数のコストの高い評価を必要とするという問題を抱えている。

手法
著者らは、Functional Tensor Train (FTT) 表現を活用した、基礎となるHJB方程式の新しいソルバーである Tensor Train Diffusion (TTD) を提案する。

1. 数学的定式化:

   • 問題は、前向きのノイズ付加プロセスを逆転させるサンプリングSDEのための最適制御（ドリフト）を見つけることとしてフレーム化される。
   • 前向きプロセスの対数密度 $V = -\log \vec{p}_Y$ は、特定のHJB PDEを満たす。
   • HJB方程式と後退確率微分方程式（BSDE）の接続を利用し、著者らはPDEの解を回帰問題のシーケンスとして再定式化する。これにより、解はバックワード・イン・タイムの反復スキームを通じて計算可能となる。

2. Functional Tensor Train (FTT) 近似:

   • ニューラルネットワークの代わりに、値関数 $V$ とその勾配をFTTを用いて近似する。この形式は、高次元関数に潜む低ランク構造を利用し、低ランクのテンソルコアの積として表現する。
   • 本手法は、TT形式と特定の単変数基底関数（例：ルジャンドル多項式、B-スプライン、フーリエモード）を組み合わせた Extended Tensor Trains (xTT) を活用する。
   • FTT係数の最適化は、交互最小二乗法 (ALS) を用いて行われる。これは、確率的勾配降下法（SGD）の不安定性を回避する回帰ベースのアプローチである。

3. アルゴリズムの特徴:

   • 反復精緻化 (Iterative Refinement): アルゴリズムは、制御ポリシーを反復的に更新する外側ループを採用している。初期の推測（例：ランジュバン動力学）から開始し、手法はサンプリング軌道を反復的に精緻化することで、訓練分布をターゲット分布にうまく適合させ、新たなモードを発見する。
   • 適応性 (Adaptivity): 本フレームワークには、以下の適応戦略が含まれている：
     • ランク選択: 特異値の減衰に基づいてTTランクを動的に調整し、精度と計算コストのバランスを取る。
     • 基底次数: 各タイムステップにおける解の滑らかさに基づいて、基底関数の多項式次数を適応させる。
     • 正則化: 導関数の学習を安定させるために、正則化の大きさを自動的に調整する。
   • ドメイン処理: 手法は、移動する近似ドメインを扱い、学習されたポリシーをアフィン線形拡張を用いて訓練ドメイン外の領域にも拡張し、サンプリング中の安定性を確保する。

主な貢献

• 新規ソルバー: 拡散サンプリング、HJB理論、およびBSDEを統合し、計算コストと精度の良好なトレードオフを実現する、HJB方程式のためのソルバーであるTTDの導入。
• 堅牢性と柔軟性: 多様な基底関数（ルジャンドル、B-スプライン、フーリエ）を取り入れることによる柔軟性の実証、および、特に主要な訓練ドメイン外の軌道を扱う際の、従来のテンソル・トレインベースのPDEソルバーに対する堅牢性の向上。
• 性能: TTDが、様々な次元の複雑なマルチモーダル分布からサンプリングできることを示し、既存の拡散ベースのサンプラーよりも速度と精度の両面で優れていることを示す。

実験結果
著者らは、いくつかの困難なベンチマークにおいてTTDを評価している：

• 高次元マルチウェル問題: $d=10$ および $d=50$ の複数のモードを持つ次元において、TTDは高忠実度のサンプリングを実現する。本手法は、拡散ベースのサンプリングにおける一般的な問題であるモード崩壊を回避する。実行時間はタイムステップの数に対して線形にスケールし、ニューラルネットワークベースの代替手法（例：DIS）よりも大幅に高速である。
• 異方性ガウスターゲット: アルゴリズムがターゲットが標準正規分布に収束するにつれてランクを適応的に減少させる能力を確認し、ランク適応メカニメントの妥当性を検証した。
• $\phi^4$ スカラー場理論: 指数関数的に多くのメタステーブルモードと強い近傍相関を持つシステムである、離散化されたギンツブルグ・ランダウ・モデルからのサンプリングに成功した。これは、標準的なMCMC手法がモード崩壊により失敗する典型的なケースである。
• 比較: TTDは、有効サンプルサイズ（ESS）、対数分散ダイバージェンス、および推定正規化定数の精度において、最先端のベースライン（PIS、DDS、DIS、および各種MCMCバリアントを含む）を一貫して上回り、しばしばより短い時間でより良い結果を達成している。

意義と主張
本論文は、TTDが、高次元のHJB方程式を解くための、高速、正確、かつ安定したニューラルネットワーク技術に代わる手段を提供すると主張している。低ランクのテンソル構造を活用することで、本手法は、従来のメッシュベースの手法に伴う「次元の呪い」や、SGDベースのニューラルソルバーの高価な学習コストおよびハイパーパラメータへの敏感さを回避する。著者らは、TTDをHJB方程式を解くための堅牢な基盤として位置づけており、ターゲット密度の評価回数が少なく、高次元サンプリング問題において優れた効率性を提供することを強調している。著者らは、非常に複雑な分布においてすべてのモードを特定する保証や、特定のターゲット密度に対するFTTの表現力の限界についても認めているが、テストされた領域における本手法の優れた性能を強調している。