Joint inference of weak lensing convergence map and cosmology with diffusion models

技術要約：拡散モデルを用いた弱重力レンズ収束マップと宇宙論の同時推論

問題提起
現在の弱重力レンズ解析は、通常、要約統計量（例：2点相関関数）に依存しているが、これらは物質分布に存在する非ガウス的な情報、特に小スケールの非線形スケールにおける情報を破棄してしまう。全フィールド・アプローチは、すべての統計的情報を捉えるための原理的な方法を提供しているが、既存の手法には以下の二極化が存在する：

1. 周辺推論（Marginal Inference）： ニューラルネットワークを用いた手法は、観測データから宇宙論パラメータ（$\theta$）を直接推論するが、基礎となる収束場（$\kappa$）については周辺化してしまい、質量マップの再構成には至らない。
2. 場（フィールド）の再構成： マス・マッピングのような手法は、ノイズを含む剪断（$\gamma_{obs}$）から収束場 $\kappa$ を再構成するが、宇宙論パラメータを固定されたものとして扱い、同時には推論しない。
3. 明示的推論（MCMC）： BORGのような手法は、ジョイント推論（$p(\theta, \kappa | \gamma_{obs})$）を行うが、これには初期条件から観測量へとつながる微分可能なフォワードモデルが必要であり、複雑な物理モデルの使用を制限し、高次元空間における計算コストも高い。

著者らは、微分可能なフォワードモデルを必要とせず、かつ反復的なMCMCサンプリングを行わずに、収束場と宇宙論パラメータを同時に再構成および推論するための**同時推論（joint inference）**手法の開発を目指している。

手法：JADE
著者らは、条件付き拡散モデルに基づく暗黙的推論フレームワークであるJADE（Joint Architecture for fielD and cosmological parameter Estimation）を導入する。

• 拡散モデルによる暗黙的推論： 尤度関数を学習する代わりに、JADEはシミュレーションから直接、結合事後分布 $p(\theta, \kappa | \gamma_{obs})$ を学習する。この手法は、ベースとなる分布（ガウスノイズ）をターゲットとなる分布へと輸送する速度場をニューラルネットワークに学習させる**フロー・マッチング（Flow Matching）**を利用している。モデルは、ノイズを含む補間体（$x_t$）から、観測された剪断場を与えられた状態で、「クリーンな」サンプル（$x_1 = [\theta, \kappa]$）を予測するように訓練される。
• Transformerアーキテクチャ： JADEの中核は、マルチモーダルなデータに適応させたピクセル空間のVision Transformerである。
  • 統一トークンシーケンス： 収束場 $\kappa$ はパッチに分解され、宇宙論パラメータ $\theta$ は同じ次元の埋め込みへと投影される。これらは単一のトークンシーケンスへと結合される。
  • 条件付け（Conditioning）： ノイズを含む観測された剪能場 $\gamma_{obs}$ もパッチ化および投影され、シーケンスに結合される。Transformerのアテンション機構により、ターゲットとなるトークン（場とパラメータ）は、すべての空間位置において条件付け入力（剪断）に注目（attend）することができる。
  • 出力： ネットワークは、デノイズされた推定値として、収束マップと宇宙論パラメータの両方を同時に出力する。
• アモルタイズ推論（Amortized Inference）： 学習後は、特定の $\gamma_{obs}$ に条件付けられた常微分方程式（ODE）を解くことで、あらゆる新しい観測に対して事後分布のサンプルを生成できる。このプロセスは「アモルタイズ（償却）」されており、これは、新たな再学習や観測ごとの最適化（MCMCのような）を必要としないことを意味する。

シミュレーション設定
本手法は、シミュレーションされた弱重力レンズデータを用いて検証された：

• 宇宙論： 6つのパラメータ（$\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$）を持つ $w$CDM モデル。
• データ生成： 大規模な学習データセット（100,000サンプル）を高速に生成するため、フルN体シミュレーションではなく、2D対数正規分布（log-normal）モデル（非線形成長の現象論的モデル）を用いて収束マップを生成した。
• 観測量： 形状ノイズを含む、LSST Year 10 (Y10) の深さの剪断場。$5 \times 5$ deg$^2$ の領域を $128 \times 128$ ピクセルでカバーしている。

