Why Your AI Model is Blind to Reality: Causal Inference

Topic Overview

What is Causal Inference in ML? It is the mathematical framework that allows machine learning models to distinguish between correlation and causation, enabling them to predict the outcome of interventions rather than just observations.

Why does it exist? Standard deep learning is fundamentally an advanced form of curve-fitting. A neural network minimizes a loss function by finding any statistical shortcut in the data. If your training data for cows contains mostly green grass, the network will learn that green pixels cause the classification “cow.” This is a spurious correlation. Causal inference exists to explicitly forbid the model from using these shortcuts by mathematically modeling the true data-generating process.

What problem does it solve? It solves the Out-of-Distribution (OOD) generalization problem. When a model relies on spurious correlations, it fails catastrophically when deployed in a new environment (e.g., a cow standing on a beach). Causal ML identifies the invariant mechanisms—the features that remain true regardless of the environment.

Why should engineers care? If you deploy a standard ML model for medical treatment, loan approval, or ad bidding, it will exploit biases in the historical data. If you change the policy (e.g., start approving loans for a demographic that was historically denied), the model’s predictions will be completely wrong because it never learned the causal effect of the input features on the outcome. Causal ML is the prerequisite for building Fair AI, Robust AI, and Offline Reinforcement Learning systems.

Prerequisites

1. Conditional Probability * Why it is needed: Causal inference is built entirely on manipulating conditional probability distributions. You must be completely fluent in reading and interpreting \(P(Y|X)\). * Where it appears: In the distinction between \(P(Y|X)\) (seeing) and \(P(Y|do(X))\) (doing), which is the core mathematical split in this lesson.

2. Bayes’ Theorem * Why it is needed: Bayes’ theorem allows us to invert probabilities (\(P(A|B)\) to \(P(B|A)\)). In causal inference, we constantly use probabilistic inversion to isolate the effect of a specific variable from the noise of its confounders. * Where it appears: In deriving adjustment formulas (like the Backdoor Criterion) where we must condition on specific upstream variables to block spurious paths.

Historical Motivation

In the 1980s, AI researcher Judea Pearl hit a wall. Bayesian networks were excellent at calculating probabilities given observed evidence (diagnosis), but they completely failed at prescriptive tasks (treatment). If a Bayesian network observed that a barometer dropped and a storm came, it learned \(P(Storm|Barometer)\). But it could not answer: “If I physically break the barometer, will it cause a storm?”

Probability theory had no mathematical operator to represent “physically manipulating” a system. Existing statistical methods (like randomized controlled trials) worked, but they were empirical, not algorithmic.

Pearl’s engineering necessity led him to invent the Structural Causal Model (SCM) and the \(do\)-operator. He combined probability theory with directed graph theory (DAGs). By drawing arrows to represent physical mechanisms, he created a mathematical way to “cut” arrows in a graph, simulating a physical intervention. Decades later, as deep learning hit the OOD generalization wall in the 2010s, the ML community realized Pearl’s causal graphs were the exact missing ingredient to make neural networks robust.

Intuitive Explanation

Imagine you are a statistician observing a room. You notice that whenever the Ice Cream Sales (\(X\)) go up, the number of Drownings (\(Y\)) also goes up. You calculate \(P(Y|X)\) and find a strong positive correlation. A standard ML model learns: High Ice Cream Sales \(\rightarrow\) High Drownings.

Now, a causal engineer walks in. The engineer understands mechanisms. They draw a graph: * Temperature (\(Z\)) causes people to buy ice cream (\(Z \rightarrow X\)). * Temperature (\(Z\)) causes people to go swimming (\(Z \rightarrow Y\)).

The engineer asks a different question: What is \(P(Y | do(X))\)? If I physically force ice cream sales to be zero (ban ice cream), how many drownings will occur?

The statistician says: “Drownings will drop to zero! The correlation is broken!” The engineer says: “Drownings will remain exactly the same. Banning ice cream does not change the temperature, and temperature is the true cause.”

Standard ML is the statistician. Causal ML is the engineer. Causal inference is the mathematical process of finding the true causes (Temperature) and ignoring the spurious correlates (Ice Cream).

Visual Understanding

To visualize this, draw a Directed Acyclic Graph (DAG) on a piece of paper.

