Causal Invariance Learning via Efficient Optimization of a Nonconvex Objective

Zhenyu Wang, Yifan Hu, Peter Bühlmann, Zijian Guo

Abstract:

Data from multiple environments offer valuable opportunities to uncover causal relationships among
variables. Leveraging the assumption that the causal outcome model remains invariant across heterogeneous environments, state-of-the-art methods attempt to identify causal outcome models by learning
invariant prediction models and rely on exhaustive searches over all (exponentially many) covariate subsets. These approaches present two major challenges: 1) determining the conditions under which the
invariant prediction model aligns with the causal outcome model, and 2) devising computationally efficient causal discovery algorithms that scale polynomially, instead of exponentially, with the number of
covariates. To address both challenges, we focus on the additive intervention regime and propose nearly
necessary and sufficient conditions for ensuring that the invariant prediction model matches the causal
outcome model. Exploiting the essentially necessary identifiability conditions, we introduce Negative
Weight Distributionally Robust Optimization (NegDRO), a nonconvex continuous minimax optimization whose global optimizer recovers the causal outcome model. Unlike standard group DRO problems
that maximize over the simplex, NegDRO allows negative weights on environment losses, which break
the convexity. Despite its nonconvexity, we demonstrate that a standard gradient method converges to
the causal outcome model, and we establish the convergence rate with respect to the sample size and
the number of iterations. Our algorithm avoids exhaustive search, making it scalable especially when
the number of covariates is large. The numerical results further validate the efficiency of the proposed
method.