跳转至

learning causal orderings for in-context tabular prediction

会议: NeurIPS 2026
arXiv: 2605.22335
领域: 因果推理
关键词: 表格基础模型, 上下文学习, 因果发现, 因果序约束注意力, 缺失值插补

一句话总结

提出了 TABORDER 框架,通过无监督的似然最大化目标端到端推断因果拓扑序,并在表格预测中引入因果序约束的列注意力机制与缺失感知不确定性建模,从而在分布偏移和机制干预下保持鲁棒预测。

研究背景与动机

表格数据上的上下文学习(In-Context Learning, ICL)和表格基础模型(Tabular Foundation Models, 如 TabPFN 等)在观测数据预测上取得了显著成功。然而,现有的表格基础模型本质上是纯统计关联驱动的,在做单元格预测或目标回归时往往利用所有可用特征作为预测源,并未显式建模特征之间的因果方向性。当系统处于稳态观测分布时,这类非因果统计关联可能非常有效;但一旦发生分布偏移或底层因果机制遭受外部干预,非因果反向关联就会发生不可控的漂移,导致预测性能大幅度退化。

解决该问题的经典思路是因果发现(Causal Discovery),旨在从观测数据中重构有向无环图(DAG)或因果拓扑序。然而传统因果发现算法(如 CAM、NOTEARS、TOPIC)多针对单一数据集求解定制的结构优化问题,计算复杂度高且与下游预测架构完全解耦;而近年来基于合成先验的有监督摊销因果发现(如 AVICI)虽能跨数据集快速推断,却极度依赖预先生成的真实因果图标签,并受限于合成先验的覆盖范围。此外,现有因果发现方法大都假定数据完全可观测,而现实表格普遍存在样本缺失,如何将因果序发现与表格缺失值填补及鲁棒预测结合仍属空白。

本文的核心切入点在于:在表格 Transformer 架构内将因果拓扑序的学习与条件分布建模深度融合,利用加性噪声模型中因果方向天然具备更低条件残差方差(更高对数似然)的信息论性质,在无监督预测目标下自发诱导因果序。核心 idea:设计 TABORDER 双分支架构,用直通松弛排序端到端推断变量因果序,通过因果序约束的单向列注意力切断反向伪相关,并利用缺失因果父节点引发的方差增量机制放大因果似然不对称性,实现无监督因果序发现与干预鲁棒预测的统一。

方法详解

整体框架

TABORDER 是一个端到端双分支 Transformer 架构,输入为一个可能包含缺失值的表格 \(D \in \mathbb{R}^{n \times d}\) 以及对应的二值缺失掩码 \(M \in \{0, 1\}^{n \times d}\)。系统由因果序推断分支与因果序约束预测分支交替堆叠而成:前者在无约束注意力下提取全局表征,并由 MLP 解码出每个变量的标量排序分数;后者根据该分数构建上三角因果注意力掩码,强制列注意力仅流向因果前驱变量,同时解耦预测均值与缺失父变量引起的方差增量。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入表格数据与缺失掩码<br/>D, M ∈ R^(n×d)"] --> B["初始单元格线性映射<br/>拼接取值与缺失标记嵌入"]
    B --> C["因果序推断分支<br/>行/列交替无约束注意力提取表征"]
    C --> D["因果序分数解码与排序<br/>按行平均得到变量排序分数 si"]
    D --> E["因果序约束注意力掩码<br/>前向硬阶跃与反向 Sigmoid 直通松弛"]
    B --> F["因果序约束预测分支<br/>行注意力无约束,列注意力受掩码单向因果约束"]
    E -.->|注入注意力掩码 As| F
    F --> G["缺失感知的自适应不确定性建模<br/>解耦条件均值预测与前驱缺失方差增量"]
    G --> H["输出各缺失单元格预测分布与插补值"]

关键设计

1. 因果序推断与可微直通松弛:端到端无监督学习拓扑变量顺序

传统的因果排序方法依赖组合搜索或贪心回归,无法与神经网络反向传播直接兼容。因果序推断分支采用无约束的 Transformer 架构,交替进行行自注意力和列自注意力操作,聚合跨样本与跨特征的统计依赖关系,生成隐状态 \(H^{\text{ord}} \in \mathbb{R}^{n \times d \times h}\)。随后,通过轻量级两层多层感知机将每个单元格状态映射为标量,并在样本维度做平均,得到每个变量 \(X_i\) 的排序标量分数 \(s_i\):

