Beyond Discrepancy: A Closer Look at the Theory of Distribution Shift

Classical learning theory operates within the statistical learning framework, in which the training and testing datasets are assumed to be drawn from the same distribution [1]. However, this assumption is rarely met in practice, where models often succeed in ever-changing real world environments rarely matching the precise conditions of their training data. This motivates the problem of distribution shift, in which a learner trains on a source distribution, with the goal of generalizing well over a distinct target distribution.

Thus far, the theory of distribution shift has consistently taken a worst-case approach, typically bounding generalization error in terms of some notion of discrepancy between the source and target distributions [2, 3, 4]. In cases where the source and target distributions are completely unrelated, or the source provides little information about the decision boundary of the target, discrepancy-based analyses correctly capture the difficulty of generalization. However, in practice, many large models appear to generalize effortlessly to target distributions with non-zero discrepancy.

Motivated by this gap, we take a closer look at the theory of distribution shift. In our setting, we consider a source distribution $\mathcal{D}_{s}$ and a target distribution $\mathcal{D}_{t}$, with the goal of building an accurate classifier over $\mathcal{D}_{t}$, primarily via training samples from $\mathcal{D}_{s}$. To accomplish this, we first select a feature map, $\hat{\phi}\in\Phi$, under which the source and target distributions are similar. To make predictions, we then use $k_{n}$-nearest neighbors ($k_{n}$-NN) inside feature space over data sampled from $\mathcal{D}_{s}$.

Instead of a worst-case, discrepancy-based approach, we study generalization under an Invariant Risk Minimization (IRM)-like assumption which we term the “Statistical IRM Assumption”. IRM assumes on the existence of a feature map $\psi^{*}$ and a classifier (over feature space) $h^{*}$ so that their composition $h^{*}\circ\psi^{*}$ achieves optimal accuracy over both source and target distributions [5]. We adapt this assumption to the nearest neighbors setting, and replace the existence of $h^{*}$ with the assumption that the some feature map $\phi^{*}\in\Phi$ maps points from the target $\mathcal{D}_{t}$ close to those from the source $\mathcal{D}_{s}$ while retaining information sufficient for optimal prediction. This property allows us to leverage the fact that nearest neighbors enjoys strong generalization properties within the support of its training distribution.

One might hope that such a condition is sufficient for generalization from source data alone. Unfortunately, the existence of a suitable feature map does not imply its identifiability - there may be many poor feature maps in $\Phi$ that appear suitable when only source data are available. We show (Theorem 4) that to guarantee generalization to the target using only source data, the source must be rich enough so that this cannot happen, i.e. that all maps that lead to optimal classification over the source distribution appropriately unify the source and target. We further exhibit a learning rule which leads to provable generalization to the target under this additional condition (Theorem 3).

We next consider the case where the learner has access to unlabeled target data in addition to labeled source data. Here, the target data provides crucial new information about which feature maps transform target data close to source data in feature space. We find that it is necessary and sufficient (Theorems 6 and 5) that all maps which both lead to optimal classification over the source, and map target data close to source data, further appropriately unify the source and target classification tasks.

When generalization is not possible with the addition of unlabeled target data, some labeled target data is needed. In this setting, the goal is to minimize the amount of labeled target data used – if large amounts of labeled target data are obtainable, we could simply use standard learning algorithms directly on the target data. We introduce a complexity measure on the embedding class, $\Phi$, which we term the distance dimension, and use it to provide an upper bound on the amount target data needed for generalization. In particular, we show that the natural procedure of minimizing the empirical risk (over $\phi\in\Phi$) on the target distribution of the source data-trained nearest neighbor classifier meets this upper bound.

As alluded to above, the theory literature has primarily studied distribution shift through the lens of discrepancy [2, 3, 7, 4, 8, 9, 10]. In the transfer learning literature, in which one considers the possibility of updating a model trained on the source distribution with a relatively small amount of target data, divergence-based analyses have also been prevalent [11, 12, 13, 14]. A notable line of work attains strong guarantees in certain cases where the divergence between source and target distributions is large, but the decision boundary on the source and target are similar by honing in on the information about the decision boundary contained in the source distribution [6, 15].

