Causal Inference for Machine Learning

Causal inference goes beyond correlation to understand the effect of interventions. While machine learning excels at prediction, causal inference answers "what if?" questions essential for decision-making.

1. Structural Causal Models (SCMs)

1.1 Definition

A structural causal model is a tuple where:

• : exogenous (unobserved) variables
• : endogenous (observed) variables
• : structural equations
• : distribution over exogenous variables

1.2 Directed Acyclic Graphs (DAGs)

Variables are nodes; causal effects are directed edges. No cycles (acyclic).

d-separation: Nodes and are d-separated by if every path between them is blocked:

• Chain: blocked if
• Fork: blocked if
• Collider: blocked if and no descendant of is in

1.3 Markov Condition

A DAG satisfies the Markov condition if each variable is independent of its non-descendants given its parents:

2. The Do-Calculus

2.1 Observational vs Interventional

• Observational: — correlation, conditioning
• Interventional: — causal effect, surgery

Do-operator removes incoming edges to and sets .

2.2 Three Rules of Do-Calculus

Rule 1 (Insertion/deletion of observations):

Rule 2 (Action/observation exchange):

Rule 3 (Insertion/deletion of actions):

where:

• : graph with edges into removed
• : graph with edges out of removed
• : descendants of not in

2.3 Backdoor Criterion

Definition: A set satisfies the backdoor criterion relative to if:

1. No node in is a descendant of
2. blocks every path between and that contains an arrow into

Backdoor adjustment formula:

2.4 Frontdoor Criterion

Definition: A set satisfies the frontdoor criterion relative to if:

1. All directed paths from to go through
2. No backdoor path from to
3. All backdoor paths from to are blocked by

Frontdoor adjustment formula:

3. Potential Outcomes Framework

3.1 Rubin Causal Model

For each unit , define potential outcomes:

• : outcome if treated ()
• : outcome if untreated ()

Individual treatment effect:

3.2 Fundamental Problem

We observe only one potential outcome per unit:

The other is counterfactual.

3.3 Average Treatment Effect (ATE)

ATE is not directly estimable because we never observe both and for the same unit.

3.4 Identifiability Assumptions

1. SUTVA (Stable Unit Treatment Value Assumption): No interference between units
2. Unconfoundedness:
3. Positivity: for all

Under these assumptions:

4. Propensity Score Methods

4.1 Propensity Score

The propensity score is the probability of treatment given covariates:

Propensity score theorem (Rosenbaum & Rubin, 1983): If unconfoundedness holds, then .

4.2 Inverse Probability Weighting (IPW)

Estimate ATE as:

Variance:

4.3 Doubly Robust Estimation

Combine outcome modeling and propensity weighting:

Doubly robust estimator:

Property: Consistent if either or is correctly specified (not necessarily both).

5. Instrumental Variables

5.1 Setup

When confounders are unobserved, use an instrument that:

1. Relevance: affects :
2. Exclusion: affects only through :
3. Exogeneity: is independent of confounders:

5.2 Wald Estimator

For binary and :

5.3 Two-Stage Least Squares (2SLS)

Stage 1: Regress on :

Stage 2: Regress on :

Then estimates the causal effect.

5.4 Local Average Treatment Effect (LATE)

The IV estimator identifies the LATE: the effect for units whose treatment status is changed by the instrument.

where compliers are those who take treatment iff .

6. Heterogeneous Treatment Effects

6.1 Conditional Average Treatment Effect (CATE)

6.2 Causal Forest

Generalized random forests adapt random forests to estimate heterogeneous treatment effects:

1. Split using treatment effect heterogeneity (not prediction accuracy)
2. Leaf predictions are local ATE estimates
3. Valid confidence intervals via influence functions

6.3 Meta-Learners

S-learner:

where is trained on .

T-learner:

Train separate models for treated and control.

X-learner:

1. Train
2. Compute imputed treatment effects
3. Train a model on the imputed effects

R-learner (Robinson):

7. Code: Causal Inference

import numpy as np
from scipy import stats
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.model_selection import cross_val_predict

class PropensityScoreEstimator:
    """Estimate propensity scores using various methods."""
    
    def __init__(self, method='logistic'):
        self.method = method
        self.model = None
    
    def fit(self, X, T):
        """Fit propensity score model."""
        if self.method == 'logistic':
            self.model = LogisticRegression(max_iter=1000)
        elif self.method == 'random_forest':
            self.model = RandomForestClassifier(n_estimators=100)
        else:
            raise ValueError(f"Unknown method: {self.method}")
        
        self.model.fit(X, T)
        return self
    
    def predict(self, X):
        """Predict propensity scores."""
        return self.model.predict_proba(X)[:, 1]


class IPWEstimator:
    """Inverse Probability Weighting estimator for ATE."""
    
    def __init__(self, propensity_estimator=None):
        self.propensity = propensity_estimator or PropensityScoreEstimator()
    
    def estimate(self, X, T, Y, clipping=0.01):
        """Estimate ATE using IPW."""
        # Fit propensity scores
        self.propensity.fit(X, T)
        e = self.propensity.predict(X)
        
        # Clip propensity scores for stability
        e = np.clip(e, clipping, 1 - clipping)
        
        # IPW estimator
        ate = np.mean(T * Y / e) - np.mean((1 - T) * Y / (1 - e))
        
        # Variance
        var_treated = np.var(T * Y / e)
        var_control = np.var((1 - T) * Y / (1 - e))
        se = np.sqrt(var_treated / np.sum(T) + var_control / np.sum(1 - T))
        
        return ate, se
    
    def confidence_interval(self, X, T, Y, alpha=0.05):
        """Compute confidence interval."""
        ate, se = self.estimate(X, T, Y)
        z = stats.norm.ppf(1 - alpha / 2)
        return ate - z * se, ate + z * se


class DoublyRobustEstimator:
    """Doubly Robust estimator for ATE."""
    
    def __init__(self, outcome_model='rf', propensity_model='logistic'):
        self.outcome_model = outcome_model
        self.propensity_model = propensity_model
    
    def estimate(self, X, T, Y):
        """Estimate ATE using doubly robust estimation."""
        n = len(Y)
        
        # Fit propensity scores
        prop_est = PropensityScoreEstimator(self.propensity_model)
        prop_est.fit(X, T)
        e = prop_est.predict(X)
        e = np.clip(e, 0.01, 0.99)
        
        # Fit outcome models
        if self.outcome_model == 'rf':
            mu1_model = RandomForestRegressor(n_estimators=100)
            mu0_model = RandomForestRegressor(n_estimators=100)
        else:
            mu1_model = LinearRegression()
            mu0_model = LinearRegression()
        
        # Fit on treated and control separately
        mask_treated = T == 1
        mask_control = T == 0
        
        mu1_model.fit(X[mask_treated], Y[mask_treated])
        mu0_model.fit(X[mask_control], Y[mask_control])
        
        # Predict counterfactuals
        mu1 = mu1_model.predict(X)
        mu0 = mu0_model.predict(X)
        
        # Doubly robust estimator
        dr = (mu1 - mu0 + 
              T * (Y - mu1) / e - 
              (1 - T) * (Y - mu0) / (1 - e))
        
        ate = np.mean(dr)
        se = np.std(dr) / np.sqrt(n)
        
        return ate, se


class InstrumentalVariableEstimator:
    """Instrumental Variable estimator (2SLS)."""
    
    def __init__(self):
        self.first_stage = None
        self.second_stage = None
    
    def estimate(self, X, Z, T, Y):
        """
        Estimate causal effect using 2SLS.
        
        Args:
            X: Covariates
            Z: Instrument
            T: Treatment
            Y: Outcome
        """
        # Stage 1: Regress T on Z (and X if needed)
        Z_aug = np.column_stack([Z, X]) if X is not None else Z.reshape(-1, 1)
        self.first_stage = LinearRegression()
        self.first_stage.fit(Z_aug, T)
        T_hat = self.first_stage.predict(Z_aug)
        
        # Stage 2: Regress Y on T_hat (and X if needed)
        T_hat_aug = np.column_stack([T_hat, X]) if X is not None else T_hat.reshape(-1, 1)
        self.second_stage = LinearRegression()
        self.second_stage.fit(T_hat_aug, Y)
        
        # Causal effect is coefficient on T_hat
        effect = self.second_stage.coef_[0]
        
        return effect
    
    def wald_estimator(self, Z, T, Y):
        """Wald estimator for binary Z and T."""
        E_Y_Z1 = np.mean(Y[Z == 1])
        E_Y_Z0 = np.mean(Y[Z == 0])
        E_T_Z1 = np.mean(T[Z == 1])
        E_T_Z0 = np.mean(T[Z == 0])
        
        return (E_Y_Z1 - E_Y_Z0) / (E_T_Z1 - E_T_Z0)


class CausalForest:
    """Simplified Causal Forest for heterogeneous treatment effects."""
    
    def __init__(self, n_estimators=100, max_depth=5):
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.trees = []
    
    def _build_tree(self, X, T, Y, depth=0):
        """Build a single tree for treatment effect estimation."""
        n = len(Y)
        
        if depth >= self.max_depth or n < 10:
            # Leaf: estimate CATE
            if np.sum(T) == 0 or np.sum(1-T) == 0:
                return {'leaf': True, 'cate': 0}
            
            # Simple difference in means
            cate = np.mean(Y[T == 1]) - np.mean(Y[T == 0])
            return {'leaf': True, 'cate': cate}
        
        # Find best split (simplified: random feature and threshold)
        best_feature = np.random.randint(X.shape[1])
        best_threshold = np.median(X[:, best_feature])
        
        left_mask = X[:, best_feature] <= best_threshold
        right_mask = ~left_mask
        
        if np.sum(left_mask) < 5 or np.sum(right_mask) < 5:
            cate = np.mean(Y[T == 1]) - np.mean(Y[T == 0])
            return {'leaf': True, 'cate': cate}
        
        left_tree = self._build_tree(X[left_mask], T[left_mask], 
                                     Y[left_mask], depth + 1)
        right_tree = self._build_tree(X[right_mask], T[right_mask],
                                      Y[right_mask], depth + 1)
        
        return {
            'leaf': False,
            'feature': best_feature,
            'threshold': best_threshold,
            'left': left_tree,
            'right': right_tree
        }
    
    def fit(self, X, T, Y):
        """Fit causal forest."""
        self.trees = []
        for _ in range(self.n_estimators):
            # Bootstrap sample
            idx = np.random.choice(len(Y), len(Y), replace=True)
            tree = self._build_tree(X[idx], T[idx], Y[idx])
            self.trees.append(tree)
        return self
    
    def _predict_tree(self, tree, X):
        """Predict using a single tree."""
        if tree['leaf']:
            return tree['cate'] * np.ones(len(X))
        
        predictions = np.zeros(len(X))
        left_mask = X[:, tree['feature']] <= tree['threshold']
        
        predictions[left_mask] = self._predict_tree(tree['left'], X[left_mask])
        predictions[~left_mask] = self._predict_tree(tree['right'], X[~left_mask])
        
        return predictions
    
    def predict(self, X):
        """Predict CATE for each sample."""
        predictions = np.array([self._predict_tree(tree, X) for tree in self.trees])
        return np.mean(predictions, axis=0)


class MetaLearner:
    """Meta-learners for heterogeneous treatment effects."""
    
    def __init__(self, base_learner='rf'):
        self.base_learner = base_learner
    
    def _get_learner(self):
        if self.base_learner == 'rf':
            return RandomForestRegressor(n_estimators=100)
        return LinearRegression()
    
    def s_learner(self, X, T, Y, X_test):
        """S-learner: treat treatment as feature."""
        XT = np.column_stack([X, T])
        
        model = self._get_learner()
        model.fit(XT, Y)
        
        # Predict under treatment and control
        XT1 = np.column_stack([X_test, np.ones(len(X_test))])
        XT0 = np.column_stack([X_test, np.zeros(len(X_test))])
        
        return model.predict(XT1) - model.predict(XT0)
    
    def t_learner(self, X, T, Y, X_test):
        """T-learner: separate models for treated and control."""
        mask_treated = T == 1
        mask_control = T == 0
        
        model_treated = self._get_learner()
        model_control = self._get_learner()
        
        model_treated.fit(X[mask_treated], Y[mask_treated])
        model_control.fit(X[mask_control], Y[mask_control])
        
        return model_treated.predict(X_test) - model_control.predict(X_test)
    
    def x_learner(self, X, T, Y, X_test, propensity_scores=None):
        """X-learner: impute treatment effects."""
        # Stage 1: fit outcome models
        mask_treated = T == 1
        mask_control = T == 0
        
        model_treated = self._get_learner()
        model_control = self._get_learner()
        
        model_treated.fit(X[mask_treated], Y[mask_treated])
        model_control.fit(X[mask_control], Y[mask_control])
        
        # Stage 2: impute treatment effects
        D1 = Y[mask_treated] - model_control.predict(X[mask_treated])
        D0 = model_treated.predict(X[mask_control]) - Y[mask_control]
        
        # Stage 3: meta-learner on imputed effects
        model_meta = self._get_learner()
        
        # For treated: use D1
        # For control: use D0
        # Simplified: just average the two estimates
        
        cate_treated = model_treated.predict(X_test) - model_control.predict(X_test)
        
        return cate_treated


class BackdoorAdjustment:
    """Estimate causal effect using backdoor adjustment."""
    
    def __init__(self, n_bins=10):
        self.n_bins = n_bins
    
    def estimate(self, X, T, Y, Z):
        """
        Estimate P(Y | do(X)) using backdoor adjustment.
        
        Args:
            X: Treatment
            T: Not used (for consistency)
            Y: Outcome
            Z: Confounders to adjust for
        """
        # Discretize Z
        Z_binned = np.zeros_like(Z)
        for j in range(Z.shape[1]):
            Z_binned[:, j] = np.digitize(Z[:, j], 
                                         np.percentile(Z[:, j], 
                                                      np.linspace(0, 100, self.n_bins + 1)[1:-1]))
        
        # Compute adjustment
        ate = 0
        for z_val in np.unique(Z_binned, axis=0):
            mask = np.all(Z_binned == z_val, axis=1)
            if mask.sum() == 0:
                continue
            
            # P(Z = z)
            p_z = mask.sum() / len(mask)
            
            # P(Y | X = 1, Z = z) - P(Y | X = 0, Z = z)
            mask_treated = mask & (X == 1)
            mask_control = mask & (X == 0)
            
            if mask_treated.sum() > 0 and mask_control.sum() > 0:
                ate += p_z * (np.mean(Y[mask_treated]) - np.mean(Y[mask_control]))
        
        return ate


class FrontdoorAdjustment:
    """Estimate causal effect using frontdoor adjustment."""
    
    def estimate(self, X, T, Y, M):
        """
        Estimate P(Y | do(X)) using frontdoor adjustment.
        
        Args:
            X: Treatment
            T: Not used
            Y: Outcome
            M: Mediators
        """
        # Step 1: P(M | do(X)) = P(M | X) (no backdoor to M)
        ate = 0
        for m_val in np.unique(M):
            mask_m = M == m_val
            
            # P(M = m | X = x)
            p_m_x1 = np.mean(M[X == 1] == m_val)
            p_m_x0 = np.mean(M[X == 0] == m_val)
            
            # P(Y | do(X = x'), M = m)
            # Use frontdoor formula
            mask_treated = mask_m & (X == 1)
            mask_control = mask_m & (X == 0)
            
            if mask_treated.sum() > 0 and mask_control.sum() > 0:
                # E[Y | M = m, X = x'] summed over x'
                e_y_m_x1 = np.mean(Y[mask_treated])
                e_y_m_x0 = np.mean(Y[mask_control])
                
                # P(X = x')
                p_x1 = np.mean(X == 1)
                p_x0 = np.mean(X == 0)
                
                ate += (p_m_x1 * (e_y_m_x1 * p_x1 + e_y_m_x0 * p_x0) - 
                       p_m_x0 * (e_y_m_x1 * p_x1 + e_y_m_x0 * p_x0))
        
        return ate


# Example usage
if __name__ == "__main__":
    # Generate synthetic data with confounding
    np.random.seed(42)
    n = 1000
    
    # Confounders
    Z = np.random.randn(n, 2)
    
    # Treatment (depends on confounders)
    prob_treat = 1 / (1 + np.exp(-Z[:, 0] - 0.5 * Z[:, 1]))
    T = np.random.binomial(1, prob_treat)
    
    # Outcome (depends on treatment and confounders)
    Y = 2 * T + Z[:, 0] + 0.5 * Z[:, 1] + np.random.randn(n) * 0.5
    
    # True ATE = 2
    
    # IPW
    ipw = IPWEstimator()
    ate_ipw, se_ipw = ipw.estimate(np.column_stack([Z]), T, Y)
    print(f"IPW ATE: {ate_ipw:.3f} (SE: {se_ipw:.3f})")
    
    # Doubly Robust
    dr = DoublyRobustEstimator()
    ate_dr, se_dr = dr.estimate(np.column_stack([Z]), T, Y)
    print(f"Doubly Robust ATE: {ate_dr:.3f} (SE: {se_dr:.3f})")
    
    # Backdoor adjustment
    bd = BackdoorAdjustment()
    ate_bd = bd.estimate(T, T, Y, Z)
    print(f"Backdoor ATE: {ate_bd:.3f}")
    
    print(f"\nTrue ATE: 2.000")

8. Summary

Causal inference provides the framework for understanding interventions:

1. SCMs and DAGs encode causal assumptions graphically
2. Do-calculus provides rules for deriving causal quantities
3. Potential outcomes define causal effects counterfactually
4. Propensity scores enable adjustment for observed confounders
5. Instrumental variables handle unobserved confounding
6. Heterogeneous treatment effects (CATE) enable personalized decisions

The key insight is that correlation is not causation — causal inference provides the tools to move from predictive models to prescriptive ones.