Title: Bridging Multicalibration and Out-of-distribution Generalization Beyond Covariate Shift
Venue: NeurIPS 2024
Accepted: yes
Keywords: Multicalibration

Robustness

Invariant Risk Minimization
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.

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

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 2 , 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.

1 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:

Calibration 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 ,

Multicalibration 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.

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.

2 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 :

P 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:

̸

X 1 , ϵ 1 , ϵ 2 are independent gaussian variables with zero mean, and variances E [ ϵ 2 1 ] = σ 2 1 and E [ ϵ 2 2 ] = σ 2 2 . For a singleton grouping function class containing h ( x, y ) = y -x 2 , the Bayes optimal predictor f ( x ) = E [ Y | x ] = σ 2 2 σ 2 1 + σ 2 2 x 1 + σ 2 1 σ 2 1 + σ 2 2 x 2 is not multicalibrated because E [ h ( X,Y )( Y -f ( X ))] = σ 2 1 σ 2 2 σ 2 1 + σ 2 2 = 0 . However, g ( x ) = x 1 is multicalibrated since E [ h ( X,Y )( Y -g ( X )) | g ( X )] = -E [ ϵ 2 ϵ 1 | X 1 ] = 0 .

While the Bayes optimal predictor is always multicalibrated for covariate-based grouping functions, it may not be multicalibrated for grouping functions that depend on the outcome. In fact, there may be no predictor that achieves multicalibration in such cases. For example, when we consider a grouping function class that includes both h ≡ 1 and h ( x, y ) = y , a multicalibrated predictor f satisfies f ( X ) = E [ Y | f ( X )] for h ≡ 1 , and E [ Y 2 | f ( X )] = E [ Y f ( X ) | f ( X )] = ( E [ Y | f ( X )]) 2 for h ( x, y ) = y . This implies Var[ Y | f ( X )] = 0 , which is impossible for regression with label noise. Therefore, we study the structure of the maximal grouping function class that allows for a multicalibrated predictor in section 4.

Most importantly, multicalibration with joint grouping functions is useful for capturing more general distribution shifts. By interpreting multicalibration error as calibration error reweighted by grouping functions, it quantifies the maximal calibration error for all subgroups associated with grouping functions in H . If grouping functions are defined on X , only the covariate distribution P ( X ) distinguishes between subgroups. In contrast, the joint distribution P ( X,Y ) differentiates subgroups for joint grouping functions. We will discuss multicalibration with covariate-based grouping functions in the next sub-section and joint grouping functions in section 3.

2.2 Multicalibration Implies Bayes Optimality under Covariate Shift

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 ) = (ˆ 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.3.2). Detailed discussion on other related works is deferred to section A in the appendix.

Assumption 2.3 (Sufficiency of Grouping Function Class (informal, see Assumption F.1)) .

(Closure under Covariate Shift) For a set of probability measures P ( X ) containing the source measure P S ( X ) , h ∈ H implies p/p S · h ∈ H for any density function p of distributions in P .

(( γ, ρ )-Weak Learning Condition) For any P ∈ P ( X ) P S ( Y | X ) ≡ { P ′ ( X ) P S ( Y | X ) : P ′ ∈ P} with the source measure P S ( Y | X ) , and every subset G ⊂ X with P ( X ∈ G ) > ρ , if the Bayes optimal predictor E P [ Y | X ] has lower risk than the constant predictor E P [ Y | X ∈ G ] by a margin γ , there exists a predictor h ∈ H that is also better than the constant predictor with the margin γ .

Theorem 2.4 (Risk Bound under Covariate Shift) . 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 , 1] with finite range m := | Range ( f ) | , consider a grouping function class H closed under affine transformation and satisfying Assumption 2.3 with ρ = γ/m . If f is γ 6 256 m 2 -approximately ℓ 2 multicalibrated w.r.t P S and H 1 := { h ∈ H : max x ∈X h ( x ) 2 ≤ 1 } , then for any target measure P T ∈ P ( X ) P S ( Y | X ) ,

3 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:

If f is α -approximately ℓ 2 multicalibrated w.r.t. H and P S , then for any measure P ∈ P ( X,Y ) ,

The 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 :

