Environment-Robust Representation Learning with Empirical Bayes¶

会议: NeurIPS 2026
arXiv: 2606.05365
领域: 自监督学习
关键词: 多环境泛化, 经验贝叶斯, 变分推断, 环境鲁棒表征, 域泛化

一句话总结¶

针对多环境预测中潜变量分布随环境漂移而生成机制稳定的问题，本文提出基于经验贝叶斯的环境鲁棒表征学习方法（EBER），通过显式建模观测条件下的环境权重来聚合环境特异性先验与后验，诱导出无需人为预设惩罚项的环境平衡变分目标，实现了未知测试环境下的最优边缘贝叶斯预测。

研究背景与动机¶

在天文天体分类、人体微生物组疾病筛查以及重症监护室（ICU）脓毒症早期预警等真实科学预测场景中，数据通常采集自多个高度异质性的环境。不同观测天区的未知天体类型比例不同、不同临床队列患者体内潜在微生物群落状态各异、不同医院科室的基础人群健康状态存在显著差异，然而观测测量过程以及潜在物理/生理状态与目标标签之间的生成机制往往保持物理不变性。现有域泛化与不变表征学习方法（如 IRM、V-REx、Fishr）大多采用启发式正则化项来显式惩罚跨环境的损失或梯度方差，试图强行学习消除环境差异的全局不变预测器。

然而，这种传统的惩罚式范式存在本质缺陷：当环境差异源于潜在状态先验分布漂移而非测量混淆时，一味压制环境差异往往会牺牲表征的表达能力，且外部预设的正则化项缺乏概率生成视角的严格解释。另一方面，直接建模环境影响的生成模型往往难以泛化到训练集之外的全新环境，因为在测试阶段新环境的标识完全未知，直接估计环境特异性条件分布在数学上不可行。若简单对训练环境采用均匀平均或经验频率分配，又会忽略具体观测样本对不同环境潜在状态的偏向性。

本文的切入角度是：虽然未知测试环境的环境特异性预测不可直接恢复，但可以通过对潜在变量和环境变量进行边缘化来估计边缘贝叶斯预测分布。核心 idea：将多环境鲁棒预测形式化为基于经验贝叶斯的层次潜变量模型，通过神经网络自适应学习观测条件下的环境权重来对各环境特异性先验进行加权聚合，使变分下界自然分解为环境特异性拟合项与跨环境平衡正则化项，无需人工设计惩罚即达成无缝外推。

方法详解¶

整体框架¶

本文构建的概率生成过程假设：环境 $e \sim p(e)$ 决定潜变量的先验分布 $z \mid e \sim p(z \mid e)$，而高维观测特征 $x$ 与标签 $y$ 分别通过跨环境稳定的物理/数据机制生成，即 $x \mid z \sim p(x \mid z)$ 与 $y \mid z \sim p(y \mid z)$。针对训练集包含的离散环境 $\mathcal{E}_{\mathrm{tr}}=\{e_1, \dots, e_m\}$，算法的目标是在面对来自未见新环境 $e' \notin \mathcal{E}_{\mathrm{tr}}$ 的新输入 $x$ 时，恢复边缘贝叶斯预测规则 $p(y \mid x) = \iint p(z \mid x, e) p(e \mid x) p(y \mid z) \, dz \, de$。

整个学习与推断流水线由四个参数化网络协同驱动：环境特异性推断网络 $g_\theta(x, y, e)$、共享标签预测器 $f_\phi^{(1)}(z)$、辅助环境条件分类器 $f_\varphi^{(0)}(x, e)$ 以及自适应环境权重预测网络 $h_\psi(x)$。在训练阶段，模型通过变分下界联合优化上述四个模块；在测试阶段，无需环境标识输入，直接通过学到的条件权重 $p_\psi(e \mid x)$ 对各训练环境诱导的潜状态先验进行自适应加权集成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入观测特征 x 与环境标识 e"] --> PostNet["1. 环境特异性推断网络与重参数化采样<br/>g_θ(x, y, e) 输出高斯分布并采样潜变量 z"]
    PostNet --> PredNet["2. 共享潜在标签预测头<br/>f_ϕ(1)(z) 输出目标分布 p_ϕ(y|z)"]
    In --> AuxNet["3. 观测级自适应环境权重估计与边缘化<br/>h_ψ(x) 预测 p_ψ(e|x) 与 f_φ(0)(x, e) 边缘化构成 EB-prior"]
    PostNet & PredNet & AuxNet --> ObjNet["4. 诱导式环境平衡变分目标与自适应外推<br/>L_e 环境特异拟合项结合自然诱导的 R(q) 跨环境散度平衡项"]
    ObjNet --> Out["输出：未知测试环境下的边缘贝叶斯预测 p(y|x)"]

