[論文レビュー] Cosmo3DFlow: Wavelet Flow Matching for Spatial-to-Spectral Compression in Reconstructing the Early Universe

Md Khairul Islam, Zeyu Xia+4 more|arXiv (Cornell University)|Feb 10, 2026

Galaxies: Formation, Evolution, Phenomena被引用数 0

ひとこと要約

Cosmo3DFlowは3D離散ウェーブレット変換とフローメッチングを組み合わせて、進化した状態から初期宇宙条件を再構成し、拡散ベース手法と比較して同等かそれ以上の忠実度を維持しつつ最大で50倍速いサンプリングを達成します。

ABSTRACT

Reconstructing the early Universe from the evolved present-day Universe is a challenging and computationally demanding problem in modern astrophysics. We devise a novel generative framework, Cosmo3DFlow, designed to address dimensionality and sparsity, the critical bottlenecks inherent in current state-of-the-art methods for cosmological inference. By integrating 3D Discrete Wavelet Transform (DWT) with flow matching, we effectively represent high-dimensional cosmological structures. The Wavelet Transform addresses the ``void problem'' by translating spatial emptiness into spectral sparsity. It decouples high-frequency details from low-frequency structures through spatial compression, and wavelet-space velocity fields facilitate stable ordinary differential equation (ODE) solvers with large step sizes. Using large-scale cosmological $N$-body simulations, at $128^3$ resolution, we achieve up to $50 imes$ faster sampling than diffusion models, combining a $10 imes$ reduction in integration steps with lower per-step computational cost from wavelet compression. Our results enable initial conditions to be sampled in seconds, compared to minutes for previous methods.

研究の動機と目的

• 進化した宇宙データからの初期条件推定の高次元性と希少性に対処する。
• スペクトルの疎性と多尺度構造を活用するウェーブレット空間で動作する生成フレームワークを開発する。
• 拡散ベースのベースラインと比較してサンプリング速度を改善し、再構成忠実度を維持または向上させる。

提案手法

• 3D Discrete Wavelet Transform (DWT) をフローメッチングと統合し、高次元宇宙場をモデル化する。
• ウェーブレット空間でフローメッチングを定式化し、ウェーブレット変換観測に条件付けられた初期条件へガウスノイズを輸送する。
• スケール認識型の条件付けとクロススケールスキップ接続を用いて多重解像ウェーブレット係数を活用する。
• 訓練中に物理的に一貫した密度統計を強制するためのオプションのパワースペクトル正規化。
• 時間埋め込みとウェーブレット空間条件付けを持つ3D U-Net速度ネットワークを用いて速度推定を行う。
• 推論はウェーブレット空間での決定論的ODE積分を介して行い、その後逆DWTで再構成場を得る。

Figure 1. Schematic comparison of data representations for the Cosmic Web. Left: The voxel grid allocates memory uniformly, forcing high-resolution processing on empty void regions. Right: The wavelet representation adapts to structure; voids are represented by a few coarse coefficients (large block

実験結果

リサーチクエスチョン

• RQ1 ウェーブレット領域のフローメッチングは、高次元の宇宙初期条件推定においてサンプリングを加速できるか、再構成品質を犠牲にしないか。
• RQ2 ウェーブレット空間モデリングはボクセル空間の拡散法と比較して宇宙密度場の疎性と多尺度構造をより適切に捉えられるか。
• RQ3 スケール認識型条件付けとクロススケール接続は再構成忠実度と物理的一貫性（例：パワースペクトル）にどのような影響を与えるか。

主な発見

Dataset	Resolution	VRMSE_Ours	VRMSE_Diffusion	CrossCorr_Ours	CrossCorr_Diffusion	PkR2_Ours	PkR2_Diffusion	TF_Ours	TF_Diffusion
Standard Latin Hypercube	128^3	0.50	0.63	0.88	0.82	0.99	0.70	0.99	0.80
Big Sobol Sequence	128^3	0.62	0.64	0.80	0.79	0.99	0.84	0.95	0.88
Non-Gaussian Local LH	128^3	0.56	0.59	0.86	0.83	1.00	1.00	0.98	0.98

• Cosmo3DFlowは128^3の10倍の統合ステップを組み合わせ、ウェーブレット圧縮による約5倍のステップごとの高速化を実現して、拡散ベースのベースラインより最大50倍速いサンプリングを実現した。
• 微細な宇宙構造を保持し、VRMSE、相互相関、パワースペクトル指標で拡散法と比較して再構成誤差を低減。
• Standard Latin Hypercube、Big Sobol、非ガウス分布データセット全体で、Cosmo3DFlowは忠実度を維持または向上させつつ物理的一貫性のある統計（パワースペクトル）を維持。
• ウェーブレット空間の速度モデリングは情報を粗いLLL帯と希薄なディテールへ集中させ、活動ボリュームの約8倍の空間圧縮を達成する効率的な表現を実現。
• スケール特異的条件付けとクロススケールスキップは、多尺度ウェーブレット構造を活用して再構成品質をさらに改善する。

Figure 2. Cosmo3DFlow targets a high-dimensional inverse problem, while achieving 50 $\times$ faster sampling. Given observations, it reconstructs samples in 5 seconds vs. $\approx$ 240 seconds for diffusion models, while preserving fine-scale structure through wavelet sparsity. Based on variance-no