Out-of-distribution generalization under random, dense distributional shifts

Distribution shift is a persistent challenge in statistics and machine learning. For example, we might be interested in understanding the determinants of positive outcomes in early childhood education. However, the relationships among learning outcomes and other variables can shift over time and across locations. This makes it challenging to gain actionable knowledge that can be used in different situations. As a result, there has been a surge of interest in robust and generalizable machine learning. Existing approaches can be classified as invariance-based approaches and adversarial approaches.

Invariance-based methods assume that the shifts only affect a sub-space of the data distribution [Pan and Yang, 2010, Baktashmotlagh et al., 2013]. For example, under covariate shift [Quinonero-Candela et al., 2009], the distribution of the covariates changes while the conditional distribution of the outcomes given the covariates stays invariant. In causal representation learning [Schölkopf et al., 2021], the goal is to disentangle the distribution into independent causal factors; based on the assumption that spurious associations can be separated from invariant associations. Both approaches assume the shift affects only a sub-space of the data. We call such shifts sparse distribution shifts.

Adversarial methods in robust machine learning consider worst-case perturbations. The idea is that some adversary can perturb the data either at the training [Huber, 1981] or deployment stage [Szegedy et al., 2014, Duchi and Namkoong, 2021]; and the goal is to learn a model that is robust to such deviations. While developing robust procedures is important, worst-case perturbations might be too conservative in many real-world applications.

Motivated by an empirical observation, in this paper we consider distributional shifts due to random and dense perturbations. As an example, consider replication studies. In replication studies, different research teams try to replicate results by following the same research protocol. Therefore, any distribution shift between the data may arise from the superposition of many unintended, subtle errors that affect every part of the distribution, which might be modeled as random. We call such shifts dense distributional shifts. The random, dense distributional shift model has recently been used to quantify distributional uncertainty in parameter estimation problems [Jeong and Rothenhäusler, 2023, Rothenhäusler and Bühlmann, 2023] and has shown success in modelling real-world temporal shifts in refugee placement [Bansak et al., 2024]. We will develop tools to measure the similarity between randomly perturbed distributions, infer parameters and predict outcomes for partially observed and shifted distributions.

Due to a superposition of many random, small changes, a data scientist might not draw samples from the target distribution $\mathbb{P}^{0}$ but from several randomly perturbed distributions $\mathbb{P}^{1},\dots,\mathbb{P}^{K}$. More specifically, we assume that a data scientist has $K$ datasets, with the $k$-th dataset being an i.i.d. sample from a perturbed distribution $\mathbb{P}^{k}$. Intuitively speaking, some of these distributions might be similar to each other, while others might be very different. For instance, one replication study might closely resemble another due to collaboration among investigators. Additionally, a dataset obtained from California might be more similar to one from Washington than to one from Wyoming.

We illustrate the scenario where multiple datasets exhibit distributional correlations using the GTEx¹¹1The Genotype-Tissue Expression (GTEx) Project was supported by the Common Fund of the Office of the Director of the National Institutes of Health, and by NCI, NHGRI, NHLBI, NIDA, NIMH, and NINDS. The dataset used for the analyses described in this manuscript is version 6 and can be downloaded in the GTEx Portal: www.gtexportal.org. data. The GTEx V6 data offers pre-processed RNA-seq gene-expression levels collected from 450 donors across 44 tissues. Treating each tissue as an individual study, we have 44 datasets. Each dataset from tissue $k$ is considered as an independent sample from a perturbed distribution $\mathbb{P}^{k}$, where $\mathbb{P}^{\bullet}$ represents a joint distribution of all gene-expression levels. Some of these datasets are likely to be correlated due to factors such as an overlap of donors and the proximity between tissues. Empirical evidence of this correlation is presented in the left side of Figure 2, where each row and column corresponds to a tissue, and each dot represents a standardized covariance between gene expression levels of a randomly selected gene-pair, among the observations in the tissue. We randomly sample 1000 gene-pairs. We can see that some tissues exhibit a clear linear relationship. In the following, we will introduce a distribution shift model that implies this linear relationship.