Nodes: Draw three circles. Label them \(Z\) (Confounding variable), \(X\) (Treatment/Input), and \(Y\) (Outcome).
Arrows (Mechanisms): Draw an arrow from \(Z\) to \(X\). Draw an arrow from \(Z\) to \(Y\). Draw an arrow from \(X\) to \(Y\).
- Rule: An arrow means “directly causes.” The arrow from \(Z\) to \(X\) means \(Z\) is a parent of \(X\).
The Spurious Path: Trace a path from \(X\) to \(Z\) to \(Y\). This is a backdoor path. It is a non-causal pipe leaking statistical correlation from \(X\) to \(Y\) via the common cause \(Z\).
Graph Surgery (\(do\)-operator): To calculate the true causal effect of \(X\) on \(Y\), take a red pen and draw a heavy slash through the arrow going into \(X\) (the \(Z \rightarrow X\) arrow).
- Visual meaning: By forcing \(X\) to be a specific value (an intervention), you sever its dependence on its natural causes. \(X\) is now an orphan node in the graph. The backdoor path is physically cut.

Mathematical Foundations

Let us formalize the graph surgery.

1. Observational vs. Interventional * Observational: \(P(Y | X=x)\). This is what we calculate using Bayes’ rule. It includes all spurious paths. * Interventional: \(P(Y | do(X=x))\). This is Pearl’s invention. It represents the distribution of \(Y\) in a manipulated world where all arrows pointing into \(X\) are deleted.

2. The Adjustment Formula (Backdoor Criterion) How do we calculate the \(do\)-operator using standard observational data without actually running an experiment? We must block the spurious backdoor paths.

If a set of variables \(\mathbf{Z}\) satisfies the “Backdoor Criterion” (it blocks all paths between \(X\) and \(Y\) that have an arrow into \(X\)), we can calculate the causal effect using this formula:

\[P(Y | do(X=x)) = \sum_{\mathbf{z}} P(Y | X=x, \mathbf{Z}=\mathbf{z}) P(\mathbf{Z}=\mathbf{z})\]

\(\sum_{\mathbf{z}}\): We iterate over every possible state of the confounders.
\(P(\mathbf{Z}=\mathbf{z})\): The marginal probability of the confounder occurring naturally in the population.
\(P(Y | X=x, \mathbf{Z}=\mathbf{z})\): The conditional probability of the outcome, given both the treatment and the specific confounder state.

Why this works mathematically: By conditioning on \(\mathbf{Z}\), we mathematically “choke” the backdoor path \(X \leftarrow Z \rightarrow Y\). We are left only with the causal path \(X \rightarrow Y\). By multiplying by \(P(\mathbf{Z})\) and summing, we re-weight the data to reflect the natural distribution of the confounders, exactly simulating the graph surgery of the \(do\)-operator.

Worked Numerical Example

Let’s compute the causal effect of a drug (\(X \in \{0, 1\}\)) on Recovery (\(Y \in \{0, 1\}\)), where Age (\(Z \in \{0 \text{ (Young)}, 1 \text{ (Old)}\}\)) is a confounder. Old people are more likely to get the drug, and old people recover less often naturally.

Step 1: Gather Observational Data We observe the following probabilities in our dataset: * \(P(Z=0) = 0.5\) (Half the population is young) * \(P(Z=1) = 0.5\) (Half is old) * \(P(Y=1 | X=1, Z=0) = 0.8\) (Young, took drug, 80% recover) * \(P(Y=1 | X=0, Z=0) = 0.4\) (Young, no drug, 40% recover) * \(P(Y=1 | X=1, Z=1) = 0.3\) (Old, took drug, 30% recover) * \(P(Y=1 | X=0, Z=1) = 0.2\) (Old, no drug, 20% recover)

A naive ML model might see that overall, \(P(Y=1|X=1) = 0.55\) and \(P(Y=1|X=0) = 0.30\). It thinks the drug is great. But this is corrupted by the fact that mostly old people took it.

Step 2: Apply the Backdoor Adjustment Formula We want to find the true causal effect: \(P(Y=1 | do(X=1))\).

\[P(Y=1 | do(X=1)) = \sum_{z \in \{0,1\}} P(Y=1 | X=1, Z=z) P(Z=z)\]

Step 3: Calculate the baseline \[P(Y=1 | do(X=0)) = (0.4 \times 0.5) + (0.2 \times 0.5) = 0.20 + 0.10 = 0.30\]

Conclusion: The true Average Treatment Effect (ATE) is \(0.55 - 0.30 = 0.25\). By explicitly summing over the confounder \(Z\), we stripped out the bias of who chose to take the drug, revealing the pure mechanical efficacy of the chemical itself.

Computational Interpretation

