Bridging Multicalibration and Out-of-distribution Generalization Beyond Covariate Shift

Authors: jiayun  wu
jiashuo  liu
peng  cui
zhiwei steven wu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 18:30:23

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.

Main paper: 
 introduction: 
 We revisit the problem of out-of-distribution generalization and establish new connections with multicalibration [17] , a criterion originating from algorithmic fairness. Multicalibration is a strengthening of calibration, which only requires a predictor f to be correct on average within each level set:E[Y -f (X) | f (X)] = 0Calibration is a relatively weak property, as it can be satisfied even by the uninformative constant predictor f (X) = E[Y ] that predicts the average outcome. More broadly, calibration provides only a marginal guarantee that does not extend to sub-populations. Multicalibration [17] mitigates this issue by requiring the calibration to hold over a family of (overlapping) subgroups H: for all h ∈ H,E[(Y -f (X)) h(X) | f (X)] = 0Multicalibration is initially studied as measure of subgroup fairness for boolean grouping functions h, with h(X) = 1 indicating X is a member of group h [17] . Subsequently, Gopalan et al. [14] and Kim et al. [20] adopt a broader class of real-valued grouping functions that can identify subpopulations through reweighting. The formulation of real-valued grouping function has enabled surprising connections between multicalibration and distribution shifts. Prior work [21, 37] studied how distribution shift affects the measure of multicalibration, with a focus on covariate shift where the relationship between X and Y remains fixed. Kim et al. [21] show that whenever the set of realvalued grouping functions H includes the density ratio between the source and target distributions, a multicalibrated predictor with respect to the source remains calibrated in the shifted target distribution.Our work substantially expands the connections between multicalibration and distribution shifts. At a high level, our results show that robust prediction under distribution shift can actually be facilitated by multicalibration. We extend the notion of multicalibration by incorporating grouping functions that simultaneously consider both covariates X and outcomes Y . This extension enables us to go beyond covariate shift and account for concept shift, which is prevalent in practice due to spurious correlation, missing variables, or confounding [31] .Our contributions. Based on the introduction of joint grouping functions, we establish new connections between our extended multicalibration notion and algorithmic robustness in the general setting of out-of-distribution generalization, where the target distribution to assess the model is different from the source distribution to learn the model.1. We first revisit the setting of covariate shift and show multicalibration implies Bayes optimality under covariate shift, provided a sufficiently rich class of grouping functions. Then, in the setting of concept shifts, we show the equivalence of multicalibration and invariance [2] , a learning objective to search for a Bayes optimal predictor E[Y |Φ(X)] under a representation over features Φ(X), even though E[Y |X] is different across target distributions. We show correspondence between an invariant representation Φ(X) and a multicalibrated predictor E[Y |Φ(X)], with a grouping function class containing all density ratios of target distributions and the source distribution.2. As part of our structural analysis of the new multicalibration concept, we investigate the maximal grouping function class that allows for a nontrivial multicalibrated predictor. For traditional covariatebased grouping functions, the Bayes optimal predictor f (X) = E[Y |X] is always multicalibrated, which is no longer the case for joint grouping functions. We show the maximal grouping function class is a linear space spanned by the density ratio of the target distributions where the predictor is invariant. As a structural characterization of distribution shift, this leads to an efficient parameterization of the grouping functions by linear combination of a spanning set of density ratios. The spanning set can be flexibly designed to incorporates implicit assumptions of various methodologies for robust learning, including multi-environment learning [39] and hard sample learning [29] .3. We devise a post-processing algorithm to multicalibrate predictors and simultaneously producing invariant predictors. As a multicalibration algorithm, we prove its convergence under Gaussian distributions of data and certify multicalibration upon convergence. As a robust learning algorithm, the procedure is plainly supervised regression with respect to models' hypothesis class and grouping function class, introducing an overhead of linear regression. This stands out from heavy optimization techniques for out-of-distribution generalization, such as bi-level optimization [12, 30] and multi-objective learning [1, 2, 24] , which typically involves high-order gradients [36] . The algorithm introduces no extra hyperparameters. This simplifies model selection, which is a significant challenge for out-of-distribution generalization since validation is unavailable where the model is deployed [15] . Under the standard model selection protocol of DomainBed [15] , the algorithm achieves superior performance to existing methods in real-world datasets with concept shift, including porverty estimation [44] , personal income prediction [7] and power consumption [32, 33] prediction. 