Remark 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 implies E e 1 [ Y | Φ] = E e 2 [ 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:

1. If there is a bijection g ⋆ : 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 -approximately ℓ 2 multicalibrated w.r.t. H and P S , where D = min e ∈ E P S ( e ) .

Remark 3.5. (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.

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.

4 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.

4.1 Maximal Grouping Function Space

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:

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:

A predictor's maximal grouping function class is spanned by density ratios of target distributions where the predictor is invariant. Correspondingly, Theorem 3.1 gives the minimal grouping function class, comprised of density ratios between target and source distributions, in order to ensure f ( x ) is an invariant predictor. In contrast, Theorem 4.2 states the maximal grouping function class for f ( x ) is exactly the linear space spanned by those density ratios. Next, we further investigate sub-structures of the maximal grouping function class. We focus on the representation learning setting of IRM.

Theorem 4.3 (Decomposition of Grouping Function Space) . Consider an absolutely continuous probability measure P S ( X,Y ) and a measurable function Φ : R d → R d Φ with d Φ ∈ Z + . We define the Bayes optimal predictor over Φ as f Φ ( x ) = E P S [ Y | Φ( x )] . We abbreviate H f Φ with H Φ . Then H Φ can be decomposed as a Minkowski sum of H 1 , Φ + H 2 , Φ .

1. If a predictor f is multicalibrated with H 1 , Φ , then R P S ( f ) ≤ R P S ( f Φ ) .

2. f Φ is an invariant predictor elicited by Φ across a set of environments E where P e (Φ , Y ) = P S (Φ , Y ) for any e ∈ E . If a predictor f is multicalibrated with H 2 , Φ , then f is also an invariant predictor across E elicited by some representation.

Remark 4.4. H 1 , Φ and H 2 , Φ contain functions defined on x, Φ( x ) , y which can both be rewritten as functions on x, y by variable substitution. Thus, H 1 , Φ , H 2 , Φ are still subspaces of grouping functions. H 1 , Φ is spanned by the density ratio of P (Φ , Y ) where the Bayes optimal predictor over Φ must be invariant on the distribution of P . H 2 , Φ is spanned by general density ratio of P ( X | Φ , Y ) .

Multicalibration w.r.t. H 1 , Φ ensures at least the accuracy of the Bayes optimal predictor on Φ , and multicalibration w.r.t. H 2 , Φ ensures at least the invariance of this predictor. However, we show in the following proposition that sizes of two subspaces are negatively correlated. When Φ is a variable selector, H 1 , Φ expands with more selected covariates while H 2 , Φ shrinks. By choosing a combination of H 1 , Φ and H 2 , Φ , we strike a balance between accuracy and invariance of the multicalibrated predictor.

Proposition 4.5 (Monotonicity) . Consider X ∈ R d which could be sliced as X = (Φ , Ψ) 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:

C is a constant value function.

4.2 Design of Grouping Function Classes

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 :

From 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:

Another design utilizing Distributionally Robust Optimization's assumption [10] is in section B.

5 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 ˆ P D ( X,Y ) , a regression oracle A solves the optimization: A F ( D ) = arg min f ∈F R ˆ P D ( 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 -˜ err 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 -˜ err 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 } .

Algorithm 1 MC-PseudoLabel

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 ) = ˜ f t +1 ( x, · ) where ˜ f t +1 does not depend on y . For joint grouping functions, since ˜ f t +1 ∈ R X×Y , we project it to models' space of R X by learning the model with ˜ f t +1 as pseudolabels. The projection substantially changes the optimization dynamics. LSBoost constantly decreases risks of models, due to R ˆ P D ( f t +1 ) = R ˆ P D ( ˜ f t +1 ) < R ˆ P D ( 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. Assume we are given a grouping function class H Φ with a latent representation Φ . If a predictor is multicalibrated w.r.t H 1 , Φ , H 2 , Φ respectively, then it is also multicalibrated w.r.t. H Φ . Therefore, we separately study the convergence with two grouping function classes. In Proposition F.26, we show the convergence for a subset of H 1 , Φ consisting of covariate-based grouping functions, which is a corollary of Globus-Harris et al.'s result. As a greater challenge, we derive convergence for H 2 , Φ when data follows multivariate normal distributions.

Theorem 5.2 (Covergence for H 2 , Φ (informal, see Theorem 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:

For an initial predictor f 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

6.1 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 11 and Equation 12 for 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.

6.2 Results

Results are shown in Table 1 for multi-environment settings and Table 2 for single-environment settings. MC-Pseudolabel achieves superior performance in all datasets with in-distribution and worst-environment validation which does not violate test data. For oracle validation, MCPseudolabel achieves comparable performances to the best method. For example, CLOvE, which also learns invariance by calibration, achieves best performance under oracle validation in PovertyMap, but it sharply degrades when target validation data is unavailable. It's

1

4

2

12

R

2

=0.50

ERM

1

3

2

14

R

2

=0.67

C-Mixup

1.25

2.00

2

10

R

2

=0.96

MC-Pseudolabel

ID Validation RMSE

Test RMSE

because CLOvE tunes its regularizer's coefficient to tradeoff with ERM risk, whose optimal value depends on the target distribution shift. In contrast, MC-Pseudolabel exhibits an advantage with in-distribution model selection. This is further supported by Figure 6.2, which shows that MC-Pseudolabel's out-of-distribution er- rors 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.

7 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.

Acknowledgments and Disclosure of Funding

Peng Cui is supported in part by National Natural Science Foundation of China (No. 62425206, 62141607).

References

Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 3438-3450, 2021. URL https://proceedings.neurips.cc/paper/2021/hash/ 1c336b8080f82bcc2cd2499b4c57261d-Abstract.html .

Martín Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. CoRR , abs/1907.02893, 2019. URL http://arxiv.org/abs/1907.02893 .

Jaroslaw Blasiok, Parikshit Gopalan, Lunjia Hu, and Preetum Nakkiran. A unifying theory of distance from calibration. In Barna Saha and Rocco A. Servedio, editors, Proceedings of the 55th Annual ACM Symposium on Theory of Computing, STOC 2023, Orlando, FL, USA, June 20-23, 2023 , pages 1727-1740. ACM, 2023. doi: 10.1145/3564246.3585182. URL https://doi.org/10.1145/3564246.3585182 .

Jaroslaw Blasiok, Parikshit Gopalan, Lunjia Hu, and Preetum Nakkiran. When does optimizing a proper loss yield calibration? In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 -16, 2023 , 2023. URL http://papers.nips.cc/paper_files/paper/2023/hash/ e4165c96702bac5f4962b70f3cf2f136-Abstract-Conference.html .

Elliot Creager, Jörn-Henrik Jacobsen, and Richard S. Zemel. Environment inference for invariant learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 2189-2200. PMLR, 2021. URL http: //proceedings.mlr.press/v139/creager21a.html .

Zhun Deng, Cynthia Dwork, and Linjun Zhang. Happymap : A generalized multicalibration method. In Yael Tauman Kalai, editor, 14th Innovations in Theoretical Computer Science Conference, ITCS 2023, January 10-13, 2023, MIT, Cambridge, Massachusetts, USA , volume 251 of LIPIcs , pages 41:1-41:23. Schloss Dagstuhl - Leibniz-Zentrum für Informatik, 2023. doi: 10. 4230/LIPICS.ITCS.2023.41. URL https://doi.org/10.4230/LIPIcs.ITCS.2023.41 .

Frances Ding, Moritz Hardt, John Miller, and Ludwig Schmidt. Retiring adult: New datasets for fair machine learning. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 6478-6490, 2021. URL https://proceedings.neurips.cc/paper/2021/hash/ 32e54441e6382a7fbacbbbaf3c450059-Abstract.html .

John C. Duchi and Hongseok Namkoong. Variance-based regularization with convex objectives. J. Mach. Learn. Res. , 20:68:1-68:55, 2019. URL http://jmlr.org/papers/v20/17-750. html .

John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics , 49(3):1378-1406, 2021.

John C Duchi, Tatsunori Hashimoto, and Hongseok Namkoong. Distributionally robust losses against mixture covariate shifts. Under review , 2(1), 2019.

John C. Duchi, Peter W. Glynn, and Hongseok Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Math. Oper. Res. , 46(3):946-969, 2021. doi: 10.1287/MOOR.2020.1085. URL https://doi.org/10.1287/moor.2020.1085 .

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017 , volume 70 of Proceedings of Machine Learning Research , pages 1126-1135. PMLR, 2017. URL http://proceedings.mlr.press/v70/finn17a.html .

Ira Globus-Harris, Declan Harrison, Michael Kearns, Aaron Roth, and Jessica Sorrell. Multicalibration as boosting for regression. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA , volume 202 of Proceedings of Machine Learning Research , pages 11459-11492. PMLR, 2023. URL https://proceedings.mlr.press/v202/globus-harris23a.html .

Parikshit Gopalan, Adam Tauman Kalai, Omer Reingold, Vatsal Sharan, and Udi Wieder. Omnipredictors. In Mark Braverman, editor, 13th Innovations in Theoretical Computer Science Conference, ITCS 2022, January 31 - February 3, 2022, Berkeley, CA, USA , volume 215 of LIPIcs , pages 79:1-79:21. Schloss Dagstuhl - Leibniz-Zentrum für Informatik, 2022. doi: 10. 4230/LIPICS.ITCS.2022.79. URL https://doi.org/10.4230/LIPIcs.ITCS.2022.79 .

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=lQdXeXDoWtI .

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016 , pages 770-778. IEEE Computer Society, 2016. doi: 10.1109/CVPR.2016.90. URL https://doi.org/10.1109/CVPR.2016.90 .

Úrsula Hébert-Johnson, Michael P. Kim, Omer Reingold, and Guy N. Rothblum. Multicalibration: Calibration for the (computationally-identifiable) masses. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 1944-1953. PMLR, 2018. URL http://proceedings.mlr.press/v80/hebert-johnson18a.html .

Dongsung Huh and Avinash Baidya. The missing invariance principle found the reciprocal twin of invariant risk minimization. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 -December 9, 2022 , 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/ 91b482312a0845ed86e244adbd9935e4-Abstract-Conference.html .

Badr Youbi Idrissi, Martín Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Bernhard Schölkopf, Caroline Uhler, and Kun Zhang, editors, 1st Conference on Causal Learning and Reasoning, CLeaR 2022, Sequoia Conference Center, Eureka, CA, USA, 11-13 April, 2022 , volume 177 of Proceedings of Machine Learning Research , pages 336-351. PMLR, 2022. URL https://proceedings. mlr.press/v177/idrissi22a.html .

Michael P. Kim, Amirata Ghorbani, and James Y. Zou. Multiaccuracy: Black-box postprocessing for fairness in classification. In Vincent Conitzer, Gillian K. Hadfield, and Shannon Vallor, editors, Proceedings of the 2019 AAAI/ACM Conference on AI, Ethics, and Society, AIES 2019, Honolulu, HI, USA, January 27-28, 2019 , pages 247-254. ACM, 2019. doi: 10.1145/3306618.3314287. URL https://doi.org/10.1145/3306618.3314287 .

Michael P Kim, Christoph Kern, Shafi Goldwasser, Frauke Kreuter, and Omer Reingold. Universal adaptability: Target-independent inference that competes with propensity scoring. Proceedings of the National Academy of Sciences , 119(4):e2108097119, 2022.

Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023 . OpenReview.net, 2023. URL https://openreview.net/pdf?id=Zb6c8A-Fghk .

Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee,

Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran S. Haque, Sara M. Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 1824 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 5637-5664. PMLR, 2021. URL http://proceedings.mlr.press/v139/koh21a.html .

Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an out-of-distribution generalization problem? arXiv preprint arXiv:2008.01883 , 2020.

David Krueger, Ethan Caballero, Jörn-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Rémi Le Priol, and Aaron C. Courville. Out-of-distribution generalization via risk extrapolation (rex). In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 5815-5826. PMLR, 2021. URL http://proceedings.mlr.press/v139/krueger21a.html .

Daniel Levy, Yair Carmon, John C. Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020. URL https://proceedings.neurips.cc/ paper/2020/hash/64986d86a17424eeac96b08a6d519059-Abstract.html .

Tian Li, Ahmad Beirami, Maziar Sanjabi, and Virginia Smith. Tilted empirical risk minimization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum? id=K5YasWXZT3O .

Tsung-Yi Lin, Priya Goyal, Ross B. Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. IEEE Trans. Pattern Anal. Mach. Intell. , 42(2):318-327, 2020. doi: 10.1109/TPAMI.2018.2858826. URL https://doi.org/10.1109/TPAMI.2018.2858826 .

Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 6781-6792. PMLR, 2021. URL http://proceedings.mlr.press/v139/liu21f.html .

Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally robust optimization with data geometry. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022 , 2022. URL http://papers.nips.cc/paper_files/paper/2022/ hash/da535999561b932f56efdd559498282e-Abstract-Conference.html .

Jiashuo Liu, Tianyu Wang, Peng Cui, and Hongseok Namkoong. On the need for a language describing distribution shifts: Illustrations on tabular datasets. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2023.

Andrey Malinin, Neil Band, Yarin Gal, Mark J. F. Gales, Alexander Ganshin, German Chesnokov, Alexey Noskov, Andrey Ploskonosov, Liudmila Prokhorenkova, Ivan Provilkov, Vatsal Raina, Vyas Raina, Denis Roginskiy, Mariya Shmatova, Panagiotis Tigas, and Boris Yangel. Shifts: A dataset of real distributional shift across multiple large-scale tasks. In Joaquin Vanschoren and Sai-Kit Yeung, editors, Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks 2021, December 2021, virtual , 2021. URL https://datasets-benchmarks-proceedings.neurips.cc/ paper/2021/hash/ad61ab143223efbc24c7d2583be69251-Abstract-round2.html .

Andrey Malinin, Andreas Athanasopoulos, Muhamed Barakovic, Meritxell Bach Cuadra, Mark J. F. Gales, Cristina Granziera, Mara Graziani, Nikolay Kartashev, Konstantinos Kyriakopoulos, Po-Jui Lu, Nataliia Molchanova, Antonis Nikitakis, Vatsal Raina, Francesco La Rosa, Eli Sivena, Vasileios Tsarsitalidis, Efi Tsompopoulou, and Elena Volf. Shifts 2.0: Extending the dataset of real distributional shifts. CoRR , abs/2206.15407, 2022. doi: 10.48550/ARXIV.2206.15407. URL https://doi.org/10.48550/arXiv.2206.15407 .

John Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang, Yair Carmon, and Ludwig Schmidt. Accuracy on the line: on the strong correlation between out-of-distribution and in-distribution generalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 7721-7735. PMLR, 2021. URL http://proceedings.mlr.press/v139/ miller21b.html .

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Alexandre Ramé, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 18347-18377. PMLR, 2022. URL https://proceedings.mlr.press/v162/rame22a.html .

Aaron Roth. Uncertain: Modern topics in uncertainty estimation, 2022.

Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations , 2019.

Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv:2104.09937 , 2021.

Aman Sinha, Hongseok Namkoong, and John C. Duchi. Certifying some distributional robustness with principled adversarial training. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018. URL https://openreview.net/forum?id=Hk6kPgZA-.

Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 2215-2227, 2021. URL https://proceedings.neurips.cc/ paper/2021/hash/118bd558033a1016fcc82560c65cca5f-Abstract.html .

Haoxiang Wang, Bo Li, and Han Zhao. Understanding gradual domain adaptation: Improved analysis, optimal path and beyond. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 22784-22801. PMLR, 2022. URL https://proceedings.mlr.press/v162/wang22n.html .

Huaxiu Yao, Yiping Wang, Linjun Zhang, James Y. Zou, and Chelsea Finn. Cmixup: Improving generalization in regression. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 -December 9, 2022 , 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/ 1626be0ab7f3d7b3c639fbfd5951bc40-Abstract-Conference.html .

Christopher Yeh, Anthony Perez, Anne Driscoll, George Azzari, Zhongyi Tang, David Lobell, Stefano Ermon, and Marshall Burke. Using publicly available satellite imagery and deep learning to understand economic well-being in africa. Nature Communications , 2020.

Hongyi Zhang, Moustapha Cissé, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018. URL https://openreview.net/forum?id=r1Ddp1-Rb .

Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Re. Correct-n-contrast: a contrastive approach for improving robustness to spurious correlations. In International Conference on Machine Learning , pages 26484-26516. PMLR, 2022.

A Related Work

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. Globus-Harris et al. also provides a boosting algorithm for both regression and multicalibration. The connection between multicalibration and distribution shift is first studied by Kim et al. [21], who proves that ℓ 1 multicalibration error remains under covariate shift, given a sufficiently large real-valued grouping function class. Kim et al. further shows that under covariate shift, a multicalibrated predictor can perform statistical inference of the average outcome of a sample batch. In contrast, we derive a robustness result for individual prediction of outcomes for ℓ 2 multicalibrated predictors. In addition, Wald et al. [41] studies the equivalence of Invariant Risk Minimization and simultaneous calibration on each environment. Our equivalence results for multicalibration can be perceived as a generalization of Wald et al.'s results beyond the multi-environment setting, by deriving a mapping between density ratios and grouping functions. We also extend the equivalence to approximately multicalibrated and approximately invariant predictors. Furthermore, we move beyond Wald et al.'s multi-objective optimization with Lagrangian regularization, by proposing a new post-processing optimization framework consisting of a series of supervised regression. Meanwhile, Blasiok et al. [4] discusses connections between calibration and post-processing, which is an equivalent expression of invariance. There are other extensions of multicalibration, such as Deng et al. [6] who generalize the term Y -f ( X ) in multicalibration's definition to a class of general functions. While our work is the first to generalize the grouping functions h to consider the outcomes.

Out-of-distribution Generalization Beyond Covariate Shift Despite abundant literature from domain generalization that focuses on image classification where covariate shift dominates, research on algorithmic robustness on regression tasks beyond covariate shift is relatively limited. The setting can be categorized according to if the source distribution is partitioned into several environments. For the multi-environment generalization setting, Invariant Risk Minimization and its variants assume that outcomes are generated by a common causal structural equation across all environments, and aims to recover such an invariant (or causal) predictor [1, 2, 18, 24, 25, 36]. Group DRO [38] is a simple but surprisingly strong technique that optimizes for the worst group risk with reweighting of environments. 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.

B Grouping Functions for Distributionally Robust Optimization

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:

By the correspondence between density ratios and grouping functions, the equivalent design of a grouping function class is given by:

We can also extend the grouping functions to consider both covariates and outcomes, such that general subgroups are incorporated into the uncertainty set:

In 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.

C Model Discretization

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.

D Computational Complexity

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.

E Experiments

E.1 An Additional Experiment: Synthetic Dataset

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 MCPseudolabel. 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:

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.

10

2

10

1

10 0

10 1

10 2

0.5

1.0

1.5

2.0

2.5

RMSE

MC-Pseudolabel

MRI

MIP

IRM

IB-IRM

GroupDRO

CLOvE

ERM

REX

IDGM

We also investigate the evolution of pseudo labels ˜ f t in Algorithm 1 to recover the dynamics of MCPseudolabel. 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.

E.2 VesselPower

In figure 4, we provide the correlation between models' in-distribution validation performance and out-of-distribution test performance across all methods.

Step=0

Step=2

Step=4

3

3

3

alpha_v = 1.25

2

2

2

alpha_v = 0.75

1

1

1

Pseudo

0

0

0

-1-

-1

-1

-2

-2

-2

-3

-2

0

2

-3

-2

0

2

-3

-2

0

2

ModelPrediction

ModelPrediction

3

3

3

2

·

alpha_v = 1.25

2

2

alpha_v = 0.75

1

1

1

Pseudo

0

01

0

-1

-1

-1

-2

-2

-2

35

-3

-2

0

2

-3

-2

0

2

-3

-2

0

2

V values

Vvalues

V values

E.3 Training Details

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

1

4

2

12

R

2

=0.50

ERM

1

3

2

12

R

2

=0.60

Tilted ERM

1

4

2

10

R

2

=0.66

JTT

1

3

2

14

R

2

=0.67

C-Mixup

1

3

2

10

R

2

=0.74

2

-DRO

1

3

2

12

R

2

=0.62

CVaR

1.25

2.00

2

10

R

2

=0.96

MC-Pseudolabel

Test RMSE

ID Validation RMSE

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 12 for 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.

E.4 Software and Hardware

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.

F Theory

F.1 Multicalibration and Bayes Optimality under Covariate Shift

Assumption F.1 (Restatement of Assumption 2.3) .

1. (Closure under Covariate Shift) For a set of probability measures P ( X ) containing the source measure P S ( X ) ,

2. (( γ, ρ )-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 ) > ρ , if

then there exists h ∈ H satisfying

LemmaF.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:

be the set of functions in H upper-bounded by 1 on X . Let m = | Range ( f ) | , γ > 0 , and α ≤ γ 3 16 m . 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 error

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 :

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 function h satisfying max x ∈X ,y ∈Y h ( x, y ) 2 ≤ B where B > 0 ,

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 30 follows from the Cauchy-Schwarz inequality.

Then we prove K 1 ( f, h, P ) ≤ √ BK 2 ( f, h, P ) . Still from the Cauchy-Schwarz inequality:

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 , 1] with finite range m := | Range ( f ) | , consider a grouping function class H closed under affine transformation satisfying Assumption 2.3 with ρ = γ/m . If f is γ 6 256 m 2 -approximately ℓ 2 multicalibrated w.r.t P S and H 's bounded subset H 1 := { h ∈ H : max x ∈X h ( x ) 2 ≤ 1 } , then for any target measure P T ∈ P ( X ) P S ( Y | X ) ,

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,

Since K 2 ( f, h, P S ) ≤ γ 6 256 m 2 for any h ∈ H 1 ,

With Lemma F.4, for any h ∈ H 1 and P T ∈ P , we have

Again, it follows from Lemma F.5 that:

which completes the proof.

F.2 Multicalibration and Invariance under Concept Shift

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:

If f is α -approximately ℓ 2 multicalibrated w.r.t. H and P S , then for any measure P ∈ P ( X,Y ) ,

Proof. For any h ( x, y ) = p ( x, y ) /p S ( x, y ) where P ∈ P , since f is α -approximately ℓ 2 multicalibrated, K 2 ( f, h, P S ) ≤ α . It follows from Lemma F.5 that K 1 ( f, h, P S ) ≤ √ α .

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] E P [ η ( f ( X ))( Y -f ( X ))] can be proved by taking η ( v ) = 2 I [ v > 0] -1 . On the other hand,

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 : [0 , 1] → [0 , 1] , there exists η : [0 , 1] → [ -1 , 1] such that g ( v ) = v + η ( v ) for any v ∈ [0 , 1] .

Theorem 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:

1. 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 ,

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α 1 / 4 -approximately ℓ 2 multicalibrated w.r.t. H and P S , where D = min e ∈ E P S ( e ) .

Proof. We first prove statement 1.

For any g : supp (Φ) → [0 , 1] , since g ∗ is a bijection, g ◦ Φ = ( g ◦ g ∗-1 )( g ∗ ◦ Φ) where g ◦ g ∗-1 ∈ [0 , 1] [0 , 1] . Since g ∗ ◦ Φ is α -approximately ℓ 2 multicalibrated, it follows from Theorem F.8 that for any e ∈ E ,

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 η : [0 , 1] → [ -1 , 1] , define β := E P e [( Y -f ∗ ( X )) η ( f ∗ ( X ))] ∈ [ -1 , 1] . Construct κ ( v ) := proj [0 , 1] ( v + βη ( v )) , where proj [0 , 1] ( · ) = max { 0 , min { 1 , ·}} .

Rearranging the inequality above gives:

Since κ ◦ g ∗ ∈ supp (Φ) → [0 , 1] , it follows from Equation 63 that:

Combining Equation 71 and Equation 72 gives E P e [( Y -f ∗ ( X )) η ( f ∗ ( X ))] ≤ √ 2 α 1 / 4 for any η : [0 , 1] → [ -1 , 1] . From Equation 53, it follows that:

From Equation 52, K 1 ( f ∗ , h, P S ) ≤ √ 2 α 1 / 4 for any h ∈ H . Further, for any h ∈ H ,

By Lemma F.5, it follows that K 2 ( f ∗ , h, P S ) ≤ √ 1 /D · K 1 ( f ∗ , h, P S ) ≤ √ 2 /Dα 1 / 4 .

F.3 Structure of Grouping Function Classes

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:

Then H f is a linear space.

Particularly for f Φ ( 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 any h 1 , h 2 ∈ H f , γ 1 h 1 + γ 2 h 2 ∈ C ( X × Y ) .

According 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, K 2 ( f, γ 1 h 1 + γ 2 h 2 , P S ) = 0 which implies γ 1 h 1 + γ 2 h 2 ∈ H f . Now we finishes the proof that H f is a linear space.

For f Φ ( x ) = E [ Y | Φ( x )] ,

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:

γh + ρ is still continuous.

We have γh ( x, y ) + ρ = γ ( h ( x, y ) + ρ ′ ) > 0 for any x ∈ X , y ∈ Y .

Thus, 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 predictor f Φ ( x ) = E P S [ Y | Φ( x )] where Φ : R d → R d Φ , d Φ ∈ Z + is a measurable function. We abbreviate H f Φ with H Φ .

Proof. First we prove:

For each v ∈ [0 , 1] and any h ∈ C ( X × Y ) ,

Equation 98 follows from Equation 87.

For any h ∈ C ( X × Y ) ,

Next, we prove

This is equivalent to saying that for any absolutely continuous probability measure P satisfying E P [ Y | f Φ ] = E P S [ Y | f Φ ] almost surely, Cov P S [ p ( X,Y ) /p S ( X,Y ) , Y | f Φ = v ] = 0 for almost every v ∈ [0 , 1] .

= 0 almost surely.

Next, we prove

̸

By 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 that

Cov P S [ γp ( X,Y ) /p S ( X,Y ) + ρ, Y | f Φ = v ] = 0 implies E P [ Y | f Φ = v ] = E P S [ Y | f Φ = v ] .

Equation 118 follows from Equation 110. So we have E 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 .

Remark 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 have E P S [ Y | f ] = f . Take Φ( x ) = f ( x ) .

Theorem F.15 (Restatement of Theorem 4.3 (first part)) . Consider an absolutely continuous probability measure P S ( X,Y ) and a predictor f Φ ( x ) = E P S [ Y | Φ( x )] where Φ : R d → R d Φ , d Φ ∈ Z + is a measurable function. H Φ can be decomposed as H Φ = H 1 , Φ + H 2 , Φ .

Remark 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 , Φ ,

Next we prove H Φ ⊃ H 1 , Φ + H 2 , Φ .

For any h ∈ H Φ , let

Then we have

Thus, h 1 (Φ( x ) , y ) ∈ H 1 , Φ .

Thus, h 2 ( x, y ) ∈ H 2 , Φ . Following a similar proof of Theorem F.12, we have

Next, we prove H 2 , Φ ⊃ span { p ( x | Φ , y ) /p S ( x | Φ , y ) : p is continuous } .

This is equivalent to saying that p ( x | Φ , y ) /p S ( x | Φ , y ) ∈ H 2 , Φ for any continuous density function p .

Thus, we have p ( x | Φ , y ) /p S ( x | Φ , y ) ∈ H 2 , Φ .

Next, we prove H 2 , Φ ⊂ span { p ( x | Φ , y ) /p S ( x | Φ , y ) : p is continuous } .

̸

By Lemma F.11, any grouping function h ∈ H 2 , Φ could be rewritten as h 2 ( x, y ) = γp ( x, y ) /p S ( x, y ) + ρ for some continuous density function p and γ = 0 . Thus, we just need to prove the statement that E P S [ γp ( X,Y ) /p S ( X,Y ) + ρ | Φ , Y ] ≡ C h 2 implies p ( X,Y ) /p S ( X,Y ) ≡

p ( X | Φ , Y ) /p S ( X | Φ , Y ) .

Thus, we have p (Φ , Y ) /p S (Φ , Y ) ≡ ( C h 2 -ρ ) /γ which is a constant. Since E P S [ p (Φ , Y ) /p S (Φ , Y )] = 1 , we have p (Φ , Y ) /p S (Φ , Y ) ≡ 1 .

Lemma F.17 (Theorem 3.2 from Globus-Harris et al. [13]) . If f is calibrated and there exists an h ( x ) such that

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 Φ ) , then

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:

From Lemma F.5, we have K 2 ( f, f Φ , P S ) ≥ K 1 ( f, f Φ , P S ) 2 > 0 . Since f Φ ∈ H 1 , Φ , it contradicts with the fact that f is multicalibrated with H 1 , Φ .

Proposition F.19 (Restatement of Theorem 4.3 (third part)) . f Φ is an invariant predictor elicited by Φ across a set of environments E where P e (Φ , Y ) = P S (Φ , Y ) for any e ∈ E . If a predictor f is multicalibrated with H 2 , Φ , then f is also an invariant predictor across E elicited by some representation.

then:

Then,

Proof. Since P e (Φ , Y ) = P S (Φ , Y ) , we have E P e [ Y | Φ] = E P S [ Y | Φ] = f Φ for every e ∈ E . Thus,

This implies f Φ is an invariant predictor elicited by Φ across E .

For any e ∈ E , we have:

Since f is multicalibrated with H 2 , Φ , it follows from Theorem F.8:

This 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 as X = (Φ , Ψ) 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:

C is a constant value function.

Remark F.21. The proposition shows that H ′ 1 , as a subspace of H 1 , evolves monotonically and in opposite direction to H 2 . If we perceive the representation Φ as a filter, gaining more information from covariates facilitates multicalibration w.r.t. H 1 (and accuracy) but hampers multicalibration w.r.t. H 2 (and invariance). With H ′ 1 and H 2 combined together, a multicalibrated predictor is searching for an appropriate level of information filter to balance the tradeoff between accuracy and invariance.

Proof. We first prove H ′ 1 , Φ ⊂ H 1 , Φ . According to Equation 122,

Since h (Φ) ⊥ Y | Φ , we have Cov [ h (Φ) , Y | Φ] = 0 .

Thus, 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 h (Λ) ∈ H ′ 1 , Φ . It follows that H ′ 1 , Φ ⊃ H ′ 1 , Λ . Similarly we have H ′ 1 ,X ⊃ H ′ 1 , Φ and H ′ 1 , Λ ⊃ H ′ 1 , ∅ .

For any h (Λ , Ω , Ψ , Y ) ∈ H 2 , Φ such that E [ h (Φ , Ψ , Y ) | Φ , Y ] = C h for any values of Φ , Y , we have

Thus, h (Λ , Ω , Ψ , Y ) ∈ H 2 , Λ . It follows that H 2 , Φ ⊂ H 2 , Λ . Similarly, we have H 2 ,X ⊂ H 2 , Φ and H 2 , Λ ⊂ H 2 , ∅ . Particularly for h ( x, y ) ∈ H 2 ,X , we have h ( X,Y ) = E [ h ( X,Y ) | X,Y ] is a constant for any values of X,Y .

Lemma F.22 (Globus-Harris et al. [13]) . Let H ⊂ X R be such a grouping function class that h ∈ H implies γh + ρ ∈ H for any h ∈ H and γ, ρ ∈ R . If H satisfies the (0 , 0) -weak learning condition in Assumption F.1, a predictor f is multicalibrated w.r.t. H if and only if f ( x ) = E [ Y | x ] almost surely.

Proposition F.23. For a measurable function Φ : R d → R d Φ , d Φ ∈ Z + , a predictor f : supp (Φ) → [0 , 1] is multicalibrated w.r.t. H 1 , Φ if and only if f is multicalibrated w.r.t. H ′ 1 , Φ .

Proof. Since H ′ 1 , Φ ⊂ H 1 , Φ , f 's multicalibration w.r.t. H 1 , Φ implies f 's multicalibration w.r.t. H ′ 1 , Φ .

On the other hand, H ′ 1 , Φ satisfies the (0 , 0) -weak learning condition with the pushforward measure on Φ , because E [ Y | Φ] ∈ H ′ 1 , Φ . It follows from Lemma F.22 that f is multicalibrated w.r.t. H ′ 1 , Φ implies f (Φ) = E [ Y | Φ] almost surely. By the definition of H 1 , Φ , f (Φ) is multicalibrated w.r.t. H 1 , Φ .

F.4 MC-PseudoLabel: An Algorithm for Extended Multicalibration

Lemma F.24. Fix a model f : X → [0 , 1] . Suppose for some v ∈ Range ( f ) there is an h ∈ H such that:

Let h ′ = v + ηh ( x, y ) for η = α E [ h ( x,y ) 2 | f ( x )= v ] .

Then:

Proof. Following [13], we have

Theorem F.25 (Restatement of Theorem 5.1) . In Algorithm 1, for α, B > 0 , if the following is satisfied:

the 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:

For each v ∈ [1 /m ] define

Then we have ∑ v ∈ [1 /m ] α v > α .

According to Lemma F.24, for each v ∈ [1 /m ] , there exists h v ∈ H such that:

Then,

which contradicts the condition in Equation 177.

The following proposition is a direct corollary from Globus-Harris et al. [13]'s Theorem 4.3.

Proposition F.26. For any distribution D supported on X × Y and Φ ∈ σ ( X ) , take the grouping function class H ⊂ H ′ 1 , Φ := { h (Φ( x )) ∈ C ( X × Y ) } and the predictor class F = R X . For any 0 < α < 1 , B > 0 and an initial predictor f 0 : X → [0 , 1] with | Range ( f 0 ) | ≥ 2 B α , then MC-Pseudolabel ( D, H , F ) halts after at most T ≤ 2 B α steps and outputs a model f ′ T -1 ( x ) that is α -approximately ℓ 2 multicalibrated w.r.t D and H B = { h ∈ H : sup h ( x, y ) 2 ≤ B } .

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:

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:

For 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 that ∣ ∣ f ( t ) ( x ) -E [ Y | Φ( x )] ∣ ∣ ≤ C x M (Σ) t , where

We 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:

However, 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. When Y ⊥ Ψ | Φ implying that Φ is sufficient for prediction, following from E [ Y | Φ , Ψ] = E [ Y | Φ] :

It 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:

It 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 .

Let c = ( c x , c y , c b ) T = ( c Φ , c Ψ , c y , c b ) T . Denote dimensions of Φ , Ψ by d Φ , d Ψ .

Let E [ Y | Φ , Ψ] = ( α ⋆ Φ ) T Φ+( α ⋆ Ψ ) T Ψ and E [Ψ | Φ , Y ] = ( β ⋆ Φ ) T Φ+( β ⋆ y ) T Y with α ⋆ Φ ∈ R d Φ , α ⋆ Ψ ∈ R d Ψ , β ⋆ Φ ∈ R d Φ × d Ψ , β ⋆ y ∈ R 1 × d Ψ .

We have:

According to Theorem F.15, E [ h ( X,Y ) | Φ , Y ] is a constant for different values of Φ , Y .

This implies:

Rearranging to:

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 .

Thus, ˜ f ( t ) -˜ f ( t ) ′ ∈ H .

On the other hand,

The 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) × d t whose columns support the solution space of Equation 195. Then H := C T ( S, T, Y ) T ∈ R d t is a random vector. According to the definition of ˜ f ( t ) ,