To model the relatedness of multiple datasets, we will assume that each $\mathbb{P}^{k}$ is a random perturbation of a common target distribution $\mathbb{P}^{0}$. Then, we adopt the random distributional perturbation model in Jeong and Rothenhäusler [2023], where dense distributional shifts emerge as the superposition of numerous small random changes, to construct perturbed distributions $\mathbb{P}^{1},\dots,\mathbb{P}^{K}$. In this random perturbation model, probabilities of events from a random probability measure $\mathbb{P}^{k}$ are slightly up-weighted or down-weighted compared to the target distribution $\mathbb{P}^{0}$.

Without loss of generality, we construct distributional perturbations for uniform distributions on $[0,1]$. A result from probability theory shows that any random variable $D$ on a finite or countably infinite dimensional probability space can be written as a measurable function $D\stackrel{{\scriptstyle d}}{{=}}h(U)$, where $U$ is a uniform random variable on $[0,1]$.²²2For any Borel-measurable random variable $D$ on a Polish (separable and completely metrizable) space $\mathcal{D}$, there exists a Borel-measurable function $h$ such that $D\stackrel{{\scriptstyle d}}{{=}}h(U)$ where $U$ follows the uniform distribution on $[0,1]$ [Dudley, 2018]. With the transformation $h$, the construction of distributional perturbations for uniform variables can be generalized to the general cases by using

Let us now construct the perturbed distribution $\mathbb{P}^{k}$ for a random variable $U$ that follows a uniform distribution under $\mathbb{P}^{0}$. Take $m$ bins $I_{j}=[(j-1)/m,j/m]$ for $j=1,\dots,m$. Let $W^{k}_{1},\dots,W^{k}_{m}$ be i.i.d. positive random variables with finite variance. We define the randomly perturbed distribution $\mathbb{P}^{k}$ by setting

Then the datasets are generated as follows:

1. Generate perturbed distributions: Start with the target distribution $\mathbb{P}^{0}$ and generate $K$ perturbed distributions $\mathbb{P}^{1},\dots,\mathbb{P}^{K}$ using random weights $W=\{\{W_{j}^{k}\}_{j=1}^{m}\}_{k=1}^{K}$.

2. Draw samples: For each $k=1,\dots,K$, draw $n_{k}$ i.i.d. samples, $\mathcal{D}_{k}=\{D_{ki}\}_{i=1}^{n_{k}}$, from the perturbed distribution $\mathbb{P}^{k}$ ($D_{ki}\stackrel{{\scriptstyle i.i.d}}{{\sim}}\mathbb{P}^{k}$), conditioned on the weights $W=\{\{W_{j}^{k}\}_{j=1}^{m}\}_{k=1}^{K}$.

Further, let $m=m(n)$ such that $\frac{m(n)}{\min(n_{1},\dots,n_{K})}$ converges to 0. This is the regime where the distributional uncertainty is of higher order than the sampling uncertainty.

In this section, we demonstrate how, under the random distributional perturbation model, we can share information across datasets from different times and locations, and infer parameters of a partially observed target distribution. Specifically, we show how the distributional CLT from the previous section allows us to identify the optimal convex combination of source distributions to best approximate the target distribution.

We consider the following domain adaptation setting throughout the section. We aim to predict a target $Y$ based on some covariates $X$, that is the data is $D=(X,Y)$. We want to estimate $\theta^{0}=\arg\min_{\theta}\mathbb{E}^{0}[\mathcal{L}(\theta,X,Y)]$ for some loss function $\mathcal{L}(\theta,X,Y)$. We consider the setting where we observe full $D=(X,Y)$ only on the source distributions $\mathbb{P}^{1},\dots,\mathbb{P}^{K}$. More specifically, for the $k$-th source data ($k\in\{1,2,\dots,K\}$), we observe $n_{k}$ i.i.d. samples $\mathcal{D}_{k}=\{D_{k1},\dots,D_{kn_{k}}\}$ drawn from the source distribution $\mathbb{P}^{k}$. Furthermore, there are $n_{0}$ i.i.d. samples from the target distribution $\mathbb{P}^{0}$, but we only observe a subset $X$ and have $\mathcal{D}_{0}=\{X_{01},\dots,X_{0n_{0}}\}$.