multicalibration and bayes optimality under covariate shift 2.1 multicalibration with joint grouping functions: 
 We consider prediction tasks where covariates are denoted by a random vector X ∈ X and the target by Y ∈ Y. Lowercase x, y denote the specific values of these random variables. Predictors are defined as real-valued functions f : X → Y. Our theoretical analysis focuses on the setting where Y = [0, 1]. In this context, we propose a new definition of ℓ 2 approximate multicalibration with joint grouping functions. Definition 2.1 (Multicalibration with Joint Grouping Functions). For a probability measure P (X, Y ) and a predictor f , let H ⊂ R X ×Y be a real-valued grouping function class. We say that f is α-approximately ℓ 2 multicalibrated w.r.t. H and P if for all h ∈ H:EQUATIONP f (X) (v) = P (f -1 (v)) is the pushforward measure. We say f is α-approximately calibrated if H includes the constant function h ≡ 1. We say f is multicalibrated (calibrated) for α = 0. If the grouping function is defined on X, which implies h(x, y 1 ) = h(x, y 2 ) for any x ∈ X and y 1 , y 2 ∈ Y, we abbreviate h(x, •) by h(x).Our definition generalizes several notions of (multi)calibration in the literature by specific choices of grouping functions. For example, K 2 (f, 1, P ) recovers the overall calibration error in the case of a constant grouping function h ≡ 1. For boolean grouping functions defined on X [17] , K 2 (f, h, P ) computes the calibration error of the subgroup with h(x) = 1. For real-valued grouping functions defined on X [14, 20] , K 2 (f, h, P ) evaluates a reweighted calibration error, whose weights h(x) are proportional to the likelihood of a sample belonging to the subgroup. Furthermore, we propose an extended domain of grouping functions defined on covariates and outcomes jointly, for which the Bayes optimal predictor E[Y |x] may not be multicalibrated, in contrast to all existing multicalibration frameworks with X-based grouping functions. Multicalibration with joint groupings thus implies a distinct learning objective from accuracy, which we will characterize as invariance in section 3. Example 2.2 (Multicalibration Does Not Imply Bayes Optimality). Consider covariates X = (X 1 , X 2 ) T and an outcome Y generated by the following structural equations:Y = X 1 + ϵ 1 . X 2 = Y + ϵ 2 .Settings of Out-of-distribution Generalization. We characterize distribution shift by an uncertainty set of absolutely continuous probability measures, denoted by P(X, Y ), where there is an accessible source measure P S ∈ P and unknown target measure P T ∈ P. We use capital letters such as P to denote a single probability measure and lowercase letters such as p to denote its probability density function. A predictor f is learned in the source distribution P S and assessed in the target distribution P T . Given a loss function ℓ : Y × Y → R, we evaluate the average risk of a predictor f w.r.t. a probability measure P , defined by R P (f ) := E P [ℓ(f (X), Y )].We focus on ℓ(ŷ, y) = (ŷ -y) 2 in our theoretical analyses.In this subsection we focus on grouping functions h(x) defined on covariates. We will prove approximately multicalibrated predictors simultaneous approaches Bayes optimality in each target distribution with covariate shift, bridging the results of Kim et al. [21] and Globus-Harris et al. [13] . To recap, Kim et al. [21] studies multicalibration under covariate shift and shows that a multicalibrated predictor remains calibrated in target distribution for a sufficiently large grouping function class. Further, it is shown that multicalibration predictors remain multicalibrated under covariate shift [21, 37] , assuming the grouping function class H is closed under some transformation by density ratios (Assumption 2.3.1). Second, Globus-Harris et al. [13] shows multicalibration implies Bayes optimal accuracy [13] , assuming H satisfies a weak learning condition (Assumption 2. EQUATIONRemark 2.5. Following prior work in multicalibration [13, 37] , we study functions f with finite cardinality, which can be obtained by discretization. 
multicalibration and invariance under concept shift: 
 Theorem 2.4 shows multicalibration implies Bayes optimal accuracy for target distributions under covariate shift. However, in practical scenarios, there are both marginal distribution shifts of covariates (X) and concept shift of the conditional distributions (Y |X). Concept shift is especially prevalent in tabular data due to missing variables and confounding [31] . In order to go beyond covariate shift, we will focus on grouping functions defined on covariates and outcomes jointly. We show that multicalibration notion w.r.t. joint grouping functions is equivalent to invariance, a criterion for robust prediction under concept shift. Extending the robustness of multicalibration to general shift is non-trivial. The fundamental challenge is that there is no shared predictor that is generally optimal in each target distribution because the Bayes optimal predictor varies for different Y |X distributions. As a first step, we show multicalibrated predictors w.r.t. joint grouping functions are robust as they are optimal over any post-processing functions in each target distribution. Theorem 3.1 (Risk Bound under Concept Shift). For a set of absolutely continuous probability measures P(X, Y ) containing the source measure P S (X, Y ), consider a predictor f : X → [0, 1]. Assume the grouping function class H satisfies the following condition:EQUATION)If f is α-approximately ℓ 2 multicalibrated w.r.t. H and P S , then for any measure P ∈ P(X, Y ),EQUATIONThe theorem shows an approximately multicalibrated predictor on the source almost cannot be improved by post-processing for each target distribution. To ensure such robustness, the grouping function class must include all density ratios between target and source measures, which are functions over X × Y. This characterization of robustness in terms of post-processing echoes with Invariant Risk Minimization (IRM) [2] , a paradigm for out-of-distribution generalization with Y |X shift. However, their analysis focuses on representation learning. Definition 3.2 (Invariant Predictor). Consider data selected from multiple environments in the set E where the probability measure in an environment e ∈ E is denoted by P e (X, Y ). Denote the representation over covariates by a measurable function Φ(x). We say that Φ elicits an α-approximately invariant predictor g * • Φ across E if there exists a function g * ∈ G := {g : supp(Φ) → [0, 1]} such that for all e ∈ E :EQUATIONRemark 3.3. (1) Predictors in G take a representation Φ extracted from the covariates as input. For a general predictor f (x), if we take Φ(x) = f (x) and g * as an identity function, Equation 5 reduces to the form of Equation 4. Therefore, f in Equation 4 is a 2 √ α-approximately invariant predictor across environments collected from the uncertainty set P. (2) We give an approximate definition of invariant predictors, which recovers the original definition [2] when α = 0. In this case, there exists a shared Bayes optimal predictor g ⋆ across environments, taking Φ as input. This impliesE e1 [Y |Φ] = E e2 [Y |Φ] almost surely for any e 1 , e 2 .IRM searches for a representation such that the optimal predictors upon the representation are invariant across environments. Motivated from causality, the interaction between outcomes and their causes are also assumed invariant, so IRM learns a representation of causal variables for stable prediction. We extend Theorem 3.1 to representation learning and prove equivalence between multicalibrated and invariant predictors. Theorem 3.4 (Equivalence of Multicalibration and Invariance). Assume samples are drawn from an environment e ∈ E with a prior P S (e) such that e∈E P S (e) = 1 and P S (e) > 0. The overall population satisfies P S (X, Y ) = e∈E P e (X, Y )P S (e) where P e (X, Y ) is the environment-specific absolutely continuous measure. With a measurable function Φ(x), define a function class H as:EQUATION1. If there is a bijection g (1) In the first statement, assuming g ⋆ is a bijection avoids degenerate cases where Φ contains redundant information. For example, every predictor g ⋆ • Φ(X) upon representation Φ equals g ⋆ (Φ) • I(X) upon representation X. Confining g ⋆ to bijections ensures some unique decomposition into predictors and representations. (2) Wald et al. [41] proves equivalence between exact invariance and simultaneous calibration in each environment. We strengthen their result to show multicalibration on a single source distribution suffices for invariance. Moreover, our results can be directly extended beyond their multi-environment setting to a general uncertainty set of target distributions, by the mapping between grouping functions and density ratios. Further, our theorem is established for both exact and approximate invariance.⋆ : supp(Φ) → [0, 1] such that g ⋆ • Φ is α-approximately ℓ 2 multicalibrated w.r.t. H and P S , then Φ elicits an 2 √ α-approximately invariant predictor g ⋆ • Φ across E . 2. If there is g ⋆ : supp(Φ) → [0, 1] such that Φ elicits an α-approximately invariant predictor g ⋆ • Φ across E , then g ⋆ • Φ is α/D-The theorem bridges approximate multicalibration with approximate invariance for out-of-distribution generalization beyond covariate shift. The equivalence property indicates that the density ratios of target and source distributions constitute the minimal grouping function class required for robust prediction in terms of invariance. 
structure of grouping function classes: 
 Section 3 inspires one to construct richer grouping function classes for stronger generalizability. However, fewer predictors are multicalibrated to a rich function class, and a multicalibrated predictor may not exist at all, as illustrated by the example in section 2.1. In this section, we first study the maximal grouping function class that is feasible for a multicalibrated predictor. Then, we will leverage our structural results to inform the design of grouping functions.We focus on continuous grouping functions defined on a compact set X × Y ⊂ R d+1 , i.e., h ∈ C(X × Y), and consider absolutely continuous probability measures supported on X × Y with continuous density functions. Our first proposition shows that the maximal grouping function class for any predictor is a linear space. Proposition 4.1 (Maximal Grouping Function Class). Given an absolutely continuous probability measure P S (X, Y ) and a predictor f : X → [0, 1], define the maximal grouping function class that f is multicalibrated with respect to:H f := {h ∈ C(X × Y) : K 2 (f, h, P S ) = 0}. (7) Then H f is a linear space.In the following, we further analyze the spanning set of maximal grouping function classes for nontrivial predictors which are at least calibrated. Theorem 4.2 (Spanning Set). Consider an absolutely continuous probability measure P S (X, Y ) and a calibrated predictor f : X → [0, 1]. Then its maximal grouping function class H f is given by: EQUATIONEQUATION= (Λ, Ω) T . Define H ′ 1,Φ := {h(Φ(x)) ∈ C(X × Y)}, with H ′ 1,Φ ⊂ H 1,Φ . H ′ 1,X and H ′ 1,Λ are similarly defined. We have:1. H ′ 1,X ⊃ H ′ 1,Φ ⊃ H ′ 1,Λ ⊃ H ′ 1,∅ = {C}. 2. {C} = H 2,X ⊂ H 2,Φ ⊂ H 2,Λ ⊂ H 2,∅ .C is a constant value function.The objective of a robust learning method can be represented by a tuple consisting of an assumption about the boundary of distribution shift and a metric of robustness. Multicalibration is equivalent to invariance as a metric of robustness, while the grouping function class provides a unifying view for assumptions over potential distribution shift. Given any uncertainty set of target distributions P, Theorem 4.2 implies an efficient and reasonable construction of grouping functions as linear combinations of density ratios from P. We implement two designs of grouping functions for the learning setting with and without environment annotations respectively.From Environments If samples are drawn from multiple environments and the environment annotations are available, we assume the uncertainty set as the union of each environment's distribution P e . This completely recovers IRM's objective, but we approach it with a different optimization technique in the next section. Taking pooled data as the source S, density ratios spanning the grouping function class are h e (x, y) = p e (x, y)/p S (x, y) = p S (e|x, y)/p S (e), where p S (e|x, y) is estimated by an environment classifier. Then a grouping function can be represented as a linear combination of h e :EQUATIONFrom Hard Samples When data contains latent sub-populations without annotations, the uncertainty set can be constructed by identifying sub-populations. Hard sample learning [27, 28, 29] suggests the risk is an indicator for sub-population structures. Samples from the minority sub-population M are more likely to have high risks. For example, JTT [29] identified the minority subgroup using a risk threshold of a trained predictor f id . We adopt a continuous grouping by assuming P S (X, Y ∈ M ) ∝ (f id (X) -Y ) 2 . We construct the uncertainty set as the union of the source S and minority sub-population M , resulting in a grouping function represented as:EQUATIONAnother design utilizing Distributionally Robust Optimization's assumption [10] is in section B. 