关键设计¶

1. 环境特异性推断网络与重参数化采样：捕获细粒度后验分布

由于在数据集中每个观测 $x$ 通常只伴随一个确定标签 $y$，直接通过训练对推断出的隐变量信息天然受标签条件制约，即服从 $p(z \mid x, y, e)$。为了高效逼近该真实后验，本文设计了均值和对角方差参数化网络 $g_\theta(x, y, e) = (\mu_\theta(x, y, e), \sigma_\theta(x, y, e))$，构建变分高斯分布： $$q(z; x, y, e) = \mathcal{N}(\mu_\theta(x, y, e), \mathrm{diag}(\sigma_\theta(x, y, e)))$$ 该网络显式接收当前环境标识作为输入，从而为每个环境保留其特定的潜变量结构与后验几何形状。在反向传播过程中，利用标准高斯重参数化技巧采样潜变量 $z = \mu_\theta + \sigma_\theta \odot \epsilon$（其中 $\epsilon \sim \mathcal{N}(0, I)$），确保梯度能够平滑穿透随机采样层传导至编码网络。

2. 共享潜在标签预测头：固化跨环境稳定的因果机制

为了严格贯彻图模型中 $p(y \mid z)$ 跨环境保持不变的因果假设，设计了一个全局共享的神经网络 $f_\phi^{(1)}(z)$，用于直接参数化条件标签分布 $\hat{p}_\phi(y \mid z)$。这一设计将标签预测完全建立在解耦后的潜在表征 $z$ 之上，杜绝了模型走捷径直接从输入特征 $x$ 中学习特定环境的虚假关联。预测头在优化过程中接收来自所有训练环境的反向传播梯度，促使潜空间中表征的几何分布真正对齐于目标的本质语义。

3. 观测级自适应环境权重估计与边缘化：构建数据驱动的经验贝叶斯先验

在测试阶段由于新环境未见且无环境标签输入，传统的经验频率均值假设（赋予各环境 $1/m$ 恒定权重）在观测特征分布偏向某些潜状态时会导致严重的渐近偏差。为此，EBER 引入了环境判别网络 $h_\psi(x)$ 直接估计后验概率 $\hat{p}_\psi(e \mid x)$，并结合辅助分类器 $f_\varphi^{(0)}(x, e)$ 估计 $\hat{p}_\varphi(y \mid x, e)$。通过贝叶斯定理得到联合后验权重 $\hat{p}(e \mid x, y) \propto \hat{p}_\varphi(y \mid x, e) \hat{p}_\psi(e \mid x)$。随后，通过对标签空间积分获得单环境先验 $\hat{q}(z; x, e) = \int q(z; x, y, e) \hat{p}_\varphi(y \mid x, e) \, dy$，最终以数据自适应加权的方式获得全局经验贝叶斯先验（EB-prior）： $$\hat{q}(z; x) = \sum_{e \in \mathcal{E}_{\mathrm{tr}}} \hat{p}_\psi(e \mid x) \, q(z; x, e)$$ 这一机制使模型在处理新样本时，能根据输入 $x$ 的特征动态度量其与各个训练环境的亲疏程度，从而生成针对该具体样本定制的潜变量先验分布。

4. 诱导式环境平衡变分目标与自适应外推：模型结构自然衍生正则化

本文严格推导了条件对数似然 $\log p_q(y \mid x)$ 的 Jensen 变分下界，证明了全局期望变分目标能够严格分解为各环境内部期望损失与显式的跨环境平衡项：$\mathcal{L} = \mathbb{E}_{p(e)}[\mathcal{L}_e] + R(q)$。其中环境平衡项定义为： $$R(q) \coloneqq \mathbb{E}_{p(x,y)}\Big[\mathbb{E}_{p(e \mid x,y)}[\mathrm{KL}(q(z; x, y, e) \,\|\, q(z; x, y))]\Big] - \mathbb{E}_{p(x)}\Big[\mathbb{E}_{p(e \mid x)}[\mathrm{KL}(q(z; x, e) \,\|\, q(z; x))]\Big]$$ 该平衡项由两部分 KL 散度构成：前项衡量各环境特异性后验与全局聚合后验的偏离程度，后项衡量各环境先验与全局先验的偏离程度。与 IRM 等方法凭空定义梯度惩罚不同，$R(q)$ 完全由潜变量图模型自身的数学一致性自然诱导而出。当环境特异性后验的发散程度超过先验本身的差异时，该项自动施加惩罚，防止潜变量表征过度过拟合于个别环境的特异性模式。

