Conditional Counterfactual Mean Embeddings: Doubly Robust Estimation and Learning Rates¶
会议: NeurIPS 2026
arXiv: 2602.04736
OpenReview: https://openreview.net/forum?id=WBrsTgkHw7
代码: https://github.com/donlap/Conditional-Counterfactual-Mean-Embeddings
领域: 因果推理
关键词: 条件反事实均值嵌入, 双重鲁棒估计, 再生核希尔伯特空间, 学习率, 异质性因果效应
一句话总结¶
本文提出将潜在结果的条件反事实分布映射到再生核希尔伯特空间的条件反事实均值嵌入(CCME)框架,通过去偏核特征构建泛函伪结果并设计两阶段元估计器及三类具体估计器,首次给出了显式依赖特征维数与索伯列夫光滑度的有限样本双重鲁棒收敛率。
研究背景与动机¶
因果推断的核心目标是评估干预或策略对潜在结果的影响。在传统潜在结果框架下,学术界和工业界大多聚焦于标量因果效应指标的估计,例如平均因果效应(ATE)或条件平均因果效应(CATE)。然而,标量汇总指标往往会掩盖潜在结果分布的全局几何形态:两个具有完全相同均值的反事实分布可能在方差、偏度、尾部风险或多峰聚集特性上存在巨大差异。例如在个性化医疗与金融风控中,一种新药或策略可能对全人群的平均疗效中性,但对特定亚群可能诱发严重的极端尾部风险;反事实分布若呈现多峰性,往往直接揭示了未观测异质性反应群体的存在。因此,为了实现精准决策与安全干预,必须对给定特定协变量条件下的完整条件反事实分布进行全貌刻画。
为了捕捉概率分布的丰富统计特征而不受限于参数化分布假设,再生核希尔伯特空间(RKHS)中的核均值嵌入(KME)方法被广泛引入,将无限维分布一一映射为希尔伯特空间中的确定向量。此前已有学者探索了边缘反事实均值嵌入(CME)的估计,但现有针对条件反事实分布(CCD)的研究非常匮乏。现有少数尝试大多依赖一步法插入估计器(plug-in estimator),这类方法直接将处理组上的回归结果作为反事实分布的估计,极易遭受由混杂因子诱发的严重选择偏差(selection bias);而在经典因果效应分析中被广泛用于偏差校正的双重鲁棒(Doubly Robust, DR)技术,其形式局限于实数标量空间,如何将其自然推广至无限维 RKHS 泛函空间并保证非参数学习率,此前在理论和算法层面均缺乏明确解答。
面对无限维泛函目标估计中维数灾难与混杂偏差并存的核心矛盾,本文的切入点在于将 Neyman 正交化与双重机器学习(DML)思想延伸至希尔伯特空间特征映射。核心 idea:通过对潜在结果的 RKHS 核特征映射进行双重鲁棒去偏构造泛函伪结果,提出支持任意回归器的两阶段 CCME 元估计器,并衍生出岭回归、深度特征与神经核三类计算权衡互补的实用估计器,证明其在索伯列夫空间下具备双重鲁棒有限样本极小极大收敛率。
方法详解¶
整体框架¶
CCME 的估计目标是给定协变量子集 \(V = \eta(X)\) 时处理组潜在结果 \(Y^1\) 的核均值嵌入 \(\mu_{Y^1|V}(v) = \mathbb{E}[\phi(Y^1) \mid V=v] \in \mathcal{H}_\mathcal{Y}\)。整个估计框架采用两阶段样本分割流水线:首先将大小为 \(2n\) 的观测数据集 \(\mathcal{D} = \{(X_i, A_i, Y_i)\}\) 等分为 \(\mathcal{D}_0\) 和 \(\mathcal{D}_1\)。在第一阶段(讨厌参数估计),利用 \(\mathcal{D}_0\) 拟合倾向得分模型 \(\widehat{\pi}(x) = P(A=1 \mid X=x)\) 以及处理组条件均值嵌入模型 \(\widehat{\mu}_0(x) = \mathbb{E}[\phi(Y) \mid X=x, A=1]\);在第二阶段(伪结果回归),利用 \(\mathcal{D}_1\) 上的样本构建希尔伯特空间值泛函去偏伪结果 \(\widehat{\xi}(Z_{1i})\),并以 \(V_{1i} = \eta(X_{1i})\) 为输入对该伪结果进行经验风险最小化回归,从而得到最终的 CCME 估计量 \(\widehat{\mu}_{Y^1|V}\);最后通过内积计算在任意查询点 \(y\) 处重构出光滑反事实概率密度 \(\widehat{\mathfrak{p}}^1(y \mid v)\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["观测数据集 D = {X_i, A_i, Y_i}"] --> B["样本均匀对半分割<br/>D_0 (讨厌参数) 与 D_1 (伪结果回归)"]
B --> C["一阶讨厌函数估计<br/>倾向得分 pi_hat 与处理组均值嵌入 mu_hat_0"]
C --> D["泛函双重鲁棒伪结果构建<br/>RKHS 空间去偏特征向量 xi_hat(Z_1i)"]
D --> E["二阶条件核回归估计器<br/>岭回归 / 深度特征 / 神经核三选一"]
E --> F["输出条件反事实均值嵌入 mu_hat_Y1|V<br/>内积重构条件反事实概率密度 p_hat(y|v)"]
关键设计¶
1. 泛函双重鲁棒伪结果构建:在无限维希尔伯特空间完成特征去偏
传统双重鲁棒方法(如 AIPW)仅针对实值标量结果设计,直接应用无法处理分布层面的丰富高阶矩信息。本文突破这一限制,直接在 RKHS 空间 \(\mathcal{H}_\mathcal{Y}\) 内构建去偏特征伪结果:
$\(\widehat{\xi}(Z) = \frac{A}{\widehat{\pi}(X)}\left(\phi(Y) - \widehat{\mu}_0(X)\right) + \widehat{\mu}_0(X)\)$
在无混杂性假设(\(Y^a \perp A \mid X\))与强重叠假设(\(\epsilon \le \pi(x) \le 1-\epsilon\))下,只要倾向得分模型 \(\widehat{\pi}\) 或处理组条件均值嵌入模型 \(\widehat{\mu}_0\) 中任意一个正确设定,经验伪结果的条件期望都能无偏识别目标分布的均值嵌入,即 \(\mathbb{E}[\xi(Z) \mid V=v] = \mu_{Y^1|V}(v)\)。该设计将 Neyman 正交性从实数轴升维至再生核赋范空间,从根本上消除了单一边际插入估计器的选择偏差。
2. 两阶段元估计器与模块化解耦:兼容任意一阶与二阶回归架构
直接在不可观测的因果目标上做端到端训练极易发生过拟合或引入第一阶段误差泄漏。本文通过样本分割将训练划分为独立解耦的两阶段:阶段 1 允许算法自由选择二分类器(如随机森林、逻辑回归、神经网络)估计 \(\widehat{\pi}\),以及任意条件均值嵌入算法拟合 \(\widehat{\mu}_0\);阶段 2 则将伪结果 \(\widehat{\xi}(Z_{1i})\) 作为目标,通过求解正则化经验风险最小化问题确定二阶泛函假说 \(\widehat{\mu}_{Y^1|V} = \arg\min_{\mu \in \mathcal{F}} \frac{1}{n}\sum_{i=1}^n \|\mu(V_{1i}) - \widehat{\xi}(Z_{1i})\|_{\mathcal{H}_\mathcal{Y}}^2 + \Lambda(\mu)\)。两阶段之间的样本独立性使得一阶误差在理论推导中可被精准控制,避免了复杂的经验过程自相关问题。
3. 三类具体估计器的计算权衡:核精确度与神经网络扩展性的统一
针对第二阶段中无限维目标优化的计算可解性,本文实例化了三种代表性估计器:
- 核岭回归估计器(Ridge Regression, RR):利用算子值核 \(\Gamma_\mathcal{V}(v, v') = k_\mathcal{V}(v, v') \mathrm{Id}_{\mathcal{H}_\mathcal{Y}}\),显式闭式解为 \(\widehat{\mu}_{\mathrm{RR}}(v) = \widehat{\mathbf{\Xi}} (\mathbf{K}_V + n\lambda_1 \mathbf{I}_n)^{-1} \mathbf{k}_V(v)\)。该估计器方差极小、无逼近误差,但求逆涉及 \(O(n^3)\) 复杂度,适合中小样本。
- 深度特征估计器(Deep Feature, DF):采用轻量神经网络 \(\psi_\theta: \mathcal{V} \to \mathbb{R}^M\) 将协变量降维映射至 \(M \ll n\) 维隐特征,将算子求逆降维至 \(M \times M\) 矩阵(\(O(M^3)\) 复杂度),利用可导 Fréchet 导数闭式解解析表达损失函数进行端到端梯度更新。
- 神经核估计器(Neural-Kernel, NK):预先在输出空间 \(\mathcal{Y}\) 上设置 \(M\) 个均匀网格锚点 \(\{\widetilde{y}_j\}_{j=1}^M\),利用神经网络直接预测网格基函数的组合权重系数 \(f_\theta(v)\),估计形式为 \(\widehat{\mu}_{\mathrm{NK}}(v) = \sum_{j=1}^M f_\theta(v)_j \phi(\widetilde{y}_j)\)。该方法完全消除了矩阵求逆和二阶正则化参数 \(\lambda_1\) 的微调,在超大样本与高维场景下具备最优的推理效率。
4. 统计与投影分解的收敛速率保证:有限样本下的极小极大最优性
论文通过将总体均方误差分解为统计误差(\(\mathbf{A}\))、网络逼近误差(\(\mathbf{B}\))、讨厌参数误差(\(\mathbf{C}\))和核投影误差(\(\mathbf{D}\))四大模块,在索伯列夫光滑度空间 \(\mathcal{W}^{s,q}(\mathcal{Y})\) 与 \(\mathcal{W}^{r,2}(\mathcal{V})\) 下建立了有限样本界。一阶误差项严格受到 \(\min\{\mathcal{R}_\pi^2(\widehat{\pi}), \mathcal{R}_{\mu_0}^2(\widehat{\mu}_0)\}\) 的约束,严格证明了收敛速率维度的双重鲁棒性。当第一阶段估计达到极小极大非参数率时,第二阶段估计器在 \(V\) 上完全收敛至非参数最优率 \(n^{-\frac{2r}{2r+d_v}}\)(在高斯核下对输出维数 \(d_y\) 仅有对数级依赖),从纯理论角度填补了条件反事实测度嵌入的收敛理论空白。
损失函数与优化策略¶
对于深度特征估计器(DF),利用最佳线性算子带入后得到的第二阶段特征学习目标为: $\(\widehat{\mathcal{L}}_{\mathrm{DF}}(\theta) = \mathrm{Tr}\left(\mathbf{K}_{\widehat{\xi}} \left(\mathbf{I}_n - \mathbf{\Psi}_\theta (\mathbf{\Psi}_\theta^\top \mathbf{\Psi}_\theta + n\lambda_1 \mathbf{I}_M)^{-1}\mathbf{\Psi}_\theta^\top\right)\right)\)$ 其中 \(\mathbf{K}_{\widehat{\xi}}\) 为伪结果样本之间的核 Gram 矩阵。对于神经核估计器(NK),目标损失函数退化为在锚点 Gram 矩阵 \(\mathbf{K}_M\) 下的二次型: $\(\widehat{\mathcal{L}}_{\mathrm{NK}}(\theta) = \frac{1}{n} \sum_{i=1}^n \left[ f_\theta(V_{1i})^\top \mathbf{K}_M f_\theta(V_{1i}) - 2 f_\theta(V_{1i})^\top \mathbf{b}_i \right]\)$ 其中 \(\mathbf{b}_i = \frac{A_{1i}}{\widehat{\pi}(X_{1i})} \mathbf{k}_i + \left(1 - \frac{A_{1i}}{\widehat{\pi}(X_{1i})}\right) \mathbf{K}_M g_{\widehat{\theta}_0}(X_{1i})\),所有运算均仅涉及 \(M \times M\) 预计算 Gram 矩阵,极大提升了梯度反传效率。
实验关键数据¶
主实验:合成复杂混杂数据下的反事实密度估计与模型错误指定评估¶
在 10 维连续混杂协变量 \(X \sim \mathcal{N}(\mathbf{1}_{10}, \mathbf{I}_{10})\)、条件子变量 \(V=(X_1,\dots,X_5)\)、强选择偏差 \(\pi(X) = 0.1 + 0.8 \cdot \mathbb{I}(X_1 \in [0,2], X_6 \ge 1.5)\) 的双峰潜在结果实验中,对比了本文 DR 与反概率加权(IPW)、插入法(PI)及一步法(One-Step)在不同样本量与错误设定下的对数均方误差(MSE)。
| 估计器架构 | 伪结果策略 | 两种讨厌模型均正确设定 (n=20k) | 倾向得分模型错误设定 (n=20k) | CME 结果模型错误设定 (n=20k) | 对抗选择偏差能力 |
|---|---|---|---|---|---|
| Ridge Regression (RR) | DR (本文) | 0.0021 (最低) | 0.0022 (保持一致) | 0.0021 (保持一致) | 完全双重鲁棒 |
| Ridge Regression (RR) | IPW | 0.0022 | 0.0084 | 0.0022 | 依赖准确倾向得分 |
| Ridge Regression (RR) | PI | 0.0035 | 0.0035 | 0.0410 (不收敛) | 依赖正确结果模型 |
| Ridge Regression (RR) | One-Step (基线) | 0.0152 | 0.0152 | 0.0520 | 严重残留选择偏差 |
| Deep Feature (DF) | DR (本文) | 0.0048 | 0.0049 | 0.0048 | 完全双重鲁棒 |
| Deep Feature (DF) | IPW | 0.0115 | 0.0240 | 0.0115 | 网络表达受加权方差影响 |
| Deep Feature (DF) | PI | 0.0051 | 0.0051 | 0.0480 (不收敛) | 依赖正确结果模型 |
| Deep Feature (DF) | One-Step (基线) | 0.0210 | 0.0210 | 0.0590 | 存在固定偏差平台 |
| Neural-Kernel (NK) | DR (本文) | 0.0065 | 0.0066 | 0.0065 | 完全双重鲁棒 |
| Neural-Kernel (NK) | IPW | 0.0142 | 0.0285 | 0.0142 | 对非稳倾向得分敏感 |
| Neural-Kernel (NK) | PI | 0.0068 | 0.0068 | 0.0510 (不收敛) | 依赖正确结果模型 |
| Neural-Kernel (NK) | One-Step (基线) | 0.0255 | 0.0255 | 0.0630 | 存在固定偏差平台 |
消融与模型配置分析:三类估计器的超参数敏感性与计算复杂度对比¶
下表汇总了各估计器在合成数据与高维半合成 MNIST 图像去噪因果效应实验中的计算资源需求、参数配置及稳定性特征。
| 估计器方案 | 矩阵求逆时间复杂度 | 网格/隐特征维数 \(M\) | 核心超参数配置 | 方差水平 (30次运行区间) | 适用典型场景 |
|---|---|---|---|---|---|
| Ridge Regression (RR) | \(O(n^3)\) | 无需网格 | \(\sigma=2.0, \lambda_0=\lambda_1=20.0\) | 极低 (确定性正规方程) | 小样本、强理论保证基准场景 |
| Deep Feature (DF) | \(O(M^3)\) | \(M=20\) (低维) / \(1000\) (MNIST) | 2层\(\times\)20隐单元, \(\lambda_1=20.0\), SGD | 中等 (正规化抑制方差) | 中大样本、高维协变量或输出 |
| Neural-Kernel (NK) | \(O(1)\) (无求逆) | \(M=20\) (预设网格) | 2层\(\times\)20隐单元, 无需\(\lambda_1\), 早停/权重衰减 | 偏高 (随机初始化与无正则) | 超大规模数据、流式实时因果推断 |
关键发现¶
- 双重鲁棒性的数值验证:当倾向得分模型错误指定时(用逻辑回归代替随机森林),IPW 变体的误差显著上升,而 DR 变体依靠正确的一阶 CME 结果模型维持了最优收敛轨迹;反之,当结果模型故意遗漏关键混杂变量 \(X_6\) 时,PI 变体完全失效,DR 则依托正确的倾向得分稳健收敛,实验与定理 1-4 的速率结论完美契合。
- 多峰分布形态重构:定性密度重构实验显示,一步法基线(One-Step)由于未能校正 \(X_1\) 与 \(X_6\) 的混杂偏倚,严重扭曲了混合伯努利潜在分布的双峰比例;而本文 DR 估计器能够高保真拟合 \(v_1\)(主峰概率 0.75)与 \(v_2\)(主峰概率 0.48)的双峰轮廓。
- 高维半合成 MNIST 图像反事实重构:在 \(784\) 维高维图像去噪因果实验中,强度偏向选择导致传统一步法过度偏向于检索高笔画强度的数字图像,而本文 DR 深度特征估计器成功抵消了选择偏差,检索出的类别经验众数图像与 Oracle 干净数据训练的基准几乎完全一致。
亮点与洞察¶
- 将 Neyman 正交化推广至赋范 Hilbert 空间:不仅停留在传统的条件平均因果效应(CATE)标量层面,首次通过核特征内积构建去偏伪结果,使得分布级别的信息完全被保存在确定性向量中。
- 完备的计算-统计谱系设计:给出的三类估计器精巧地覆盖了从闭式高精度(RR)、特征学习降维(DF)到无求逆超高效(NK)的完整计算谱系,适应从精准小样本到工业级大数据的不同需求。
- 理论推导首次突破无穷维速率瓶颈:通过精心分解局部 Rademacher 复杂度与索伯列夫嵌入定理,克服了算子逼近与非参数核估计中的复合交织误差,给出了关于因果变量维度与光滑度指标显式依赖的紧确阶数。
局限与展望¶
- 条件重叠假设敏感性:理论与算法均强依赖于正重叠假设 \(\pi(X) \ge \epsilon > 0\);在极端选择偏倚或高维混杂场景下,倾向得分倒数可能引发伪结果范数膨胀,未来可探索剪枝(trimming)或稳定权重机制。
- 网格点依赖与维数瓶颈:神经核估计器(NK)在输出空间 \(\mathcal{Y}\) 上依赖准均匀网格点覆盖,这在标量或低维连续潜在结果(\(d_y \le 3\))下极其高效,但在高维连续输出(如医学多模态图像/高维时序)中网格点规模会呈指数级爆炸,需要进一步结合生成模型或隐空间锚点采样。
相关工作与启发¶
- vs Muandet et al. (2021) / Park et al. (2021): 后者提出了边缘反事实均值嵌入(CME)或无偏校正的一步法插值估计器,但缺乏对选择混杂的倾向得分加权保护;本文通过双重鲁棒机制,在保证分布刻画能力的同时彻底消除了选择偏差。
- vs Kennedy et al. (2023) / Martinez-Taboada & Kennedy (2024): 后者基于余弦级数或 Stein 差异进行双重鲁棒反事实概率密度函数(PDF)估计;本文基于 RKHS 均值嵌入,无需对密度形式做特定级数基展开,且产物可直接下游复用于条件独立性检验、核两样本检验等多样化任务。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将双重鲁棒因果去偏与条件反事实均值嵌入在 RKHS 中系统融合,构建了完整的理论与估计器体系。
- 实验充分度: ⭐⭐⭐⭐☆ 覆盖了双峰合成实验的全面错误指定测试、定性轮廓拟合及 784 维 MNIST 半合成实验,理论与实证紧密对应。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密,算法设计由浅入深,多角度计算权衡论述非常透彻。
- 价值: ⭐⭐⭐⭐⭐ 为异质性因果效应的全分布评估提供了通用且具备严格非参数收敛保证的基础工具。