mc-pseudolabel: an algorithm for extended multicalibration: 
 In this section, we introduce an algorithm for multicalibration with respect to joint grouping functions. Simultaneously, the algorithm also provides a new optimization paradigm for invariant prediction under distribution shift. The algorithm, called MC-PseudoLabel, post-processes a trained model by supervised learning with pseudolabels generated by grouping functions. As shown in Algorithm 1, given a predictor function class F and a dataset D with an empirical distribution PD (X, Y ), a regression oracle A solves the optimization:A F (D) = arg min f ∈F R PD (f ).We take as input a model f 0 , possibly trained by Empirical Risk Minimization. f 0 has a finite range following conventions of prior work in multicalibration [13] . For continuous predictors, we discretize the model output and introduce a small rounding error (see section C). For each iteration, the algorithm performs regression with grouping functions on each level set of the model. The prediction of grouping functions rectify the uncalibrated model and serves as pseudolabels for model updates.Since we regress Y with grouping functions defined on Y , a poor design of groupings violating Theorem 4.2 can produce trivial outputs. For example, if grouping functions contain h(x, y) = y, then err t-1 -ẽ rr t never decreases and the algorithm outputs f 0 , because there does not exist a multicalibrated predictor. However, the algorithm certifies a multicalibrated output if it converges. Theorem 5.1 (Certified Multicalibration). In Algorithm 1, for α, B > 0, if err t-1 -ẽ rr t ≤ α B , the output f ′ t-1 (x) is α-approximately ℓ 2 multicalibrated w.r.t. H B = {h ∈ H : sup h(x, y) 2 ≤ B}.f ′ t ← Round(f t ; m) := arg min v∈[1/m] |f t (x) -v|;6:err t = E x,y∼D [(f ′ t (x) -y) 2 ]; 7:foreach v ∈ [1/m] do 8: D t v ← D|f ′ t (x) = v; 9: h t v (x, y) ← A H (D t v ); {Regression on level sets with grouping functions.} 10:end for11: ft+1 (x, y) ← v∈[ 1 m ] 1 {f ′ t (x)=v} • h t v (x, y); {Generate pseudolabels.} 12: ẽ rr t+1 ← E x,y∼D [( ft+1 (x, y) -y) 2 ];13:D t+1 ← (D x , ft+1 (D));14:f t+1 (x) ← A F (D t+1); {Update the model with pseudolabels.} 15:t ← t + 1; 16: until err t-1 -ẽ rr t stops decreasing. Ensure:f ′ t-1 (x).MC-PseudoLabel reduces to LSBoost Globus-Harris et al. [13] , a boosting algorithm for multicalibration if H only contains covariate-based grouping functions. In this case, Line 14 of Algorithm 1 reduces to f t+1 (x) = ft+1 (x, •) where ft+1 does not depend on y. For joint grouping functions, since ft+1 ∈ R X ×Y , we project it to models' space of R X by learning the model with ft+1 as pseudolabels. The projection substantially changes the optimization dynamics. LSBoost constantly decreases risks of models, due toR PD (f t+1 ) = R PD ( ft+1 ) < R PD (f t ).The projection step disrupts the monotonicity of risks, implying that MC-Pseudolabel can output a predictor with a higher risk than input. This is because multicalibration with joint grouping functions implies balance between accuracy and invariance, as is discussed in Theorem 4.3. The convergence of LSBoost relies on the monotonicity of risks, which is not applicable to MC-Pseudolabel. We study the algorithm's convergence in the context of representation learning. F.27)). Consider X ∈ R d with X = (Φ, Ψ) T . Assume that (Φ, Ψ, Y ) follows a multivariate normal distribution N d+1 (µ, Σ)where the random variables are in general position such that Σ is positive definite. For any distribution D supported on X × Y, take the predictor class F = R X and the grouping function class H as a subset of H 2,Φ which is defined in Equation 10:EQUATION)For an initial predictorf 0 (x) = E[Y |x], run MC-Pseudolabel(D, H, F) without rounding, then f t (x) converges pointwise to E[Y |Φ(x)] as t → ∞, with a convergence rate of O(M (Σ) t ) where 0 ≤ M (Σ) < 1.MC-Pseudolabel is also an optimization paradigm for invariance. Certified multicalibration in Theorem 5.1 also implies certified invariance. Furthermore, MC-Pseudolabel introduces no extra hyperparameters to tradeoff between risks and robustness. Both certified invariance and light-weighted hyperparameters simplify model selection, which is challenging for out-distribution generalization because of unavailable validation data from target distributions [15] . MC-Pseudolabel has lightweighted optimization consisting of a series of supervised regression. It introduces an overhead to Empirical Risk Minimization by performing regression on level sets. However, the extra burden is linear regression by designing the grouping function class as linear space. Furthermore, regression on different level sets can be parallelized. Computational complexity is further analyzed in section D.6 Experiments 
settings: 
 We benchmark MC-Pseudolabel on real-world regression datasets with distributional shift. We adopt two experimental settings. For the multi-environment setting, algorithms are provided with training data collected from multiple annotated environments. Thereafter, the trained model is assessed on new environments. For the single-environment setting, algorithms are trained on a single source distribution. There could be latent sub-populations in training data, but environment annotations are unavailable. The trained model is assessed on a target dataset with distribution shift from the source. The grouping function class is implemented according to Equation 11and Equation 12for the multi-environment and single-environment setting respectively.Datasets We experiment on PovertyMap [44] and ACSIncome [7] for the multi-environment setting, and VesselPower [33] for the single-environment setting. As the only regression task in WILDS [23] , a popular benchmark for in-the-wild distribution shift, PovertyMap performs poverty index estimation for different spatial regions by satellite images. Data are collected from both urban and rural regions, by which the environment is annotated. The test dataset also covers both environments, but is collected from different countries. The primary metric is Worst-U/R Pearson, the worst Pearson correlation of prediction between rural and urban regions. The other two datasets are tabular, where natural concept shift (Y |X shift) is more common due to existence of missing variables and hidden confounders [31] . ACSIncome [7] performs personal income prediction with data collected from US Census sources across different US states. The task is converted to binary classification by an income threshold, but we take raw data for regression. Environments are partitioned by different occupations with similar average income. VesselPower comes from Shifts [32, 33] , a benchmark focusing on regression tasks with real-world distributional shift. The objective is to predict power consumption of a merchant vessel given navigation and weather data. Data are sampled under different time and wind speeds, causing distribution shift between training and test data.Baselines For the multi-environment setting, baselines include ERM (Empirical Risk Minimization); methods for invariance learning which mostly adopts multi-objective optimization: IRM [2] , MIP [24] , IB-IRM [1] , CLOvE [41] , MRI [18] , REX [25] , Fishr [36] ; an alignment-based method from domain generalization: IDGM [39] ; and Group DRO [38] . Notably, CLOvE learns a calibrated predictor simultaneously on all environments, but it is optimized by multi-objective learning with a differentiable regularizer for calibration. For the singe-environment setting, baselines include reweighting based techniques: CVaR [26] , JTT [29] , Tilted-ERM [27] ; a Distributionally Robust Optimization method χ 2 -DRO [8] ; and a data augmentation method C-Mixup [43] . Other methods are not included because of specification in classification [45, 46] or exposure to target distribution data during training [19, 22] . For all experiments, we train an Oracle ERM with data sampled from target distribution.Implementation We implement the predictor with MLP for ACSIncome and VesselPower, and Resnet18-MS [16] for PovertyMap, following WILDS' default architecture. We follow DomainBed's protocol [15] for model selection. Specifically, we randomly sample 20 sets of hyperparameters for each method, containing both the training hyperparameters and extra hyperparameters from the robust learning algorithm. We select the best model across hyperparameters based on three model selection criteria, including in-distribution validation on the average of training data, worst-environment validation with the worst performance across training environments, and oracle validation on target data. Oracle validation is not recommended by DomainBed, which suggests limited numbers of access to target data. The entire run is repeated with different seeds for three times to measure standard errors of performances. Specifically for PovertyMap, we perform 5-fold cross validation instead of three repeated experiments, following WILDS' setup.Results are shown in Table 1 for multi-environment settings and which shows that MC-Pseudolabel's out-of-distribution errors strongly correlates with in-distribution errors. The experiment spans across different hyperparameters and seeds with the same model architecture on VesselPower. The phenomenon, known as accuracy-on-the-line [34] , is well known for a general class of models under covariate shift. However, Liu et al. [31] shows accuracy-on-the-line does not exist under concept shift, which is the case for ERM and C-Mixup. This introduces significant challenge for model selection. However, MC-Pseudolabel recovers the accuracy to the line. 
