Causality-Encoded Diffusion Models for Interventional Sampling and Edge 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.

• The paper introduces CEDM which combines diffusion models with DAG causal structures, enabling parent-local conditioning and principled do-interventions.
• It achieves dimension reduction by scaling recovery rates with the maximum local parent-child dimension, thereby enhancing statistical and computational efficiency.
• CEDM-based Inference (CEDMI) robustly tests edge presence using resampling-based hypothesis tests, outperforming baseline methods in both simulations and biological networks.

Summary of "Causality-Encoded Diffusion Models for Interventional Sampling and Edge Inference" (2604.21843)

Motivation and Formulation

Standard diffusion models (DMs) offer state-of-the-art performance for generative modeling but are fundamentally causality-agnostic. They lack the ability to encode directed acyclic graph (DAG)-based causal structures of multivariate systems, which precludes principled intervention, counterfactual estimation, and edge-level inference. The paper introduces Causality-Encoded Diffusion Models (CEDM), an iterative generative construction that intertwines DMs with known causal DAG factorization, enforcing topological ordering and parent-local conditioning during reverse diffusion. Each node's reverse SDE is conditioned only on its parent nodes, replacing the global score function with a collection of parent-local conditional scores. This design aligns generative sampling with the causal mechanisms specified by structural causal models (SCM).

Concretely, CEDM operates by:

• Decomposing the observed data distribution into nodewise factors respecting DAG structure;
• Training conditional diffusion models per node, guided only by its parents;
• Sampling by sequentially running nodewise conditional reverse SDEs in topological order;
• Facilitating do-intervention sampling by freezing injected variables and propagating interventions downstream.

  Figure 1: Simulation results on the Chain, Hub, Random, and Sachs structures, showing squared MMD between model-generated and true interventional samples.

Statistical Guarantees and Dimension Reduction

A central theoretical contribution is demonstrating that recovery rates for both observational and interventional distributions scale with the maximum local parent-child dimension $q = \max_j (d_j + r_j)$, rather than the global ambient dimension $d$. This structural locality enables principled dimension reduction, improving statistical and computational efficiency especially for sparse graphs. The paper delivers explicit bounds on total variation error and expectation recovery for both observational and do-interventional laws, with rates $\mathcal{O}(n^{-\beta / (q+2\beta)})$, where $\beta$ captures score function regularity.

The theoretical analyses leverage nodewise score approximation, denoising score matching, and KL-divergence bounds, linking the expressivity and estimation errors of neural conditional score networks to the structure of the underlying SCM. The intervention-specific distribution shift coefficient $\mathcal{T}(\mathbf{x}_{\mathcal{S}})$ quantifies support mismatch and governs error amplification under poorly-supported interventions.

Edge Inference via CEDMI

Building on the generative simulation capabilities of CEDM, the authors propose CEDM-based Inference (CEDMI), a resampling-based hypothesis test for edge presence in a working DAG. The procedure encodes the null hypothesis by deleting candidate edges, resamples the relevant child node under do-interventions holding parents fixed, and quantifies residual conditional dependence using an aggregated rank-based multivariate coefficient (MCODEC), an extension of Azadkia-Chatterjee's statistic.

CEDMI computes p-values by comparing observed conditional dependence to synthetic null replicates, achieving asymptotic control of Type I error through sample-splitting and exchangeability, with strong consistency under alternatives.

Figure 2: Empirical rejection rates at level 0.05, contrasting CEDMI and baseline conditional independence tests under chain and fork null configurations.

Empirical Validation

Extensive simulations demonstrate CEDM's superiority over causality-agnostic diffusion baselines (e.g., RePaint) and graph-based variational autoencoders (VACA) in recovering interventional distributions, particularly under sparse causal structures (Chain and Hub DAGs), shown via squared MMD metrics. In inference, CEDMI attains near-nominal test size and strong statistical power, outpacing classifier-based, kernel-based, and generative-model-based conditional independence tests, especially in high-dimensional nonlinear regimes with non-additive noise.

Figure 3: Empirical power at level 0.05 for edge tests, with power increasing in signal strength and sample size.

Application: Flow Cytometry Signalling Network

CEDMI is applied to the canonical flow cytometry dataset of Sachs et al., comprising 11 protein/phosphorylation nodes, comparing literature-based consensus and empirically reconstructed Bayesian networks. Focusing on edges disputed between the two networks, the authors construct a super-DAG and formally test the inclusion of specific linkages using CEDMI, NNLSCIT, CDCIT, and CCIT. The methods consistently reject several consensus-network edges, while offering robust support for $\text{Erk} \to \text{Akt}$, which aligns with established cross-talk findings in mammalian signal transduction literature.

Figure 4: Rejection rates for edge tests across disputed linkages in the flow cytometry network for CEDMI and baselines as a function of sample size.

Implications and Future Directions

Practically, CEDM provides a tractable, statistically efficient approach for generative modeling and causal intervention in high-dimensional graphical systems, particularly suitable for complex networks with known or hypothesized structure. Theoretical guarantees reveal how structural properties drive complexity, not just geometry or ambient dimension. CEDMI broadens the scope of edge inference beyond linear/Gaussian or additive-noise settings, enabling robust hypothesis testing in arbitrary nonlinear causal DAGs.

Theoretically, the locality of learning via parent-conditioning interfaces diffusion models with the SCM/causal factorization paradigm. Future developments may explore DAG discovery by coupling structure learning with local score estimation, formal uncertainty quantification, and fully Bayesian treatment. Given the reliance on known or working graphs, further work is needed to characterize robustness to misspecification and extend theory to adaptive, structure-learning regimes.

Conclusion

This paper advances the field of causal generative modeling by integrating diffusion processes with explicit DAG structure. The CEDM framework aligns generative simulation and intervention with the semantics of SCMs, enhances distribution recovery through structural dimension reduction, and yields a powerful statistical tool for directed edge inference. Empirical results affirm both generative accuracy under intervention and calibrated inference capabilities, validated in synthetic and real biological networks. CEDM and CEDMI thus provide a foundation for structurally-informed deep generative and inferential methods, with wide-ranging applicability and future potential in scalable, causal AI systems.