损失函数 / 训练策略¶

在训练阶段，算法采用端到端的微型批次梯度上升（Algorithm 1）。批次目标函数为基于单个样本蒙特卡洛采样的经验变分下界： $$\widehat{\mathcal{L}}(\theta, \phi, \varphi, \psi) = \frac{1}{|\mathcal{D}|} \sum_{e \in \mathcal{E}_{\mathrm{tr}}} \sum_{i=1}^{n_e} \left[ \log p_\phi(y_i^e \mid z_i) - \mathrm{KL}\left(\hat{q}(z; x_i^e, y_i^e) \,\|\, \hat{q}(z; x_i^e)\right) \right]$$ 在极低信噪比的极端合成实验中，可选择性引入有监督环境分类辅助损失 $\lambda_{\mathrm{env}} \frac{1}{|\mathcal{D}|} \sum \log \hat{p}_\psi(e \mid x_i^e)$（$\lambda_{\mathrm{env}}=1$），而在所有真实数据场景中该权重均设为 0，完全依靠变分下界端到端驱动。测试期预测则通过封闭积分的加权形式高效完成：$\hat{p}(y \mid x) = \sum_{e \in \mathcal{E}_{\mathrm{tr}}} \hat{p}_\psi(e \mid x) \int \hat{p}_\phi(y \mid z) q(z; x, e) \, dz$。

实验关键数据¶

主实验¶

论文在两套受控模拟环境（参数化合成、Colored MNIST）以及三项涵盖天文学、宏基因组学与重症监护的真实世界跨环境基准上全面评估了 EBER。对比基线包括标准经验风险最小化（ERM）以及三大学术界主流环境平衡/不变学习算法 IRM、V-REx 与 Fishr。测试指标包括新环境上的负对数似然（NLL）、分类准确率（Accuracy）以及重症场景下的不平衡指标 AUROC 和 AUPRC。

数据集 / 任务	评估指标	EBER (本文)	最佳基线	最佳基线名称	提升幅度
参数化模拟数据	NLL ↓	0.319 ± 0.069	0.640 ± 0.026	V-REx	-0.321 (NLL显著降低)
参数化模拟数据	Accuracy ↑	0.928 ± 0.003	0.582 ± 0.011	Fishr	+0.346 (+59.5% 相对提升)
Colored MNIST	NLL ↓	0.519 ± 0.077	0.572 ± 0.007	IRM	-0.053
Colored MNIST	Accuracy ↑	0.787 ± 0.071	0.747 ± 0.017	IRM	+0.040 (+5.4%)
SDSS 类星体-恒星分类	NLL ↓	0.221 ± 0.019	0.382 ± 0.009	V-REx	-0.161 (p < 0.001)
SDSS 类星体-恒星分类	Accuracy ↑	0.921 ± 0.007	0.897 ± 0.001	ERM	+0.024 (p < 0.001)
MicrobiomeHD 结直肠癌筛查	NLL ↓	0.850 ± 0.112	0.903 ± 0.087	V-REx	-0.053 (p < 0.01)
MicrobiomeHD 结直肠癌筛查	Accuracy ↑	0.591 ± 0.027	0.588 ± 0.011	IRM	+0.003 (p < 0.01)
PhysioNet 脓毒症预警 (ICU)	AUROC ↑	0.740 ± 0.008	0.724 ± 0.012	IRM	+0.016 (显著超越其他基线)
PhysioNet 脓毒症预警 (ICU)	AUPRC ↑	0.179 ± 0.008	0.169 ± 0.012	IRM	+0.010 (+5.9%)

消融实验¶

为了严格厘清自适应环境权重 $p(e \mid x)$ 与环境特异性后验估计各自对跨环境外推的贡献，论文设计了 7 种系统性消融变体。以下在分支线索强度 $\rho = 0.15$ 的合成任务上展现关键消融对比：