In the above equation, k ( t ) ∈ R .

Since ˜ α ( t +1) Ψ = k ( t ) α ( t ) Ψ + c ( t ) Ψ = 0 , we have c ( t ) Ψ = -k ( t ) α ( t ) Ψ . Substituting into Equation 194:

Substituting into Equation 202:

Then we have f ( t +1) = E [ ˜ f ( t +1) | Φ , Ψ] = (˜ α ( t +1) Φ ) T Φ+(˜ α ( t +1) y ) T E [ Y | Φ , Ψ] .

This is equivalent as:

Combining the two equations above, we have:

In the above equation, K ( t ) = ∏ 0 ≤ u<t k ( u ) .

Define ˆ f ( t ) = (ˆ α ( t ) Φ ) T Φ+(ˆ α ( t ) Ψ ) T Ψ

Thus,

In 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 .

Since, (Σ yy -Σ y Φ Σ -1 ΦΦ Σ Φ y ) and (Σ ΨΨ -Σ ΨΦ Σ -1 ΦΦ Σ ΦΨ ) are both Schur complements of Σ 's principal submatrix, they are still positive definite. Thus,

It follows that β ⋆ y α ⋆ Ψ ≥ 0 . So we just have to show β ⋆ y α ⋆ Ψ < 1 .

