SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Furthermore, note that given the feature extractor  $\varphi$ , the NW mechanism  $f$  is a nonlearnable classifier, whereas  $w$  is learned in the IRM setting. Thus, our proposed objective can be interpreted as learning invariant features  $\varphi$ , where the fixed classifier constraint is satisfied by construction. This avoids the need to approximate the complex bilevel optimization problem with a regularizer which assumes convexity and requires computing the Hessian. Essentially,  $f$  enforces invariance through the manipulation of the support set, providing a more intuitive and computationally simpler objective to optimize.
In the Experiments section, we compare IRM against a variant of our algorithm where we freeze the learned representations and finetune a linear classifier on top using the same training data. We find that our algorithm performs better than IRM on all datasets, suggesting that it captures invariant representations better than IRM.
We compare against several popular and competitive baseline algorithms: empirical risk minimization (ERM) [52], invariant risk minimization (IRM) [1], deep CORAL [48], Fish [46], LISA [64], and CLOvE [54]. When available, results on baselines are pulled from their respective papers. Details on baseline algorithms are provided in the Appendix.
5.2 Datasets
We experiment on 3 real-world domain generalization tasks. Two are from the WILDS benchmark [23], and the third is a challenging melanoma detection task. Details on the datasets are summarized in Table 1, and further information is provided in the Appendix.
1. The Camelyon-17 dataset [4] comprises of microscopic images of stained tissue patches from different hospitals, where the label corresponds to presence of tumor tissue in patches and the environment is the hospital where the patch comes from.
2. The melanoma dataset is from the International Skin Imaging Collaboration (ISIC) archive<sup>5</sup>. The ISIC dataset comprises of dermoscopic images of skin lesions from different hospitals, where the label corresponds to whether or not the lesion is diagnosed as melanoma and the environment is the hospital where the image comes from. There is significant less positive examples and negative examples, with this label imbalance varying across environments (see Appendix).
3. The Functional Map of the World (FMoW) dataset [6] comprises of RGB satellite images, where the label is one of 62 building or land use categories, and the environment represents the year the image was taken and its geographical region.
5.3 Experimental Setup
For each model variant, we train 5 separate models with different random seeds, and perform model selection on an out-of-distribution (OOD) validation set. For WILDS datasets, we follow all hyperparameters, use model selection techniques, and report metrics as specified by the benchmark. This includes using a DenseNet-121 backbone initialized with pretrained ImageNet weights as  $\varphi$  and no random augmentations for both datasets. Similarly for ISIC, we use a pretrained ResNet-50 backbone as  $\varphi$  with no augmentations, and perform model selection on an OOD validation set. Due to significant label imbalance, we report F1-score instead of average accuracy.
For NW algorithms, we refer to models which balance classes (i.e. modeling Eq. (3)) as  $\mathrm{NW}^{\mathrm{B}}$ , and models which additionally condition on environment (i.e. modeling both Eq. (3) and Eq. (4)) as  $\mathrm{NW}_{\mathrm{e}}^{\mathrm{B}}$ . For  $\mathrm{NW}_{\mathrm{e}}^{\mathrm{B}}$  models, we train explicit and implicit variants. For all NW algorithms, we perform evaluation on all inference modes. In addition, for completeness, we experiment on a variant where we freeze the feature extractor and finetune a linear probe on top of the learned representations on the same training data  $\mathcal{D}_{tr}$ , which we refer to as "Probe". As an example, the implicit variant of  $\mathrm{NW}_{\mathrm{e}}^{\mathrm{B}}$  is trained on Eq. (8), where the support set is balanced across classes (B) and conditioned on an environment (e).
We set  $N_{c} = 8$  for Camelyon-17 and ISIC and  $N_{c} = 1$  for FMoW. An analysis of this hyperparameter is provided in the Appendix. The query batch size  $N_{q}$  is set to 8 for all NW experiments. For Random and Cluster inference modes, we set  $k = 3$ . This was chosen based on prior work [55], where  $k = 3$  was shown to balance good error rate performance with computational efficiency. For explicit variants, we tune  $\lambda$  for all datasets via grid search on a held-out validation set. Full hyperparameter details are provided in the Appendix.
All training and inference is done on an Nvidia A6000 GPU and all code is written in Pytorch.
Table 2 shows the main results. We find that on Camelyon-17 and ISIC datasets,  $\mathrm{NW_e^B}$  with Full mode outperforms all baselines and variants we consider. In addition,  $\mathrm{NW_e^B}$  variants typically have lower variance across random seeds as compared to baselines. For FMoW,  $\mathrm{NW_e^B}$  with Ensemble mode performs around  $2\%$  lower than the best performing baseline, CLoVE. We observe that the most computationally-efficient inference mode, Cluster, performs comparably to Full mode for  $\mathrm{NW_e^B}$  models, and is in fact the highest-performing model for ISIC. Thus, we conclude that the support set can be an efficient replacement for Full.
For ISIC, we find that almost all  $\mathrm{NW}^{\mathrm{B}}$  modes (except Random) perform  $\sim 3\%$  better than ERM. This may be attributed to balancing classes across environments, which we suspect has added benefit for highly imbalanced tasks. In contrast, this boost is less apparent for Camelyon-17, which has relatively balanced classes. As an ablation, we compare  $\mathrm{NW}^{\mathrm{B}}$  against an NW variant without class-balancing in the Appendix.  $\mathrm{NW}_{\mathrm{e}}^{\mathrm{B}}$  further improves over  $\mathrm{NW}^{\mathrm{B}}$  by  $\sim 7\%$ . Exploring further, we compare  $\mathrm{NW}^{\mathrm{B}}$  against an ERM variant with balanced classes per environment, which we denote  $\mathrm{ERM}^{\mathrm{B}}$ . This
<sup>5</sup>https://www.isic-archive.com
Our code is available at https://github.com/alanqrwang/nwhead.
achieves  $63.0 \pm 2.5$ , which is on-par with  $\mathrm{NW^B}$ . This is expected as the theoretical assumptions are the same for both models.
Comparing implicit to explicit variants of  $\mathrm{NW_e^B}$ , we do not find much difference in explicitly enforcing Eq. (1), although we do observe significantly lower variances across model runs. Generally, we find the slight performance gain of explicit training to not be worth the computational overhead of doubling the number of support set forward passes per gradient step and tuning the hyperparameter  $\lambda$ .
While not the highest-performing, we highlight the consistent  $1 - 5\%$  improvement of Probe models over IRM, indicating that  $\mathrm{NW_e^B}$  may be better at capturing invariant features. However, other nonparametric inference modes still outperform Probe, possibly indicating that the learned features are more suitable for NW-style classifiers.
There are several advantages of the proposed NW approach over previous works. First, the implicit training strategy in Eq. (8) has no hyperparameter to tune, while remaining competitive with and often outperforming state-of-the-art baselines which all require tuning a hyperparameter coefficient in the regularized loss. Second, the NW head enables interpretability by interrogating nearest neighbors in the feature space. Since these neighbors directly contribute to the model's prediction (Eq. (2)), interrogation enables a user to see what is driving the model's decision-making. This not only allows for greater model transparency, but also enables interrogating the quality of the invariant features. We explore this capability in Section H in the Appendix. Note this this degree of transparency is not present in parametric baselines. Lastly, from an intuitive standpoint, we believe our non-parametric approach to enforcing invariance across environments is more natural than baseline methods, since an environment is encoded by manipulating the support set to contain real samples only from that environment. Other baseline methods resort to proxy methods to enforce invariance [54, 48, 27].
One important limitation of our method is computational (see Appendix for analysis of runtimes). The proposed approach requires pairwise comparisons, which scales quadratically with sample size. Practically, this means passing a separate support mini-batch in addition to a query mini-batch at every training iteration. This limitation is compounded for explicit variants, in which two sets of support sets must be drawn independently. Future work may explore more computationally-efficient training strategies. At inference time, Full, Cluster, and Ensemble modes are expensive procedures which require computing features for the entire support set, although precomputing features can mitigate this. However, we argue that in high-risk, safety-critical domains like medical imaging, high inference throughput may not be as important as performance, interpretability, and robustness.
We expect the proposed approach to work well with tasks that have several (and diverse) sets of examples per label class in each environment. If this is not the case, as in the FMoW dataset, the resulting model will be sub-optimal. In particular, in the extreme case where no example is present for a specific class in a given environment, constructing a support set with labels that cover the ground truth label of the query images will not always be possible. This will, in turn, impact performance.
We presented a nonparametric strategy for invariant representation learning based on the Nadaraya-Watson (NW) head. In the NW head, the prediction is made by comparing the learned representations of the query to the elements of a support set that consists of labeled data. We demonstrated two possible ways of manipulating the support set, and demonstrate how this corresponds to encoding different assumptions from a causal perspective. We validated our approach on three challenging and real-world datasets.
We believe there are many interesting directions of further research. First, our treatment is restricted to classification tasks. Future work may explore an extension to the regression setting. Second, it can be interesting to explore adaptation to the test domain, given additional information. For example, reweighting the occurrence of samples per label could provide improved results given knowledge about the edge  $E \rightarrow Y$  in the test distribution. One can further envision implementing the proposed method in settings where there are previously unseen test time labels/tasks. Finally, we are interested in replacing the fixed similarity function with a learnable kernel.
Funding for this project was in part provided by the NIH grant R01AG053949, and the NSF CAREER 1748377 grant.
[29] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Learning to generalize: Meta-learning for domain generalization, 2017.
[30] Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C. Kot. Domain generalization with adversarial feature learning. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5400-5409, 2018. doi: 10.1109/CVPR.2018.00566.
[31] Zachary C. Lipton, Yu-Xiang Wang, and Alex Smola. Detecting and correcting for label shift with black box predictors, 2018.
[32] Mingsheng Long, Yue Cao, Jianmin Wang, and Michael I. Jordan. Learning transferable features with deep adaptation networks, 2015.
[33] Yu. A. Malkov and D. A. Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs, 2018.
[34] Lucas Mansilla, Rodrigo Echeveste, Diego H. Milone, and Enzo Ferrante. Domain generalization via gradient surgery, 2021.
[35] E. A. Nadaraya. On estimating regression. Theory of Probability & Its Applications, 9(1): 141-142, 1964. doi: 10.1137/1109020. URL https://doi.org/10.1137/1109020.
[36] Timothy Nguyen, Zhourong Chen, and Jaehoon Lee. Dataset meta-learning from kernel ridge-regression, 2021.
[37] Nicolas Papernot and Patrick McDaniel. Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning, 2018. URL https://arxiv.org/abs/1803.04765.
[38] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer, 2018. URL https://arxiv.org/abs/1802.05751.
[39] Judea Pearl. Causality. Cambridge University Press, 2 edition, 2009. doi: 10.1017/CBO9780511803161.
[40] Jonas Peters, Peter Buhlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals, 2015.
[41] Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning, 2018.
[42] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization, 2021.
[43] Veronica Rotemberg, Nicholas Kurtansky, Brigid Betz-Stablein, Liam Caffery, Emmanouil Chousakos, Noel Codella, Marc Combalia, Stephen Dusza, Pascale Guitera, David Gutman, et al. A patient-centric dataset of images and metadata for identifying melanomas using clinical context. Scientific data, 8(1):1-8, 2021.
[44] Bernhard Schoelkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning, 2012.
[45] A Scope, AA Marghoob, CS Chen, JA Lieb, MA Weinstock, AC Halpern, and SONIC Study Group. Dermoscopic patterns and subclinical melanocytic nests in normal-appearing skin. British Journal of Dermatology, 160(6):1318-1321, 2009.
[46] Yuge Shi, Jeffrey Seely, Philip H. S. Torr, N. Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization, 2021.
[47] Jake Snell, Kevin Swersky, and Richard S. Zemel. Prototypical networks for few-shot learning, 2017. URL https://arxiv.org/abs/1703.05175.
[48] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation, 2016.
[68] Quanshi Zhang, Yu Yang, Haotian Ma, and Ying Nian Wu. Interpreting cnns via decision trees, 2018. URL https://arxiv.org/abs/1802.00121.
[69] Fan Zhou, Zhuqing Jiang, Changjian Shui, Boyu Wang, and Brahim Chaib-draa. Domain generalization via optimal transport with metric similarity learning. Neurocomputing, 456: 469-480, oct 2021. doi: 10.1016/j.neucom.2020.09.091. URL https://doi.org/10.1016/j.neucom.2020.09.091.
[70] Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Deep domain-adversarial image generation for domain generalisation, 2020.
[71] Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1-20, 2022. doi: 10.1109/tpami.2022.3195549. URL https://doi.org/10.1109/2Ftpami.2022.3195549.
A Description of Baselines
- Empirical risk minimization (ERM) [52] minimizes the sum of errors across domains and examples.
- Invariant risk minimization (IRM) [1] learns a feature representation such that the optimal linear classifier on top of that representation matches across domains. For WILDS datasets, we pull baseline performance from [54]. For ISIC, we use the implementation from [16].
- Deep CORAL [48] penalizes differences in the means and covariances of the feature distributions (i.e., the distribution of last layer activations in a neural network) for each domain. For WILDS datasets, we pull baseline performance from [54]. For ISIC, we use the implementation from [16].
- Fish targets domain generalization by maximizing the inner product between gradients from different domains. For WILDS datasets, we pull baseline performance from the original paper. For ISIC, we use the implementation from [16].
- LISA augments the set of training data by randomly performing two types of mixup-style [67] interpolations: intra-label (same label, different domain) and inter-label (same domain, different label). For WILDS datasets, we pull baseline performance from the original paper, and train their implementation on the ISIC dataset.
- CLOvE [54] finds an invariant classifier by enforcing the classifier to be calibrated across all training domains. While the original paper proposes several model variants leveraging this idea, we report their best-performing variant, which starts with a trained CORAL model and finetunes the weights using a regularized cross-entropy loss. The regularizer aggregates Maximum Mean Calibration Error (MMCE) [28] over all training domains. For WILDS datasets, we pull baseline performance from [54]. As their implementation is not publicly-available, we implement it for ISIC.
B Description of Datasets
Representative examples of the 3 datasets are shown in Fig. 4.
- Camelyon-17. We use Camelyon-17 from the WILDS benchmark [4, 23], which provides 450,000 lymph-node scans sampled from 5 hospitals. Camelyon-17 is a medical image classification task where the input  $x$  is a  $96 \times 96$  image and the label  $y$  is whether there exists tumor tissue in the image. The environment denotes the hospital that the patch was taken from. The training dataset is drawn from the first 3 hospitals, while out-of-distribution validation and out-of-distribution test datasets are sampled from the 4th hospital and 5th hospital, respectively.
- ISIC. The melanoma dataset is from the International Skin Imaging Collaboration (ISIC) archive<sup>7</sup>. Data from the archive are collected by different organizations at different points in time [7–9, 17, 43, 45, 50]. There are about 70k data samples in total. In particular, the resized input image  $x$  is a  $224 \times 224$  image and a binary target label  $y$  denotes whether the image exhibit is melanoma or not. The environment is the hospital from which the image was collected<sup>8</sup>. We follow a similar setup to Camelyon-17. The training dataset is drawn from the first 3 hospitals, while out-of-distribution validation and out-of-distribution test datasets are sampled from the 4th hospital and 5th hospital, respectively. For preprocessing, we filter out datapoints that are not specifically categorized as "benign" or "malignant" (e.g. "indeterminate"). The OOD validation dataset is from the "Barcelona1" site indicator and the OOD test dataset is the "Vienna1" site indicator.
- FMoW. The FMoW dataset is from the WILDS benchmark [6, 23], a satellite image classification task which includes 62 classes and 80 domains (16 years x 5 regions). Concretely, the input  $x$  is a  $224 \times 224$  RGB satellite image, the label  $y$  is one of the 62 building or land use categories, and the environment represents the year that the image was taken as well
<sup>7</sup>https://www.isic-archive.com
<sup>8</sup>Hospitals are Hospital Clinic of Barcelona, Medical University of Vienna, University of Queensland Diamantina Institute, Memorial Sloan Kettering Cancer Center, University of Sydney Melanoma Diagnostic Centre, and University of Pittsburgh Medical Center.
as its corresponding geographical region – Africa, the Americas, Oceania, Asia, or Europe. The train/test/Validation splits are based on the time when the images are taken. Specifically, images taken before 2013 are used as the training set. Images taken between 2013 and 2015 are used as the validation set. Images taken after 2015 are used for testing.
Hospital 1
Hospital 2
Hospital 4
(a) Camelyon-17
Hospital 5
Hospital 3
(b) ISIC
Hospital 4
Region 3
(c) FMoW
Figure 4: Representative images for datasets, separated by domain. Each row depicts a separate class. For FMoW, for simplicity, we show 2 classes out of 62 and only images before 2013.
Region 5
C Hyperparameter Details
Table 3 shows hyperparameter settings for all datasets, where NW-specific hyperparameters are below the midline. For all models, we use pretrained ImageNet weights. For  $\lambda$ , we perform a grid-search over the values  $\{0.01, 0.1, 1\}$ . Fig. 5 depicts  $\mathrm{NW_e^B}$  performance vs  $N_{c}$  for Camelyon-17 and ISIC datasets. We find that performance is relatively insensitive to  $N_{c}$  above  $\sim 5$  examples per class.
In the original NW head paper [55], the authors experiment with a temperature (i.e. bandwidth) hyperparameter  $\tau$ . In this work, we set  $\tau = 1$  for all experiments. The reason for this is that we optimize both the feature extractor and classifier end-to-end, and the kernel used in the classifier is dependent on the features that the feature extractor learns (unlike [25], e.g.). Thus, we let the feature extractor to optimize the bandwidth on its own. Note this same approach is taken in prior works [47].
Table 3: Hyperparameter settings for various datasets.
Hyperparameter	Camelyon-17	ISIC	FMoW
Learning rate	1e-4	5e-5	1e-4
Weight decay	1e-4	0	1e-2
Scheduler	None	None	StepLR
Batch size	32	8	8
Architecture	DenseNet-121	ResNet-50	DenseNet-121
Optimizer	SGD	Adam	Adam
Maximum Epoch	10	5	60
Nq	8	8	8
Nc	8	8	1
Ns	Nc×2=16	Nc×2=16	Nc×62=62
λ	0.01	0.01	0.1
k	3	3	3
Figure 5:  $\mathrm{NW_e^B}$  performance vs  $N_{c}$  for Camelyon-17 and ISIC datasets. Full mode. Performance is relativity insensitive to  $N_{c}$  above  $\sim 5$  examples per class.
D Table of Runtimes
Table 4 shows approximate runtimes for various datasets during training and inference. All experiments are performed on a GPU.
Table 4: Approximate runtimes for various algorithms. Training time is time to complete maximum epochs as specified in Table 3, and does not include validation. Inference time is time to evaluate the entire test set. Averaged over all training runs.
Algorithm	Camelyon-17	ISIC	FMoW
Training	ERM	7 hr	1 hr	22 hr
NW	14 hr	2 hr	40 hr
Inference	ERM	10 min	2 min	10 min
NW, Random	15 min	3 min	20 min
NW, Full	2 hr	15 min	1 hr
NW, Ensemble	2 hr	15 min	1 hr
NW, Cluster	2.2 hr	17 min	1.1 hr
NW, Probe	10 min	2 min	10 min
E ID vs. OOD Performance
With most invariant learning methods, prior works observe a tradeoff between in-distribution and out-of-distribution generalization. In Table 5, we show in-distribution (ID) and out-of-distribution (OOD) results for Camelyon-17, which provides an ID validation set. We observe that there is a tradeoff between ID and OOD performance, similar to prior work.
Table 5: In-distribution vs. out-of-distribution generalization performance on Camelyon-17.
In-distribution	Out-of-distribution
ERM	93.2±5.2	70.3±6.4
NWb, Full	96.1±1.0	72.0±6.7
NWb, Full	92.8±2.0	80.0±2.7
F Imbalanced ISIC Experiments
As ISIC exhibits significant label imbalance where the positive class is much less represented than the negative class (see Fig. 6), we experiment with an NW variant without support-set label balancing as an ablation. For both variants, we set  $N_{c} = 8$ . To train the imbalanced variant, we sample a mini-batch support set by first sampling one image per class (to guarantee both classes are represented in the support at least once), and then sampling the rest of the images randomly from the dataset.
To characterize the performance of both variants across class imbalances, we change the prevalence of  $y = 0$  in the test set by removing negative class images until the desired prevalence is achieved. Note the default prevalence is  $\sim 0.85$ . Then, we compute the accuracy over this manipulated test set (note, prior work has shown that F1-score is not a good metric for comparing classifiers with different label imbalances [3]).
Fig 7 shows results. We observe that at high prevalence where the proportion of negative classes matches in training and test, the imbalanced variant outperforms the balanced variant, whereas the opposite is true for low prevalence. This makes sense because at high prevalence, the class imbalance is similar for the training and test domains; thus, a model which overpredicts the negative class is usually right. On the other hand, this fails in test sets where the prevalence is flipped (i.e. low prevalence). These results suggest that  $\mathrm{NW}^{\mathrm{B}}$  is a more robust classifier in the presence of label shift.
Figure 6: Number of datapoints separated by class for Camelyon-17 and ISIC datasets. There is significant label imbalance for the ISIC dataset.