Question:
You are a highly experienced, conscientious, and fair academic reviewer, please help me review this paper. The review should be organized into nine sections:
1. Summary: A summary of the paper in 100-150 words.
2. Strengths/Weaknesses/Questions: The Strengths/Weaknesses/Questions of paper, which should be listed in bullet points, with each point supported by specific examples from the article where possible.
3. Soundness/Contribution/Presentation: Rate the paper's Soundness/Contribution/Presentation, and match this score to the corresponding description from the list below and provide the result. The possible scores and their descriptions are:
    1 poor
    2 fair
    3 good
    4 excellent
4. Rating: Give this paper an appropriate rating, match this rating to the corresponding description from the list below and provide the result. The possible Ratings and their descriptions are:
    1 strong reject
    2 reject, significant issues present
    3 reject, not good enough
    4 possibly reject, but has redeeming facets
    5 marginally below the acceptance threshold
    6 marginally above the acceptance threshold
    7 accept, but needs minor improvements
    8 accept, good paper
    9 strong accept, excellent work
    10 strong accept, should be highlighted at the conference
5. Paper Decision: It must include the Decision itself(Accept or Reject) and the reasons for this decision, based on the criteria of originality, methodological soundness, significance of results, and clarity and logic of presentation.

Here is the template for a review format, you must follow this format to output your review result:
**Summary:**
Summary content

**Strengths:**
- Strength 1
- Strength 2
- ...

**Weaknesses:**
- Weakness 1
- Weakness 2
- ...

**Questions:**
- Question 1
- Question 2
- ...

**Soundness:**
Soundness result

**Presentation:**
Presentation result

**Contribution:**
Contribution result

**Rating:**
Rating result

**Paper Decision:**
- Decision: Accept/Reject
- Reasons: reasons content

Please ensure your feedback is objective and constructive. The paper is as follows:

# Learning Invariant Representations with a

Nonparametric Nadaraya-Watson Head

 Alan Q. Wang

Cornell University &Minh Nguyen

Cornell 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 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.

## 1 Introduction

Machine learning models often fail when there is significant distribution shift. The goal of domain generalization is to be able to perform well with new distributions [21, 56, 71]. In this work, we are interested in settings where multiple domains/environments are available during training and we have access to environment indicators. A popular way to tackle domain generalization in this setting is to learn representations that are invariant across environments [20, 41, 60]. The hope is that such representations will work well in, or are transportable to, unseen environments. This invariance is often encoded via constraints on a learned predictor which aligns its behavior across environments; often, these conditions are derived using causal reasoning and/or by making assumptions about the data-generating process [39].

In a parametric setting, almost all existing methods enforce these constraints by training a single model and adding a regularizer on top of a standard predictive loss [1, 5, 18, 20, 27, 48, 54, 58]. Most notably, invariant risk minimization (IRM) enforces the representations to be such that the optimal classifier on top of those representations is the same across all environments. Other examples include enforcing the layer activations of the predictor to be aligned across environments [48], enforcing the predictor to be calibrated across environments [54], and enforcing gradients of the predictor to be aligned across environments [46]. Often, optimizing these constraints demand approximations or relaxations that undermine the efficacy of the approach [22].

In this work, we take a different approach using a nonparametric strategy based on a recently-proposed Nadaraya-Watson (NW) head [55]. Instead of computing the class probability directly from an input query, the NW head makes a prediction by comparing the learned representations of the query to the elements of a support set that consists of labeled data. Thus, the NW prediction is computed _relative to other real datapoints_ in the support set, with the support set providing a degree of flexibility notpossible with parametric models. In particular, one can manipulate it during training in a way which restricts the types of comparisons that the model can make.

In this work, we manipulate the support set during training to encode causal assumptions for the purposes of learning invariant representations. Specifically, restricting the support set to be drawn from a single environment precludes the possibility of using environment-specific features to make a prediction for a given query. We show that this setup is causally-motivated and relates to existing causal frameworks. Furthermore, we show that this training strategy leads to competitive to superior results compared to state-of-the-art parametric baselines.

Our contributions are as follows:

* We present causally-motivated assumptions for domain generalization which justify our modeling and training strategy.
* We present a novel approach to invariant representation learning using the nonparametric Nadaraya-Watson head, which can account for causal assumptions by manipulating a support set. In particular, we propose a training strategy which, unlike competing baselines, has _no invariance hyperparameter to tune_.
* We validate our approach on several datasets and demonstrate competitive results compared to state-of-the-art parametric baselines.

## 2 Related Works

### Domain Generalization and Invariant Representations

Domain generalization seeks to make models robust to unseen environments and is an active area of research [21; 56; 71]. One line of work augments or synthetically-generates additional training images to increase robustness of learned features to unseen environments [59; 63; 64; 65; 70]. In particular, LISA uses a mixup-style [67] interpolation technique to generate augmented images, which the authors demonstrate improves out-of-distribution robustness [64]. Another line of work broadly seeks to align features across distributions. Deep CORAL aligns correlations of layer activations in deep neural networks [48], and other works minimize the divergence of feature distributions with different distance metrics such as maximum mean discrepancy [51; 32], an adversarial loss [14; 30], and Wasserstein distance [69]. Still other works approach the problem from the perspective of the gradients and optimization [12; 29; 34; 46; 57]. For example, Fish aligns the gradients from different domains [46].

One can also achieve domain generalization via learning invariant representations, which often requires reasoning about the data-generating process from a causal perspective to arrive at appropriate constraints [39]. Invariant causal prediction (ICP) formulates the problem from a feature selection perspective, where the goal is to select the features which are direct causal parents of the label [40]. Invariant Risk Minimization (IRM) can be viewed as an extension of ICP designed for deep, nonlinear neural networks. The IRM objective can be summarized as finding the representation \(\varphi\) such that the optimal linear classifier's parameters \(w^{*}\) on top of this representation is the same across all environments [1]. This bi-level program is highly non-convex and difficult to solve. To find an approximate solution, the authors consider a Lagrangian form, whereby the sub-optimality with respect to the constraint is expressed as the squared norm of the gradients of each of the inner optimization problems. Follow-up works analyzing IRM have raised theoretical issues with this objective and presented some practical concerns [16; 22; 42]. Various flavors of IRM have also been proposed by introducing different regularization terms [27; 54; 58].

### Nonparametric Deep Learning

Nonparametric models in deep learning have received much attention in previous work. Deep Gaussian Processes [10], Deep Kernel Learning [62], and Neural Processes [24] build upon Gaussian Processes and extend them to representation learning. Other works have generalized \(k\)-nearest neighbors [37; 49], decision trees [68], density estimation [13], and more general kernel-based methods [15; 36; 66] to deep networks and have explored the interpretability that these frameworks provide. Closely-related but orthogonal to nonparametric models are attention-based models, most notably self-attention mechanisms popularized in Transformer-based architectures in natural language processing [53] and, more recently, computer vision [11, 19, 38]. Nonparametric transformers apply attention in a nonparametric setting [24].

Recently, Wang et al. proposed the NW head [55], an extension of the classical NW model [2, 35, 61] to deep learning. In the NW head, the prediction is a weighted average of labels from a support set. The weights are computed from distances between the query and support features. The NW head can yield better calibration and interpretability, with similar accuracy compared to the dominant approach of using a parametric classifier with fully-connected layers. In this work, we leverage the NW head to encode causal assumptions via the support set. The interpretability and explainability benefits of the NW head carry over in this work; while not of primary focus, we explore these properties in the Appendix.

## 3 Preliminaries

**Problem Setting.** Let \(X,Y\) denote a datapoint and its corresponding discrete class, and \(E\) denote the environment (or domain) where \(X,Y\) originates.1 That is, the elements of the training dataset \(\mathcal{D}_{tr}=\{x_{i},y_{i},e_{i}\}_{i=1}^{N}\) are drawn first by sampling the discrete random variable \(e_{i}\sim P(E)\), and then sampling \(x_{i},y_{i}\sim P(X,Y\mid E=e_{i}):=P_{e_{i}}(X,Y)\). Our goal is to learn classifiers that will generalize to new, unseen environments.

Footnote 1: We assume the support of \(E\), \(\mathrm{supp}(E)\), is finite.

**Assumptions.** We assume there exists a pair of latent causal parents of \(X\): an environment-independent ("content") factor \(Z_{C}\) and an environment-dependent ("style") factor \(Z_{S}\).2 We as

