Learning Invariant Representations with a Nonparametric Nadaraya-Watson Head
Authors: Alan Q. Wang
Minh Nguyen
Mert R. Sabuncu
Venue: NeurIPS 2023 poster (2023)
Keywords: Invariant representations
causality
domain generalization
TL;DR: We propose a nonparametric method for learning invariant representations using a Nadaraya-Watson prediction head.

Abstract: 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.

Learning Invariant Representations with a Madtsay Head Nonparametric Nadaraya- Watson
Alan Q. Wang Cornell University
Minh Nguyen Comell University
Mert R. Sabuncu Cornell University
Abstract
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 repre- sentations which are invariant across the different distributions, with the hope that these representations will be transportable to unseen domains. this work, we present nonparametric strategy for learning invariant representations based on the recently-pre accnntlyproposse Nadaraya- Watson (NW) head. The NW head makes prediction by comparing the learned representations of the query to the elements of 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 single environment encourages the model to learn invariant features that do not depend the environment. We present a causally motivvaedd setup for our modeling and training strategy and validate on three challenging real-world domain generalization tasks in computer vision.
Introduction
Machine learning models often fail when there is significant distribution shift. The goal of domain generalization to be able to perform well with new distributions [6] In this work, we are interested in settings where multiple domains/environments 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
available during training and have
representations will work well in, or are transportable to, unseen environments. This invariance is often encoded via constraints on 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 regularizer top of a standard predictive loss [ 20 27 48 [54] (8). 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 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 Nadaraya- Malarayya-Watoon (NW) head [55]. Instead of computing the class probability directly from input query, the NW head makes a prediction by comparing the learned representations of the query to the elements of 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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-speciff features to make 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 state-of-thi of-thh--rttttttt parametric baselines.
Our contributions are as follows:
We present causally- motiiveted assumptions for domain generalization which justify our modeling and training strategy.
We present a novel approach invariant representation learning using the nonparametric Nadaraya-Wann Watson head, which can account for causal assumptions by manipulating a support set. In partieular, 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-an art parametric baselines.
2 Related Works
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 7T]. One line of work augments or synthetically-general additional training images
63 65 70 In particular, to increase robustness of learned features to unseen environments [59
LISA uses a mixup-style [67] interpolation technique to generate augmented images, which the authors demonstrate improves out-of-distnn distiitbibuutoonnn robustness [64] Another line of work broadly seeks 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 metries such as maximum mean discrepancy [2] an adversarial loss [4 (30) and Wasserstein distance [69] Still other works approach the problem from the perspective of the 46 57. For example, Fish aligns the gradients from different
gradients and optimization [12
domains [46].
One can also achieve domain generalization via learning invariant representations, which often requires reasoning about the lata-generating process from a causal perspective to arrive at appropriate constraints (39. Invariant causal prediction (ICP) formulates the problem from 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 &o such that the optimal linear classifier's parameters w= on top of this representation is the same across all environments []] 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 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
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 decision trees (68], density estimation [3] and more general kernel-based 66 to deep networks and have explored the interpretability that these frameworks provide. Closely-related but orthogonal to nonparametrie models are attention- models, most notably self-attention mechanisms popularized in measffommrrbbbsee architectures in natural language
neighbors 37
methods 15
36

---

Environment' 2
Environment's 2
Eavironment
Environment
Fnvirunment
Environment
Favirunment
Environment
Environment
c)
Figure I: Illustration of proposed approach. Support set of labeled datapoints (square/triangle) from 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 function of distances to labeled datapoints in the feature space (visualized as dotted arrows). b) Balancing comparisons across labels for all environments models p XX). c) Conditioning on single environment models P. XY)
processing 53 and, more recently, computer vision attention nonparametric setting [24
11
19
38
Nonparametric transformers apply
Recently, Wang et al. proposed the NW head 55 an extension of the classical NW model
deep learning. In the NW head, the prediction is weighted average of labels from 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 the dominant approach of using parametric classifier with fully-connected layers. In this work, we leverage the NW head 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.
3 Preliminaries
E
Zs
E
Zs
X
X
Zc
Z.c
(a)
(b)
Figure 2: a) Causal Directed Acyclic Graph (DAG) we consider in this work. Solid nodes are observed and dashed nodes unobserved. We assume anti-causal setting where label Y causes X, and X has causal parents: "style" features, Zs. which are influenced by the environment E; and pnviroomeentindepeddeene "content" features of X, Ze, which are causally influenced by the label Y. E potentially influences Y. Both E and Y have direct influence on style features Zs. b) Same DAG a) with intervention Y. We note that 1 E Zo and Y y E Zs.
Problem Setting. Let X, denote a datapoint and corresponding discrete class, and denote the environment (or domain) where X, Y originates That is, the elements of the training dataset D.r = {x,, Vı, e, N are drawn first by sampling the discrete random variable e, P(E), and then sampling 24, Vi P(X,Y E = e,) == Pe, (X, Y). Our goal is to learn classifiers that will generalize to new, unseen environments.
Assumptions. We assume there exists pair of latent causal parents of X: environment- independent ("content")) factor Zç and an environment-dependerr ("style") factor Z 2 We as-
Wee assume the support of F,, supp F), is finite.
This is sometimes referred to style-content causal decomposition 21 42
60

---

