Title: Achievable distributional robustness when the robust risk is only partially identified
Venue: NeurIPS 2024
Accepted: yes
Keywords: distributional robustness

domain generalization

causal inference

partial identification
Abstract: In safety-critical applications, machine learning models should generalize well under worst-case distribution shifts, that is, have a small robust risk. Invariance-based algorithms can provably take advantage of structural assumptions on the shifts when the training distributions are heterogeneous enough to identify the robust risk. However, in practice, such identifiability conditions are rarely satisfied – a scenario so far underexplored in the theoretical literature. In this paper, we aim to fill the gap and propose to study the more general setting of partially identifiable robustness. In particular, we define a new risk measure, the identifiable robust risk, and its corresponding (population) minimax quantity that is an algorithm-independent measure for the best achievable robustness under partial identifiability. We introduce these concepts broadly, and then study them within the framework of linear structural causal models for concreteness of the presentation. We use the introduced minimax quantity to show how previous approaches provably achieve suboptimal robustness in the partially identifiable case. We confirm our findings through empirical simulations and real-world experiments and demonstrate how the test error of existing robustness methods grows increasingly suboptimal as the proportion of previously unseen test directions increases.

Achievable distributional robustness when the robust risk is only partially identified

Abstract

In safety-critical applications, machine learning models should generalize well under worst-case distribution shifts, that is, have a small robust risk. Invariancebased algorithms can provably take advantage of structural assumptions on the shifts when the training distributions are heterogeneous enough to identify the robust risk. However, in practice, such identifiability conditions are rarely satisfied - a scenario so far underexplored in the theoretical literature. In this paper, we aim to fill the gap and propose to study the more general setting of partially identifiable robustness . In particular, we define a new risk measure, the worst-case robust risk, and its corresponding (population) minimax quantity that is an algorithmindependent measure for the best achievable robustness under partial identifiability. We introduce these concepts broadly, and then study them within the framework of linear structural causal models for concreteness of the presentation. We use the introduced minimax quantity to show how previous approaches provably achieve suboptimal robustness in the partially identifiable case. We confirm our findings through empirical simulations and real-world experiments and demonstrate how the test error of existing robustness methods grows increasingly suboptimal as the proportion of previously unseen test directions increases.

1 Introduction

The success of machine learning methods typically relies on the assumption that the training and test data follow the same distribution. However, this assumption is often violated in practice. For instance, this can happen if the test data are collected at a later time or using a different measuring device. Without further assumptions on the test distribution, generalization under distribution shift is impossible. However, practitioners often have partial information about the set of possible "shifts" that may occur during test time, inducing a set of feasible test distributions that the model should generalize to. We refer to the resulting set as the robustness set . When a probabilistic model for these possible test distributions is available or estimable, one may aim for good performance on a 'typical' held-out distribution using a probabilistic framework. When no extra information is available or estimable, one possibility is to find a model β that has a small risk R ( β ; P ) on the hardest feasible test distribution. More formally, we aim to achieve a small robust risk defined by

Fanny Yang

where P rob ( θ ⋆ ) corresponds to the robustness set which depends on some true parameter θ ⋆ . In fact, this worst-case robustness aligns with security and safety-critical applications, where a small robust risk is necessary to confidently guarding against possible malicious attacks.

Causality-oriented robustness [11, 31, 45] on the other hand is based on the idea that some structural parameters (like a graphical structure of the model) remain invariant across distributions, while others may vary. For a given set of training distributions, certain sets of varying parameters induce robust risks that are identifiable. Similarly, for a given set of varying parameters, heterogeneous enough training distributions may identify the robust risk.

In practice, robustness methods aiming at minimizing a pre-defined robust risk often suffer from ineffectiveness. For adversarial robustness for example, it is known that when the perturbations during training and test time differ, the robust risks resulting from adversarial training and standard training may be comparable (see, e.g. [52, 26]). Similarly, invariance-based methods are often shown to be less robust than vanilla empirical risk minimization that ignores multi-environment information. Theoretically, besides being effective only for very specific data-generating models [1], invariance-based methods generally are bound to fail when the heterogeneity of the training data is not enough for a given set of possible test shifts. Even though this issue of non-identifiability has been pointed out previously [25, 40], prior work so far was primarily satisfied with such a binary statement - whether identifiability is given or not. We believe that the non-identifiable scenario warrants a more detailed discussion. In particular, we aim to formalize how to quantify the best possible robustness for this partially identifiable setting. In particular, we extend the discussion of invariance-based methods to include the partially identifiable setting, where not only the causal parameter, but the robust risk (1) is not determinable using training data either 1 . Specifically, we aim to discuss the following question:

What is the optimal worst-case performance any model can have for given structural relationships between test and training data and how do existing methods comparatively perform in such settings?

When the robust risk is not identifiable from training data, we obtain a whole set of possible objectives that includes the true robust risk. In this case, we are interested in the best achievable robustness for any algorithm that we capture in a quantity called the worst-case robust risk :

Note that R rob ( β ) is well-defined even when the standard robust risk is not identifiable - it takes the supremum over the robust risks induced by all possible true model parameters θ ⋆ that are consistent with the given set of training data distributions. Furthermore, the minimal value of the identifiable robust risk corresponds to the optimal worst-case performance in the partially identifiable setting. Spiritually, this minimax population quantity is reminiscent of the algorithm-independent limits in classical statistical learning theory [55]. 2 Even though our partial identifiability framework can be

evaluated for arbitrary modeling assumptions on the distribution shift (such as covariate/label shift, DRO, etc.), we present it in a concrete linear setting for clarity of the exposition. Specifically, the setting is motivated by structural causal models (SCMs) with unobserved confounding (cf. Section 2), similar to the setting of IV (instrumental variables) and anchor regression [41, 42]. Concurrently with our work, [24] proposed a framework for partial transportability which is conceptually related to our notion of worst-case robust risk. However, their approach leverages graphical assumptions, i.e., a priori knowledge about the structure of causal models, whereas our focus is a more agnostic multi-environment setting. Additionally, we do not assume a causal data generation process.

The worst-case robust risk (2) not only represents a notion of algorithm-independent optimality for any combination of training and test shifts. In the linear setting in Section 2, we also show theoretically and empirically that the ranking and optimality of different robustness methods change drastically in identifiable vs. partially identifiable settings. The same can be observed in experiments on real-world data. Our experimental results strongly indicate that evaluation and benchmarking on partially identifiable settings are important for determining the effectiveness of robustness methods. Finally, while the worst-case robust predictor derived in the paper is only provably optimal for the linear setting, experiments on real-world data in Section 4 suggest that our estimator may significantly improve upon other invariance-based methods in more realistic scenarios.

2 Setting

In this section, we state the concrete distributional setting on which we introduce our partial identifiability framework. In particular, we consider a data generating process, motivated by structural causal models (SCMs), that allows for hidden confounding, i.e., spurious correlations between the covariates X and the target Y . We describe the structure of the distribution shifts occurring in the training and test environments, which is reminiscent of interventions in causal models. Finally, we introduce our framework for distributional robustness that allows for partial identifiability and define the worst-case robust risk - for any given model, it corresponds to the maximum robust risk among all possible robust risks induced by the training distributions.

2.1 Data distribution and a model of additive environmental shifts

Data generating process (DGP). We first describe the data-generating mechanism that underlies the distributions of all environments e ∈ E that may occur during train or test time. For each environment e ∈ E , we observe the random vector ( X e , Y e ) ∼ P X,Y e consisting of input covariates X e ∈ R d and a target variable Y e ∈ R which satisfy the following data generating process:

where A e ∈ R d , ( η, ξ ) ∈ R d +1 are random vectors A e ∼ P A e , ( η, ξ ) ∼ P η,ξ with finite first and second moments and for which A e ⊥ ⊥ ( η, ξ ) for all e ∈ E .

Invariant mechanism. Note how in this setting, apart from β ⋆ , the distribution P η,ξ of the noise vector ( η, ξ ) remains constant across environments. Without loss of generality, we assume that the noise ( η, ξ ) has mean zero. Note how this linear setting is, in general, more challenging than the standard linear regression setting where η ⊥ ⊥ ξ : due to possible dependencies between η and ξ (induced by, e.g., hidden confounding/spurious features ), classical estimators, such as the ordinary least squares, are biased away from the true parameter β ⋆ . Denote by Σ ⋆ := Cov(( η, ξ )) the joint ( Σ ⋆ Σ ⋆ )

Σ

⋆

η,ξ

(

η,ξ

σ

⋆

ξ

)

and which we assume to be full-rank. We then denote the concatenation of these two invariant parameters by θ ⋆ := ( β ⋆ , Σ ⋆ ) ∈ Θ ⊂ R d × R ( d +1) × ( d +1) - the parameter that remains invariant across all environments.

Structure of the distribution shifts. Note that in the DGP, the distribution shifts between P X,Y e are induced solely by changes in the distribution of the variable A e , whose mean and covariance matrix we denote by E [ A e ] = µ e and Cov[ A e ] = Σ e , respectively. In general, we allow for degenerate shifts, i.e. the covariance Σ e can be singular. We remark that although the additive shift structure in Equation (3) allows the joint distribution P X,Y,A e = P A e × P X,Y | A to change solely via P A e , our distribution shift setting is more general than covariate shift: due to the noise variables η and ξ being

⊤

2

covariance of the noise vector

(

η, ξ

)

, which can be written in block form as

Σ

⋆

=

η

potentially dependent, both the marginal P X e and the conditional distribution P Y | X e can change across environments.

Training and test-time environments. Throughout the paper, we assume that we are given the collection of training distributions P θ ⋆ , E train = { P X,Y θ ⋆ ,e } e ∈E train , where E train denotes the index set of training environments. We omit θ ⋆ when it is clear from the context. Further, for ease of exposition, we assume that E train contains a reference (unshifted) environment e = 0 with A 0 = 0 a.s. In Appendix B, we discuss how our results apply if this condition is not met. During test time, we expect to observe a new, previously unseen distribution P X,Y test which is induced by the DGP (3) and a shift random variable A test ∼ P A test , with corresponding finite mean µ test and covariance Σ test .

Even though we do not have access to P X,Y test during training, the practitioner might have some information about the possible shift distributions P A test that may occur during test time. As an example, we may only have information about the maximum possible magnitude and direction of the test-time mean shift E [ A test ] . In this work, we assume that we are given an upper bound on the second moment of the shift variable, represented by a positive semidefinite (PSD) matrix M test ⪰ 0 such that

In practice, there may be different degrees of knowledge of the feasible set of shifts - when no knowledge is available, one can always choose the most "conservative" bound M test with the range equal to R d × d and large eigenvalues. The more information is available, the smaller the feasible set of test distributions would become. On the other hand, when the test distribution P X test of X is available during training (as in the domain adaptation setting [47]), one can directly compute the optimal shift upper bound via M test = E [ X test X test ⊤ ] . In existing literature, M test is often proportional to the pooled first or second moment of the training shifts, for instance M test = γ ∑ e ∈ E train w e µ e µ ⊤ e in discrete anchor regression [41] or M test = γ ∑ e ∈ E train w e ( µ e µ ⊤ e + Σ e ) in causality-oriented robustness with invariant gradients [45]. Here, w e are the weights representing the probability of a datapoint being sampled from the environment e . As will become apparent in the next sections, our population-level results are not impacted by the distribution of the environment variable, which we thus omit in the following.