What happens when the computer executes a causal inference algorithm, like Inverse Probability Weighting (IPW), which is the computational equivalent of the adjustment formula?

Data Structures: We have standard tensors: \(\mathbf{X} \in \mathbb{R}^N\), \(\mathbf{Y} \in \mathbb{R}^N\), and a vector of confounders \(\mathbf{Z} \in \mathbb{R}^{N \times d_z}\).

The Compute Steps: 1. Propensity Model: We train a secondary neural network (or logistic regression) \(g(\mathbf{Z})\) to predict \(P(X=1 | \mathbf{Z})\). This outputs a propensity score \(e_i\) for every sample \(i\). 2. Weight Calculation: We compute the inverse probability weight for every sample: \(w_i = \frac{X_i}{e_i} + \frac{1 - X_i}{1 - e_i}\). * Why? If an old person (\(Z=1\)) took the drug (\(X=1\)), but the propensity model says old people rarely take the drug (\(e_i = 0.1\)), this sample is highly valuable and gets a massive weight (\(w_i = 1 / 0.1 = 10\)). It creates a “synthetic” population where the drug was assigned randomly. 3. Weighted Loss: We train the main prediction model \(f(\mathbf{X})\) using a weighted loss function: \(\mathcal{L} = \sum_{i} w_i \cdot \text{Loss}(f(\mathbf{X}_i), \mathbf{Y}_i)\).

Complexity & Scaling: * The computational complexity is \(\mathcal{O}(N \cdot d_z)\) for the propensity network, plus \(\mathcal{O}(N)\) for the weighting. * The Bottleneck: Numerical instability. If \(e_i\) is very close to 0 or 1, the weight \(w_i\) explodes to infinity. In production code, weights are strictly clamped (e.g., \(e_i \in [0.01, 0.99]\)) to prevent NaN gradients from destroying the main network.

Implementation Perspective

How do we inject causality directly into a Deep Learning architecture without manually specifying propensity scores? We use Invariant Risk Minimization (IRM).