For any fixed non-negative weights $\beta_{1},\ldots,\beta_{K}$ with $\sum_{k=1}^{K}\beta_{k}=1$, one can consider weighted empirical risk minimization, that is, one can set

In the following, we want to study the excess risk

This will give us insights into how to choose the weights $\beta$ in an optimal fashion, and allow us to conduct statistical inference.

The proof of Lemma 1 is provided in Appendix 10, along with the additional regularity assumptions under which the consistency assumption stated in Lemma 1 holds.

By this lemma, the optimal $\beta$ in terms of out-of-distribution error is

A priori, $\beta^{*}$ is unknown to the researcher. In the following, we will describe estimation and inference for $\beta^{*}$. Our estimation strategy is based on the following observation. Using the distributional CLT, for any function $\phi(X)$ and any $\beta$,

We can use this fact that the covariance term of the limiting distribution is inflated by the factor $\beta^{\intercal}\Sigma^{W}\beta$ regardless of the test functions to estimate $\beta^{*}$. Consider different test functions of $X$, that is $\phi_{\ell}(X)$, $\ell=1,\dots,L$. For example, we may use individual covariates as test functions, $\phi_{\ell}=X_{\ell}$. Then, we can estimate $\beta^{*}$ by solving

and set the final estimate of $\theta^{0}$ as

We will discuss the accuracy of this estimate in the following section. More specifically, we will derive confidence intervals for $\beta^{*}$ and $\theta^{0}$.

In this section, we demonstrate the effectiveness and robustness of our method when the training dataset is obtained from multiple sources. We focus on the ACS Income dataset [Ding et al., 2021] where the goal is to predict the logarithmic value of individual income based on tabular census data.

Similar as in Shen et al. [2023], we consider a scenario where we initially have a limited training dataset from California (CA), and subsequently, we obtain additional training dataset from Puerto Rico (PR). We iterate through the target state among the rest of 49 states. Note that there are substantial economic, job market, and cost-of-living disparities between CA and PR. Most of the states are more similar to CA than PR. Therefore, increasing the training dataset with data sourced from PR will lead to substantial performance degradation of the prediction algorithm on the target state.

We compare three methods in our study. In the first method, we fit the XGBoost on the training dataset in the usual manner, with all samples assigned equal weights. In the second method, we employ sample-specific importance weights. The importance weights are calculated as follows. First, pool the covariates $X$ of the training dataset (CA + PR) and the target dataset. Let $A_{i}$ be the label which is 1 if the $i$-th sample is from the target dataset and 0 otherwise. Then, the importance weight for the $i$-th sample is obtained as

where we estimate $\hat{P}(A_{i}|X=X_{i})$ using XGBoost on the pooled covariates data and estimate $\hat{P}(A_{i}=1)$ as the ratio of the sample size of the target data to the sample size of the pooled data.

In the third method, based on our approach, the XGBoost is fitted on the training dataset, but this time samples receive different weights depending on whether they originate from CA or PR. We utilize the occupation-code feature to construct test functions. The test function $\phi_{\ell}(D_{i})$ is defined as whether $i$-th sample has the $\ell$-th occupation code. The occupation code is a categorical variable representing around 500 different occupation categories. The data dictionary at ACS PUMS documentation³³3https://www.census.gov/programs-surveys/acs/microdata/documentation.2018.html provides the full list of occupation codes. The number of test functions in total is therefore around 500. Then, we estimate distribution-specific-weights $(w_{\text{CA}},w_{\text{PR}})$ by running

Note that $\hat{\mathbb{E}}^{\text{state}}[\phi_{\ell}]$ is the proportion of samples of the srtate with the $\ell$-th occupation code.

The results with the target states of Texas (TX) and Washington (WA) are given in Figure 1. Similar patterns are observed for other target states, aligning with either TX or WA behaviors. Results for other target states can be found in the Appendix 13.

For the target TX, as we include the PR data, the Mean Squared Error (MSE) initially drops for all the methods. However, as more data from PR is added, the MSE starts to increase for both methods with equal weights and importance weights. In contrast, our method demonstrates robustness, with the MSE decreasing after including PR data and staying consistently low, even when the PR source becomes dominant in the training dataset. Turning to the target WA, for both methods with equal weights and importance weights, adding data sourced from PR results in a straightforward increase in MSE. In contrast, our method assigns weights close to 0 for samples originating from PR, preventing significant performance degradation.