Much of the attention towards the selection of feature representations has been devoted to problem of “domain generalization”, wherein the learner tries to generalize to a large of set testing environments using samples from a smaller set of source environments, which provide training data [16, 17]. As mentioned above, the IRM literature hinges on the existence of a feature map, $\psi^{*}$ and a classifier $h^{*}$ whose composition achieves optimal accuracy over both source and target distributions [5]. Another line of work considers a different assumption, namely the existence of some suitable feature map $\psi^{*}$ for which the conditional distributions of transformed features $\psi^{*}(x)$ given a label are shared across all environments [18, 19, 20].

An important part of this work considers the case where some relatively small amount of labeled target data is available to the learner, and can be exploited in the determination of a suitable feature map. In the theory literature, this setting is most closely explored by the work on ‘’few-shot representation learning” [21, 22], where the goal is to use data on a set of source tasks to learn a low dimensional representation that connects tasks together, allowing for generalization to a related target task without too many extra samples.

In considering the case where the learner has access to unlabeled target samples, we enter the ‘’unsupervised domain adaptation” setting. Here, one often uses unlabeled target data to find some feature space under which source and target supports align [23, 24]. The literature has shown that unlabeled data has provable utility in certain common settings, e.g. under covariate shift [25, 26]. Unsupervised domain adaptation has also been studied through the lens of discrepancy [9, 10].

Let the instance space $(\mathcal{X},d_{\mathcal{X}})$ be a compact metric space, and $\mathcal{Y}$ be a finite label set. A data distribution $\mathcal{D}=(\mu,\eta)$ over $\mathcal{X}\times\mathcal{Y}$ is defined by a Borel measure $\mu$ over $\mathcal{X}$, and a conditional probability function $\eta(y|x):=\Pr_{(X,Y)\sim\mathcal{D}}[Y=y\mid X=x]$.

We assume our distributions satisfy some measure-theoretic regularity conditions. In particular, we assume our Borel measures are open measures, and that the Lebesgue Differentiation Theorem always holds. See Appendix B for details to this end.

For a classifier $h:\mathcal{X}\to\mathcal{Y}$, we define its risk $R(h,\mathcal{D})$ over $\mathcal{D}$ as the probability it misclassifies, i.e. we define $R(h,\mathcal{D}):=\Pr_{(X,Y)\sim\mathcal{D}}[h(X)\neq Y]$. The classifier with the lowest possible risk is called the Bayes optimal classifier, defined as $g_{\mathcal{D}}(x)=\operatorname*{arg\,max}_{y\in\mathcal{Y}}\eta(y|x)$.

Generalizing from source data in a feature space induced by some $\phi\in\Phi$ is only possible if $\Phi$ contains a map that appropriately unifies the classification tasks on $\mathcal{D}_{s}$ and $\mathcal{D}_{t}$. We use this section to motivate and define some desirable properties of feature maps vis a vis this goal, and to introduce the Statistical IRM Assumption, formalizing our requirement for the existence of quality maps in $\Phi$.

Our investigation of the identifiability of realizing feature maps relies on one further notion – one of embedding classes with bounded complexity. To this end, we introduce a complexity measure on $\Phi$ which will play a key role in each of the settings we consider. We begin with an intermediate definition.

Distance comparers are a natural tool for our analysis – all nearest-neighbor computations inside the feature space $\mathcal{Z}$ can be expressed in such terms. This observation gives rise to a natural complexity measure for the determination of a suitable feature map, which we term the distance dimension.

In providing upper bounds, it will be important that the distance dimension be finite. We note that it is easily bounded for the two important classes of feature maps mentioned in Section 3.

We first study the possibility of constructing a classifier that generalizes to $\mathcal{D}_{t}$ using only labeled samples from $\mathcal{D}_{s}$. In this setting, a learner $L$ takes input $S\sim D_{s}^{n}$, and outputs a classifier $L(S):\mathcal{X}\to\mathcal{Y}$, with the goal of achieving a small risk on $\mathcal{D}_{t}$.