Figure 1: Illustration of proposed approach. Support set of labeled datapoints (square/triangle) from 3 environments lie in 3 regions in the feature space. Black circle denotes query datapoint with unknown label. a) The NW head models \(P(Y|X)\) by making predictions as a function of distances to labeled datapoints in the feature space (visualized as dotted arrows). b) Balancing comparisons across labels for all environments models \(P^{B}(Y|X)\). c) Conditioning on a single environment models \(P_{e}(Y|X)\).

Figure 2: a) Causal Directed Acyclic Graph (DAG) we consider in this work. Solid nodes are observed and dashed nodes are unobserved. We assume an anti-causal setting where label \(Y\) causes \(X\), and \(X\) has 2 causal parents: “style” features, \(Z_{S}\), which are influenced by the environment \(E\); and environment-independent “content” features of \(X\), \(Z_{C}\), which are causally influenced by the label \(Y\). \(E\) potentially influences \(Y\). Both \(E\) and \(Y\) have direct influence on style features \(Z_{S}\). b) Same DAG as a) with an intervention on \(Y\). We note that \(Y\mathchoice{\mathrel{\hbox{\hbox to 0.0pt{\kern 2.999954pt\vrule height 6.299904pt wid th 1px\hss}\hbox{$\perp$}\kern 2.499968pt}}}{\mathrel{\hbox{\hbox to 0.0pt{ \kern 2.999954pt\vrule height 6.299904pt width 1px\hss}\hbox{$\perp$} \kern 2.499968pt}}}{\mathrel{\hbox{\hbox to 0.0pt{\kern 2.099968pt\vrule height 6.299904pt wid th 1px\hss}\hbox{$\perp$}\kern 2.099968pt}}}{\mathrel{\hbox{\hbox to 0.0pt{ \kern 1.499977pt\vrule height 6.299904pt width 1px\hss}\hbox{$\perp$} \kern 2.099968pt}}}E\mid Z_{C}\) and \(Y\mathchoice{\mathrel{\hbox{\hbox to 0.0pt{\kern 2.999954pt\vrule height 6.299904pt wid th 1px\hss}\hbox{$\perp$}\kern 2.499968pt}}}{\mathrel{\hbox{\hbox to 0.0pt{ \kern 2.999954pt\vrule height 6.299904pt width 1px\hss}\hbox{$\perp$} \kern 2.499968pt}}}{\mathrel{\hbox{\hbox to 0.0pt{\kern 1.499977pt\vrule height 6.299 904pt width 1px\hss}\hbox{$\perp$}\kern 2.099968pt}}}{\mathrel{\hbox{ \hbox to 0.0pt{\kern 1.499977pt\vrule height 6.299904pt width 1px\hss}\hbox{$ \perp$}\kern 2.099968pt}}}E\mid Z_{S}\).

sume the causal mechanism that generates \(X\) from (\(Z_{C}\), \(Z_{S}\)) is injective, so that, in principle, it is possible to recover the latent features from the observations; i.e. there exists a function \(g\) such that \(g(X)=(Z_{C},Z_{S})\). We further assume that \(g\) can be disentangled into \(g_{C}\) and \(g_{S}\), such that \((Z_{C},Z_{S})=g(X)=(g_{C}(X),g_{S}(X))\). The causal graph is depicted in Fig. 2a. Finally, we assume that if any \(X=x\) has a non-zero probability in one environment, it has a non-zero probability in all environments.

**Motivation.** The motivating application in this work is image classification, where each realization of \(E\) might represent a different site, imaging device, or geographical region where \(X,Y\) are collected. For example, in medical imaging, different hospitals (\(E\)) may collect images (\(X\)) attempting to capture the presence of some disease (\(Y\)), but may differ in their imaging protocols which lead to differences in the image style features \(Z_{S}\) (e.g. staining, markings, orientation). In addition, we allow \(Z_{S}\) to be influenced by the label itself (for example, positive examples are more likely to be marked by a doctor or have specific staining than negative examples). Finally, the prevalence of \(Y\) may be influenced by \(E\) (for example, the prevalence of a disease may be higher in a certain hospital).

The goal is to find an estimator for \(Y\) which relies only on the direct causal links between \(Y\) and \(X\) and not on any spurious associations between \(E\) and \(X\), as these may change in a new, unseen environment. That is, we seek an estimator which relies only on \(Z_{C}\) and which is independent of \(E\) or \(Z_{S}\).

First, we note the direct causal dependence \(E\to Y\). For example, a model can exploit this association by learning to over-predict majority classes in a certain environment. One way to remove the direct dependence is by intervening on \(Y\), thus removing incoming edges to \(Y\). This essentially corresponds to matching the environment-specific prior on \(Y\) between environments, and results in the intervened graph in Fig. 2b.3 Let us refer to any distribution which follows the DAG in Fig. 2b as \(P^{B}_{e}(X,Y)\).

Footnote 3: This may be interpreted as making the model robust to label shift, see [31; 44].

Second, we observe that there is a potential non-causal association flow between \(E\) and \(Y\) through the colliders \(X\) and \(Z_{S}\), when either one of these are conditioned on (i.e. are observed). An estimator which relies on \(Z_{S}\) potentially leaks information from \(E\), and this is unstable in new environments. Reading d-separation on this intervened graph, we infer that \(Y\perp\!\!\!\perp E\mid Z_{C}=g_{C}(X)\) and \(Y\perp\!\!\!\perp E\mid Z_{S}=g_{S}(X)\), that is:

\[P^{B}_{e}(Y\mid g_{C}(X))=P^{B}_{e^{\prime}}(Y\mid g_{C}(X))\;\;\forall e,e^{ \prime}\in E. \tag{1}\]

In words, this assumption states that the probability of \(Y\) given the environment-invariant parts of \(X\) is the same across any environment \(e\in E\).

Thus, we seek an estimator that 1) enables interventions on \(Y\) such that the direct dependence \(E\to Y\) can be removed, and 2) can further encode the assumption in Eq. (1).

## 4 NW Head for Invariant Prediction

Given a datapoint \(x\), support set \(\mathcal{S}=\{x_{i},y_{i}\}_{i=1}^{N_{x}}\), and parameters \(\phi\), the NW head estimates \(P(Y=y\mid X=x)\) by outputting a prediction formulated as a weighted sum of support set labels, where the weights are some notion of similarity in the feature space [55]:

\[\hat{P}(Y=y\mid X=x;\phi):=f_{\phi}(x,\mathcal{S})=\frac{\sum_{i=1}^{N_{x}} \exp\left\{s(\phi(x),\phi(x_{i}))\right\}\vec{y_{i}}}{\sum_{j=1}^{N_{x}}\exp \left\{s(\phi(x),\phi(x_{j}))\right\}}. \tag{2}\]

Here, \(\vec{y}\) is the one-hot encoded version of \(y\) and \(s(\cdot,\cdot)\) is a similarity/kernel function that captures the similarity between pairs of features. In this work, we set \(s\) as the negative Euclidean distance. A graphical depiction is shown in Fig. 3.

Manipulating the support set can encode certain causal assumptions. We consider two types of manipulations:

1. Balancing classes in \(\mathcal{S}\), denoted \(\mathcal{S}^{B}\) (see Fig. 1b). This can be interpreted as an intervention on \(Y\), and removes the dependence on \(E\to Y\), i.e.: \[\hat{P}^{B}(Y=y\mid X=x;\phi):=f_{\phi}(x,\mathcal{S}^{B}).\] (3)2. Conditioning \(\mathcal{S}\) on a single environment, denoted \(\mathcal{S}_{e}\) (see Fig. 1c). This can be interpreted as conditioning the probability estimate on \(E=e\), i.e.: \[\hat{P}_{e}(Y=y\mid X=x;\phi):=f_{\phi}(x,\mathcal{S}_{e}).\] (4)

Note that both balancing and conditioning can be achieved simultaneously, which we denote \(\mathcal{S}_{e}^{B}\).

### Objective and Optimization Details

Given a dataset of samples \(\mathcal{D}_{tr}=\{x_{i},y_{i},e_{i}\}_{i=1}^{N}\), we wish to leverage the NW head as a conditional estimator for \(Y\) conditioned on \(Z_{C}=g_{C}(X)\), where \(g_{C}(X)\) is characterized by Eq. (1). This necessitates an optimization over both \(\phi\) and the space of functions \(g_{C}\). Thus, we solve the following constrained maximum likelihood over \(\phi\) and \(g_{C}\):