conclusion: 
 To conclude, we establish a new optimization framework for out-of-distribution generalization through extended multicalibration with joint grouping functions. While the current algorithm focuses on regression, there is potential for future work to extend our approach to general forms of tasks, particularly in terms of classification. 
 
 Appendix: Multicalibration Multicalibration is first proposed by Hébert-Johnson et al. [17] with binary grouping functions. Kim et al. [20] and Gopalan et al. [14] extend the grouping functions to realvalued functions. Globus-Harris et al. [13] shows that with a sufficiently rich class of real-valued grouping functions, multicalibration actually implies accuracy. There are also meta-learning methods [12] that handles multi-environment generalization with bi-level optimization. For the single environment setting, Distributionally Robust Optimization optimizes for the worst-case risk in an uncertainty set of distributions centering around the source distribution [8, 9, 11, 26, 40] . Another branch of research is targeted at mitigating spurious correlation with an assumption of simplicity bias, which utilizes a simple model to discover latent sub-populations and then correct the biased predictor by sample reweighting [27, 28, 29] , retraining on a subgroupbalanced dataset or a small batch from target distribution [19, 22, 46] , or perform Invariant Risk Minimization on discovered subgroups [5] . Data augmentation is a prevalent technique to enhance algorithmic robustness for vision tasks. Quite a lot of these methods are tailored for classification. For example, Mixup [45] interpolates between features of samples with the same label. The approach is extended to regression settings by C-Mixup [43] . Pseudolabelling is a common technique for out-of-distribution generalization, but typically adopted in a setting with exposure to unlabelled samples from target distribution, known as domain adaptation [42] . However, MC-Pseudolabel generate pseudolabels for the source distribution itself.Distributionally Robust Optimization assumes the target distribution to reside in an uncertainty set P of distributions centering around the source distribution P S . For example, Duchi et al. [10] formulates the uncertainty set as arbitrary subgroups that has a proportion of at least α 0 ∈ (0, 1). Duchi et al. only consider subgroups of covariates: P(X) = {P (X) : there exists a probability measure P ′ (X),EQUATIONEQUATIONBy the correspondence between density ratios and grouping functions, the equivalent design of a grouping function class is given by:EQUATIONWe can also extend the grouping functions to consider both covariates and outcomes, such that general subgroups are incorporated into the uncertainty set:EQUATIONIn the case of grouping functions defined on x and y jointly, the grouping function class is not closed under affine transformation and is not a linear space spanned by density ratios, which suggests that a perfectly multicalibrated solution might not exist. However, approximately multicalibrated predictors can still be pursued.For continuous predictors, we take a preprocessing step to discretize the model to as many bins as possible such that the rounding error is negligible while still ensuring enough samples in individual bins. Specifically, we equally split the outcomes of predictors to bins with equal intervals from the minimum to maximum of model output. We start from a minimum bin number m = 10, and keeps increasing m as long as 90% of the samples reside in a bin with at least 30 samples. When the criterion is violated, we stop increasing m and select it as the final bin number. The model discretization procedure is fixed across all experiments.We assume that the predictor's outcomes are uniformly distributed. Denote the average bin size by N b , which is a constant around 30 in our implementation. The bin number is given by m = N/N b where N is the sample size. For neural networks, N represents the batch size. The overhead of MC-Pseudolabel compared to Empirical Risk Minimization is linear regression on each bin, whose sample complexity is O(N 3 b ) with OLS. Please note that an individual linear regression for around 30 samples is extremely cheap. A non-parallel implementation of regression on every bin scales linearly with the bin number m, so the overall complexity is O(N 2 b N ). However, since the regression on each bin is independent, we adopt a multi-processing implementation. Denote the number of jobs by J, the overall time cost of MC-Pseudolabel is O(N 2 b N/J). As a comparison, OLS on N samples has a computational complexity of O(N 3 ).In conclusion, the complexity of MC-Pseudolabel scales linearly with sample size (or batch size for neural networks). Counterintuitively, increasing the bin number m (and thus decreasing the bin size) actually decreases the computational complexity. This is because linear regression scales cubically with sample size, so decreasing the sample size in each bin is preferred to decreasing the bin number.We start from a multi-environment synthetic dataset with a multivariate normal distribution corresponding to Theorem 5.2. In this experiment, we examine the optimization dynamics of MC-Pseudolabel. The data generation process is inspired by Arjovsky et al. [2] . The covariates can be sliced into X = (S, V ) T with S ∈ R 9 and V ∈ R, where S is the causal variable for Y and V is the spurious variable. The data is generated by the following structural equations:EQUATION)EQUATIONEQUATION)The covariate V is spuriously correlated with Y because the coefficient α V (e) depends on the specific environment e. We set α V (E) = 1.25, 0.75 respectively for two training environments while α V (E) extrapolates to -1 during testing. A robust algorithm is supposed to bypass the spurious variable and output a predictor f (X) = α T S S in order to survive the test distribution where the correlation between V and Y is negated.The predictor class for this dataset is linear models, and the environment classifier is implemented by MLP with a single hidden layer. In this experiment, we fix the training hyperparameters for the base linear model, and perform grid search over the extra hyperparameter introduced by robust learning methods. Baselines except for ERM and Group DRO share a hyperparameter which is the regularizer's coefficient, and Group DRO introduces a temperature hyperparameter. We search over their hyperparameter space and report RMSE metric on the test set in Figure 2 . Most baselines exhibit a U-turn with an increasing hyperparameter, and the minimum point varies across methods. The sensitivity of hyperparameters implies the dependence on a strong model selection criterion, such as oracle model selection on target distribution. However, the dashed line for MC-Pseudolabel's error is tangent to all the U-turns of baselines, indicating a competitive performance of MC-Pseudolabel both with and without oracle model selection. We also investigate the evolution of pseudo labels ft in Algorithm 1 to recover the dynamics of MC-Pseudolabel. The first row of Figure 3 demonstrates how pseudolabelling results in a multicalibrated predictor. It shows that pseudolabels for two environments deviate from model prediction at Step 0, but the gap quickly converges at Step 4, implying multicalibrated prediction. The second row provides insight about how pseudolabelling contributes to an invariant predictor. We observe that the curve of two environments are gradually merging because the pseudolabel introduces a special noise to the original label such that the correlation between the pseudolabel and spurious variable V is weakened. As a result, the predictor will depend on the causal variable S which is relatively more strongly correlated with pseudolabels.In figure 4 , we provide the correlation between models' in-distribution validation performance and out-of-distribution test performance across all methods. Datasets. In the ACSInome dataset, we focus on predicting personal income for California residents across four occupation fields: science professions (SCI), protective services (PRT), education (EDU), and military (MIL). While average incomes are similar across these fields, the correlation between income and usual hours worked per week (WKHP) varies significantly, as shown in Table 3 . This spurious correlation between income and WKHP introduces a concept shift across occupations. Therefore, we train the model on environments comprising the SCI, PRT, and EDU occupations and evaluate it on MIL. We follow the standard setup for PovertyMap as specified in the WILDS benchmark [23] . A natural spatial distribution shift occurs between data collected from urban and rural regions, both of which are represented in the training and test sets. However, the test data is sourced from different countries than the training data. The primary evaluation metric is the worst Pearson correlation between predictions in urban and rural environments.We use the synthetic split of VesselPower where the target power labels are generated through the simulation of a physics model based on read data input features. The dataset include a single environment training set, a dev-in set for in-distribution validation, and a dev-out set for evaluation. Distribution shifts arise due to variation in time and wind speeds. We report RMSE of prediction in megawatts (MW), rather than the kilowatts (KW) as used in the dataset's original paper [33] .Model Selction. We follow DomainBed's protocol [15] for model selection. Specifically, we randomly sample 20 sets of hyperparameters for each method, containing both the training hyperparameters of base models in Table 4 and extra hyperparameters from the robust learning algorithm in Table 5 . We select the best model across hyperparameters based on three model selection criteria. In-distribution (ID) validation selects the model with the best metric on the average of an in-distribution validation dataset, which is sampled from the same distribution as the training data. Worst-environment (Worst) validation selects the best model by the worst performance across all environments in the in-distribution validation dataset. Worst validation is applicable only to the multi-environment setting. Oracle validation selects the best model by an out-of-distribution validation dataset sampled from the target distribution of test data. Oracle validation leaks the test distribution, so it is not recommended by DomainBed. However, most robust learning methods relies on out-of-distribution validation, so Domainbed suggests limited numbers of access to target data when using Oracle validation. Though MC-Pseudolabel already performs well under ID and Worst validation, we still report its performance under Oracle validation to compare the limit of robust learning methods regardless of model selection.Following DomainBed, the entire model selection procedure is repeated with different seeds for three times to measure standard errors of performances. Thus, we have totally 60 runs per method per dataset. Specifically for PovertyMap, we follow WILDS' setup [23] and perform 5-fold cross validation instead of three repeated experiments. For each fold of the dataset, we conduct the model selection procedure four times across three seeds, summing up to a total of 12 experiments. And we report the average and standard error of performances across 5 folds. Thus, the standard error measures both the difficulty disparity across folds and the model's instability.Grouping Functions. The grouping function class of MC-Pseudolabel is implemented according to Equation 11 and Equation 12for the multi-environment and single-environment setting respectively. For the multi-environment setting, the environment classifier p(e|x, y) is implemented as MLP with a single hidden layer of size 100 for tabular datasets including Simulation and ACSIncome. For PovertyMap, the environment classifier is implemented by Resnet18-MS with the same architecture as the predictor, except that the label y is fed into the last fully-connected layer. For the singleenvironment setting of VesselPower, the identification model f id is implemented as a Ridge regression model. Our experiments are based on the architecture of PyTorch [35] . Each experiment with a single set of hyperparameters is run on one NVIDIA GeForce RTX 3090 with 24GB of memory, taking at most 15 minutes. 1. (Closure under Covariate Shift) For a set of probability measures P(X) containing the source measure P S (X),EQUATION2. ((γ, ρ)-Weak Learning Condition) For any P ∈ P(X)P S (Y |X) ≡ {P ′ (X)P S (Y | X) : P ′ ∈ P} with the source conditional measure P S (Y |X) and every measurable set G ⊂ X satisfying P (X ∈ G) > ρ, ifEQUATIONthen there exists h ∈ H satisfyingEQUATION)Lemma F.2 (Globus-Harris et al. [13] ). Fix any distribution P ∈ P(X, Y ), any model f : X → [0, 1], and any class of real valued functions H that is closed under affine transformation. Let:H 1 = {h ∈ H : max x∈X h(x) 2 ≤ 1}be the set of functions in H upper-bounded by 1 on X . Let m = |Range(f )|, γ > 0, and α ≤ γ 3 16m . Then if H satisfies the (γ, γ m )-weak learning condition and f is α-approximately ℓ 2 multicalibrated with respect to H 1 and P , then f has squared errorE P [(f (x) -y) 2 ] ≤ E P [(f * (x) -y) 2 ] + 3γ, where f * (x) = E P [Y |x].Definition F.3. For a probability measure P (X, Y ) and a predictor f : X → [0, 1], let H ⊂ R X ×Y be a function class. We say that f is α-approximately ℓ 1 multicalibrated w.r.t. H and P if for all h ∈ H:EQUATIONEQUATION)Lemma F.4 ( Roth [37] ). Suppose P S , P T ∈ P have the same conditional label distribution, and suppose f is α-approximately ℓ 1 multicalibrated with respect to P S and H. If H satisfies Equation 21, then f is also α-approximately ℓ 1 multicalibrated with respect to P T and H: Lemma F.5. For a predictor f : X → [0, 1] and a grouping functionEQUATION)Remark F.6. The lemma is extended from Roth [37] 's result for B = 1.Proof. First we prove K 2 (f, h, P ) ≤ √ BK 1 (f, h, P ). For any v ∈ [0, 1],EQUATION)Equation 30 follows from the Cauchy-Schwarz inequality.EQUATION)Then we prove K 1 (f, h, P ) ≤ BK 2 (f, h, P ). Still from the Cauchy-Schwarz inequality:EQUATION)Theorem F.7 (Restatement of Theorem 2.4). For a source measure P S (X, Y ) and a set of probability measures P(X) containing P S (X), given a predictor f : X → [0 EQUATION)where f * (x) = E P T [Y |x] is the optimal regression function in each target distribution.Proof. For |h(x)| ≤ 1 and f ∈ [0, 1], according to Lemma F.5,EQUATIONEQUATION)With Lemma F.4, for any h ∈ H 1 and P T ∈ P, we haveEQUATION)Again, it follows from Lemma F.5 that:EQUATIONWith Lemma F.2, we haveEQUATIONwhich completes the proof.Theorem F.8 (Restatement of Theorem 3.1). For a set of absolutely continuous probability measures P(X, Y ) containing the source measure P S (X, Y ), consider a predictor f : X → [0, 1]. Assume the grouping function class H satisfies the following condition:EQUATION)If f is α-approximately ℓ 2 multicalibrated w.r.t. H and P S , then for any measure P ∈ P(X, Y ),EQUATIONProof. For any h(x, y) = p(x, y)/p S (x, y) whereEQUATION)Thus, we have K 1 (f, 1, P ) ≤ √ α for any P ∈ P. We will prove an equivalent form of ℓ 1 calibration error:K 1 (f, 1, P ) = sup η:[0,1]→[-1,1] η(v)E P Y -v f (X) = v dP f (X) (v) (53) = sup η:[0,1]→[-1,1] E P [η(f (X))(Y -f (X))]. (54) K 1 (f, 1, P ) ≤ sup η:[0,1]→[-1,1] E P [η(f (X))(Y -f (X))] can be proved by taking η(v) = 2I[v > 0] -1.On the other hand,EQUATION)Actually the right hand side of Equation 53 resembles smooth calibration [3] , which restricts η to Lipschitz functions. Based on smooth calibration, Blasiok et al. [4] shows that approximately calibrated predictors cannot be improved much by post-processing. In the above we present a similar proof for ℓ 1 calibration error.For any g :EQUATIONTheorem F.9 (Restatement of Theorem 3.4). Assume samples are drawn from an environment e ∈ E with a prior P S (e) such that e∈E P S (e) = 1 and P S (e) > 0. The overall population satisfies P S (X, Y ) = e∈E P e (X, Y )P S (e) where P e (X, Y ) is the environment-specific absolutely continuous measure. For a representation Φ ∈ σ(X) over features, define a function class H as:EQUATION1. If there exists a bijection g * : supp(Φ) → [0, 1] such that g * • Φ is α-approximately ℓ 2 multicalibrated w.r.t. H and P S , then for any e ∈ E ,EQUATION)2. For Φ ∈ σ(X), if there exists g * : supp(Φ) → [0, 1] such that Equation 63 is satisfied for any e ∈ E , then g * • Φ is 2/Dα Proof. We first prove statement 1.For any g : 1] . Since g * • Φ is α-approximately ℓ 2 multicalibrated, it follows from Theorem F.8 that for any e ∈ E ,supp(Φ) → [0, 1], since g * is a bijection, g • Φ = (g • g * -1 )(g * • Φ) where g • g * -1 ∈ [0, 1] [0,EQUATION)Then we give a proof of statement 2, which is inspired by Blasiok et al. [4] .For simplicity let f * := g * • Φ. For any e ∈ E and any η :EQUATION)Rearranging the inequality above gives:EQUATION) Since κ • g * ∈ supp(Φ) → [0, 1], it follows from Equation 63 that:EQUATION)Combining Equation 71and Equation 72gives 53, it follows that:E Pe [(Y -f * (X))η(f * (X))] ≤ √ 2α 1/4 for any η : [0, 1] → [-1, 1]. From EquationEQUATIONFrom Equation 52, K 1 (f * , h, P S ) ≤ √ 2α 1/4 for any h ∈ H. Further, for any h ∈ H,= p e (x, y)EQUATION≤ p e (x, y) p e (x, y)P S (e) (76)EQUATIONBy Lemma F.5, it follows thatK 2 (f * , h, P S ) ≤ 1/D • K 1 (f * , h, P S ) ≤ 2/Dα 1/4 .In this subsection, we focus on Euclidean space where X ⊂ R d is compact and measurable for some d ∈ Z + and Y = [0, 1]. Grouping functions are assumed to be continuous, i.e., h ∈ C(X × Y). We consider absolutely continuous probability measures with continuous density functions. Proposition F.10 (Restatement of Proposition 4.1). Consider an absolutely continuous probability measure P S (X, Y ) and a predictor f : X → [0, 1], define the maximal grouping function class that f is multicalibrated with respect to:EQUATIONThen H f is a linear space.Particularly forf Φ (x) = E[Y |Φ(x)] where Φ : R d → R dΦ , d Φ ∈ Z + is a measurable function,we abbreviate H fΦ with H Φ . Then any finite subset G ⊂ H Φ implies span{1, G} ⊂ H Φ , where 1 denotes a constant function.Proof. For any γ 1 , γ 2 ∈ R and anyEQUATIONEQUATIONEQUATIONEQUATIONAccording to Lemma F.5, K 1 (f, h, P S ) = 0 is equivalent as K 2 (f, h, P S ) = 0 for bounded h. Thus,EQUATIONEQUATION) Thus, K 2 (f Φ , 1, P S ) = 0 which implies 1 ∈ H Φ . Since H Φ is a linear space, G ⊂ H Φ implies span{1, G} ⊂ H Φ .Lemma F.11. For any absolutely continuous probability measure P (X, Y ) with a continuous density function p, and any grouping function h ∈ C(X × Y), there exists γ ̸ = 0 and ρ ∈ R such that γh + ρ ∈ C(X × Y), and it is density ratio between some absolutely continuous probability measure P ′ and P , i.e., dP ′ (x, y) = (γh(x, y) + ρ)dP (x, y), where P ′ also has a continuous density function.Proof. Since h is bounded, there exists ρ ′ ∈ R such that h(x, y) + ρ ′ > 0 for any x ∈ X , y ∈ Y. Define:EQUATIONEQUATION)γh + ρ is still continuous.We have γh(x, y) + ρ = γ(h(x, y) + ρ ′ ) > 0 for any x ∈ X , y ∈ Y.EQUATIONThus, P ′ (X, Y ) is an absolutely continuous probability measure.Its density function p ′ = (γh + ρ)p is continuous.Theorem F.12. Consider an absolutely continuous probability measure P S (X, Y ) and a predictorEQUATION= span p(x, y) p S (x, y)EQUATIONProof. First we prove:EQUATIONFor each v ∈ [0, 1] and any h ∈ C(X × Y),EQUATIONEQUATIONEQUATIONEquation 98 follows from Equation 87.EQUATIONEQUATIONEQUATIONNext, we proveEQUATIONThis is equivalent to saying that for any absolutely continuous probability measure P satisfyingEQUATIONEQUATIONEQUATIONEQUATION= 0 almost surely.Next, we proveEQUATIONBy Lemma F.11, any grouping function h ∈ H Φ could be rewritten as h(x, y) = γp(x, y)/p S (x, y) + ρ for some continuous density functions p and γ ̸ = 0. Thus, we just need to prove the statement thatEQUATIONEquation 118 follows from Equation 110. So we haveE P [Y |f Φ = v] = E P S [Y |f Φ = v] if Cov P S [γp(X, Y )/p S (X, Y ) + ρ, Y |f Φ = v] = 0.Corollary F.13 (Restatement of Theorem 4.2). Consider an absolutely continuous probability measure P S (X, Y ) and a calibrated predictor f .EQUATION= span p(x, y) p S (x, y)EQUATIONRemark F.14.E P [Y |f ] = E P S [Y |f ] almost surely implies E P [Y |f ] = f , which is equivalent as R P (f ) = inf g:[0,1]→[0,1] R P (g • f ), since we adopt square error.Proof. Since f is calibrated, we haveEQUATIONApply Theorem F.12 and the proof is complete.Theorem F.15 (Restatement of Theorem 4.3 (first part)). Consider an absolutely continuous probability measure P S (X, Y ) and a predictor fEQUATIONEQUATIONRemark F.16. H 1,Φ contains functions defined on supp(Φ) × Y which can be rewritten as functions on X × Y by variable substitution. Thus, H 1,Φ is still a set of grouping functions. For H 2,Φ , C h is a constant depending on h.Proof. First we prove H Φ ⊂ H 1,Φ + H 2,Φ .For any h 1 ∈ H 1,Φ and h 2 ∈ H 2,Φ ,EQUATIONEQUATIONEQUATION)Next we proveH Φ ⊃ H 1,Φ + H 2,Φ .For any h ∈ H Φ , letEQUATIONThen we haveEQUATIONEQUATIONThus, h 1 (Φ(x), y) ∈ H 1,Φ .EQUATIONThus, h 2 (x, y) ∈ H 2,Φ . Following a similar proof of Theorem F.12, we have H 1,Φ = span p(Φ,EQUATIONEQUATIONLemma F.17 (Theorem 3.2 from Globus-Harris et al. [13] ). If f is calibrated and there exists anEQUATIONthen:EQUATION)Proposition F.18 (Restatement of Theorem 4.3 (second part)). If a predictor f is multicalibrated with H 1,Φ , then R P S (f ) ≤ R P S (f Φ ).Proof. We prove by contradiction. If R P S (f ) > R P S (f Φ ), thenEQUATION) Let α v = E (f (X) -Y ) 2 -(h(X) -Y ) 2 |f (X) = v .Since f is multicalibrated with H 1,Φ , f is calibrated. It follows from Lemma F.17:EQUATION)Then,EQUATIONEQUATIONFrom Lemma F.5, we have K 2 (f, f Φ , P S ) ≥ K 1 (f, f Φ , P S ) 2 > 0. Since f Φ ∈ HEQUATIONThis implies that f is an invariant predictor across E elicited by f .Proposition F.20 (Restatement of Proposition 4.5). Consider X ∈ R d which could be sliced asX = (Φ, Ψ) T and Φ = (Λ, Ω) T . Define H ′ 1,Φ := {h(Φ(x)) ∈ C(X × Y)}, with H ′ 1,Φ ⊂ H 1,Φ . H ′1,X and H ′ 1,Λ are similarly defined. We have: 1. H ′ 1,X ⊃ H ′ 1,Φ ⊃ H ′ 1,Λ ⊃ H ′ 1,∅ = {C}. 2. {C} = H 2,X ⊂ H 2,Φ ⊂ H 2,Λ ⊂ HEQUATIONThus, H ′ 1,Φ ⊂ H 1,Φ . Next, we prove the two arguments in the proposition. 1. For any h(Λ) ∈ H ′ 1,Λ , since Φ = (Λ, Ω), h(Λ) is also a function of Φ. Thus, we have EQUATIONEQUATIONX → [0, 1]. Suppose for some v ∈ Range(f ) there is an h ∈ H such that: E[h(x, y)(y -v)|f (x) = v] > α Let h ′ = v + ηh(x, y) for η = α E[h(x,y) 2 |f (x)=v] . Then: E[(f (x) -y) 2 -(h ′ (x, y) -y) 2 |f (x) = v] > α 2 E[h(x, y) 2 |f (x) = v] .Proof. Following [13] , we haveE[(f (x) -y) 2 -(h ′ (x, y) -y) 2 |f (x) = v] = E[(v -y) 2 -(v + ηh(x, y) -y) 2 |f (x) = v] = E[v 2 -2vy + y 2 -(v + ηh(x, y)) 2 + 2y(v + ηh(x, y)) -y 2 |f (x) = v] = E[2ηyh(x, y) -2ηvh(x, y) -η 2 h(x, y) 2 |f (x) = v] = E[2ηh(x, y)(y -v) -η 2 h(x, y) 2 |f (x) = v] > 2ηα -η 2 E[h(x, y) 2 |f (x) = v] = α 2 E[h(x, y) 2 |f (x) = v] .Theorem F.25 (Restatement of Theorem 5.1). In Algorithm 1, for α, B > 0, if the following is satisfied:EQUATIONthe output f ′ t-1 (x) is α-approximately ℓ 2 multicalibrated w.r.t. H B = {h ∈ H : sup h(x, y) 2 ≤ B}.Proof. We prove by contradiction. Assume that f t-1 is not α-approximately multicalibrated with respect to H B . Then there exists h ∈ H B such that:v∈[1/m] P (f t-1 (x) = v) E h(x, y)(y -v) f t-1 (x) = v 2 > α. For each v ∈ [1/m] define α v := P (f t-1 (x) = v) E h(x, y)(y -v) f t-1 (x) = v 2 .Then we have v∈[1/m] α v > α.According to Lemma F.24, for each v ∈ [1/m], there exists h v ∈ H such that:EQUATIONThen,E (f t-1 (x) -y) 2 -( ft (x) -y) 2 = v∈[1/m] P (f t-1 (x) = v)E (f t-1 (x) -y) 2 -( ft (x) -y) 2 |f t-1 (x) = v = v∈[1/m] P (f t-1 (x) = v)E (f t-1 (x) -y) 2 -(h t v (x, y) -y) 2 |f t-1 (x) = v > α B ,which contradicts the condition in Equation 177.The following proposition is a direct corollary from Globus-Harris et al. [13] 's Theorem 4. Theorem F.27 (Restatement of Theorem 5.2). Consider X ∈ R d with X = (Φ, Ψ) T . Assume that (Φ, Ψ, Y ) follows a multivariate normal distribution N d+1 (µ, Σ) where the random variables are in general position such that Σ is positive definite. We partition Σ into blocks:EQUATION)For any distribution D supported on X × Y, take the the predictor class F = R X and the grouping function class H as a subset of H 2,Φ , which is defined in Equation 124:EQUATIONFor an initial predictor f (0) (x) = E[Y |x], run MC-Pseudolabel(D, H, F) without rounding, then there exists some constant C x depending on x and some constant M (Σ) depending on Σ, such thatEQUATIONEQUATIONEQUATIONWe have 0 ≤ M (Σ) < 1.Remark F.28. f (t) and E[Y |Φ(x)] are both linear. Thus, the convergence of the functions f (t) is equivalent as the convergence of their coefficients. The theorem essentially states that the coefficients of f(t) converges to those of E[Y |Φ(x)] at a rate of O(M (Σ) t ). Remark F.29. E[Y |Φ(x)] is multicalibrated with respect to H. Furthermore, any calibrated predictor on Φ, denoted by g(Φ), is multicalibrated with respect to H. This is because:EQUATIONHowever, E[Y |Φ] is the most accurate predictor among all multicalibrated g(Φ).Remark F.30. The convergence rate M (Σ) does not depend on the dimension d of covariates. WhenEQUATIONIt follows that M (Σ) = 0 and the algorithm will converge in one step.On the other hand, when Y and Ψ are linearly dependent given Φ such that Σ is singular, which violates positive definiteness, following from the proof below:EQUATIONIt follows that M (Σ) = 1 and the algorithm can't converge.So the convergence rate depends on the singularity of the problem. Since the algorithm converges to a predictor that does not depend on Ψ, stronger the "spurious" correlation between Y and Ψ given Φ in the distribution D, the algorithm takes longer to converge.Proof. Without loss of generality, assume µ = 0.EQUATIONEQUATIONAccording to Theorem F.15,EQUATIONThis implies:EQUATIONRearranging to:EQUATION) Let f (t) = (α (t) Φ ) T Φ + (α (t) Ψ ) T Ψ and f (t) = (α (t) Φ ) T Φ + (α (t) Ψ ) T Ψ + (α (t) y ) T Y . We claim that α(t) Ψ = 0. Otherwise, consider f (t) ′ = (α (t) Φ ) T Φ + ( α(t) Ψ ) T E[Ψ|Φ, Y ] + ( α(t) y ) T Y . E[ f (t) -f (t) ′ |Ψ, Y ] = E[(α (t) Ψ ) T Ψ|Φ, Y ] -(α (t) Ψ ) T E[Ψ|Φ, Y ] = 0. (196) Thus, f (t) -f (t) ′ ∈ H.On the other hand,EQUATIONThe inequality follows from the fact that E[ΨΨ T |Φ, Y ] -E[Ψ|Φ, Y ]E[Ψ T |Φ, Y ]is the covariance matrix of Ψ|Φ, Y , which is positive definite because Σ is positive definite. The inequality contradicts with the definition of f (t) . Thus,f (t) = (α (t) Φ ) T Φ + ( α(t) y ) T Y .Define a matrix C ∈ R (d+1)×dt whose columns support the solution space of Equation 195. Then H := C T (S, T, Y ) T ∈ R dt is a random vector. According to the definition of f (t) ,EQUATIONIn the above equation,EQUATIONΨ . Substituting into Equation 194:EQUATIONSubstituting into Equation 202:EQUATIONThen we have f(t+1) = E[ f (t+1) |Φ, Ψ] = (α (t+1) Φ ) T Φ + ( α(t+1) y ) T E[Y |Φ, Ψ].This is equivalent as: EQUATIONCombining the two equations above, we have:EQUATIONThus,EQUATIONIn the above equation, K (t) = 0≤u<t k (u) .Define f (t) = (α EQUATIONIn the following we show α(t) Ψ → 0. Since α(t) Ψ = (α ⋆ Ψ β ⋆ y ) t α ⋆ Ψ , and α ⋆ Ψ β ⋆ y ∈ R dΨ×dΨ has exactly one nonzero eigenvalue β ⋆ y α ⋆ Ψ , we just have to show |β ⋆ y α ⋆ Ψ | < 1.EQUATIONSince, (Σ yy -Σ yΦ Σ -1 ΦΦ Σ Φy ) and (Σ ΨΨ -Σ ΨΦ Σ -1 ΦΦ Σ ΦΨ ) are both Schur complements of Σ's principal submatrix, they are still positive definite. Thus,EQUATIONEQUATION)It follows that β ⋆ y α ⋆ Ψ ≥ 0. So we just have to show β ⋆ y α ⋆ Ψ < 1. Since det(Σ) > 0, by applying row addition on Σ, we have:EQUATION)It follows that:EQUATIONEQUATIONRearranging to:EQUATIONEQUATIONThus, 0 ≤ β ⋆ y α ⋆ Ψ < 1.In the following, we show α(t) Φ → Σ -1 ΦΦ Σ Φy . By Equation 212 and |β ⋆ y α ⋆ Ψ | < 1, we have:EQUATIONEquation 225 follows from the fact that Subsequently,EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONf (t) = E[ f (t) |Φ, Ψ] → E[Y |Φ].Both our theory and algorithm focuses on the bounded regression setting. The definition of extended multicalibration does not depend on the risk function. However, the analysis of the maximal grouping function class as a linear space assumes a continuous probability distribution of observations, implying a continuous target domain. The convergence of MC-Pseudolabel is also established in a regression setting. All experiments are performed on regression tasks. As most algorithms for out-of-distribution generalization are set up with classification problems, we fill the gap for regression and leave an extension to general risk functions for future work.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: There are no negative social impacts for research conducted in this paper. Guidelines:• The answer NA means that there is no societal impact of the work performed.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 