We now provide an example based on structural causal models (SCM) that falls under the aforementioned distrubtion shift setting.

Example 1 . Consider the structural causal model and its induced graph in Figure 1. In this model, the variable Z is a soft intervention on the covariates X . Additionally, the exogenous noise vector ( ε X , ε Y , ε H ) and the intervention variable Z are mutually independent. This model is the basis of multiple causality-oriented robustness works, e.g. [41, 45]. Let β ⋆ := B ⊤ Y X and ξ := B Y H H + ε Y . Then, from (5), we obtain Y = B Y X X + ( B Y H H + ε Y ) = X ⊤ β ⋆ + ξ . Suppose that I -B is invertible and let C := ( I -B ) -1 with entries C XX , C XY , etc. Define A := C XX Z and η := C XX ε X + C XY ε Y + C XH ε H . Then, we can write X = A + η . Since shifts in distribution of Z induce shifts in the distribution of A , a collection of interventions { Z e } e ∈E train translates into a collection of additive shifts { A e } e ∈E train and gives rise to training distributions varying with the environment e . In summary, our DGP Equation (3) includes the classical setting of causality-oriented robustness as depicted in Figure 1.

2.2 The robust risk

Our goal is to find an estimator using the training data that has a small risk, in this paper exclusively the expected square loss R ( β ; P ) := E P [( Y -β ⊤ X ) 2 ] , over the robustness set. In our setting, given

a test shift upper bound M test defined in Equation (4), the robustness set corresponds to

yielding the corresponding robust risk that reads

Wecall the minimizer of the robust risk β rob θ ⋆ := arg min β ∈ R d R rob ( β ; θ ⋆ , M test ) the robust predictor . For the squared loss and linear model, the robust risk can be computed in closed form and solely depends on M test and the invariant parameters θ ⋆ = ( β ⋆ , Σ ⋆ ) , and not on other properties of the distributions:

This observation motivates us to define an equivalence relation between two data-generating processes that holds whenever they induce the same robust risk for any model β and shift upper bound M test . Specifically, observe that DGP 1 and DGP 2 induce the same robust risks for all M test and β iff β ⋆ 1 = β ⋆ 2 and P η,ξ 1 ∼ = P η,ξ 2 , where ∼ = denotes the equivalence of distributions based on equality of their first and second moments. Thus, in the following, we treat our data-generating process as uniquely defined by θ ⋆ up to this equivalence relation.

In practice, the model parameters θ ⋆ typically cannot be identified from the training distributions, and the robust risk R rob can only be computed for specific combinations of training and test shifts, studied, e.g., in [41, 45]. In the next section, we describe concepts that allow us to reason about robustness in the case when the robust risk is only partially identifiable.

2.3 Partially identifiable robustness framework

We start by formally introducing the general notion of partial identifiability for the robust risk. The following notion of observational equivalence of parameters is reminiscent of the corresponding notion in the econometrics literature [16]:

Definition 1 (Observational equivalence) . Two model parameter vectors θ 1 = ( β 1 , Σ 1 ) and θ 2 = ( β 2 , Σ 2 ) are observationally equivalent with respect to a set of shift distributions { P A e : e ∈ E train } 3 if they induce the same set P θ, E train of training distributions over the observed variables ( X e , Y e ) as described in Section 2.1, i.e.

By observing P θ ⋆ , E train , we can identify the model parameters θ ⋆ up to the observationally equivalent set defined as

In general, observationally equivalent set is not a singleton, that is, θ ⋆ is not identifiable from the collection of training environments P θ ⋆ , E train . However, prior work has exclusively considered test shifts M test that still allow identifiability of the robust risk nonetheless, depicted in Figure 2a and discussed again in Section 3.2. In this work we argue for analyzing the more general partially identifiable setting, where set-identifiability of the invariant parameter θ ⋆ only allows us to compute a superset of the robustness set

and correspondingly, a set of robust risks {R rob ( β ; θ, M test ) : θ ∈ Θ eq } . In this case, we would still like to achieve the 'best-possible' robustness, that is the test shift robustness for the 'hardest-possible' parameters that could have induced the observed training distributions.

model parameter space model parameter space

𝜃

𝜃

𝜃

space of robust risks of space of robust risks of

𝑅𝑟𝑜𝑏(𝛽; 𝜃, 𝛾𝑀 𝑠𝑒𝑒𝑛 )

𝜃

𝑅𝑟𝑜𝑏(𝛽; 𝜃, 𝛾𝑀 𝑠𝑒𝑒𝑛 )

model parameter space

Θ

set of robust risks of

𝛽

𝜃 ⋆

𝛩𝑒𝑞

model parameter space

Θ

set of robust risks of

𝜃 ⋆

𝛩𝑒𝑞

𝑅𝑟𝑜𝑏 (𝛽; 𝜃

⋆

, 𝑀

𝑠𝑒𝑒𝑛 )

𝑅𝑟𝑜𝑏 (𝛽; 𝜃

⋆

, 𝑀

𝑢𝑛𝑠𝑒𝑒𝑛 )

(a) Identifiable robust risk

model parameter space

Θ

set of robust risks of

𝛽

𝜃 ⋆

𝛩𝑒𝑞

model parameter space

Θ

set of robust risks of

𝛽

𝜃 ⋆

𝛩𝑒𝑞

𝑅𝑟𝑜𝑏 (𝛽; 𝜃

⋆

, 𝑀

𝑠𝑒𝑒𝑛 )

𝑅𝑟𝑜𝑏 (𝛽; 𝜃

⋆

, 𝑀

𝑢𝑛𝑠𝑒𝑒𝑛 )

(b) Partially identifiable robust risk

𝛽

Definition 2 (Worst-case robust risk and the minimax quantity) . For the data model in Equation (3) , the worst-case robust risk is defined as

The optimal robustness on test shifts bounded by M test given training data P θ ⋆ , E train is described by the minimax quantity

When the minimizer of Equation (10) exists, we call it the worst-case robust predictor defined by

In the next sections, we explicitly compute these quantities for the linear setting of Section 2. This will allow us to compare the best achievable robustness in the partially identified case with the guarantees of prior methods in this setting.

3 Theoretical results for the linear setting

We now compute the worst-case robust risk (9) and derive a lower bound for the minimax quantity (10) in the linear additive shift setting of Section 2. We then compare the worst-case robust risk of some existing robustness methods and ordinary least squares (OLS) with the minimizer of the worst-case robust risk both theoretically and empirically.

3.1 Minimax robustness results for the linear setting

The degree to which the model parameters θ ⋆ in the linear additive shift setting (3) can be identified depends on the number of environments and the total rank of the moments of the additive shifts. For structural causal models, this is well-studied, for instance, in the instrumental variable (IV) regression literature [4, 9]. As we show in Proposition 1, the true parameter β ⋆ can only be identified along the directions of the training-time mean and variance shifts µ e and Σ e . Therefore, if not enough shift directions are observed, β ⋆ is merely set-identifiable , leading to set-identifiability of the robust prediction model (8). More formally, we denote by S the subspace consisting of all additive shift directions seen during training :

and by S ⊥ its orthogonal complement. The definition of the space S induces an orthogonal decomposition of the true parameter β ⋆ = β S + β S ⊥ . The identifiable part β S then uniquely defines a set of identified model parameters that reads

that can be computed from the training distributions. For the following results, we assume a similar decomposition of the test shift upper bound M test which is essentially a decomposition into "seen" and "unseen" directions.

Θ

Θ

𝛽

𝛽

Assumption 3.1 (Structure of M test ) . We assume that M test = γM seen + γ ′ RR ⊤ , where γ, γ ′ ≥ 0 , M seen is a PSD matrix satisfying range M seen ⊂ S and R is a semi-orthogonal matrix satisfying range R ⊂ S ⊥ .

In the next proposition, we show that the model parameters and robust predictor can be identified up to a neighborhood around θ S .

Proposition 1 (Identifiability of model parameters and robust predictor) . Suppose that the set of training and test distributions is generated according to Section 2.1 and Assumption 3.1 holds. Then,

(a) the model parameters generating the training distribution (3) can be identified up to the following observationally equivalent set :

(b) the robust predictor β rob θ as defined in Equation (8) is identified up to the set

where B rob Θ eq = { β rob θ : θ ∈ Θ eq } .

The proof of Proposition 1 is provided in Appendix F.1. Proposition 1 implies two well-known settings: If we observe a rich enough set P E train of training environments such that S = R d , the model parameters are uniquely identified, corresponding to the setting of full-rank instruments [4]. From a dual perspective, for a given set of training environments, the robust predictor is identifiable whenever the test shifts are in the same direction as the training shifts , i.e. range M test ⊂ S and R = 0 - this holds even when the invariant parameters are not identifiable and S ̸ = R d . This is the setting considered e.g. in anchor regression [41] and discussed again in Section 3.2 and Appendix C.

̸

So far, we have described how the identifiability of the robust prediction model depends on the structure of both the training environments (via the space S ) and the test environments (via M test ). Wenowaim to compute the smallest achievable robust loss for the general partially identifiable setting, which allows for R = 0 . In particular, we provide a lower bound on the best-possible achievable distributional robustness formalized by the minimax quantity (10). First observe that without further assumptions on the parameter space Θ , the observationally equivalent set is unbounded, and the worst-case robust risk (9) can be infinite. The following boundedness assumption allows us to provide a fine-grained analysis of robustness in a partially identified setting.

Assumption 3.2 (Boundedness of the causal parameter) . There exists a constant C > 0 such that any parameter β in the DGP (3) is norm-bounded by C , i.e. ∥ β ∥ 2 ≤ C .

Furthermore, we define C ker = √ C 2 -∥ β S ∥ 2 , the maximum norm of the non-identified part of the linear parameter β ⋆ . Finally, recall that the reference distribution P X,Y θ ⋆ , 0 is observed and hence identifiable.

Theorem 3.1. Assume that the training and test data follow the data-generating mechanism in Section 2.1 and M test satisfies Assumption 3.1 for some M seen , R with range M seen ⊂ S , range R ⊂ S ⊥ . Further, let Assumption 3.2 hold with parameter C . The worst-case robust risk (9) is then given by

The minimax quantity in Equation (10) is lower bounded as follows:

where γ ′ th = ( κ (Σ ⋆ η )+1) ∥ RR ⊤ Σ S η,ξ ∥ C ker 4 . Moreover, for small unseen shifts

̸

We prove Theorem 3.1 in Appendix F.2. First, in the case of no new test shifts where γ ′ = 0 (as it appears in prior work [41, 45]) we can plug in the robust risk Equation (7) into Equation (16) to observe the following: as the strength γ of the shift grows, the optimal robust risk saturates. On the other hand, if γ ′ = 0 , i.e., the test shift contains new directions w.r.t. to the training data, the best achievable robustness M (Θ eq , M test ) grows linearly with γ ′ . Further note that for γ ′ ≥ γ ′ th , we have a tight expression for the minimax quantity and the worst-case robust predictor β rob Θ eq can be explicitly computed (cf. Appendix F.2) and is orthogonal to the space range R of non-identifiable test shift directions. In other words, for large shifts in non-identified directions, the optimal robust model would "abstain" from prediction in those directions. For smaller γ ′ , β rob Θ eq gradually utilizes less information in the non-identified directions, thus interpolating between maximum predictive power (OLS) and robustness w.r.t. new directions (abstaining). The model β rob Θ eq is a population quantity that is identifiable from the collection of training distributions . When only finite samples are available, we discuss in Appendix D how we can compute the worst-case robust estimator by minimizing an empirical loss function that can be computed from multi-environment data.

