Bridging Multicalibration and Out-of-distribution Generalization Beyond Covariate Shift
Jiayun Wu, Jiashuo Liu, Peng Cui, Zhiwei Steven Wu
TL;DR
This work links multicalibration to out-of-distribution generalization by extending the calibration notion to joint grouping functions over covariates and outcomes, enabling robustness beyond covariate shift. It establishes theoretical connections to invariance (IRM) under concept shift, characterizes the maximal grouping function space as a linear span of density ratios, and decomposes grouping functions to balance accuracy and invariance. A post-processing algorithm, MC-PseudoLabel, is proposed to achieve extended multicalibration and invariant predictions with lightweight optimization and convergence guarantees in Gaussian regimes. Empirical evaluation on real-world distribution-shift datasets demonstrates strong performance advantages and robust generalization under distribution shifts. The framework offers a practical pathway for robust regression under OOD conditions and suggests avenues for extending to classification tasks.
Abstract
We establish a new model-agnostic optimization framework for out-of-distribution generalization via multicalibration, a criterion that ensures a predictor is calibrated across a family of overlapping groups. Multicalibration is shown to be associated with robustness of statistical inference under covariate shift. We further establish a link between multicalibration and robustness for prediction tasks both under and beyond covariate shift. We accomplish this by extending multicalibration to incorporate grouping functions that consider covariates and labels jointly. This leads to an equivalence of the extended multicalibration and invariance, an objective for robust learning in existence of concept shift. We show a linear structure of the grouping function class spanned by density ratios, resulting in a unifying framework for robust learning by designing specific grouping functions. We propose MC-Pseudolabel, a post-processing algorithm to achieve both extended multicalibration and out-of-distribution generalization. The algorithm, with lightweight hyperparameters and optimization through a series of supervised regression steps, achieves superior performance on real-world datasets with distribution shift.