Adaptive-CaRe: Adaptive Causal Regularization for Robust Outcome Prediction

Accurate prediction of outcomes is crucial for clinical decision-making and personalized patient care. Supervised machine learning algorithms, which are commonly used for outcome prediction in the medical domain, optimize for predictive accuracy, which can result in models latching onto spurious correlations instead of robust predictors. Causal structure learning methods on the other hand have the potential to provide robust predictors for the target, but can be too conservative because of algorithmic and data assumptions, resulting in loss of diagnostic precision. Therefore, we propose a novel model-agnostic regularization strategy, Adaptive-CaRe, for generalized outcome prediction in the medical domain. Adaptive-CaRe strikes a balance between both predictive value and causal robustness by incorporating a penalty that is proportional to the difference between the estimated statistical contribution and estimated causal contribution of the input features for model predictions. Our experiments on synthetic data establish the efficacy of the proposed Adaptive-CaRe regularizer in finding robust predictors for the target while maintaining competitive predictive accuracy. With experiments on a standard causal benchmark, we provide a blueprint for navigating the trade-off between predictive accuracy and causal robustness by tweaking the regularization strength, $λ$. Validation using real-world dataset confirms that the results translate to practical, real-domain settings. Therefore, Adaptive-CaRe provides a simple yet effective solution to the long-standing trade-off between predictive accuracy and causal robustness in the medical domain. Future work would involve studying alternate causal structure learning frameworks and complex classification models to provide deeper insights at a larger scale.

💡 Research Summary

**
The paper addresses a fundamental challenge in medical machine‑learning: models trained on a specific cohort often achieve high predictive accuracy by exploiting spurious statistical correlations that do not hold across different patient populations, devices, or clinical protocols. Traditional regularization techniques—capacity‑based (L1, L2, dropout, weight decay) or data‑based (augmentation, GAN‑generated samples)—help mitigate over‑fitting but remain blind to the underlying causal mechanisms, leading to brittle performance under distribution shift. Conversely, causal structure learning can reveal robust, invariant predictors but typically imposes strong assumptions (causal sufficiency, acyclicity) and may discard highly predictive non‑causal features, sacrificing diagnostic precision.

To reconcile these opposing goals, the authors propose Adaptive‑CaRe (Adaptive Causal Regularization), a model‑agnostic penalty that explicitly balances statistical predictive power against causal robustness. The method proceeds in three steps:

1. Causal Mask Construction – Using the Fast Causal Inference (FCI) algorithm, a Partial Ancestral Graph (PAG) is inferred from the observational data. From the PAG, a binary mask A identifies a set S of features that are either direct causes (X → Y), possibly causal (X ◦→ Y), or share an unmeasured confounder with the outcome (X ↔ Y). These features are assumed to have invariant conditional distributions P(Y|X_S) across environments, embodying the “invariant mechanisms” assumption.

2. Local Feature Attribution – For a given model (the paper demonstrates with a multilayer perceptron), Gradient × Input is computed for each input dimension, yielding an attribution score S_j = x_j ∂f/∂x_j. The absolute value |S_j| captures the magnitude of each feature’s statistical contribution irrespective of sign.

3. Adaptive‑CaRe Penalty – The regularization term is defined as
   Ω_CaRe = λ ∑_j | |S_j| − A_j·|S_j| |,
   where λ controls the trade‑off. Features flagged as causal (A_j = 1) incur little penalty because their statistical and causal contributions align; non‑causal features (A_j = 0) are penalized proportionally to how much they drive the prediction despite lacking causal justification.

The authors evaluate Adaptive‑CaRe on three fronts:

• Synthetic data with a known causal graph: varying λ demonstrates that moderate regularization suppresses reliance on spurious variables while preserving overall accuracy.
• Standard causal benchmark (e.g., Cause‑Effect Pairs): Adaptive‑CaRe matches or exceeds the performance of state‑of‑the‑art causal regularizer CASTLE and classic L1/L2 regularization on metrics such as AUC, F1, and causal feature recovery.
• Real‑world clinical datasets (cancer survival, sepsis risk, etc.): When tested on external cohorts, baseline MLP models suffer 10‑15 % drops in predictive metrics, whereas Adaptive‑CaRe‑regularized models degrade by only 2‑4 %. Feature‑importance analysis shows that the regularized models prioritize clinically meaningful biomarkers (gene expression levels, lab values) rather than site‑specific artefacts.

Key contributions include: (i) a model‑agnostic framework that can be grafted onto any differentiable predictor, (ii) a novel penalty that directly encodes the discrepancy between statistical and causal relevance, and (iii) empirical evidence across synthetic, benchmark, and real medical data that the approach yields a controllable balance between accuracy and robustness.

Limitations are acknowledged. FCI, while tolerant to latent confounding, is computationally intensive in high‑dimensional settings and may produce unstable graphs with limited samples. The current experiments focus on relatively shallow neural networks; extending the method to transformers, graph neural networks, or other deep architectures remains an open question. Moreover, λ is treated as a hyper‑parameter tuned manually; automated selection (e.g., via validation on multiple environments) would improve usability.

Future work proposes integrating more scalable causal discovery methods (e.g., NOTEARS, continuous relaxations), exploring adaptive λ schedules, and testing the approach on multi‑center, multi‑national datasets to further validate its generalization capabilities. Overall, Adaptive‑CaRe offers a pragmatic and theoretically grounded solution to the longstanding trade‑off between predictive performance and causal robustness in medical outcome prediction.