Paper: Removing Spurious Concepts from Neural Network Representations via Joint Subspace Estimation (ICLR-2024)

Reviewer comment (Writing):
The organization and clarity of this paper needs improvement to enhance its readability. For example, section 2 seems to be a discussion on the problem setting, but it's not clear what exactly is the data distribution model that the paper focus on upon reading the section. The assumptions are not properly and formally stated. A more structured discussion, with clear definitions and assumptions, would make it easier to understand the focus of this paper.

Author response:
We have tried to improve the manuscript's readability, by more clearly structuring the different topics that we need to properly discuss and assess JSE. In Section 2 the focus is on the data distribution model (eq. 1) and the general approach and assumptions of (linear) concept removal. The assumptions that are specific to JSE can be found in Section 3. We have moved the discussion of OOD generalization from Section 2 to Section 4, as we realized this only becomes relevant in the context of the experiments we present in the manuscript.

Reviewer comment (Evaluation):
It seems that ERM outperforms many of the baselines in both the toy data and real data. Given that the baselines are intended to enhance the model's performance on the worst-group accuracy, it would be beneficial if the authors could shed light on why ERM seems to have a superior performance than the presented baselines.

Author response:
Existing concept-removal methods tend to remove the spurious concept, but also other features that are correlated with the concept [1]. Because these main-task features are removed, a classifier that is trained on the transformed embeddings tends to do much worse than standard ERM. The purpose of Figure 3 (in the revised manuscript) is to show that JSE does not remove main-task features, as opposed to other concept-removal methods. To improve the clarity of this message, we moved the ERM results to the figure where JSE is compared to instance reweighing methods. In that figure, ERM can be seen as the default baseline.

Reviewer comment (Experiments):
There are some methods in the previous line of work that are mentioned in Appendix E, but I'm confused which method exactly is implemented as the baseline. By group weighted (GW) ERM, the authors cited papers on last layer retraining (DFR) as well as paper discussing group/class-balanced baselines that are not related to last layer retraining. These methods share similar concepts but have different methodologies and performances. I would appreciate a clearer explanation/citation of the exact compared baseline method.

Author response:
We thank the reviewer for raising this issue. We have added a description of the method in the main-text of the paper. In short, the last layer of the neural network is trained, and in each batch there is an equal probability to encounter a sample from each group – e.g. the probability of encountering a land bird on land is equal to that of a land bird on water. This is equivalent to the so-called 'RWG' baseline implemented in [2], but only applied to the last-layer - this why it was cited in the original manuscript.

Reviewer comment (Evaluation):
In appendix E, it seems that the proposed method is not doing better than the baseline on the real datasets, but significantly outperforms it on the toy data. Is there a reasoning on why the toy data distribution specifically suits the proposed method, and would we be able to find similar distributions in real data?

Author response:
In the original manuscript we tried to clarify this in the appendix. In the revised manuscript, a reasoning is presented at the beginning of Section 4.3. JSE removes the spurious embedding, which makes it robust to changes at the more fundamental level of the conditional distribution $ p(y\_{\mathrm{mt}}| z\_{\mathrm{sp}}) .$ The other methods aim for robustness against changes in $ p(y\_{\mathrm{mt}}| y\_{\mathrm{sp}}).$ This explains why JSE outperforms the other methods for the Toy dataset, because in this setting the conditional distribution $p(y\_{\mathrm{mt}}| z\_{\mathrm{sp}})$ itself is changed for the OOD data.

Reviewer comment (Experiments):
The considered datasets typically have a very large value of $\rho$. For example, Waterbirds have $\rho=0.95$. The experiments of this paper stops at $\rho=0.9$, which does not seem like a conventional setting. It would be better to show improvement on the original dataset as well.

Author response:
We chose the range of $\rho$ to ensure consistency across all datasets, but we agree with the reviewer that in the specific case of Waterbirds $\rho=0.95$ should be included. We have done so in the revised manuscript in Section 4.3. We did not do this for the comparison with other concept-removal methods (Section 4.2), because at lower spurious correlations it was already clear that JSE outperforms them.

Reviewer comment (Theory):
unclear description about the orthogonality assumption. How to guarantee this assumption at any scene? Why linear subspaces are orthogonal? Why linear?

Author response:
*Why linear*: this is in line with the so-called linear subspace hypothesis, for which there is existing empirical evidence (see for instance [1,2]). Even if there is a non-linear representation of the concept, a linear classifier cannot use this information for the main-task if the concept is binary (see [3]).

