Achievable distributional robustness when the robust risk is only partially identified

We are glad that the reviewer appreciated the main idea of our paper. In the following, we hope to clarify the real-world motivation for our study and some theoretical assumptions.

On real-world applications The main motivation for our paper was expanding the existing assumptions in distribution shift literature to allow for a statement in a more realistic setting, particularly in safety-critical applications. We wanted to account for the fact that distribution shifts are mostly bounded in the real world and occur in “realistic” directions. Contrary to many existing results, in the real world, we rarely have access to enough training environments to achieve point identification of the “best” robust predictor. Instead, we aim to study the “best” model if only a few training environments are given, but the distribution shift (in potentially new directions) is not too large. In the following, we provide a toy example, motivated by medical applications, that illustrates our setting.

Suppose we are conducting a long-term medical study, where data is collected from the same group of patients over the years to predict a health parameter $Y$, e.g., cholesterol level in the blood, from a group of covariates $X = (X_1, …, X_n)$ such as age, blood pressure, physical activity, resting pulse, BMI, etc. We are given data $(X^e, Y^e)$ for $e \in E_{train}$ from multiple past studies, where $\mathcal{E}_{train} = \{2010, 2015\}$ are the years in which the studies were conducted. We assume that the data $(X, Y)$ are generated by an underlying causal model in which not every variable is observed (confounded setting). In this example, we can identify the causal effect of the covariate $X_1$, age, on the cholesterol level since age has a mean shift of 5 years across the studies. Suppose, however, that the distribution of $X_2$, physical activity, remains relatively stable across the studies. The causal effect of $X_2$ on $Y$ is then not identifiable. We now want to train a model that generalizes best on the data $X^{2020}$ collected in $2020$. The age variable shifts again by 5 years; however, the physical activity variable now also shifts a bit (e.g., due to COVID-19), which is a shift we have not observed previously. In this situation, both the causal parameter and the robust predictor are only partially identifiable.

Assumption on $M_{test}$: Often, practitioners will have some information on the strength and direction of the distribution shift (for instance, we might know that only the resting pulse will shift across studies, and at most by 20). This knowledge is what we tried to formalize in our assumption on $M_{test}$ (lines 134-137). The parameter gamma corresponds to the maximum strength of the shift, whereas the subspace $M$ corresponds to the expected direction of the shift (if one doesn’t know, one can take $M$ to be the whole covariate space). Equations (3) and (4) bound the second moment of the test distribution shift by $\gamma \Pi_{M}$, effectively bounding the mean and variance of the shift, as well as constraining it geometrically to the subspace $M$.

We thank the reviewer for appreciating the contribution of the paper and precisely summarizing its main idea.

Bias of the structural causal model The reviewer is correct that the causal effect estimate would be biased when the cross-correlation between the noises $\xi$ and $\eta$ is not 0. We would now like to argue that this is a feature rather than a bug. First, we would like to clarify that this paper considers the prediction performance, or MSE, during test time. Therefore, having a biased causal effect estimate does not necessarily affect the test performance, which is still optimal for $\gamma=0$ (no shift). Furthermore, allowing for latent confounding (expressed via cross-correlation of $\eta$ and $\xi$) only makes our model more general and the problem much more challenging to solve. The presence of confounding creates a tradeoff between predictive power and robustness (see, e.g., [2]). Furthermore, we would like to highlight that even IRM might not always produce unbiased estimates of the causal coefficient. For example, in [1], the synthetic data experiment shows that IRM and OLS have a bias of similar magnitude in a confounded setting with homoskedastic noise.

Q2 (bidirectional edges) Normally, a bidirectional edge $X \leftrightarrow Y$ between two nodes indicates latent confounding. In Figure 1, we explicitly list the latent confounder $H$, resulting in the notation $X \leftrightarrow H \leftrightarrow Y$, to express that the latent confounder can be both an ancestor and descendant of $X$ and/or $Y$. More precisely, we allow for the following scenarios: $X \rightarrow H \rightarrow Y$, $X \leftarrow H \rightarrow Y$, $X \rightarrow H \leftarrow Y$ (the fourth one is excluded due to the acyclicity constraint).

Q2 (equivalence of Figure 1 and the SCM) To see the equivalence between the SCM and DAG, it is helpful to consider the three configurations discussed before. When $X \leftarrow H \rightarrow Y$, the causal coefficient $\beta^\star$ equals the weight of the path $X \rightarrow Y$. In this setting, $\eta$ and $\xi$ are correlated via $H$. When $X \rightarrow H \rightarrow Y$, the causal coefficient $\beta^\star$ equals the sum of the weights along the paths $X \rightarrow Y$ and $X \rightarrow H \rightarrow Y$. In this setting, $\eta$ and $\xi$ are independent. When $X \rightarrow H \leftarrow Y$, the causal coefficient $\beta^\star$ equals the weight of the path $X \rightarrow Y$. In this setting, $\eta$ and $\xi$ are independent.