Since det(Σ) > 0 , by applying row addition on Σ , we have:

It follows that:

Rearranging to:

Thus, 0 ≤ β ⋆ y α ⋆ Ψ < 1 .

In the following, we show ˆ α ( t ) Φ → Σ -1 ΦΦ Σ Φ y . By Equation 212 and | β ⋆ y α ⋆ Ψ | < 1 , we have:

Equation 225 follows from the fact that ( I d Ψ -α ⋆ Ψ β ⋆ y ) α ⋆ Ψ = (1 -β ⋆ y α ⋆ Ψ ) α ⋆ Ψ .

Define γ ⋆ Φ = Σ -1 ΦΦ Σ Φ y such that E [ Y | Φ] = ( γ ⋆ Φ ) T Φ . We have

Thus, ˆ α ( t ) Φ → Σ -1 ΦΦ Σ Φ y . Subsequently, ˆ f ( t ) → (Σ -1 ΦΦ Σ Φ y ) T Φ = E [ Y | Φ] . The convergence ratio is M (Σ) = β ⋆ y α ⋆ Ψ .

We have:

where

Define ˆ ˜ f ( t ) = ( ˆ ˜ α ( t ) Φ ) T Φ+( ˆ ˜ α ( t ) y ) T Y , where

Thus, ˆ ˜ f ( t ) → E [ Y | Φ] . Since ˜ f ( t ) = K ( t ) ˆ ˜ f ( t ) , we have ˜ f ( t ) = E [ Y | ˆ ˜ f ( t ) ] → E [ Y | Φ] . Subsequently, f ( t ) = E [ ˜ f ( t ) | Φ , Ψ] → E [ Y | Φ] .