*Why orthogonal*: as stated in the paper, this is an assumption. We assume that the two concepts have a distinct representation in the embeddings – e.g. part of the embeddings are related to the background, and part are related to the bird type. Orthogonality is an intuitive way to formalize this distinctiveness (see for instance [4]). See also the general comment about the orthogonality assumption (point 4) and the revised manuscript.

Reviewer comment (Novelty):
The novelty of the proposed jointly subspace estimation is limited and unclear.

Author response:
We respectfully maintain a different perspective than the reviewer. Existing concept removal methods do not jointly determine the subspaces of main-task and spurious concepts. This makes them vulnerable to spurious correlations – as illustrated by our experiments. Furthermore, in our revised manuscript we have emphasized the advantages of concept-removal methods (interpretability, transfer learning, fairness/sensitive biases) as opposed to, for example, instance reweighing methods.

Reviewer comment (Theory):
Why logistic regression can separate spurious from main-task concept?

Author response:
In principle, any linear classifier can be used. We chose logistic regression for convenience and because it allows for fast optimization, but if one believes that the best linear separation can be made by a support vector machine (SVM) or other generalized linear model, then they can use these to determinine the subspaces.

An additional argument is that we apply logistic regression to the embeddings of the final layer. The decision boundaries of the original DNN are in fact linear in this space. Furthermore, the softmax output activation functions correspond to a logistic regression model from the embeddings to the outputs.

Reviewer comment (Experiments):
One of my concerns is regarding the experiment settings and baselines. The benchmark datasets used in this paper (Waterbirds, CelebA, and MultiNLI) are also used to evaluate debiased learning algorithms (GroupDRO, DFR, JTT, etc.). Since the settings are the same, JSE should be compared with these debiased learning methods. Alternatively, another option can be made by designing and showing experiment settings that differentiate concept removal methods from debiased learning algorithms.

Author response:
See the general comment (point 1) and the revised manuscript (in particular Section 4.3). On top of an empirical comparison, the revised manuscript also contains additional arguments in favor of concept removal methods (explainability, transfer learning, fairness/undesirable biases, robustness against feature shifts instead of label shifts).

Reviewer comment (Evaluation):
The author leveraged PCA to reduce the computational cost, but it would bring out the information loss.

Author response:
We understand the concern of the reviewer. At 100 components, a large percentage of the variation in the data ($>80\%$) is preserved (across all datasets). Furthermore, on realistic datasets the dimensions of the estimated subspaces were never found to be larger than 10, which is well below the 100 dimensions after PCA was conducted. It seems unlikely that the removed dimensions contain significant amounts of information about one of the concepts.

The choice for using PCA was for computational convenience. We are currently running experiments of JSE without PCA and will add the results before the end of the rebuttal phase.

Reviewer comment (Experiments):
Is it possible to train and classify another waterbird dataset in which waterbirds are on the sand background, and landbirds are on the forest background (without $y_{sp}$ for both training and validation dataset) if we leverage the subspace information from the original waterbirds dataset?

Author response:
We thank the reviewer for this useful suggestion. We believe this is certainly possible. Concept-removal methods can be used for transfer learning. In the current context, one could project out the background features (found with JSE) from the embedding space and do last-layer re-training on the new dataset.

It may even be possible that re-training is not even needed. The Grad-CAM figure in the manuscript suggests that after JSE the model is actually looking at the correct input features (the bird, not the background). Presented with a new dataset, this should still be the case. We have included transfer learning as one of the potential benefits of concept-removal methods, but consider experiments using JSE for transfer learning as beyond the scope of the present manuscript.

Reviewer comment (Theory):
Does the orthogonality assumption always hold for every dataset and trained model?

Author response:
We see that at high spurious correlation JSE has more difficulty with the text dataset (MultiNLI) than with the vision data, although it is still competitive with instance reweighing techniques and outperforming the other concept-removal methods. In the manuscript we elaborate on this and suggest that during the finetuning of BERT the main-task concept and spurious concept become overlapping in the [CLS] embedding, in line with observations from previous work [1]. This is potentially a sign that the orthogonality assumption does not always hold, or that the dimensions of the subspaces is so small that it is difficult to identify orthogonal subspaces within the subspaces. See also the general comment about the orthogonality assumption and the revised manuscript.

Reviewer comment (Evaluation):
Could you compare the computational cost of JSE with INLP?