While one might hope that the Statistical IRM assumption alone is sufficient for generalization to the target, this is unfortunately false. In fact, we have already seen an example of this phenomenon in Figure 1(c). Here, we essentially argued that $\phi_{y}$ was a realizing feature map: it preserves the source risk, unifies the classification tasks on source and target, and maps all target points close to source points. However, we argued that $\phi_{x}$ and $\phi_{y}$ were statistically indistinguishable in this setting, leaving the learner in need of more information.

To the end of a general characterization of learnability in this setting, recall our discussion of Figure 1(a). Here, the projection $\phi_{y}$ could be determined as realizing the Statistical IRM assumption given that it was the only map in $\Phi$ that preserved the source distribution – it’s clear from the figure that that $\phi_{x}$ does not preserve the source, and so cannot possibly realize the Statistical IRM assumption. It is vital that this reasoning could be carried out with source data alone.

More generally, by the Statistical IRM Theorem, it is sufficient for generalization from source data alone that the learner be able to identify $\phi^{*}$ from source data alone. Such a realizing feature map must of course satisfy all three requirements of Assumption 1. However, note that only one of these requirements, namely source-preservation, depends on the source distribution alone – the others, namely contraction and Bayes-unification, are defined in terms of the target distribution. As such, only source-preservation can be tested using source data.

That said, if the learner can be assured that all source-preserving feature maps realize the Statistical IRM assumption, i.e. $\Phi^{*}=\mathcal{S}(\Phi)$, it can identify realizing feature maps by identify source-preserving feature maps. We formalize this intuitive idea with the following theorem, which shows that PAC guarantees for target generalization are obtainable when the additional condition $\Phi^{*}=\mathcal{S}(\Phi)$ holds.

We relegate specification of the learning rule to the Appendix. It is founded on minimizing the empirical risk on the source data over feature maps in $\Phi$, but further leverages the knowledge that source-preserving feature maps induce separated distributions over feature space. After selecting a candidate feature map which empirically matches the requirements for source preservation, it uses $k_{n}$-NN in the implied feature space to make predictions.

In light of the discussion above, the condition that $\Phi^{*}=\mathcal{S}(\Phi)$ is intuitively necessary as well – without it, there may be some feature map which is source-preserving, but which e.g. fails to Bayes-unify the source and the target. It’s simple to see that attempting to generalize to the target via such a feature space could be catastrophic, and thus that blindly choosing between source-preserving feature maps will eventually lead the learner astray. On the other hand, not classifying through a feature space subjects the learner to the standard pitfalls of out of distribution generalization. We formalize these ideas via the following hardness result.

Thus, in the case that some feature maps preserve the source but do not realize the Statistical IRM assumption, there is always some problem nearly identical problem instance where the Statistical IRM assumption is realized by $\Phi$, buy which causes a given learner to have unbounded sample complexity (for some choice of $\epsilon$ and $\delta$).

We now consider the less restrictive unlabeled target data are also available. Here, a learner $L$ takes input $S\sim\mathcal{D}_{s}^{n}$ and $U\sim\mu_{t}^{m}$, and outputs a classifier $L(S,U):\mathcal{X}\to\mathcal{Y}$.

The story given additional access to unlabeled data is similar to source-only setting: the Statistical IRM assumption alone is insufficient for guaranteeing successful generalization when additional unlabeled target data are available. In other words, the combination of labeled source and unlabeled target is generally insufficient for identifying a $\phi\in\Phi$ that realizes the Statistical IRM assumption.

For a simple example to this end, we return to panel (c) of Figure 1. For the source and target distributions shown, it is evident that no amount of labeled data from the source and unlabeled data from target will allow us decide whether we should project data onto the $x$-axis or the $y$-axis. This is because the only difference between them is the manner in which $\mathcal{D}_{t}$ is labeled. By contrast, the example depicted by Figure 1 panel (b) illustrates a case in which the additional unlabeled data from $\mathcal{D}_{t}$ proves sufficient: because projecting onto the $x$-axis fails to map target points close source points, we can conclude that $\phi$ must be the projection onto the $y$-axis.