\[s_i(D) = \frac{1}{n} \sum_{r=1}^n f_{\text{ord}}\left(H^{\text{ord}}_{r,i}\right)\]

标量 \(s_i\) 越小代表该变量在因果拓扑序中的位置越靠前。在前向计算中,模型将分数转换为严格的二值因果注意力掩码 \(A^s_{i,j} = \mathbb{I}(s_j \le s_i)\),使得变量 \(i\) 只能关注到排序分小于等于自己的前驱变量;在反向传播中,采用温度为 \(\tau\) 的 Sigmoid 连续松弛 \(\hat{A}^s_{i,j} = \sigma\left(\frac{s_i - s_j}{\tau}\right)\) 替代阶跃函数,构建直通估计器(Straight-Through Estimator)。该设计使得整个因果排序过程完全可微,梯度可直接从下游预测损失反传至排序分数。

2. 因果序约束预测注意力:切断反向统计伪相关

常规表格 Transformer 的全双向特征注意力容易过度利用下游后代变量与目标之间的反向统计依赖(Anti-Causal Correlation),这在因果干预发生时会导致剧烈预测崩溃。为了消除该缺陷,因果序约束预测分支虽然保留行自注意力以在样本间共享上下文线索,但在列自注意力计算中强加由前序分支传递来的因果掩码 \(A^s\)。

该掩码保证了每个变量的隐表示 \(H^{\text{pred}}_{r,i}\) 只能由其在因果拓扑序中的直接或间接父节点决定,从而在架构层面上严格诱导了多元联合分布的自回归因果因子分解:

\[p_\theta(x^{(r)} \mid D, \pi) = \prod_{k=1}^d p_\theta\left(x^{(r)}_{\pi(k)} \;\middle|\; x^{(r)}_{\pi(1:k-1)}, D\right)\]

若模型推断出错误的因果逆序,将导致下游变量无法获得真实因果前驱的充分条件信息,进而承受更高的残差方差与更低的对数似然;只有当推断出的拓扑序完全契合真实因果图的拓扑排序时,联合预测似然才能达到最大化。

3. 缺失感知的自适应不确定性建模:利用因果父节点缺失放大似然不对称性

在现实表格中,待预测特征的因果父节点常常自身也处于缺失状态,导致真实条件方差动态膨胀。TABORDER 假设加性噪声模型 \(X_i = f(X_{\text{pa}(i)}) + N_i\)(其中 \(N_i \sim \mathcal{N}(0, \sigma_i^2)\)),并通过双头解码网络分别预测基础条件均值 \(\hat{\mu}_i^{(r)}\) 和基线条件方差 \(\hat{\sigma}_i^2\)。针对第 \(r\) 行中前驱因果变量可能缺失的现象,模型显式引入了可学习的非负方差增量参数 \(\delta_{i,j} \in \mathbb{R}^+\):

\[\hat{\sigma}^{2(r)}_i = \hat{\sigma}^2_i + \sum_{j=1}^d \delta_{i,j} \cdot \mathbb{I}(s_j < s_i) \cdot m^{(r)}_j\]

该设计的深刻理论意义在于放大因果方向的似然差异(Theorem 1)。由于残差方差的负对数函数具有严格凸性,考虑因果前驱缺失的方差增量模型所获得的对数似然增益在正向因果关系 \(X \to Y\) 中严格高于反向关系 \(Y \to X\),充分必要条件为正向信噪比 \(\frac{\text{Var}(f(X))}{\text{Var}(N)}\) 超过反向统计比率 \(\frac{\text{Var}(\mathbb{E}[X \mid Y])}{\mathbb{E}[\text{Var}(X \mid Y)]}\)。这一机制使样本缺失非但不会干扰因果发现,反而在部分观测下显著拉大了正确因果序与错误因果序之间的似然差距。

损失函数 / 训练策略

TABORDER 采用端到端无监督的最大对数似然准则进行训练,完全不需要有向无环图的真值监督。在训练阶段,从数据生成分布中随机采样表格 \(D \sim \mathcal{P}\),并以 \(q = 20\%\) 的概率随机将单元格掩码为缺失。目标函数为最大化全部被掩码单元格在因果高斯条件分布下的对数似然:

\[\mathcal{L} = \frac{1}{|M|} \sum_{m^{(r)}_i = 1} \log \left[ p_{\mathcal{N}}\left( x^{*(r)}_i \;\middle|\; \hat{\mu}^{(r)}_i, \hat{\sigma}^{2(r)}_i \right) \right]\]

优化过程中,因果序推断与因果约束预测两组 Transformer 的所有参数联合更新。掩码前向偏置 \(\beta\) 从 \(-5.0\) 退火至 \(-20.0\),反向松弛温度 \(\tau\) 从 \(1.0\) 退火至 \(0.1\)。模型特征嵌入维度设为 128,推断分支包含 6 个重复块,预测分支包含 4 个重复块,使用 Adam 优化器在 100,000 个合成数据集上完成预训练。

实验关键数据

主实验

论文在 Sachs 真实单细胞信号通路基准(含 11 个变量,真值有向图对应细胞信号传导途径)以及真实世界表格基准 OpenML CC18(33 个分类数据集)和 CTR23(23 个回归数据集)上展开了因果拓扑序推断与缺失插补的主实验评估。拓扑偏离度指标 \(d_{\text{TOP}}\) 越小表示推断排序与真值图结构的不一致边越少;在缺失值插补中以平均排名衡量算法相对优劣。

数据集 / 基准 评估任务与指标 TABORDER (本文) 经典最优基线 (CAM / HYPERIMPUTE) 对比差值 / 表现
Sachs (cd3cd28 观测态) 因果拓扑序偏离度 \(d_{\text{TOP}}\) (↓) 0.206 (违背3条边) CAM: 0.353 / TOPIC: 0.412 \(d_{\text{TOP}}\) 显著降低超 40%
Sachs (全条件平均) 平均拓扑偏离度 \(d_{\text{TOP}}\) (↓) 0.26 CAM: 0.56 / NOTEARS: 0.54 远优于所有传统因果发现算法
合成非线性高斯过程 (ANM) 因果拓扑序偏离度 \(d_{\text{TOP}}\) (↓) ~0.18 TOPIC: ~0.17 / AVICI: ~0.20 与专业因果结构学习算法相当
CTR23 表格插补 (缺失率 \(q \ge 0.4\)) 插补平均排名 Avg. Rank (↓) 1.8 (最优) HYPERIMPUTE: ~2.8 / MISSFOREST: ~3.8 高缺失率下夺得第一
CC18 表格插补 (缺失率 \(q \ge 0.4\)) 插补平均排名 Avg. Rank (↓) 2.2 (最优) HYPERIMPUTE: ~2.6 / MISSFOREST: ~3.5 高缺失率下持续超越树模型与神经基线

消融实验

消融实验重点验证了:(1) 在不同噪声结构下模型的拓扑序鲁棒性;(2) 错误因果拓扑序强制替换对负对数似然(NLL)的影响;(3) 面对因果机制干预时相较于黑盒基础模型与梯度提升树的抗扰动能力。

模型配置 / 实验场景 关键评估指标 实验说明与对比分析
加性噪声 (Additive Noise, GP) \(d_{\text{TOP}} = 0.202\) 符合理论设定的标准加性高斯分布场景
异方差噪声 (Heteroskedastic Noise, GP) \(d_{\text{TOP}} = 0.190\) 噪声方差依赖父节点,模型拓扑序表现依旧稳健
乘性噪声 (Multiplicative Noise, GP) \(d_{\text{TOP}} = 0.232\) (非加性GP为 0.172) 显著偏离加性噪声假设,非加性版本具备更好鲁棒性
测试时强制替换为随机错误因果序 \(\text{NLL}\) 从 0.08 飙升至 0.28+ 似然值随 \(d_{\text{TOP}}\) 增大呈单调急剧恶化,证明因果序对预测至关重要
链式结构 \(X \to Y \to Z\) 干预机制迁移 预测 \(Y\) 的 MSE 维持在 ~0.08 TABPFN 与 XGBoost 因依赖反向 \(Z \to Y\) 关联,MSE 暴增至 0.70+

