Causal Invariance Learning via Efficient Nonconvex Optimization
Zhenyu Wang, Yifan Hu, Peter Bühlmann, Zijian Guo
TL;DR
This work addresses learning the direct causal drivers of an outcome from multi-environment observational data by enforcing invariance of the causal outcome model. It introduces NegDRO, a continuous nonconvex minimax formulation allowing negative weights to enforce risk invariance across environments, which avoids combinatorial subset searches. Under additive interventions, it derives concrete identification conditions ensuring the invariant model recovers the causal model, proves a benign optimization landscape where stationary points are near the true causal predictor, and provides a gradient-based algorithm with non-asymptotic convergence guarantees. It further shows that even with limited additive interventions, NegDRO can identify $\beta^*$ and outperforms existing invariant-learning methods, with strong scalability to high-dimensional covariates. Overall, the approach offers a theoretically grounded, computationally efficient path to causal discovery in heterogeneous observational data, with practical implications for fields like marketing and epidemiology.
Abstract
Identifying the causal relationship among variables from observational data is an important yet challenging task. This work focuses on identifying the direct causes of an outcome and estimating their magnitude, i.e., learning the causal outcome model. Data from multiple environments provide valuable opportunities to uncover causality by exploiting the invariance principle that the causal outcome model holds across heterogeneous environments. Based on the invariance principle, we propose the Negative Weighted Distributionally Robust Optimization (NegDRO) framework to learn an invariant prediction model. NegDRO minimizes the worst-case combination of risks across multiple environments and enforces invariance by allowing potential negative weights. Under the additive interventions regime, we establish three major contributions: (i) On the statistical side, we provide sufficient and nearly necessary identification conditions under which the invariant prediction model coincides with the causal outcome model; (ii) On the optimization side, despite the nonconvexity of NegDRO, we establish its benign optimization landscape, where all stationary points lie close to the true causal outcome model; (iii) On the computational side, we develop a gradient-based algorithm that provably converges to the causal outcome model, with non-asymptotic convergence rates in both sample size and gradient-descent iterations. In particular, our method avoids exhaustive combinatorial searches over exponentially many subsets of covariates found in the literature, ensuring scalability even when the dimension of the covariates is large. To our knowledge, this is the first causal invariance learning method that finds the approximate global optimality for a nonconvex optimization problem efficiently.