3.2 Theoretical analysis of existing finite robustness methods

We now evaluate existing finite robustness methods in our partial identifiability framework and characterize their (sub)optimality in different scenarios. A spiritually similar systematic comparison of domain adaptation methods is presented in [12], however, in our setting, the robust risk is not identifiable from data. Concretely, we compare discrete anchor regression [41] and pooled OLS estimators 5 with the minimax quantity in Theorem 3.1. We consider the same scenario as in discrete anchor regression, which is a the specific case of the setting in Equation (3), where for each environment e , A e is just a mean shift with variance 0 . In addition, discrete anchor regression assumes that the environment variable E ∈ E train follows a probability distribution with P [ E = e ] = w e . The discrete anchor setting then corresponds to setting a test shift upper bound M test = γM anchor for some γ > 0 (cf. Equation (4)) with M anchor = ∑ e ∈E train w e µ e µ ⊤ e . The (oracle) discrete anchor regression estimator minimizes the robust risk and reads

The pooled ordinary least squares (OLS) estimator β OLS corresponds to β anchor with γ = 1 . We observe that the test shifts bounded by γM anchor are fully contained in the space of identified directions S , since S = range ∪ e ∈E train µ e µ ⊤ e = range M anchor . Thus, according to Proposition 1, the robust risk and robust predictor β anchor are identifiable for all γ > 0 . In the next corollary, we compute worst-case robust risk of both β anchor and β OLS with respect to the more general shifts bounded by M test := γM anchor + γ ′ RR ⊤ , thus possibly including unseen shifts consisting of additional unseen shifts range R ⊂ S ⊥ .

Corollary 3.2 (Worst-case robust risk of the anchor regression estimator) . Assume that the test shift upper bound is given by M test := γM anchor + γ ′ RR ⊤ . Let P X,Y train = ∑ e w e P X,Y e be the pooled training distribution. Then the general worst-case robust risk is given by

Furthermore, the the anchor and OLS predictor, respectively, it holds that there exists constants c 1 , c 2 , c 3 independent of γ, γ ′ such that

In contrast, the best achievable robustness reads

Observe that the worst-case robust risk in the extended anchor regression setting is equal to the anchor regression risk with an additional non-identifiability penalty term γ ′ ( C ker + ∥ R ⊤ β ∥ 2 ) 2 . The anchor regression estimator is optimal in the limit of vanishing unseen shifts but (for any γ ) significantly deviates 6 from the best achievable robustness for larger unseen shifts γ ′ ≥ γ ′ th . Moreover, in case of completely new shifts ( γ = 0 ), pooled OLS and the anchor estimator achieve the same rate in γ ′ , showcasing how finite robustness methods can perform similarly to empirical risk minimization if the assumptions on the robustness set are not met. We provide additional performance comparisons for the more general shift in Appendix C and the proof of the corollary in Appendix F.3.

4 Experimental results

In this section, we provide empirical evidence of our theoretical conclusions in Sections 3.1 and 3.2. In particular, we compare the prediction performance of multiple existing robustness methods to the (estimated) minimax robustness in identifiable and partially identifiable settings. We observe that both on synthetic and real-world data, in the partially identified setting, empirical risk minimization and invariance-based robustness methods not only have significantly sub-optimal test loss, but also perform more similarly, thereby aligning with our theoretical results in Section 3.2. This stands in contrast to the identifiable setting, where the anchor predictor is optimal up to finite-sample effects. Furthermore, we observe that even though the minimizer of the worst-case robust risk is optimal only for the linear causal setting in Section 2.1, it surprisingly outperforms existing methods in a real-world experiment.

Experiments on synthetic Gaussian data We simulate Gaussian covariates according to Equation (3) with multiple environments differing by linearly independent randomly selected mean shifts. For a randomly sampled collection of mean shifts, we evaluate a proxy for the worst-case robust risk by picking the most adversarial ( β ⋆ , Σ ⋆ η ) for the shifts, and then computing its robust risk (7). We describe the full details of the data generation and loss evaluation in Appendix E.1. We consider two shift scenarios: in the identifiable case in (see Figure 2a), the test environment is only perturbed by bounded shifts in training directions with increasing strength γ , as considered in prior work [41, 45]. In the non-identifiable case (see Figure 2b), the test environment is perturbed by a mixture of training shifts and shifts in previously unobserved directions, where γ is fixed and γ ′ varies (cf. Assumption 3.1). We compute the empirical minimizers ˆ β OLS , ˆ β anchor and ˆ β rob Θ eq of the OLS, anchor regression and worst-case robust losses, respectively, and compare their worst-case robust risk (mean squared error) in Figure 3. In the identifiable setting - Figure 3 (left) - the robust risk is asymptotically constant across γ for both robust methods, while the error for the OLS, or vanilla ERM, estimator increases linearly. In contrast, in the second, partially identified, setting - Figure 3 (right) - all estimators exhibit linearly increasing test errors; however the slopes of the anchor and OLS estimator are much steeper and lead to larger errors than the empirical minimizer of (15) that closely matches the analytic theoretical lower bound.

Identifiable case (

γ′

=

0)

0.0

2.5

5.0

7.5

10.0

0.50

0.75

1.00

1.25

γ

(

β

;

Θ

eq

,

Θ

test

)

Partially identifiable case (fixed

γ ≠

0)

0.0

0.5

1.0

1.5

2.0

0

5

10

15

γ′

Methods:

Worst-case Rob.

Anchor

OLS

Lower bound

Real-world data experiments We evaluate the performance of OOD methods using single-cell gene expression data from [38], consisting of d = 622 genes across observational and interventional environments. As in [44], we focus on 28 genes active in the observational environment. For each gene j = 1 , . . . , 28 , we define the target variable Y := X j and select the three genes most strongly

̸

correlated with Y as covariates. This yields 28 prediction problems indexed by j , each consisting of data from an observational environment O and three interventional environments I j 1 , I j 2 , I j 3 representing the gene knockout on a single covariate. For each prediction problem, we consider three training datasets D j 1 , D j 2 , D j 3 , obtained by combining data from O with a single interventional environment I j 1 , I j 2 , I j 3 , respectively. For each training dataset D jk , k = 1 , 2 , 3 , we evaluate the mean-squared error (MSE) at test time using four datasets consisting of varying proportions of unseen shifts (e.g., ' 33% unseen directions' in Figure 4 represents a test dataset with 67% observations sampled from I jk and 33% from I jℓ with ℓ = k ). Hence, for each prediction problem predicting a gene j , we evaluate on 12 configurations (three training and four test datasets). 7 Figure 4 illustrates the test MSE of the worst-case robust estimator (Worst-case Rob.) alongside anchor regression, invariant causal prediction (ICP), DRIG, and OLS, as a function of perturbation strength s . 8 For a given proportion of unseen shifts, s controls the distance of the test data points from the observational mean, acting as a proxy for shift strengths γ and γ ′ . 9 We observe that the performance ranking of the robustness methods significantly varies with the proportion of new test shift directions. As expected, when no new shift directions are present at test time (0%), anchor regression and DRIG are optimal, since they protect against shifts observed at training time. However, as soon as some unseen directions are present, their performance becomes inferior to OLS/ERM and the gap to the worst-case robust predictor (in the linear setting described in Section 2) grows with the proportion of unseen shifts. Further, while the MSE of the previous invariant methods increases significantly with the strength of the test shift s , the test loss of the worst-case robust predictor remains relatively stable.

0% unseen directions

33% unseen directions

67% unseen directions

100% unseen directions

0.25

0.50

0.75

0.25

0.50

0.75

0.25

0.50

0.75

0.25

0.50

0.75

0.1

0.2

0.3

0.4

0.5

Perturbation strength s

Test MSE

Methods:

Worst-case Rob.

Anchor

DRIG

ICP

OLS

5 Conclusion and future directions

This paper introduces the worst-case robust risk - a quantity that is well-defined even in settings where the usual robust risk is not computable from training distributions, and in identifiable scenarios [41, 45] reduces to the conventional robust risk. We instantiate our general framework for linear models with additive distribution shifts and compute tight lower bounds for this setting. Further, we demonstrate how i) the benefits of invariance-based robustness methods strongly decrease in the partially identifiable setting; and ii) this suboptimality increases with perturbation strength and proportion of previously unobserved test shifts.

The main limitation of our paper is its reliance on a linear setting to explicitly compute the worst-case robust risk and estimate the minimax quantity. However, we expect that the results and intuition developed in this paper can be extended to linear shifts in a lower-dimensional latent space via a suitable parametric or non-linear map [51, 10]. Important future directions include extending our results to more general shift models, non-linear functional relationships and the classification setting. Further, a potential use of our work is in the field of active intervention selection (e.g, [57, 20]). By computing the most adversarial model parameter for a given estimator, e.g., OLS, we can obtain an intervention which minimizes the worst-case robust risk of the estimator on the next unseen shift.

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. Advances in Neural Information Processing Systems , 34:3438-3450, 2021.

Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning , pages 145-155. PMLR, 2020.

Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R. Varshney. Empirical or invariant risk minimization? A sample complexity perspective. In International Conference on Learning Representations , 2021.

Takeshi Amemiya. Advanced Econometrics . Harvard University Press, 1985.

Joshua D Angrist, Guido W Imbens, and Donald B Rubin. Identification of causal effects using instrumental variables. Journal of the American Statistical Association , 91(434):444-455, 1996.

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

Aharon Ben-Tal, Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science , 59(2):341-357, 2013.

Andrew Bennett, Nathan Kallus, and Tobias Schnabel. Deep generalized method of moments for instrumental variable analysis. Advances in Neural Information Processing Systems , 32, 2019.

Roger J Bowden and Darrell A Turkington. Instrumental variables . Number 8. Cambridge university press, 1990.

Simon Buchholz, Goutham Rajendran, Elan Rosenfeld, Bryon Aragam, Bernhard Schölkopf, and Pradeep Ravikumar. Learning linear causal representations from interventions under general nonlinear mixing. Advances in Neural Information Processing Systems , 36, 2024.

Peter Bühlmann. Invariance, causality and robustness. Statistical Science , 35(3):404-426, 2020.

Yuansi Chen and Peter Bühlmann. Domain adaptation under structural causal models. Journal of Machine Learning Research , 22(261):1-80, 2021.

Mathieu Chevalley, Yusuf Roohani, Arash Mehrjou, Jure Leskovec, and Patrick Schwab. Causalbench: A large-scale benchmark for network inference from single-cell perturbation data. arXiv preprint arXiv:2210.17283 , 2022.

Rune Christiansen, Niklas Pfister, Martin Emil Jakobsen, Nicola Gnecco, and Jonas Peters. A causal framework for distribution generalization. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(10):6614-6630, 2021.

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

Jean-Marie Dufour and Cheng Hsiao. Identification , pages 65-77. Palgrave Macmillan UK, London, 2010.

Jianqing Fan, Cong Fang, Yihong Gu, and Tong Zhang. Environment invariant linear least squares. arXiv preprint arXiv:2303.03092 , 2023.

