Cross-Fitting for Neural Posterior Estimation¶

会议: NeurIPS 2026
arXiv: 2610.26733
代码: https://github.com/jeff-regier/cf-npe
领域: 科学计算
关键词: 神经网络后验估计, 模拟推断, 交叉拟合, 样本分割, 后验校准

一句话总结¶

针对神经后验估计（NPE）联合训练均值与方差时样本内残差趋零所导致的后验严重过度自信问题，本文提出基于样本分割的交叉拟合神经后验估计（CF-NPE），将后验均值回归与样本外残差条件密度估计两阶段解耦，从根本上消除了均值欠拟合与方差过拟合的权衡，在多项高维任务中实现了近乎完美的后验校准与显著更高的泛化对数似然。

研究背景与动机¶

基于模拟的推断（Simulation-Based Inference, SBI）在现代天文学、物理学与生命科学等涉及复杂黑盒模拟器的前沿领域扮演着核心角色。其中，神经后验估计（Neural Posterior Estimation, NPE）通过在先验与模拟器联合生成的全数据样本上拟合摊销条件密度估计器（如高斯编码器或归一化流），在测试时实现快速、无需重复运行模拟器的后验推断。然而，NPE 的最大似然训练目标对在训练点上分配任意高密度的参数配置存在病态偏好，而典型的条件密度估计器具备极高的函数拟合容量。这使得多轮训练（multi-epoch）下的 NPE 极易陷入严重的过度自信（extreme overconfidence），其学习到的后验分布在测试数据上表现出灾难性的方差塌缩。

常规深度回归任务中的过度拟合通常被证明是良性的（benign overfitting），隐式正则化能确保泛化误差平稳；但 NPE 的过度拟合是灾难性的——网络容量甚至可能将预测标准差推向浮点数下限。目前从业者普遍依赖验证集对数似然的早停（early stopping）或权重衰减来避免这种发散。然而，由于均值与尺度（方差）是在同一批训练样本上联合参数化并沿同一优化轨迹同步更新的，均值网络的拟合必然会降低样本内残差。优化器轨迹因此强行绑定了均值与方差：当早停阻止方差继续塌缩时，后验均值尚未达到其最佳估计精度；若继续训练均值，方差则迅速过拟合。这种由样本内残差引起的均值欠拟合与方差过拟合的权衡，使得在现有轨迹内做一维早停搜索无法找到真正优质的后验近似。

面对这一困境，本文没有尝试去修改目标函数的启发式权重或梯度流，而是从统计学样本分割的经典思想中汲取灵感。本文的核心 idea 是：提出交叉拟合神经后验估计（CF-NPE），利用 K 折交叉拟合将后验均值估计与高阶矩条件密度估计解耦为两个独立阶段，仅在样本外（out-of-fold）残差上训练第二阶段密度估计器，从而彻底消除两阶段间的样本依赖，使训练集残差的期望二阶矩与测试集严格一致。

方法详解¶

整体框架¶

CF-NPE 的整体流程由两阶段构成。第一阶段将规模为 \(N\) 的模拟数据集均匀划分为 \(K\) 个互斥的折（folds），针对每一个折 \(k\)，使用其余 \(K-1\) 折数据训练后验均值回归网络（可选结合 \(M\) 个随机初始化的模型做平均集成）。第二阶段在全数据集上收集每个样本对应的样本外残差，并以输入特征和预测均值为条件，拟合一个以零为中心的条件残差密度估计器（如对角高斯或神经样条流）。推理阶段将测试样本送入均值回归网络得到中心预测，再叠加条件残差分布完成完整后验推断。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["模拟数据集 D = {(x_i, z_i)}<br/>均匀划分 K 折"] --> B["K 折样本分割与均值回归<br/>训练 K 组独立或集成回归器"]
    B --> C["样本外残差构建<br/>生成无偏样本外残差元组"]
    C --> D["条件残差密度估计<br/>以均值为条件拟合以零为中心的分布"]
    D --> E["测试期解耦后验组合<br/>均值回归叠加残差密度生成后验"]

关键设计¶

