技術要約：拡散サンプリングによるデータ変換の反転

問題定義

本論文は、一般的なリー群における**変換の反転（transformation inversion）**の問題に取り組んでいる。多くの機械学習や科学モデリングのシナリオにおいて、観測データ $\tilde{x}$ は、分布内サンプル $x$ に対して未知の変換 $g$（リー群 $G$ の元）を適用することで生成される（すなわち $\tilde{x} = g \cdot x$）。目標は、分布外の観測値を元のデータ分布へとマッピングする逆変換 $g^{-1}$ を復元することである。

これは、未知の摂動変換が介在し、解が一般に一意ではないため、**ブラインド逆問題（blind inverse problem）**として定式化される。既存の手法は、データ空間上で動作するものや、特定の変換ファミリー（アフィン変換や回転群など）に特化したものが多いが、本研究は、非コンパクトかつ非アーベル型で、双定値計量（bi-invariant metric）を持たないグループ（例：画像のホモグラフィ $PGL(3, \mathbb{R})$ や、偏微分方程式におけるリー点対称性）を含む、一般的なリー群作用を対象としている。

強調されている主要なアプリケーションは、**テスト時の等変性（test-time equivariance）**である。学習済みニューラルネットワークは、幾何学的変換に対して頑健性に欠けることが多い。著者らは、入力データを学習分布へと戻すために、テスト時にこれらの変換を反転させる手法を提案しており、これにより、特殊な等変アーキテクチャや再学習を必要とせずに、頑健性を向上させる。

手法：変換反転エネルギー拡散 (TIED)

提案手法である TIED は、データ空間上のエネルギー関数 $E$ によって定義されるボルツマン分布として、変換の事後分布 $p(g | \tilde{x})$ をモデル化する。エネルギーはデータ事前分布 $p_x(x)$ から導出される（ここで $E(x) = -\log p_x(x)$）。事後密度は以下のように与えられる：
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

この事後分布からサンプリングするために、TIED はリー群上で直接行われる拡散プロセスを導入する。この手法は、エネルギー地形の険しくマルチモーダルな性質（多くの場合、ニューラルネットワークから導出される）と、多様体上でのサンプリングにおける幾何学的制約という2つの主要な課題を克服している。

1. リー群における自明化された拡散 (Trivialized Diffusion on Lie Groups)

任意の群の元における接空間で更新を行う代わりに、TIED は**自明化（trivialization）**を利用する。更新はリー代数 $\mathfrak{g}$（単位元における接空間）内で行われ、指数写像を通じて現在の群の元へと輸送される。

• 順方向プロセス： リー代数における微小な摂動を抽出し、右乗法によって輸送することで、群上に分散拡大型確率微分方程式（variance-exploring SDE）を定義する：
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  ここで $w^{\mathfrak{g}}_t$ はリー代数におけるブラウン運動である。
• 逆方向プロセス： 事後分布からのサンプルを回収するために、逆時間 SDE が導出される。決定的なことに、ドリフト項には自明化されたスコア $\nabla_g \log p_t(g_t)$ が含まれ、これはリー代数の元である。

2. 自明化されたターゲット・スコア恒等式

核心となる理論的貢献は、新しい自明化されたターゲット・スコア恒等式である。標準的なスコアマッチングはクリーンなサンプルへのアクセスを必要とするが、この反転の設定ではそれらは利用できない。代わりに、TIED はノイズ分布のスコアを、クリーンなエネルギー関数の勾配を用いて表現する。
著者らは、一般的なリー群（双定値計量を必要としない）において、自明化されたスコアがノイズカーネル $k_t$ に関するモンテカルロ期待値を通じて記述できることを証明した：
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
ここで $\lambda(w)$ は反転に伴うハール測量の変化を考慮している。これにより、自動微分を通じて計算されるエネルギー関数 $E$ の勾配のみを使用して、スコアを推定することが可能になる。

3. 実装の詳細

サンプリングアルゴリズム（Algorithm 1）は、逆 SDE を離散化する。各ステップにおいて、以下の処理を行う：

