Causality-Encoded Diffusion Models for Interventional Sampling and Edge Inference

1. Problem Statement

Standard diffusion models are powerful nonparametric estimators capable of learning complex joint distributions via score matching and reverse-time stochastic differential equations (SDEs). However, they suffer from two critical limitations in causal inference contexts:

1. Causality-Agnostic Nature: They learn the joint law without encoding directional asymmetries (causal structure). Consequently, they cannot inherently distinguish between observationally equivalent graphs or support principled interventional queries (e.g., $p(x | do(X_j = x_j))$).
2. Curse of Dimensionality: Theoretical convergence rates for standard diffusion models typically scale with the ambient dimension ($d$) of the data. In high-dimensional settings with sparse causal structures, this leads to inefficient estimation and slow convergence.
3. Inference Challenges: Existing methods for testing directed edges in Directed Acyclic Graphs (DAGs) often rely on restrictive assumptions (e.g., Gaussian linearity) or lack robustness in general nonlinear, non-additive settings.

The paper aims to bridge this gap by integrating known causal structures (DAGs) into the diffusion framework to enable interventional sampling and develop a robust edge inference procedure.

2. Methodology: Causality-Encoded Diffusion Models (CEDM)

The authors propose CEDM, a framework that enforces causal factorization into the generative process.

A. Causal Factorization in Reverse Dynamics

Instead of learning a single joint score function $\nabla \log p_t(x)$, CEDM learns a set of local conditional score functions consistent with a given DAG $G = (V, E)$:
$$\nabla_{x_j} \log p_t(x_j | x_{pa(j)})$$
where $pa(j)$ denotes the parents of node $j$.

• Training: The model is trained using Denoising Score Matching on these conditional scores. The loss function minimizes the error between the neural network's predicted score and the true conditional score of the noisy data.
• Sampling (Algorithm 1): To generate samples, the model traverses the DAG in topological order. For each node $j$, it evolves a reverse-time SDE conditioned only on the already-generated parent values $x_{pa(j)}$. This ensures the generated samples respect the causal factorization $p(x) = \prod p(x_j | x_{pa(j)})$.

B. Interventional Sampling (Do-Intervention)

To sample from an interventional distribution $p(x | do(X_S = x^*_S))$, the authors modify the sampling process (Algorithm 2):

1. Freeze Intervened Nodes: Nodes in the intervention set $S$ are fixed at values $x^*_S$.
2. Propagate Effects: The reverse diffusion process evolves only for non-intervened nodes, using the frozen values of their parents (which may include intervened nodes) as inputs.
3. Result: This effectively "cuts" incoming edges to $S$ in the generative process, simulating the truncated factorization required for causal interventions without needing closed-form identification formulas.

C. Edge Inference (CEDMI)

Building on the interventional sampler, the authors develop CEDMI, a resampling-based hypothesis test for directed edges.

• Hypothesis: Test if a candidate edge $(k, j)$ exists in the true DAG.
  • $H_0$: The edge is absent ($k \not\to j$).
  • $H_a$: The edge is present.
• Procedure:
  1. Construct a "null graph" by removing the candidate edge $(k, j)$ from a working DAG.
  2. Train a CEDM on this null graph.
  3. Resampling: For observed data points, intervene on the parents of $j$ (excluding $k$) and use the CEDM to generate synthetic $X_j$ values. This creates a null distribution where $X_j \perp X_k | X_{pa(j)\setminus k}$ by construction.
  4. Test Statistic: Use a Multivariate Conditional Dependence Coefficient (MCODEC), an extension of the Azadkia-Chatterjee coefficient, to measure the dependence between $X_j$ and $X_k$ given the other parents.
  5. P-value: Compare the observed statistic against the distribution of statistics computed from the synthetic null samples.

3. Key Contributions

1. Methodological Innovation (CEDM & CEDMI):

   • Introduced a score-based generative framework that strictly respects DAG factorization, enabling exact recovery of observational distributions and principled interventional sampling.
   • Developed a resampling-based edge test (CEDMI) that requires minimal assumptions (nonlinear, non-additive) and uses a novel multivariate dependence coefficient.

2. Statistical Guarantees (Dimension Reduction):

   • Theoretical Breakthrough: Proved that the convergence rate for distribution recovery is governed by the maximum local dimension ($q = \max_j (d_j + |pa(j)|)$) rather than the ambient dimension ($d$).
   • This implies that for sparse graphs, the statistical complexity scales with the size of local neighborhoods, not the total number of variables, effectively mitigating the curse of dimensionality.
   • Established asymptotic control of Type I error for the edge test under sample splitting.

3. Empirical Validation:

   • Demonstrated superior performance in recovering interventional distributions compared to baselines (standard diffusion, VACA) across synthetic and real-world data.
   • Showed that CEDMI achieves near-nominal Type I error control (calibration) and high statistical power, outperforming existing conditional independence tests (e.g., CCIT, RCIT, CDCIT).

4. Results

A. Distribution Recovery (Simulations)

• Setup: Tested on Chain, Hub, Random, and Sachs signaling network structures with varying sample sizes ($N=500$ to $5000$).
• Metric: Squared Maximum Mean Discrepancy (MMD) between generated samples and true interventional distributions.
• Finding: CEDM consistently achieved the lowest MMD. The performance gap was most pronounced in sparse structures (Chain, Hub) and complex biological networks (Sachs), confirming the theoretical advantage of local dimension scaling. Standard diffusion models failed significantly in "Hub" and "Sachs" scenarios where conditioning on the wrong variables led to distributional shifts.

B. Edge Inference (Simulations)

• Calibration: CEDMI maintained rejection rates close to the nominal 0.05 level under null hypotheses, whereas competitors like RCoT and RCIT were often anti-conservative (high false positives).
• Power: CEDMI demonstrated superior power in detecting true edges, especially at moderate sample sizes and weaker signal strengths, compared to NNLSCIT, CDCIT, and CCIT.

C. Real-World Application: Flow Cytometry

• Dataset: Sachs et al. (2005) protein signaling network (11 proteins, 1755 cells).
• Task: Resolve disputed edges between a literature-based "Consensus" network and a data-driven Bayesian network.
• Findings:
  • CEDMI supported the edge Erk $\to$ Akt, which is present in the data-driven graph but absent in the consensus. This aligns with recent biological literature suggesting ERK-AKT cross-talk.
  • CEDMI rejected edges present in the consensus but absent in the data-driven graph (e.g., specific links to PKC), suggesting the consensus network may contain spurious edges for this specific dataset.
  • Results were stable across subsample sizes and consistent with permutation tests where applicable.

5. Significance and Impact

• Bridging Generative AI and Causal Inference: The paper provides a rigorous framework for using diffusion models not just for generation, but for causal reasoning, specifically handling interventions ($do$-calculus) which standard diffusion models cannot do.
• Scalability: By proving that convergence depends on local parent-child dimensions, the method offers a scalable solution for high-dimensional causal discovery, making it applicable to fields like genomics and proteomics where $d$ is large but causal graphs are sparse.
• Robust Inference: The proposed edge test (CEDMI) offers a flexible, assumption-light alternative to traditional causal discovery methods, capable of handling complex nonlinearities without relying on Gaussian assumptions.
• Practical Utility: The application to flow cytometry data demonstrates the method's ability to generate testable biological hypotheses and resolve conflicts between established domain knowledge and data-driven findings.

In summary, this work establishes that encoding causal structure into diffusion models transforms them from flexible density estimators into powerful tools for causal simulation and inference, offering both theoretical guarantees on dimensionality reduction and practical advantages in complex real-world datasets.