Decision: The paper provides new insights on multicalibration in the context of concept shifts. The reviewers unanimously acknowledge the notable technical contribution of the paper, however, the presentation can be improved for the final version. I recommend accept.

Reviews: review: paper_summary: paper_summary: This work studies an extension of multicalibration to include grouping functions of both covariates and labels. They show that just as multicalibration with respect to covariate density functions guarantees robustness to covariate shift, that extended multicalibration can imply robustness to concept shift. They further show a connection between extended multicalibrated predictors and invariant predictors. They study necessary and sufficient conditions on (covariate-dependent) grouping functions to achieve approximately Bayes optimal predictions on all target distributions, and introduce a boosting style algorithm that converges to an extended multicalibrated predictor under certain distributional assumptions. The usefulness of this new algorithm is empirically validated on PovertyMap, ACSIncome, and VesselPower.
strength_weakness: strength_weakness: There are some questions regarding the feasibility of the MC-PseudoLabel that I did not see addressed in the work that would be helpful in understanding its usefulness. First, there's a known correspondence between multicalibration with respect to a hypothesis class H and weak agnostic learning for H, which unfortunately limits the grouping function classes for which we can efficiently obtain multicalibrated predictors. Does a similar correspondence also hold for extended multicalibration, or is a stronger weak learning condition required (alongside distributional assumptions)? I'm also curious about the necessary assumptions on the data for convergence of MC-PseudoLabel. On what kinds of distributions will the algorithm fail to convergence for H_{2, \Phi}?