Author response:
Let’s presume that the spurious concept is represented by a $d\_{\mathrm{sp}}$-dimensional subspace, and the main-task concept by a $d\_{\mathrm{mt}}$-dimensional subspace. INLP involves estimating $d\_{\mathrm{sp}}$ logistic regressions, whereas JSE involves estimating $d\_{\mathrm{mt}}$ *  $d\_{\mathrm{sp}}$ logistic regressions. On realistic datasets the dimensions of the estimated subspaces were never found to be larger than 10.

Reviewer comment (Experiments):
Why $V_{sp}^\perp$ is used to train a last layer instead of $V_{mt}$? Is there an ablation study?

Author response:
In our original submission (as well as the updated version) there is an ablation study on this very question in Section C.2 of the appendix.

Reviewer comment (Theory):
The method assumes the pixel corresponding to the main concept and the spurious concept doesn't overlap. This may not always hold true — for example, if the main concept is the shape and the spurious concept is colour, the pixels can overlap.

Author response:
Thank you for raising this. JSE does not need the assumption of non-overlapping main-task and spurious input features. In fact, the sets of pixels corresponding to the main-task and spurious concept can even be identical. As long as the correlation between the pixel values and the respective labels is different, the DNN is able to learn distinct representations in the embedding space. The latter is in fact an assumption, which we explicitly state in Section 2 and which is based on existing literature [1,2,3]. We have corrected the text in the manuscript accordingly.

Reviewer comment (Theory):
Definition of spurious concept in **eqn 1** is not correct.  
> label $y_mt$ and the spurious features $x_sp$ are independent  
They are independent but correlated in the training data (spurious **correlations**).

Author response:
We respectfully hold a different view than the reviewer on this. Since $x\_{\mathrm{sp}}$ and $y\_{\mathrm{mt}}$ are correlated, they are necessarily dependent random variables. To clarify, $x\_{\mathrm{sp}}$ is not assumed to be independent of $y\_{\mathrm{mt}}.$ There is a causal relationship between $y\_{\mathrm{mt}}$ and $x\_{\mathrm{mt}},$ and because there is a dependency between $x\_{\mathrm{sp}}$ and $x\_{\mathrm{mt}},$ $x\_{\mathrm{sp}}$ is not independent of $y\_{\mathrm{mt}}.$

Regarding the second statement, thank you for raising this issue. We understand the confusion, as we were referring to only the conditional distributions in eq. 1. Other conditional distributions indeed will change due to a change in the joint of $x\_{\mathrm{mt}}$ and $x\_{\mathrm{sp}}$. We now refer explicitly to the conditional distributions of eq. 1.

Reviewer comment (Theory):
The statement "while the conditional and marginal distributions are same" is incorrect. If both the conditional and marginal distributions are the same, the joint distribution will be the same too.

Author response:
We respectfully hold a different view than the reviewer on this. Since $x\_{\mathrm{sp}}$ and $y\_{\mathrm{mt}}$ are correlated, they are necessarily dependent random variables. To clarify, $x\_{\mathrm{sp}}$ is not assumed to be independent of $y\_{\mathrm{mt}}.$ There is a causal relationship between $y\_{\mathrm{mt}}$ and $x\_{\mathrm{mt}},$ and because there is a dependency between $x\_{\mathrm{sp}}$ and $x\_{\mathrm{mt}},$ $x\_{\mathrm{sp}}$ is not independent of $y\_{\mathrm{mt}}.$

Regarding the second statement, thank you for raising this issue. We understand the confusion, as we were referring to only the conditional distributions in eq. 1. Other conditional distributions indeed will change due to a change in the joint of $x\_{\mathrm{mt}}$ and $x\_{\mathrm{sp}}$. We now refer explicitly to the conditional distributions of eq. 1.

Reviewer comment (Experiments):
Another major issue is the use of pre-trained ResNet50. Since it is trained on a large ImageNet dataset, it can extract features related to both main and spurious concepts.

Author response:
Our method indeed assumes the presence of both spurious and main-task features in the embedding space. Previous work notes that many large neural networks learn both spurious and main-task features, even when trained on data with a spurious correlation [2]. If this holds, then our method should also work for a fine-tuned Resnet50 model (pre-trained on Imagenet data).

Reviewer comment (Writing):
How do you define main concept and spurious concept? For example, in CelebA, gender classification is a much harder problem than blond vs. non-blond hair.

Author response:
CelebA is a standard dataset for the problem of OOD generalization in the presence of spurious correlations. We followed the default definition of main-task and spurious concept, where gender is defined as the spurious concept.

The hypothesis tests of JSE can actually account for a difference in predictability between the spurious or main-task concept. We show this in Appendix D.