Justin Frake, Anthony Gibbs, Brent Goldfarb, Takuya Hiraiwa, Evan Starr, and Shotaro Yamaguchi. From perfect to practical: Partial identification methods for causal inference in strategic management research. Available at SSRN 4228655 , 2023.

Charlie Frogner, Sebastian Claici, Edward Chien, and Justin Solomon. Incorporating unlabeled data into distributionally robust learning. Journal of Machine Learning Research , 22(56):1-46, 2021.

Juan L Gamella and Christina Heinze-Deml. Active invariant causal prediction: Experiment selection through stability. Advances in Neural Information Processing Systems , 33:1546415475, 2020.

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2021.

Jason Hartford, Greg Lewis, Kevin Leyton-Brown, and Matt Taddy. Deep IV: A flexible approach for counterfactual prediction. In International Conference on Machine Learning , pages 1414-1423. PMLR, 2017.

Martin Emil Jakobsen and Jonas Peters. Distributional robustness of k-class estimators and the pulse. The Econometrics Journal , 25(2):404-432, 2022.

Kasra Jalaldoust, Alexis Bellot, and Elias Bareinboim. Partial transportability for domain generalization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024.

Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics , pages 4069-4077. PMLR, 2021.

Daniel Kang, Yi Sun, Tom Brown, Dan Hendrycks, and Jacob Steinhardt. Transfer of adversarial robustness between perturbation types. arXiv preprint arXiv:1905.01034 , 2019.

Lucas Kook, Beate Sick, and Peter Bühlmann. Distributional anchor regression. Statistics and Computing , 32(3):39, 2022.

David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (REx). In International Conference on Machine Learning , pages 5815-5826. PMLR, 2021.

Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally robust optimization with data geometry. Advances in Neural Information Processing Systems , 35:33689-33701, 2022.

Sara Magliacane, Thijs Van Ommen, Tom Claassen, Stephan Bongers, Philip Versteeg, and Joris M Mooij. Domain adaptation by using causal inference to predict invariant conditional distributions. Advances in Neural Information Processing Systems , 31, 2018.

Nicolai Meinshausen. Causality from a distributional robustness point of view. In 2018 IEEE Data Science Workshop (DSW) , pages 6-10. IEEE, 2018.

Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming , 171(1):115-166, 2018.

Krikamol Muandet, Arash Mehrjou, Si Kai Lee, and Anant Raj. Dual instrumental variable regression. Advances in Neural Information Processing Systems , 33:2710-2721, 2020.

Judea Pearl. Causality: Models, Reasoning and Inference . Cambridge University Press, USA, 2nd edition, 2009.

Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society Series B: Statistical Methodology , 78(5):947-1012, 2016.

Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms . 2017.

Lei S Qi, Matthew H Larson, Luke A Gilbert, Jennifer A Doudna, Jonathan S Weissman, Adam P Arkin, and Wendell A Lim. Repurposing CRISPR as an RNA-guided platform for sequence-specific control of gene expression. Cell , 152(5):1173-1183, 2013.

Joseph M Replogle, Reuben A Saunders, Angela N Pogson, Jeffrey A Hussmann, Alexander Lenail, Alina Guna, Lauren Mascibroda, Eric J Wagner, Karen Adelman, Gila Lithwick-Yanai, et al. Mapping information-rich genotype-phenotype landscapes with genome-scale Perturb-seq. Cell , 185(14):2559-2575, 2022.

Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. Journal of Machine Learning Research , 19(36):1-34, 2018.

Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. In International Conference on Learning Representations , 2021.

Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, and Jonas Peters. Anchor regression: Heterogeneous data meet causality. Journal of the Royal Statistical Society Series B: Statistical Methodology , 83(2):215-246, 2021.

Sorawit Saengkyongam, Leonard Henckel, Niklas Pfister, and Jonas Peters. Exploiting independent instruments: Identification and distribution generalization. In International Conference on Machine Learning , pages 18935-18958. PMLR, 2022.

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

Christoph Schultheiss and Peter Bühlmann. Assessing the overall and partial causal wellspecification of nonlinear additive noise models. Journal of Machine Learning Research , 25(159):1-41, 2024.

Xinwei Shen, Peter Bühlmann, and Armeen Taeb. Causality-oriented robustness: exploiting general additive interventions. arXiv preprint arXiv:2307.10299 , 2023.

Yuge Shi, Jeffrey Seely, Philip Torr, Siddharth N, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. In International Conference on Learning Representations , 2022.

Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference , 90(2):227-244, 2000.

Rahul Singh, Maneesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. Advances in Neural Information Processing Systems , 32, 2019.

Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations , 2018.

Elie Tamer. Partial identification in econometrics. Annu. Rev. Econ. , 2(1):167-195, 2010.

Nikolaj Thams, Michael Oberst, and David Sontag. Evaluating robustness to dataset shift via parametric robustness sets. Advances in Neural Information Processing Systems , 35:1687716889, 2022.

Florian Tramèr and Dan Boneh. Adversarial training and robustness for multiple perturbations. In NeurIPS , pages 5858-5868, 2019.

Aad W Van der Vaart. Asymptotic statistics , volume 3. Cambridge University Press, 2000.

Chuanlong Xie, Haotian Ye, Fei Chen, Yue Liu, Rui Sun, and Zhenguo Li. Risk variance penalization. arXiv preprint arXiv:2006.07544 , 2020.

Bin Yu. Assouad, Fano, and le Cam. In Festschrift for Lucien Le Cam: Research papers in probability and statistics , pages 423-435. Springer, 1997.

Yi Yu, Tengyao Wang, and Richard J Samworth. A useful variant of the Davis-Kahan theorem for statisticians. Biometrika , 102(2):315-323, 2015.

Jiaqi Zhang, Louis Cammarata, Chandler Squires, Themistoklis P Sapsis, and Caroline Uhler. Active learning for optimal intervention design in causal models. Nature Machine Intelligence , 5(10):1066-1075, 2023.

Appendix

The following sections provide deferred discussions, proofs and experimental details.

Table of contents

A Extended related work

To put our work into context, first, we discuss relevant distributional robustness literature organized according to structural assumptions on the desired robustness set. Second, we summarize existing views on partial identifiability in the causality and econometrics literature and how our findings connect to their perspective.

No structural assumptions on the shift. DRO: Distributionally robust optimization (DRO) tackles the problem of domain generalization when the robustness set is a ball around the training distribution w.r.t. some probability distance measure, e.g., Wasserstein distance [49, 32] or f -divergences [7, 15]. Considering all test distributions in a discrepancy ball can lead to overly conservative predictions, and therefore, alternatives have been proposed in, e.g., the Group DRO literature [43, 19, 29]. However, these methods cannot protect against perturbations larger than the ones seen during training time and do not provide a clear interpretation of the perturbations class [45].

Structural assumptions on the shift. Robustness from the lens of causality takes a step further, by assuming a structural causal model [34] generating the observed data ( X,Y ) . Infinite robustness methods: The motivation of causal methods for robustness is that the causal function is worst-case optimal to predict the response under interventions of arbitrary direction and strength on the covariates [31, 11]. For this reason, causal models achieve what we call infinite robustness . Depending on the assumptions of the SCM, there are different ways to achieve infinite robustness. When there are no latent confounders, several works [35, 17, 30, 39, 2, 6, 46, 54, 28, 1] aim to identify the causal parents and achieve infinite robustness by exploiting the heterogeneity across training environments. In the presence of latent confounders, it is possible to achieve infinite robustness by identifying the causal function with, e.g., the instrumental variable method [5, 22, 48, 8, 33]. There are different limitations to infinitely robust methods. First, the identifiability conditions of the causal parents and/or causal function are often challenging to verify in practice. Second, ERM can outperform these methods when the interventions (read shifts) at test time are not arbitrarily strong or act directly on the response or latent variable [3, 21]. Finite robustness methods: In real data, shifts of arbitrary direction and strength in the covariates are unrealistic. Thus, different methods [41, 23, 27, 45, 14] trade off robustness against predictive power to achieve what we call finite robustness . The main idea of finite robustness methods is to learn a function that is as predictive as possible while protecting against shifts up to some strength in the directions that are observed during training time. These methods, however, only provide robustness guarantees that depend on the heterogeneity of the training data and do not offer insights into the limits of algorithm-independent robustness under shifts in new directions.

Partial identifiability : The problem of identification is at the center of the causal and econometric literature [36, 4]. It studies the conditions under which the (population) training distribution uniquely determines the causal parameters of the underlying SCM. Often, the training distribution only offers partial information about the causal parameters and, therefore, determines a set of observational equivalent parameters. This setting is known as partial or set identification and is used in causality and econometrics to learn intervals within which the true causal parameter lies [50]. In this work, we borrow the notion of partial identification to study the problem of distributional robustness when the robustness set itself is only partially identified.

B Extension to the general additive shift setting

We discuss how our setting changes when we relax the assumptions on the existence of the reference environment. We consider the data-generating process in Equation (3), where E train = [ m ] , m ∈ N . If no environment e exists with µ e = 0 and Σ e = 0 , we first pick an arbitrary distribution P X,Y ref as the reference environment 10 . We denote Σ ′ η := Σ ⋆ η +Σ ref .

First, we show we can express the space S of training additive shift directions defined in Equation (12) in the general case. We center all distributions by µ ref to obtain centered distributions ˜ Pr that

E X ∼ ˜ P e [ X ] = 0 . With respect to the arbitrary reference environment, we now define

We now consider test shifts with respect to the environment P X,Y ref 11 . We define the test shift upper bound M test = γM seen + γ ′ RR ⊤ , where range M seen ⊂ S and range R ⊂ S ⊥ . Again, we can decompose the parameter β ⋆ as β ⋆ = β S + β S ⊥ . The projection β S of the causal parameter onto the relative training shifts induces the following observationally equivalent parameters corresponding to the reference distribution:

Again, θ S can be identified from the training distributions and is referred to as the identified model parameters . The following adapted version of Proposition 1 shows that assuming shifts on P X,Y ref , the robust prediction model is only identifiable if the test shifts are in the direction of the relative training shifts:

Proposition 2 (Identifiability of reference distribution parameters and robust prediction model) . Suppose that the set of training and test distributions is generated according to Equations (3) and (4) . Then, θ S is observationally equivalent to θ ⋆ and computable from training distributions. Furthermore, it holds that

(a) the model parameters generating the reference distribution can be identified up to the following observationally equivalent set :

rob

(b) the robust prediction model β as defined in Equation (8) is identified up to the set

The proof is analogous to Appendix F.1. A version of Theorem 3.1 for perturbations on the reference environment follows accordingly.

C Comparison to finite robustness methods continued

C.1 The setting of continuous anchor regression [41]

In this section, we evaluate the worst-case robust risk of the continuous anchor regression estimator. In the continuous anchor regression setting, during training we observe the distribution according to the process X = MA + η ; Y = β ⋆ ⊤ X + ξ , where A is an observed q -dimensional anchor variable with mean 0 and covariance Σ A and M ∈ R d × q is a known matrix. Note that in this setting, we do not have a reference environment, but, since the anchor variable is observed, the distribution of the additive shift MA is known. The test shifts are assumed to be bounded by M test = γM Σ A M ⊤ . Since range M test ⊂ S = range M , no new directions are observed during test time, in other words, R = 0 . Thus, both the corresponding robust loss and the anchor regression estimator can be determined from training data. It holds that