G Limitations

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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction accurately outline our research questions, and faithfully reflect the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss limitations in section G and analyze computational complexity of our algorithm in section D.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Full assumptions, formal statements of theories and proofs can be found in section F.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experiment settings and implemention of methods are described in section 6.1. Further training details can be found in section E.3. Codes are available in supplementary materials.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: All the datasets in this paper are public with citations (see section 6.1). Code is provided in additional supplemental materials.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Experiment settings are described in section 6.1. Further training details can be found in section E.3.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Standard errors are reported for all results. The factors of variability include random drawing of hyperparameters and seeds, and cross validation on dataset splits. See section 6.1 for details.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer:[Yes]

Justification: Refer to section E.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in this paper conforms with the NeurIPS Code of Ethics.

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.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no safeguards risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All datasets in this paper are cited. Codes are credited with original licenses provided.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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$. \n\nIn 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.\n\nNext, 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).\n\nLastly, experiments are conducted which demonstrate the usefulness of multicalibration in achieving distributional robustness for simple neural networks on the PovertyMap and ACSIncome datasets.\n\nSTRENGTHS:\nOriginality: To my knowledge, the work is very original, especially the consideration of the maximal grouping function space. While extending subgroups to be defined on the label as well as feature is a natural generalization, the authors have unique function space decomposition results (section 4), which are interesting in their own right and whose proofs utilize classical results from linear algebra. Substantial theory is developed through most of the paper, mostly surrounding the inclusion of particular density ratios in the subgroup function classes. I think the results taken together are very original and interesting!\n\nSignificance: understanding what is possible with families of multicalibration post-processing algorithms is a very important and interesting direction, and combining this with out of distribution generalization is an area which is quite interesting.\n\nWEAKNESSES:\nI have not checked the proofs thoroughly.\n\nI 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.\n\nW1: 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.\n\nW2: 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.\n\nW3: 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.\n\nAs 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.\nThere 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?\nAlso, how does Kim et al. [22] relate to theorem 2.3? These should all be explicitly discussed.\n\nAs 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.\n\nW4: 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?\n\nMinor comments\n35: 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.\n\n41-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. \n\nIn line 120, there is some motivation for covariate shift and concept shift. I believe this discussion should go far earlier, perhaps in the introduction.\n\nQUESTIONS:\nQ1. 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.\n\nQ2. 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. \n\nQ3. 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.\n\nQ4. 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?\n\nQ5. 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.\n\nQ6. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-19"}