[1]: Arjovsky M, Bottou L, Gulrajani I, Lopez-Paz D. Invariant risk minimization. arXiv preprint arXiv:1907.02893. 2019 Jul 5.

[2]: Rothenhäusler D, Meinshausen N, Bühlmann P, Peters J. Anchor regression: Heterogeneous data meet causality. Journal of the Royal Statistical Society Series B: Statistical Methodology. 2021 Apr;83(2):215-46.

We thank the reviewer for pointing out the novelty of the direction of our study. We are happy to fill in on the missing details below:

Q1 (real-world example of partial identifiability) Indeed we can give a toy example that illustrates our abstract setting in Section 3.1: suppose that we are conducting a long-term medical study, where data is collected from the same group of patients over the years to predict a health parameter $Y$, e.g., cholesterol level in blood, from a group of covariates $X = (X_1, …, X_n)$ such as age, blood pressure, physical activity, resting pulse, BMI etc. We assume that the data $(X, Y)$ are generated by an underlying causal model, in which not every variable is observed (i.e. there might exist latent confounding). During training time, we are given data $(X^e, Y^e)$ for $e \in E_{train}$ from multiple past studies, where $E_{train} = \{2010, 2015\}$ are the years in which the studies were conducted. We now want to train a model which generalizes best on data $X^{2020}$ collected in 2020.

In such a setting, we can expect that the distribution shift between the data in 2010, 2015 and 2020 is not entirely arbitrary: for example, the covariates such as age ($X_1$) might shift similarly between 2010 to 2015 and 2015 to 2020, whereas other covariates such as physical activity ($X_2$), BMI ($X_3$) or resting pulse ($X_4$) may exhibit unseen, but bounded shifts (e.g. due to a disease outbreak such as Covid). This is an example of a structured distribution shift. We observe that we can identify the causal parameter $\beta^\star$ in the direction of $X_1$, since there is a shift of age across the training environments. However, the causal parameter is not identifiable for $X_2$ if the distribution of physical activity has not shifted in previous studies. This renders the causal parameter, and (since the test data shifts in the direction of $X_2$) the robust risk of the problem, partially identifiable. In this setting, only $X_1$ is guaranteed to give an invariant prediction, but an estimator which just uses the age might not have enough predictive power. Instead, we suggest to also utilize the spurious correlations between $X_2,...,X_4$ and $Y$ to predict the target variable, but penalize the predictor in their directions based on the strength of the expected test shift (e.g., we do not expect the average resting pulse to shift by more than 10).

Q2 (empirical estimation of S and R) Indeed, empirical estimation of the subspaces $S$ and $R$ is an important part of the practical application of our method. We will add to the manuscript that the space $S$ can be computed from the second moments of the observed distribution via $S := \text{span} \cup_{e \in E_{train}} [(\Sigma_e - \Sigma_0) + (\mu_e - \mu_0)(\mu_e - \mu_0)^\top]$, and can be estimated via empirical means and covariances of the training distributions. Results on the consistency of estimation of the eigenvectors when the dimension is fixed and eigenvalues are finite are given, for example, by [1]. For the test shift, there’s either prior information on their direction available via a subspace $M$, which we can incorporate by setting $R = \Pi_{S^{\perp}} M$, or, if not, we may take the conservative estimate of $R$ being the orthogonal complement of $S$. The downside of this choice is that the robustness requirements might be more restrictive than necessary.

Q3a (active intervention selection) Thanks a lot for this comment. Although we plan to have a longer elaboration of this topic in the full manuscript, we briefly discuss here how our approach might be utilized for active intervention selection. In Equation (31), the identifiable robust risk is given by a supremum over the possible true causal parameters from the observationally equivalent set. The argmax of the supremum (which depends on the estimator) utilizes partial identifiability to find the “most adversarial” direction for the test shift, along which the given estimator will suffer the highest test risk. One can actively sample the next dataset by performing an intervention along this direction, which will maximally decrease the identifiable robust risk (31) of a given estimator (e.g., OLS). This procedure can then be repeated.

Q3b connections to listed literature In [2] (causal change attribution), the main difference seems to be that the causal graph is known and the objective is, instead of robust prediction, causal change attribution. A conceptual similarity seems to be that [2] allows for recovery of the target parameter under partial misspecification. However, in our framework this would mean that although some components of the model might be misspecified (or, in our words, underidentified), the target parameter is still identifiable – thus, we would call the setting in [2] identifiable under the listed assumptions, similarly to the case of anchor regression. [3] is fairly related to our work motivation-wise: our model allows for $Y|X$-shifts in addition to $X$-shifts, consistent with the observation that $Y|X$-shifts frequently occur in tabular data and cannot be neglected.

