SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alan Q. Wang Cornell University
Minh Nguyen Cornell University
Mert R. Sabuncu Cornell University
Machine learning models will often fail when deployed in an environment with a data distribution that is different than the training distribution. When multiple environments are available during training, many methods exist that learn representations which are invariant across the different distributions, with the hope that these representations will be transportable to unseen domains. In this work, we present a nonparametric strategy for learning invariant representations based on the recently-proposed Nadaraya-Watson (NW) head. The NW head makes a prediction by comparing the learned representations of the query to the elements of a support set that consists of labeled data. We demonstrate that by manipulating the support set, one can encode different causal assumptions. In particular, restricting the support set to a single environment encourages the model to learn invariant features that do not depend on the environment. We present a causally-motivated setup for our modeling and training strategy and validate on three challenging real-world domain generalization tasks in computer vision.
Machine learning models often fail when there is significant distribution shift. The goal of domain generalization is to be able to perform well with new distributions [21, 56, 71]. In this work, we are interested in settings where multiple domains/environments are available during training and we have access to environment indicators. A popular way to tackle domain generalization in this setting is to learn representations that are invariant across environments [20, 41, 60]. The hope is that such representations will work well in, or are transportable to, unseen environments. This invariance is often encoded via constraints on a learned predictor which aligns its behavior across environments; often, these conditions are derived using causal reasoning and/or by making assumptions about the data-generating process [39].
In a parametric setting, almost all existing methods enforce these constraints by training a single model and adding a regularizer on top of a standard predictive loss [1, 5, 18, 20, 27, 48, 54, 58]. Most notably, invariant risk minimization (IRM) enforces the representations to be such that the optimal classifier on top of those representations is the same across all environments. Other examples include enforcing the layer activations of the predictor to be aligned across environments [48], enforcing the predictor to be calibrated across environments [54], and enforcing gradients of the predictor to be aligned across environments [46]. Often, optimizing these constraints demand approximations or relaxations that undermine the efficacy of the approach [22].
In this work, we take a different approach using a nonparametric strategy based on a recently-proposed Nadaraya-Watson (NW) head [55]. Instead of computing the class probability directly from an input query, the NW head makes a prediction by comparing the learned representations of the query to the elements of a support set that consists of labeled data. Thus, the NW prediction is computed relative to other real datapoints in the support set, with the support set providing a degree of flexibility not
possible with parametric models. In particular, one can manipulate it during training in a way which restricts the types of comparisons that the model can make.
In this work, we manipulate the support set during training to encode causal assumptions for the purposes of learning invariant representations. Specifically, restricting the support set to be drawn from a single environment precludes the possibility of using environment-specific features to make a prediction for a given query. We show that this setup is causally-motivated and relates to existing causal frameworks. Furthermore, we show that this training strategy leads to competitive to superior results compared to state-of-the-art parametric baselines.
- We present causally-motivated assumptions for domain generalization which justify our modeling and training strategy.
- We present a novel approach to invariant representation learning using the nonparametric Nadaraya-Watson head, which can account for causal assumptions by manipulating a support set. In particular, we propose a training strategy which, unlike competing baselines, has no invariance hyperparameter to tune.
- We validate our approach on several datasets and demonstrate competitive results compared to state-of-the-art parametric baselines.
2.1 Domain Generalization and Invariant Representations
Domain generalization seeks to make models robust to unseen environments and is an active area of research [21, 56, 71]. One line of work augments or synthetically-generates additional training images to increase robustness of learned features to unseen environments [59, 63-65, 70]. In particular, LISA uses a mixup-style [67] interpolation technique to generate augmented images, which the authors demonstrate improves out-of-distribution robustness [64]. Another line of work broadly seeks to align features across distributions. Deep CORAL aligns correlations of layer activations in deep neural networks [48], and other works minimize the divergence of feature distributions with different distance metrics such as maximum mean discrepancy [51, 32], an adversarial loss [14, 30], and Wasserstein distance [69]. Still other works approach the problem from the perspective of the gradients and optimization [12, 29, 34, 46, 57]. For example, Fish aligns the gradients from different domains [46].
One can also achieve domain generalization via learning invariant representations, which often requires reasoning about the data-generating process from a causal perspective to arrive at appropriate constraints [39]. Invariant causal prediction (ICP) formulates the problem from a feature selection perspective, where the goal is to select the features which are direct causal parents of the label [40]. Invariant Risk Minimization (IRM) can be viewed as an extension of ICP designed for deep, nonlinear neural networks. The IRM objective can be summarized as finding the representation  $\varphi$  such that the optimal linear classifier's parameters  $w^{*}$  on top of this representation is the same across all environments [1]. This bi-level program is highly non-convex and difficult to solve. To find an approximate solution, the authors consider a Lagrangian form, whereby the sub-optimality with respect to the constraint is expressed as the squared norm of the gradients of each of the inner optimization problems. Follow-up works analyzing IRM have raised theoretical issues with this objective and presented some practical concerns [16, 22, 42]. Various flavors of IRM have also been proposed by introducing different regularization terms [27, 54, 58].
2.2 Nonparametric Deep Learning
Nonparametric models in deep learning have received much attention in previous work. Deep Gaussian Processes [10], Deep Kernel Learning [62], and Neural Processes [24] build upon Gaussian Processes and extend them to representation learning. Other works have generalized  $k$ -nearest neighbors [37, 49], decision trees [68], density estimation [13], and more general kernel-based methods [15, 36, 66] to deep networks and have explored the interpretability that these frameworks provide. Closely-related but orthogonal to nonparametric models are attention-based models, most notably self-attention mechanisms popularized in Transformer-based architectures in natural language
Figure 1: Illustration of proposed approach. Support set of labeled datapoints (square/triangle) from 3 environments lie in 3 regions in the feature space. Black circle denotes query datapoint with unknown label. a) The NW head models  $P(Y|X)$  by making predictions as a function of distances to labeled datapoints in the feature space (visualized as dotted arrows). b) Balancing comparisons across labels for all environments models  $P^{B}(Y|X)$ . c) Conditioning on a single environment models  $P_{e}(Y|X)$ .
processing [53] and, more recently, computer vision [11, 19, 38]. Nonparametric transformers apply attention in a nonparametric setting [24].
Recently, Wang et al. proposed the NW head [55], an extension of the classical NW model [2, 35, 61] to deep learning. In the NW head, the prediction is a weighted average of labels from a support set. The weights are computed from distances between the query and support features. The NW head can yield better calibration and interpretability, with similar accuracy compared to the dominant approach of using a parametric classifier with fully-connected layers. In this work, we leverage the NW head to encode causal assumptions via the support set. The interpretability and explainability benefits of the NW head carry over in this work; while not of primary focus, we explore these properties in the Appendix.
Figure 2: a) Causal Directed Acyclic Graph (DAG) we consider in this work. Solid nodes are observed and dashed nodes are unobserved. We assume an anti-causal setting where label  $Y$  causes  $X$ , and  $X$  has 2 causal parents: "style" features,  $Z_{S}$ , which are influenced by the environment  $E$ ; and environment-independent "content" features of  $X$ ,  $Z_{C}$ , which are causally influenced by the label  $Y$ .  $E$  potentially influences  $Y$ . Both  $E$  and  $Y$  have direct influence on style features  $Z_{S}$ . b) Same DAG as a) with an intervention on  $Y$ . We note that  $Y \perp E \mid Z_{C}$  and  $Y \not\perp E \mid Z_{S}$ .
Problem Setting. Let  $X, Y$  denote a datapoint and its corresponding discrete class, and  $E$  denote the environment (or domain) where  $X, Y$  originates. That is, the elements of the training dataset  $\mathcal{D}_{tr} = \{x_i, y_i, e_i\}_{i=1}^N$  are drawn first by sampling the discrete random variable  $e_i \sim P(E)$ , and then sampling  $x_i, y_i \sim P(X, Y \mid E = e_i) \coloneqq P_{e_i}(X, Y)$ . Our goal is to learn classifiers that will generalize to new, unseen environments.
Assumptions. We assume there exists a pair of latent causal parents of  $X$ : an environment-independent ("content") factor  $Z_{C}$  and an environment-dependent ("style") factor  $Z_{S}$ .<sup>2</sup> We as-
<sup>1</sup>We assume the support of  $E$ ,  $\operatorname{supp}(E)$ , is finite.
This is sometimes referred to as style-content or causal decomposition [21, 42, 60].
sume the causal mechanism that generates  $X$  from  $(Z_{C}, Z_{S})$  is injective, so that, in principle, it is possible to recover the latent features from the observations; i.e. there exists a function  $g$  such that  $g(X) = (Z_{C}, Z_{S})$ . We further assume that  $g$  can be disentangled into  $g_{C}$  and  $g_{S}$ , such that  $(Z_{C}, Z_{S}) = g(X) = (g_{C}(X), g_{S}(X))$ . The causal graph is depicted in Fig. 2a. Finally, we assume that if any  $X = x$  has a non-zero probability in one environment, it has a non-zero probability in all environments.
Motivation. The motivating application in this work is image classification, where each realization of  $E$  might represent a different site, imaging device, or geographical region where  $X, Y$  are collected. For example, in medical imaging, different hospitals ( $E$ ) may collect images ( $X$ ) attempting to capture the presence of some disease ( $Y$ ), but may differ in their imaging protocols which lead to differences in the image style features  $Z_{S}$  (e.g. staining, markings, orientation). In addition, we allow  $Z_{S}$  to be influenced by the label itself (for example, positive examples are more likely to be marked by a doctor or have specific staining than negative examples). Finally, the prevalence of  $Y$  may be influenced by  $E$  (for example, the prevalence of a disease may be higher in a certain hospital).
The goal is to find an estimator for  $Y$  which relies only on the direct causal links between  $Y$  and  $X$  and not on any spurious associations between  $E$  and  $X$ , as these may change in a new, unseen environment. That is, we seek an estimator which relies only on  $Z_{C}$  and which is independent of  $E$  or  $Z_{S}$ .
First, we note the direct causal dependence  $E \to Y$ . For example, a model can exploit this association by learning to over-predict majority classes in a certain environment. One way to remove the direct dependence is by intervening on  $Y$ , thus removing incoming edges to  $Y$ . This essentially corresponds to matching the environment-specific prior on  $Y$  between environments, and results in the intervened graph in Fig. 2b. Let us refer to any distribution which follows the DAG in Fig. 2b as  $P_{e}^{B}(X,Y)$ .
Second, we observe that there is a potential non-causal association flow between  $E$  and  $Y$  through the colliders  $X$  and  $Z_{S}$ , when either one of these are conditioned on (i.e. are observed). An estimator which relies on  $Z_{S}$  potentially leaks information from  $E$ , and this is unstable in new environments. Reading d-separation on this intervened graph, we infer that  $Y \perp E \mid Z_{C} = g_{C}(X)$  and  $Y \not\perp E \mid Z_{S} = g_{S}(X)$ , that is:
$$
P _ {e} ^ {B} (Y \mid g _ {C} (X)) = P _ {e ^ {\prime}} ^ {B} (Y \mid g _ {C} (X)) \forall e, e ^ {\prime} \in E. \tag {1}
$$
In words, this assumption states that the probability of  $Y$  given the environment-invariant parts of  $X$  is the same across any environment  $e \in E$ .
Thus, we seek an estimator that 1) enables interventions on  $Y$  such that the direct dependence  $E\to Y$  can be removed, and 2) can further encode the assumption in Eq. (1).
4 NW Head for Invariant Prediction
Given a datapoint  $x$ , support set  $S = \{x_i, y_i\}_{i=1}^{N_s}$ , and parameters  $\phi$ , the NW head estimates  $P(Y = y \mid X = x)$  by outputting a prediction formulated as a weighted sum of support set labels, where the weights are some notion of similarity in the feature space [55]:
$$
\hat {P} (Y = y \mid X = x; \phi) := f _ {\phi} (x, \mathcal {S}) = \frac {\sum_ {i = 1} ^ {N _ {s}} \exp \left\{s \left(\phi (x) , \phi \left(x _ {i}\right)\right) \right\} \vec {y _ {i}}}{\sum_ {j = 1} ^ {N _ {s}} \exp \left\{s \left(\phi (x), \phi \left(x _ {j}\right)\right) \right\}}. \tag {2}
$$
Here,  $\vec{y}$  is the one-hot encoded version of  $y$  and  $s(\cdot ,\cdot)$  is a similarity/kernel function that captures the similarity between pairs of features. In this work, we set  $s$  as the negative Euclidean distance. A graphical depiction is shown in Fig. 3.
Manipulating the support set can encode certain causal assumptions. We consider two types of manipulations:
1. Balancing classes in  $S$ , denoted  $S^B$  (see Fig. 1b). This can be interpreted as an intervention on  $Y$ , and removes the dependence on  $E \to Y$ , i.e.:
$$
\hat {P} ^ {B} (Y = y \mid X = x; \phi) := f _ {\phi} (x, \mathcal {S} ^ {B}). \tag {3}
$$
This may be interpreted as making the model robust to label shift, see [31, 44].
Figure 3: A depiction of the NW head on a tumor detection task. The NW head computes Euclidean distances  $s(\cdot, \cdot)$  between query and support features, and uses the distances to weight the support labels. Colored squares represent labels. Diagram displays two different support sets. Top is unconditional support, where support data is drawn from the training data without knowledge of environment information. Bottom is an example of a manipulated support where all support data is drawn from a fixed environment (note similarity in color). Such a support set precludes the possibility of using environment-specific features to make a prediction.
2. Conditioning  $S$  on a single environment, denoted  $S_{e}$  (see Fig. 1c). This can be interpreted as conditioning the probability estimate on  $E = e$ , i.e.:
$$
\hat {P} _ {e} (Y = y \mid X = x; \phi) := f _ {\phi} (x, S _ {e}). \tag {4}
$$
Note that both balancing and conditioning can be achieved simultaneously, which we denote  $S_e^B$ .
4.1 Objective and Optimization Details
Given a dataset of samples  $\mathcal{D}_{tr} = \{x_i,y_i,e_i\}_{i = 1}^N$ , we wish to leverage the NW head as a conditional estimator for  $Y$  conditioned on  $Z_{C} = g_{C}(X)$ , where  $g_{C}(X)$  is characterized by Eq. (1). This necessitates an optimization over both  $\phi$  and the space of functions  $g_{C}$ . Thus, we solve the following constrained maximum likelihood over  $\phi$  and  $g_{C}$ :
$$
\underset {\phi , g _ {C}} {\operatorname {a r g m a x}} \sum_ {i = 1} ^ {N} \log \hat {P} _ {e _ {i}} ^ {B} \left(y _ {i} \mid g _ {C} \left(x _ {i}\right); \phi\right) \tag {5}
$$
$$
\text {s . t .} \hat {P} _ {e} ^ {B} \left(y _ {i} \mid g _ {C} \left(x _ {i}\right); \phi\right) = \hat {P} _ {e ^ {\prime}} ^ {B} \left(y _ {i} \mid g _ {C} \left(x _ {i}\right); \phi\right), \forall i \in \{1, \dots , N \}, \forall e, e ^ {\prime} \in E.
$$
Note that Eq. (1) implies that  $P_{e}^{B}(y_{i}\mid g_{C}(x_{i})) = P^{B}(y_{i}\mid g_{C}(x_{i}))$ . Thus, the objective is equivalent to unconstrained maximum likelihood under the assumption in Eq. (1).
Instead of solving for  $g_{C}$  explicitly, we let both  $\phi$  and  $g_{C}$  be related by the composition  $\varphi = \phi \circ g_{C}$ , and set  $\varphi$  to be the learnable mapping of the NW head, i.e. a neural network. Then, the objective becomes:
$$
\underset {\varphi} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {N} L \left(f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e _ {i}} ^ {B}\right), y _ {i}\right) \tag {6}
$$
$$
\begin{array}{l} \text {s . t .} f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e} ^ {B}\right) = f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e ^ {\prime}} ^ {B}\right), \forall i \in \{1, \dots , N \}, \forall e, e ^ {\prime} \in E, \end{array}
$$
where  $L$  is the cross-entropy loss. To make the objective tractable, we consider two possible variants:
1. Explicit. Solve the optimization problem explicitly via a Lagrangian formulation:
$$
\underset {\varphi} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {N} L \left(f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e _ {i}} ^ {B}\right), y _ {i}\right) + \lambda \sum_ {e, e ^ {\prime} \in E} \sum_ {i = 1} ^ {N} \left\| f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e} ^ {B}\right) - f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e ^ {\prime}} ^ {B}\right) \right\| _ {2} ^ {2}. \tag {7}
$$
where  $\lambda > 0$  is a hyperparameter.
2. Implicit. Relax the optimization problem into the following unconstrained problem:
$$
\underset {\varphi} {\operatorname {a r g m i n}} \sum_ {e \in E} \sum_ {i = 1} ^ {N} L \left(f _ {\varphi} \left(x _ {i}, \mathcal {S} _ {e} ^ {B}\right), y _ {i}\right). \tag {8}
$$
In this formulation, the constraint will be approximately satisfied in the sense that the model will be encouraged to predict the ground truth for a given image, which is identical across all environments. In practice, how well the solution satisfies the constraint will depend on model capacity, the data sample, and optimization procedure.
4.2 Optimization Details
During training, the support set  $S$  is drawn stochastically from the training set  $D_{tr}$ , and all queries and support datapoints are passed through the feature extractor  $\varphi$ . For computational efficiency, instead of sampling a unique support mini-batch at the query-level, we sample a unique support at the mini-batch level. Thus, if  $N_{q}$  and  $N_{s}$  are the query and support mini-batch sizes respectively, the effective mini-batch size is  $N_{q} + N_{s}$ , instead of  $N_{q}N_{s}$ . For the implicit variant, we sample one support set for a given mini-batch of queries, forward pass through the NW head, and compute the loss in Eq. (8). For the explicit variant, we sample two support sets for a given mini-batch of queries, perform two independent forward passes through the NW head for each support set, and compute the loss in Eq. (7). As discussed in prior work [55], the support batch size is a hyperparameter analogous and orthogonal to the query batch size.
A technical point is that the set of labels in the support mini-batch must cover the set of labels in the query mini-batch. Thus, in our implementation, for  $S^B$ , we cycle through all classes and randomly draw  $N_c$  examples per class to include in the support. For tasks with a large number of classes, one can subsample from the total number of classes, so long as the sampled classes cover the set of query classes.
4.3 Inference modes
Similar to how the support set can be manipulated during training, we can also design different inference strategies corresponding to different configurations of the support set at test-time. We explore several different inference modes which are possible under the NW framework:
1. Random. Sample uniformly at random over the dataset, such that each class is represented  $k$  times.
2. Full. Use the entire balanced training set.
3. Ensemble. Given the set of balanced features computed from Full mode, partition the computed features for all training datapoints by environment, compute the softmax predictions with respect to each environment balanced across labels, and average the predictions.
4. Cluster. Given the set of balanced features computed from Full mode, perform  $k$ -means clustering on the features of the training datapoints for each class. These  $k$  cluster centroids are then used as the support features for each class. This can be viewed as a distillation of the full training set for efficient inference, with the caveat that the support set no longer corresponds to observed datapoints.
While Full, Ensemble, and Cluster require computing features for the entire support set, in practice these features and centroids can be precomputed. In our experiments, we find that Cluster mode can be a sufficient replacement to Full mode, while being computationally cheaper. These inference modes can be used interchangeably and flexibly. As an example, consider a workflow which would involve using Cluster mode to perform efficient inference, and then using Full mode on a select few (potentially problematic) test queries to understand model behavior.
4.4 Connections to Prior Work
Our assumptions in Eq. (1) are common across many works related to learning invariant predictors [40, 27, 54, 41, 26]. Representatively, under the binary classification setting, the IRM objective finds a representation function  $\varphi$  which elicits an invariant predictor across environments  $E$  such that for all  $h$  that has a non-zero probability for  $\varphi(X)$  in any (and all) environment(s):
$$
\mathbb {E} _ {e} [ Y \mid \varphi (X) = h ] = \mathbb {E} _ {e ^ {\prime}} [ Y \mid \varphi (X) = h ], \forall e, e ^ {\prime} \in E.
$$
Eq. (1) can be viewed as a generalization of this equality to multi-class settings. $^4$ .
4This equality has also been called "sufficiency invariance" [60].
Table 1: Summary of Datasets.
Dataset	# Classes	Env	# Envs	Architecture	Metric
Camelyon-17	2	Hospital	3	DenseNet-121	Average acc.
ISIC	2	Hospital	3	ResNet-50	F1-score
FMoW	62	Region	5	DenseNet-121	Worst-region acc.
Table 2: Metric average ± standard deviation for all datasets (\%). Higher is better. **Bold** is best and **underline** is second-best. Implicit / Explicit.
Algorithm	Camelyon-17	ISIC	FMoW
ERM [52]	70.3±6.4	58.2±2.9	32.6±1.6
IRM [1]	70.9±6.8	57.9±1.0	31.3±1.2
CORAL [48]	72.4±4.4	59.1±2.2	31.7±1.0
Fish [46]	74.7±7e-2	64.4±1.7	34.6±0.0
LISA [64]	77.1±6.5	64.8±2.3	35.5±1.8
CLOvE [54]	79.9±3.9	66.2±2.2	40.1±0.6
NWb, Random	71.7±5.3	56.7±1.4	31.1±0.8
NWb, Full	72.0±6.7	61.9±3.5	31.6±0.9
NWb, Cluster	70.6±6.9	61.4±2.3	31.3±0.9
NWb, Ensemble	71.9±6.0	63.9±3.8	32.2±1.0
NWb, Probe	69.2±7.4	59.7±2.5	29.9±1.5
NWb, Random	74.8±8.4 / 75.3±3.2	57.5±1.9 / 55.0±0.9	31.2±0.7 / 30.9±0.5
NWb, Full	80.0±2.7 / 79.7±1.9	69.6±2.3 / 70.0±1.0	35.0±0.7 / 34.6±0.4
NWb, Cluster	78.6±2.5 / 79.0±1.4	71.1±1.7 / 71.0±1.0	33.9±0.6 / 34.0±0.3
NWb, Ensemble	79.5±2.6 / 79.6±1.9	69.5±2.2 / 69.8±0.8	37.8±0.9 / 38.2±0.4
NWb, Probe	75.3±7.3 / 75.8±8.3	61.4±3.1 / 63.4±2.8	33.9±1.5 / 32.7±1.4