主な結果

1. 同時再構成： モデルは、ノイズを含む剪断を条件として、$\kappa$ と $\theta$ の両方の事後分布サンプルを正常に生成した。視覚的確認および統計的テストにより、再構成された収束マップは真の構造を回復しつつ、小スケールにおける残留不確実性を正しく反映していることが確認された。
2. 統計的一貫性：
   • パワースペクトル： 事後サンプルの自己および相互パワースペクトルは、サンプリングされた宇宙論において独立に描かれたシミュレーションと一致しており、ビンおよびスケールにわたる平均相対誤差は約6.7%であった。
   • 1点PDF： 回復された分布は期待通りの対数正規形状を示し、非ガウス統計を捉えていることが確認された。
   • 残差： 簡約カイ二乗統計量（$\chi^2_{red} \approx 1.002$）は、観測と事後サンプル間の残差が、予想通り注入された形状ノイズによって支配されていることを裏付けた。
3. 宇宙論的制約： JADEから得られた周辺宇宙論事後分布（$p(\theta | \gamma_{obs})$）は、信頼領域および縮退の方向において、従来のNUTS（No-U-Turn Sampler）MCMCチェーンと極めて良好な一致を示した。
4. 較正（Calibration）：
   • MIRAスコア： ジョイント事後分布は $0.635 \pm 0.017$ のMIRA（Mass In Random Areas）較正スコアを達成した（最適値は0.667）。これは、わずかな過剰信頼を伴うものの、良好な較正であることを示している。周辺宇宙論事後分布のスコアは $0.659 \pm 0.0139$ であった。
   • TARP： ランダムな点を用いた精度テスト（Tests of Accuracy with Random Points）により、信頼領域がパラメータ空間全体にわたって正しい頻度論的な被覆（coverage）を達成していることが確認された。
5. 効率性： 単一のジョイント事後サンプルを生成するのに、約 0.2秒（256 ODEステップ）しかかからない。1,000個のサンプルを抽出するには約3分を要し、これはベンチマークとなったMCMCチェーン（約218万回の勾配評価を必要とした）と比較して大幅に低い計算コストである。

意義と主張
本論文は、拡散ベースの暗黙的推論が、フルフィールドレベルの宇宙論解析への実行可能な経路であることを確立したと主張している。主な意義は、質量マップの再構成と宇宙論パラメータの推定を、単一の非微分的なフレームワーク内で統合した点にある。

手法の利点に関する主な主張は以下の通りである：

• 微分可能なフォワードモデルを必要としない： フィールドレベルの推論を行うMCMCとは異なり、JADEはフォワードモデルの解析的な勾配を必要としない。そのため、一度学習データさえあれば、理論上、任意の複雑なシシミュレータ（例：フルN体や流体シミュレーション）を使用することが可能である。
• アモルタイゼーション： モデルは一度訓練すれば、あらゆる新しい観測に対して即座に適用でき、観測ごとのサンプリングを回避できる。
• 同時推論： 単一の条件付き拡散モデルを用いて、全フィールドレベルの事後分布 $p(\theta, \kappa | \gamma_{obs})$ を同時に推論する初めての手法である。

限界と今後の課題
著者らは、現在のデモンストレーションが、N体シミュレーションよりも高速ではあるが物理的に完全ではない対数正規シミュレーションに依存していることを指摘している。より現実的で高コストなシミュレータへのスケールアップには、宇宙論パラメータ空間における学習データの密度に関する課題がある。また、アモルタイズ手法であるため、データモデル（例：サーベイマスク、ノイズモデル）が変更された場合には、JADEを再訓練する必要がある（これは、推論時に適応できる明示的推論手法とは異なる）。著者らは、今後の課題として、初期条件と宇宙論パラメータの同時推論へのこのアーキテクチャの拡張、および、より高解像度の場をより効率的に扱うための潜在的拡散（latent diffusion）アプローチの探索を挙げている。