# Optimization-based Causal Estimation from Heterogenous Environments

arXiv: https://arxiv.org/abs/2109.11990 · HTML: https://mlsanity.com/paper/2109.11990/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 78 references that are in the map.
- First posted on arXiv 24 Sep 2021.

## Builds on (most cited first, 30 of 37)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Causality for Machine Learning](https://mlsanity.com/paper/1911.10500/) (2019)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [Learning De-biased Representations with Biased Representations](https://mlsanity.com/paper/1910.02806/) (2019)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization](https://mlsanity.com/paper/2109.02934/) (2021)
- [Invariant Causal Prediction for Nonlinear Models](https://mlsanity.com/paper/1706.08576/) (2017)
- [Invariance, Causality and Robustness](https://mlsanity.com/paper/1812.08233/) (2018)
- [Anchor regression: heterogeneous data meets causality](https://mlsanity.com/paper/1801.06229/) (2018)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Best Subset Selection via a Modern Optimization Lens](https://mlsanity.com/paper/1507.03133/) (2015)
- [Meta-Learning without Memorization](https://mlsanity.com/paper/1912.03820/) (2019)
- [A causal framework for distribution generalization](https://mlsanity.com/paper/2006.07433/) (2020)
- [Invariant Causal Prediction for Sequential Data](https://mlsanity.com/paper/1706.08058/) (2017)
- [Out of Distribution Generalization in Machine Learning](https://mlsanity.com/paper/2103.02667/) (2021)
- [Stable Prediction across Unknown Environments](https://mlsanity.com/paper/1806.06270/) (2018)
- [Causal Discovery from Changes](https://mlsanity.com/paper/1301.2312/) (2013)
- [DecAug: Out-of-Distribution Generalization via Decomposed Feature Representation and Semantic Augmentation](https://mlsanity.com/paper/2012.09382/) (2020)
- [Causal Discovery and Forecasting in Nonstationary Environments with State-Space Models](https://mlsanity.com/paper/1905.10857/) (2019)
- [Nonlinear Invariant Risk Minimization: A Causal Approach](https://mlsanity.com/paper/2102.12353/) (2021)
- [Learning Causal Structures Using Regression Invariance](https://mlsanity.com/paper/1705.09644/) (2017)
- [Out-of-distribution Prediction with Invariant Risk Minimization: The Limitation and An Effective Fix](https://mlsanity.com/paper/2101.07732/) (2021)

## Built on by (most cited first, 16 of 16)

- [Probable Domain Generalization via Quantile Risk Minimization](https://mlsanity.com/paper/2207.09944/) (2022)
- [Towards LLMs Robustness to Changes in Prompt Format Styles](https://mlsanity.com/paper/2504.06969/) (2025)
- [Environment Invariant Linear Least Squares](https://mlsanity.com/paper/2303.03092/) (2023)
- [Data Augmentations for Improved (Large) Language Model Generalization](https://mlsanity.com/paper/2310.12803/) (2023)
- [Causal Optimal Transport of Abstractions](https://mlsanity.com/paper/2312.08107/) (2023)
- [An Invariant Learning Characterization of Controlled Text Generation](https://mlsanity.com/paper/2306.00198/) (2023)
- [Causality Pursuit from Heterogeneous Environments via Neural Adversarial Invariance Learning](https://mlsanity.com/paper/2405.04715/) (2024)
- [Conditionally Invariant Representation Learning for Disentangling Cellular Heterogeneity](https://mlsanity.com/paper/2307.00558/) (2023)
- [Causal Invariance Learning via Efficient Nonconvex Optimization](https://mlsanity.com/paper/2412.11850/) (2024)
- [Causal Effect Identification in Heterogeneous Environments from Higher-Order Moments](https://mlsanity.com/paper/2506.11756/) (2025)
- [Bayesian Invariance Modeling of Multi-Environment Data](https://mlsanity.com/paper/2506.22675/) (2025)
- [Multi-environment Topic Models](https://mlsanity.com/paper/2410.24126/) (2024)
- [Addressing Label Shift in Distributed Learning via Entropy Regularization](https://mlsanity.com/paper/2502.02544/) (2025)
- [Class-Conditional Distribution Balancing for Group Robust Classification](https://mlsanity.com/paper/2504.17314/) (2025)
- [Interpretable Clustering with Adaptive Heterogeneous Causal Structure Learning in Mixed Observational Data](https://mlsanity.com/paper/2509.04415/) (2025)
- [Hessian Rank Constraint for Learning Structure of Nonlinear Latent Variable Models](https://mlsanity.com/paper/2609.24241/) (2026)

## Abstract (arXiv)

This paper presents a new optimization approach to causal estimation. Given data that contains covariates and an outcome, which covariates are causes of the outcome, and what is the strength of the causality? In classical machine learning (ML), the goal of optimization is to maximize predictive accuracy. However, some covariates might exhibit a non-causal association with the outcome. Such spurious associations provide predictive power for classical ML, but they prevent us from causally interpreting the result. This paper proposes CoCo, an optimization algorithm that bridges the gap between pure prediction and causal inference. CoCo leverages the recently-proposed idea of environments, datasets of covariates/response where the causal relationships remain invariant but where the distribution of the covariates changes from environment to environment. Given datasets from multiple environments-and ones that exhibit sufficient heterogeneity-CoCo maximizes an objective for which the only solution is the causal solution. We describe the theoretical foundations of this approach and demonstrate its effectiveness on simulated and real datasets. Compared to classical ML and existing methods, CoCo provides more accurate estimates of the causal model and more accurate predictions under interventions.