[论文解读] Diffusion Models for Causal Discovery via Topological Ordering

Pedro L. Sánchez, Xiao Liu+2 more|arXiv (Cornell University)|Oct 12, 2022

Bayesian Modeling and Causal Inference被引用 4

一句话总结

该论文提出 DiffAN，一种可扩展的因果发现方法，利用去噪扩散概率模型（DPMs）估计加性噪声模型中数据对数似然的海森矩阵，以实现拓扑排序。通过一次训练神经分数模型，并利用无需重新训练的“落叶分数”机制高效更新，DiffAN 实现了对高达 500 个变量和 10⁵ 个样本的高维数据集的因果发现，在可扩展性方面优于先前方法，同时在小样本数据集上保持了最先进水平的准确性。

ABSTRACT

Discovering causal relations from observational data becomes possible with additional assumptions such as considering the functional relations to be constrained as nonlinear with additive noise (ANM). Even with strong assumptions, causal discovery involves an expensive search problem over the space of directed acyclic graphs (DAGs). \emph{Topological ordering} approaches reduce the optimisation space of causal discovery by searching over a permutation rather than graph space. For ANMs, the \emph{Hessian} of the data log-likelihood can be used for finding leaf nodes in a causal graph, allowing its topological ordering. However, existing computational methods for obtaining the Hessian still do not scale as the number of variables and the number of samples increase. Therefore, inspired by recent innovations in diffusion probabilistic models (DPMs), we propose \emph{DiffAN}\footnote{Implementation is available at \url{https://github.com/vios-s/DiffAN} .}, a topological ordering algorithm that leverages DPMs for learning a Hessian function. We introduce theory for updating the learned Hessian without re-training the neural network, and we show that computing with a subset of samples gives an accurate approximation of the ordering, which allows scaling to datasets with more samples and variables. We show empirically that our method scales exceptionally well to datasets with up to $500$ nodes and up to $10^5$ samples while still performing on par over small datasets with state-of-the-art causal discovery methods. Implementation is available at https://github.com/vios-s/DiffAN .

研究动机与目标

• 解决高维数据集（变量和样本数量众多）中因果发现的可扩展性瓶颈。
• 克服大规模设置下基于核方法的海森矩阵估计在计算上的不可行性。
• 通过去噪扩散模型学习分数函数，实现在拓扑排序中的高效计算。
• 开发一种方法，在每次叶节点移除后无需重新训练即可更新分数函数，从而保持效率。
• 在小样本数据集上实现最先进性能的同时，实现对最多 500 个节点和 10⁵ 个样本的大规模问题的可扩展性。

提出的方法

• 训练去噪扩散概率模型（DPM），以学习数据分布的分数函数，从而通过反向传播实现高效的海森矩阵计算。
• 利用 DPM 的分数及其完整的雅可比矩阵，计算对数似然的海森矩阵，以识别因果图中的叶节点。
• 提出“落叶分数”框架，通过解析方式在移除叶节点后更新分数函数，而无需重新训练神经网络。
• 对数据子样本应用更新后的分数函数进行海森矩阵评估，实现在样本规模上恒定时间复杂度，从而实现可扩展性。
• 通过迭代识别并移除基于海森矩阵对角元素的叶节点，利用拓扑排序重建因果DAG。
• 利用DPM能够泛化到分布外样本点并支持子采样的特性，这与基于核的方法不同。

实验结果

研究问题

• RQ1 去噪扩散模型能否有效用于估计数据对数似然的海森矩阵以实现因果发现？
• RQ2 在每次叶节点移除后，能否在不重新训练神经网络的情况下高效更新分数函数？
• RQ3 所提出的方法是否能够扩展到包含最多 500 个变量和 10⁵ 个样本的数据集，而现有方法在这些规模下失效？
• RQ4 该方法在小样本数据集上的性能与最先进因果发现算法相比如何？
• RQ5 在海森矩阵计算中使用子采样是否能在显著降低计算成本的同时保持准确性？

主要发现

• DiffAN 可扩展至包含最多 500 个变量和 10⁵ 个样本的数据集，而现有方法如 SCORE 在 64GB 内存机器上超过 2,000 个样本时即失效。
• 该方法在小样本数据集（如 50 个变量）上保持了与最先进方法相当的性能，展示了具有竞争力的准确性。
• 落叶分数机制使得海森矩阵计算在样本规模上具有恒定复杂度，从而在每次叶节点移除后实现高效更新。
• 使用子样本进行海森矩阵评估可实现准确的拓扑排序，验证了近似策略的有效性。
• 基于 DPM 的分数估计方法能够泛化到样本外点，而基于核的方法则需要完整的核矩阵。
• 该方法避免了对神经网络架构的约束，而这是先前基于深度学习的因果发现方法所不具备的。