Again, the pooled OLS estimator corresponds to β anchor with γ = 1 . Similar to the discrete anchor case, in case the test shifts are given by M new = γM Σ A M ⊤ + γ ′ RR ⊤ , the worst-case robust risk (9) is given by

and for the best worst-case robustness of the anchor estimator it holds

The above results follow by analogy with Appendix F.3.

C.2 Distributionally robust invariant gradients (DRIG) [45]

DRIG [45] introduce a more general additive shift framework, where a collection of additive shifts A e is given with moments ( µ e , Σ e ) . For each environment e , we observe data ( X e , Y e ) distributed according to the equations X e = A e + η ; Y e = β ⋆ ⊤ X e + ξ , where the noise is distributed like in Equation (3). This DGP arises from the structural causal model assumption as described in Figure 1. DRIG consider more a more general intervention setting, additionally allowing additive shifts of Y and hidden confounders H . However, their identifiability results can only be shown for the case of interventions on X , and since identifiability of the causal parameter is a crucial part of our analysis, we only consider shifts on the covariates. DRIG assumes existence of a reference environment e = 0 with µ 0 = 0 and for which it is required that the second moment of the reference environment is dominated by the second moment of the training mixture:

This assumption allows [45] to derive the DRIG estimator which is robust against test shifts upper bounded by M DRIG := γ ∑ e ∈ [ m ] w e (Σ e -Σ 0 + µ e µ ⊤ e ) . The following lemma allows us to make further statements about M DRIG :

Lemma C.1. Let A and B be positive semidefinite matrices such that B ⪯ A . Then it holds that range B ⊂ range A .

̸

Proof. It suffices to show that ker A ⊂ ker B . ( ker A ⊂ ker B implies that range A = (ker A ) ⊥ ⊂ (ker B ) ⊥ = range B .) Consider x ∈ ker A , x = 0 . Then it holds that x ⊤ ( A -B ) x = x ⊤ Ax -x ⊤ Bx = 0 -x ⊤ Bx ≥ 0 , from which it follows that x ⊤ Bx = 0 and thus x ∈ ker B .

Because of the assumption Σ 0 ⪯ ∑ e ∈ [ m ] w e (Σ e + µ e µ ⊤ e ) , by Lemma C.1 it follows that range Σ 0 ⊂ ∪ e ≥ 1 range (Σ e + µ e µ ⊤ e ) and thus

Hence, the robustness directions achievable by DRIG in the "dominated reference environment" setting are the same as the ones under the assumption Σ 0 = 0 .

Again, we observe that the test shifts bounded by γM DRIG are fully contained in the space of identified directions S . If the test shifts are instead bounded by M new := γM DRIG + γ ′ RR ⊤ , including some unseen directions range R ⊂ S ⊥ , the robust risk in the DRIG setting is only partially identified. The worst-case robust risk (9) is given by

and again, the DRIG estimator is optimal for infinitesimal shifts γ ′ and suboptimal for larger γ ′ :

The above results follow by plugging M new with M := M DRIG into the proof of Corollary 3.2 in Appendix F.3..

D Empirical estimation of the worst-case robust predictor

In this section, we discuss how to compute the worst-case robust loss and its minimizer from finitesample multi-environment training data. We first describe the finite-sample setting and provide a high-level algorithm. We then discuss some parts of the algorithm in more detail. Finally, we show that the empirical worst-case robust loss is consistent under certain assumptions. For simplicity, in this section, we assume that M test = γSS ⊤ + γ ′ RR ⊤ , where γ, γ ′ ≥ 0 , SS ⊤ is a semi-orthogonal matrix satisfying range SS ⊤ ⊂ S and R is a semi-orthogonal matrix satisfying range R ⊂ S ⊥ . However, the results and strategies in this section can be easily applied to more general M test .

D.1 Computing the worst-case robust loss

Algorithm 1 Computation of the worst-case robust loss

1: Input: Multi-environment data D := ∪ e ∈E train D e , test shift strength γ > 0 , test shift directions M ∈ R d × d , causal parameter upper bound C > 0 .

2: Step 1: Estimate the training shift directions ˆ S ( D ) , its orthogonal complement ˆ S ⊥ ( D ) , and the identified causal parameter ˆ β S .

3: Step 2: Estimate the identified and non-identified test shift directions ˆ S , ˆ R and their projections ˆ S ˆ S ⊤ and ˆ R ˆ R ⊤ .

4: Step 3: Estimate the norm ˆ C ker of the non-identified causal parameter.

5: Step 4: Compute the worst-case robust loss function

6: Return: worst-case robust predictor and the estimated minimax "hardness" of the problem:

Training data. We observe data from m + 1 training environments indexed by E ∈ E train = { 0 , ..., m } , where E = 0 represents the reference environment. We impose a discrete probability distribution P E on the training environment E ∈ E train , resulting in the joint distribution ( X,Y,E ) ∼ P X,Y | E × P E . For each environment E = e , we observe the samples D e := { ( X e,i , Y e,i ) } n e i =1 , where ( X e,i , Y e,i ) are independent copies of ( X e , Y e ) ∼ P X,Y | E = e . Then, the resulting dataset is D := ∪ e ∈E train D e with n := n 0 + · · · + n m . Furthermore, for each environment E = e , we define the weights w e := n e /n .

Computation of the worst-case robust loss. In Algorithm 1, we present a high-level scheme for computing the worst-case robust loss from multi-environment data, which consists of multiple steps. First, nuisance parameters related to the training and test shift directions are estimated, which we describe in more detail below. Afterwards, the three terms of the loss are computed: the (squared) loss L ref ( β ; D 0 ) on the reference environment is computed as

The invariance penalty term L inv ( β ; ˆ β S , ˆ S ˆ S ⊤ , γ ) (which increasingly aligns any estimator β in the direction of the estimated invariant causal predictor ˆ β S as γ →∞ ) can be computed as following in the linear SCM setting:

Finally, the non-identifiability penalty term L id ( β ; ˆ C ker , ˆ R ˆ R ⊤ , γ ) can be computed as follows:

The non-identifiability term, with increasing γ , penalizes any predictor β towards zero on the subspace R of non-identified test shift directions. In total, the worst-case robust loss (in the linear SCM setting) equals

where we suppress dependence on C and γ and only leave the dependence on the nuisance parameters.

Choice/Estimation of nuisance parameters. We now provide more details on the empirical estimation of the nuisance parameters ˆ S , ˆ S, ˆ R , and ˆ β S .

The constant C corresponds to the upper bound on the norm of the true causal parameter β ⋆ . Thus, the practitioner chooses C in advance to ensure that (with high probability) ∥ β ⋆ ∥ 2 ≤ C .

The training shift directions ˆ S can be computed via

where for e ∈ E train , the matrix Cov( X e ) is the empirical covariance matrix estimated within the training environment E = e , and µ e ∈ R d is the empirical mean of the covariates within the training environment E = e . Additionally, we compute the orthogonal complement ˆ S ⊥ ( D ) of the space ˆ S ( D ) 12 .

The decomposition of the test shift directions M into identified and non-identified shift directions (and their corresponding projection matrices) can be computed as follows. Let Π ˆ S and Π ˆ S ⊥ denote the projection matrices on ˆ S ( D ) and ˆ S ⊥ ( D ) , respectively. Consider the singular value decompositions Π ˆ S M = U ˆ S Σ ˆ S V ⊤ ˆ S and Π ˆ S ⊥ M = U ˆ S ⊥ Σ ˆ S ⊥ V ⊤ ˆ S ⊥ Then, define

The subspaces range (Π ˆ S M ) and range (Π ˆ S ⊥ M ) are minimal subspaces contained in ˆ S and ˆ S ⊥ , respectively, such that range ( M ) ⊂ range (Π ˆ S M ) ⊕ range (Π ˆ S ⊥ M ) . The matrices ˆ S ˆ S ⊤ and ˆ R ˆ R ⊤ are their corresponding projection matrices.

The identified causal parameter ˆ β S (approximately) equals the true causal parameter β ⋆ on the space of training shift directions ˆ S . As conjectured in the anchor regression literature [41, 45, 23] (see, for example, the discussion right after Theorem 3.4 in [23] and Appendix H.3 therein) for γ →∞ , the estimators β γ anchor and β γ DRIG converge to the causal parameter β ⋆ on S . Thus, the identified causal parameter can be estimated as

for the setting of mean or mean+variance shifts, respectively.

D.2 Consistency of the worst-case robust predictor

For any estimator β ∈ R d and given the estimated nuisance parameters ˆ φ := ( ˆ S ˆ S ⊤ , ˆ R ˆ R ⊤ , ˆ β S ) , we define the sample worst-case robust risk as

Correspondingly, we define the estimator of the worst-case robust predictor by

where B ⊆ R d is some compact set whose interior contains the true parameter β rob Θ eq .

To show the consistency of (21), we first require consistency of the nuisance parameter estimators, which we state as an assumption.

Assumption D.1. The estimated nuisance parameters ˆ φ := ( ˆ S ˆ S ⊤ , ˆ R ˆ R ⊤ , ˆ β S ) are consistent, that is, for n →∞ ,

where for any matrix A ∈ R m × q , ∥ A ∥ F = √ tr ( A ⊤ A ) denotes the Frobenius norm, and SS ⊤ , RR ⊤ are the corresponding population projection matrices onto Π S M , Π S ⊥ M respectively.

Depending on the assumptions of the data-generating process, Assumption D.1 can be shown to hold. For example, in the anchor regression setting [41], the consistency of the projection matrices ˆ S ˆ S ⊤ , ˆ R ˆ R ⊤ , and Π ˆ S holds if the dimension of S is known (due to the mean shift structure). The proof relies on the Davis-Kahan theorem (see, for example, [56]) and the consistency of the covariance matrix estimator. Moreover, in the anchor regression setting, it is conjectured that the estimator β ∞ anchor converges to its population counterpart (as discussed right after Theorem 3.4 in [23] and Appendix H.3 therein) which implies that ˆ β S := Π ˆ S β ∞ anchor consistently estimates β S = Π S β ⋆ .

Under the assumption of the consistency of the nuisance parameter estimators, we can now show that (21) is a consistent estimator of the worst-case robust predictor.

Proposition 3. Consider the estimator ˆ β rob Θ eq of the worst-case robust predictor defined in (21) . Suppose the optimization problem is over a compact set B ⊆ R d whose interior contains the true minimizer β rob Θ eq . Moreover, suppose Assumption D.1 holds. Finally, assume that the covariance matrix E [ X 0 X ⊤ 0 ] ≻ 0 with bounded eigenvalues and E [ Y 2 0 ] < ∞ . Then, ˆ β rob Θ eq is consistent, i.e., as n, n 0 →∞ it holds that

D.3 Proof of Proposition 3

For ease of notation define β 0 := β rob Θ eq and ˆ β := ˆ β rob Θ eq . For any parameter of interest β ∈ B and nuisance parameters φ = ( P S , P R , b ) , define the function

Using (22), the robust identifiable risk and its sample version defined in (20) can be written, respectively as

Our goal is to show that ˆ β P → β 0

. First, we show that the minimum of the loss is well-separated. Lemma D.1. Suppose that E [ X 0 X ⊤ 0 ] ≻ 0 . Then, for all δ > 0 , it holds that

Fix δ > 0 . From the well-separation of the minimum from Lemma D.1, there exists ε > 0 such that

