CauScale:大规模神经因果发现
机器学习
2026-02-10 v1 人工智能
机器学习
摘要
因果发现是推动科学 AI 和数据分析等数据驱动领域的关键任务,但现有方法在扩展到大型图时面临显著的时间和空间效率瓶颈。为解决这一挑战,我们提出 CauScale,一种用于高效因果发现的神经网络架构,能够将推理扩展到最高达 1000 个节点的图上。CauScale 通过压缩数据嵌入的压缩单元提高了时间效率,通过采用绑定注意力权重来避免维护轴特定注意力图,从而提高了空间效率。为保持高水平的因果发现准确性,CauScale 采用双流设计:数据流从高维观测中提取关系证据,而图流整合统计图先验并保留关键结构信号。CauScale 在训练中成功扩展到 500 节点图,Prior work 由于空间限制无法实现。跨不同图规模和因果机制的测试数据,CauScale 在分布内数据上实现 99.6% 的平均精确率 (mAP),在分布外数据上实现 84.4%,同时在推理速度上相对于先前方法实现 4-13000 倍的加速。我们的项目页面位于 https://github.com/OpenCausaLab/CauScale。
引用
@article{arxiv.2602.08629,
title = {CauScale: Neural Causal Discovery at Scale},
author = {Bo Peng and Sirui Chen and Jiaguo Tian and Yu Qiao and Chaochao Lu},
journal= {arXiv preprint arXiv:2602.08629},
year = {2026}
}