[论文解读] Field Level Neural Network Emulator for Cosmological N-body Simulations

Drew Jamieson, Yin Li+4 more|arXiv (Cornell University)|Jun 9, 2022

Astronomy and Astrophysical Research被引用 6

一句话总结

本文提出了一种可微分的、基于场的神经网络模拟器，用于宇宙学N体模拟，能够准确预测不同 $Ω_{\mathrm{m}}$ 宇宙学参数下的非线性结构形成。通过使用基于风格的卷积网络，从线性初始条件建模粒子位移和速度，该方法在 $k \sim 1\ \mathrm{Mpc}^{-1}\,h$ 尺度下实现了亚百分之一的精度，优于COLA方法和基准神经网络，同时实现了快速、可微分的场级推理。

ABSTRACT

We build a field level emulator for cosmic structure formation that is accurate in the nonlinear regime. Our emulator consists of two convolutional neural networks trained to output the nonlinear displacements and velocities of N-body simulation particles based on their linear inputs. Cosmology dependence is encoded in the form of style parameters at each layer of the neural network, enabling the emulator to effectively interpolate the outcomes of structure formation between different flat $Λ$CDM cosmologies over a wide range of background matter densities. The neural network architecture makes the model differentiable by construction, providing a powerful tool for fast field level inference. We test the accuracy of our method by considering several summary statistics, including the density power spectrum with and without redshift space distortions, the displacement power spectrum, the momentum power spectrum, the density bispectrum, halo abundances, and halo profiles with and without redshift space distortions. We compare these statistics from our emulator with the full N-body results, the COLA method, and a fiducial neural network with no cosmological dependence. We find our emulator gives accurate results down to scales of $k \sim 1\ \mathrm{Mpc}^{-1}\, h$, representing a considerable improvement over both COLA and the fiducial neural network. We also demonstrate that our emulator generalizes well to initial conditions containing primordial non-Gaussianity, without the need for any additional style parameters or retraining.

研究动机与目标

• 开发一种快速、可微分且精确的场级宇宙学N体模拟模拟器，以捕捉大尺度结构的完整非线性演化。
• 通过实现全场推理而无需数据压缩，克服传统汇总统计方法因压缩数据而损失信息的局限性。
• 通过风格调制将宇宙学参数依赖性（特别是 $\Omega_{\mathrm{m}}$）直接嵌入神经网络架构，实现在不重新训练的情况下跨宇宙学参数插值。
• 在非线性区域（$k \sim 1\ \mathrm{Mpc}^{-1}\,h$）实现高精度，同时保持计算效率和可微分性，以支持下游推理任务。
• 在无需额外微调或重新训练的情况下，实现对非高斯初始条件的泛化。

提出的方法

• 模拟器使用两个独立的3D U-Net类卷积神经网络，从线性初始密度场预测粒子位移和速度。
• 通过风格调制编码 $\Omega_{\mathrm{m}}$ 的宇宙学依赖性：在每一层学习风格向量调制卷积核权重，实现跨宇宙学参数的插值。
• 位移网络采用组合损失函数进行训练：$\log L_{\delta} + \lambda \log L_{\mathbf{\Psi}}$，其中 $L_{\delta}$ 为粒子密度的MSE（通过CIC赋值实现），$L_{\mathbf{\Psi}}$ 为位移场的MSE，$\lambda = 3$。
• 速度网络采用多尺度损失函数：$\log L_{\mathbf{v}(\mathbf{q})} + \log L_{\mathbf{v}(\mathbf{x})} + \log L_{{\mathbf{v}}^{2}(\mathbf{x})}$，其中最后一项最小化质量加权速度弥散的残差，以正确建模红移空间扭曲。
• 为在GPU内存限制下处理 $512^3$ 的模拟场，模型采用 $128^3$ 的裁剪推理方式，并使用周期性填充以保持平移等变性。
• 通过 map2map 框架实现高效数据加载，优化I/O操作，防止在12 TB数据训练过程中出现GPU饥饿。

实验结果

研究问题

• RQ1 可微分神经网络模拟器是否能在不依赖汇总统计量的情况下，准确预测跨一系列宇宙学参数的非线性大尺度结构形成？
• RQ2 场级模拟器在多大程度上能准确再现关键宇宙学统计量，如功率谱、三阶谱、晕丰度和红移空间扭曲？
• RQ3 该模拟器是否能在不重新训练或增加额外风格参数的情况下，泛化至具有原初非高斯性的初始条件？
• RQ4 与基准神经网络相比，通过风格调制引入宇宙学依赖性在泛化能力和精度方面有多大提升？
• RQ5 该模拟器在非线性区域中保持亚百分之一精度的最小尺度（$k$）是多少？

主要发现

• 该模拟器在密度功率谱和位移功率谱中均实现了 $k \sim 1\ \mathrm{Mpc}^{-1}\,h$ 尺度下的亚百分之一精度，显著优于COLA方法和无宇宙学依赖性的基准神经网络。
• 通过使用二阶速度损失 ($L_{{\mathbf{v}}^{2}}$)，模型准确再现了红移空间扭曲中的“手指状效应”（Fingers of God），该损失项通过最小化质量加权速度弥散的残差来实现。
• 该模拟器在无需任何重新训练或额外风格参数的情况下，对具有原初非高斯性 ($f_{\mathrm{NL}} \neq 0$) 的初始条件表现出良好的泛化能力，表明其对非高斯初始条件具有鲁棒性。
• 可微分架构支持快速场级推理，支持无需似然函数的推理和参数约束，完整保留信息内容，避免了汇总统计量固有的数据压缩损失。
• 该模型在平坦 $\Lambda\mathrm{CDM}$ 宇宙学中广泛 $\Omega_{\mathrm{m}}$ 值范围内保持高精度，通过基于风格的调制实现了有效插值。
• 通过 map2map 框架，12 TB 数据的训练成为可能，该框架优化了数据加载，防止GPU利用率不足，实现了对 $512^3$ 场的高效大规模训练。