Therefore,

We now want to prove convergence the four terms in (24) and (25). For this, we use the following statements proved in Appendix D.4.

Lemma D.2. Suppose B ⊆ R d is a compact set. Moreover, assume that the covariance matrix E [ X 0 X ⊤ 0 ] ≻ 0 with bounded eigenvalues and E [ Y 2 0 ] < ∞ . Then, as n, n 0 →∞ it holds that

Lemma D.3. As n →∞ , it holds that

The two terms in (24) converge to 0 by Lemma D.2 and Lemma D.3, respectively. The first term in (25) equals 0 since ˆ β minimizes β ↦→L n ( β, ˆ φ ) . Finally, we observe that

since we have that

where the first term converges in probability by Lemma D.3, and the second term converges in probability by Lemma D.2. This implies that the second term in (25) converges to zero. Since δ > 0 was arbitrary, it follows that ˆ β P → β 0 .

D.4 Proof of auxiliary lemmas

D.4.1 Proof of Lemma D.1

By definition,

Since E [ X 0 X ⊤ 0 ] ≻ 0 , the first term is strongly convex in β . Moreover, the second and third terms are convex in β . Therefore, L ( β, φ 0 ) is strongly convex in β . Since L ( β, φ 0 ) is also continuous in β , it follows that there exists a unique global minimum. Let β 0 denote the global minimizer of L ( β, φ 0 ) . By the fact that L ( β 0 , φ 0 ) is a global minimum, and by definition of strong convexity, there exists a positive constant m> 0 such that, for all β ∈ B ,

Fix δ > 0 . Then, by (29), for all β ∈ B such that ∥ β -β 0 ∥ 2 > δ it holds that

Since the inequality holds for all β ∈ B such that ∥ β -β 0 ∥ 2 > δ , we conclude that

Since δ > 0 was arbitrary, the claim follows.

D.4.2 Proof of Lemma D.2

Recall that for any β ∈ B

To show the result, we must establish that the class of functions { g β,φ 0 : β ∈ B} is Glivenko-Cantelli. From [53], a set of sufficient conditions for being a Glivenko-Cantelli class is that (i) B is compact, (ii) β ↦→ g β,φ 0 ( x, y ) is continuous for every ( x, y ) , and (iii) β ↦→ g β,φ 0 is dominated by an integrable function. By assumption, (i) holds. Moreover, by (22), it follows that β ↦→ g β,φ 0 is continuous

for all ( x, y ) and thus (ii) holds. We now show that (iii) holds. Since B is compact we have that sup β ∈B ∥ β ∥ 2 = C 1 < ∞ . For fixed γ > 0 , and all ( x, y ) , we have that

where K < ∞ is a finite constant not depending on ( x, y ) . Furthermore, we have that

since E [ Y 2 ] < ∞ and E [ X 0 X ⊤ 0 ] has bounded eigenvalues by assumption. From (30) and (31), it follows that (iii) holds.

D.4.3 Proof of Lemma D.3

For fixed γ > 0 , we have that

We can upper bound (32) as follows,

where (34) follows from the Cauchy-Schwarz inequality and that ∥ A ∥ 2 ≤ ∥ A ∥ F , the constants C 1 , C 2 < ∞ in (35) follow from compactness of B , and the convergence in probability follows from Assumption D.1. Furthermore, we can upper bound (33) as follows,

By Assumption D.1, ( I ) converges in probability to zero. Regarding ( II ) , we have

where the inequality follows from Cauchy-Schwarz and that ∥ A ∥ 2 ≤ ∥ A ∥ F , and the convergence in probability follows from Assumption D.1 along with the compactness of B . It remains to upper bound ( III ) . We have that

The inequality in (36) follows from the compactness of B , the fact that ˆ R ˆ R ⊤ has bounded eigenvalues, and that | √ x - √ y | ≤ | x -y | 1 / 2 for all x, y ≥ 0 . The inequality in (37) follows from CauchySchwarz and that ∥ A ∥ 2 ≤ ∥ A ∥ F . The convergence in probability follows form Assumption D.1 and the compactness of B .

E Details on finite-sample experiments

In this section, we provide more details of the data generation for our synthetic finite-sample experiments as well as data processing for the real-world data experiments.

E.1 Synthetic experiments

For the synthetic experiments, we generate a random SCM which satisfies our assumptions. For d = 15 , we randomly sample the joint covariance Σ ⋆ of ( η, ξ ) , fixing its total variance and the eigenvalues. We consider 7 environments including the reference environment, and for each environment except the reference, we randomly generate mean shifts µ e of fixed norm 1 . Since we have 6 non-zero random Gaussian mean shifts, it holds a.s. that dim S = 6 . We then randomly generate an "initial guess" for β ⋆ ∈ R d of fixed norm C = 10 . Now, with respect to the space S of the identifiable directions induced by the mean shifts, we choose the most "adversarial" causal parameter β ⋆ adv which is equal to β ⋆ on S , but on S ⊥ has the opposite direction of the noise OLS estimator Σ ⋆ η -1 Σ ⋆ η,ξ . We ensure that ∥ β ⋆ adv ∥ 2 = C . Note that under the observed shifts, β ⋆ and β ⋆ adv are observationally equivalent. We complete β ⋆ adv to the set θ adv of observationally equivalent model parameters and generate the multi-environment training data according to θ adv and the collection of mean shifts.

For Figure 3 (left), we define the test shift upper bound as M anchor = γ 1 7 ∑ e µ e µ ⊤ e . We vary γ from 0 to 10 , and for each γ , we compute the oracle anchor regression estimator by minimizing the discrete anchor regression loss with the correct γ . Additionally, we compute the pooled OLS estimator and the worst-case robust predictor β rob Θ eq as described in Appendix D. Finally, we generate test data with a Gaussian additive shift A test ∼ N (0 , M anchor ) . We evaluate the loss of β OLS , β anchor and β rob Θ eq on this test environment and include the population lower bound.

For Figure 3 (right), we define the test shift upper bound as M new = γ 1 7 ∑ e µ e µ ⊤ e + γ ′ RR ⊤ , where R is a 2-dimensional subspace of the space S ⊥ . We fix the magnitude γ of the 'seen' test shift

directions at γ = 40 and set vary γ ′ from 0 to 2 to showcase the effect of small unseen shifts compared to large identified shifts. We compute the oracle anchor regression estimator by minimizing the discrete anchor regression loss. Additionally, we compute the pooled OLS estimator and the worst-case robust predictor β rob Θ eq as described in Appendix D, for which we use the oracle γ ′ , given M anchor and empirical estimates of the spaces S , S ⊥ , R .

Finally, we generate test data with a Gaussian additive shift A test ∼ N (0 , M new ) . We evaluate the loss of β OLS , β anchor and β rob Θ eq on this test environment, plot the resulting test losses for different estimators and include the population lower bound.

E.2 Real-world data experiments

0

2

4

0

1

2

3

4

ENSG00000089009

ENSG00000112306

Environment

ENSG00000089009

Obs.

(a)

0

2

4

0

1

2

3

4

ENSG00000089009

ENSG00000112306

Setting

test

train

(b)

0

2

4

0

1

2

3

4

ENSG00000089009

ENSG00000112306

Setting

test

train

(c)

̸

We consider the K562 dataset from [38] and perform the preprocessing as done in [13]. The resulting dataset consists of n = 162 , 751 single-cell observations over d = 622 genes collected from observational and several interventional environments. The interventional environments arise by knocking down a single gene at a time using the CRISPR interference method [37]. Following [44], we select only always-active genes in the observational setting, resulting in a smaller dataset of 28 genes. For each gene j = 1 , . . . , 28 , we set Y := X j as the target variable and select the three genes X k 1 , . . . , X k 3 most strongly correlated with Y (using Lasso), resulting in a prediction problem over Y, X k 1 , . . . , X k 3 . Given this prediction problem, we construct the training and test datasets as follows. Let O denote the 10,691 observations collected from the observational environment, and let I i denote the observations collected from the interventional environment where the gene k i was knocked down. We will denote by I i,s the s × 100 percent of datapoints in I i that are closest to the mean of gene k i in the observational environment O . For example, I i, 0 . 1 consists of the 10% of datapoints in I i closest to the observational mean of gene k i . Thus, the parameter s ∈ [0 , 1] acts as a proxy for the strength of the shift. Denote by I ∗ i,s a random sample of I i,s of a certain size. For each i ∈ { 1 , 2 , 3 } , we fit the methods on the training data D train i := O ∪ I ∗ i, 1 , with |I ∗ i, 1 | = 20 . Figure 5(a) illustrates an example of training data D train i . Having fitted the methods on D train i , we evaluate them on test datasets constructed as follows. For each shift strength s ∈ { 0 . 1 , . . . , 0 . 9 } and proportion π ∈ { 0 , . 33 , . 67 , 1 } , define the test dataset D π,s consisting of π observations from ∪ ℓ = i I ℓ,s and 1 -π (out-of-training) observations from I i,s . An example of a test dataset for different shift strengths s and previously unseen directions (i.e., π = 1 ) is shown in Figure 5(b-c). We compare our method Worst-case Rob., defined as the minimizer of the empirical worst-case robust risk (20), with anchor regression [41], invariant causal prediction (ICP) [35], Distributional Robustness via Invariant Gradients (DRIG) [45], and OLS (corresponding to vanilla ERM). We use the following parameters for Worst-case Rob.: γ = 50 , C ker = 1 . 0 , and M = Id. For anchor regression and DRIG, we select γ = 50 . For ICP, we set the significance level for the invariance tests to α = 0 . 05 .

These numerical experiments are computationally light and can be run in ≈ 5 minutes on a personal laptop. 13

F Proofs

F.1 Proof of Proposition 1

For every environment e ∈ E train , we observe the first moments E ( X e ) and E ( Y e ) , and second moments E ( X e X ⊤ e ) , E ( Y 2 e ) and E ( X e Y e ) . Since it holds by assumption that µ 0 = 0 and Σ 0 = 0 , we have that E ( X 0 X ⊤ 0 ) = Σ ⋆ η , and so we can identify Σ ⋆ η uniquely. Furthermore, it holds that

By taking the difference between Equation (39) and Equation (38), we can identify (Σ e + µ e µ ⊤ e ) β ⋆ . Thus, the parameter β ⋆ is identifiable on the subspace S defined in Equation (12) and is not identifiable on its orthogonal complement S ⊥ . Thus, for any vector α ∈ S ⊥ , the vector β = β ⋆ + α is consistent with the data-generating process. It remains to compute the covariance parameters induced by an arbitrary ˜ β := β ⋆ + α , for α ∈ S ⊥ . For every environment e ∈ E train , the second mixed moment between X e and Y e has to satisfy the following equality

from which it follows that ˜ Σ η,ξ := Σ ⋆ η,ξ -Σ ⋆ η α . By computing E ( Y 2 e ) and inserting ˜ β = β ⋆ + α and ˜ Σ η,ξ , we similarly obtain

Thus, we obtain the following set of observationally equivalent model parameters consistent with P θ ⋆ , E train :

Since the observationally equivalent set is identifiable from the training distribution, but model parameters β ⋆ , Σ ⋆ η,ξ , ( σ ⋆ ξ ) 2 are not, it is helpful to re-express the observationally equivalent set through identifiable quantities. For this, we note that the "identifiable linear predictor" β S = β ⋆ -β S ⊥ induces an observationally equivalent model given by

