Causality in AI and Counterfactual Reasoning

Every time I talk about causal inference in genomics, people ask, ‘But how?’ How do we move from observing correlations in massive genomic datasets to uncovering true cause-and-effect relationships? How can we reliably predict what would happen if we edited a gene, silenced a pathway, or introduced a new compound into a cellular environment?

These questions strike at the heart of modern biology and require moving beyond surface-level observations to deeper causal understanding. In this blog, we will focus on the mathematical foundations of causality and counterfactual reasoning, exploring tools like Pearl’s causal calculus and structural equation modeling, along with practical coding examples. In a future blog, we can look at their implications in genomics.

I have provided comprehensive working code at the end of the article.

Pearl’s Causal Calculus and AI Pipelines

Judea Pearl’s causal calculus provides a formal framework for reasoning about causality. At its core is the do-calculus, a mathematical tool for analyzing interventions. It allows us to manipulate joint probability distributions under causal assumptions.

The key operator is do(⋅), representing external intervention. For example, P(Y ∣ do(X = x)) computes the probability of Y when X is set to x independently of its natural causes. This distinction is crucial for separating correlation from causation. The computation of P(Y ∣ do(X = x)) relies on structural causal models (SCMs), which are formalized as systems of equations:

Here, Pa(Xi) are the parents of Xi in the causal graph, Ui represents unobserved noise variables, and fi​ encodes deterministic or stochastic causal relationships. Directed acyclic graphs (DAGs) visualize these relationships, with nodes representing variables and edges signifying causal dependencies.

A classic example involves computing the causal effect of a treatment X on an outcome Y, given a confounder Z. Using the backdoor adjustment formula:

This equation integrates over Z, ensuring the effect of X on Y is isolated from spurious correlations mediated by Z.

In AI pipelines, embedding do-calculus requires careful integration. Loss functions may incorporate causal constraints derived from SCMs, ensuring that learned representations respect the causal structure. This integration can be extended to reinforcement learning, where agents optimize policies that consider causal effects of actions on long-term outcomes.

Structural Equation Modeling and Counterfactual Prediction

Structural equation modeling (SEM) is a cornerstone of counterfactual reasoning. Counterfactuals are hypothetical scenarios: “What would Y have been if X were different?” Mathematically, counterfactual reasoning involves modifying the SCM to reflect the intervention and solving the resulting equations.

Consider the observed data X=x and Y=y. To compute the counterfactual Yx′, the steps are:

1) Abduction: Infer the values of noise variables U consistent with the observation:

2) Action: Modify the SCM to reflect the intervention do(X=x′). This replaces the structural equation for X with X=x′.

3) Prediction: Solve the modified SCM to compute Yx′:

For example, suppose X represents treatment, Y is the outcome, and U represents unobserved confounders. If the original equations are X=U_1, Y = 2X + U_2​, and we observe X=1, Y=2, we infer U_1=1, U_2=0. Under the counterfactual X=2, the predicted outcome is Y = 2(2) + 0 = 4.

Counterfactual reasoning extends naturally to probabilistic settings. The counterfactual distribution P(Yx′ ∣ X=x, Y=y) quantifies the likelihood of alternate outcomes under hypothetical interventions. Computing this involves integrating over noise variables U consistent with the observed data.

Applications in Fairness, Explainability, and Dynamic Systems

Causal inference underpins critical applications in AI. Counterfactual reasoning provides a principled approach to ensuring fairness, enhancing explainability, and optimizing decisions in dynamic systems.

Fairness in Decision-Making. In fairness, counterfactuals test whether decisions depend on protected attributes. For instance, consider a hiring algorithm trained on X (qualifications), A (age), and Y (hiring decision). To ensure fairness with respect to A, we evaluate:

If P(Y ∣ do(A=a′), X=x) ≠ P(Y ∣ do(A=a), X=x), the decision is causally influenced by age, necessitating mitigation strategies.

Explainability. Counterfactuals also enhance explainability by offering actionable insights. For example, a counterfactual explanation for a rejected loan might state: “Had your income been $5,000 higher, your loan would have been approved.” This is computed by simulating the outcome under a modified income variable:

Dynamic Decision-Making in Autonomous Systems. Autonomous systems operate in dynamic environments where counterfactual reasoning predicts outcomes of hypothetical actions. For example, self-driving cars can simulate alternate trajectories to assess safety. This requires combining causal inference with reinforcement learning, where causal models constrain policy optimization.

Challenges and Future Directions

Causal inference in AI faces several theoretical and computational challenges. First, discovering causal structures from data remains an open problem. While algorithms like PC and FCI provide tools for learning DAGs, they are computationally expensive for high-dimensional data and sensitive to unmeasured confounders. Second, non-identifiability is a fundamental issue. Causal queries often rely on unverifiable assumptions, such as the absence of hidden confounders. Third, integrating causality into deep learning requires hybrid frameworks. Neural networks excel at learning complex correlations but lack explicit causal reasoning. Embedding causal constraints into neural architectures or training on counterfactual samples generated from SCMs offers promising directions.