Typos/suggested edits:

Abstract

“in existence” -> “in the presence”

Section 1

“flexibly designed to incorporates” -> “incorporate”

“Simultaneoulsy producing invariant predictors” -> “produce”

“porverty estimation”

Section 2

“which is learned in the source distribution” -> “which are learned”

2.2 

“simultaneous approaches” -> “simultaneously approach”

“We show that multicalibration notion w.r.t.” -> “We show that our/this multicalibration notion”


Section 5

M(Sigma) is defined in the Appendix, but not in Theorem 5.2 where is first appears
questions: questions: Repeated from weaknesses:

1. Is there a correspondence between weak agnostic learning and extended multicalibration?

2. On what kinds of distributions will the algorithm fail to convergence for H_{2, \Phi}?
limitations: limitations: The authors largely addressed the limitations (with the exception of possible computational limitations mentioned in Questions).
paper_summary: paper_summary: The paper explores multicalibration in the context of concept shifts, theoretically demonstrating the equivalence of multicalibration and invariance while providing a structural analysis of multicalibration. It introduces a novel algorithm that simplifies model selection and improves performance on real-world datasets.
strength_weakness: strength_weakness: - The empirical work is limited and not that convincing.
- Although technical solid, some assumptions might be too strong to achieve in reality.
- The complexity and detailed mathematical theories/assumptions would make the paper challenging for a broader audience to easily digest.
questions: questions: - It seems that the setting you called "out-of-distribution generalization beyond covariate shift" is basically concept drift. I wonder why you choose to call that.
- In theorem 4.2 and 4.3, you consider an "absolutely continuous probability measure" without stating what measure it is with respect to.
limitations: limitations: See weakness.
paper_summary: paper_summary: The authors explore an extension of multicalibration which includes joint grouping functions; groups that depend on both x and y. They show that multicalibration confers robustness to distribution shift problems. The authors then develop an optimization framework that post-processes a model to multicalibrate it. Finally, empirical demonstration of the robustness properties are given through experimentation.
strength_weakness: strength_weakness: A few questions are given below. I have nothing else to add.
questions: questions: Line 9 of Algorithm 1 is unclear to me. I had thought A finds low risk models f(x) on an empirical dataset {(x,y)}. What does it mean for a grouping function to have low risk on a dataset? What is risk here?