配置 / 变体	训练阶段权重	测试阶段权重	后验分布形式	NLL ↓	Accuracy ↑	关键结论说明
EBER 完整模型 (Var 1)	学习估计	学习估计	学习估计	0.301 ± 0.071	0.935 ± 0.002	完整模型性能逼近理论最优 Bayes-x (0.231 / 0.935)
均匀测试权重 (Var 2)	学习估计	均匀权重 (1/m)	学习估计	0.656 ± 0.015	0.614 ± 0.036	仅在测试期替换为均匀权重，准确率暴跌 32.1%
Oracle 测试权重 (Var 3)	学习估计	真实后验	学习估计	0.292 ± 0.071	0.935 ± 0.001	证实学到的后验推断网络已高度精准
全 Oracle 权重 (Var 4)	真实后验	真实后验	学习估计	0.242 ± 0.004	0.935 ± 0.001	揭示剩余性能缺口主要源自权重估计精度
均匀权重+Oracle后验 (Var 7)	均匀权重	均匀权重	真实后验	0.674 ± 0.009	0.589 ± 0.048	关键反证：即便拥有完美后验，均权聚合依然失效
全均匀基线 (Var 6)	均匀权重	均匀权重	学习估计	0.664 ± 0.009	0.571 ± 0.022	退化至与经典 ERM / IRM 相当的较差水准

关键发现¶

自适应条件权重的决定性作用：消融实验（Var 2 与 Var 7）决定性地证明，单纯依靠高质量的隐变量推断网络不足以应对分布偏移。若在聚合时使用粗糙的均匀权重，模型将退化出系统性偏差；只有配合依赖于输入 $x$ 的自适应环境权重 $\hat{p}(e \mid x)$，才能正确实现以 $m^{-1/2}$ 速率收敛至真实边缘贝叶斯预测。
环境数与样本规模扩展性：随着训练环境数量从 $m=12$ 增加到 $m=36$，EBER 的 NLL 从 0.254 进一步收敛至 0.252，极度逼近理论贝叶斯最优限度（0.234）。同时，每个环境的样本量达到 $n=1000$ 时，EBER 几乎完全复现了 Bayes-x 的判别精度。
真实天文学特征的自适应选择：在 SDSS 类星体实验中，模型将测试集样本自适应赋予最大权重的训练环境恰好是色指数分布与测试分布重合度最高的银道纬度区间（$[72^\circ, 144^\circ)$），验证了 $\hat{p}(e \mid x)$ 在物理层面上捕获了观测域的流形重叠性。

亮点与洞察¶

将环境平衡目标内生化：不同于以往启发式添加人为加权正则化项的不变学习工作，EBER 证明了环境平衡项 $R(q)$ 可以从标准的层次生成模型变分下界中精确推导展开，实现了经验贝叶斯理论与域泛化实践的优美统一。
破解未知环境外推的信息死锁：在测试期目标环境完全未知的强约束下，巧妙运用全概率公式将无法求解的环境特异分布转化为对已知训练环境的动态软插值，为跨域迁移提供了通用数学框架。
高维鲁棒表征学习的方法论启示：这一机制可以直接迁移至多模态临床表型整合或多源遥感遥测等存在严重批次效应（Batch Effects）的场景，通过将批次作为隐式环境先验进行解耦与边缘化，实现对全新仪器或医院的零样本适应。

局限与展望¶

计算复杂度随训练环境数量线性上升：在批次前向推断中，模型必须针对每个训练环境分别计算 $q(z; x, e)$ 和 $h_\psi(x)$，导致训练用时从 $m=12$ 时的 36 秒激增至 $m=36$ 时的 173 秒。未来工作可探索结合 VampPrior 伪观测混合技术，将大规模离散环境压缩至少量紧凑的原型先验中。
超高维稀疏特征下的维度诅咒：敏感度分析表明，当特征维度达到 $d=1000$ 且有效信号稀疏时，所有无监督/自监督环境识别与表征学习方法的精度均出现下滑，如何在超高维噪声下保持环境权重的准确估计仍需引入更强的稀疏归纳偏置。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从经验贝叶斯变分视角推导出自然诱导的环境平衡目标，理论自洽而深刻。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖解析推导、参数化模拟、抗混淆图像测试以及三大经典跨领域科学真实数据集，并配有完备的显著性检验与消融。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨规范，图模型与变分下界分解推导清晰明了。
价值: ⭐⭐⭐⭐⭐ 为解决具有潜在状态漂移的多环境科学发现与鲁棒临床决策提供了极具复用价值的概率表征范式。