sume the causal mechanism that generates X from (Ze, Zş) is injective, so that, in principle, it possible recover the latent features from the observations; i.c. there exists function g such that g(X) = (7c, 7:s). We further assume that g can be disentangled into 9C and Ss, such that (Zc, Zs) = ( = (gc(X) (X); gs gx))). 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 non-zero probability in all environments.
Motivation. The motivating application in this work is image classification, where each realization of E might represent 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 Zg (e.g. staining, markings, orientation). In addition, we allow Zs to be influenced by the label itself (for example, positive examples are more likely to marked by 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 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 and X, as these may change in new, unseen environment. That is, we seek an estimator which relies only on Zo and which is independent of E or Zs.
First, we note the direct causal dependence E + Y. For example, a model can exploit this association by learning to over-predict majority classes in certain environment. One way remove the direct dependence is by intervening on Y, thus removing incoming edges to Y. This essentially corresponds to matching the environment-speciffe prior on Y between environments, and results in the intervened graph in Fig. 2
Let us refer to any distribution which follows the DAG in Fig. 2p as P (X
Second, we observe that there is potential non-causal association flow between E and Y through the colliders and Zs, when either one of these are conditioned on (i.e. are observed). An estimator which relies on Zs potentially leaks information from E, and this is unstable in new environments. Reading d-separation on this intervened graph, we infer that Y 1 E Ze = (X) and Y L E Zs (X), that is:
(x)))) Ve, E E. Pa (Y (c(X)) = (Y
(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 E.
Thus, we seek estimator that 1) enables interventions on Y such that the direct dependence E - Y can be removed, and 2) can further encode the assumption in Eq. M.
NW Head for Invariant Prediction
Given datapoint x, support set S (I,, yi = and parameters o, the NW head estimates P(Y = y X = x) by outputting a prediction formulated as a weighted sum of support set labels, where the weights some notion of similarity in the feature space (55):
Enn exp (s(p(z), ((.)))
P(Y = y X
= x; b) = fo(o,
En exp (s((()), o(z;)))
Here, y is the one-hot encoded version of and s(,, -) is similarity/kemell 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 (see Fig. [p). This can be interpreted as an intervention Y, and removes the dependence on F. Y, i.e.:
p = y X 2; Φ) = (0 Sa).
This may be interpreted as making the model robust to label shift,
31 44

---

Support
Query
Envisumnnnt
Bevaoranen
Uneverement
NW Head
Mi x
Softmax
M02x
All Eavinnamns
Figure 3: A depiction of the NW head on a tumor detection task. The NW head computes Euclidean distances .(..)) between query and support features, and uses the distances 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 fixed environment (note similarity in color). Such support precludes the possibility of using environment--peccffc features to make prediction.
2. Conditioning S on a single environment, denoted S. (see Fig. [p). This can be interpreted conditioning the probability estimate E c, i.e.:
P. = = x; の) == (o(x,
Note that both balancing and conditioning can be achieved simultaneously, which denote S
4.1 Objective and Optimization Details
Given dataset of samples Der {zx,, St, Ei 1.W t=1 wish leverage the NW head conditional estimator for Y conditioned on Zc = 9c (X), where 9c(X) is characterized by Eq. ( This necessitates an optimization over both and the space functions gc. Thus, solve the following constrained maximum likelihood over á and gc:
Σ
log PB (ye (c(x:))
argmax 9.9G
(5)
2=1
s.t. PB (1 go (x:): Φ) P (y ())) Vi E fI,..., N), Ve,e' E.
Note that Eq. ① implies that P. (y: (c(x))) = PB (y: (x;)). Thus, the objective equivalent unconstrained maximum likelihood under the assumption in Eq. ①
Instead of solving for ge explicitly, we let both & and gc be related by the composition & = áo gc. and set y to be the learnable mapping of the NW head, i.e. neural network. Then, the objective becomes:
argmin
(6)
s.t. (Z,, SB) fa, Vi (L, Ve,e' E E, where L is the cross-entropy loss. To make the objective tractable, consider possible variants:
1. Explicit. Solve the optimization problem explicitly via Lagrangian formulation:
N argmiin L(Ja(s, S S.),)) +
Σ
[a, fo(x,, s
e
where 0 is a hyperparameter.
2. Implicit. Relax the optimization problem into the following unconstrained problem:
N
argmin
ΣV
S:)
(8)
6EE
5

---

In this formulation, the constraint will be approximately satisfied in the sense that the model will be encouraged to predict the ground truth for 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 Der. and all queries and support datapoints passed through the feature extractor (p 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₂ and N₂ the query and support mini-batch sizes respectively, the effective mini-batch size is N₃ + N,, instead of N, For the implicit variant, we sample one support set for 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 given mini-batch of queries, perform two independent forward passes through the NW head for each support set, and compute the loss in Eq. 77 As discussed in prior work [55]. the support batch size is hyperparameter analogous and orthogonal to the query batch size.
technical point is that the set of labels in the support mini-bater must cover the set of labels in the query mini-batch. Thus, in our implementation, for S cycle through all classes and randomly draw N₂ examples per class to include in the support. For tasks with a large number of classes, one subsample from the total number of classes, long 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 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 com- puted 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 distillation of the full training for efficient inference, with the caveat that the support set 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. our experiments, we find that Cluster mode can be sufficient replacement to Full mode, while being computationally cheaper. These inference modes can 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 A0 14 41 26). Representatively, under the binary classification setting, the IRM objective finds representation function y which elicits an invariant predictor across environments E such that for all that has non-zero probability for xxX) in any (and all) environment(s):
27
E. (xX) h] E.
p(X)= = h), Ve, E.
Eq. 11p can be viewed generalization of this equality multi-class settings
"This equality has also been called "sufficiency invariance"
60

---

Table 1: Summary of Datasets.
Dataset
# Classes
Env
# Envs
Architecture
Metric
Camelyon-17
2
Hospital
DenseNet-121
Average acc.
ISIC
2
Hospital
ResNet-50 DenseNet-122
Fi-score
FMoW
62
Region
Worst-region acc.
Table 2: Metric average + standard deviation for all datasets (%). underline second-best. Implicit / Explicit.
Higher is better. Bold is best and
Algorithm
Camelyon-17
ISIC
FMoW
ERM [52
70.3
58.2
32.6+16
IRM 1
70.9
57.9+1a
31.3+12
CORAL 48
72.4
59. =22
31.7+10
Fish 46
74.7=762
64.4
34.6100
LISA 64
77. 16.5
64.8
35.5
CLOvE [54
79.9+3.9
66.2
40.1
NW Random
71.725.
56.7+14
31.
NW Full
72.026.
61.9
31.6+09
NW Cluster
70.616.9
61.4123
31.3
NW Ensemble
71.9
63. 23.9
32.2+10
NW8 Probe
69.2274
59. 39.7222
29.9
NW Random
74.8 75.31132
57.5 55.0a09
31.2 /0.9
NW Full
80.0 /9..2/77777.. 79.7.119
69.6 70.0
35.0107 34.6.10.
NW Cluster
78.6#25 79.0+14
71.1 71.0+10
33.94a6 34.040a
NW Ensemble
79.5 79,6+199
69.5 69,84ax
37.8
/tt.////..000 38.2+04
NW Probe
75.3 75.8
61.4 / 63.4
33.9 32.7
Furthermore, note that given the feature extractor y, the NW mechanism f is nonlearnable classifier, whereas is learned in the IRM setting. Thus, proposed objective can be interpreted as learning invariant features ip, where the fixed classifier constraint is satisfied by construction. This avoids the need to approximate the complex bilevel optimization problem with regularizer which assumes convexity and requires computing the Hessian. Essentially, f enforces invariance through the manipulation of the support set, providing more intuitive and computationally simpler objective to optimize.
In the Experiments section, compare IRM against variant of our algorithm where freeze the learned representations and finetune a linear classifier on top using same training data. We find that our algorithm performs better than IRM on all datasets, suggesting that it captures invariant representations better than IRM.
5
Experiments and Results
5.1
Baselines
We compare against several popular and competitive baseline algorithms: empirical risk minimization
(ERM) (52)
invariant risk minimization (IRM) [I] 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 bench- mark (23), and the third challenging melanoma detection task. Details on the datasets are summarized in Table I and further information is provided in the Appendix.

---

1. The Camelyon-11 17 dataset A comprises of microscopic images of stained tissue patches from different hospitals, where the label corresponds 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 The ISIC dataset comprises of dermoscopic images of skin lesions from different hospitals, where the label corresponds to whether not the lesion is diagnosed 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,
model selection techniques, and report metrics specified by the benchmark. This includes using a DenseNet-122 121 backbone initialized with pretrained ImageNet weights as and random augmentations for both datasets. Similarly for ISIC, backbone as with no augmentations, and perform model selection on an OOD validation set. Due significant label imbalance, we report Fl-score instead of average accuracy.
pretrained ResNet-50
For NW algorithms, we refer to models which balance classes (i.e. modeling Eq. (3) as NW and models which additionally condition environment (i.e. modeling both Eq. and Eq. 4 as NW For NW models, we train explicit and implicit variants. For all NW algorithms, we perform evaluation on all inference modes. In addition, for completeness, we freeze the feature extractor and finetune a linear probe on top of the learned representations on the same training data Dar, which we refer as "Probe". As an example, the implicit variant of NW is trained on Eq. (8), where the support set is balanced across classes (B) and conditioned on an environment (e).
experiment
variant where
We set N₂ = for 1amelyon- and ISIC and N, for FMoW. An analysis this hyperparameter is provided in the Appendix. The query batch size Ng is set to 8 for all NW experiments. For Random and Cluster inference modes, we set &o = 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 À for all datasets via grid search on a held-out validation set. Full hyperparameter details provided in the Appendix.
All training and inference done on an Nvidia A6000 GPU and all code is written in Pytorch
5.4 Results
Table 2 shows the main results. We find that on Camelyon- 17 and ISIC datasets, NW with Full mode outperforms all baselines and variants we consider. In addition, NWE variants typically have lower variance across random seeds as compared to baselines. For FMoW, NW! with Ensemble mode performs around 2% lower than the best performing baseline, CLoVE. We observe that the most computationally-efiiii models, and is in fact the highest-performet can be an efficient replacement for Full.
inference mode, Cluster, performs comparably to Full mode for NW! model for ISIC. Thus, conclude that the support set
For ISIC, we find that almost all NW modes (except Random) perform ~ 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 NW against an NW variant without blassbbalnncing in the Appendix. NWB further improves over NWB by 7%. Exploring further, we compare NW against an ERM variant with balanced classes per environment, which we denote ERM This
http:://www..siiarrrhhiieee
"Our code is available at httpp:s://// //tg://///ttbbb www.coom//aaa. honm/alllngrvaa/uuhedd
8

---

achieves 63.0 + 2.5, which is on on-par with NW This is expected as the theoretical assumptions are the same for both models.
Comparing implicit to explicit variants of NW we do not find much difference in explicitly enforcing Eq. (T). although we do observe significantly lower variances across model runs. Generally, 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 A.
While not the over IRM, indicating that N may be better at capturing invariant features. However, other non- parametric inference modes still outperform Probe, possibly indicating that the learned features are more suitable for NW-style classifiers.
we highlight the consistent 1-5% improvement Probe models
Discussion and Limitations
There are several advantages of the proposed NW approach over previous works. First, the implicit training strategy in Eq. (8) has no hyperparameter tune, while remaining competitive with and often outperforming state-of-the-at 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 see what is driving the model's 's perision-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 present in parametric baselines. Lastly, from an intuitive standpoint, we believe our parammamtttic approach to enforeing 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
the Appendix. Note this this degree of transparency not
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 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-effiii 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 important performance, interpretability, and robustness.
We expect the proposed approach 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 specific class in a given environment, constructing 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.
7 Conclusion
We presented nonparametric strategy for invariant representation learning based on the Nadaraya- Watson (NW) head. the NW head, the prediction is made by comparing the learned representations of the query to the elements of 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 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 classification tasks. Future work may explore an extension to the regression setting. Second, 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 - 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 learnable kernel.

---

Acknowledgements
Funding for this project was in part provided by the NIH grant RO1AG053949, and the NSF CAREER 1748377 grant.
References
[1] Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk mini- mization, 2019. URL https://arxiv... wttp:/iw...g//b///990.///// 190.0098933333311
[2] Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statistics). Springer-Verlag, Berlin, Heidelberg, 2006. ISBN 0387310738.
[3] Jan Brabec, Tomáš Komárek, Vojtech Franc, and Lukáš Machlica. On model evaluation under non-constant class imbalance, 2020.
[4] Péter Bándi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, Quanzheng Li, Farhad Ghazvinian Zanjani. Svitlana Zinger, Keisuke Fukuta, Daisuke Komura, Vlado Ovtcharov. Shenghua Cheng, Shaoqun Zeng, Jeppe Thagaard, Anders B. Dahl, Huangjing Lin, Hao Chen, Ludwig Jacobsson, Martin Hedlund, Melih Çetin, Eren HalIIı, Hunter Jackson, Richard Chen, Fabian Both, Jörg Franke, Heidi Küsters- Vandevelde, Willem Vreuls, Peter Bult, Bram van Ginneken, Jeroen van der Laak, and Geert Litjens. From detection of individual metas- tases classification of lymph node status the patient level: The camelyon17 challenge. IEEE Transactions on Medical Imaging. 38(2):550-500 560, 2019. doi: 10.1 10.1109//L.22118888777777
[5] Mathieu Chevalley, Charlotte Bunne, Andreas Krause, and Stefan Bauer. Invariant causal mechanisms through distribution matching, 2022.
[6] Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world, 2018.
[7] Noel Codella, Veronica Rotemberg, Philipp Tschandl, M Emre Celebi, Stephen Dusza, David Gutman, Brian Helba, Aadi Kalloo, Konstantinos Liopyris, Michael Marchetti, et al. Skin lesion analysis toward melanoma detection 2018: A challenge hosted by the international skin imaging collaboration (isic). Eprint arxiv: 1902.03368 2019.
[8] Noel CF Codella, David Gutman, M Emre Celebi, Brian Helba, Michael A Marchetti, Stephen W Dusza, Aadi Kalloo, Konstantinos Liopyris, Nabin Mishra, Harald Kittler, et al. Skin lesion analysis toward melanoma detection: A challenge the 2017 international symposium biomedical imaging (isbi), hosted by the international skin imaging collaboration (isic). In Proceedings of ISBI, pages 168-172. IEEE, 2018.
[9] Mare Combalia, Noel CF Codella, Veronica Rotemberg, Brian Helba, Veronica Vilaplana. Ofer Reiter, Cristina Carrera, Alicia Barreiro, Allan C Halpern, Susana Puig, et al. Bcn20000: Dermoscopic lesions in the wild. Eprint 110: 02288 2019.
[10] Andreas Damianou and Neil D. Lawrence, Deep gaussian processes. In Carlos M. Carvalho and Pradeep Ravikumar, editors, Proceedings of the Sixteenth International Conference Artificial Intelligence and Statistics, volume 31 of Proceedings of Machine Learning Research,
pages 207-215, Scottsdale, Arizona, USA, 29 Apr-01 May 2013. PMLR. URL https //proceedings. mlr. pwwps/vv.l1 satess/v31/damiannooo dwouianninn..mmmmmmmmmmmmmm html
[11] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2020. URL htttt://aaxxii hrwpiiiworr////ee httpg/abs/200.... 11929
[12] Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Do- main generalization via model-agnostic learning of semantic features. In H. Wal- lach, Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett, edi- tors, Advances in Neural Information Processing Systems, volume 32. Curran Associates,
10

---

Inc.. 2019. URL httpo///tttt////r/coooiing.gaase
pdf
[13] Rasool Fakoor, Pratik Chaudhari, Jonas Mueller, and Alexander J. Smola. Trade: Transformers for density estimation, 2020. UtL/Atttss///radr..or//////////p 02441
[14] Yaroslav Ganin, Evgeniya Ustinova, Hama Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-aaveessrial training of neural networks, 2016.
[15] J. Ghosh and A. Nag. An Overview of Radial Basis Function Networks, pages 1-36. Physica- Verlag HD, Heidelberg, 2001. ISBN 978-3-7908 1826-0. doi: 10.1007/978-3 790---888----0-00 1826-0_ I_ URL https://doi. org
1826-0_1
[16] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization, 2020.
[17] David Gutman, Noel CF Codella, Emre Celebi, Brian Helba, Michael Marchetti, Nabin Mishra, and Allan Halpern. Skin lesion analysis toward melanoma detection: A challenge at the international symposium on biomedical imaging (isbi) 2016, hosted by the international skin imaging collaboration (isic). Eprint arXiv:: 1605.01397 2016.
[18] Christina Heinze- Deeezee-Dmmlle and Nicolai Meinshausen. Conditional variance penalties and domain shift robustness, 2019.
[19] Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention, 2021. URL https://arxiv org/abs/2103 02220 03206
[20] Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts, 2022. URL. httppt:st////xiit org/ abs/ 20..:2227.0020 01603
[21] Jean Kaddour, Aengus Lynch, Qi Liu, Matt J. Kusner, and Ricardo Silva. Causal machine learning: A survey and open problems, 2022. URI.https:/ //axiisv//axxv.... http//aaavvv.rraaaa/a/// http/abb/2200.... 2206 15475
[22] Pritish Kamath, Akilesh Tangella, Danica J. Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance?, 2021.
[23] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness. Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson. Sergey Levine, Chelsea Finn, and Percy Liang. Wilds: A benchmark of in-the-wile distribution shifts, 2021.
[24] Jannik Kossen, Neil Band, Clare Lyle, Aidan N. Gomez, Tom Rainforth, and Yarin Gal. Self- attention between datapoints: Going beyond individual input-output pairs in deep learning, 2021. URL //tp:://///xiii http:::://axxvv..rrre. htce//bs//200 /206. 22111020864444
[25] Nikita Kotelevskii, Aleksandr Artemenkov, Kirill Fedyanin, Fedor Noskov, Alexander Fishkov, Artem Shelmanov, Artem Vazhentsev, Aleksandr Petiushko, and Maxim Panov. Nonparametric uncertainty quantification for single deterministic neural network, 2022.
[26] Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an out-of-distribution generalization problem ?, 2021.
[27] David Krueger, Ethan Caballero, Joern-eenrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang. Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex), 2021.
[28] Aviral Kumar, Sunita Sarawagi, and Ujjwal Jain. Trainable calibration measures for neu- ral networks from kernel mean embeddings. Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Pro- ceedings of Machine Learning Research, pages 2805-2814, PMLR, 10-15 Jul 2018. URL. https:///roceee
mlr. wwwpps/eev.0.. vatss/80//oaa 2t000//mmr118... html

---

[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: 11110//VVR.2218.000555
[31] Zachary C. Lipton, Yu- 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 generaliza- tion via gradient surgery, 2021.
[35] E. A. Nadaraya. On estimating regression. Theory of Probability & Its Applications, 9(1): 14I-142, 1964. doi: 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, inter- pretable and robust deep learning, 2018.
04765
[38] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Eukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer, 2018. URL hRtpss///aax org/ /ttp////wiw..rr/aaa//////000// httg//bs/1802.05777
[39] Judea Pearl. Causality. Cambridge University Press, 2 edition, 2009. doi: 10.1017/
CBO97805 11803166.
[40] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals, 2015.
[41] Mateo Rojas- 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 minimiza- tion, 2021.
[43] Veronica Rotemberg, Nicholas Kurtansky, Brigid Betz-Stablein, Liam Caffery, Emmanouil Chousakos, Noel Codella, Marc Combalia, Stephen Dusza, Pascale Guitera, David Gutman, al. A patient-centric dataset of images and metadata for identifying melanomas using clinical context. Scientific data, 8(1):188, 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 Halpem, and SONIC Study Group. Dermoscopic patterns and subclinical melanoeytic nests in normal-appearing skin. 2009.
British Journal of Dermatology, (3006):1318-133)
[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 battps://arie org/ abs /1///703.. .01755
[48] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation,
2016.
12

---

[49] Mohammad Reza Taesiri, Giang Nguyen, and Anh Nguyen. Visual correspondence-bue based explanations improve Al robustness and human-Al team accuracy. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL.itttps:/openreviie wwtp:e/oouummdddUV//H//YYY
[50] Philipp Tschandl, Cliff Rosendahl, and Harald Kittler. The ham 10000 dataset, large collection of multi-source dermatoscopic images of common pigmented skin lesions. Scientific data, 5(1):
1-9, 2018.
[51] Eric Tzeng, Judy Hoffman, Ning Zhang, Kate Saenko, and Trevor Darrell. Deep domain confusion: Maximizing for domain invariance, 2014.
[52] V.N. Vapnik. An overview of statistical learning theory. IEEE Transactions on Neural Networks, 10(5):988-99))) 1999. doi: 10. 19.1009977.886440.
[53] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017. URL https://arxxi
org/abs/ httaaabss/10....0 http///10.0077666.
[54] Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit, On calibration and out-of-domain generalization, 2021. (tRpsttts://
21bp.//22..//.///22...... 10395
[55] Alan Q. Wang and Mert R. Sabuncu. A flexible nadaraya head can offer explainable and calibrated classification. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
URL
[56) Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip S. Yu. Generalizing to unseen domains: survey on domain generalization, 2022.
[57] Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization, 2023.
[58] Yixin Wang and Michael Jordan. Desiderata for representation learning: A causal perspective, 2022.
[59] Yufei Wang, Haoliang Li, and Alex C. Kot. Heterogeneous domain generalization via domain mixup. In ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 3622-3626, 2020. doi: 10. N011000CCASSP44066200000663333
[60] Zihao Wang and Victor Veitch. The causal structure of domain invariant supervised representa- tion learning, 2022. URL /attps//arriiv.oe htrg/abb/22008 06987
[61] G. S. Watson. Smooth regression analysis. Sankhya: The Indian Journal of Statistics, Series A (26):399-372, 1964.
[62] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P. Xing. Deep kernel learning, 2015. URL https://arxx wwterrvvv...or.ooooo aatp//1bb/1111.. .22222
[63] Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang. Adversarial domain adaptation with domain mixup, 2019.
[64] Huaxiu Yao, Yu Wang. Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation, 2022.
[65] Xiangyu Yue, Yang Zhang, Sicheng Zhao, Alberto Sangiovanni-Vincetellii Boqing Gong, Domain randomization and pyramid consistency: Simulation-to-real generaliza- tion without accessing target domain data, 2022.
Vincentelli, Kurt Keutzer, and
[66] Dequan Zhang, Ning Zhang, Nan Ye, Jianguang Fang, and Xu Han. Hybrid learning algorithm of radial basis function networks for reliability analysis. IEEE Transactions on Reliability, 70 (3)8887-900, 2021. doi: 2111000.22200000000
[67] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization, 2018.
13

---

[68] Quanshi Zhang, Yu Yang, Haotian Ma, and Ying Nian Wu. Interpreting emns via decision trees, 2018. URLhttpps /ttp///axvvvv... www.aaawwww....a....rrra. arg/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.10 www.:euuum....000.... URL pttttps:/ddoi.or//100/
2F neucom. 2020. 09. 091
[70] Kaiyang Zhou, Yongxin Yang. Timothy Hospedales, and Tao Xiang. Deep domain-adversanii 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. 00.11099ppam.2222331955444 URL httph::t/////tt org/ 10. 1109%/2Ftpamm 2022. 3195549
14

---

A Description of Baselines
Empirical risk minimization (ERM) 52 minimizes the sum of errors across domains and
examples.
Invariant risk minimization (IRM) learns a feature representation such that the optimal linear classifier 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 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 [[]]
LISA augments the set of training data by randomly performing two types of mixup- style 66 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, report their best-perlorming variant, which starts with 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 [44 As their implementation is not publicly-available, implement it for ISIC.
B Description of Datasets
Representative examples of the 3 datasets are shown in Fig.
Camelyon-17, We use Camelyon-17 from the WILDS benehmark 4 (3], which provides 450,000 lymph-node scans sampled from 5 hospitals. Camelyon-17 is medical image classification task where the input æ is a 96 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 drawn from the first 3 hospitals, while out-of-distribution validation and -uttof-distribution test datasets are sampled from the 4th hospital and 5th hospital, respectively.
ISIC. The melanoma dataset from the International Skin Imaging Collaboration (ISIC) archive Data from the archive are collected by different organizations at different points in time 9
45 50). There are about 70k data samples in total. In particular, the resized input image x: is 224 x 224 image and binary target label y denotes whether the image exhibit is melanoma or not. The environment is the hospital from which the image was collected [ We follow 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" (c.g. "indeterminate"). The OOD validation dataset is from the "Barcelonal" site indicator and the OOD test dataset is the "Viennal" site indicator.
FMoW. The FMoW dataset from the WILDS benchmark 6 23 a satellite image classi- fication task which includes 62 classes and 80 domains (16 years 5 regions). Concretely, the input x is 224 224 RGB satellite image, the label y one of the 62 building or land use categories, and the environment represents the year that the image was taken well
"ptts://www.isie-rrchii archive.com
'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.
15

---

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
Hospital
Hospital
Hospital
Hospital
(a) Camelyon-1 17
Hospital
Hospital
Hospital
Hospital
Hospital 5
(b) ISIC
Region
Region
Region
Region
Region
(c) FMoW
Figure 4: Representative images for datasets, separated by domain. Each row depicts separate class.
For FMoW, for simplicity, we show 2 classes out and only images before 2013.
16

---

C Hyperparameter Details
Table shows hyperparameter settings for all datasets, where NW-specific hyperparameters are below the midline. For all models, we use pretrained ImageNet weights. For A, we perform a grid-search over the values (0.01, 0.1, 1). Fig. depicts NW: performance vs N₂ for Camelyon-17 and ISIC datasets. We find that performance is relativity insensitive to N above ~ 5 examples per class.
In the original NW head paper [55], the authors experiment with temperature (i.e. bandwidth) hyperparameter т. In this work, we set r = 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), c.g.). Thus, we let the feature extractor optimize the bandwidth on its own. Note this same approach taken in prior works [47
Table 3: Hyperparameter settings for various datasets.
Hyperparameter
Camelyon-17
ISIC
FMoW
Learning rate
le-4
5e-5
le-4
Weight decay
le-4
0
le-2
Scheduler
None
None
StepLR
Batch size
32
8
8
Architecture
DenseNet-122
ResNet-50
DenseNet-122
Optimizer
SGD
Adam
Adam
Maximum Epoch
10
5
60
N
8
8
8
N
8
8
N
No x 2 = 16
N₂ x 2 = 16
N, x 62 = 62
0.01
0.01
0.1
k
3
3
3
Camelyon-17
ISIC
70.0
69.5
Average accuracy
69.0
F1 score
78
68.5
77
68.0
67.5
128
Nc
Nc
Figure 5: NWB performance vs No for Camelyon-17 17 and ISIC datasets. Full mode. Performance is relativity insensitive N. above 5 examples per class.
17

---

D Table of Runtimes
Table shows approximate runtimes for various datasets during training and inference. All experi- ments are performed on GPU.
Table 4:
Approximate runtimes for various algorithms. Training time is time to complete maximum epochs as specified in Table and does not include validation. Inference time time evaluate the entire test set. Averaged over all training runs.
Algorithm
Camelyon-17 17
ISIC
FMoW
Training
ERM
7 hr 14 hr
1 hr
22 hr
NW
2 hr
40
Inference
ERM
10 min
2 min
10 min
NW, Random
15 min
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
min
10 min
E ID vs. OOD Performance
With most invariant learning methods, prior works observe tradeoff between in-distribution and out-of-distribution generalization. In Table we show in-distribution (ID) and out-of-distribution (OOD) results for Camelyon-17, 17, which provides ID validation set. We observe that there is tradeoff between ID and OOD performance, similar to prior work.
Table 5: In-distribution vs. out-of-distribution generalization performance on Camelyon-17. 17.
In-distribution
Out-of-distribttion
ERM NW Full
23.2
70.3
96. 1
72.0 72.0+67
NW! Full
92.8+20
80.0#27
18

---

Imbalanced ISIC Experiments
As ISIC exhibits significant label imbalance where the positive class is much less represented than the negative class (see Fig. [ we experiment with an NW variant without support-set label balancing as an ablation. For both variants, we N₂ To train the imbalanced variant, we sample mini-batch support set by first sampling one image per class (to guarantee both classes are represented the support at least once), and then sampling the rest the images randomly from dataset.
To characterize the performance of both variants across class imbalances, we change the prevalance of y 0 in the test set by removing negative class images until the desired prevalence is achieved. Note the default prevalence is ~ 0.85. Then, we compute the accuracy over this manipulated test set (note, prior work has shown that Fl-score label imbalances (3).
a good metric for comparing classifiers with different
Fig 7 shows results. We observe that 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 high prevalence, the class imbalance is similar for the training and test domains; thus, a model which overpredicts the negative class usually right. On the other hand, this fails in test where the prevalence is flipped (i.e. low prevalence). These results suggest that NW8 a more robust classifier in the presence of label shift.
Camelyon-17
ISIC
7000
E=0
E=0
60000
E=1
6000
E=1
E=2
E=2
50000
5000
Datapoints
Datapoints
40000
4000
30000
3000
20000
2000
10000
1000
0
0=0
y=1
y=0
y=1
y=0
y=1
y=0
y=1
y=0
y=1
y=o
y=1
Figure 6: Number of datapoints separated by class for Camelyon-17 and ISIC datasets. significant label imbalance for the ISIC dataset.
There is
Random
Full
Cluster
Ensemble
0.9
0.8
Belanced
Accuracy
0.6
0.4
0.50
0.75
0.25
0.50
0.75
0.25
0.50
0.75
0.25
0.50
0.75
Prevalence Y=
Prevalence Y=
Prevalence Y=0
Prevalence of Y=0
Figure 7: Accuracy of NW (imbalanced) and NW8 (balanced) models over varying prevalence of y = 0 for ISIC dataset. At low prevalence where the prevalence differs the most from training domains, we observe that model performance is higher for BW The default prevalence is 6705/7818 0.8576, which is the right-most value in the x-axis.
19

---

FMoW G Nearest-neighbor Inference Modes
In addition to Random, Full, Cluster, and Ensemble inference modes, additionally experiment with two inference modes based on nearest neighbors: k-NN and HNSW. HNSW (Hierarchical
Navigable Small Worlds) is fast approximate nearest neighbor algorithm 33.. We choose k 20 based on prior work 225] HNSW is about two times faster in total runtime
GPU than full k-NN.
Overall, observe that k-NN and HNSW perform nearly identically for all the datasets and variants. Additionally, both modes perform better in terms of mean performance on Camelyon-17, and perform on par with the best-performing modes for ISIC (Cluster) and FMoW (Ensemble). However, they generally have higher variances across model runs, We suspect that less total samples used in the support (20 as compared more than 1000 for Full) and the fact that all classes are guaranteed to be represented in the support may lead to less stability across model runs.
Table 6: Metric average + standard deviation for all datasets (%). Higher is better. Implicit Explicit.
Algorithm
17
ISIC
NW Full
72.0267
61.9+35
31.6 31.6109
NW Cluster
70.6269
61.4+23
31.32a0
NW Ensemble
71.966.
63.9
32.2
NW k-NN
72.5
64.2122
32.5
NWB HNSW
72.5232
64.2226
32.5+14
NW Full
80.0#27 79.7+19
69.6 77.04+0
35.0+a. / 34.6+00
NW
0. Cluster
78.6 79.0
71.1 ( /7.00
33.9 34.0a03
NW Ensemble
79.5 79.6119
69.5 69.8
37.8109 / 38.2.00
NW k-NN
80.9 80.2
/0..0z4
37.9+23 / 38.3
NW HNSW
80.9 80.2
70.7
37.9 / 38.2+
H Interpretability of NW Head
In this section, we provide both visual and quantitative exploration of the interpretability capabilities of the NW head. In Fig. 8 we show several query images and its 8 nearest neighbors in the feature space, for both NW and NW! variants. The colored border around each neighbor indicates which training environment the image comes from. Interestingly, we notice that the neighbors for NW! come from variety of environments (note a variety of colored borders), while the neighbors for NW8 are less diverse.
Fig. D quantifies this phenomenon, depicting a normalized histogram of the environments from which the top 20 nearest neighbors originate the training dataset for Camelyon-17, averaged over all queries in the test set. From these results, it is clear that NW leverages support images from wider variety of environments make its prediction, suggesting that captures more invariant representations.
20

---

Query
Query
0
Query
Figure 8: Visualization of 3 query images and their 8 nearest neighbors in the feature space, for both NW? (top) and NW (bottom). Labels shown in the bottom right corner, and colored borders of neighbors indicate which of the training environments the image comes from. For both variants, we observe visual similarity between the query images and the nearest neighbors. However, observe the neighbors for NW tend to lack diversity in the environments from which they originate. In contrast, neighbors for NW tend to come from wider variety of environments, suggesting that NW captures more invariant representations.
NW8
NW!
0.4
Average proportion of NNs
0.3
0.2
0.1
0.0
0.0
E=1
E=2
E=3
E=1
E=2
E=3
Figure Normalized histogram of the environments from which the top 20 nearest neighbors originate in the training dataset for Camelyon-17 17, averaged over all queries in the test set. We observe more balanced proportion for ΝW indicating that the model relies more evenly across all 3 environments to make its prediction, and further suggesting that representations are more invariant than NWB
21


Reviews:
[{"review_id": "on3e0Me2Gd", "reviewer": "Reviewer_d2Ui", "summary": "This works proposes a causally motivated new setting for domain generalization building on the existing nonparametric Nadaraya-Watson head on top of a learned neural encoder. More precisely, it assumes that data inputs are causally generated from style latent features, which are environment-dependent, and content latent features, which are environment-independent after class-wise balancing. The authors hence try to learn these content features with a neural network. Predictions are made using the classical Nadaraya-Watson classifier on the learned representations of some support points. Training is carried with gradient-based maximum-likelihood, both with and without an additional penalty that encourages environment-independence of representations. At inference, fixed support points from the training set are chosen following one out of four strategies. The proposed method is compared to well-known domain generalization approaches in three datasets.", "questions": "Question above", "limitations": "Limitations of the method are cited and discussed in the appendix.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "### Originality\n\nThe Nadaraya-Watson head is not novel, nor the causal framework proposed for domain generalization. Hence, the originality of the paper lies in using this nonparametric method for this purpose, motivated by the fact that the choice of the support set provides “a degree of flexibility not possible with parametric models”. This seems enough to say that the work is original and potentially significant for the subfield.\n\n### Clarity\n\nI found the paper very well-written and structured. Figures are clean and the method description is clear.\n\n### Quality\n\nThe paper presents interesting results in a rather fair setting, with error bars and well explained experimental protocol.\n\n### Significance\n\nThis work addresses the very important and difficult question of domain generalization with a very creative approach: non-parametric learning coupled with neural representations.", "weaknesses": "### Clarity\n\nThe only remark I would have is that the difference between NWb and NWBe-implicit could be better explained and highlighted in the paper (it took me a while to understand the subtle difference between them).\n\nTypos:\n\n- l. 255 “and”\n\n\n### Quality\n\nDespite the very positive points listed above, I still have a few concerns (from the more important to the least):\n\n1. My main remark here concerns the pros and cons of the method, which are not very well highlighted in my opinion. While the idea is very interesting and it is easy to understand why having access to support points can help for domain generalization, I don’t think this is very well illustrated in the experiments. Results from table 2 show that NWbe is equivalent to CLOvE in Camelyon-17, worse than it in FMoW and (given the high uncertainties) only slightly better than it on ISIC (the only dataset on which CLOvE was reimplemented, according to the Appendix). Given the important computational caveats of the NWbe method (2x longer training and 8-16x longer inference than ERM) and additional overhead (double or triple batches, additional hyperparameters like number of cluster, strategy, etc), I find important to highlight and illustrate what benefits the method has compared to other existing methods which lead to comparable results. For example, it is mentioned in the conclusion that the method has interpretability advantages but this is not shown in the experiments.\n2. Another point that surprised me was the fact that the implicit training (without regularization) led to comparable or better results than the regularized version. How do you explain that? I believe this point needs more investigation as it could indicate that the reason your method works is not exactly the ones explained (learning of environment-independent content latent features). It could be interesting for example to evaluate for both implicit and explicit models the invariance of learned neural embeddings to environment changes.\n3. Is it fair to compare IRM to your representations + MLP given that IRM is trained with a simple linear classifier?\n4. In figure 7 of the appendix, the rightmost accuracies for NWb (close to 90%) are substantially higher than those reported in table 2 (around 60%). What is the reason?\n5. Also regarding this figure 7, I am curious why the balanced plot increases quite fast as the test set becomes more and more imbalanced. It is a bit surprising given that the method consists in class-balancing at training isn’t it ? Have you investigated why?\n6. Also regarding balancing, I think that an additional baseline corresponding to ERM with class-balancing could be useful to better understand what is bringing the deltas in performance between ERM and NWb in table 2: is it the balancing or the NW head or both?\n7. How was the number of clusters $k$ set? I see in appendix it was set to 3, but an analysis of its impact or at least a comment explaining why such hyperparameter is not a sensitive one could be a plus.\n\n### Significance\n\nDespite the very interesting method and important questions addressed, I think the paper experimental part still needs a little work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The only remark I would have is that the difference between NWb and NWBe-implicit could be better explained and highlighted in the paper.\n\nWe plan to add additional clarifying statements in our revised manuscript. To be clear, the difference between them is that in $NW^B_e$-implicit, we not only balance across classes (B), but also condition the support on an environment (e). The (-implicit) indicates that we don’t explicitly regularize the domain invariance (Eq. 7), but enforce it implicitly through the training process (Eq. 8).\n\n> My main remark here concerns the pros and cons of the method, which are not very well highlighted in my opinion. [...]\n\nThank you for raising this important question. We have addressed it in the general response, and kindly refer the reviewer there.\n\n> Another point that surprised me was the fact that the implicit training (without regularization) led to comparable or better results than the regularized version. How do you explain that?\n\nIn Implicit (Eq. 8), the constraint will be approximately satisfied in the sense that the model will be encouraged to predict the identical ground truth across all environments for an input. As an example, consider a classification task with 3 environments, and consider 3 images from the 3 environments each with label 0. Given infinite data and infinite model capacity, the model will predict a label 0 for each of these images, thus achieving invariance across the 3 domains.\n\nThis analysis is borne out in the experiments. For the simpler datasets (Camelyon-17 and ISIC), Implicit performs equally well with higher variances compared to Explicit. In contrast, for FMoW, Explicit outperforms Implicit on Ensemble mode. This could be because the model does not have enough capacity to capture the invariance in an implicit setting.\n\n> Is it fair to compare IRM to your representations + MLP given that IRM is trained with a simple linear classifier?\n\nIRM and NW-probe ($NW^B_e$-learned representations + finetuned MLP) have exactly the same architecture. In IRM, both the feature extractor weights ($\\varphi$) and the classifier weights ($w$) are trained simultaneously. In contrast, NW-probe is trained in two stages. In the first stage, we train the feature extractor weights ($\\varphi$) only using an $NW^B_e$ training strategy (Eq. (7) or (8)). In the second stage, we freeze $\\varphi$ and train a linear probe ($w$) on top of the frozen representations. \n\nNW-probe provides value as an academic exercise, as it can be directly compared with IRM; the causal assumptions and constraints are theoretically identical. Interestingly, we find that it outperforms IRM in the tasks we try, suggesting that $NW^B_e$ training strategy learns better invariant representations.\n\n> In figure 7 of the appendix, the rightmost accuracies for NWb (close to 90%) are substantially higher than those reported in table 2 (around 60%). What is the reason?\n\nThe reason for this difference is that we report F1 scores in Table 2, but we report accuracies in Figure 7. As we mention in Lines 525-6, [2] finds that F1 score is not a good metric for comparing models with different label imbalances. In practice, we find graphs akin to Figure 7 but using F1 scores are not as striking as using accuracies.\n\n> Also regarding this figure 7, I am curious why the balanced plot increases quite fast as the test set becomes more and more imbalanced. It is a bit surprising given that the method consists in class-balancing at training isn’t it ? Have you investigated why?\n\nNote that in Fig. 7, a prevalence value of 0.25 indicates that 25% of samples in training are Y=0. Thus, the plot sweeps from low prevalence values for Y=0 (~25%) to high prevalence values (>75%). Note that the test-set prevalence is ~85%, so the right-most side of Fig. 7 shows performance when the train and test set prevalences are close.\n\nWe agree with the reviewer that the Balanced (orange) plot increases as the prevalence of label Y=0 increases, although it doesn’t increase as quickly as Imbalanced (blue). Thus, while $NW^B$ mitigates label shift to an extent, increasing the prevalence of Y=0, thereby matching the train/test prevalences, leads to improved performance. However, we observe that the highest gap between the orange and blue lines are at low prevalence values. Thus, when prevalence values between train and test are maximally different, $NW^B$ proves to be a good robust predictor.\n\n\n> Also regarding balancing, I think that an additional baseline corresponding to ERM with class-balancing could be useful…\n\nIn response to this comment, we have trained an ERM variant with balanced classes per environment for ISIC, which we denote $ERM^B$, and present the results below. \n\n|      | F1 score |  \n| ----------- | ----------- |  \n| $ERM$ | 58.2 (2.9)|  \n| $NW^B$, Ensemble|63.9 (3.8)|  \n| $ERM^B$|63.0 (2.5)|\n\nWe find that performance is on-par with $NW^B$. This is expected as the theoretical assumptions are the same for both models; that is, removing dependence of E on Y via class balancing. We plan to add these results to our revised manuscript.\n\n> How was the number of clusters set?\n\nA value of $k=3$ was chosen as it seemed to balance good error rate performance (see row 1 of Figure 2 in [1]) with computational efficiency. In [1], accuracy is shown to be stable for varying values of $k$. We plan to add this detail to our revised manuscript.\n\n> Despite the very interesting method and important questions addressed, I think the paper experimental part still needs a little work.\n\nWe hope that our additional results and explanations regarding experimental details provide a clearer picture of our work.\n\n[1] Alan Q. Wang and Mert R. Sabuncu. A flexible nadaraya-watson head can offer explainable and calibrated classification. Transactions on Machine Learning Research, 2023.   \n[2] Jan Brabec, Tomáš Komárek, Vojtech Franc, and Lukáš Machlica. On model evaluation under non-constant class imbalance, 2020.   "}, {"review_id": "LsOJkUZWbH", "reviewer": "Reviewer_48ug", "summary": "The authors apply a recently proposed method for similarity-based prediction to the problem of invariant learning. This paper builds off of the Nadaraya-Watson architecture where predictions for a test input are derived via nearest proximity according to a learned kernel. The proposed method uses a NW head for each training environment, plus some additional regularization to encourage invariance over an assumed causal graph.\n", "questions": "* What is the role of the support set in OOD generalization? Are training samples always used to make predictions at test time? Empirical evidence suggests OOD generalization is possible, but since the method relies on domain-labeled training examples to make predictions, is there some implicit assumption about how domains are related? For example in the IRM paper they discuss the relationship between domains geometrically in terms of “linear general position”.\n* How is the kernel bandwidth (called \"temperature\" in the NW head paper) chosen?  This seems like a critical hyperparameter but I didn't see any discussion of it.\n* Why is it more effective to regularize a non-parametric predictor (Eqn 7) than to directly regularize a standard classification loss?\n* With most invariant learning methods there is a tradeoff between in-distribution and out-of-distribution generalization. Do you see that pattern with your method? Table 2 only reports OOD generalization if I understand correctly\n\nMisc comments:\n* The proposed method “leverages the NW head as a conditional estimator for Y conditioned on Z” [143–144]. The workshop paper “Towards Environment-Invariant Transfer...\" by Eyre et al (https://openreview.net/forum?id=c4l4HoM2AFf) may be of interest because they also use kernels to estimate similar statistics of the learned representations.\n", "limitations": "Limitations are addressed [Sec 6].\n", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Regularization-based approaches to invariant learning are known to entail optimization difficulties [https://proceedings.mlr.press/v162/zhang22u.html]. The non-parametric prediction approach taken by the authors provides a unique and interesting alternative.\n\nThe proposed method shows promise on relevant datasets. I was especially interested in the probe variant of NW-training, which indicates a more invariant internal representation than IRM.\n", "weaknesses": "While the proposed approach is novel and interesting, after reading the paper I was still left with questions as to how the method is implemented, and why it works (see questions below). Also, I think the paper could benefit from some simple theoretical analysis that demonstrates (even in a simplified setting) when we expect NW-based training will discover invariant features. I also feel that the proposed approach, which uses causal graphs to motivate independances on subsets of the training data, is very related to recently proposed MMD-based approaches [https://proceedings.mlr.press/v151/makar22a, https://arxiv.org/abs/2209.09423]. A discussion on how the proposed method (and its use of non-parameteric prediction) differs from these papers would be useful.\n\nThe fact that so many different variants of the NW method are tried makes me a bit wary of claims like \"NW^B does X percent better than ERM\" [line 256]. Not all the NW^B flavors outperform ERM.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> [...]. A discussion on how the proposed method differs from these papers would be useful.\n\nWe agree with the reviewer and reiterate that our causal setup and DAG in Fig. 2b is not novel, and has been proposed in many works within this literature (see also [1-2]). As mentioned in the general response, our primary contribution in this work is manipulation of the support set in an NW head during training for learning invariant and robust representations, and for the broader purpose of domain generalization. To the best of our knowledge, ours is the first work to approach domain generalization in a non-parametric manner, as well as the first work to experiment with manipulating the support set to inject prior information into the training process. \n\n> The fact that so many different variants of the NW method are tried makes me a bit wary of claims like \"NW^B does X percent better than ERM\" [line 256]. Not all the NW^B flavors outperform ERM.\n\nTo clarify, we only make the claim that NW^B outperforms ERM on the imbalanced ISIC dataset, where balancing theoretically offers a potential OOD advantage. For other more balanced datasets (Camelyon-17 and FMoW), NW^B does not improve over ERM, as one might expect.\n\nTo demonstrate this further, we have trained an ERM variant with balanced classes per environment for ISIC, which we denote $ERM^B$, and present the results below. \n\n|      | F1 score |  \n| ----------- | ----------- |  \n| $ERM$ | 58.2 (2.9)|  \n| $NW^B$, Ensemble|63.9 (3.8)|  \n| $ERM^B$|63.0 (2.5)|\n\nWe find that performance is on-par with $NW^B$, as is expected.\n\n> What is the role of the support set in OOD generalization? Are training samples always used to make predictions at test time?...is there some implicit assumption about how domains are related?\n\nThe reviewer is correct in that the support set is composed of the training set in order to make predictions at test time. Note that we only rely on domain-labeled training examples to make predictions in Ensemble inference mode, and not any other inference mode (Random, Full, and Cluster). For Camelyon and ISIC datasets, the use of domain information at prediction doesn’t make much difference in optimal performance (e.g. compare $NW^B_e$ performance for both Full and Ensemble).\n\nTheoretically, we do not make any assumptions about how domains are related other than assuming that if any X has a non-zero probability in one environment, it has a non-zero probability in all environments (this is akin to the positivity assumption in causal inference). We agree with the reviewer that this is an important question that, while we believe is out of scope for our current manuscript, should and has been further explored in other works (e.g. see [1])\n\n> How is the kernel bandwidth (called \"temperature\" in the NW head paper) chosen? This seems like a critical hyperparameter but I didn't see any discussion of it.\n\nThe bandwidth hyperparameter (i.e. temperature) was set to 1 for all experiments. Thank you for pointing out this omission – we have added this detail in the revised manuscript. \n\nWe did not tune this hyperparameter. 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 [6], e.g.). Thus, we let the feature extractor to figure out the bandwidth on its own (note the same approach is taken in prior works, see [5]). Nevertheless, we take the reviewers point and believe the exploration of kernel bandwidths (and other kernels) and its relation to optimization procedures might be an important future work (as we allude to in the Conclusion section, lines 296-297).\n\n> Why is it more effective to regularize a non-parametric predictor (Eqn 7) than to directly regularize a standard classification loss?\n\nWe note that the first term in Eq. 7 is a standard cross-entropy classification loss. The second term in Eq. 7 enforces invariance across environments by forcing predictions from two different environments to be the same for a given input. We argue that our non-parametric regularization approach to enforcing invariance is more natural and intuitive than baseline methods because an environment is encoded by manipulating the support set to contain real samples only from that environment. Other baseline methods use proxy methods to achieve this, for example by aligning calibration [2], intermediate layer activations [3], or variances [4] across environments.\n\nHowever, as we mention in our general response, the result we find the most compelling in this work is the Implicit variant (NWbe-Implicit), which is competitive with and often outperforms state-of-the-art baselines while requiring no hyperparameter to tune. Such an implicit approach is difficult to achieve with a parametric strategy. \n\n> With most invariant learning methods there is a tradeoff between in-distribution and out-of-distribution generalization. Do you see that pattern with your method?\n\nBelow, we show in-distribution (ID) and out-of-distribution (OOD) results for Camelyon-17, which provides an ID validation set.\n\n\n|      | ID | OOD |  \n| ----------- | ----------- | ----------- | \n| $ERM$|93.2 (5.2)|70.3 (6.4)|\n| $NW^B$, Full|96.1 (1.0) |72.0 (6.7)| \n| $NW^B_e$, Full|92.8 (2.0)| 80.0 (2.7)|  \n\nWe observe that there is a tradeoff between ID and OOD performance, similar to prior work. We plan to add these results to our revised manuscript.\n\n[1] Rosenfeld et al. The risks of invariant risk minimization, 2021.   \n[2] Wald et al. On calibration and out-of-domain generalization, 2021   \n[3] Sun et al. Deep coral: Correlation alignment for deep domain adaptation, 2016    \n[4] Krueger et al. Out-of-distribution generalization via risk extrapolation (rex), 2021.    \n[5] Snell et al. Prototypical networks for few-shot learning, 2017.    \n[6] Kotelevskii et al. Nonparametric Uncertainty Quantification for Single Deterministic Neural Network. 2022.   "}, {"review_id": "To6p2w0l4F", "reviewer": "Reviewer_xs7Q", "summary": "The authors address an important problem of reliability in deep learning, given data collected from different sources (environments).\nFor this, the authors propose a method that allows the separation of style and content of input objects, which ensures stable behavior in different environments. The authors develop a causality graph and show how to embed causal assumptions into the model.\nFor a concrete model, the authors use a nonparametric method based on the Nadaraya-Watson (NW) head. For predictions, they leverage the NW head, which uses learned representations and a support set of labeled data. \nThey conduct experiments on several datasets with real-world environment variability, and experiments show the competitive performance of the approach.", "questions": "1. It is well known in the literature related to the methods which work with embeddings (see, for example, [3, 4]) that there might be a feature collapse. So that for the small perturbation in the input, the change in the output is arbitrary. Did you use any techniques to prevent it, like spectral normalizations?\n2. Typically, in kernel methods, there is a scale parameter (e.g., bandwidth) that is very important for performance and requires careful selection. As I see from Eq. 2, you simply use negative Euclidean distance. Is it safe? Might there be a case that embedding from one environment a scattered wider than from another, and it is required to use different scale parameters to fit them well?\n3. In [2], authors also use NW head to a potentially big dataset, and to ease computations, they use approximate nearest neighbors algorithms. It would be interesting to apply this option here as well since the \"random\" inference scheme might be too noisy, \"full\" too expensive, and others too rough.\n\n\n[3] Liu, Jeremiah, et al. \"Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.\" Advances in Neural Information Processing Systems 33 (2020): 7498-7512.\n\n[4] van Amersfoort, Joost, et al. \"On feature collapse and deep kernel learning for single forward pass uncertainty.\" arXiv preprint arXiv:2102.11409 (2021).\n\n\n\n\nEdit: I would like to thank the authors for their answers during rebuttal period.", "limitations": "Authors discussed limitations in the main part of the paper.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses an important problem and provides an interesting solution based on Nadaraya-Watson kernel regression, which is interesting.\n2. In general, I find the paper well-written, well-structured, supplied with nice illustrations and easy to follow. \n3. The choice of datasets for experiments is interesting and challenging.", "weaknesses": "1. The overall novelty and the extent of the contribution made by this paper are unclear. The use of the Nadaraya-Watson (NW) head on top of a feature extractor has been previously used [1, 2]. Furthermore, it's not clear whether the manipulation of the support set can be considered as a novel idea. It would be beneficial for the authors to explicitly state the specific contributions of their paper in the introduction section.\n2. There appears to be a missing reference to Figure 3 in the paper.\n3. No code available, so I can not check if experiments are reproducible.\n\n[1] Alan Q. Wang and Mert R. Sabuncu. A flexible nadaraya-watson head can offer explainable and calibrated classification. Transactions on Machine Learning Research, 2023.\n\n[2] Kotelevskii, Nikita, et al. \"Nonparametric Uncertainty Quantification for Single Deterministic Neural Network.\" Advances in Neural Information Processing Systems 35 (2022): 36308-36323.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The overall novelty and the extent of the contribution made by this paper are unclear…\n\nWe have enumerated our specific contributions in the general response. To reiterate here, to the best of our knowledge, we believe that the manipulation of the support set for learning invariant features is a novel approach to invariant learning and the broader goal of domain generalization. We are not aware of any prior work that approaches invariant representation learning from the perspective of a manipulation of support set from a non-parametric perspective.\n\nIn our revised manuscript, we plan to explicitly add a Contributions subsection in the Introduction section. \n\n> There appears to be a missing reference to Figure 3 in the paper.\n\nThank you for pointing out this omission – we will add this missing reference in our revised manuscript.\n\n> No code available, so I can not check if experiments are reproducible.\n\nTo maintain anonymity in the initial submission, we withheld a link to our code. We intend to make it available upon acceptance. Note that our code and dataset relies heavily on the NW head repository [1] and WILDS benchmark [6], thus facilitating easy reproduction and comparison. Both repositories are publicly available.\n\n> It is well known in the literature related to the methods which work with embeddings (see, for example, [3, 4]) that there might be a feature collapse. [...] Did you use any techniques to prevent it, like spectral normalizations?\n\nIn our experiments, we have not observed the feature collapse phenomenon. Other works have found similar results; for example, the authors in [2] find little difference when adding spectral normalization (see Table 6 in [2]). \n\nThe feature collapse phenomenon seems to arise in the context of Gaussian processes and its extension to deep networks, Deep Kernel Learning. These methods require complicated training schemes that minimize the ELBO, which can be decomposed into a data fit term and a complexity penalty term. From [4], the authors associate feature collapse with this complexity term. However, in the NW head, training is performed with a standard maximum-likelihood loss and thus is relatively simpler than GP-based methods. In particular, there is no complexity term which can cause feature collapse. Put another way, theoretically speaking, the NW head departs from standard parametric classification only from an architectural standpoint, and not from a training/inference/optimization standpoint. \n\n> Typically, in kernel methods, there is a scale parameter (e.g., bandwidth) that is very important for performance and requires careful selection….\n\nWe did not tune the bandwidth in our 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 [2], e.g.). Effectively, we allow the feature extractor to optimize the bandwidth during training (note the same approach is taken in prior works, see [5]). Nevertheless, we take the reviewers point and believe the exploration of kernel bandwidths (and other kernels) might be an important future work (as we allude to in the Conclusion section Lines 296-297).\n\nWe tend to disagree with the reviewer’s point that the embeddings from one environment might be scattered wider than another environment, since our optimization (Eqs. 7 or 8) enforces (either explicitly or implicitly) the embeddings from one environment to match the embeddings from another. Our additional results in Figure 9 shed some light on this hypothesis from an empirical perspective – we observe that the nearest neighbors of query images in the feature space come from all 3 training environments evenly. This indicates that the model relies more evenly across all 3 environments to make its prediction, and further suggests that representations are invariant across environments. \n\n> In [2], authors also use NW head to a potentially big dataset, and to ease computations, they use approximate nearest neighbors algorithms. It would be interesting to apply this option… \n\nThank you for this interesting suggestion. Prompted by this, we have added these results in Table 5 in the corresponding PDF. In it, we show additional results for K-NN and HNSW (Hierarchical Navigable Small Worlds), a fast approximate nearest neighbor algorithm [7].\n\nOverall, we observe that k-NN and HNSW perform nearly identically for all the datasets and variants we tried. Additionally, both modes perform better in terms of mean performance on Camelyon-17, and otherwise perform about as well as the best-performing modes for ISIC (Cluster) and FMoW (Ensemble). However, they generally have higher variances across model runs. We suspect this is because there are less total samples used in the support (20 vs. more than 1000 for Full), and also not all classes are guaranteed to be represented in the support, leading to more unstable results.\n\n\n[1] Wang et al. A flexible nadaraya-watson head can offer explainable and calibrated classification, 2023.  \n[2] Kotelevskii et al. Nonparametric Uncertainty Quantification for Single Deterministic Neural Network., 2022  \n[3] Liu et al. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness, 2020  \n[4] van Amersfoort et al. On feature collapse and deep kernel learning for single forward pass uncertainty, 2021  \n[5] Snell et al. Prototypical networks for few-shot learning, 2017    \n[6] Koh et al. Wilds: A benchmark of in-the-wild distribution shifts, 2021.  \n[7] Malkov et al. Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs, 2016.  "}, {"review_id": "WSUcIqN49z", "reviewer": "Reviewer_f7MY", "summary": "This paper proposes a novel algorithm for learning invariant representations using a fixed head that is the sum of similarities of the query features with a set of support features, weighted by one-hot-encoded class label (in other words, the head predicts the class whose support features most align with the query features). Given this fixed head and a training dataset, the proposed method entails optimizing for a representation that solves a constrained maximum likelihood estimation problem. Experimental results evaluating the proposed technique on three benchmark datasets against a variety of baselines are provided.", "questions": "1. The statement \"the objective is equivalent to unconstrained maximum likelihood under the assumption in Eq. 1\" is not clear to me, since Eq. 1 only applies for some particular g_C. \n\n2. The random inference mode does not make sense to me: how can the sampling be uniform across the dataset if each class is represented k times?\n\n3. How are the features for NW Probe trained?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper does a good job of formulating the problem and describing the proposed approach. The writing is generally clear.\n\n2. The motivating application to medical image classification in different domains is strong, as this is an important and relevant unsolved problem.\n\n3. The proposed approach is intriguing as it removes the bilevel complexity from the IRM problem, as it fixes a head rather than learning one.\n\n4. The empirical results are rigorous and promising, as they show state-of-the-art or competitive performance by certain variations of the proposed method on three benchmark datasets.\n\n5. The authors clearly acknowledge the limitations of the proposed method.", "weaknesses": "1. Practically speaking, the proposed method is computationally expensive because it requires computing pairwise comparisons between the representations support and query samples during training and inference, as the authors acknowledge. However, the support batch size is small during training, and the experiments show that using the Cluster method for inference can allow for doing inference with a small number of support samples without much drop in performance, meaning this is not a huge issue. Still, the empirical results are not overly impressive, as the performance of the best variant essentially matches the baseline [50]. It would be helpful to compare the computational costs of the proposed approach with [50], or run further experiments demonstrating a more clear advantage.\n\n2. The motivating application is strong, and from this it is clear why we want estimators that satisfy conditions 1) and 2) in lines 127 and 128. It is also clear that the proposed estimator with the NW head satisfies these conditions. But it is not clear why we need the NW head in order to satisfy these; perhaps a simpler estimator that does not require using a support set of images for every query image could also satisfy 1) and 2).\n\n3. The Introduction and/or Related Works would benefit from a more detailed discussion of the criticisms of IRM, especially relating to how the proposed method can alleviate these criticisms. Granted, IRM is compared with in Section 4.4, but this section does not discuss whether the proposed method addresses the concerns with IRM posed by [15,21,39]. \n\n---------\n\nPost-rebuttal: I have raised my score, please see comment below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Practically speaking, the proposed method is computationally expensive because… [...] It would be helpful to compare the computational costs of the proposed approach with [50], or run further experiments demonstrating a more clear advantage.\n\nWe agree with the reviewer that computational cost is the primary limitation of our method. The computational costs of parametric approaches like [50] are very similar in practice to ERM, which we provide details of in Table 4. \n\nWhile we agree that empirical results of our proposed method are not substantially improved over [50], we believe our non-parametric regularization approach to enforcing invariance (whether explicit or implicit) is more natural and intuitive than baseline methods because an environment is encoded by manipulating the support set to contain real samples only from that environment. Other baseline methods use proxy methods to achieve this, for example by aligning calibration [2], intermediate layer activations [3], or variances [4] across environments.\n\nFor other advantages of our method, we refer the reader to our general response.\n\n> The motivating application is strong, and from this it is clear why we want estimators that satisfy conditions 1) and 2) in lines 127 and 128. [...]. But it is not clear why we need the NW head in order to satisfy these; perhaps a simpler estimator…\n\nSatisfying 1) requires intervening on Y in the causal DAG. In practice, this can be achieved by matching the environment-specific prior for Y (e.g. by sampling balanced mini-batches during training). The approach we take is by balancing classes in the support set, which we view as a minor contribution that fits well with our theme of support set manipulation.\n\nSatisfying 2) essentially defines the invariant representation learning task and many prior works seek to satisfy it. Generally, this is achieved by regularizing the standard classification loss with constraints which align representations across environments (e.g. aligning correlations of layer activations or aligning model calibration). With the NW head, however, this can be achieved naturally, by manipulating the support set.\n\n> The Introduction and/or Related Works would benefit from a more detailed discussion of the criticisms of IRM, especially relating to how the proposed method can alleviate these criticisms. Granted, IRM is compared with in Section 4.4, but this section does not discuss whether the proposed method addresses the concerns with IRM posed by [15,21,39].\n\nPrompted by this suggestion, we have elaborated further in Section 4.4 discussing how the proposed method addresses the concerns with IRM. We provide an overview of these elaborations here.\n\nThe IRM objective seeks a representation such that the optimal (parametric) linear classifier’s parameters on top of this representation is the same across all environments. The main concern raised by the aforementioned works (foremost in [21]) is the tractable version of this objective (IRMv1 in the IRM paper). This tractable objective assumes convexity and is complex, requiring computation of the Hessian. The NW head side-steps these issues by eliminating the need to learn the linear classifier parameters at all – instead, the classifier is *fixed by construction* since it has no learnable parameters. Essentially, we can enforce invariance through design of the support set, providing a much more intuitive and computationally simpler objective to optimize.\n\nNote that from an empirical standpoint, our experiments with NW-probe also address this issue. NW-probe can be directly compared against IRM, as both have exactly the same architecture and number of learnable parameters; additionally, the causal assumptions and the constraints imposed are theoretically identical. Interestingly, we find that NW-probe outperforms IRM in the datasets we try, suggesting that our $NW^B_e$ training strategy does learn better invariant representations.\n\n> The statement \"the objective is equivalent to unconstrained maximum likelihood under the assumption in Eq. 1\" is not clear to me, since Eq. 1 only applies for some particular g_C.\n\nApologies for the confusion; our purpose in including this statement was to highlight the similarity between our objective and that of maximum likelihood estimation. Specifically, the particular $\\phi$ and $g_C$ that satisfies the constraint in Eq. 5 (which is derived from Eq. 1) allows us to essentially drop the conditioning on $e_i$, which is the subscript for the probability estimator being maximized. Thus, we could equivalently write the objective as:\n\n$$ \n   \\arg\\max_{\\phi, g_C}  \\sum_{i=1}^N \\log \\hat{P}^B (y_i \\mid g_C(x_i); \\phi)\n$$\n$$\ns.t. \\ \n    \\hat{P}^B_e (y_i \\mid g_C(x_i); \\phi) = \\hat{P}^B_{e'} (y_i \\mid g_C(x_i); \\phi), \\ \\ \\forall i \\in \\{1, ..., N\\}, \\ \\forall e, e'\\in E.\n$$\n\n\nThis formulation highlights the fact that our objective is simply maximum likelihood estimation with an additional invariance constraint.\n\n> The random inference mode does not make sense to me: how can the sampling be uniform across the dataset if each class is represented k times?\n\nEach class is represented k times, but which images from each class are present in the support set is chosen uniformly at random from the full training set.\n\n> How are the features for NW Probe trained?\n\nNW-probe is trained in two stages. In the first stage, we train the feature extractor weights ($\\varphi$) only using an $NW^B_e$ training strategy (Eq. (7) or (8)). In the second stage, we freeze $\\varphi$ and train a linear probe ($w$) on top of the frozen representations. "}]

Author rebuttal:
We thank the reviewers for their positive and constructive feedback. Here, we provide a general response to all reviewers, and provide a point-by-point response to each reviewer’s comment/questions below their corresponding review.

### Contributions

As raised by reviewers 48ug and ​​xs7Q, we enumerate the contributions of our paper:
1. Our primary contribution is manipulation of the support set in an NW head during training for learning invariant representations, and for the broader purpose of domain generalization. To the best of our knowledge, ours is the first work to approach domain generalization in a non-parametric manner, as well as the first work to experiment with manipulating the support set to inject prior information into the training process. 
2. Another (minor) contribution is that we motivate label shift and IRM simultaneously through d-separation in a causal DAG. To the best of our knowledge, ours is the first to model label shift as an intervention on the Y node in the DAG in Fig. 2a. In our experiments, we find that both lead to a performance boost. Our non-parametric method enforces both these constraints in one model.

We clarify that the NW head itself is not a novel contribution, nor is our causal setup related to IRM (Fig. 2b), which has been proposed in many prior works. 

### Advantages of the NW Head

As raised by reviewers d2Ui and f7MY, we enumerate the advantages of NW over baselines and refer the reader to Section 6 of our paper for a discussion of its disadvantages:
1. The Implicit variant (NWbe-Implicit) has no hyperparameter to tune. This variant is competitive with and often outperforms state-of-the-art baselines which all require tuning a hyperparameter coefficient in the regularized loss. 
2. 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. See Figs. 8 and 9 in the included PDF and “Additional Changes” below. Note this is not possible with parametric baselines, see [1].
3. Intuitively, we believe our non-parametric approach to enforcing invariance across environments is more natural and intuitive than baseline methods because 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 [3-5].

### Additional Changes

1. Prompted by reviewer d2Ui, in Figs. 8 and 9 in the included PDF, we provide both a visual and quantitative exploration of the interpretability capabilities of the NW head. In Fig. 8, we show several query images and its 8 nearest neighbors in the feature space, for both $NW^B$ and $NW^B_e$ variants. We add a colored border around each neighbor to indicates which training environment it comes from. Interestingly, we notice that the neighbors for $NW^B_e$ come from a variety of environments (note a variety of colored borders), while the neighbors for $NW^B$ are less diverse. Fig. 9 quantifies this phenomenon. Thus, $NW^B_e$ leverages support images from a wider variety of environments to make its prediction, suggesting that it captures more invariant representations.

2. Prompted by reviewer xs7Q, we have experimented with 2 additional inference modes: k-NN and HNSW (Hierarchical Navigable Small Worlds), a fast approximate nearest neighbor algorithm [6]. These results are shown in Table 5 in the included PDF. We choose $k=20$ based on prior work [2]. HNSW is about 2x faster in total runtime on a GPU than full k-NN; we plan to add these findings to our computational results in Table 4. Overall, we observe that k-NN and HNSW perform nearly identically for all the datasets and variants. Additionally, both modes perform better in terms of mean performance on Camelyon-17, and perform on par with the best-performing modes for ISIC (Cluster) and FMoW (Ensemble). However, they generally have higher variances across model runs. We suspect that less total samples used in the support (20 vs. more than 1000 for Full) and the fact that not all classes are guaranteed to be represented in the support may lead to more unstable results.

### Further Directions

Finally, we would like to highlight the future directions that our work may inspire. In general, training-time manipulation of a nonparametric support provides a novel means of imposing prior knowledge. Intuitively, every manipulation *creates a new classification problem*. In this work, we explored 2 such sources of prior knowledge: environment/domain knowledge, and assumption of balanced labels. We show that encoding these priors in a non-parametric way results in comparable to superior results compared to methods which encode priors via a regularized loss, and in certain cases obviates the need for tuning the associated hyperparameter.

Further research may explore other ways to leverage this new degree of flexibility. As we mention in the Conclusion, an interesting way to exploit any prior knowledge of class imbalances might be to upweight the occurrence of samples for more prevalent labels during training. 

[1] Wang et al. A flexible nadaraya-watson head can offer explainable and calibrated classification. Transactions on Machine Learning Research, 2023.   
[2] Kotelevskii et al. Nonparametric Uncertainty Quantification for Single Deterministic Neural Network, 2022.   
[3] Wald et al. On calibration and out-of-domain generalization, 2021   
[4] Sun et al. Deep coral: Correlation alignment for deep domain adaptation, 2016    
[5] Krueger et al. Out-of-distribution generalization via risk extrapolation (rex), 2021.  
[6] Malkov et al. Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs, 2016.    

Meta-review: Summary:

The authors address an important problem of reliability in deep learning, given data collected from different sources (environments). For this, the authors propose a method that allows the separation of style and content of input objects, which ensures stable behavior in different environments. The authors develop a causality graph and show how to embed causal assumptions into the model. For a concrete model, the authors use a nonparametric method based on the Nadaraya-Watson (NW) head. For predictions, they leverage the NW head, which uses learned representations and a support set of labeled data. They conduct experiments on several datasets with real-world environment variability, and experiments show the competitive performance of the approach. Experimental results evaluating the proposed technique on three benchmark datasets against a variety of baselines are provided.

Strengths:

1 - The work is original and potentially significant for the subfield.

2 - The paper is very well-written and structured. Figures are clean and the method description is clear.

3 - The paper presents interesting results in a rather fair setting, with error bars and well explained experimental protocol.

4 - This work addresses the very important and difficult question of domain generalization with a very creative approach: non-parametric learning coupled with neural representations.

5 - Regularization-based approaches to invariant learning are known to entail optimization difficulties. The non-parametric prediction approach taken by the authors provides a unique and interesting alternative.

6 - The proposed method shows promise on relevant datasets. I was especially interested in the probe variant of NW-training, which indicates a more invariant internal representation than IRM.

7 - The paper addresses an important problem and provides an interesting solution based on Nadaraya-Watson kernel regression, which is interesting.

8 - The choice of datasets for experiments is interesting and challenging.

9 - The paper does a good job of formulating the problem and describing the proposed approach. The writing is generally clear.

10 - The motivating application to medical image classification in different domains is strong, as this is an important and relevant unsolved problem.

11 - The proposed approach is intriguing as it removes the bilevel complexity from the IRM problem, as it fixes a head rather than learning one.

12 - The empirical results are rigorous and promising, as they show state-of-the-art or competitive performance by certain variations of the proposed method on three benchmark datasets.

13 - The authors clearly acknowledge the limitations of the proposed method.

Weaknesses:

1 - The pros and cons of the method are not very well highlighted in my opinion.

2 - The paper experimental part still needs a little work.

3 - After reading the paper one is still left with questions as to how the method is implemented, and why it works.

4 - The fact that so many different variants of the NW method are tried makes one a bit wary of claims like "NW^B does X percent better than ERM" .

5 - The overall novelty and the extent of the contribution made by this paper are unclear.

6 - No code available, so I can not check if experiments are reproducible.

7 - The proposed method is computationally expensive because it requires computing pairwise comparisons between representations support and query samples.

8 - The empirical results are not overly impressive.

9 - It is not clear why we need the NW head.

10 - The Introduction and/or Related Works would benefit from a more detailed discussion of the criticisms of IRM.

Decision:

All reviewers vote for acceptance. I, therefore, decide to accept the paper and encourage the authors to use the feedback provided to improve the paper for its camera ready version.
Decision: Accept (poster)