In this section, we illustrate our methods using the GTEx data. We will also run model diagnostics to evaluate model fit. The GTEx V6 dataset provides RNA-seq gene-expression levels collected from 450 donors across 44 tissues.

Treating each tissue as a separate study, we have 44 different datasets. As shown in Figure 2, some tissues exhibit higher correlations than others. To demonstrate our method, we select 5 different tissues (same as in Figure 2): Adipose subcutaneous, Adipose visceral omentum, Brain cortex, Brain frontal cortex, and Brain cerebellum. Let brain cortex be our target tissue (the third tissue in Figure 2). The code and the data are available as R package dlm at https://github.com/yujinj/dlm.

For our test functions, we randomly sample 1000 gene-pairs and define a test function as the standardized product of gene-expression levels of a pair. In total, we have 1000 test functions. Most genes are known to be uncorrelated with each other, with the 90th percentile range of correlations between two different test functions being about $[-0.06,0.06]$. We employ group Lasso to estimate the sparse inverse covariance matrix. As the estimated inverse covariance matrix is close to a diagonal matrix (details provided in the Appendix 13), we do not perform a whitening transformation.

The R package dlm (https://github.com/yujinj/dlm) contains a function dlm that performs distributional linear regressions:

    dlm(formula, test.function, data, whitening)

Here, the formula parameter specifies which model we want to fit. The test.function parameter represents considered test functions. The data parameter passes a list of datasets. The whitening parameter is a boolean indicating whether one wants to perform a whitening transformation. Additional details and descriptions of the function can be found on the GitHub page.

Below is the summary output of the dlm function with a model formula “Brain cortex $\sim$ Adipose subcutaneuous + Adipose visceral omentum + Brain frontal cortex + Brain cerebellum” with 1000 test functions defined as standardized products of randomly selected gene-pairs. Note that the summary output of the dlm function closely resembles that of the commonly used lm function.

Call:dlm(formula = Brain_Cortex ~            Adipose_Subcutaneous +            Adipose_Visceral_Omentum +            Brain_Frontal_Cortex_BA9 +            Brain_Cerebellum,    test.function = phi,    data = GTEx_data, whitening = FALSE)Residuals:     Min       1Q   Median       3Q      Max-0.56838 -0.09436 -0.00604  0.08739  0.40917Coefficients:                           Estimate Std. Error t value Pr(>|t|)Adipose_Subcutaneous     -0.0001295  0.0304036  -0.004   0.9966Adipose_Visceral_Omentum  0.0462641  0.0250888   1.844   0.0655 .Brain_Frontal_Cortex_BA9  0.7846112  0.0184325  42.567  < 2e-16 ***Brain_Cerebellum          0.1692543  0.0217128   7.795 1.61e-14 ***---Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1Residual standard error: 0.1377 on 997 degrees of freedomMultiple R-squared:  0.5088,Adjusted R-squared:  0.5073F-statistic: 344.3 on 3 and 997 DF,  p-value: < 2.2e-16

The coefficient represents the estimated weight for each perturbed dataset. The sum of coefficients equals to one. The estimated weight for the brain-frontal-cortex dataset is close to 1: 0.7846112. On the other hand, the estimated weights for adipose-subcutaneous and adipose-visceral-omentum datasets are close to 0: -0.0001295 and 0.0462641. This supports the observations made in Figure 2, where the brain-cortex dataset (target) exhibits a higher correlation with the brain-frontal-cortex dataset than others.

Let us discuss the summary output in more detail. All tests have in common that the probability statements are with respect to distributional uncertainty, that is, the uncertainty induced by the random distributional perturbations.

1. 1.

   (Distributional t-test). The summary output provides a t-statistic and a p-value for each estimated weight. Based on the output, the estimated weight for the brain-frontal-cortex dataset is highly significant (with a t-statistic 42.567 and a p-value less than $2\cdot 10^{-16}$). We have enough evidence to believe that the true optimal weight for the brain-frontal-cortex dataset is nonzero, and thus we conclude that the brain-frontal-cortex dataset is important for predicting the target dataset. On the other hand, the estimated weights for the adipose-subcutaneous and adipose-visceral-omentum datasets are not statistically significant (with t-statistics -0.004 and 1.844 and p-values 0.9966 and 0.0655). We conclude those datasets are not as important as the brain-frontal-cortex dataset for predicting the target dataset.

2. 2.

   (Distributional F-test). The summary output provides a F-statistic and a corresponding p-value for a null hypothesis: the optimal weights are uniform ($\beta^{*}=(0.25,0.25,0.25,0.25)$ in our example). This means that different datasets are equally informative in explaining the target dataset. The F-statistic is given as 344.3, which follows the F-distribution with degrees of freedom 3 and 997 under the null hypothesis. The resulting p-value is less than $2.2\cdot 10^{-16}$. Therefore, we reject the null hypothesis. From Figure 2 and the estimated weights, we can see that the brain-frontal-cortex and brain-cerebellum datasets are more informative in explaining the target dataset than the adipose-subcutaneous and adipose-visceral-omentum datasets.

3. 3.

   (R-squared). Usually, the coefficient of determination (R-squared) measures how much of the variation of a target variable is explained by other explanatory variables in a regression model. In our output, R-squared measures how much of the unexplained variance of a target dataset, compared to considering uniform weights, is explained by considering our proposed weights. Specifically, it is calculated as

   $$R^{2}=1-\frac{\text{RSS with estimated weights}}{\text{RSS with uniform weights}}.$$

   R-squared is close to 0 if the estimated weights are close to uniform weights.

Finally, we should consider diagnostic plots to evaluate the appropriateness of the fitted distribution shift model. This can be evaluated using a distributional Tukey-Anscombe and a distributional QQ-Plot. In contrast to conventional residual plots and QQ-Plots, each point in these plot corresponds to a mean of a test function instead of a single observation. Figure 3 shows that neither the QQ-Plot nor the TA plot indicate a deviation from the modelling assumptions.

In many practical settings, there is a distribution shift between the training and the target distribution. Existing distribution shift models fall into two categories. One category, which we term “sparse distribution shift”, assumes that the shift affects specific parts of the data generation process while leaving other parts invariant. Existing methods attempt to re-weight training samples to match the target distribution or learn invariant representations of features. The second category considers worst-case distributional shifts. More specifically, they consider the worst-case behavior of a statistical functional over a fixed neighborhood of the model based on the distributional distance. This often requires knowledge of the strength and shape of the perturbations.

In contrast, we consider random dense distributional shifts, where the shift arises as the superposition of many small random changes. We have found that the random perturbation model can be useful in modeling empirical phenomena observed in Figure 2. Moreover, we see that even when the overlap assumptions seem to be violated (when the reweighing approach fails), or when the distribution shift appears large in Kullback-Leibler divergence, the random perturbation model may still be appropriate and useful. Under the random distributional perturbation model, we establish foundations for transfer learning and generalization.

Our method shares methodological similarities with synthetic controls. While synthetic controls are typically applied to panel data under the assumption of a linear factor model, our procedure can be applied to model distribution shifts of any type of data (discrete, continuous, ordinal) with or without time structure. Furthermore, our method justifies the linearity in synthetic control methods under random distributional shifts, rather than assuming a linear factor model. The random distributional shift assumption can be evaluated based on plots in Figure 3. Additionally, we propose a generalization of our procedure to empirical risk minimization.

In practice, perturbations may involve a combination of sparse and dense shifts. In such cases, hybrid approaches (sparse + dense) may be appropriate. We view hybrid models as an important generalization and leave it for future work.

A companion R package, dlm, is available at https://github.com/yujinj/dlm. Our package performs distributional generalization under the random perturbation model. Users can replicate the results in Section 5 with the data and code provided in the package.

The authors are grateful for insightful discussions with Naoki Egami, Kevin Guo, Ying Jin, Hongseok Namkoong, and Elizabeth Tipton. This work was supported by Stanford University’s Human-Centered Artificial Intelligence (HAI) Hoffman-Yee Grant, and by the Dieter Schwarz Foundation.