\[\operatorname*{argmax}_{\phi,g_{C}}\sum_{i=1}^{N}\log\hat{P}_{e_{ i}}^{B}(y_{i}\mid g_{C}(x_{i});\phi) \tag{5}\] \[\text{s.t. }\hat{P}_{e}^{B}(y_{i}\mid g_{C}(x_{i});\phi)=\hat{P}_{e ^{\prime}}^{B}(y_{i}\mid g_{C}(x_{i});\phi),\;\;\forall i\in\{1,...,N\},\; \forall e,e^{\prime}\in E.\]

Note that Eq. (1) implies that \(P_{e}^{B}(y_{i}\mid g_{C}(x_{i}))=P^{B}(y_{i}\mid g_{C}(x_{i}))\). Thus, the objective is equivalent to unconstrained maximum likelihood under the assumption in Eq. (1).

Instead of solving for \(g_{C}\) explicitly, we let both \(\phi\) and \(g_{C}\) be related by the composition \(\varphi=\phi\circ g_{C}\), and set \(\varphi\) to be the learnable mapping of the NW head, i.e. a neural network. Then, the objective becomes:

\[\operatorname*{argmin}_{\varphi}\sum_{i=1}^{N}L(f_{\varphi}(x_{i },\mathcal{S}_{e_{i}}^{B}),y_{i}) \tag{6}\] \[\text{s.t. }f_{\varphi}(x_{i},\mathcal{S}_{e}^{B})=f_{\varphi}(x_{i },\mathcal{S}_{e}^{B}),\;\;\forall i\in\{1,...,N\},\;\forall e,e^{\prime}\in E,\]

where \(L\) is the cross-entropy loss. To make the objective tractable, we consider two possible variants:

1. **Explicit.** Solve the optimization problem explicitly via a Lagrangian formulation: \[\operatorname*{argmin}_{\varphi}\sum_{i=1}^{N}L(f_{\varphi}(x_{i },\mathcal{S}_{e_{i}}^{B}),y_{i})+\lambda\sum_{e,e^{\prime}\in E}\sum_{i=1}^{N }\|f_{\varphi}(x_{i},\mathcal{S}_{e}^{B})-f_{\varphi}(x_{i},\mathcal{S}_{e^{ \prime}}^{B})\|_{2}^{2}.\] (7) where \(\lambda>0\) is a hyperparameter.
2. **Implicit.** Relax the optimization problem into the following unconstrained problem: \[\operatorname*{argmin}_{\varphi}\sum_{e\in E}\sum_{i=1}^{N}L(f_{ \varphi}(x_{i},\mathcal{S}_{e}^{B}),y_{i}).\] (8)

Figure 3: A depiction of the NW head on a tumor detection task. The NW head computes Euclidean distances \(s(\cdot,\cdot)\) between query and support features, and uses the distances to weight the support labels. Colored squares represent labels. Diagram displays two different support sets. Top is unconditional support, where support data is drawn from the training data without knowledge of environment information. Bottom is an example of a manipulated support where all support data is drawn from a fixed environment (note similarity in color). Such a support set precludes the possibility of using environment-specific features to make a prediction.

In this formulation, the constraint will be approximately satisfied in the sense that the model will be encouraged to predict the ground truth for a given image, which is identical across all environments. In practice, how well the solution satisfies the constraint will depend on model capacity, the data sample, and optimization procedure.

### Optimization Details

During training, the support set \(\mathcal{S}\) is drawn stochastically from the training set \(\mathcal{D}_{tr}\), and all queries and support datapoints are passed through the feature extractor \(\varphi\). For computational efficiency, instead of sampling a unique support mini-batch at the query-level, we sample a unique support at the mini-batch level. Thus, if \(N_{q}\) and \(N_{s}\) are the query and support mini-batch sizes respectively, the effective mini-batch size is \(N_{q}+N_{s}\), instead of \(N_{q}N_{s}\). For the implicit variant, we sample one support set for a given mini-batch of queries, forward pass through the NW head, and compute the loss in Eq. (8). For the explicit variant, we sample two support sets for a given mini-batch of queries, perform two independent forward passes through the NW head for each support set, and compute the loss in Eq. (7). As discussed in prior work [55], the support batch size is a hyperparameter analogous and orthogonal to the query batch size.

A technical point is that the set of labels in the support mini-batch must cover the set of labels in the query mini-batch. Thus, in our implementation, for \(\mathcal{S}^{B}\), we cycle through all classes and randomly draw \(N_{c}\) examples per class to include in the support. For tasks with a large number of classes, one can subsample from the total number of classes, so long as the sampled classes cover the set of query classes.

### Inference modes

Similar to how the support set can be manipulated during training, we can also design different inference strategies corresponding to different configurations of the support set at test-time. We explore several different inference modes which are possible under the NW framework:

1. **Random.** Sample uniformly at random over the dataset, such that each class is represented \(k\) times.
2. **Full.** Use the entire balanced training set.
3. **Ensemble.** Given the set of balanced features computed from Full mode, partition the computed features for all training datapoints by environment, compute the softmax predictions with respect to each environment balanced across labels, and average the predictions.
4. **Cluster.** Given the set of balanced features computed from Full mode, perform \(k\)-means clustering on the features of the training datapoints for each class. These \(k\) cluster centroids are then used as the support features for each class. This can be viewed as a distillation of the full training set for efficient inference, with the caveat that the support set no longer corresponds to observed datapoints.

While Full, Ensemble, and Cluster require computing features for the entire support set, in practice these features and centroids can be precomputed. In our experiments, we find that Cluster mode can be a sufficient replacement to Full mode, while being computationally cheaper. These inference modes can be used interchangeably and flexibly. As an example, consider a workflow which would involve using Cluster mode to perform efficient inference, and then using Full mode on a select few (potentially problematic) test queries to understand model behavior.

### Connections to Prior Work

Our assumptions in Eq. (1) are common across many works related to learning invariant predictors [40; 27; 54; 41; 26]. Representatively, under the binary classification setting, the IRM objective finds a representation function \(\varphi\) which elicits an invariant predictor across environments \(E\) such that for all \(h\) that has a non-zero probability for \(\varphi(X)\) in any (and all) environment(s):

\[\mathbb{E}_{e}[Y\mid\varphi(X)=h]=\mathbb{E}_{e^{\prime}}[Y\mid\varphi(X)=h], \;\forall e,e^{\prime}\in E.\]

Eq. (1) can be viewed as a generalization of this equality to multi-class settings.4.

Footnote 4: This equality has also been called “sufficiency invariance” [60].

Furthermore, note that given the feature extractor \(\varphi\), the NW mechanism \(f\) is a nonlearnable classifier, whereas \(w\) is learned in the IRM setting. Thus, our proposed objective can be interpreted as learning invariant features \(\varphi\), where the _fixed classifier constraint is satisfied by construction_. This avoids the need to approximate the complex bilevel optimization problem with a regularizer which assumes convexity and requires computing the Hessian. Essentially, \(f\) enforces invariance through the manipulation of the support set, providing a more intuitive and computationally simpler objective to optimize.

In the Experiments section, we compare IRM against a variant of our algorithm where we freeze the learned representations and finetune a linear classifier on top using the same training data. We find that our algorithm performs better than IRM on all datasets, suggesting that it captures invariant representations better than IRM.

## 5 Experiments and Results

### Baselines

We compare against several popular and competitive baseline algorithms: empirical risk minimization (ERM) [52], invariant risk minimization (IRM) [1], deep CORAL [48], Fish [46], LISA [64], and CLOvE [54]. When available, results on baselines are pulled from their respective papers. Details on baseline algorithms are provided in the Appendix.

### Datasets

We experiment on 3 real-world domain generalization tasks. Two are from the WILDS benchmark [23], and the third is a challenging melanoma detection task. Details on the datasets are summarized in Table 1, and further information is provided in the Appendix.

