SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Cian Eastwood\*1,2 Shashank Singh\*1 Andrei L. Nicolicioiu\*1 Marin Vlastelica' Julius von Kugelgen1,3 Bernhard Scholkopf
<sup>1</sup> Max Planck Institute for Intelligent Systems, Tübingen
$^{2}$  University of Edinburgh  $^{3}$  University of Cambridge
To avoid failures on out-of-distribution data, recent works have sought to extract features that have an invariant or stable relationship with the label across domains, discarding "spurious" or unstable features whose relationship with the label changes across domains. However, unstable features often carry complementary information that could boost performance if used correctly in the test domain. In this work, we show how this can be done without test-domain labels. In particular, we prove that pseudo-labels based on stable features provide sufficient guidance for doing so, provided that stable and unstable features are conditionally independent given the label. Based on this theoretical insight, we propose Stable Feature Boosting (SFB), an algorithm for: (i) learning a predictor that separates stable and conditionally-independent unstable features; and (ii) using the stable-feature predictions to adapt the unstable-feature predictions in the test domain. Theoretically, we prove that SFB can learn an asymptotically-optimal predictor without test-domain labels. Empirically, we demonstrate the effectiveness of SFB on real and synthetic data.
Machine learning systems can be sensitive to distribution shift [26]. Often, this sensitivity is due to a reliance on "spurious" features whose relationship with the label changes across domains, ultimately leading to degraded performance in the test domain of interest [21]. To avoid this pitfall, recent works on domain or out-of-distribution (OOD) generalization have sought predictors which only make use of features that have a stable or invariant relationship with the label across domains, discarding the spurious or unstable features [45, 1, 35, 15]. However, despite their instability, spurious features can often provide additional or complementary information about the target label. Thus, if a predictor could be adjusted to use spurious features optimally in the test domain, it would boost performance substantially. That is, perhaps we don't need to discard spurious features at all but rather use them in the right way.
As a simple but illustrative example, consider the ColorMNIST or CMNIST dataset [1]. This transforms the original MNIST dataset into a binary classification task (digit in 0–4 or 5–9) and then: (i) flips the label with probability 0.25, meaning that, across all 3 domains, digit shape correctly determines the label with probability 0.75; and (ii) colorizes the digit such that digit color (red or green) is a more informative but spurious feature (see Fig. 1a). Prior work focused on learning an invariant predictor that uses only shape and avoids using color—a spurious feature whose relationship with the label changes across domains. However, as shown in Fig. 1b, the invariant predictor is suboptimal test domains where color can be used in a domain-specific manner to improve performance. We thus ask: when and how can such informative but spurious features be safely harnessed without labels?
*Equal contribution. Correspondence to c.eastwood@ed.ac.uk or shashankssingh44@gmail.com.
(b)
Figure 1: Invariant (stable) and spurious (unstable) features. (a) Illustrative images from CMNIST [1]. (b) CMNIST accuracies (y-axis) over test domains of decreasing color-label correlation (x-axis). The 'Oracle' uses both invariant (shape) and spurious (color) features optimally in the test domain, boosting performance over an invariant model (orange region). We show how this can be done without test-domain labels. (c) Generally, invariant models use only the stable component  $X_{S}$  of  $X$ , discarding the spurious or unstable component  $X_{U}$ . We prove that predictions based on  $X_{S}$  can be used to safely harness a sub-component of  $X_{U}$  (dark-orange region).
Our main contribution lies in answering this question, showing when and how it is possible to safely harness spurious or unstable features without test-domain labels. In particular, we prove that predictions based on stable features provide sufficient guidance for doing so, provided that stable and unstable features are conditionally independent given the label (see Fig. 1c).
Structure and contributions. The rest of this paper is organized as follows. We first discuss related work in § 2, providing context and high-level motivation for our approach. In § 3, we then explain how stable and unstable features can be extracted, how unstable features can be harnessed with test-domain labels, and the questions/challenges that arise when trying to harness unstable features without test-domain labels. In § 4, we present our main theoretical contributions which provide precise answers to these questions, before using these insights to propose the Stable Feature Boosting (SFB) algorithm in § 5. In § 6, we present our experimental results, before ending with a discussion and concluding remarks in § 7. Our contributions can be summarized as follows:
- Algorithmic: We propose Stable Feature Boosting (SFB), the first algorithm for using invariant predictions to safely harness spurious features without test-domain labels.
- Theoretical: SFB is grounded in a novel theoretical result (Thm 4.4) giving sufficient conditions for provable test-domain adaptation without labels. Under these conditions, Thm 4.6 shows that, given enough unlabeled data, SFB learns the Bayes-optimal adapted classifier in the test domain.
- Experimental: Our experiments on synthetic and real-world data demonstrate the effectiveness of SFB—even in scenarios where it is unclear if its assumptions are fully satisfied.
Domain generalization, robustness and invariant prediction. A fundamental starting point for work in domain generalization is the observation that certain "stable" features, often direct causes of the label, may have an invariant relationship with the label across domains [45, 1, 67, 55, 40, 77, 14]. However, such stable or invariant predictors often discard highly informative but unstable information. Rothenhausler et al. [51] show that we may need to trade off stability and predictiveness, while Eastwood et al. [15] seek such a trade-off via an interpretable probability-of-generalization parameter. The current work is motivated by the idea that one might avoid such a trade-off by changing how unstable features are used at test time, rather than discarding them at training time.
Test-domain adaptation without labels (unsupervised domain adaptation). In the source-free and test-time domain adaptation literature, it is common to adapt to new domains using a model's own pseudo-labels [20, 36, 39, 71, 30]—see Rusak et al. [52] for a recent review. In contrast, we: (i) use one (stable) model to provide reliable/robust pseudo-labels and another (unstable) model to adapt domain-specific information; and (ii) propose a bias correction step that provably ensures an accurate, well-calibrated unstable model  $(\operatorname*{Pr}[Y|X_U])$  as well as an optimal joint/combined model  $(\operatorname*{Pr}[Y|X_S,X_U])$ . Beyond this literature, Bui et al. [12] propose a meta-learning approach for exploiting unstable/domain-specific features. However, they use unstable features in the same way in the test domain, which, by definition, is not robust and can degrade performance. Sun et al. [63] share the goal of exploiting unstable features to go "beyond invariance". However, in contrast to our approach, they require labels for the unstable features (rarely available) and only address label shifts.
Table 1: Comparison with related work. *QRM [15] uses a hyperparameter  $\alpha \in [0,1]$  to balance the probability of robust generalization and using more information from  $X$ .
Method	Components of X Used	Robust	No test-domain labels
Stable	Complementary	All
ERM [65]	✓	✓	✓	✗	✓
IRM [1]	✓	✗	✗	✓	✓
QRM [15]	✓	✓*	✓*	✓*	✓
DARE [50]	✓	✓	✓	✓	✗
ACTIR [31]	✓	✓	✗	✓	✗
SFB (Ours)	✓	✓	✗	✓	✓
Test-domain adaptation with labels (few-shot fine-tuning). Fine-tuning part of a model using a small number of labeled test-domain examples is a common way to deal with distribution shift [16, 17, 13]. More recently, it has been shown that simply retraining the last layer of an ERM-trained model outperforms more robust feature-learning methods on spurious correlation benchmarks [50, 32, 74]. Similar to our approach, Jiang and Veitch [31] separate stable and conditionally-independent unstable features and then adapt their use of the latter in the test domain. However, in contrast to our approach, theirs requires test-domain labels. In addition, they assume data is drawn from an anti-causal generative model, which is strictly stronger than our "complementarity" assumption (see § 4).
Table 1 summarizes related work while App. H discusses further related work.
3 Problem Setup: Extracting and Harnessing Unstable Features
Setup. We consider the problem of domain generalization (DG) [8, 42, 24] where predictors are trained on data from multiple training domains and with the goal of performing well on data from unseen test domains. More formally, we consider datasets  $D^{e} = \{(X_{i}^{e},Y_{i}^{e})\}_{i = 1}^{n_{e}}$  collected from  $m$  different training domains or environments  $\mathcal{E}_{\mathrm{tr}} := \{E_1,\dots ,E_m\}$ , with each dataset  $D^{e}$  containing data pairs  $(X_{i}^{e},Y_{i}^{e})$  sampled i.i.d. from  $\mathbb{P}(X^e,Y^e)$ . The goal is then to learn a predictor  $f(X)$  that performs well on a larger set  $\mathcal{E}_{\mathrm{all}} \supset \mathcal{E}_{\mathrm{tr}}$  of possible domains.
Average performance: use all features. The first approaches to DG sought predictors that perform well on average over domains [8, 42] using empirical risk minimization (ERM, [66]). However, predictors that perform well on average provably lack robustness [43, 49], potentially performing quite poorly on large subsets of  $\mathcal{E}_{\mathrm{all}}$ . In particular, minimizing the average error leads predictors to make use of any features that are informative about the label (on average), including "spurious" or "shortcut" [21] features whose relationship with the label is subject to change across domains. In test domains where these feature-label relationships change in new or more severe ways than observed during training, this usually leads to significant performance drops or even complete failure [73, 4].
Worst-case or robust performance: use only stable features. To improve robustness, subsequent works sought predictors that only use stable or invariant features, i.e., those that have a stable or invariant relationship with the label across domains [45, 1, 47, 70, 58]. For example, Arjovsky et al. [1] do so by enforcing that the classifier on top of these features is optimal for all domains simultaneously. We henceforth use stable features and  $X_{S}$  to refer to these features, and stable predictors to refer to predictors which use only these features. Analogously, we use unstable features and  $X_{U}$  to refer to features with an unstable or "spurious" relationship with the label across domains. Note that  $X_{S}$  and  $X_{U}$  partition the components of  $X$  which are informative about  $Y$ , as depicted in Fig. 1c, and that formal definitions of  $X_{S}$  and  $X_{U}$  are provided in § 4.
3.1 Harnessing unstable features with labels
A stable predictor  $f_{S}$  is unlikely to be the best predictor in any given domain. As illustrated in Fig. 1b, this is because it excludes unstable features  $X_{U}$  which are informative about  $Y$  and can boost performance if used in an appropriate, domain-specific manner. Assuming we can indeed learn a stable predictor with prior methods, we start by showing how  $X_{U}$  can be harnessed with test-domain labels.
2We drop the domain superscript  $e$  when referring to random variables from any environment.
Boosting the stable predictor. We describe boosted joint predictions  $f^{e}(X)$  in domain  $e$  as some combination  $C$  of stable predictions  $f_{S}(X)$  and domain-specific unstable predictions  $f_{U}^{e}(X)$ , i.e.,  $f^{e}(X) = C(f_{S}(X),f_{U}^{e}(X))$ . To allow us to adapt only the  $X_{U} - Y$  relation, we decompose the stable  $f_{S}(X) = h_{S}(\Phi_{S}(X))$  and unstable  $f_{U}^{e}(X) = h_{U}^{e}(\Phi_{U}(X))$  predictions into feature extractors  $\Phi$  and classifiers  $h$ .  $\Phi_S$  extracts stable components  $X_{S} = \Phi_{S}(X)$  of  $X$ ,  $\Phi_U$  extracts unstable components  $X_{U} = \Phi_{U}(X)$  of  $X$ ,  $h_{S}$  is a classifier learned on top of  $\Phi_S$  (shared across domains), and  $h_{U}^{e}$  is a domain-specific unstable classifier learned on top of  $\Phi_U$  (one per domain). Putting these together,
$$
f ^ {e} (X) = C \left(f _ {S} (X), f _ {U} ^ {e} (X)\right) = C \left(h _ {S} \left(\Phi_ {S} (X)\right), h _ {U} ^ {e} \left(\Phi_ {U} (X)\right)\right) = C \left(h _ {S} \left(X _ {S}\right), h _ {U} ^ {e} \left(X _ {U}\right)\right), \tag {3.1}
$$
where  $C:[0,1]\times [0,1]\to [0,1]$  is a combination function that combines the stable and unstable predictions. For example, Jiang and Veitch [31, Eq. 2.1] add stable  $p_S$  and unstable  $p_U$  predictions in logit space, i.e.,  $C(p_S,p_U) = \sigma (\log (p_S) + \log (p_U))$ . Since it is unclear, a priori, how to choose  $C$ , we will leave it unspecified until Thm. 4.4 in § 4, where we derive a principled choice.
Adapting with labels. Given a new domain  $e$  and labels  $Y^{e}$ , we can boost performance by adapting  $h_{U}^{e}$ . Specifically, letting  $\ell : \mathcal{Y} \times \mathcal{Y} \to \mathbb{R}$  be a loss function (e.g., cross-entropy) and  $R^{e}(f) = \mathbb{E}_{(X,Y)}[\ell(Y,f(X))|E=e]$  the risk of predictor  $f: \mathcal{X} \to \mathcal{Y}$  in domain  $e$ , we can adapt  $h_{U}^{e}$  to solve:
$$
\min  _ {h _ {U}} R ^ {e} \left(C \left(h _ {S} \circ \Phi_ {S}, h _ {U} \circ \Phi_ {U}\right)\right) \tag {3.2}
$$
3.2 Harnessing unstable features without labels
We now consider the main question of this work—can we reliably harness  $X_U$  without test-domain labels? We could, of course, simply select a fixed unstable classifier  $h_U^e$  by relying solely on the training domains (e.g., by minimizing average error), and hope that this works for the test-domain  $X_U - Y$  relation. However, by definition of  $X_U$  being unstable, this is clearly not a robust or reliable approach—the focus of our efforts in this work, as illustrated in Table 1. As in § 3.1, we assume that we are able to learn a stable predictor  $f_S$  using prior methods, e.g., IRM [1] or QRM [15].
From stable predictions to robust pseudo-labels. While we don't have labels in the test domain, we do have stable predictions. By definition, these are imperfect (i.e., noisy) but robust, and can be used to form pseudo-labels  $\widehat{Y}_i = \arg \max_j(f_S(X_i))_j$ , with  $(f_{S}(X_{i}))_{j}\approx \operatorname *{Pr}[Y_{i} = j|X_{S}]$  denoting the  $j^{\mathrm{th}}$  entry of the stable prediction for  $X_{i}$ . Can we somehow use these noisy but robust pseudo-labels to guide our updating of  $h_U^e$ , and, ultimately, our use of  $X_{U}$  in the test domain?
From joint to unstable-only risk. If we simply use our robust pseudo-labels as if they were true labels—updating  $h_U^e$  to minimize the joint risk as in Eq. (3.2)—we arrive at trivial solutions since  $f_S$  already predicts its own pseudo-labels with  $100\%$  accuracy. For example, if we follow [31, Eq. 2.1] and use the combination function  $C(p_S, p_U) = \sigma(\logit(p_S) + \logit(p_U))$ , then the trivial solution  $\logit(h_U^e(\cdot)) = 0$  achieves  $100\%$  accuracy (and minimizes cross-entropy; see Prop. D.1 of App. D). Thus, we cannot minimize a joint loss involving  $f_S$ 's predictions when using  $f_S$ 's pseudo-labels. A sensible alternative is to update  $h_U^e$  to minimize the unstable-only risk  $R^e(h_U^e \circ \Phi_U)$ .
More questions than answers. While this new procedure could work, it raises questions about when it will work, or, more precisely, the conditions under which it can be used to safely harness  $X_U$ . We now summarise these questions before addressing them in § 4:
1. Does it make sense to minimize the unstable-only risk? In particular, when can we minimize the unstable-predictor risk alone or separately, and then arrive at the optimal joint predictor? This cannot always work; e.g., for independent  $X_S$ ,  $X_U \sim$  Bernoulli(1/2) and  $Y = X_S$  XOR  $X_U$ ,  $Y$  is independent of each of  $X_S$  and  $X_U$  and hence cannot be predicted from either alone.
2. How should we combine predictions? Is there a principled choice for the combination function  $C$  in Eq. (3.1)? In particular, is there a  $C$  that correctly weights stable and unstable predictions in the test domain? As  $X_U$  could be very strongly or very weakly predictive of  $Y$  in the test domain, this seems a difficult task. Intuitively, correctly weighting stable and unstable predictions requires them to be properly calibrated: do we have any reason to believe that, after training on  $f_S$ 's pseudo-labels,  $h_U^e$  will be properly calibrated in the test domain?
3. Can the student outperform the teacher? Stable predictions likely make mistakes—indeed, this is the motivation for trying to improve them. Is it possible to correct these mistakes with  $X_U$ ? Is it
possible to learn an unstable "student" predictor that outperforms its own supervision signal or "teacher"? Perhaps surprisingly, we show that, for certain types of features, the answer is yes. In fact, even a very weak stable predictor, with performance just above chance, can be used to learn an optimal unstable classifier in the test domain given enough unlabeled data.
4 Theory: When Can We Safely Harness Unstable Features Without Labels?
Suppose we have already identified a stable feature  $X_S$  and a potentially unstable feature  $X_U$  (we will return to the question of how to learn/extract  $X_S$  and  $X_U$  themselves in § 5). In this section, we analyze the problem of using  $X_S$  to leverage  $X_U$  in the test domain without labels. We first reduce this to a special case of the so-called "marginal problem" in probability theory, i.e., the problem of identifying a joint distribution based on information about its marginals. In the special case where two variables are conditionally independent given a third, we show this problem can be solved exactly. This solution, which may be of interest beyond the context of domain generalization/adaptation, motivates our test-domain adaptation algorithm (Alg. 1), and forms the basis of Thm. 4.6 which shows that Alg. 1 converges to the best possible classifier given enough unlabeled data.
We first pose a population-level model of our domain generalization setup. Let  $E$  be a random variable denoting the environment. Given an environment  $E$ , we have that the stable feature  $X_{S}$ , unstable feature  $X_{U}$  and label  $Y$  are distributed according to  $P_{X_{S},X_{U},Y|E}$ . We can now formalize the three key assumptions underlying our approach, starting with the notion of a stable feature, motivated in § 3:
Definition 4.1 (Stable and Unstable Features).  $X_S$  is a stable feature with respect to  $Y$  if  $P_{Y|X_S}$  does not depend on  $E$ ; equivalently, if  $Y$  and  $E$  are conditionally independent given  $X_S$  ( $Y \perp E|X_S$ ). Conversely,  $X_U$  is an unstable feature with respect to  $Y$  if  $P_{Y|X_U}$  depends on  $E$ ; equivalently, if  $Y$  and  $E$  are conditionally dependent given  $X_U$  ( $Y \not\perp E|X_U$ ).
Next, we state our complementarity assumption, which we will show justifies the approach of separately learning the relationships  $X_{S} - Y$  and  $X_{U} - Y$  and then combining them:
Definition 4.2 (Complementary Features).  $X_S$  and  $X_U$  are complementary features with respect to  $Y$  if  $X_S \perp \perp X_U | (Y, E)$ ; i.e., if  $X_S$  and  $X_U$  share no redundant information beyond  $Y$  and  $E$ .
Finally, to provide a useful signal for test-domain adaptation, the stable feature needs to help predict the label in the test domain. Formally, we assume:
Definition 4.3 (Informative Feature).  $X_S$  is said to be informative of  $Y$  in environment  $E$  if  $X_S \nsubseteq Y|E$ ; i.e.,  $X_S$  is predictive of  $Y$  within the environment  $E$ .
We will discuss the roles of these assumptions after stating our main result (Thm. 4.4) that uses them. To keep our results as general as possible, we avoid assuming a particular causal generative model, but the above conditional (in)dependence assumptions can be interpreted as constraints on such a causal model. App. D.2 formally characterizes the set of causal models that are consistent with our assumptions and shows that our setting generalizes those of prior works [49, 68, 31, 69].
Reduction to the marginal problem with complementary features. By Defn. 4.1, we have the same stable relationship  $P_{Y|X_S,E} = P_{Y|X_S}$  in training and test domains. Now, suppose we have used the training data to learn this stable relationship and thus know  $P_{Y|X_S}$ . Also suppose that we have enough unlabeled data from test domain  $E$  to learn  $P_{X_S,X_U|E}$ , and recall that our ultimate goal is to predict  $Y$  from  $(X_S,X_U)$  in test domain  $E$ . Since the rest of our discussion is conditioned on  $E$  being the test domain, we omit  $E$  from the notation. Now note that, if we could express  $P_{Y|X_S,X_U}$  in terms of  $P_{Y|X_S}$  and  $P_{X_S,X_U}$ , we could then use  $P_{Y|X_S,X_U}$  to optimally predict  $Y$  from  $(X_S,X_U)$ . Thus, our task thus becomes to reconstruct  $P_{Y|X_S,X_U}$  from  $P_{Y|X_S}$  and  $P_{X_S,X_U}$ . This is an instance of the classical "marginal problem" from probability theory [28, 29, 19], which asks under which conditions we can recover the joint distribution of a set of random variables given information about its marginals. In general, although one can place bounds on the conditional distributions  $P_{Y|X_U}$  and  $P_{Y|X_S,X_U}$ , they cannot be completely inferred from  $P_{Y|X_S}$  and  $P_{X_S,X_U}$  [19]. However, the following section demonstrates that, under the additional assumptions that  $X_S$  and  $X_U$  are complementary and  $X_S$  is informative, we can exactly recover  $P_{Y|X_U}$  and  $P_{Y|X_S,X_U}$  from  $P_{Y|X_S}$  and  $P_{X_S,X_U}$ .
4.1 Solving the marginal problem with complementary features
We now present our main result which shows how to reconstruct  $P_{Y|X_S,X_U}$  from  $P_{Y|X_S}$  and  $P_{X_S,X_U}$  when  $X_S$  and  $X_U$  are complementary and  $X_S$  is informative. To simplify notation, we assume the label  $Y$  is binary and defer the multi-class extension to App. C.
Theorem 4.4 (Solution to the marginal problem with binary labels and complementary features). Consider three random variables  $X_{S}$ ,  $X_{U}$ , and  $Y$ , where (i)  $Y$  is binary ( $\{0,1\}$ -valued), (ii)  $X_{S}$  and  $X_{U}$  are complementary features for  $Y$  (i.e.,  $X_{S} \perp X_{U}|Y$ ), and (iii)  $X_{S}$  is informative of  $Y$  ( $X_{S} \nsubseteq Y$ ). Then, the joint distribution of  $(X_{S}, X_{U}, Y)$  can be written in terms of the joint distributions of  $(X_{S}, Y)$  and  $(X_{S}, X_{U})$ . Specifically, if  $\widehat{Y} | X_{S} \sim \text{Bernoulli}(\Pr[Y = 1| X_{S}])$  is a pseudo-label<sup>3</sup> and
$$
\epsilon_ {0} := \Pr [ \widehat {Y} = 0 | Y = 0 ] \quad a n d \quad \epsilon_ {1} := \Pr [ \widehat {Y} = 1 | Y = 1 ] \tag {4.1}
$$
are the accuracies of the pseudo-labels on classes 0 and 1, respectively. Then, we have:
$$
\epsilon_ {0} + \epsilon_ {1} > 1, \tag {4.2}
$$
$$
\Pr [ Y = 1 | X _ {U} ] = \frac {\Pr [ \widehat {Y} = 1 | X _ {U} ] + \epsilon_ {0} - 1}{\epsilon_ {0} + \epsilon_ {1} - 1}, \quad a n d \tag {4.3}
$$
$$
\Pr [ Y = 1 \mid X _ {S}, X _ {U} ] = \sigma \left(\operatorname {l o g i t} \left(\Pr [ Y = 1 \mid X _ {S} ]\right) + \operatorname {l o g i t} \left(\Pr [ Y = 1 \mid X _ {U} ]\right) - \operatorname {l o g i t} \left(\Pr [ Y = 1 ]\right)\right). \tag {4.4}
$$