Controllable Generative Sandbox for Causal Inference¶

会议: NeurIPS 2026
arXiv: 2603.03587
论文: https://arxiv.org/abs/2603.03587
代码: https://github.com/zhangqiecho/causalmix
领域: 因果推理
关键词: 因果推断基准、可控生成式沙箱、异质性因果效应、潜变量混杂、贝叶斯高斯混合先验

一句话总结¶

提出了统一真实表格数据保真度与显式因果可控性的变分生成沙箱框架 CausalMix，通过贝叶斯高斯混合先验解耦多模态特征分布，并引入重叠度对数密度比正则与 CATE / 未观测混杂复合约束，为因果评估、超参调优及统计功效分析提供了高保真且真值已知的基准环境。

研究背景与动机¶

因果推断的核心瓶颈在于“因果推断根本问题”——每个研究个体在现实中仅能观测到一种潜在结局，反事实结局的不可观测性使得评估与验证因果估计器必须依赖反事实真值已知的合成数据。在现代临床流行病学与观察性研究中，合成数据扮演着不可替代的沙箱角色，不仅用于因果森林、双重机器学习等前沿估计器的压力测试与超参数调优，还能在保护隐私的前提下辅助前瞻性试验的样本量与功效测算。然而，传统参数化模拟器往往基于简化的结构方程，脱离真实临床表格数据中普遍存在的混合变量类型（连续、二值、无序分类）、多模态聚类以及非线性特征依赖；而基于实证重采样的半合成基准或 Plasmode 模拟又极难灵活调节干预重叠度、异质性效应尺度与未观测混杂偏倚。

现有深度生成因果基准（如基于 WGAN 或流模型的生成器）虽然提升了边际与条件分布的拟合保真度，却普遍陷入了“实证保真度”与“因果可控性”难以兼得的困境。这些模型要么将异质性因果效应和重叠度调节退化为拟合极值之间的启发式插值，完全缺乏对隐藏混杂的显式建模机制；要么在引入因果函数约束时，因潜空间假设过于单一（如标准高斯先验）以及连续变量仅采用均方误差优化，导致无法捕捉表格数据的多模态异质分布，且缺乏约束在生成阶段实际落盘的稳定性保障。

本文的切入点是打破数据分布建模与结构因果控制的割裂，在非参数结构因果模型（SCM）指导下重构模块化条件变分自编码器。核心 idea：构建 CausalMix 变分生成框架，在潜空间利用自适应贝叶斯高斯混合先验捕捉复杂表格多模态依赖，并结合变量特异性解码头与对数密度比重叠度正则、CATE 复合惩罚与方差正则，在统一优化目标下实现实证分布逼真度与因果机制全因子显式控制的无损解耦。

方法详解¶

整体框架¶