One exciting frontier is counterfactual generative adversarial networks (CGANs). These models generate counterfactual data samples by simulating interventions within an SCM. Mathematically, the generator G learns a mapping:

where z is a noise vector, and x′ represents the counterfactual intervention. Training such models requires adversarial losses that enforce consistency with observed causal relationships.

Causality also has profound implications for scientific discovery. By formalizing causal hypotheses and testing them against data, researchers can disentangle complex relationships in fields ranging from genomics to economics. In AI, this ability to reason about interventions and counterfactuals will be critical for building systems that do more than predict, they will explain, adapt, and optimize in fundamentally human ways.

Python Libraries for Causal Inference

There are several Python libraries and frameworks that provide tools for causal inference and counterfactual reasoning, including implementations of Pearl’s causal calculus and structural equation modeling. Below are some of the most popular libraries, along with examples to get you started.

1. Causal Inference

A Python library designed for basic causal inference tasks, particularly for treatment effect estimation.

Example: Backdoor Adjustment

from causalinference import CausalModel

# Simulated data
import numpy as np
np.random.seed(42)
N = 1000
X = np.random.binomial(1, 0.5, N)  # Treatment
Z = np.random.normal(0, 1, N)      # Confounder
Y = 2*X + Z + np.random.normal(0, 1, N)  # Outcome

# Create a causal model
data = np.column_stack((Y, X, Z))
cm = CausalModel(Y=data[:, 0], D=data[:, 1], X=data[:, 2:])

# Run causal inference
cm.est_via_ols()
print(cm.summary_stats)
print(cm.estimates)

Output:

Summary Statistics:

We have two groups:

• Control group (X=0, no treatment): Average outcome (Y) is 0.09.
• Treated group (X=1, treatment applied): Average outcome (Y) is 2.13.

The raw difference between these groups is about 2.04. At first glance, this suggests the treatment might increase Y.

Treatment Effects:

After controlling for the confounder (Z), the treatment effect (how much Y changes because of X) is about 2.08. This number, called the Average Treatment Effect (ATE), tells us the effect of the treatment across everyone in the data.

Acronyms: Average Treatment Effect (ATE), Average Treatment effect for Controls (ATC), and Average Treatment effect for the Treated (ATT).

The effect is statistically significant (p-value = 0.000), meaning we can confidently say the treatment has a real impact on the outcome.

2. DoWhy

A comprehensive framework for causal inference that integrates Pearl’s do-calculus and supports structural causal modeling and counterfactual analysis.

Example: Structural Causal Model and Counterfactual Reasoning

import dowhy
from dowhy import CausalModel
import pandas as pd
import numpy as np

# Simulate data
np.random.seed(42)
data = pd.DataFrame({
    'X': np.random.binomial(1, 0.5, 1000),
    'Z': np.random.normal(0, 1, 1000),
})
data['Y'] = 2 * data['X'] + data['Z'] + np.random.normal(0, 1, 1000)  # Calculate 'Y'

# Define the causal model
model = CausalModel(
    data=data,
    treatment='X',
    outcome='Y',
    common_causes=['Z']
)

# View causal graph
model.view_model()

# Estimate treatment effect
identified_estimand = model.identify_effect()
estimate = model.estimate_effect(identified_estimand, method_name="backdoor.linear_regression")
print("Treatment Effect Estimate:")
print(estimate)

# Estimate effect with intervention
estimate_with_intervention = model.estimate_effect(
    identified_estimand,
    method_name="backdoor.linear_regression",
    control_value=0,
    treatment_value=1
)
print("Effect Estimate with Intervention:")
print(estimate_with_intervention)

Output:

3. EconML

A Microsoft library focused on combining machine learning with causal inference techniques.

Example: Treatment Effect Estimation with Machine Learning

from econml.dml import LinearDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LogisticRegression
import numpy as np
import matplotlib.pyplot as plt

# Simulated data
np.random.seed(42)
N = 1000
X = np.random.normal(0, 1, (N, 2))  # Features
T = np.random.binomial(1, 0.5, N)   # Treatment (discrete: 0 or 1)
Y = 5 * T + X[:, 0] + np.random.normal(0, 1, N)  # Outcome

# Define a causal model
model = LinearDML(
    model_y=RandomForestRegressor(),  # Outcome model
    model_t=RandomForestClassifier(),  # Treatment model for discrete treatment
    discrete_treatment=True  # Specify that treatment is discrete
)

# Fit the model
model.fit(Y, T, X=X)

# Estimate treatment effect
treatment_effect = model.effect(X)
print("Estimated treatment effect:", treatment_effect[:5])

