Boosted Control Functions: Distribution generalization and invariance in confounded models

TL;DR

This paper introduces the Boosted Control Function (BCF), a new invariant target for prediction under distributional shifts with hidden confounding, supported by a theoretical framework and practical estimation algorithm.

Contribution

It proposes a novel invariance concept and the BCF, enabling distribution generalization in confounded models, along with the SIMDG framework and the ControlTwicing algorithm for estimation.

Findings

BCF achieves worst-case optimality under shifts.

ControlTwicing outperforms existing methods on datasets.

Theoretical guarantees support the invariance properties.

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 nonparametric machine-learning techniques and study its generalization performance on synthetic and real-world datasets compared to robust and empirical risk minimization approaches.