\begin{table}
\begin{tabular}{l l l l l} \hline \hline _Dataset_ & _\# Classes_ & _Env_ & _\# Envs_ & _Architecture_ & _Metric_ \\ \hline Camelyon-17 & 2 & Hospital & 3 & DenseNet-121 & Average acc. \\ ISIC & 2 & Hospital & 3 & ResNet-50 & F1-score \\ FMoW & 62 & Region & 5 & DenseNet-121 & Worst-region acc. \\ \hline \hline \end{tabular}
\end{table}
Table 1: Summary of Datasets.

\begin{table}
\begin{tabular}{l l l l} \hline \hline _Algorithm_ & _Camelyon-17_ & _ISIC_ & _FMoW_ \\ \hline ERM [52] & 70.3\(\pm\)6.4 & 58.2\(\pm\)2.9 & 32.6\(\pm\)1.6 \\ IRM [1] & 70.9\(\pm\)6.8 & 57.9\(\pm\)1.0 & 31.3\(\pm\)1.2 \\ CORAL [48] & 72.4\(\pm\)4.4 & 59.1\(\pm\)2.2 & 31.7\(\pm\)1.0 \\ Fish [46] & 74.7\(\pm\)7e-2 & 64.4\(\pm\)1.7 & 34.6\(\pm\)0.0 \\ LISA [64] & 77.1\(\pm\)6.5 & 64.8\(\pm\)2.3 & 35.5\(\pm\)1.8 \\ CLOvE [54] & 79.9\(\pm\)3.9 & 66.2\(\pm\)2.2 & **40.1\(\pm\)**0.6 \\ \hline NW\({}^{\text{B}}\), Random & 71.7\(\pm\)5.3 & 56.7\(\pm\)1.4 & 31.1\(\pm\)0.8 \\ NW\({}^{\text{B}}\), Full & 72.0\(\pm\)6.7 & 61.9\(\pm\)3.5 & 31.6\(\pm\)0.9 \\ NW\({}^{\text{B}}\), Cluster & 70.6\(\pm\)6.9 & 61.4\(\pm\)2.3 & 31.3\(\pm\)0.9 \\ NW\({}^{\text{B}}\), Ensemble & 71.9\(\pm\)6.0 & 63.9\(\pm\)3.8 & 32.2\(\pm\)1.0 \\ NW\({}^{\text{B}}\), Probe & 69.2\(\pm\)7.4 & 59.7\(\pm\)2.5 & 29.9\(\pm\)1.5 \\ \hline NW\({}^{\text{B}}_{\text{e}}\), Random & 74.8\(\pm\)8.4 / 75.3\(\pm\)3.2 & 57.5\(\pm\)1.9 / 55.0\(\pm\)0.9 & 31.2\(\pm\)0.7 / 30.9\(\pm\)0.5 \\ NW\({}^{\text{B}}_{\text{e}}\), Full & **80.0\(\pm\)**2.7 / 79.7\(\pm\)1.9 & 69.6\(\pm\)2.3 / 70.0\(\pm\)1.0 & 35.0\(\pm\)0.7 / 34.6\(\pm\)0.4 \\ NW\({}^{\text{B}}_{\text{e}}\), Cluster & 78.6\(\pm\)2.5 / 79.0\(\pm\)1.4 & **71.1\(\pm\)**1.7 / 71.0\(\pm\)1.0 & 33.9\(\pm\)0.6 / 34.0\(\pm\)0.3 \\ NW\({}^{\text{B}}_{\text{e}}\), Ensemble & 79.5\(\pm\)2.6 / 79.6\(\pm\)1.9 & 69.5\(\pm\)2.2 / 69.8\(\pm\)0.8 & 37.8\(\pm\)0.9 / 38.2\(\pm\)0.4 \\ NW\({}^{\text{B}}_{\text{e}}\), Probe & 75.3\(\pm\)7.3 / 75.8\(\pm\)8.3 & 61.4\(\pm\)3.1 / 63.4\(\pm\)2.8 & 33.9\(\pm\)1.5 / 32.7\(\pm\)1.4 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Metric average \(\pm\) standard deviation for all datasets (%). Higher is better. **Bold** is best and underline is second-best. Implicit / Explicit.

