CauScale：大规模神经因果发现 CauScale: Neural Causal Discovery at Scale
通过压缩数据嵌入和共享注意力权重，实现高效的大规模因果发现
前置知识
因果图模型（Causal Graphical Model）
因果图模型是一种用于表示变量间因果关系的概率图模型。它由联合分布 $P_X$ 和有向无环图 $G=(V,E)$ 组成，其中节点对应变量，边表示直接因果影响。分布满足马尔可夫性质：$p(x_1,...,x_n) = \prod_{j=1}^n p(x_j | PA_j)$，其中 $PA_j$ 是节点 $j$ 的父节点集。这种表示允许我们通过干预来修改特定变量的条件分布，从而研究因果效应。
本文的目标是从观测数据中学习这种因果图结构，理解因果图模型是理解整个方法的基础
摊销因果发现（Amortized Causal Discovery）
摊销因果发现是一种零样本推理方法：先在大量模拟数据上预训练一个神经网络模型，使其学习从数据集到因果图的映射。测试时，模型可以直接对新数据集进行预测，无需针对每个数据集重新优化。这种方法避免了传统方法中为每个数据集求解独立优化问题的高计算成本，特别适合大规模或实时推理场景。
CauScale 采用摊销方法范式，但解决了现有摊销方法（如 AVICI）在大规模图上的内存瓶颈问题
轴向注意力（Axial Attention）
轴向注意力是一种高效的注意力机制，将标准的全注意力分解为沿不同轴向的顺序注意力。在处理三维数据（如 $m \times n \times d$ 的数据嵌入）时，先沿样本维度做注意力，再沿节点维度做注意力。这种方法将计算复杂度从 $O((mn)^2)$ 降低到 $O(m^2n + mn^2)$，在保持表达能力的同时显著减少了计算量。
CauScale 的数据流和图流都使用轴向注意力作为核心计算单元，理解其机制是理解效率优化的关键
图先验（Graph Prior）
图先验是从数据中计算的统计信息，用于为模型提供关于变量间关系的初始线索。在本文中，图先验定义为数据的逆协方差矩阵：$\rho = (E[(D-\mu)(D-\mu)^\top])^{-1}$。逆协方差矩阵的非零元素对应偏相关关系，即在控制其他变量后两个变量仍然相关的程度，这为因果结构提供了有价值的线索。
图先验是 CauScale 的重要输入，消融实验证明它对模型性能有显著贡献
研究动机
现有因果发现方法在处理大规模图时面临严重的效率瓶颈。约束-based 方法（如 PC 和 FCI）依赖大量条件独立性检验，检验次数在最坏情况下呈指数增长。基于分数的方法（如 NOTEARS 和 RL-BIC）虽然避免了显式组合搜索，但需要为每个数据集求解新的连续优化问题，计算成本仍然很高。最近的摊销方法（如 AVICI）通过预训练模型实现了零样本推理，但其注意力机制的内存复杂度与变量数量呈不良的缩放关系，导致在大规模图上出现内存压力。具体而言，AVICI 在处理 $n=100$ 个节点的 SERGIO-GRN 数据时就会出现内存溢出错误，这严重限制了其在实际应用中的可扩展性。
本文的目标是本文的目标是开发一种能够高效处理大规模因果发现的神经架构，具体目标包括：（1）将训练扩展到包含 500 个节点的图，这是现有方法无法达到的规模；（2）实现对包含 1000 个节点的图的推理，这是训练规模的两倍；（3）在保持高因果发现准确率的同时，实现比现有方法快 4 到 13000 倍的推理速度；（4）在分布外（OOD）数据上保持鲁棒的性能，目标是在 OOD 机制下仍能达到 50% 以上的 mAP。
与已有工作不同的是，本文的独特切入角度在于同时解决时间和空间效率问题，而不是只关注其中一个方面。现有方法要么优化时间效率但忽略内存（如一些基于优化的方法），要么优化内存但牺牲时间效率（如 SEA 依赖经典估计器和大量子批次采样）。CauScale 通过两个关键设计填补了这一空白：一是使用缩减单元（Reduction Unit）压缩数据嵌入以节省计算时间，二是采用绑定注意力权重（Tied Attention Weights）避免维护轴特定的注意力图以大幅减少内存占用。更重要的是，本文认识到在压缩数据嵌入时可能丢失关键的因果信号，因此设计了数据-图块（Data-Graph Block）来在压缩前将关系证据提炼到图流中，从而在效率和准确性之间取得平衡。
核心方法
CauScale 的核心思想可以用'压缩但不丢信息'来类比。想象你要分析一个包含 1000 个变量的大规模数据集，传统方法需要为每个样本维护完整的变量关系，这会导致内存爆炸。CauScale 的做法是：先让数据流和图流并行工作，数据流从原始观测中提取关系证据，图流维护变量间的结构信息；然后通过缩减单元逐步压缩数据流的表示，但在这个压缩之前，数据-图块会将关键的关系信息提炼到图流中保存。这样，随着网络层数的增加，数据流的表示越来越紧凑（计算量减少），但关键信息通过图流得以保留。技术路线上，模型接收观测数据 $D \in \mathbb{R}^{m \times n \times 2}$（包含观测值和干预指示）和图先验 $\rho \in \mathbb{R}^{n \times n}$，通过线性层编码为初始嵌入，然后经过交替堆叠的数据-图块和缩减单元处理，最终由预测头输出概率邻接矩阵 $\hat{G} \in \mathbb{R}^{n \times n}$。
CauScale 最核心的创新是其'两流一压缩'的设计理念。两流指的是数据流（Data Stream）和图流（Graph Stream）的并行处理：数据流通过轴向注意力从高维观测中提取变量间的成对关系，图流则整合统计图先验并维护关键的结构信号。一压缩指的是缩减单元（Reduction Unit）：每隔 $k$ 个数据-图块，沿样本维度对数据嵌入进行平均池化，将其长度缩减 $r$ 倍。关键在于，这种压缩不是简单地丢弃数据，而是在压缩之前，数据-图块中的数据到图层（Data2Graph Layer）会将数据流中的关系证据提炼成图消息 $\omega^{D \to G}_b \in \mathbb{R}^{n \times n}$，注入到图流中。这样，即使数据流被压缩，关键的因果信号仍然保留在图流的表示中。与 AVICI 相比，这种设计将注意力图的内存从 $O(RHC^2)$ 降低到 $O(HC^2)$，同时通过压缩将样本维度的计算量减少到基线的 36.60%。
方法步骤详情
CauScale 的方法步骤如下：（1）输入编码：将观测数据 $D \in \mathbb{R}^{m \times n \times 2}$ 和图先验 $\rho \in \mathbb{R}^{n \times n}$ 通过线性层分别编码为数据嵌入 $h^D \in \mathbb{R}^{m \times n \times d}$ 和图嵌入 $h^G \in \mathbb{R}^{n \times n \times d}$，其中 $d$ 是嵌入维度。（2）数据-图块处理：每个块包含三层——数据层使用轴向注意力更新数据嵌入；数据到图层通过两个 PoolingFFN 模块将数据嵌入压缩为节点级嵌入 $u^{D \to G}, v^{D \to G} \in \mathbb{R}^{n \times d}$，然后计算外积得到图消息 $\omega^{D \to G}_b = u^{D \to G} (v^{D \to G})^\top$；图层将图消息与之前的图嵌入拼接，通过线性投影和图轴向注意力更新图嵌入。（3）缩减操作：每 $k$ 个块后，将数据嵌入沿样本维度分成大小为 $r$ 的块，对每个块进行平均池化，将样本数从 $m$ 缩减到 $m/r$。（4）预测：最终的图嵌入经过层归一化后，输入成对预测头：对每对无序节点 $\{i,j\}$，拼接两个方向的嵌入，通过 FFN 输出三个边状态（无边、$i \to j$、$j \to i$）的 logits，再通过 softmax 得到概率。
技术新颖性
CauScale 的技术新颖性体现在三个方面。首先是缩减单元的设计：与直接对原始数据采样不同，CauScale 在网络中间层压缩数据嵌入，此时原始输入已被转换为更信息丰富的表示，减少了信息损失。而且，压缩在数据-图块之后执行，确保关键关系信号已被提炼到图流中。其次是绑定注意力权重的应用：借鉴自蛋白质语言模型 MSA Transformer 的技术，CauScale 在轴向注意力中共享注意力权重，避免为每个轴存储单独的注意力图，将行轴注意力的内存从 $O(RHC^2)$ 降低到 $O(HC^2)$。最后是数据-图块的设计：它不仅更新数据流和图流，还在两者之间建立信息桥梁——数据到图层将数据流的关系证据提炼到图流，图层则将图信息注入数据流，这种双向信息流动是保持压缩后准确率的关键。
实验结果
CauScale 在多个实验设置下展现了卓越的性能。在合成数据集上，当 $n=100, |E|=400$ 时，CauScale 在线性机制下达到 99.6% 的 mAP、15.2 的 SHD、100.0% 的 AUC 和 100.0% 的 OA，显著优于所有基线方法；在神经网络非加性噪声机制下达到 89.0% mAP，比第二好的 SEA-gies（51.2%）高出 37.8 个百分点。在分布外（OOD）测试中，CauScale 在 sigmoid 机制下达到 84.4% mAP，在多项式机制下达到 50.3% mAP，均优于 SEA（72.7% 和 36.2%）和 SDCD（61.9% 和 41.9%）。在大规模图测试中（$n=1000, |E|=2000$），CauScale 在线性机制下达到 96.6% mAP，而 AVICI 仅为 0.2%。效率方面，CauScale 的推理时间在 $n=1000$ 时仅为 0.8288 秒，比 NOTEARS（10896 秒）快 13000 倍，比 SEA-gies（218 秒）快 200 倍，比 AVICI（3.34 秒）快 4 倍。在 SERGIO-GRN 数据集上，当 $n=200, |E|=400$ 时，CauScale 达到 39.2% mAP，而 AVICI 由于内存溢出无法运行，SEA-gies 仅为 1.2%。消融实验证明，移除缩减单元会导致内存溢出错误，移除图流会导致最显著的性能下降，验证了各组件的必要性。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 合成数据因果发现（线性，n=100） | mAP (%) | 99.6 | SEA-gies: 92.1 | +7.5 |
| 合成数据因果发现（NN，n=100） | mAP (%) | 89.0 | SDCD: 65.7 | +23.3 |
| 大规模图因果发现（线性，n=1000） | mAP (%) | 96.6 | SDCD: 54.1 | +42.5 |
| SERGIO-GRN 因果发现（n=200） | mAP (%) | 39.2 | SEA-gies: 1.2 | +38.0 |
| 推理效率（n=1000） | 推理时间 (s) | 0.8288 | AVICI: 3.34 | 4× 加速 |
局限与改进
本文存在以下局限性：首先，CauScale 的训练数据完全基于模拟，包括合成 SCM 数据和 SERGIO-GRN 基因调控网络数据。虽然实验涵盖了多种图结构（Erdős-Rényi、Scale-Free、Stochastic Block Model）和因果机制（线性、神经网络、sigmoid、多项式），但这些仍然无法完全覆盖真实世界数据的复杂性。论文承认，对模拟训练数据的依赖可能限制模型在数据生成过程严重偏离假设 SCM 框架的场景中的泛化能力。其次，当图规模超过训练分布（$n=1000$，训练时最大 $n=500$）且因果机制也为 OOD 时，模型会产生循环，需要后处理的循环打破策略。在最极端的多项式 OOD 情况下，平均需要移除 103 条边，这可能影响预测图的质量。第三，模型假设因果充分性（无未观测的共同原因），在存在潜在混杂因素时性能会下降。实验显示，当隐藏 64 个高外度节点时（$n=100$），线性机制的 mAP 从 99.6% 大幅下降。第四，模型对图先验的质量有一定依赖，特别是在大规模图上。当用于计算先验的样本数较少时，性能会下降。
独立分析的弱点
基于对论文的深入分析，我发现以下几个值得改进的弱点。首先，缩减单元采用固定的缩减因子 $r$ 和间隔 $k$，这可能不是最优的。在不同的网络层，数据嵌入的信息密度可能不同，早期层可能需要保留更多细节，而后期层可以更激进地压缩。改进方向是设计自适应的缩减策略，根据每一层的信息量动态调整压缩率。其次，绑定注意力权重虽然节省内存，但也限制了模型沿不同轴学习独立注意力模式的能力。在因果发现问题中，样本维度和节点维度的注意力模式可能有本质差异。改进方向是探索部分绑定或分组绑定的策略，在效率和表达能力之间取得更好的平衡。第三，当前的循环打破策略是启发式的（移除最低置信度的边），这可能不是最优的。可以设计端到端的方法，在训练时就鼓励无循环的预测，或者使用更复杂的后处理算法（如基于反馈弧集的方法）。第四，模型对干预数据的利用相对简单（仅作为二进制指示），可以探索更精细的干预信息编码方式，例如区分不同类型的干预或建模干预强度。
未来方向
论文和基于其成果可延伸的未来研究方向包括以下几个方面。首先，作者明确指出，在真实世界观测数据集上的评估是未来工作的重要方向。这需要收集或整理来自生物信息学、流行病学、经济学等领域的真实因果发现基准数据集。其次，可以探索将 CauScale 与其他因果发现范式结合，例如利用 Granger 因果关系处理时间序列数据，或结合功能因果模型处理非高斯数据。第三，当前模型需要预先计算图先验，这在某些场景下可能不可行。可以研究如何将先验的计算整合到端到端的学习框架中，或者开发不需要显式先验的变体。第四，可以探索 CauScale 在因果推断任务中的应用，例如在学习到的因果图上进行因果效应估计或反事实推理。第五，鉴于模型展现了对 OOD 图结构的良好泛化能力但对 OOD 机制泛化较弱，可以在训练时引入更多样的噪声分布和因果机制函数来增强鲁棒性。
复现评估
从复现性角度来看，CauScale 具有良好的可复现条件。代码已在 GitHub 上开源（https://github.com/OpenCausaLab/CauScale），这对于学术研究和实际应用都是重要优势。数据方面，论文使用了公开的 SERGIO-GRN 模拟器生成基因表达数据，合成数据的生成过程也有详细描述。算力需求方面，训练需要 8 块 NVIDIA H200 GPU（每块 141GB 内存），合成数据两阶段训练分别需要 37 小时和 2.75 小时，SERGIO-GRN 训练需要 44 小时。推理效率很高，在单 GPU 上处理 1000 节点图仅需不到 1 秒。复现难度中等：虽然代码开源且训练流程清晰，但大规模训练需要相当的计算资源。对于资源有限的研究者，可以考虑在较小规模（如 $n \leq 100$）上复现，或者使用论文提供的预训练模型进行推理。值得注意的是，论文提供了详细的超参数设置（10 层、128 维嵌入、16 个注意力头）和训练策略，这降低了调参的难度。
论文图表