From this reparameterization, we infer the final form of the observationally equivalent set :

Therefore, Equation (13) follows. To find the robust predictor β rob , we write down the robust loss with respect to M test and any θ α from the observationally equivalent set :

inserting α ∈ S ⊥ and rearranging, Equation (14) follows.

F.2 Proof of Theorem 3.1

We structure the proof as follows: first, we quantify the non-identifiability of the robust risk by explicitly computing its supremum over the observationally equivalent set of the model parameters (referred to as the worst-case robust risk). Second, we derive a lower bound for the worst-case robust risk by considering two cases depending on how a predictor β interacts with the possible test shifts M test .

Computation of the worst-case robust risk. For any model-generating parameter θ = ( β, Σ) it holds that the robust risk of the model Equation (3) under test shifts M test ⪰ 0 is given by

We recall that the observationally equivalent set of model parameters after observing the multienvironment training data Equation (3) is given by

where S is the span of identified directions defined in Equation (12). Moreover, we recall that by Assumption 3.2, for any causal parameter β it should hold that ∥ β ∥ 2 = ∥ β S + α ∥ 2 ≤ C , which translates into the following constraint for the parameter α :

Inserting Equation (40) in Equation (9), we obtain

where θ α is a short notation for ( β S + α, Σ ⋆ η , Σ S η,ξ -Σ ⋆ η α, ( σ S ξ ) 2 -2 α ⊤ Σ S η,ξ + α ⊤ Σ ⋆ η α ) . We now compute the supremum explicitly in case M test has the form M test = γM seen + γ ′ RR ⊤ , where M seen is a PSD matrix with range M ⊆ S and R is a semi-orthogonal matrix with range R ⊆ S ⊥ . For any α ∈ S ⊥ , we write down the robust loss as

The first term is the robust risk of β under test shift M test and the identified model-generating parameter θ S , thus it does not depend on α . By the structure of M test , we obtain that

If γ ′ = 0 , i.e., the test shifts consist only of the identified directions, we have f ( α ) = 0 , independently of α , and thus

This implies the first statement of the theorem.

̸

We now consider the case where R = 0 , i.e., RR ⊤ is a non-degenerate projection. Our goal is to maximize f ( α ) subject to constraints α ∈ S ⊥ , ∥ α ∥ 2 ≤ C ker . Let ˜ R be an orthonormal extension of R such that range ( R | ˜ R ) = S ⊥ . Then, we can parameterize α ∈ S ⊥ as α = ( R | ˜ R )( w ˜ w ) and the corresponding Lagrangian reads

Differentiating with respect to w, ˜ w yields

After differentiating w.r.t. λ , we obtain γ ′ γ ′ -λ = ± C ker ∥ R ⊤ β ∥ 2 . By inserting in the objective function and comparing, we obtain the value of the worst-case robust risk :

Putting together the two cases and simplifying, we obtain

where R rob ( β ; θ S , γM seen ) is the robust risk of the estimator β w.r.t. the "identified" test shift γM and the identified model parameter θ S , whereas R 0 ( β, θ S ) is the risk of β on the reference environment e = 0 .

Derivation of the lower bound for the worst-case robust risk. Now that we have explicitly computed the worst-case robust risk, we devote ourselves to the computation of the lower bound for its best possible value

̸

In this part, we will only consider the case R = 0 , since the case R = 0 corresponds to the (discrete) anchor regression-like setting, where both the robust risk and its minimizer are uniquely identifiable, and computable from training data. We will distinguish between two cases.

Case 1: ∥ R ⊤ β ∥ 2 = 0 . In this case, β is fully located in the orthogonal complement of R , which consists of S and ˜ R (the orthogonal complement or R in S ⊥ ). We will denote (the basis of) this subspace by S tot = S ⊕ ˜ R . Thus, S tot is the "total" stable subspace consisting of identified directions in S and non-identified, but unperturbed directions ˜ R . We will parameterize β as β = S tot w . Thus, we are looking to solve the optimization problem

Setting the gradient to zero yields the asymptotic worst-case robust estimator

which corresponds to the loss value of

As we observe, this quantity grows linearly in γ ′ . However, as γ →∞ , the quantity saturates and is upper-bounded by ( σ S ξ ) 2 .

̸

Case 2: ∥ R ⊤ β ∥ 2 = 0 . Since for ∥ R ⊤ β ∥ 2 = 0 , the objective function is differentiable, we compute its gradient to be

̸

This equation is, in general, not solvable w.r.t. β in closed form. Instead, we provide the limit of the optimal value of the function when the strength of the unseen shifts is small, i.e. γ ′ → 0 . We know that for γ ′ = 0 , the minimizer of the worst-case robust risk is given by the anchor estimator

Instead, we lower bound the non-differentiable term 2 γ ′ C ker ∥ R ⊤ β ∥ by the scalar product 2 γ ′ C ker ⟨ R ⊤ β, R ⊤ β anchor ⟩ / ∥ β anchor ∥ and expect it to be tight for small γ ′ . After inserting this lower bound in Equation (41) we obtain the minimizer of the lower bound of form

We can now lower bound ∥ RR ⊤ β LB ∥ as

Thus, the γ ′ -rate of the worst-case robust risk of β LB is at least γ ′ ( C ker + ∥ RR ⊤ (Σ ⋆ η + γM ) -1 Σ ⋆ η,ξ ∥ ) 2 + O ( γ ′ 2 ) , from which the claim for small γ ′ follows. For Section 3.2, the lower bound directly implies optimality of the worst-case robust risk of the anchor estimator when the strength of the unseen shifts γ ′ is small. Additionally. if γ = 0 , i.e. only unseen test shifts occur, we conclude that the OLS and anchor estimators have the same rates.

Lower bound γ ′ th for γ ′ . Finally, we want to derive a lower bound on the shift strength γ ′ such that for all γ ′ ≥ γ ′ th Case 1 of our proof is valid, i.e. it holds that β rob Θ eq is given by the closed form "abstaining" estimator (44). For this, we find γ ′ th such that for all γ ′ ≥ γ ′ th zero is contained in the subdifferential of R rob ( β rob Θ eq ; Θ eq , M test ) at β rob Θ eq . Then the KKT conditions are met, and β rob Θ eq is the unique minimizer of the worst-case robust risk due to strong convexity of the objective. We compute the subdifferential to be

Since β rob Θ eq is the minimizer of R rob ( β ; θ S , γM seen ) under the constraint R ⊤ β = 0 , the gradient is zero in R ⊥ and it remains to show that

Via an upper bound on the projected gradient, we derive the stricter condition

where κ (Σ ⋆ η ) is the condition number of the covariance matrix.

F.3 Proof of Corollary 3.2

To obtain a new formulation for the worst-case robust risk, we start with (46) and expand

where we have used that the pooled second moment of X equals to Σ ⋆ η + ∑ e w e ( µ e µ ⊤ e ) = Σ ⋆ η + γM anchor -( γ -1) M anchor . This reformulation shows that the worst-case robust risk is equal to the anchor population loss (cf. [41]) with an additional non-identifiability penalty term γ ′ ( C ker + ∥ R ⊤ β ∥ 2 ) 2 .

We now want to evaluate the rates of the anchor and OLS estimators in terms of the magnitude γ ′ of unseen shift directions. We observe that only the non-identifiability term depends on γ ′ , whereas the second term only depends on γ . First, we compute the closed-form anchor regression estimator, which reads

Since β OLS equals to the anchor estimator with γ = 1 , we obtain

The claim of the corollary now follows by computing ∥ RR ⊤ β anchor ∥ and ∥ RR ⊤ β OLS ∥ and observing that the rest of the terms is constant on γ ′ .

Comparing to the lower bound (45) for the minimax quantity for the case of γ ′ → 0 , we observe that the anchor estimator is optimal (achieves the minimax rate) in the limit γ ′ → 0 . Additionally, if γ = 0 (only new shifts occur during test time), anchor and OLS have identical rates in γ ′ and, in particular, OLS (corresponding to vanilla empirical risk minimization) is minimax-optimal in the limit of small unseen shifts.

or

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: We state our contributions relative to prior work in the abstract, in Section 1, and in Appendix A.

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: In the abstract and in Section 1, we highlight the setting that we consider. We explicitly describe the assumptions in Section 2 and summarize the limitations in Section 5. 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: Appendix F contains proofs of all results appearing in the main paper. Appendix B, Appendix C, and Appendix D are self-contained and contain derivations and proof of the results.

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: Appendix E provides all the necessary information to reproduce the experimental results presented in Section 3.2. We provide details on empirical estimation of the proposed loss function in Appendix D.

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: [No]

Justification: While we do not provide the code, the paper provides all necessary information on reproducing the experiment in Appendices D and E.

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: We specify all details to understand the experimental results in Section 3.2 and Appendix E.

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:

In the numerical experiment, shown in Figure 3, we provide the average test MSE and its 5% and 95%-quantiles over 100 repetitions for each method.

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: The numerical experiment described in Section 3.2 is computationally very light and can be run on a personal laptop in a few minutes. We describe this in Appendix E.

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: We have read the NeurIPS Code of Ethics and confirm that our work conforms to it in all aspects.

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: [Yes]

Justification: Even if our work addresses the theoretical limits of distributional robustness, we mention in the abstract and in Section 1 that the topic of distributional robustness is central to safety-critical applications.

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 work develops a theoretical framework and considers synthetic experiments. Therefore, explicit safeguards do not seem applicable at this stage.

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: In the numerical experiment in Section 4, we cite the existing work that we compare to our framework and the dataset used. In running the numerical experiment, we reimplemented all the methods (including existing ones) for ease of comparison.

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: At this stage, 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 experiments or 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 experiments or 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:
comments: SUMMARY:
This paper proposed a general framework of partially identifiable robustness to evaluate robustess in scenarios where the training distributions are not heterogeneous enough to identify the robust risk. They define 'the identifiable robust risk' and its correspondig minimax quantity. They show previous approaches achieve suboptimal robustness in this scenario.  Finally, they propose the empirical minimizer of the identifiable robust risk and show that it outperforms existing methods in finite-sample experiments.

STRENGTHS:
The paper is crealy written. The idea of establishing partially identifiability to fill the gap between an "all-or-nothing" view on robustness in the SCM framework is interesting.

WEAKNESSES:
I am not an expert in this field so I cannot point out many weaknesses. Maybe one thing is the lack of motivation from real-world application, such as in what applications or when the proposed framework is usefull in practice?