1. K 折样本分割与均值回归：解耦后验一阶矩并释放良性过拟合潜力
传统 NPE 在同一个网络中联合预测后验均值 \(\mu_\phi(x)\) 与尺度 \(\sigma_\phi(x)\)，导致两者的优化步调发生恶性耦合。CF-NPE 在第一阶段将包含 \(N\) 个模拟样本的数据集 \(\mathcal{D} = \{(x_i, z_i)\}_{i=1}^N\) 均匀随机划分为 \(K\) 个等大子集 \(\mathcal{D}_1, \dots, \mathcal{D}_K\)。对每个折 \(k \in \{1, \dots, K\}\)，算法仅使用排除当前折的数据 \(\mathcal{D}_{-k} = \mathcal{D} \setminus \mathcal{D}_k\) 训练后验均值回归器 \(\hat{\mu}_{-k}\)。为了提高均值预测的数值稳定性并平滑优化随机性，每个折可独立训练 \(M \ge 1\) 个不同随机初始化的网络并取其平均。第一阶段完全采用均方误差（MSE）损失函数进行训练：

\[\mathcal{L}_{\text{MSE}} = \frac{1}{|\mathcal{D}_{-k}|} \sum_{i \in \mathcal{D}_{-k}} \| z_i - \hat{\mu}_{-k}(x_i) \|^2\]

因为最小化 MSE 在总体极限下严格收敛于条件期望 \(\mathbb{E}[z|x]\)，此阶段的早停仅服务于均值预测精度的最大化，不受任何方差塌缩的干扰，从而充分受益于神经网络在纯回归任务中的良性过拟合特性。

2. 样本外残差构建：消除样本内乐观偏差并保证二阶矩无偏
标准 NPE 产生过度自信的本质在于样本内残差低估了泛化残差误差。论文从理论上严格证明了样本内预测的乐观性定理：对于任意预测均值 \(c_i\)，测试样本与训练样本的期望平方残差之差严格等于两倍的协方差项：