comments: 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.

STRENGTHS:
The authors provide compelling reasons why multicalibration should be considered in modeling and model evaluation frameworks. In doing so, they provide a calibration framework that confers robustness on post-processed models. The proposed method appears significantly easier to use than methods that leverage actor-critic, domain confusion, and other modeling tricks to achieve robustness.

Section 4 shows that the condition "the grouping function class must include all density ratios between target and source measures" is achievable in practice.

The paper is well written and easy to follow.

WEAKNESSES:
A few questions are given below. I have nothing else to add.

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:
The authors have adequately addressed limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-08

comments: 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.

STRENGTHS:
- The paper provides rigorous and solid theoretical work for its claims.
- MC-Pseudolabel algorithm offers a practical tool that appears to integrate well with their frameworks.

WEAKNESSES:
- 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:
- 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:
See weakness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis is a very interesting work that as far as I know if the first to considering multicalibration with respect to grouping functions that depend on labels. This gives a new approach to handling concept shift, and a new perspective on invariant predictors. I'm personally looking forward to follow-up work. \n\nI found the paper very well-structured and easy to follow.\n\nWEAKNESSES:\nThere 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}?\n\nTypos/suggested edits:\n\nAbstract\n\n“in existence” -> “in the presence”\n\nSection 1\n\n“flexibly designed to incorporates” -> “incorporate”\n\n“Simultaneoulsy producing invariant predictors” -> “produce”\n\n“porverty estimation”\n\nSection 2\n\n“which is learned in the source distribution” -> “which are learned”\n\n2.2 \n\n“simultaneous approaches” -> “simultaneously approach”\n\n“We show that multicalibration notion w.r.t.” -> “We show that our/this multicalibration notion”\n\n\nSection 5\n\nM(Sigma) is defined in the Appendix, but not in Theorem 5.2 where is first appears\n\nQUESTIONS:\nRepeated from weaknesses:\n\n1. Is there a correspondence between weak agnostic learning and extended multicalibration?\n\n2. On what kinds of distributions will the algorithm fail to convergence for H_{2, \\Phi}?\n\nLIMITATIONS:\nThe authors largely addressed the limitations (with the exception of possible computational limitations mentioned in Questions).", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-14"}