[Paper Review] Modeling Causal Mechanisms with Diffusion Models for Interventional and Counterfactual Queries

Patrick Chao, Patrick Blöbaum+2 more|arXiv (Cornell University)|Feb 2, 2023

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper introduces Diffusion-Based Causal Models (DCM), a novel framework that leverages denoising diffusion models to learn structural causal models from observational data and a known causal graph. By modeling each node's structural equation via a diffusion process with latent encodings as proxies for exogenous noise, DCM enables accurate sampling under interventions and counterfactual inference, outperforming state-of-the-art methods by up to an order of magnitude on nonadditive noise models.

ABSTRACT

We consider the problem of answering observational, interventional, and counterfactual queries in a causally sufficient setting where only observational data and the causal graph are available. Utilizing the recent developments in diffusion models, we introduce diffusion-based causal models (DCM) to learn causal mechanisms, that generate unique latent encodings. These encodings enable us to directly sample under interventions and perform abduction for counterfactuals. Diffusion models are a natural fit here, since they can encode each node to a latent representation that acts as a proxy for exogenous noise. Our empirical evaluations demonstrate significant improvements over existing state-of-the-art methods for answering causal queries. Furthermore, we provide theoretical results that offer a methodology for analyzing counterfactual estimation in general encoder-decoder models, which could be useful in settings beyond our proposed approach.

Motivation & Objective

• To develop a flexible and scalable framework for answering observational, interventional, and counterfactual queries in structural causal models using only observational data and a causal graph.
• To leverage the expressive power of diffusion models to learn causal mechanisms by treating latent codes as proxies for unobserved exogenous noise.
• To provide a unified method for both intervention and counterfactual inference through a deterministic, autoencoder-style diffusion process (DDIM).
• To theoretically analyze the accuracy of counterfactual estimates in encoder-decoder models, with bounds applicable beyond diffusion models.
• To empirically validate DCM’s superiority over existing state-of-the-art methods across diverse structural equation types and real-world data.

Proposed method

• Model each node in the causal graph as a conditional diffusion model that takes parent values as input and learns to encode observations into a latent space and decode back to the variable space.
• Use Denoising Diffusion Implicit Models (DDIM) as the underlying diffusion architecture, enabling deterministic sampling and efficient inference via a learned reverse process.
• Train the DCM in topological order, ensuring that each node’s diffusion model is conditioned on its parents’ values, mimicking the structural equation model (SCM) generation process.
• For interventional queries, apply do-interventions by fixing the intervened variable’s value and sampling from the corresponding diffusion model.
• For counterfactual queries, perform abduction by encoding the observed outcome into the latent space, then apply action (intervention on the latent) and prediction (reconstruction) steps.
• Leverage the latent representation as a proxy for exogenous noise, enabling the model to capture complex, nonparametric structural equations.

Figure 1 : Left: Nonlinear setting (NLIN), Right: Nonadditive setting (NADD). Box plots of observational, interventional, and counterfactual queries of the ladder and random SCMs over $20$ random initializations of the model and training data.

Experimental results

Research questions

• RQ1 Can diffusion models effectively learn causal mechanisms from observational data and a causal graph to support all three types of causal queries?
• RQ2 How does the performance of DCM compare to existing state-of-the-art methods in interventional and counterfactual inference, especially under nonadditive noise?
• RQ3 What theoretical guarantees can be provided for counterfactual estimation in encoder-decoder models like DCM?
• RQ4 Can DCM generalize to real-world data, such as fMRI data, with complex, nonlinear dependencies?
• RQ5 Does the latent space in DCM serve as a valid proxy for exogenous noise, enabling accurate abduction and counterfactual reasoning?

Key findings

• DCM outperforms existing state-of-the-art methods by up to an order of magnitude on interventional and counterfactual queries in nonadditive noise models, with a 10.42 ± 2.77 MSE on the Y graph in the NADD setting.
• On the semi-synthetic Sachs dataset, DCM achieves the lowest counterfactual MSE (0.72 ± 0.19 × 10⁻²) in the NLIN setting, significantly outperforming VACA (17.71 ± 9.23 × 10⁻²) and CAREFL (7.30 ± 0.95 × 10⁻²).
• In the NADD setting on the Sachs dataset, DCM achieves a counterfactual MSE of 1.99 ± 2.49 × 10⁻², outperforming VACA (7.30 ± 11.57 × 10⁻²) and CAREFL (8.30 ± 8.34 × 10⁻²).
• DCM achieves the lowest observational MMD (0.16 ± 0.11) on the Ladder graph in the NLIN setting, indicating strong distributional fidelity.
• Theoretical analysis shows that counterfactual estimation error in DCM can be bounded under reasonable assumptions, providing a general framework applicable to other encoder-decoder models.
• Empirical results on synthetic datasets with various structural equation types (linear, nonlinear, nonadditive) consistently show DCM’s superiority in interventional and counterfactual queries across all graph structures tested.