\[\mathbb{E}[(z_i' - c_i)^2 \mid x_{1:N}] - \mathbb{E}[(z_i - c_i)^2 \mid x_{1:N}] = 2\,\text{Cov}(z_i, c_i \mid x_{1:N})\]

在标准 NPE 中，由于 \(c_i\) 由全量数据拟合，该协方差项严格大于零且随训练轮次增加而急剧发散。CF-NPE 通过交叉拟合设计，对每个样本 \(i\) 采用其所在折 \(k(i)\) 之外训练的均值网络来计算残差 \(r_i = z_i - \hat{\mu}_{-k(i)}(x_i)\)。由于 \(\hat{\mu}_{-k(i)}\) 在训练时完全未见过 \(z_i\)，条件协方差项被恒等置零：\(\text{Cov}(z_i, \hat{\mu}_{-k(i)}(x_i) \mid x_{1:N}) = 0\)。此时样本内期望平方残差与样本外完全吻合，且其期望严格等于真实后验方差与均值估计误差平方之和：

\[\mathbb{E}[(z_\ell - c_\ell)^2 \mid x] = \Sigma^\star_{\ell\ell}(x) + (c_\ell - \mu^\star_\ell(x))^2\]

随着第一阶段均值网络逐步逼近真实均值 \(\mu^\star(x)\)，第二阶段所针对的目标二阶矩精确收敛于真实的条件后验方差 \(\Sigma^\star(x)\)。

3. 条件残差密度估计：在中心化残差上拟合高阶矩与自动校准
第二阶段利用构建的三元组集合 \(\mathcal{T} = \{(x_i, \hat{\mu}_{-k(i)}(x_i), r_i)\}_{i=1}^N\) 拟合条件残差密度估计器 \(q_\psi(r \mid x, m)\)。其训练目标是最大化残差的对数似然：

\[\mathcal{L}(\psi) = \frac{1}{N} \sum_{i=1}^N \log q_\psi\left(r_i \mid x_i, \hat{\mu}_{-k(i)}(x_i)\right)\]

与原始响应变量 \(z_i\) 相比，样本外残差 \(r_i\) 具有一个至关重要的结构优势：其均值在先验和结构上天然严格以零为中心。当采用对角高斯编码器并显式构造零均值高斯分布 \(q_\psi(r \mid x, m) = \mathcal{N}(r; 0, \text{diag}(V_\psi(x, m)))\) 时，论文给出了极具分量的理论保证：在训练集对数似然极大值点处，尺度参数满足一阶驻点条件，必然导致约化卡方统计量严格等于 1（\(\chi_d^2(q; P_{\text{train}}) = 1\)）。再结合样本外协方差置零的性质，训练集的校准性直接无损迁移至测试集。当采用神经样条流（Neural Spline Flow, NSF）时，对称初始化与零均值基分布使得均值在整个优化路径上始终保持在最优值附近，早停机制能纯粹根据高阶矩的拟合质量决定停止点。

4. 测试期解耦后验组合：单折适配与多重混合扩展
完成两阶段独立训练后，CF-NPE 组合输出完整的后验分布。基础版 CF-NPE 选取任意一折（例如第 1 折）的均值网络与训练好的条件残差密度网络拼接：\(q_\phi(z \mid x) = q_\psi(z - \hat{\mu}_{-1}(x) \mid x, \hat{\mu}_{-1}(x))\)。进阶版 CF-NPE-mixture 则在第二阶段使用不同的随机初始化和验证折独立训练 \(K\) 个残差密度网络 \(\psi_1, \dots, \psi_K\)，并在推理时构建等权混合后验：

\[q_\phi(z \mid x) = \frac{1}{K} \sum_{k=1}^K q_{\psi_k}\left(z - \hat{\mu}_{-k}(x) \mid x, \hat{\mu}_{-k}(x)\right)\]

这一混合方式在提升后验表达能力的同时，与标准 NPE 简单集成导致的过度弥散（方差膨胀、欠自信）截然不同，它能始终保持校准统计量稳定在 1.0 附近。

一个完整示例¶

以合成的高维高斯-高斯问题（\(d = 200, N = 10,000\)）为例，设划分折数 \(K = 10\)： 1. 数据切分：将 10,000 对观测样本均匀分成 10 份，每份 1,000 个样本。 2. 均值网络训练：对折 \(k = 1\)，取出其余 9,000 个样本，用 1 层隐藏层、宽度 256、学习率 \(2.38 \times 10^{-3}\) 的 MLP 最小化 MSE。早停点选在验证集 MSE 最小处（测试 MSE 达到 0.224，逼近贝叶斯极限 0.200）。对其余 9 折重复此操作。 3. 计算残差：将第 1 折的 1,000 个测试样本输入 \(\hat{\mu}_{-1}\)，得到 1,000 个残差向量 \(r_i = z_i - \hat{\mu}_{-1}(x_i)\)；汇集 10 个折得到 10,000 个样本外残差。 4. 残差流拟合：以 \((x_i, \hat{\mu}_{-k(i)}(x_i))\) 作为调节器输入，用神经样条流最大化 \(r_i\) 的似然。由于残差均值天然在 0 附近，NSF 集中精力拟合方差，其约化卡方统计量 \(\chi_d^2\) 稳定在 1.000。 5. 推理输出：对于新的观测 \(x_{\text{test}}\)，直接输出 \(q(z \mid x_{\text{test}}) = q_{\text{NSF}}(z - \hat{\mu}_{-1}(x_{\text{test}}) \mid x_{\text{test}}, \hat{\mu}_{-1}(x_{\text{test}}))\)，测试对数似然高达 \(-133.92\) nats，远超标准 NPE 的 \(-169.66\) nats。

实验关键数据¶

主实验¶

论文在三个合成基准问题（高维高斯-高斯问题 \(d=200\)、病态漏斗问题 \(d=200\)、高斯混合问题 \(d=20\)）以及宇宙学强引力透镜源重建实际任务（\(d=256, 16\times 16\) 像素网格）上进行了广泛评测，主要指标包括测试集对数似然（Test LL，nats，越高越好）、约化卡方后验校准统计量（\(\chi_d^2\)，理想值为 1.0）以及后验均值均方误差（MSE，越低越好）。

数据集 / 问题	后验族	指标	标准 NPE (Std-NPE)	CF-NPE (M=1)	CF-NPE (M=3)	真实后验 / 理论极限
Gaussian–Gaussian (\(d=200\))	Gaussian	Test LL ↑	\(-141.30 \pm 0.11\)	\(-133.92 \pm 0.13\)	\(-128.42 \pm 0.10\)	\(-122.84\)
Gaussian–Gaussian (\(d=200\))	Gaussian	\(\chi_d^2\) (→1)	\(1.095 \pm 0.001\)	\(0.994 \pm 0.002\)	\(0.993 \pm 0.002\)	\(1.000\)
Gaussian–Gaussian (\(d=200\))	Gaussian	MSE ↓	\(0.234\)	\(0.224\)	\(0.212\)	\(0.200\)
Gaussian–Gaussian (\(d=200\))	NSF	Test LL ↑	\(-169.66 \pm 0.54\)	\(-133.92 \pm 0.13\)	\(-128.31 \pm 0.10\)	\(-122.84\)
Gaussian–Gaussian (\(d=200\))	NSF	\(\chi_d^2\) (→1)	\(0.979 \pm 0.004\)	\(1.000 \pm 0.001\)	\(1.000 \pm 0.001\)	\(1.000\)
Gaussian–Gaussian (\(d=200\))	NSF	MSE ↓	\(0.306\)	\(0.224\)	\(0.212\)	\(0.200\)
Funnel Problem (\(d=200\))	NSF	Test LL ↑	\(-58.98 \pm 5.84\)	\(38.09 \pm 0.67\)	\(83.82 \pm 0.73\)	\(178.68\)
Funnel Problem (\(d=200\))	NSF	\(\chi_d^2\) (→1)	\(0.995 \pm 0.011\)	\(1.068 \pm 0.005\)	\(1.107 \pm 0.006\)	\(1.000\)
Funnel Problem (\(d=200\))	NSF	MSE ↓	\(0.242\)	\(0.045\)	\(0.031\)	\(0.017\)
Mixture Problem (\(d=20\))	NSF	Test LL ↑	\(-3.02 \pm 0.03\)	\(-0.81 \pm 0.02\)	\(-0.63 \pm 0.03\)	\(0.07\)
Mixture Problem (\(d=20\))	NSF	\(\chi_d^2\) (→1)	\(0.815 \pm 0.004\)	\(1.001 \pm 0.005\)	\(1.008 \pm 0.010\)	\(1.000\)
Mixture Problem (\(d=20\))	NSF	MSE ↓	\(0.104\)	\(0.092\)	\(0.090\)	\(0.087\)
强引力透镜重建 (\(d=256\))	NSF	Test LL [IQR] ↑	\(2.59\,[-0.67, 3.40]\)	\(60.05\,[58.82, 63.03]\)	\(83.21\,[80.27, 84.61]\)	—
强引力透镜重建 (\(d=256\))	NSF	\(\chi_d^2\) (→1)	\(1.196 \pm 0.005\)	\(1.025 \pm 0.021\)	\(0.985 \pm 0.020\)	\(1.000\)
强引力透镜重建 (\(d=256\))	NSF	MSE ↓	\(0.124 \pm 0.001\)	\(0.050 \pm 0.001\)	\(0.043 \pm 0.000\)	—

消融实验¶

在 Gaussian–Gaussian 基准（\(d=200, N=10,000\)）上，作者针对交叉拟合核心机制、折数 \(K\) 以及已有替代修复策略进行了系统消融（均采用 \(M=3\)）。

配置 / 变体	后验族	Test LL ↑	\(\chi_d^2\) (→1)	MSE ↓	说明与结论
CF-NPE (完整模型, K=10)	Gaussian	\(-128.42 \pm 0.10\)	\(0.993 \pm 0.002\)	\(0.212\)	样本外残差 + 交叉拟合，基准完整配置
样本内残差 (in-fold residuals)	Gaussian	\(-129.45 \pm 0.12\)	\(1.150 \pm 0.002\)	\(0.212\)	均值不变，方差使用样本内残差，后验显著过度自信
单次均值拟合 (single mean fit)	Gaussian	\(-136.01 \pm 0.14\)	\(1.094 \pm 0.009\)	\(0.227\)	无交叉拟合（单折80/20划分），均值误差与过自信双双恶化
温度缩放 (Std-NPE + Temp Scaling)	Gaussian	\(-140.88 \pm 0.11\)	\(1.002 \pm 0.001\)	\(0.234\)	能校正卡方统计量，但无法改进均值，对数似然几乎无提升
替代损失 (Std-NPE + \(\beta\)-NLL)	Gaussian	\(-140.41 \pm 0.11\)	\(1.082 \pm 0.001\)	\(0.233\)	加权损失仍作用于同批样本，仍严重过度自信
截断梯度 (Std-NPE + stop-gradient)	Gaussian	\(-141.09 \pm 0.12\)	\(1.090 \pm 0.001\)	\(0.235\)	改变梯度路径无法消除样本内残差趋零的物理根源
折数 \(K=2\)	Gaussian	\(-132.93 \pm 0.11\)	—	—	最小折数依然显著击败 Std-NPE (\(-141.30\))
折数 \(K=4\)	Gaussian	\(-129.50 \pm 0.10\)	—	—	性能随有效训练样本量提升迅速上升
折数 \(K=16\)	Gaussian	\(-128.18 \pm 0.11\)	—	—	收益逐步饱和，K=10 具备极佳的性价比
CF-NPE (完整模型, K=10)	NSF	\(-128.31 \pm 0.10\)	\(1.000 \pm 0.001\)	\(0.212\)	神经样条流实现完全校准
样本内残差 (in-fold residuals)	NSF	\(-129.55 \pm 0.10\)	\(1.164 \pm 0.001\)	\(0.212\)	破坏样本外性质导致卡方大幅偏离至 1.164

关键发现¶

样本分割是校准的核心决定因素：单纯在样本内残差上改动损失函数（\(\beta\)-NLL、截断梯度）或事后温度缩放均无法闭合与 CF-NPE 的性能差距。消融实验中一旦改用样本内残差（in-fold residuals），后验均值虽然完全相同，但方差迅速发生欠弥散，卡方统计量从 0.993 恶化至 1.150。
样本效率产生质的飞跃：在数据量扩展实验中，CF-NPE（\(M=3\)）在 \(N=10,000\) 样本预算下的测试集对数似然（\(-128.42\)）甚至大幅超越了标准 NPE 在 \(N=30,000\) 模拟样本下的表现（\(-129.37\)），将昂贵的黑盒模拟器调用开销降低了 3 倍以上。
高维空间的优势极度放大：在潜变量维度扫描中，当潜变量维度 \(d=2\) 时，标准 NPE 与 CF-NPE 表现接近无差异；而随着维度升至 \(d=200\)（高斯问题），对数似然增益激增至 \(+12.88\) nats；强引力透镜（\(d=256\)）中增益更达到惊人的 \(+80.6\) nats。这证明该方法在复杂高维科学逆问题中具有压倒性优势。

亮点与洞察¶

解耦两阶段优化动态，破除恶性权衡：传统 NPE 将均值与方差捆绑优化导致均值尚未收敛而方差已经塌缩；CF-NPE 巧妙地用交叉拟合使两阶段优化目标完全独立，让均值网络享受到深度回归特有的良性过拟合红利。
样本外残差的数学性质严谨闭环：通过严格的协方差置零推导（Proposition 3.1 & 3.2），保证了样本外残差二阶矩与真实后验方差的一致性，并在高斯假设下给出了极大似然处 \(\chi_d^2 = 1\) 的闭式校准保证。
高并行度与计算架构友好：在显存未饱和的合成实验中，借助 PyTorch 的 vmap 原语，\(K\) 折与 \(M\) 重复网络可完全并行训练；在引力透镜大任务中，即便串行计算，增加的计算开销也远小于昂贵物理模拟器本身的运行开销。

局限与展望¶

计算时间增加：当单模型训练即可占满单张 GPU 算力时，CF-NPE 需要训练 \(K\) 折甚至 \(K \times M\) 个均值网络，训练耗时较单模型标准 NPE 增加接近一个数量级（例如强引力透镜任务从 7 分钟增至约 55 分钟）。
仅针对一阶与高阶矩的单次解耦：CF-NPE 仅解耦了一阶矩（均值）与后续高阶矩的估计，但在第二阶段中，方差与三阶及以上高阶矩仍然由标准 NPE 联合建模，高阶矩之间可能仍残留微妙的权衡。
当前局限于单轮摊销推断：论文仅探讨了一次性模拟的单轮摊销后验估计，尚未拓展至多轮序列化引导提议分布（Sequential NPE / SNPE / APT）的场景。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将计量经济学与因果推断中的交叉拟合思想引入基于模拟的推断，精准抓住样本内残差塌缩的理论根源]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 3 个理论玩具/病态基准及真实宇宙学强引力透镜物理任务，包含超参数网格调优与全面的统计诊断指标]
写作质量: ⭐⭐⭐⭐⭐ [论证逻辑严密，定理表述清晰，动机直观且实验数据详实]
价值: ⭐⭐⭐⭐⭐ [直接解决了高维科学模拟推断长期存在的后验过度自信难题，极具工程实用与理论指导价值]