1. ノイズ $w$ をノイズカーネル $k_t$ からサンプリングする。
2. 変換された候補に対するエネルギー $E$ を計算する。
3. モンテカルロ恒等式を用いて、自明化されたスコアを推定する（安定性のために log-sum-exp トリックを用いる）。
4. 指数写像を用いて群の元を更新する。

本手法は、VAEから導出されたものや、学習済み分類器の予測信頼度から直接導出されるものなど、様々なエネルギー定義をサポートしている。

主要な貢献

1. TIED サンプラー： 一般的なリー群上でデータ変換を反転させるための、新しい拡散サンプラー。更新が多様体上に留まることを保証し、リー代数の知識のみを用いて非コンパクトかつ非線形な作用を扱う。
2. 理論的枠組み： 変換の反転が、学習済みモデルのテスト時等変性として機能することを証明。著者らは、変換反転の事後分布が $G$-等変であることを示し、このアプローチの妥当性を保証している。
3. 新しいスコア恒等式： 双定値計量を必要とせずにリー群で機能する、ターゲット・スコア恒等式の一般化。これにより、クリーンなデータサンプルではなく、エネルギーの勾配を用いた効率的なスコアベース・サンプリングが可能になる。
4. 実証的検証： 画像ホモグラフィ（$PGL(3, \mathbb{R})$）および PDE 対称性（$SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$）に関する実験により、強力なベースラインに対する一貫した性能向上を示した。

実験結果

論文では、主に3つのタスクで TIED を評価している。

1. $SO(10)$ 上の合成サンプリング： TIED は、高次元リー群（$dim=45$）上のマルチモーダルなエネルギー分布からサンプリングすることに成功し、自明化されたキネティック・ランジュバン・サンプラーよりも大幅に速く収束した。
2. 画像分類 (MNIST)：
   • 設定： 未知のアフィン変換（$Aff(2, \mathbb{R})$）およびホモグラフィ変換（$PGL(3, \mathbb{R})$）を受けた MNIST 画像に対して、学習済み ResNet18 をテスト。
   • ベースライン： 最適化ベースの手法（LieLAC, FoCal, ITS）およびサンプリング手法（Kinetic Langevin）と比較。
   • 結果： TIED はすべてのベースラインを大幅に上回った。例えば、分類器ベースのエネルギーを用いたアフィン変換下では、生の ResNet18 の 55.48% や Kinetic Langevin の 74.72% に対し、TIED は 82.64% の精度を達成した。また、一部の設定では特殊な等変ネットワーク（affConv）をも上回った。
3. PDE 解法：
   • 設定： 未知のリー点対称性変換を受ける 1次元熱伝導方程式および Burgers 方程式に対して、DeepONet ニューラルオペレータをテスト。
   • 結果： TIED はすべてのケースにおいて最も低い相対 $L_2$ 誤差を達成し、LieLAC や FoCal を上回った。$SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$ 対称性を持つ1次元熱伝導方程式において、TIED は誤差を 0.043 まで低減させた（LieLAC は 0.066、補正なしの DeepONet は 0.690）。

意義と主張

著者らは、TIED が学習済みモデルをテスト時に任意のリー群変換に対して等変にするための、トレーニングフリーのパイプラインを提供すると主張している。エネルギーベースモデル（学習済みモデル自体で近似可能）を活用することで、本手法は特殊な等変アーキテクチャや高コストな再学習を回避できる。

その意義は以下の点にある：

• 汎用性： 非コンパクトで双定値計量を持たないものを含む、幅広いクラスのリー群に適用可能であり、これらは従来の拡散法や最適化法にとって困難であった。
• 頑健性： 変換された入力を訓練分布へと効果的に復元し、学習済みネットワークの分布外の幾何学的変換に対する汎化性能を大幅に向上させる。
• 効率性： 拡散ベースのアプローチと自明化されたスコア恒等式の組み合わせにより、従来の MCMC や険しいエネルギー地形における反復最適化と比較して、より効率的で安定したサンプリングプロセスを提供する。

本論文は、このアプローチが**等変性のための推論時スケーリング（inference-time scaling for equivariance）**の一形態であることを結論付けており、分布外の変換への対処が極めて重要となる大規模基盤モデルの時代において、実用的なソリューションを提示している。