As in source-only setting, identifying a feature map realizing the Statistical IRM assumption requires testing the three conditions of Assumption 1. Understanding the utility of additional unlabeled target data is to realize that it allows the learner to not only test which feature maps preserve the source, but also which maps contract $\mathcal{D}_{s}$ and $\mathcal{D}_{t}$. On the other hand, it is insufficient to determine which feature maps Bayes-unify, as this notion intrinsically depends on labeling under $\mathcal{D}_{t}$. This motivates a similar sufficient condition for learnability as we saw in the previous section –namely, that all feature maps which both preserve the source and contract the source and target further Bayes unify. The following theorem shows that this is indeed a sufficient condition for learnability in this setting.

The learning rule is specified in detail in the Appendix. It proceeds by first selecting a feature map which both empirically preserves the source, and maps each unlabeled target in $U$ point close to some source point in $S$. As above, it uses $k_{n}$-NN in the selected feature space to make predictions.

In accordance with the intuition developed above, the condition that $\Phi^{*}=\mathcal{S}(\Phi)\cap\mathcal{C}(\Phi)$ is necessary. The issue of course is that without access to labeled target data, testing whether a feature map Bayes-unifies is impossible. We formalize this via the following hardness result.

Theorem 6 shows that no combination of embedding class $\Phi$ and learner $L$ can circumvent the impossibility of testing Bayes unification with unlabeled target data. For any embedding class and learning algorithm, one can always find a pair of source and target distributions on which $\Phi$ realizes the Statistical IRM Assumption, but on which the learning algorithm will fail.

The discussion above implies that even under the Statistical IRM assumption, there are many situations where label target data is required for generalization. In such cases, we would hope that we can exploit the information encoded in the Statistical IRM Assumption to achieve generalization through labeled source data and a small amount of labeled target data. In this section we show that the Statistical IRM assumption allows for significant convergence rate speed-ups in many settings.

Recall that the Statistical IRM theorem states that given a realizing feature map, generalization to the target can be accomplished purely through source data – the limitation of a lack is labeled target data is the difficulty in identifying such a feature map. This inspires the strategy of allocating all of the labeled target data towards determining a realizing feature map.

In this spirit, we analyze the natural scheme of constructing a classifier by composing nearest-neighbors trained solely on source data with the map $\phi\in\Phi$ that minimizes the empirical risk over $T\sim\mathcal{D}_{t}^{m}$, finding that the number of target examples required for guarantees can be controlled in terms of the “distance dimension” of the class $\partial(\Phi)$.

Thus, the amount of labeled target data required for generalization when $\Phi$ realizes the Statistical IRM assumption can be largely controlled through our complexity measure on the class $\Phi$. We say “largely” given that $m$, the amount of data required from $\mathcal{D}_{t}$, has a logarithmic dependence on $n$, the amount of data drawn from $\mathcal{D}_{s}$. This implies a near distributional-independence between source and target in the sample complexity.

We note that the above margin assumptions are not required for the analysis leading to Theorem 7. Thus – comparing e.g. to rates of convergence under the canonical Tsybakov noise assumption in $\mathcal{X}=\mathbb{R}^{d}$, under which nonparametric classifiers necessarily incur rates of $\tilde{\Omega}(m^{-1/1+d})$ – the guarantees of Theorem 7 represent significant convergence rate speed-ups over naively training a non-parametric classifier with target data in many cases where the distance dimension $\partial(\Phi)$ is polynomial in the dimension of the instance space [28].

In this work, we study the problem of distribution shift under a variant of the IRM assumption, wherein it is known that a feature map in a class $\Phi$ unifies classification on source and target. We investigate the identifiability of such maps, characterizing learnability in settings where worst-case approaches indicate that learning should be impossible or expensive.

Our work suggests that the study of IRM-like assumptions is a promising direction for shedding light on new situations where guaranteeing generalization under distribution shift is possible. It also highlights that a primary issue in learning under IRM-like assumptions may be the statistical identifiability of suitable feature maps.

Acknowledgements: This work was supported by the National Science Foundation under the following grants: NSF CIF-2402817, SaTC-2241100, CCF-2217058, and ARO-MURI W911NF2110317.

RB was also partially supported by the German Research Foundation through the Cluster of Excellence “Machine Learning - New Perspectives for Science" (EXC 2064/1 number 390727645)