IRM forces the neural network weights \(\mathbf{w}\) to learn features that have the exact same causal predictive power across different environments \(e\) (e.g., different hospitals, different lighting conditions).

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalIRMClassifier(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        # A feature extractor (does not need to be causal)
        self.phi = nn.Linear(in_features, 64)
        # A linear classifier head (WE FORCE THIS TO BE CAUSAL)
        self.w = nn.Linear(64, 1, bias=False) 

    def forward(self, x):
        # Extract features
        features = F.relu(self.phi(x))
        # Return logits and the features (needed for penalty)
        return self.w(features), features

def irm_loss(model, x_env, y_env, penalty_weight=1.0):
    """
    x_env: [Batch, Features] from a specific environment
    y_env: [Batch, 1] binary labels
    """
    logits, features = model(x_env)
    
    # 1. Standard Empirical Risk (Binary Cross Entropy)
    error = F.binary_cross_entropy_with_logits(logits, y_env)
    
    # 2. The IRM Penalty
    # We want the classifier w to perform optimally on this environment.
    # If w is relying on spurious features, it CANNOT be optimal for all environments.
    # We create a dummy classifier w_hat scaled by 1e-1
    dummy_w = torch.linalg.norm(model.w.weight) * 1e-1
    dummy_logits = (features / dummy_w) @ model.w.weight.T
    
    # Penalty: How much does the error change if we scale w?
    # If w is causal, scaling it doesn't change the relative ranking (penalty -> 0)
    penalty = error - F.binary_cross_entropy_with_logits(dummy_logits, y_env)
    
    # 3. Total Loss
    return error + penalty_weight * (penalty ** 2)

Explanation: Standard ERM minimizes error. IRM minimizes error + penalty**2. The penalty forces the gradient of the error with respect to the classifier weights to be zero. If a feature (like green pixels) works in Environment A but fails in Environment B, the penalty will be massive. The network is mathematically forced to drop the green pixel feature and find the invariant causal feature (the shape of the cow).

Where It Appears In AI

AI Component	Exact Usage
Recommender Systems (Transformers)	Removing “Popularity Bias.” A user’s historical clicks (\(Z\)) cause both the exposure to an item (\(X\)) and the click on the item (\(Y\)). Causal graphs deconfound the recommendation to find items the user actually likes, not just items they were bombarded with.
NLP / LLMs	Deconfounded Prompting. When extracting knowledge, the prompt format (\(Z\)) causes both the model’s generated reasoning (\(X\)) and the final answer (\(Y\)). Causal tracing isolates the actual knowledge in the MLP layers from the formatting artifacts in the attention layers.
Computer Vision (CNNs/ViTs)	Domain Generalization. Using causal DAGs to separate “style” (spurious, environment-dependent) from “content” (causal, environment-invariant) in image generation and classification.
Offline RL	Policy Evaluation. The history of states (\(Z\)) confounds the action taken (\(X\)) and the reward received (\(Y\)). Causal inference is required to answer “What will be the reward if I deploy a new policy?”
Fairness (Tabular MLPs)	Path-Specific Effects. Using causal graphs to separate the “direct effect” of a protected attribute (e.g., gender) on an outcome from the “indirect effect” (e.g., gender \(\rightarrow\) job experience \(\rightarrow\) hiring).

Deep Dive Into Real Models

Causal Representation Learning in VAEs Standard VAEs disentangle latent variables \(\mathbf{z}\) to maximize the ELBO, but they have no guarantee of finding causal variables. * The Causal VAE modification: The decoder is structured as a Structural Causal Model (SCM). If the true generative process is \(Z_1 \rightarrow Z_2 \rightarrow X\), the decoder enforces this directional graph. * Interventional Data Augmentation: During training, the model applies the \(do\)-operator inside the latent space. It samples \(Z_1\), intervenes to set \(Z_2\) to a random value (cutting the \(Z_1 \rightarrow Z_2\) arrow), and passes it to the decoder. If the decoder fails, the gradients push the encoder to learn true causal factors.

Why Standard LLMs are not Causal A Transformer calculates \(P(Y|X)\) where \(X\) is the prompt. If you prompt: “My friend took a bullet to the chest. What is their heart rate?”, the LLM outputs “high.” But if you ask the causal question \(P(Y|do(X))\): “If I surgically intervene to stop a bullet in a phantom model, what happens to the heart rate?”, the LLM hallucinates because its attention mechanism is purely correlational. It lacks an internal SCM of human anatomy to perform graph surgery.

Failure Modes

1. Unmeasured Confounding (The Fatal Flaw) * Cause: The Backdoor Adjustment formula requires you to measure and include all confounders \(\mathbf{Z}\) in your dataset. If there is a hidden variable causing both \(X\) and \(Y\) that you didn’t log, the math is fundamentally broken. * Symptoms: The model’s causal effect estimates vary wildly depending on which variables you include in \(\mathbf{Z}\) (sign flip fallacy). * Solution: You cannot solve this with code. You must go back to domain experts and gather more features, or use instrumental variables (a variable that affects \(Y\) only through \(X\)).

2. Collider Bias (Selection Bias) * Cause: Conditioning on a variable that is caused by both \(X\) and \(Y\). For example, \(X\) = talent, \(Y\) = effort, Collider = “Hired at Google”. If you train a model only on Google employees, you mathematically induce a negative correlation between talent and effort (because you restricted the dataset to people who must have at least one). * Symptoms: A model that works on the general population fails completely on a specific subpopulation (e.g., hospitalized patients). * Solution: Never condition on descendants of the treatment and outcome in your DAG without careful mathematical correction (front-door criterion).

3. Violation of Positivity * Cause: In propensity score weighting, if for some value of \(\mathbf{Z}\), the probability of treatment is exactly 0 or 1 (e.g., no old people were given the drug in the data), the weight becomes undefined (division by zero). * Symptoms: NaN losses, exploding variances in the causal effect estimate. * Solution: Trim the data. Drop subpopulations where \(0.01 < e(\mathbf{z}) < 0.99\) is violated. Accept that you cannot make causal claims for those subgroups.

Engineering Insights

Tradeoffs: Accuracy vs. Robustness Causal models (like IRM) almost always have lower in-distribution accuracy than standard ERM models. Why? Because standard ERM cheats. It uses every available spurious correlation to get 99% accuracy on the validation set. A causal model restricts its hypothesis space, maybe achieving only 95% on the validation set, but it will maintain that 95% when deployed to a new country. The engineering tradeoff is strictly sacrificing local performance for guaranteed OOD robustness.

The Computational Cost of DAG Discovery As an engineer, you will be tempted to use algorithms (like PC or GES) to learn the causal graph from data automatically. Do not do this for high-dimensional deep learning. Learning a DAG is an NP-hard problem. The number of possible graphs scales super-exponentially with the number of nodes. In production ML, causal graphs must be hardcoded as priors by domain experts. The neural network optimizes within the constraints of that human-specified graph.

Interview Questions

Beginner: Question: What is the mathematical difference between \(P(Y|X)\) and \(P(Y|do(X))\)? Answer: \(P(Y|X)\) is observational; it measures the probability of \(Y\) given that we see \(X\). It includes the effect of any confounders that cause both \(X\) and \(Y\). \(P(Y|do(X))\) is interventional; it represents the probability of \(Y\) if we physically force \(X\) to a value, severing all incoming causal arrows to \(X\). It isolates the direct causal effect.

Intermediate: Question: How does the Backdoor Adjustment formula fix the problem of confounding? Answer: It blocks the non-causal “backdoor” path between \(X\) and \(Y\) by conditioning on a set of confounders \(\mathbf{Z}\). By calculating \(P(Y|X, \mathbf{z})\) for every state of \(\mathbf{z}\), and then re-weighting by the natural distribution \(P(\mathbf{z})\), we simulate a randomized controlled trial where \(X\) and \(\mathbf{Z}\) are independent.

Advanced: Question: In Invariant Risk Minimization (IRM), why do we penalize the gradient of the loss with respect to the classifier weights? Answer: If a feature is spurious (e.g., watermarks in images), its predictive power will change across environments. Therefore, the optimal classifier weight for that feature will be non-zero in Environment A, but zero in Environment B. The gradient of the loss with respect to that weight will be non-zero. By forcing this gradient to be zero across all environments, IRM mathematically guarantees that the feature being used has an invariant relationship with the target, which is the definition of a causal feature.

Research-Level: Question: Standard ML relies on i.i.d. (independent and identically distributed) data. Why is causal inference fundamentally at odds with the i.i.d. assumption? Answer: Causal inference is entirely about interventions, which by definition violate the i.i.d. assumption. An intervention \(do(X)\) changes the data-generating process, meaning the distribution of the interventional data \(P(Y|do(X))\) is fundamentally different from the observational distribution \(P(Y|X)\). Causal ML requires models that can generalize across different data distributions (environments), whereas standard i.i.d. ML assumes the test distribution perfectly matches the train distribution.

Research Perspective

Current Limitations: The “Causal Hierarchy” (Pearl’s Ladder) shows we are mostly stuck on Level 1 (Association/Standard ML) and just reaching Level 2 (Intervention). Level 3 (Counterfactuals: “What would have happened if I had acted differently?”) requires individual-level causal models, which is severely underdeveloped in deep learning.

Modern Improvements & Research Directions: 1. Causal Representation Learning: The biggest open problem. How can we take raw pixels/video and output a set of latent variables that map 1-to-1 with true causal mechanisms, without human supervision? (e.g., decoupling object identity from lighting). 2. Neuro-Symbolic Causal AI: Combining the pattern recognition of neural networks with the logical, graph-surgery capabilities of symbolic AI (like Pearl’s do-calculus) in a single differentiable architecture. 3. LLMs as Causal Reasoners: Research into whether scaling LLMs implicitly learns world models (causal graphs), or if they must be explicitly prompted with DAG structures to perform accurate counterfactual reasoning.

Connections

Prerequisites: * Conditional Probability (The base math) * Bayesian Networks (The graphical foundation) * ELBO / VAEs (Where causal constraints are injected)

Enables: * Out-of-Distribution (OOD) Generalization * Algorithmic Fairness (Counterfactual fairness) * Offline Reinforcement Learning

Used by: * CausalGANs (Generating data that respects causal constraints) * Invariant Risk Minimization (IRM) * Targeted Maximum Likelihood Estimation (TMLE) in biostats

Extended by: * Do-Calculus (The complete ruleset for graph surgery) * Structural Causal Models (SCMs) (The equations governing the nodes) * Counterfactual Reasoning (Level 3 of the Causal Ladder)

Related Concepts: * Spurious Correlations * Domain Adaptation / Transfer Learning * Instrumental Variables (A tool for when unmeasured confounding exists)

Final Mental Model

Standard Machine Learning is a Passive Observer sitting in a room, drawing curves through data points. It can tell you what usually happens together.

Causal Inference in ML hands the Passive Observer a Scalpel and a Soldering Iron. It allows the engineer to look at the web of connections, cut the wires that represent spurious correlations (the Backdoor Criterion), and force a variable to be a specific value (the \(do\)-operator) to see what truly breaks. It transforms AI from a system that asks “What do I see?” into a system that asks “What makes this happen?”

Neural Nexus

Search This Blog