Learning Invariant Representations with a Nonparametric Nadaraya-Watson Head

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.

Strengths

1. The paper does a good job of formulating the problem and describing the proposed approach. The writing is generally clear. 2. The motivating application to medical image classification in different domains is strong, as this is an important and relevant unsolved problem. 3. The proposed approach is intriguing as it removes the bilevel complexity from the IRM problem, as it fixes a head rather than learning one. 4. 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. 5. 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. 2. 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). 3. 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]. --------- Post-rebuttal: I have raised my score, please see comment below.

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. 2. 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? 3. How are the features for NW Probe trained?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

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.

Strengths

1. The paper addresses an important problem and provides an interesting solution based on Nadaraya-Watson kernel regression, which is interesting. 2. In general, I find the paper well-written, well-structured, supplied with nice illustrations and easy to follow. 3. 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. 2. There appears to be a missing reference to Figure 3 in the paper. 3. No code available, so I can not check if experiments are reproducible. [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. [2] Kotelevskii, Nikita, et al. "Nonparametric Uncertainty Quantification for Single Deterministic Neural Network." Advances in Neural Information Processing Systems 35 (2022): 36308-36323.

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? 2. 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? 3. 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. [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. [4] van Amersfoort, Joost, et al. "On feature collapse and deep kernel learning for single forward pass uncertainty." arXiv preprint arXiv:2102.11409 (2021). Edit: I would like to thank the authors for their answers during rebuttal period.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Authors discussed limitations in the main part of the paper.

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.

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

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

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”. * 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. * Why is it more effective to regularize a non-parametric predictor (Eqn 7) than to directly regularize a standard classification loss? * 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 Misc comments: * 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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Limitations are addressed [Sec 6].

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.

Strengths

### Originality The 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. ### Clarity I found the paper very well-written and structured. Figures are clean and the method description is clear. ### Quality The paper presents interesting results in a rather fair setting, with error bars and well explained experimental protocol. ### Significance This 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 The 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). Typos: - l. 255 “and” ### Quality Despite the very positive points listed above, I still have a few concerns (from the more important to the least): 1. 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. 2. 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. 3. Is it fair to compare IRM to your representations + MLP given that IRM is trained with a simple linear classifier? 4. 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? 5. 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? 6. 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? 7. 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. ### Significance Despite the very interesting method and important questions addressed, I think the paper experimental part still needs a little work.

Questions

Question above

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

4 excellent

Contribution

2 fair

Limitations

Limitations of the method are cited and discussed in the appendix.

acknowledging author rebuttal

I have read the author response and other reviews. Thank you to the authors for acknowledging my concerns and helping me to better understand the implementation (temperature selection, etc.). I don't think we are on the same page about the ISIC results. Even if I look at just that column, I see some $NB^B$ flavors that do better than ERM, and some that do worse. I really appreciate the completeness of the experiment design w.r.t. trying different baselines and variants of the proposed method. I'm just not totally comfortable with how the results are described. Correct me if I'm missing something. I think the paper will indeed be stronger by discussing more related works. But the submission and rebuttal as-is, especially given that a theoretical analysis of when NW can succeed is out of scope, is not enough to move the needle for me. I maintain my score.

Answer to authors

I would like to thank the authors for their detailed and thoughtful rebuttal. My concerns were well addressed as long as the proposed clarifications are added to the manuscript. Just to make sure my comment concerning IRM is clear, I think that saying you “finetune a” and **linear classifier** instead of a "**fully connected classifier**“ in l.205 would be clearer and avoid doubts concerning fairness.