1. The Camelyon-17 dataset [4] comprises of microscopic images of stained tissue patches from different hospitals, where the label corresponds to presence of tumor tissue in patches and the environment is the hospital where the patch comes from.
2. The melanoma dataset is from the International Skin Imaging Collaboration (ISIC) archive5. The ISIC dataset comprises of dermoscopic images of skin lesions from different hospitals, where the label corresponds to whether or not the lesion is diagnosed as melanoma and the environment is the hospital where the image comes from. There is significant less positive examples and negative examples, with this label imbalance varying across environments (see Appendix). Footnote 5: [https://www.isic-archive.com](https://www.isic-archive.com)
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.

### Experimental Setup

For each model variant, we train 5 separate models with different random seeds, and perform model selection on an out-of-distribution (OOD) validation set. For WILDS datasets, we follow all hyperparameters, use model selection techniques, and report metrics as specified by the benchmark. This includes using a DenseNet-121 backbone initialized with pretrained ImageNet weights as \(\varphi\) and no random augmentations for both datasets. Similarly for ISIC, we use a pretrained ResNet-50 backbone as \(\varphi\) with no augmentations, and perform model selection on an OOD validation set. Due to significant label imbalance, we report F1-score instead of average accuracy.

For NW algorithms, we refer to models which balance classes (i.e. modeling Eq. (3)) as \(\text{NW}^{\text{B}}\), and models which additionally condition on environment (i.e. modeling both Eq. (3) and Eq. (4)) as \(\text{NW}^{\text{B}}_{\text{e}}\). For \(\text{NW}^{\text{B}}_{\text{e}}\) models, we train explicit and implicit variants. For all NW algorithms, we perform evaluation on all inference modes. In addition, for completeness, we experiment on a variant where we freeze the feature extractor and finetune a linear probe on top of the learned representations on the same training data \(\mathcal{D}_{tr}\), which we refer to as "Probe". As an example, the implicit variant of \(\text{NW}^{\text{B}}_{\text{e}}\) is trained on Eq. (8), where the support set is balanced across classes (B) and conditioned on an environment (e).

We set \(N_{c}=8\) for Camelyon-17 and ISIC and \(N_{c}=1\) for FMoW. An analysis of this hyperparameter is provided in the Appendix. The query batch size \(N_{q}\) is set to 8 for all NW experiments. For Random and Cluster inference modes, we set \(k=3\). This was chosen based on prior work [55], where \(k=3\) was shown to balance good error rate performance with computational efficiency. For explicit variants, we tune \(\lambda\) for all datasets via grid search on a held-out validation set. Full hyperparameter details are provided in the Appendix.

All training and inference is done on an Nvidia A6000 GPU and all code is written in Pytorch.6.

Footnote 6: Our code is available at [https://github.com/alanqrwang/nwhead](https://github.com/alanqrwang/nwhead).

### Results

Table 2 shows the main results. We find that on Camelyon-17 and ISIC datasets, \(\text{NW}^{\text{B}}_{\text{e}}\) with Full mode outperforms all baselines and variants we consider. In addition, \(\text{NW}^{\text{B}}_{\text{e}}\) variants typically have lower variance across random seeds as compared to baselines. For FMoW, \(\text{NW}^{\text{B}}_{\text{e}}\) with Ensemble mode performs around \(2\%\) lower than the best performing baseline, CLoVE. We observe that the most computationally-efficient inference mode, Cluster, performs comparably to Full mode for \(\text{NW}^{\text{B}}_{\text{e}}\) models, and is in fact the highest-performing model for ISIC. Thus, we conclude that the support set can be an efficient replacement for Full.

For ISIC, we find that almost all \(\text{NW}^{\text{B}}\) modes (except Random) perform \(\sim 3\%\) better than ERM. This may be attributed to balancing classes across environments, which we suspect has added benefit for highly imbalanced tasks. In contrast, this boost is less apparent for Camelyon-17, which has relatively balanced classes. As an ablation, we compare \(\text{NW}^{\text{B}}\) against an NW variant without class-balancing in the Appendix. \(\text{NW}^{\text{B}}_{\text{e}}\) further improves over \(\text{NW}^{\text{B}}\) by \(\sim 7\%\). Exploring further, we compare \(\text{NW}^{\text{B}}\) against an ERM variant with balanced classes per environment, which we denote \(\text{ERM}^{\text{B}}\). Thisachieves \(63.0\pm 2.5\), which is on-par with \(\text{NW}^{\text{B}}\). This is expected as the theoretical assumptions are the same for both models.

Comparing implicit to explicit variants of \(\text{NW}^{\text{B}}_{\text{e}}\), we do not find much difference in explicitly enforcing Eq. (1), although we do observe significantly lower variances across model runs. Generally, we find the slight performance gain of explicit training to not be worth the computational overhead of doubling the number of support set forward passes per gradient step and tuning the hyperparameter \(\lambda\).

While not the highest-performing, we highlight the consistent \(1\)-\(5\%\) improvement of Probe models over IRM, indicating that \(\text{NW}^{\text{B}}_{\text{e}}\) 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.

## 6 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 to tune, while remaining competitive with and often outperforming state-of-the-art baselines which all require tuning a hyperparameter coefficient in the regularized loss. Second, the NW head enables interpretability by interrogating nearest neighbors in the feature space. Since these neighbors directly contribute to the model's prediction (Eq. (2)), interrogation enables a user to see what is driving the model's decision-making. This not only allows for greater model transparency, but also enables interrogating the quality of the invariant features. We explore this capability in Section H in the Appendix. Note this this degree of transparency is not present in parametric baselines. Lastly, from an intuitive standpoint, we believe our non-parametric approach to enforcing invariance across environments is more natural than baseline methods, since an environment is encoded by manipulating the support set to contain real samples only from that environment. Other baseline methods resort to proxy methods to enforce invariance [54; 48; 27].

One important limitation of our method is computational (see Appendix for analysis of runtimes). The proposed approach requires pairwise comparisons, which scales quadratically with sample size. Practically, this means passing a separate support mini-batch in addition to a query mini-batch at every training iteration. This limitation is compounded for explicit variants, in which two sets of support sets must be drawn independently. Future work may explore more computationally-efficient training strategies. At inference time, Full, Cluster, and Ensemble modes are expensive procedures which require computing features for the entire support set, although precomputing features can mitigate this. However, we argue that in high-risk, safety-critical domains like medical imaging, high inference throughput may not be as important as performance, interpretability, and robustness.

We expect the proposed approach to work well with tasks that have several (and diverse) sets of examples per label class in each environment. If this is not the case, as in the FMoW dataset, the resulting model will be sub-optimal. In particular, in the extreme case where no example is present for a specific class in a given environment, constructing a support set with labels that cover the ground truth label of the query images will not always be possible. This will, in turn, impact performance.

## 7 Conclusion

We presented a nonparametric strategy for invariant representation learning based on the Nadaraya-Watson (NW) head. In the NW head, the prediction is made by comparing the learned representations of the query to the elements of a support set that consists of labeled data. We demonstrated two possible ways of manipulating the support set, and demonstrate how this corresponds to encoding different assumptions from a causal perspective. We validated our approach on three challenging and real-world datasets.

We believe there are many interesting directions of further research. First, our treatment is restricted to classification tasks. Future work may explore an extension to the regression setting. Second, it can be interesting to explore adaptation to the test domain, given additional information. For example, reweighting the occurrence of samples per label could provide improved results given knowledge about the edge \(E\to Y\) in the test distribution. One can further envision implementing the proposed method in settings where there are previously unseen test time labels/tasks. Finally, we are interested in replacing the fixed similarity function with a learnable kernel.

## Acknowledgements

Funding for this project was in part provided by the NIH grant R01AG053949, and the NSF CAREER 1748377 grant.

## References

* Arjovsky et al. [2019] Martin Arjovsky, Leon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization, 2019. URL [https://arxiv.org/abs/1907.02893](https://arxiv.org/abs/1907.02893).
* Bishop [2006] Christopher M. Bishop. _Pattern Recognition and Machine Learning (Information Science and Statistics)_. Springer-Verlag, Berlin, Heidelberg, 2006. ISBN 0387310738.
* Brabec et al. [2020] Jan Brabec, Tomas Komarek, Vojtech Franc, and Lukas Machlica. On model evaluation under non-constant class imbalance, 2020.
* Bandi et al. [2019] Peter Bandi, 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 Ceetin, Eren Halcu, Hunter Jackson, Richard Chen, Fabian Both, Jorg Franke, Heidi Kusters-Vandevelde, Willem Vreuls, Peter Bult, Bram van Ginneken, Jeroen van der Laak, and Geert Litjens. From detection of individual metastases to classification of lymph node status at the patient level: The camelyon17 challenge. _IEEE Transactions on Medical Imaging_, 38(2):550-560, 2019. doi: 10.1109/TMI.2018.2867350.
* Chevalley et al. [2022] Mathieu Chevalley, Charlotte Bunne, Andreas Krause, and Stefan Bauer. Invariant causal mechanisms through distribution matching, 2022.
* Christie et al. [2018] Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world, 2018.
* Codella et al. [2019] 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.
* Codella et al. [2018] 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 at the 2017 international symposium on biomedical imaging (isbi), hosted by the international skin imaging collaboration (isic). In _Proceedings of ISBI_, pages 168-172. IEEE, 2018.
* Combalia et al. [2019] Marc 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 arXiv:1908.02288, 2019.
* Damianou and Lawrence [2013] Andreas Damianou and Neil D. Lawrence. Deep gaussian processes. In Carlos M. Carvalho and Pradeep Ravikumar, editors, _Proceedings of the Sixteenth International Conference on 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.press/v31/damianou13a.html](https://proceedings.mlr.press/v31/damianou13a.html).
* Dosovitskiy et al. [2020] 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 [https://arxiv.org/abs/2010.11929](https://arxiv.org/abs/2010.11929).
* Dou et al. [2020] Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alche-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems_, volume 32. Curran Associates,Inc., 2019. URL [https://proceedings.neurips.cc/paper_files/paper/2019/file/2974788b53f73e7950e8aa49f3a306db-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2019/file/2974788b53f73e7950e8aa49f3a306db-Paper.pdf).
* Fakoor et al. [2020] Rasool Fakoor, Pratik Chaudhari, Jonas Mueller, and Alexander J. Smola. Trade: Transformers for density estimation, 2020. URL [https://arxiv.org/abs/2004.02441](https://arxiv.org/abs/2004.02441).
* Ganin et al. [2016] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, Francois Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks, 2016.
* Ghosh and Nag [2001] 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-7908-1826-0_1. URL [https://doi.org/10.1007/978-3-7908-1826-0_1](https://doi.org/10.1007/978-3-7908-1826-0_1).
* Gulrajani and Lopez-Paz [2020] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization, 2020.
* Gutman et al. [2016] 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.
* Heinze-Deml and Meinshausen [2019] Christina Heinze-Deml and Nicolai Meinshausen. Conditional variance penalties and domain shift robustness, 2019.
* Jaegle et al. [2021] 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.03206](https://arxiv.org/abs/2103.03206).
* Jiang and Veitch [2022] Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts, 2022. URL [https://arxiv.org/abs/2207.01603](https://arxiv.org/abs/2207.01603).
* Kaddour et al. [2022] Jean Kaddour, Aengus Lynch, Qi Liu, Matt J. Kusner, and Ricardo Silva. Causal machine learning: A survey and open problems, 2022. URL [https://arxiv.org/abs/2206.15475](https://arxiv.org/abs/2206.15475).
* Kamath et al. [2021] Pritish Kamath, Akilesh Tangella, Danica J. Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance?, 2021.
* Koh et al. [2021] 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-wild distribution shifts, 2021.
* Kossen et al. [2021] 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 [https://arxiv.org/abs/2106.02584](https://arxiv.org/abs/2106.02584).
* Kotelevskii et al. [2022] 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.
* Koyama and Yamaguchi [2021] Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an out-of-distribution generalization problem?, 2021.
* Krueger et al. [2021] David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex), 2021.
* Kumar et al. [2018] Aviral Kumar, Sunita Sarawagi, and Ujjwal Jain. Trainable calibration measures for neural networks from kernel mean embeddings. In Jennifer Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_, volume 80 of _Proceedings of Machine Learning Research_, pages 2805-2814. PMLR, 10-15 Jul 2018. URL [https://proceedings.mlr.press/v80/kumar18a.html](https://proceedings.mlr.press/v80/kumar18a.html).

* Li et al. [2017] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Learning to generalize: Meta-learning for domain generalization, 2017.
* Li et al. [2018] Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C. Kot. Domain generalization with adversarial feature learning. In _2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition_, pages 5400-5409, 2018. doi: 10.1109/CVPR.2018.00566.
* Lipton et al. [2018] Zachary C. Lipton, Yu-Xiang Wang, and Alex Smola. Detecting and correcting for label shift with black box predictors, 2018.
* Long et al. [2015] Mingsheng Long, Yue Cao, Jianmin Wang, and Michael I. Jordan. Learning transferable features with deep adaptation networks, 2015.
* Malkov and Yashunin [2018] Yu. A. Malkov and D. A. Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs, 2018.
* Mansilla et al. [2021] Lucas Mansilla, Rodrigo Echeveste, Diego H. Milone, and Enzo Ferrante. Domain generalization via gradient surgery, 2021.
* Nadaraya [1964] E. A. Nadaraya. On estimating regression. _Theory of Probability & Its Applications_, 9(1):141-142, 1964. doi: 10.1137/1109020. URL [https://doi.org/10.1137/1109020](https://doi.org/10.1137/1109020).
* Nguyen et al. [2021] Timothy Nguyen, Zhourong Chen, and Jaehoon Lee. Dataset meta-learning from kernel ridge-regression, 2021.
* Papernot and McDaniel [2018] Nicolas Papernot and Patrick McDaniel. Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning, 2018. URL [https://arxiv.org/abs/1803.04765](https://arxiv.org/abs/1803.04765).
* Parmar et al. [2018] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer, 2018. URL [https://arxiv.org/abs/1802.05751](https://arxiv.org/abs/1802.05751).
* Pearl [2009] Judea Pearl. _Causality_. Cambridge University Press, 2 edition, 2009. doi: 10.1017/CBO9780511803161.
* Peters et al. [2015] Jonas Peters, Peter Buhlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals, 2015.
* Rojas-Carulla et al. [2018] Mateo Rojas-Carulla, Bernhard Scholkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning, 2018.
* Rosenfeld et al. [2021] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization, 2021.
* Rotemberg et al. [2021] Veronica Rotemberg, Nicholas Kurtansky, Brigid Betz-Stablein, Liam Caffery, Emmanouil Chousakos, Noel Codella, Marc Combalia, Stephen Dusza, Pascale Guitera, David Gutman, et al. A patient-centric dataset of images and metadata for identifying melanomas using clinical context. _Scientific data_, 8(1):1-8, 2021.
* Schoelkopf et al. [2012] Bernhard Schoelkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning, 2012.
* Scope et al. [2009] A Scope, AA Marghoob, CS Chen, JA Lieb, MA Weinstock, AC Halpern, and SONIC Study Group. Dermoscopic patterns and subclinical melanocytic nests in normal-appearing skin. _British Journal of Dermatology_, 160(6):1318-1321, 2009.
* Shi et al. [2021] Yuge Shi, Jeffrey Seely, Philip H. S. Torr, N. Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization, 2021.
* Snell et al. [2017] Jake Snell, Kevin Swersky, and Richard S. Zemel. Prototypical networks for few-shot learning, 2017. URL [https://arxiv.org/abs/1703.05175](https://arxiv.org/abs/1703.05175).
* Sun and Saenko [2016] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation, 2016.

* Teasiri et al. [2022] Mohammad Reza Teasiri, Giang Nguyen, and Anh Nguyen. Visual correspondence-based explanations improve AI robustness and human-AI team accuracy. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, _Advances in Neural Information Processing Systems_, 2022. URL [https://openreview.net/forum?id=UavQ9HYye6n](https://openreview.net/forum?id=UavQ9HYye6n).
* Tschandl et al. [2018] Philipp Tschandl, Cliff Rosendahl, and Harald Kittler. The ham10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. _Scientific data_, 5(1):1-9, 2018.
* Tzeng et al. [2014] Eric Tzeng, Judy Hoffman, Ning Zhang, Kate Saenko, and Trevor Darrell. Deep domain confusion: Maximizing for domain invariance, 2014.
* Vapnik [1999] V.N. Vapnik. An overview of statistical learning theory. _IEEE Transactions on Neural Networks_, 10(5):988-999, 1999. doi: 10.1109/72.788640.
* Vaswani et al. [2017] 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://arxiv.org/abs/1706.03762](https://arxiv.org/abs/1706.03762).
* Wald et al. [2021] Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization, 2021. URL [https://arxiv.org/abs/2102.10395](https://arxiv.org/abs/2102.10395).
* Wang and Sabuncu [2023] Alan Q. Wang and Mert R. Sabuncu. A flexible nadaraya-watson head can offer explainable and calibrated classification. _Transactions on Machine Learning Research_, 2023. ISSN 2835-8856. URL [https://openreview.net/forum?id=iEq61hG403](https://openreview.net/forum?id=iEq61hG403).
* Wang et al. [2022] Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip S. Yu. Generalizing to unseen domains: A survey on domain generalization, 2022.
* Wang et al. [2023] Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization, 2023.
* Wang and Jordan [2022] Yixin Wang and Michael I. Jordan. Desiderata for representation learning: A causal perspective, 2022.
* 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_, pages 3622-3626, 2020. doi: 10.1109/ICASSP40776.2020.9053273.
* Wang and Veitch [2022] Zihao Wang and Victor Veitch. The causal structure of domain invariant supervised representation learning, 2022. URL [https://arxiv.org/abs/2208.06987](https://arxiv.org/abs/2208.06987).
* Watson [1964] G. S. Watson. Smooth regression analysis. _Sankhya: The Indian Journal of Statistics_, Series A (26):359-372, 1964.
* Wilson et al. [2015] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P. Xing. Deep kernel learning, 2015. URL [https://arxiv.org/abs/1511.02222](https://arxiv.org/abs/1511.02222).
* Xu et al. [2019] Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang. Adversarial domain adaptation with domain mixup, 2019.
* Yao et al. [2022] Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation, 2022.
* Yue et al. [2022] Xiangyu Yue, Yang Zhang, Sicheng Zhao, Alberto Sangiovanni-Vincentelli, Kurt Keutzer, and Boqing Gong. Domain randomization and pyramid consistency: Simulation-to-real generalization without accessing target domain data, 2022.
* Zhang et al. [2021] 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):887-900, 2021. doi: 10.1109/TR.2020.3001232.
* Zhang et al. [2018] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization, 2018.

* Zhang et al. [2018] Quanshi Zhang, Yu Yang, Haotian Ma, and Ying Nian Wu. Interpreting cnns via decision trees, 2018. URL [https://arxiv.org/abs/1802.00121](https://arxiv.org/abs/1802.00121).
* Zhou et al. [2021] Fan Zhou, Zhuqing Jiang, Changjian Shui, Boyu Wang, and Brahim Chaib-draa. Domain generalization via optimal transport with metric similarity learning. _Neurocomputing_, 456:469-480, oct 2021. doi: 10.1016/j.neucom.2020.09.091. URL [https://doi.org/10.1016%2Fj.neucom.2020.09.091](https://doi.org/10.1016%2Fj.neucom.2020.09.091).
* Zhou et al. [2020] Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Deep domain-adversarial image generation for domain generalisation, 2020.
* Zhou et al. [2022] Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. _IEEE Transactions on Pattern Analysis and Machine Intelligence_, pages 1-20, 2022. doi: 10.1109/tpami.2022.3195549. URL [https://doi.org/10.1109%2Ftpami.2022.3195549](https://doi.org/10.1109%2Ftpami.2022.3195549).

Description of Baselines

* Empirical risk minimization (ERM) [52] minimizes the sum of errors across domains and examples.
* Invariant risk minimization (IRM) [1] learns a feature representation such that the optimal linear classifier on top of that representation matches across domains. For WILDS datasets, we pull baseline performance from [54]. For ISIC, we use the implementation from [16].
* Deep CORAL [48] penalizes differences in the means and covariances of the feature distributions (i.e., the distribution of last layer activations in a neural network) for each domain. For WILDS datasets, we pull baseline performance from [54]. For ISIC, we use the implementation from [16].
* Fish targets domain generalization by maximizing the inner product between gradients from different domains. For WILDS datasets, we pull baseline performance from the original paper. For ISIC, we use the implementation from [16].
* LISA augments the set of training data by randomly performing two types of mixup-style [67] interpolations: intra-label (same label, different domain) and inter-label (same domain, different label). For WILDS datasets, we pull baseline performance from the original paper, and train their implementation on the ISIC dataset.
* CLOvE [54] finds an invariant classifier by enforcing the classifier to be calibrated across all training domains. While the original paper proposes several model variants leveraging this idea, we report their best-performing variant, which starts with a trained CORAL model and finetunes the weights using a regularized cross-entropy loss. The regularizer aggregates Maximum Mean Calibration Error (MMCE) [28] over all training domains. For WILDS datasets, we pull baseline performance from [54]. As their implementation is not publicly-available, we implement it for ISIC.

## Appendix B Description of Datasets

Representative examples of the 3 datasets are shown in Fig. 4.

* **Camelyon-17.** We use Camelyon-17 from the WILDS benchmark [4, 23], which provides 450,000 lymph-node scans sampled from 5 hospitals. Camelyon-17 is a medical image classification task where the input \(x\) is a 96 x 96 image and the label \(y\) is whether there exists tumor tissue in the image. The environment denotes the hospital that the patch was taken from. The training dataset is drawn from the first 3 hospitals, while out-of-distribution validation and out-of-distribution test datasets are sampled from the 4th hospital and 5th hospital, respectively.
* **ISIC.** The melanoma dataset is from the International Skin Imaging Collaboration (ISIC) archive7. Data from the archive are collected by different organizations at different points in time [7, 8, 9, 17, 43, 45, 50]. There are about 70k data samples in total. In particular, the resized input image \(x\) is a 224 x 224 image and a binary target label \(y\) denotes whether the image exhibit is melanoma or not. The environment is the hospital from which the image was collected 8. We follow a similar setup to Camelyon-17. The training dataset is drawn from the first 3 hospitals, while out-of-distribution validation and out-of-distribution test datasets are sampled from the 4th hospital and 5th hospital, respectively. For preprocessing, we filter out datapoints that are not specifically categorized as "benign" or "malignant" (e.g. "indeterminate"). The OOD validation dataset is from the "Barcelona1" site indicator and the OOD test dataset is the "Vienna1" site indicator. Footnote 7: [https://www.isic-archive.com](https://www.isic-archive.com)
* **FMoW.** The FMoW dataset is from the WILDS benchmark [6, 23], a satellite image classification task which includes 62 classes and 80 domains (16 years x 5 regions). Concretely, the input \(x\) is a 224 x 224 RGB satellite image, the label \(y\) is one of the 62 building or land use categories, and the environment represents the year that the image was taken as well 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.

Figure 4: Representative images for datasets, separated by domain. Each row depicts a separate class. For FMoW, for simplicity, we show 2 classes out of 62 and only images before 2013.

Hyperparameter Details

Table 3 shows hyperparameter settings for all datasets, where NW-specific hyperparameters are below the midline. For all models, we use pretrained ImageNet weights. For \(\lambda\), we perform a grid-search over the values {0.01, 0.1, 1}. Fig. 5 depicts NW\({}_{\text{e}}^{\text{B}}\) performance vs \(N_{c}\) for Camelyon-17 and ISIC datasets. We find that performance is relativity insensitive to \(N_{c}\) above \(\sim 5\) examples per class.

In the original NW head paper [55], the authors experiment with a temperature (i.e. bandwidth) hyperparameter \(\tau\). In this work, we set \(\tau=1\) for all experiments. The reason for this is that we optimize both the feature extractor and classifier end-to-end, and the kernel used in the classifier is dependent on the features that the feature extractor learns (unlike [25], e.g.). Thus, we let the feature extractor to optimize the bandwidth on its own. Note this same approach is taken in prior works [47].

Figure 5: NW\({}_{\text{e}}^{\text{B}}\) performance vs \(N_{c}\) for Camelyon-17 and ISIC datasets. Full mode. Performance is relativity insensitive to \(N_{c}\) above \(\sim 5\) examples per class.

\begin{table}
\begin{tabular}{l l l l} \hline \hline _Hyperparameter_ & _Camelyon-17_ & _ISIC_ & _FMoW_ \\ \hline Learning rate & 1e-4 & 5e-5 & 1e-4 \\ Weight decay & 1e-4 & 0 & 1e-2 \\ Scheduler & None & None & StepLR \\ Batch size & 32 & 8 & 8 \\ Architecture & DenseNet-121 & ResNet-50 & DenseNet-121 \\ Optimizer & SGD & Adam & Adam \\ Maximum Epoch & 10 & 5 & 60 \\ \hline \(N_{q}\) & 8 & 8 & 8 \\ \(N_{c}\) & 8 & 8 & 1 \\ \(N_{s}\) & \(N_{c}\times 2=16\) & \(N_{c}\times 2=16\) & \(N_{c}\times 62=62\) \\ \(\lambda\) & 0.01 & 0.01 & 0.1 \\ \(k\) & 3 & 3 & 3 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Hyperparameter settings for various datasets.

Table of Runtimes

Table 4 shows approximate runtimes for various datasets during training and inference. All experiments are performed on a GPU.

## Appendix E ID vs. OOD Performance

With most invariant learning methods, prior works observe a tradeoff between in-distribution and out-of-distribution generalization. In Table 5, we show in-distribution (ID) and out-of-distribution (OOD) results for Camelyon-17, which provides an ID validation set. We observe that there is a tradeoff between ID and OOD performance, similar to prior work.

\begin{table}
\begin{tabular}{l l l l l} \hline \hline  & _Algorithm_ & _Camelyon-17_ & _ISIC_ & _FMoW_ \\ \hline Training & ERM & 7 hr & 1 hr & 22 hr \\  & NW & 14 hr & 2 hr & 40 hr \\ \hline Inference & ERM & 10 min & 2 min & 10 min \\  & NW, Random & 15 min & 3 min & 20 min \\  & NW, Full & 2 hr & 15 min & 1 hr \\  & NW, Ensemble & 2 hr & 15 min & 1 hr \\  & NW, Cluster & 2.2 hr & 17 min & 1.1 hr \\  & NW, Probe & 10 min & 2 min & 10 min \\ \hline \hline \end{tabular}
\end{table}
Table 4: Approximate runtimes for various algorithms. Training time is time to complete maximum epochs as specified in Table 3, and does not include validation. Inference time is time to evaluate the entire test set. Averaged over all training runs.

\begin{table}
\begin{tabular}{l l l} \hline \hline  & _In-distribution_ & _Out-of-distribution_ \\ \hline ERM & 93.2\(\pm\)5.2 & 70.3\(\pm\)6.4 \\ NW\({}^{\text{B}}\), Full & 96.1\(\pm\)1.0 & 72.0\(\pm\)6.7 \\ NW\({}^{\text{B}}_{\text{e}}\), Full & 92.8\(\pm\)2.0 & 80.0\(\pm\)2.7 \\ \hline \hline \end{tabular}
\end{table}
Table 5: In-distribution vs. out-of-distribution generalization performance on Camelyon-17.

Imbalanced ISIC Experiments

As ISIC exhibits significant label imbalance where the positive class is much less represented than the negative class (see Fig. 6), we experiment with an NW variant without support-set label balancing as an ablation. For both variants, we set \(N_{c}=8\). To train the imbalanced variant, we sample a mini-batch support set by first sampling one image per class (to guarantee both classes are represented in the support at least once), and then sampling the rest of the images randomly from the dataset.

To characterize the performance of both variants across class imbalances, we change the prevalance of \(y=0\) in the test set by removing negative class images until the desired prevalence is achieved. Note the default prevalence is \(\sim 0.85\). Then, we compute the accuracy over this manipulated test set (note, prior work has shown that F1-score is not a good metric for comparing classifiers with different label imbalances [3]).

Fig 7 shows results. We observe that at high prevalence where the proportion of negative classes matches in training and test, the imbalanced variant outperforms the balanced variant, whereas the opposite is true for low prevalence. This makes sense because at high prevalence, the class imbalance is similar for the training and test domains; thus, a model which overpredicts the negative class is usually right. On the other hand, this fails in test sets where the prevalence is flipped (i.e. low prevalence). These results suggest that \(\text{NW}^{\text{B}}\) is a more robust classifier in the presence of label shift.

Figure 6: Number of datapoints separated by class for Camelyon-17 and ISIC datasets. There is significant label imbalance for the ISIC dataset.

Figure 7: Accuracy of NW (imbalanced) and \(\text{NW}^{\text{B}}\) (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 \(\text{NW}^{\text{B}}\). The default prevalence is 6705/7818 = 0.8576, which is the right-most value in the x-axis.

Nearest-neighbor Inference Modes

In addition to Random, Full, Cluster, and Ensemble inference modes, we additionally experiment with two inference modes based on nearest neighbors: k-NN and HNSW. HNSW (Hierarchical Navigable Small Worlds) is a fast approximate nearest neighbor algorithm [33]. We choose \(k=20\) based on prior work [25]. HNSW is about two times faster in total runtime on a GPU than full k-NN.

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 as compared to more than 1000 for Full) and the fact that not all classes are guaranteed to be represented in the support may lead to less stability across model runs.

## Appendix H Interpretability of NW Head

In this section, 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 \(\text{NW}^{\text{B}}\) and \(\text{NW}^{\text{B}}_{\text{e}}\) variants. The colored border around each neighbor indicates which training environment the image comes from. Interestingly, we notice that the neighbors for \(\text{NW}^{\text{B}}_{\text{e}}\) come from a variety of environments (note a variety of colored borders), while the neighbors for \(\text{NW}^{\text{B}}\) are less diverse.

Fig. 9 quantifies this phenomenon, depicting a normalized histogram of the environments from which the top 20 nearest neighbors originate in the training dataset for Camelyon-17, averaged over all queries in the test set. From these results, it is clear that \(\text{NW}^{\text{B}}_{\text{e}}\) leverages support images from a wider variety of environments to make its prediction, suggesting that it captures more invariant representations.

\begin{table}
\begin{tabular}{l l l l} _Algorithm_ & _Camelyon-17_ & _ISIC_ & _FMoW_ \\ \hline \(\text{NW}^{\text{B}}\), Full & 72.0\(\pm\)6.7 & 61.9\(\pm\)3.5 & 31.6\(\pm\)0.9 \\ \(\text{NW}^{\text{B}}\), Cluster & 70.6\(\pm\)6.9 & 61.4\(\pm\)2.3 & 31.3\(\pm\)0.9 \\ \(\text{NW}^{\text{B}}\), Ensemble & 71.9\(\pm\)6.0 & 63.9\(\pm\)3.8 & 32.2\(\pm\)1.0 \\ \(\text{NW}^{\text{B}}\), k-NN & 72.5\(\pm\)3.2 & 64.2\(\pm\)2.6 & 32.5\(\pm\)1.4 \\ \(\text{NW}^{\text{B}}\), HNSW & 72.5\(\pm\)3.2 & 64.2\(\pm\)2.6 & 32.5\(\pm\)1.4 \\ \hline \(\text{NW}^{\text{B}}_{\text{e}}\), Full & 80.0\(\pm\)2.7 / 79.7\(\pm\)1.9 & 69.6\(\pm\)2.3 / 70.0\(\pm\)1.0 & 35.0\(\pm\)0.7 / 34.6\(\pm\)0.4 \\ \(\text{NW}^{\text{B}}_{\text{e}}\), Cluster & 78.6\(\pm\)2.5 / 79.0\(\pm\)1.4 & 71.1\(\pm\)1.7 / 71.0\(\pm\)1.0 & 33.9\(\pm\)0.6 / 34.0\(\pm\)0.3 \\ \(\text{NW}^{\text{B}}_{\text{e}}\), Ensemble & 79.5\(\pm\)2.6 / 79.6\(\pm\)1.9 & 69.5\(\pm\)2.2 / 69.8\(\pm\)0.8 & 37.8\(\pm\)0.9 / 38.2\(\pm\)0.4 \\ \(\text{NW}^{\text{B}}_{\text{e}}\), k-NN & 80.9\(\pm\)10.3 / 80.2\(\pm\)4.9 & 70.7\(\pm\)7.4 / 71.0\(\pm\)7.2 & 37.9\(\pm\)2.3 / 38.3\(\pm\)1.1 \\ \(\text{NW}^{\text{B}}_{\text{e}}\), HNSW & 80.9\(\pm\)10.3 / 80.2\(\pm\)5.0 & 70.7\(\pm\)7.4 / 71.0\(\pm\)7.2 & 37.9\(\pm\)2.2 / 38.2\(\pm\)1.1 \\ \hline \end{tabular}
\end{table}
Table 6: Metric average \(\pm\) standard deviation for all datasets (%). Higher is better. Implicit / Explicit.

Figure 8: Visualization of 3 query images and their 8 nearest neighbors in the feature space, for both \(\text{NW}_{\text{e}}^{\text{B}}\) (top) and \(\text{NW}^{\text{B}}\) (bottom). Labels are shown in the bottom right corner, and colored borders of neighbors indicate which of the 3 training environments the image comes from. For both variants, we observe visual similarity between the query images and the nearest neighbors. However, we observe the neighbors for \(\text{NW}^{\text{B}}\) tend to lack diversity in the environments from which they originate. In contrast, neighbors for \(\text{NW}_{\text{e}}^{\text{B}}\) tend to come from a wider variety of environments, suggesting that \(\text{NW}_{\text{e}}^{\text{B}}\) captures more invariant representations.

Figure 9: Normalized histogram of the environments from which the top 20 nearest neighbors originate in the training dataset for Camelyon-17, averaged over all queries in the test set. We observe a more balanced proportion for \(\text{NW}_{\text{e}}^{\text{B}}\), indicating that the model relies more evenly across all 3 environments to make its prediction, and further suggesting that representations are more invariant than \(\text{NW}^{\text{B}}\).

Response:
## Review

### Summary

This paper addresses the important problem of reliability in deep learning when dealing with data collected from diverse environments. The authors propose a novel method that separates style and content in input objects to ensure stable behavior across different settings. Building on the Nadaraya-Watson (NW) framework, the method employs a causality graph to embed causal assumptions into the model. Predictions are made using a nonparametric approach based on learned representations and a support set of labeled data. The authors conduct experiments across several benchmark datasets that demonstrate competitive performance, although some areas for improvement are identified in the experimental design and clarity of results.

### Strengths

- The work is original and potentially significant for the subfield.
- The paper is well-written and structured, with clean figures and clear method descriptions.
- The experimental protocol is well explained, and results are presented with appropriate error bars.
- The approach creatively addresses the challenging issue of domain generalization through non-parametric learning coupled with neural representations.
- The method shows promise on relevant datasets, particularly the probe variant of NW-training indicating a more invariant internal representation than IRM.
- The motivating application to medical image classification across different domains is strong and relevant.

### Weaknesses

- The pros and cons of the method are not clearly highlighted, leaving some claims unsubstantiated.
- Some experimental aspects need further clarification, particularly regarding the computational efficiency of the method.
- The overall novelty and contribution of the paper are not clearly articulated, particularly in relation to existing literature.
- There is a lack of reproducibility due to the absence of code, which raises concerns about the validity of experimental results.
- The introduction and related works section would benefit from a more detailed discussion of criticisms surrounding IRM and how the proposed method addresses these issues.

### Questions

- What role does the support set play in out-of-distribution (OOD) generalization, and how are the domains related?
- How is the kernel bandwidth (temperature) chosen, and what impact does it have?
- Why is it more effective to regularize a non-parametric predictor than a standard classification loss?
- What justifies the use of the NW head in achieving invariance, and could simpler alternatives suffice?
- Can you clarify the equivalence of the objective to unconstrained maximum likelihood under specific assumptions?

### Soundness

**Score:** 3

**Description:** Good. The foundational concepts and methods proposed are generally solid, but some claims require further validation and clarity.

### Presentation

**Score:** 3

**Description:** Good. While the paper is mostly well-structured and clear, there are areas where clarity could be improved, particularly regarding methodology and experimental results.

### Contribution

**Score:** 2

**Description:** Fair. Although the approach is interesting, the extent of its contribution relative to existing methods in the literature is not fully established.

### Rating

**Score:** 5

**Description:** Borderline accept: The paper is technically solid and presents a novel approach, but there are significant areas for improvement which should be addressed before final acceptance.

### Paper Decision

**Decision:** Accept (poster)

**Reasons:** The paper presents an original approach to a significant problem in deep learning, with promising experimental results. However, there are areas for improvement in clarity, detail, and the presentation of results that should be addressed in the final version. The strengths of the work outweigh the weaknesses, leading to a recommendation for acceptance.