[1]: Anderson TW. Asymptotic theory for principal component analysis. The Annals of Mathematical Statistics. 1963 Mar;34(1):122-48.

[2]: Quintas-Martinez V, Bahadori MT, Santiago E, Mu J, Janzing D, Heckerman D. Multiply-Robust Causal Change Attribution. arXiv preprint arXiv:2404.08839. 2024 Apr 12.

[3]: Liu J, Wang T, Cui P, Namkoong H. On the need for a language describing distribution shifts: Illustrations on tabular datasets. Advances in Neural Information Processing Systems. 2024 Feb 13;36.

We thank the reviewer for appreciating the overall direction and clarity of our paper, as well as originality of our work. We are grateful for the constructive comments and will incorporate the minor points in the revised version of the paper. We now respond to some of the major points and questions.

Linearity of the model We acknowledge that understanding which shift directions are “learned” and which are “unobserved” during training time for more general, nonlinear models is an important direction for future work. We further recognize the linearity of our setting as a limitation. We would like to clarify that the main contribution of our paper is not primarily to propose a new method for domain generalization (that can then also be tested in non-linear settings), but to introduce and make a first step towards quantifying the limits of robustness in a partially identifiable setting. On this note, we expect that the results and intuition developed in this paper for the linear case can be utilized beyond linear models, since realistic distribution shifts can be often reduced to linear shifts in a lower-dimensional latent space via a suitable parametric or non-linear map [5, 6].

Evaluation of the method in real-world settings Even though the method is not necessarily the focus of our paper, we agree that our work could benefit from a more thorough evaluation on real-world data. To this end, in the attached PDF, we present preliminary results of experiments on real-world single-cell gene expression data [7]. The dataset consists of single-cell observations of over 622 genes from observational and several interventional environments, arising through a knockdown of each unique gene. We pre-select always-active genes in the observational setting, resulting in a smaller dataset of 28 genes. We measure the performance of our method, the identifiable robust predictor (Rob-ID), compared to other algorithms [10, 11, 12], as follows. We select each gene once as the target variable $Y$ and select the three genes most strongly correlated with $Y$ (using Lasso), resulting in a smaller dataset. Given this dataset, we test on all combinations of training (observational + one shift) and test environments (other shift).
We train the algorithms on the training environments and evaluate them on the test environments using mean square error (MSE). The results in Figures 2(a) and 2(b) indicate that Rob-ID outperforms existing distributional robustness methods, particularly for larger shifts in new directions.

Failure cases of IRM as motivation for our work We agree that while the cited papers include non-identifiability as failure cases (Kamath et al. ‘21 Section 4 and Rosenfeld et al. ‘20), they also discuss other possible reasons for failure. We can clarify this in the revised version by adding that multiple failure cases have been discussed in the literature (citing these papers), one of which is non-identifiability, the focus of this work.

Application of existing work for partially identifiable case We agree that prior methods can be applied and evaluated in the non-identifiable setting even though prior analysis so far has focused on the case where the robust risk was identifiable (including your described case when the ”uncertainty set which is restricted to dimensions that shifted between the training environments”). In fact, one of the goals of our paper was exactly to evaluate prior methods in this new setting.

Existing formalisms for spurious correlations We thank the reviewer for pointing us to relevant and interesting works [3] and [4]. Although these works describe cases where the optimal robust predictor cannot be recovered, they seem to be binary negative statements in that they do not provide a quantification of this failure nor a proposed method for robustness in this failure case.

Q1 (upper bounds) In our result, for $\gamma > \gamma_{th}$, the value of the inf sup is exact, thus, in Theorem 3.1, Case (b) line 2 it is both a lower and an upper bound (we will clarify this in the main text). For small $\gamma$, the loss of the anchor regression estimator is a tight upper bound – which for the practitioner means that if the shifts in unknown directions are expected to be very small, the anchor estimator is optimal.

Q2 (connections to invariance objectives and regularizers) Thanks for the suggestion, we will include a more intuitive interpretation of the individual terms in the revised manuscript. In particular, we will discuss the population objective (13) that corresponds to the empirical objective in (19) when we replace $S$, $\beta^S$, $R$ by their empirical estimates. The objective consists of three parts: 1) loss on the reference environment (without any distribution shift), 2) the term $||S^\top (\beta^S - \beta)||$ that “aligns” $\beta$ in the directions of the true causal parameter $\beta^\star$ projected on S, and 3) the term $(C_{ker} + || R^\top \beta ||_2)^2$ that shrinks $\beta$ in the directions unseen during training. Since $\beta^\star$ is the optimal invariant predictor (although unknown), the second term can be interpreted as aligning the estimator towards the true invariant predictor along the observed shifts (and hence corresponds to inducing invariance across training environments, spiritually corresponding to the invariant literature as the reviewer pointed out).

For citations, please refer to the general rebuttal out of space reasons.