🤖 AI Summary
This paper addresses the challenge of robust out-of-distribution prediction under distributional shift and latent confounding. To this end, it proposes a novel inference objective—Boosted Control Function (BCF)—grounded in a newly defined strong invariance principle. It establishes, for the first time, strong invariance for nonlinear, non-identifiable structural functions; develops a rigorous BCF theoretical framework with minimax-optimal guarantees under worst-case distribution shifts; and introduces SIMDGs, a unified model integrating causal inference and distribution generalization. Leveraging the ControlTwicing algorithm alongside flexible machine learning estimators—including neural networks and tree-based models—the approach achieves significant improvements over empirical risk minimization on both synthetic and real-world benchmarks. Experimental results demonstrate robust cross-distribution prediction performance, validating BCF’s strong adaptability to latent confounding and non-stationarity.
📝 Abstract
Modern machine learning methods and the availability of large-scale data have significantly advanced our ability to predict target quantities from large sets of covariates. However, these methods often struggle under distributional shifts, particularly in the presence of hidden confounding. While the impact of hidden confounding is well-studied in causal effect estimation, e.g., instrumental variables, its implications for prediction tasks under shifting distributions remain underexplored. This work addresses this gap by introducing a strong notion of invariance that, unlike existing weaker notions, allows for distribution generalization even in the presence of nonlinear, non-identifiable structural functions. Central to this framework is the Boosted Control Function (BCF), a novel, identifiable target of inference that satisfies the proposed strong invariance notion and is provably worst-case optimal under distributional shifts. The theoretical foundation of our work lies in Simultaneous Equation Models for Distribution Generalization (SIMDGs), which bridge machine learning with econometrics by describing data-generating processes under distributional shifts. To put these insights into practice, we propose the ControlTwicing algorithm to estimate the BCF using flexible machine-learning techniques and demonstrate its generalization performance on synthetic and real-world datasets compared to traditional empirical risk minimization approaches.