QUESTIONS:
1. On line 135-137, the $M_{\text{test}}$ is defined as  $M_{\text{test}}=\gamma\Pi_{M}$. What's the reasoning behind this choice?

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the optimal minimax risk of a robust predictor when the robustness set is partially observable, under a structural causal model with hidden confounders. By decomposing the test covariance matrix of latent parameters into a component spanned by the training distributions and the orthogonal component, it is shown that the robust predictor is only identifiable if the test shifts are in the direction of the training shifts. While for unobserved test shifts, the best achievable minimax risk grows linearly with shift scale. The theory is applied to show sub-optimality of OLS and anchor regression with partially observable test shifts.\n\nSTRENGTHS:\nI am not familiar with the field of causal inference, but it appears from the paper that it has made contributions as the first result for distributional robustness or invariant causal prediction when the robustness set is partially identifiable. Particularly, it shows that infinite robustness is impossible under this scenario, and finite robustness methods can show performances reduced to ERM. This echoes empirical evidence and provides a possible theoretical explanation for reported failure of distributional robustness methods under wild environments.\n\nWEAKNESSES:\nThe structural causal model in Eq 2 and the resulting explicit solution in Eq 6 shows that the model is biased even without distribution shifts. To see this, take $\\gamma =0$ in Eq 6 and the predictor does not reduce to $\\beta^\\star$. The model is unbiased only when the cross covariance between $\\eta$ and $\\xi$ vanishes, which implies no hidden confounder and only covariate shift. In contrast, the classic approach for invariant causal prediction [1] produces unbiased estimation beyond covariate shift.\n\n[1] Arjovsky, M., Bottou, L., Gulrajani, I., & Lopez-Paz, D. (2019). Invariant risk minimization. arXiv preprint arXiv:1907.02893.\n\nQUESTIONS:\n1. Typo: L137, the sub-space.\n2. In Fig 1, bidirectional edges don't make sense for causal graphs, and it's not explained either. Moreover, is there exact equivalence between Fig 1 and the SCM in Eq 2?\n\nLIMITATIONS:\nThe author has addressed limitations with regarding to the model, e.g., linear structure and additive noises.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new framework for distributionally robustness under the linear causal setting. Specifically, the authors minimize the so-called identifiable robust risk, which corresponds to the maximum of the robust risk for parameters in the observationally equivalent set. Under such partially identifiable robustness framework, they discuss the lower bound of the risk and show that some estimated identifiable robust predictor can achieve that lower bound risk, where the corresponding empirical version can approximate the lower bound risk well. They also validate the better performance with finite sample numerical experiments.\n\nSTRENGTHS:\n-\tThe motivating illustration and results are very clear;\n-\tThe direction the authors study is interesting in bridging the structure and non-structure of distributional robustness.\n\nWEAKNESSES:\n-\tSome paragraph in the training and testing data part can be reorganized as some formal assumptions, such as the assumption of M_{test} and linear assumption as well as the discussion in Section 3.2 in terms of the definition of S and M. This gives a clear picture on what some key conditions (and possible relaxations) are in this paper.\n-\tThere are a few missing details (see Questions as follows), which I hope the authors can explain them a bit more.\n\nQUESTIONS:\n-\tCan the authors highlight a bit more on the (approximated) real-world example when the predictor is partially identifiable? Since it may be a bit unfamiliar to the general DRO audience for these notations and their relevance in practice.\n-\tCan the authors elaborate more on how the empirical estimations for the space of the training shift direction $\\hat S$ and $\\hat R$ are computed, as compared with the current version in Appendix D? From my preliminary understanding, this space determination is important for subsequent estimation.\n-\tThe potential utility of **active intervention selection** is also interesting, and it seems aligned with some recent relevant work on distribution shifts. For example, in causal explanation [1] and incorporating specific features in improving distributional robustness [2]. I am wondering if the authors can provide more details on how their own intervention is and the connections with these existing literature, which can be potentially incorporated in the main body and appendix.\n\n[1] Quintas-Martinez, Victor, et al. \"Multiply-Robust Causal Change Attribution.\" arXiv preprint arXiv:2404.08839 (2024).\n\n[2] Liu, Jiashuo, et al. \"On the Need for a Modeling Language Describing Distribution Shifts: Illustrations on Tabular Datasets.\" arXiv preprint arXiv:2307.05284 (2023).\n\nLIMITATIONS:\nThe authors discuss the main limitation of this paper, which I think is reasonable compared with existing literature.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies a linear Structural Causal Model (SCM) for prediction under distribution shifts due to an an additive term to covariates that changes at test time, and an unobserved confounder between the covariates and the label. The key difference between the proposed analysis and those presented in other works (e.g. those in anchor regression, invariant causal prediction etc.) is that the shift is bounded in its strength. Another difference is that the optimal robust predictor w.r.t the entire uncertainty set is not identifiable. Hence, instead of the optimal robust error, the paper studies the optimal error that can be achieved from observable data. This corresponds to the robust error over an uncertainty set that is generated by all SCM parameters that could have produced the training data, which is in general a larger set than the uncertainty set we are interested in (i.e. of bounded additive shifts to the covariates).\n\nOnce the problems is set up, a parametric form for the set of parameters that can generate the training data is derived, along with the corresponding set of possible robust predictors. Then under a mild assumption on boundedness of the ground truth regression parameters, the following are derived: 1) a closed form is derived for the robust loss in the case where test distributions only induce shifts in directions that have been observed at training. The loss grows linearly with the strength of the shift; 2) a lower bound on the achievable risk from observed data, which is tight (and thus can be learned from observed data) for large enough shifts.\n\nFurther analysis and simulations with Gaussian data are done to compare the risk of the derived estimator with that of ERM and anchor regression (which does not exploit the boundedness of shifts to achieve better prediction). The results verify that the risk obtained by the proposed empirical estimator are close the lower bound given in the theorem, and outperform the two baselines as shifts grow large.\n\nSTRENGTHS:\n* Overall I enjoyed reading the paper, it is clear and written with care. Generally, I also like the direction of formalizing bounded shifts and unidentifiable settings in detail. Finally, the analysis is well performed and easy to follow.\n\n* The work is original in formally analyzing bounded distribution shifts, where even in population the optimal robust predictor might still have non-zero projection on directions that shift at test time. \\\nSmall note on this: I think that for finite sample guarantees boundedness assumptions on the strength of the shift must be made, and they are made in works that give sample complexity results. I believe the reason is that with unbounded shifts, any small weight on a shifting feature can be magnified unboundedly to yield a large robust error. With finite samples, it is usually impossible to guarantee strictly zero weights on the shifting directions. Therefore it might be worthwhile clarifying that considering the population loss is an important component for the analysis.\n\n* Beyond the points mentioned above, the bounds on the robust risk, closed form solutions, and the formalism used in the work, can be useful for future theoretical analyses. In terms of practical significance, the approach derived from these results to explicitly account for bounded shifts might be useful in the future, if it is generalized beyond linear models.\n\nWEAKNESSES:\n* As mentioned above and the authors mention in describing the limitations of their work, it is restricted to linear models. Intuitively, a considerable challenge in learning robust models is to identify the ``directions” that shift between domain. Under linear models it is rather straightforward, and the method proposed in the paper depends on linearity of the model in order to find these directions. This is unlike some other methods in the domain generalization literature, where certain formal results are provided on linear models but the methods can easily be tested in the non-linear case.\n*  Even within the realm of linear models, I am not throughly convinced that the method is useful in real world problems. To make this more convincing, it might have been nice to run simulations over non-synthetic data and some more variants of methods. E.g. using several methods that learn invariant models when hyper-parameter tuning is performed with the objective proposed in this work may be of interest. That is to see whether boundedness of the shift should be taken into account during training, or is it enough to simply use it for model selection. Yet the most significant drawback is still of course, the limitation to linear models.\n\nSome other/smaller comments:\n* In line 60 prior work is cited to claim \"that even minor violations of the identifiability assumptions can cause invariance based methods to perform equally or worse than ERM\". However, I am not sure that the failures portrayed in these works are strictly due to identifiability violations. At least not the ones alluded to in this paper, which is the lack of heterogeneity in the training environments. In Kamath et al. 21, the predictor is identifiable but the problem is specifically in the IRMv1 objective, see also [1] who point out an objective that solves this issue. In Rosenfeld et al. 20, only one failure is due to not having enough environments (and arguably, this is since they consider the uncertainty set that includes shifts in all directions. I will also touch on this in the next minor comment), and the second one is due to non-linearity.\n* Regarding analysis under lack of heterogeneity or unidentifiable optimal robust predictor: I think that the part of the analysis that touches upon the insufficient heterogeneity of the environments in order to identify the robustness set/robust predictor can be slightly reframed. If I am not mistaken, even in works that give results about identifiable robustness sets (e.g. when the number of environments is linear in the number of shifting features, or stronger ones like [2]), it is most likely possible to draw guarantees about robust risks, but only with respect to a smaller uncertainty set which is restricted to dimensions that shifted between the training environments. That is since the methods still enforce constraints which restrict weights in certain directions. While it is true that most of these prior works did not consider unidentifiable uncertainty sets, it might be related to choice of presentation, and not strictly because the methods are technically limited in that sense. Further, in lines 168-169 it is claimed that “prior work only considers scenarios where the robustness set and hence also the robust prediction model are still identifiable”. I am not sure this is entirely true. There are different formalisms that were considered for spurious correlations, such as those in [3]. In their setting, when the association is not what they call “purely spurious”, then the robust predictor cannot necessarily be recovered. Also [4] study a similar setting where due to similar reasons no guarantee can be given on identifying the robust predictor, but only a lower bound on the error is given.\n* It might be worth mentioning in lines 112-113 that in case $H$ in figure 1 is a selection variable (I assume it can be since both edges are bi-directional) then the shift reduces to covariate shift, if I am not mistaken.\n* Personally, I found the notation in the introduction which uses $\\mathrm{shift}\\in{\\mathrm{shift set}}$ etc. somewhat redundant and confusing. The more formal notation in section 2 onwards was much easier to understand, so in my opinion it might be worth considering to go straight ahead into that notation.\n\n[1] Wald Y, Feder A, Greenfeld D, Shalit U. On calibration and out-of-domain generalization. Advances in neural information processing systems. 2021 Dec 6;34:2215-27.\n\n[2] Chen Y, Rosenfeld E, Sellke M, Ma T, Risteski A. Iterative feature matching: Toward provable domain generalization with logarithmic environments. Advances in Neural Information Processing Systems. 2022 Dec 6;35:1725-36.\n\n[3] Veitch V, D'Amour A, Yadlowsky S, Eisenstein J. Counterfactual invariance to spurious correlations in text classification. Advances in neural information processing systems. 2021 Dec 6;34:16196-208.\n\n[4] Puli AM, Zhang LH, Oermann EK, Ranganath R. Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations. In International Conference on Learning Representations.\n\nQUESTIONS:\n* It could’ve been nice to derive upper bounds on the robust loss in addition to the lower bounds proved in the theorem. That is, assuming these bounds can be calculated from observable data, I’d imagine that most practitioners would be more interested in an upper bound as it gives a strong guarantee on the risk worst possible risk. Do you think this is a reasonable goal, and do you perhaps have any insights on this? \n* Could you perhaps discuss the differences between the empirical optimization problem you derived in the paper and those derived in prior work? For instance, looking at eq. 19 in appendix D, it seems like there is a term that penalizes $\\|\\| \\hat{S}^\\top(\\hat{\\beta}^{\\mathcal{S}} - \\hat{\\beta}) \\|\\|$ which I read as the weights in the “invariant directions” should be the same across the environments. This seems quite similar to ICP, IRM etc. Then the other term limits the projection on the shifting directions, but it might be useful for readers to have a short discussion on the conceptual similarities between this and penalized versions of other methods. Also, there might be a typo in that equation, $\\mathrm{arg}\\min_{\\beta\\in{\\mathbb{R}^d}}$ should be $\\mathrm{arg}\\min_{\\hat{\\beta}\\in{\\mathbb{R}^d}}$ instead.\n\nLIMITATIONS:\nThe authors have adequately discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-29"}