CausalMix 将联合观测分布分解为 $p(X, T, Y) = p(T) p(X \mid T) p(Y \mid X, T)$ 三阶段序贯生成结构，包含处理分配模型、前处理协变量生成器 $G_{X,\theta}$ 以及后处理潜在结局生成器 $G_{Y,\theta}$。在模型前向流程中，首先通过伯努利模型抽取二值干预 $T'$，随后将 $T'$ 作为条件送入协变量条件变分自编码器采样生成高维混合变量 $X'$，最后将 $(X', T')$ 联合输入结局生成器并行解码潜在结局对 $(Y'(0), Y'(1))$，并由一致性假设得到观测结果 $Y' = T' Y'(1) + (1 - T') Y'(0)$。

整个架构在训练时直接施加三维显式因果控制函数集 $\Psi = \{\alpha(X), \tau(X), \kappa(X, T)\}$，分别精确控制协变量重叠度、条件平均处理效应（CATE）异质性以及未观测混杂偏倚；并在潜空间拟合狄利克雷过程贝叶斯高斯混合模型，使得生成采样能够精准还原复杂真实分布的聚类几何形态。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    In["输入观测数据 (X, T, Y)<br/>元数据预处理与类型识别"] --> Enc["条件编码器参数化<br/>输出潜高斯后验分布均值与方差"]
    Enc --> MixPrior["贝叶斯高斯混合先验拟合<br/>狄利克雷过程自适应多模态潜空间"]
    MixPrior --> DecHeads["多数据类型特异性解码头<br/>高斯连续 / 伯努利二值 / Softmax离散"]
    DecHeads --> OvReg["对数密度比重叠度正则<br/>对齐预设倾向评分重叠度 α(X)"]
    OvReg --> CausalReg["因果效应与混杂复合正则<br/>对齐预设 CATE τ(X) 与混杂偏倚 κ(X,T)"]
    CausalReg --> Out["生成可控因果沙箱数据<br/>带真实反事实 (X', T', Y'(0), Y'(1))"]

关键设计¶

1. 贝叶斯高斯混合先验拟合：消除单模态高斯失真并保留复杂潜空间几何 标准条件变分自编码器通常假设各向同性高斯先验 $p(z) = \mathcal{N}(0, I)$，但在具有多重亚群、非线性关联与聚类特性的异质性临床表格数据中，后验潜变量编码均值往往呈现强烈的多峰分布。强行从单模态标准高斯先验中采样，会严重扭曲潜变量空间的固有拓扑结构，导致生成的协变量在密度边缘出现伪聚类或不真实的组合空洞。为此，CausalMix 采用后验拟合机制，在 VAE 优化收敛后冻结编码器权重，提取所有训练样本的潜均值 $\{\mu_\theta(\mathcal{O}_i)\}_{i=1}^N$，并拟合带有狄利克雷过程（Dirichlet Process）截断折棍构造的贝叶斯高斯混合模型（BGMM）： $$p_{\mathrm{BGMM}}(z) = \sum_{k=1}^K \pi_k \mathcal{N}(z \mid \mu_k, \Sigma_k)$$ 该设计将高斯成分上限设为潜空间维度，利用变分推断自适应确定稀疏有效的聚类分支 $K$。在生成阶段，模型直接从 $p_{\mathrm{BGMM}}(z)$ 抽样并结合平均后验方差噪声进行解码，既无须对编码器施加破坏凸性优化的对抗判别器，又完整保留了混合表格数据的真实拓扑分布。

2. 多数据类型特异性解码头：统一连续变量异方差高斯似然与分类交叉熵建模 真实观察性研究包含连续生理指标、二值合并症诊断与多分类用药史，传统因果生成模型常将所有特征扁平化为连续数值并以均方误差（MSE）作为重建损失，容易引发数值尺度不均、方差坍缩以及梯度尺度失配问题。CausalMix 设计了结构解耦的多头解码架构：对于二值变量输出 logit 并应用伯努利对数似然；对于多分类变量输出维度对应的类别 logit 并应用 Softmax 交叉熵；对于连续变量则由解码头并行预测条件均值 $\mu_i$ 与截断对数方差 $\log \sigma_i^2$，采用高斯负对数似然（Gaussian NLL）进行联合优化： $$\log p_\theta(X_i \mid z, T'=t) = -\frac{1}{2}\left[\log(2\pi) + \log \sigma_i^2 + \frac{(X_i - \mu_i)^2}{\sigma_i^2}\right]$$ 该设计使模型不仅能够自适应学习连续特征的局部位置与离散程度，还能在处理异方差噪声与受限值域时维持平滑稳健的梯度反向传播。

3. 对数密度比重叠度正则：精准操控倾向评分重叠度与实证正值性假设边界 在因果推断方法论评测中，检验估计器对“正值性（Positivity）违反”或极端倾向评分区域的外推鲁棒性是核心命题。传统模拟器难以精细调控两组协变量的重叠程度。CausalMix 显式定义协变量在不同干预组下的条件密度比 $\alpha(x) = P(X=x \mid T=0) / P(X=x \mid T=1)$，并通过多头解码器输出的特征条件似然独立累加，推导出生成器诱导的特征级对数密度比： $$\log \alpha_\theta(X) = \log p_\theta(X \mid T'=0) - \log p_\theta(X \mid T'=1)$$ 在模型训练前处理生成器 $G_{X,\theta}$ 时，引入均方误差对齐惩罚项 $\mathcal{L}_\alpha(\theta) = \mathbb{E}_X[(\log \alpha_\theta(X) - \log \alpha(X))^2]$。通过调整目标函数 $\log \alpha(X)$（如常数 0 对应完全重叠，非零常数或协变量依赖函数对应渐进性正值性破坏），研究者能够以闭式连续梯度精确诱导从强重叠到局部不可辨识的任意因果场景。

4. 因果效应与混杂复合正则：Huber 均值对齐结合方差正则稳定异质性表达 当用户预先设定低维或弱非线性的因果函数时，单纯依靠均方误差（MSE）惩罚极易在复杂神经网络优化过程中受到异常值干扰，或者导致因果效应在微观个体层面过度平滑和离散坍缩。CausalMix 为 CATE 目标残差 $\Delta\tau_\theta(X) = \tau_\theta(X) - \tau(X)$ 和混杂偏倚残差 $\Delta\kappa_\theta(X, T) = \kappa_\theta(X, T) - \kappa(X, T)$ 构建了加权 Huber 复合结构与方差正则： $$\mathcal{L}_\tau^{\mathrm{mean}}(\theta) = \mathbb{E}_X\left[\lambda_\tau^{\mathrm{mse}} (\Delta\tau_\theta)^2 + \lambda_\tau^{\mathrm{sl1}} \mathrm{SmoothL1}(\Delta\tau_\theta)\right], \quad \mathcal{L}_\tau^{\mathrm{var}}(\theta) = \mathrm{Var}[\Delta\tau_\theta]$$ 二次项牢牢锚定总体平均处理效应（ATE）的无偏对齐，Smooth L1 项抑制高维协变量边缘区域极端预测值的扰动，而方差惩罚项 $\mathrm{Var}[\Delta\tau_\theta]$ 强力约束个体预测值偏离真实因果超曲面的杂散方差。三者协同优化，确保在忠实保留预设效应异质性的同时，完全杜绝了模型自主幻觉出的虚假个体方差。

损失函数 / 训练策略¶

CausalMix 的总训练目标由数据保真度 ELBO 与各项因果刚性惩罚项线性加权构成： $$\mathcal{L}(\theta) = \mathcal{L}_{\mathrm{VAE}}(\theta) + \lambda_\alpha \mathcal{L}_\alpha(\theta) + \lambda_\tau \mathcal{L}_\tau^{\mathrm{mean}}(\theta) + \lambda_\tau^{\mathrm{var}} \mathcal{L}_\tau^{\mathrm{var}}(\theta) + \lambda_\kappa \mathcal{L}_\kappa^{\mathrm{mean}}(\theta) + \lambda_\kappa^{\mathrm{var}} \mathcal{L}_\kappa^{\mathrm{var}}(\theta)$$ 其中 $\mathcal{L}_{\mathrm{VAE}}(\theta) = -\mathbb{E}_{q_\theta(z)}[\log p_\theta(\mathcal{O} \mid z)] + \lambda_{\mathrm{KL}} \mathrm{KL}(q_\theta(z) \parallel p(z))$。在具体实现中，前处理模块 $G_{X,\theta}$ 与后处理模块 $G_{Y,\theta}$ 分别独立优化。因果刚性权重 $\lambda_\tau$ 与 $\lambda_\kappa$ 设为 $10^3$，重叠度权重 $\lambda_\alpha$ 设为 $10^1 \sim 10^2$；优化采用 Adam 优化器（学习率 $10^{-3}$），并依托 PyTorch Lightning 在 80/20 验证集划分上实时监控验证集因果残差与 VAE 重建损失以触发早停机制。

实验关键数据¶

主实验：mCRPC 队列复杂因果场景（Scenario 3）下的数据保真度与因果真实性评估¶

在最具挑战性的前列腺癌真实队列场景 3 中（包含 4,098 名患者、18 个基线混合特征，设置饱和非线性 CATE、双向协变量依赖未观测混杂及病史驱动的重叠度破坏），对比 CausalMix 采用 BGMM 先验与传统标准各向同性高斯先验的生成表现：

评估维度	具体指标	优化方向	CausalMix (BGMM 先验)	标准高斯先验基线	性能差异与优势
连续变量边际拟合	归一化 Wasserstein 距离	↓ 越低越好	0.075 (±0.001)	0.679 (±0.008)	降低 88.9% 边际误差
连续变量边际分布	KS 补数 (KSComplement)	↑ 越高越好	0.954 (±0.001)	0.720 (±0.005)	拟合优度提升 +0.234
离散变量边际拟合	TV 补数 (TVComplement)	↑ 越高越好	0.961 (±0.001)	0.898 (±0.002)	离散频数还原更精确
连续-连续变量相关	Pearson 相关相似度	↑ 越高越好	0.984 (±0.002)	0.883 (±0.009)	协方差结构还原度极高
全变量互信息依赖	对称不确定性 (SU) 相似度	↑ 越高越好	0.991 (±0.002)	0.990 (±0.000)	保持优秀非线性关联
全局联合分布拟合	归一化能量距离 (Energy Dist)	↓ 越低越好	0.011 (±0.000)	0.080 (±0.002)	联合空间距离降低 86.2%
双样本分类器判别	C2ST 分数 (1 - 2\|AUC-0.5\|)	↑ 越高越好	0.625 (±0.008)	0.333 (±0.010)	分类器难以区分真实与生成数据
因果效应异质性对齐	CATE 相关系数 $\rho(\tau_\theta, \tau)$	↑ 越高越好	0.956 (±0.003)	0.931 (±0.005)	高度还原复杂非线性饱和趋势
重叠度重构误差	解码器对数密度比 MSE	↓ 越低越好	0.004 (±0.000)	0.300 (±0.091)	误差降低接近两个数量级

消融与基准评测：不同 CATE 估计器在校准沙箱中的精度、置信区间覆盖率与耗时对比¶

在 Scenario 2（真实 ATE = 0.056，弱线性异质性，无未观测混杂）校准环境下，运行 50 次独立蒙特卡洛复现，评测十种代表性 CATE 估计器的综合表现：

估计器家族	具体方法	ATE RMSE	Mean ATE SE	ATE 覆盖率 (95% 目标)	Mean CATE 覆盖率	单次运行耗时 (秒)
贝叶斯树模型	Bayesian Causal Forest (BCF)	0.010	0.010	0.90	0.91	258.7
自适应森林	Causal Forest	0.009	0.009	0.94	0.87	25.1
双重机器学习	Standard Linear DML	0.009	0.009	0.96	0.90	2.2
双重机器学习	Lasso DML	0.009	0.008	0.88	0.37 (严重失真)	229.9
双重机器学习	GBR DML	0.010	0.113	1.00 (SE虚高)	0.82	268.1
双重稳健学习	Standard Linear DR-Learner	0.010	0.046	1.00	0.92	2.3
双重稳健学习	Lasso DR-Learner	0.010	0.008	0.88	0.40 (严重失真)	234.5
元学习器	Lasso X-Learner	0.024	0.036	1.00	0.82	118.4
元学习器	GBR X-Learner	0.024	0.085	1.00	0.82	160.4

关键发现¶

BGMM 先验是解锁高难度复杂分布生成的命脉：在简单场景 1 下两种先验差异微弱，但在场景 3（高维多模态、饱和效应与病史相关重叠度）中，标准高斯先验的连续边际 Wasserstein 距离大幅恶化至 0.679，解码器重叠度 MSE 激增至 0.300；而 BGMM 依然保持了 0.075 的超低距离与 0.004 的重叠度 MSE，表明潜变量混合分布对支撑多模态条件生成至关重要。
群体 ATE 准确性不等于可靠的个体 CATE 不确定性推断：Lasso DML 与 Lasso DR 估计器虽然能够获得顶尖的 ATE RMSE（0.009），但其 CATE 置信区间覆盖率断崖式下跌至 0.37 与 0.40，原因是 Lasso 过度收缩抹杀了细粒度异质性；相比之下，BCF 与 Linear DR-Learner 展现出优异的 CATE 覆盖率（0.91~0.92）。
因果森林超参优化规律：在超参网格实验中，树的棵数（1000~4000）仅线性增加计算时间（从 11s 增至 42s），对估计精度影响极小；而最小叶节点尺寸（min.node.size）是核心瓶颈——叶节点设为 20 时在 PEHE（0.027）与 CATE 覆盖率（0.818）之间达到黄金平衡点，避免了叶节点过小导致的过拟合方差膨胀与过大导致的欠拟合过度平滑。
功效分析揭示异质性探测与效应修饰因子排名的样本量鸿沟：在 mCRPC 前瞻性功效模拟中，探测心血管疾病（CVD）修饰效应存在（达到 90% power）仅需约 2,000 例样本；但要让因果森林在变量重要性排名中将 CVD 稳定识别为 Top-1 驱动因子，则必须将样本量提升至 5,000 例以上。

亮点与洞察¶

将因果识别假设违背转化为可参数化、可导的设计旋钮：通过将倾向评分重叠度转化为特征对数密度比惩罚，将未观测混杂形式化为反事实潜在结局间的条件差值，使因果学者首次能够在神经网络反向传播中精确操控正值性破坏与不可忽略性偏倚程度。
后验无损结合 BGMM 先验与因果约束：巧妙地将潜空间的多模态拟合解耦至 VAE 训练之后，用狄利克雷过程自动探测表格聚类，彻底规避了端到端 GMM-VAE 常见的后验坍缩与训练不稳定性。
隐私保护与实证真实性的良好折中：DCR（最邻近记录距离）隐私分析显示，BGMM 先验合成数据的受保护比例稳定在 0.60~0.72，距离比中位数均大于 1.0，证明生成的高保真因果沙箱不会机械记忆真实患者敏感信息，可作为合规安全的开放基准发布。

局限与展望¶

低维解析设定的表达局限：因果控制函数 $\Psi$ 依赖研究者预先指定的函数形式，虽然直观且易于收敛，但在真实高维强非线性交互机制下可能遗漏深层因果拓扑结构。
小样本场景下的拟合稳定性瓶颈：作为深层神经网络架构，CausalMix 依赖充足的观测数据来稳定学习混合分布解码器；当表格样本量极小且协变量维度极高时，复合因果惩罚可能出现优化超平面震荡。
静态干预与完整观测假设：当前框架仅支持静态点暴露与无缺失表格，尚未直接建模随时间变化的动态干预序列、时变混杂以及临床纵向随访中的非随机缺失机制。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次在深度变分框架中严密统一了复杂多模态表格数据拟合与全因子因果控制机制。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖三类渐进因果场景、十种前沿估计器基准评测、超参调优网格及前瞻性真实临床功效分析。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严谨，从统计理论推导、神经架构构建到真实临床案例落地环环相扣。
价值: ⭐⭐⭐⭐⭐ 为整个因果推断领域提供了一个急需的、兼具真实复杂性与因果确定性的高质量方法论评估沙箱。