With regard to Algorithm 1, the authors state "The prediction of grouping functions rectify the uncalibrated model and serves as pseudolabels for model updates." Can this statement be expanded, and details be provided?
limitations: limitations: The authors have adequately addressed limitations.
paper_summary: paper_summary: The authors propose utilizing multicalibration algorithms to achieve distributional robustness w.r.t. both concept and covariate shift. They do this by allowing subgroups to not only be a function of the features X (as is standard in multigroup fairness), but also a function of the label Y. There are numerous theoretical and experimental results. The first result (thm 3.2) states that under closure properties of the hypothesis class $\mathcal{H}$ w.r.t. possible distribution shifts, and a weak-learnability assumption on $\mathcal{H}$, multicalibration w.r.t. a particular set of subgroups $\mathcal{H}_1$ guarantees that the resulting predictor will achieve good error (relative to the bayes optimal predictor) under any possible covariate shifted distribution $P_T$. 

In theorem 3.1., the authors show that an L2-multicalibrated predictor f cannot be significantly improved by any post-processing function in order to achieve robustness to concept shift. The authors then show an equivalence between approximate multicalibration and a notion called approximate invariance from prior work (theorem 3.4). Invariance is a property of a representation $\Phi$ which says that there exists a predictor $g^*$ based only on $\Phi$ which achieves good error across all environments. The equivalence between the two notions is up to considering subgroup functions defined by the density ratios between the target environments and the source distribution, similar to the result in theorem 3.1.

