Causality-Encoded Diffusion Models for Interventional Sampling and Edge Inference

TL;DR

This paper introduces a causality-encoded diffusion framework that integrates known causal graphs into diffusion models, enabling interventional sampling, causal inference, and edge testing with theoretical guarantees.

Contribution

It develops a novel diffusion-based method that incorporates causal structures, supporting interventional analysis and edge inference with convergence guarantees.

Findings

Improved recovery of interventional distributions over baselines.

Theoretical guarantees for distribution estimation and edge testing.

Successful application to flow cytometry data for causal inference.

Abstract

Standard diffusion models are flexible estimators of complex distributions, but they do not encode causal structures and therefore do not by themselves support causal analysis. We propose a causality-encoded diffusion framework that incorporates a known directed acyclic graph by training conditional diffusion models consistent with the graph factorisation. The resulting sampler approximately recovers the observational distribution and enables interventional sampling by fixing intervened variables while propagating effects through the graph during reverse diffusion. Building on this interventional simulator, we develop a resampling-based test for directed edges that generates null replicates under a candidate graph. We establish convergence guarantees for observational and interventional distribution estimation, with rates governed by the maximum local dimension rather than the ambient dimension, and prove asymptotic control of type I error for the edge test. Simulations show improved interventional distribution recovery relative to baselines, with near-nominal size and favourable power in inference. An application to flow cytometry data demonstrates practical utility of the proposed method in assessing disputed signalling linkages.