Interpretable Neural Marked Statistics for Cosmological Inference

技术摘要：用于对比宇宙学推断的可解释等变标记

问题陈述
在超越功率谱（power spectrum）的范畴内恢复宇宙学信息，是即将到来的大规模结构（LSS）巡天（如 SPHEREx、DESI、Euclid）的核心目标。虽然功率谱对于高斯场是最优的，但晚期物质密度表现出非高斯信号，这是仅靠二点统计量无法捕捉的。传统的提取方法面临显著障碍：

• 高阶相关函数： 在构型空间中，沿着 $n$ 点相关函数的阶梯上升计算成本极高，且微扰理论在非线性机制下会失效。
• 场级推断： 基于模拟的推断和场级神经摘要绕过了显式的统计量，但往往以牺牲可解释性为代价，且需要高保真度的模拟。
• 经典标记统计（Classical marked statistics）： 现有的方法通过空间标记 $M(x)$ 对密度场进行重加权，从而将高阶相关性折叠进二点统计层面。然而，这些标记通常是手工设计的，具有固定的函数形式（例如平滑密度响应），或者仅针对单一的标度模型进行优化，限制了其在参数空间中的适应性和可解释性。

方法论
作者提出了一个框架，用学习到的、与宇宙学无关且具有可解释性的神经标记方案取代了手工设计的标记。该框架由两个主要部分组成：

1. 可解释的学习标记架构：

   • 标记 $M(x)$ 被构建为一个从物质密度场 $\delta(x)$ 导出的平移等变标量响应。
   • 球面谐波滤波： 密度场在傅里叶空间中通过具有球面谐波感知滤波器（$\ell \in \{0, 1, 2\}$）的通道进行分解。这些滤波器产生旋转不变的标量通道：
     • $\ell=0$：有符号单极响应（$E_0$）。
     • $\ell=1$：具有不变振幅的类矢量响应（$E_1$）。
     • $\ell=2$：源自无迹对称张量的两个标量不变量（用于区分形状性质的 $E_2$ 振幅强度和 $I_3$ 纵向/扁平形状区分）。
   • 处理： 这些不变通道由独立的感知机（MLP）进行处理，并通过加法结合一个交叉相互作用项。输出通过 softplus 激活函数以确保正定性。
   • 初始化： 网络初始化使得 $M(x) \approx 1$，这意味着训练从未标记场开始。这使得学习到的标记可以分解为不同的形态贡献。

2. 对比学习目标：

   • 该框架采用对比学习目标（InfoNCE），将标记后的摘要与宇宙学参数 $\theta$ 在共享潜空间中进行对齐。
   • 摘要向量： 摘要包括未标记场的功率谱（$P_{\delta\delta}$）、标记场的功率谱（$P_{\Delta\Delta}$），以及交叉谱（$P_{\delta\Delta}$）和相关系数。
   • 残差化（Residualization）： 为了确保标记学习的是互补信息，而非仅仅复现功率谱，标记后的摘要嵌入（$z_M$）会对未标记的嵌入（$z_\delta$）进行残差化处理，从而产生一个正交分量 $z_\perp$。
   • 损失函数： 损失函数通过学习到的马哈拉诺比斯（Mahalanobis）度量将 $z_\perp$ 与参数嵌入 $z_\theta$ 对齐，同时将其与从先验中采样的负样本（真实的、全局合成的、局部合成的负样本）分离。
   • 训练策略： 未标记分支经过预训练并被冻结；随后使用对比损失对标记网络和参数嵌入器进行微调。

核心贡献

• 标记统计的泛化： 本文引入了一种神经架构，它泛化了标记统计的概念，超越了固定的函数形式，转而使用一种在构型空间中仍可被检查的、学习到的变换。
• 通过等变性实现的可解释性： 通过强制执行局部 $SO(3)$-等变性并使用球面谐波滤波器，该模型提取了具有明确物理意义的特征（单极、偶极、四极响应）。
• 对比对齐： 使用带有残差化的对比目标，确保学习到的标记专门针对与标准功率谱正交的信息，从而打破了传统统计量无法解决的简并问题。

结果
该方法在 Quijote N-body 模拟集（$z=0$）上进行了评估，涉及 5 个变化的宇宙学参数（$\Omega_m, \Omega_b, h, n_s, \sigma_8$）。

• 费舍尔信息约束： 在 $k_{\max} = 0.2 \, h\text{Mpc}^{-1}$ 时，学习到的标记相比经典标记，使 $\sigma_8$ 的边际约束收紧了 2.9 倍，使 $\Omega_m$ 收紧了 1.8 倍。至关重要的是，它打破了经典的 $\Omega_m - \sigma_8$ 简并，旋转了联合后验轮廓，从而实现了对这两个参数的同时改进。
• 泛化能力（MSE）： 在拉丁超立方体测试集上，学习到的标记比最佳经典标记降低了 1.45 倍 的参数均方误差（MSE）。它展示了鲁棒的泛化能力，特别是在 $\sigma_8$ 方面，尽管在特定的 $(p, \delta_s)$ 配置下，经典标记在 $\Omega_m$ 方面的稳定性略高。
• 潜几何（Latent Geometry）： 对学习到的嵌入进行主成分分析（PCA）显示，前两个主成分几乎完美地与 $\sigma_8$ 和 $\Omega_m$ 的方向对齐，证实了对比目标成功恢复了宇宙学信息的主要轴线。
• 内省分析： 对训练后的标记进行分析表明，各向同性（$\ell=0$）通道在响应中占据主导地位，有效地恢复了一个非线性重加权的密度，而各向异性通道则在结构边界处提供次要修正。

意义
本文声称，这种方法为更强大且可解释的摘要统计量开启了大门。通过用具有物理动机的学习架构取代手工设计的变换，该方法弥合了高约束能力的场级神经网络与传统统计量的可解释性之间的鸿沟。通过在构型空间中检查学习到的标记，研究人员可以理解哪些局部形态特征（例如空洞、纤维状结构、各向异性结构）携带了最丰富的宇宙学信号。作者将此定位为灵活推断工具迈出的一步，旨在降低模拟成本并提高对银河系巡天的科学回报，并具有扩展到其他参数（如原初非高斯性）的潜力。