Next, the authors demonstrate that (joint) multicalibration cannot always be achieved for all jointly defined subgroups. Thus, they are interested in the maximal feasible set of jointly defined groups for which we can achieve multicalibration towards. These groups turn out to lie in a linear space (Prop. 4.1.) and have a well-defined spanning set which looks similar to an invariant predictor (Thm. 4.2). Finally, they show that the subgroup function space can be decomposed into an invariant and multicalibrated part (Thm. 4.3).

Lastly, experiments are conducted which demonstrate the usefulness of multicalibration in achieving distributional robustness for simple neural networks on the PovertyMap and ACSIncome datasets.
strength_weakness: strength_weakness: I have not checked the proofs thoroughly.

I believe the paper has much room for improvement in terms of clarity, organization, and discussion. In isolation, these may be minor, but taken together, the writing poses a major limitation to the greater community understanding the work (as well as its position in the literature). Throughout, I have asked explicit questions which I would like the authors to address, and labeled them as Q1, Q2, etc.

W1: Firstly, I believe the introduction requires substantial re-writing. Currently, it makes sense to someone who is familiar with both the multicalibration and distributional robustness literature (most people are not). For example, the first sentence of the paper mentions both out-of-distribution generalization and multicalibration, without motivating either of these communities. (Q1) Why might we want to connect these notions? What do we stand to gain from applying multicalibration, and what does multicalibration address that the (substantial) work on OOD generalization has not considered? Perhaps these questions are answered implicitly elsewhere in the work (and I missed them), but I believe that the authors should do a better job of framing the problem they are trying to solve in the context of the literature.

W2: I believe that an explicit discussion of related work in the introduction can be very helpful. I understand that the authors defer all related work discussion to the appendix, but as someone familiar with the multicalibration literature but not distributional robustness, the main paper was hard to follow.

W3: Little interpretation of results. As this is predominantly a theoretical paper, I understand that it can be challenging to concisely state and provide intuitive explanations of each result. However, I believe that the paper can be substantially improved with some additional intuition. Even understanding the theorem statements in section 4 took substantial effort from my part.

As another example, in Section 2.2: line 101-106: it seems implied that these assumptions are from previous work, but it would be good to state explicitly. (Assuming that is the case). Further, any further discussion on the assumptions would be appreciated. I believe it could be useful to state the assumptions in 2.2 as two separate assumptions, in order to make references to them in the following theorem statement more clear.
There is no discussion after the statement of the theorem 2.3. What is the class H_1, and why does it help us achieve robustness to distribution shift? How does it relate to the two assumptions stated previously?
Also, how does Kim et al. [22] relate to theorem 2.3? These should all be explicitly discussed.

As a positive example, the authors do provide some discussion in lines 175-178. Expanding and including similar discussion throughout the paper can help make the ideas more parseable to the unfamiliar reader.

W4: Section 4 is extremely difficult to parse. This also connects to W3, as there is little interpretation of the results. What are the implications of the grouping functions being a linear space? Why is the decomposition in theorem 4.3 a Minkowski sum?  Also, how do these results connect to the algorithm developed in section 5?

Minor comments
35: The authors use the term “joint” grouping functions, but have only implicitly described them in the previous paragraph. It would be clearer to concretely mention that you call these subgroups which depend on both X and Y joint grouping functions.

41-43: The definition of invariance / IRM is not very clear here. I think it is ok to defer a formal description to later, and just mention that you provide an equivalence between the invariant learning framework and multicalibration. 

In line 120, there is some motivation for covariate shift and concept shift. I believe this discussion should go far earlier, perhaps in the introduction.
questions: questions: Q1. How does section 5 relate to the rest of the work? Why do we need another multicalibration algorithm? What problem is MC-PseudoLabel solving which existing OOD / domain adaptation algorithms do not already solve? I believe that these are critical and should be stated explicitly.

Q2. Line 83: “We say that f is $\alpha$-approximately calibrated if $h \equiv 1$.” What does this mean? I think you mean if $\mathcal{H}$ includes the constant function $h\equiv 1$, then we can say that f is $\alpha$ approximately calibrated. 

Q3. Line 18: “Multicalibration is a strengthening  of calibration, which requires a predictor f to be correct on average within each level set:” For clarity, perhaps you could introduce calibration first in an isolated way, without mentioning multicalibration? The discussion here could potentially be misconstrued by the reader.

Q4. Line 80: Definition 2.1. Outside of extending to joint subgroups h(X) -> h(X, Y), how does this definition relate to existing multicalibration definitions?

Q5. Theorem  3.1: In the appendix proof of this theorem, the authors state a similarity to Blasiok et al. [5] w.r.t. the theorem statement. How similar is the proof, and what complications arise from needing to consider joint grouping functions? This should be in the main paper, especially since others may notice a resemblance between the results as well.

Q6. Line 180: “Section 3 inspires one to construct richer grouping function classes for stronger generalizability”. Why is this the case, or can you expand on what this means exactly? Is the idea that due to theorem 3.4, we understand that multicalibration w.r.t. Sufficiently rich joint grouping functions gives us a nice property, but we want to efficiently achieve this which may be difficult in practice?
limitations: limitations: The authors have stated the main limitation of the work, which is that most OOD generalization papers deal with classification, not regression. However, I see this as a feature, since regression can also be a useful consideration in practice.

Comments: comment: Thank you for the rebuttal. It addresses most of my concerns. I'll maintain my score for now.
Thanks to the authors for the helpful response to my questions! I will keep my score.
I acknowledge and thank the authors for the response. In hindsight, and taking into account other reviewers insights, I believe my given score (3: reject) to be overly pessimistic. In my review, I agree that the technical contribution is quite strong. I also believe that the experiments are impressive and very comprehensive for a mainly theoretical paper. I have therefore increased my score to a 5.

I still believe that the paper has much room for improvement in terms of presentation. In particular, I spent quite a long time understanding the main theorems / propositions from the paper (mainly section 4). I would appreciate either informal statements of the theorems, or expanded discussion before and after theorem statements. At the moment, this discussion seems to occur in different sections and not immediately preceding or succeeding the statements.