关键发现

  • 因果序推断的无监督精准性:在合成数据与 Sachs 真实蛋白质通路数据集上,TABORDER 仅依赖最大对数似然目标便重构出与真实 DAG 吻合度极高的拓扑序(Sachs 基准上拓扑偏离度低至 0.206,仅违背 3 条边),表现大幅超越了包括 CAM、NOTEARS、AVICI、SCORE 等在内的专业因果发现算法。
  • 高缺失率下的插补优势:在缺失率 \(q \ge 40\%\) 的极端挑战下,传统基于迭代或局部回归的插补方法(如 MICE、KNN)误差急剧累积,而 TABORDER 凭借因果前驱注意力约束与方差自适应增量机制,在 CTR23 和 CC18 上均夺得平均排名第一。
  • 面对因果机制干预的绝对鲁棒性:在 \(X \to Y \to Z\) 的因果链测试中,当对下游变量 \(Z\) 引入机制干预时,TabPFN 和 XGBoost 由于吸收了 \(Z\) 对 \(Y\) 的伪关联导致预测崩溃(MSE > 0.7),而 TABORDER 严格依据因果序仅允许由原因 \(X\) 预测结果 \(Y\),在干预前后预测误差几无波动。

亮点与洞察

  • 无监督因果序驱动的结构化上下文预测:打破了因果发现与监督预测割裂的传统格局,无需真值 DAG 监督,直接通过自回归因果因子分解的对数似然最大化诱导因果拓扑序,理论严谨且具备完全端到端的训练特性。
  • 以缺失作为因果方向辨识的放大镜:巧妙揭示并严格证明了部分观测和样本缺失在因果识别中的积极作用——父节点缺失引发的自适应方差增量通过詹森不等式放大了正反向似然差异,颠覆了“缺失必然恶化因果发现”的直觉偏见。
  • 真实生物系统干预的拓扑敏感性解析:将模型应用于 Sachs 单细胞受试数据时,模型不仅精准捕获了参考状态拓扑序,且在面对不同分子抑制剂干预时,下游靶点(如 PIP2 扰动)表现为极度局部的拓扑序微调,而枢纽上游靶点(如 PKC 抑制)则自发呈现全局拓扑序重排,具备极高的科学解释力。

局限与展望

  • 连续实值特征限制与复杂机制依赖:当前架构与理论证明主要锚定在实值连续特征与加性噪声模型(ANM)框架下,尚未直接原生支持高基数类别变量与离散混合类型表格。
  • 拓扑序向稀疏图结构精化的缺失:TABORDER 输出的是因果拓扑序而非完全稀疏的有向无环图(DAG),若要剔除拓扑序中的非直接因果边,目前仍需依赖后处理稀疏回归剪枝。
  • 预训练零样本迁移能力仍需拓展:纯合成数据预训练的权重在现实真实表格上直接零样本推断时无法全面压制深度调优的传统插补器,必须经过目标领域的微调训练,未来亟需设计更丰富多样的高真实感合成因果先验分布。

相关工作与启发

  • vs TabPFN / TabICL: TabPFN 开启了表格上下文学习先河,但属于纯关联驱动模型,特征间无向注意力容易在分布偏移与机制干预下引入脆弱的反向关联;TABORDER 显式引入因果序注意力掩码,从机理上杜绝了对非因果统计线索的过度拟合。
  • vs CAM / TOPIC / NOTEARS: 经典因果发现方法采用独立的数据集级数值优化或非线性回归搜索,计算开销随特征数剧烈增长(如 CAM 在特征数达到 50 时无法扩展);TABORDER 前向推断接近实时,并在单次推理中同步完成因果序输出与缺失值条件分布预测。
  • vs AVICI: AVICI 采用有监督摊销学习,必须在包含大量真实因果图标签的合成数据集上训练,受先验偏差影响显著;TABORDER 坚持无监督自回归似然准则,无须任何图标签即可泛化到真实未知领域。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [将无监督因果序推断与因果序约束列注意力有机统一于表格上下文学习模型中,立意极为新颖]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖合成数据、真实世界生物因果基准 Sachs 及 OpenML 56 个回归与分类数据集,兼顾因果序评估、插补与干预鲁棒性]
  • 写作质量: ⭐⭐⭐⭐⭐ [理论推导严密自洽,定理证明清晰,实验图表层次分明]
  • 价值: ⭐⭐⭐⭐⭐ [为构建具备真正因果推理与干预鲁棒性的下一代表格基础模型指明了关键技术范式]