plt.hist(treatment_effect, bins=30, edgecolor='k')
plt.xlabel("Treatment Effect")
plt.ylabel("Frequency")
plt.title("Distribution of Estimated Treatment Effects")
plt.show()

ATE = treatment_effect.mean()
print("Average Treatment Effect (ATE):", ATE)

Output:

4. Pyro (Probabilistic Programming for SCMs)

A probabilistic programming framework for Bayesian causal modeling and counterfactual analysis.

Example: Structural Causal Model in Pyro

import pyro
import pyro.distributions as dist
from pyro.infer import MCMC, NUTS
import torch

pyro.clear_param_store()

class BernoulliMixtureNormal(dist.TorchDistribution):
    arg_constraints = {}
    support = dist.constraints.real
    has_rsample = False

    def __init__(self, p, mean0, mean1, scale):
        super().__init__()
        self.p = p
        self.mean0 = mean0
        self.mean1 = mean1
        self.scale = scale
        self.normal0 = dist.Normal(self.mean0, self.scale)
        self.normal1 = dist.Normal(self.mean1, self.scale)

    def log_prob(self, value):
        lp0 = self.normal0.log_prob(value)
        lp1 = self.normal1.log_prob(value)
        max_lp = torch.maximum(lp0, lp1)
        return max_lp + torch.log(self.p * torch.exp(lp1 - max_lp) + (1 - self.p) * torch.exp(lp0 - max_lp))

    def sample(self, sample_shape=torch.Size()):
        mask = torch.bernoulli(self.p.expand(sample_shape))
        return mask * self.normal1.sample(sample_shape) + (1 - mask) * self.normal0.sample(sample_shape)

def scm(intervention=None):
    Z = pyro.sample("Z", dist.Normal(0, 1))
    Z = torch.clamp(Z, -3, 3)
    if intervention and "X" in intervention:
        X = intervention["X"]
        raw_Y = 3 * X + Z
        scaled_Y = torch.tanh(raw_Y)
        Y = pyro.sample("Y", dist.Normal(scaled_Y, 0.2))
    else:
        p = torch.sigmoid(Z)
        mean_0 = torch.tanh(Z)
        mean_1 = torch.tanh(Z + 3.0)
        Y = pyro.sample("Y", BernoulliMixtureNormal(p, mean_0, mean_1, 0.2))
    return Y

def summarize_samples(samples):
    mean = samples.mean().item()
    ci_low = torch.quantile(samples, 0.025).item()
    ci_high = torch.quantile(samples, 0.975).item()
    return mean, ci_low, ci_high

# Observational inference (no intervention)
nuts_kernel = NUTS(scm)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=300)
mcmc.run()
observational_samples = mcmc.get_samples()["Y"]
obs_mean, obs_ci_low, obs_ci_high = summarize_samples(observational_samples)

# Interventional inference (X=1)
def scm_with_intervention():
    return scm(intervention={"X": torch.tensor(1.0)})

nuts_kernel_intervention = NUTS(scm_with_intervention)
mcmc_intervention = MCMC(nuts_kernel_intervention, num_samples=1000, warmup_steps=300)
mcmc_intervention.run()
interventional_samples = mcmc_intervention.get_samples()["Y"]
int_mean, int_ci_low, int_ci_high = summarize_samples(interventional_samples)

print(f"Observational Y: Mean={obs_mean:.2f}, 95% CI=({obs_ci_low:.2f}, {obs_ci_high:.2f})")
print(f"Interventional Y (X=1): Mean={int_mean:.2f}, 95% CI=({int_ci_low:.2f}, {int_ci_high:.2f})")

# Ensure that the observational and interventional samples are 1D tensors
obs_samples = observational_samples.flatten()
int_samples = interventional_samples.flatten()

# Visualization
plt.figure(figsize=(10, 6))
plt.hist(obs_samples.numpy(), bins=30, alpha=0.5, label="Observational", color="blue")
plt.hist(int_samples.numpy(), bins=30, alpha=0.5, label="Interventional (X=1)", color="green")
plt.axvline(obs_mean, color="blue", linestyle="--", label=f"Obs Mean: {obs_mean:.2f}")
plt.axvline(int_mean, color="green", linestyle="--", label=f"Int Mean: {int_mean:.2f}")
plt.legend()
plt.title("Observational vs Interventional Y Distributions")
plt.xlabel("Y")
plt.ylabel("Frequency")
plt.show()

Output:

How to Interpret?

These results suggest that when we passively observe the system, the outcome Y hovers around a lower mean with wide uncertainty. In contrast, forcing X to 1 raises Y’s expected value and confines it to a range that does not dip into negative territory.

This implies a causal effect: setting X to 1 appears to shift Y’s distribution toward more positive values relative to what is seen under mere observation. It indicates that the action of making X equal to one has a direct and meaningful impact on Y’s outcome that is not just a byproduct of correlation or confounding.

Comparison Across Libraries: