Research Pipeline Utilities for Causal Inference

Introduction

Chapters 1–9 developed the theory, methodology, and validation of Double Machine Learning for causal inference. This chapter sketches research pipeline utilities for organizing model versions, diagnostics, and retraining decisions. The companion repo code is not a production deployment framework; it is a book companion that illustrates the additional requirements a real deployment would need to satisfy.

The production paradox

Standard machine learning production systems monitor prediction accuracy against ground truth. A deployed image classifier can be evaluated against labeled images. A recommendation system can be evaluated by click-through rates. The feedback loop is direct: prediction $\rightarrow$ observation $\rightarrow$ accuracy.

Causal inference breaks this paradigm fundamentally:

This creates a production monitoring challenge without parallel in standard ML. Instead of monitoring prediction accuracy, we must monitor the conditions under which our estimates remain valid:

1. Overlap violations (positivity breakdown).
2. Treatment distribution shift (assignment mechanism change).
3. Nuisance model degradation (first-stage fit quality).
4. Effect stability (is the treatment effect itself changing?).

Production monitoring: standard ML vs. causal inference

Model serialization for DML

DML models have a unique architecture requiring specialized serialization.

Components to serialize

A deployment-oriented DML model would need to serialize:

1. Nuisance models for each cross-fitting fold:
   • Propensity model: $\hat{e}(X) \approx P(T=1 \mid X)$.
   • Outcome model: $\hat{m}(X) \approx \mathbb{E}[Y \mid X]$.
2. Cross-fitting configuration (K folds, temporal blocking).
3. Feature transformer (scaling, encoding).
4. HAC parameters for time series standard errors.
5. Training metadata (data snapshot, hyperparameters).

@dataclass
class DMLModelVersion:
    """Versioned DML model with full serialization support."""

    version_id: str  # Unique hash-based identifier
    created_at: str
    model_type: str  # "double_ml", "temporal_plr_dml", etc.
    n_folds: int

    # Serialized nuisance models by fold
    nuisance_models: Dict[int, Tuple[bytes, bytes]]

    # Feature specification
    feature_names: List[str]
    treatment_name: str
    outcome_name: str

    # Reproducibility
    hyperparameters: Dict[str, Any]
    metrics: Dict[str, float]  # R² propensity, R² outcome

    @classmethod
    def create(
        cls,
        model_type: str,
        nuisance_models: Dict[int, Tuple[Any, Any]],
        feature_names: List[str],
        treatment_name: str,
        outcome_name: str,
        hyperparameters: Dict[str, Any],
        **kwargs,
    ) -> "DMLModelVersion":
        """Create versioned model with content-based hash."""
        # Serialize nuisance models per fold
        serialized = {}
        for fold_idx, (prop, out) in nuisance_models.items():
            serialized[fold_idx] = (
                pickle.dumps(prop),
                pickle.dumps(out),
            )

        # Generate version ID from content hash
        timestamp = datetime.utcnow().isoformat()
        version_id = f"dml-{compute_hash(...)[:12]}"

        return cls(version_id=version_id, ...)

Model registry operations

A DML-aware model registry supports:

• Registration: store versioned models with full metadata.
• Retrieval: load a specific version for inference.
• Promotion: candidate $\rightarrow$ active workflow.
• Rollback: revert to a previous active version.
• Lineage: track which data/config produced which model.

# Initialize registry
registry = DMLModelRegistry("./models/dml_registry")

# After training
version_id = registry.register(model_version)

# Demo promotion workflow
registry.promote_to_staging(version_id)

# After validation passes
registry.promote_to_production()

# If issues detected
registry.rollback(to_version="dml-abc123")

Persistence format

Models are persisted as directories containing:

dml-abc123xyz/
├── metadata.json          # Version info, hyperparameters
├── propensity_fold_0.pkl  # Serialized propensity model
├── outcome_fold_0.pkl     # Serialized outcome model
├── propensity_fold_1.pkl
├── outcome_fold_1.pkl
├── ...
└── feature_transformer.pkl  # Optional scaler/encoder

This structure enables atomic updates (complete version or nothing), efficient fold-level updates (not the full model), human-readable metadata inspection, and Git-compatible versioning (track metadata.json).

Inference APIs

Deployed inference for causal models differs from standard ML prediction.

Batch vs. real-time inference

Inference patterns for causal models

API design for CATE estimation

A hypothetical deployed CATE API would need to provide:

class CATEInferenceAPI:
    """Hypothetical API for treatment effect estimation."""

    def estimate_cate(
        self,
        features: Dict[str, Any],
        return_confidence: bool = True,
    ) -> CATEResponse:
        """
        Estimate conditional treatment effect.

        Args:
            features: Covariate values for the unit
            return_confidence: Include confidence intervals

        Returns:
            CATEResponse with effect estimate and metadata
        """
        X = self._transform_features(features)

        # Compute CATE using averaged nuisance models
        cate = self._compute_cate(X)

        # Compute confidence interval
        ci_lower, ci_upper = self._compute_ci(X, cate)

        # Check overlap (propensity extremity)
        propensity = self._predict_propensity(X)
        overlap_warning = self._check_overlap(propensity)

        return CATEResponse(
            cate=cate,
            ci_lower=ci_lower,
            ci_upper=ci_upper,
            propensity=propensity,
            overlap_warning=overlap_warning,
            model_version=self._current_version,
        )

Response structure

A complete CATE response includes uncertainty and validity flags:

@dataclass
class CATEResponse:
    cate: float           # Point estimate
    ci_lower: float       # 95% CI lower
    ci_upper: float       # 95% CI upper
    propensity: float     # P(T=1|X) for this unit
    overlap_warning: bool # True if propensity extreme
    model_version: str    # For reproducibility
    timestamp: str        # Inference timestamp

Causal-specific monitoring

The heart of production causal inference is monitoring the conditions for valid inference, not prediction accuracy.

Overlap violations (positivity)

The positivity assumption requires $0 < P(T=1 \mid X) < 1$ for all $X$. In practice, we need propensity scores bounded away from 0 and 1:

def check_overlap_violations(
    propensity_scores: np.ndarray,
    clip_min: float = 0.01,
    clip_max: float = 0.99,
) -> MonitoringResult:
    """Check for positivity violations."""

    too_low = propensity_scores < clip_min
    too_high = propensity_scores > clip_max
    violation_rate = (np.sum(too_low) + np.sum(too_high)) / len(propensity_scores)

    if violation_rate >= 0.10:  # >10% violations
        level = AlertLevel.CRITICAL
        message = f"CRITICAL: {violation_rate:.1%} overlap violations"
    elif violation_rate >= 0.05:
        level = AlertLevel.WARNING
        message = f"WARNING: {violation_rate:.1%} overlap violations"
    else:
        level = AlertLevel.OK
        message = f"Overlap OK: {violation_rate:.1%} violations"

    return MonitoringResult(
        check_name="overlap_violations",
        level=level,
        value=violation_rate,
        message=message,
    )

Treatment distribution shift

Unlike standard feature drift, treatment distribution shift indicates the assignment mechanism may have changed. This is causal-specific:

• Standard ML: features changed $\rightarrow$ model may underfit.
• Causal inference: treatment assignment changed $\rightarrow$ identification may fail.

def check_treatment_shift(
    treatment_current: np.ndarray,
    treatment_baseline: np.ndarray,
) -> MonitoringResult:
    """Detect treatment distribution shift."""

    # For binary treatment: proportion difference
    # For continuous: KS statistic
    if is_binary(treatment_current):
        p_current = np.mean(treatment_current)
        p_baseline = np.mean(treatment_baseline)
        metric = abs(p_current - p_baseline)
    else:
        metric, _ = ks_2samp(treatment_current, treatment_baseline)

    if metric >= 0.10:
        level = AlertLevel.CRITICAL
        message = "Treatment assignment mechanism may have changed"
    # ...

Nuisance model degradation

DML validity requires well-fitting nuisance models. We monitor nuisance $R^2$, not treatment effect accuracy (which is unknowable):

def check_nuisance_degradation(
    r2_propensity: float,
    r2_outcome: float,
    warning_threshold: float = 0.50,
    critical_threshold: float = 0.30,
) -> MonitoringResult:
    """Check nuisance model fit quality."""

    min_r2 = min(r2_propensity, r2_outcome)

    if min_r2 < critical_threshold:
        level = AlertLevel.CRITICAL
        message = f"Nuisance fit poor (R²={min_r2:.2f}). DML may be biased."
    elif min_r2 < warning_threshold:
        level = AlertLevel.WARNING
        message = f"Nuisance fit degraded (R²={min_r2:.2f})"
    else:
        level = AlertLevel.OK

    return MonitoringResult(...)

Effect stability

Large changes in estimated treatment effects over time warrant investigation:

def check_effect_stability(
    current_effect: float,
    baseline_effect: float,
    current_se: Optional[float] = None,
    baseline_se: Optional[float] = None,
) -> MonitoringResult:
    """Monitor treatment effect stability over time."""

    if abs(baseline_effect) < 1e-10:
        relative_change = float("inf") if abs(current_effect) > 1e-10 else 0.0
    else:
        relative_change = abs(current_effect - baseline_effect) / abs(baseline_effect)

    if relative_change >= 0.50:  # 50% change
        level = AlertLevel.CRITICAL
        message = f"Effect changed {relative_change:.0%}: investigate cause"
    elif relative_change >= 0.20:  # 20% change
        level = AlertLevel.WARNING
    # ...

Effect changes may indicate genuine time-varying effects (valid, requires dynamic modeling), model degradation (problematic), selection/survivorship bias in new data (problematic), or a changed treatment definition (problematic).

Retraining triggers

Standard ML retraining triggers focus on prediction accuracy degradation. Causal inference requires different triggers aligned with identification assumptions.

Causal-specific triggers

Retraining triggers: standard ML vs. causal inference

Scheduler design

class RetrainScheduler:
    """Intelligent retraining scheduler for DML pipelines."""

    def __init__(self, config: RetrainSchedulerConfig):
        self.config = config
        self.monitor = CausalMonitor()

    def evaluate_retrain_need(
        self,
        monitoring_results: List[MonitoringResult],
        n_samples: int,
    ) -> Optional[RetrainTrigger]:
        """Evaluate if retraining is needed."""

        # Check cooldown (avoid excessive retraining)
        if self.is_in_cooldown():
            return None

        # Check minimum sample threshold
        if n_samples < self.config.min_samples:
            return None

        # Map monitoring results to triggers
        for result in monitoring_results:
            if result.level == AlertLevel.CRITICAL:
                return RetrainTrigger(
                    trigger_type=self._result_to_trigger_type(result),
                    severity=AlertLevel.CRITICAL,
                    reason=result.message,
                )

        # Check scheduled retraining
        return self.check_scheduled_retrain()

Trigger types

class TriggerType(Enum):
    SCHEDULED = "scheduled"           # Time-based
    OVERLAP_VIOLATION = "overlap"     # Positivity breakdown
    TREATMENT_SHIFT = "treatment"     # Assignment mechanism
    NUISANCE_DEGRADATION = "nuisance" # First-stage fit
    EFFECT_INSTABILITY = "effect"     # Effect magnitude change
    COVARIATE_SHIFT = "covariate"     # Feature distribution
    MANUAL = "manual"                 # Human-initiated

End-to-end research pipeline

This section assembles the companion components into a reproducible research pipeline.

Pipeline architecture

Figure: research DML pipeline architecture.

Data Ingestion (FRED + Insurance)
        |
        v
Feature Engineering
        |
        v
DML Estimation (Cross-fitting)  <───────────────┐
        |                                        │
        v                                        │
Causal Monitoring                                │
        |                                        │
        v                                        │
   Retrain?  ── yes ──>  Retrain Pipeline ───────┘
        |
        | no
        v
Model Serving

InsuranceDMLPipeline implementation

class InsuranceDMLPipeline:
    """Research DML pipeline for insurance pricing analysis."""

    def __init__(self, config: PipelineConfig):
        self.config = config
        self._registry = DMLModelRegistry(config.model_registry_path)
        self._monitor = CausalMonitor()
        self._scheduler = RetrainScheduler(monitor=self._monitor)

    def fit(
        self,
        X: np.ndarray,
        T: np.ndarray,
        Y: np.ndarray,
        baseline_ate: Optional[float] = None,
    ) -> PipelineResult:
        """
        Fit DML pipeline to data.

        Stages:
        1. Data preparation (scaling)
        2. Cross-fitted nuisance estimation
        3. ATE estimation with orthogonalization
        4. HAC standard errors (time series)
        5. Monitoring checks
        6. Model versioning
        """
        # Prepare data
        X_scaled, T, Y = self._prepare_data(X, T, Y)

        # Create cross-validator
        cv = self._create_cross_validator(len(T))

        # Fit nuisance models per fold
        propensity_scores = np.zeros(len(T))
        outcome_predictions = np.zeros(len(T))
        nuisance_models = {}

        for fold_idx, (train_idx, test_idx) in enumerate(cv.split(X_scaled)):
            prop_model, out_model = self._fit_fold(
                X_scaled, T, Y, train_idx, test_idx
            )
            propensity_scores[test_idx] = self._predict_propensity(
                prop_model, X_scaled[test_idx]
            )
            outcome_predictions[test_idx] = out_model.predict(X_scaled[test_idx])
            nuisance_models[fold_idx] = (prop_model, out_model)

        # Compute ATE via orthogonalization
        ate, ate_se = self._compute_ate(
            Y, outcome_predictions, T, propensity_scores
        )

        # Run monitoring
        monitoring_results = self._monitor.run_all_checks(
            propensity_scores=propensity_scores,
            r2_propensity=self._compute_r2(propensity_scores, T),
            r2_outcome=self._compute_r2(outcome_predictions, Y),
            current_effect=ate,
            baseline_effect=baseline_ate,
        )

        # Create and register model version
        version = DMLModelVersion.create(
            model_type="double_ml",
            nuisance_models=nuisance_models,
            feature_names=self.config.feature_columns,
            treatment_name=self.config.treatment_column,
            outcome_name=self.config.outcome_column,
            hyperparameters=self._get_hyperparameters(),
        )
        version_id = self._registry.register(version)

        return PipelineResult(
            ate=ate,
            ate_se=ate_se,
            ate_ci_lower=ate - 1.96 * ate_se,
            ate_ci_upper=ate + 1.96 * ate_se,
            nuisance_metrics={"r2_propensity": ..., "r2_outcome": ...},
            monitoring_results=monitoring_results,
            model_version=version_id,
        )

Workflow orchestration

For a hardened deployment, integrate with Airflow or Prefect:

# airflow_dags/dml_retrain.py

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from datetime import datetime, timedelta

default_args = {
    "owner": "causal-ml",
    "retries": 1,
    "retry_delay": timedelta(minutes=5),
}

with DAG(
    dag_id="dml_retrain_pipeline",
    default_args=default_args,
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False,
) as dag:

    load_data = PythonOperator(
        task_id="load_data",
        python_callable=load_inference_data,
    )

    run_monitoring = PythonOperator(
        task_id="run_monitoring",
        python_callable=run_causal_monitoring,
    )

    evaluate_retrain = BranchPythonOperator(
        task_id="evaluate_retrain",
        python_callable=evaluate_retrain_trigger,
    )

    retrain_model = PythonOperator(
        task_id="retrain_model",
        python_callable=retrain_dml_pipeline,
    )

    validate_model = PythonOperator(
        task_id="validate_model",
        python_callable=validate_new_model,
    )

    promote_staging = PythonOperator(
        task_id="promote_staging",
        python_callable=promote_to_staging,
    )

    skip_retrain = PythonOperator(
        task_id="skip_retrain",
        python_callable=lambda: None,
    )

    # Define dependencies
    load_data >> run_monitoring >> evaluate_retrain
    evaluate_retrain >> [retrain_model, skip_retrain]
    retrain_model >> validate_model >> promote_staging

Insurance pricing case study

We conclude with a complete example applying the research/demo pipeline to insurance competitor pricing.

Problem setup

An insurance company wants to estimate how competitor pricing affects policy retention:

• Treatment: competitor price difference (continuous).
• Outcome: retention indicator (binary).
• Confounders: customer demographics, policy features, macroeconomic conditions.
• Time series: quarterly observations with autocorrelation.

Pipeline configuration

from sklearn.ensemble import GradientBoostingClassifier, GradientBoostingRegressor
from dml_ts.production import InsuranceDMLPipeline, PipelineConfig

config = PipelineConfig(
    n_folds=5,
    model_registry_path="./models/insurance_dml",

    # Nuisance models
    propensity_model=GradientBoostingClassifier(
        n_estimators=100,
        max_depth=4,
        random_state=42,
    ),
    outcome_model=GradientBoostingRegressor(
        n_estimators=100,
        max_depth=4,
        random_state=42,
    ),

    # Time series configuration
    use_hac=True,
    time_column="quarter",

    # Feature specification
    feature_columns=[
        "age", "income", "tenure", "coverage_amount",
        "gdp_growth", "unemployment", "interest_rate",
    ],
    treatment_column="competitor_price_diff",
    outcome_column="retained",
)

pipeline = InsuranceDMLPipeline(config)

Training and monitoring

# Load data
X, T, Y = load_insurance_data()

# Fit pipeline
result = pipeline.fit(X, T, Y)

print(f"ATE: {result.ate:.4f}")
print(f"95% CI: [{result.ate_ci_lower:.4f}, {result.ate_ci_upper:.4f}]")
print(f"Nuisance R² (propensity): {result.nuisance_metrics['r2_propensity']:.3f}")
print(f"Nuisance R² (outcome): {result.nuisance_metrics['r2_outcome']:.3f}")

# Check monitoring
for check in result.monitoring_results:
    if not check.is_ok():
        print(f"ALERT: {check.message}")

# Mark as active in the demo registry
pipeline.promote_to_production()

Deployment-oriented inference

# Evaluate retraining need with new data
X_new, T_new = load_new_quarter_data()

results, trigger = pipeline.evaluate_retrain_need(
    X_new, T_new,
    X_baseline=X[:len(X_new)],
    T_baseline=T[:len(X_new)],
)

if trigger is not None:
    print(f"Retrain triggered: {trigger.reason}")
    # Initiate retraining workflow
else:
    print("Model stable - no retraining needed")

    # Predict propensity for new customers
    propensity = pipeline.predict_propensity(X_new)

    # Flag high-risk customers (extreme propensity)
    high_risk = (propensity < 0.05) | (propensity > 0.95)
    print(f"High-risk customers: {np.sum(high_risk)} ({np.mean(high_risk):.1%})")

Exercises

Exercise 10.1: monitoring dashboard

Design a monitoring dashboard for a production DML system. What visualizations would you include? Consider:

1. Propensity score distribution over time.
2. Treatment distribution shift tracking.
3. Nuisance model $R^2$ trends.
4. Effect estimate stability with confidence bands.

Exercise 10.2: retraining threshold calibration

The default overlap violation threshold of 10% for CRITICAL alerts may be too aggressive or lenient depending on the application. Describe a procedure for calibrating this threshold using:

1. Monte Carlo simulation with known treatment effects.
2. Historical data from past model deployments.

Exercise 10.3: A/B test integration

Propose how to integrate randomized A/B tests with a production DML system. How can occasional randomization provide:

1. Validation of DML estimates.
2. Ground truth for effect stability monitoring.
3. Retraining triggers when observational estimates diverge from experimental.

Exercise 10.4: multi-treatment extension

Extend the production pipeline to handle multiple treatments simultaneously. What additional monitoring is needed when:

1. Treatments are mutually exclusive (one of K).
2. Treatments can be combined ($2^K$ combinations).
3. Treatment assignment is continuous multivariate.

Summary

This chapter developed research/demo infrastructure patterns for causal inference systems:

The companion dml_ts.production module currently provides demo utilities:

• DMLModelVersion, DMLModelRegistry: model versioning.
• CausalMonitor: overlap, treatment shift, nuisance, effect monitoring.
• RetrainScheduler: causal-specific retraining triggers.
• InsuranceDMLPipeline: end-to-end integration.

The next and final appendix provides a roadmap for implementing this book’s methodology in Julia, offering performance advantages for large-scale causal inference.