Removing Spurious Concepts from Neural Network Representations via Joint Subspace Estimation

We thank reviewer YHWA for carefully reading our manuscript, and we are grateful for the useful questions and remarks.

unclear description about the orthogonality assumption. How to guarantee this assumption at any scene? Why linear subspaces are orthogonal? Why linear?

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.

The novelty of the proposed jointly subspace estimation is limited and unclear.

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.

Why logistic regression can separate spurious from main-task concept?

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.

[1] Tolga Bolukbasi, Kai-Wei Chang, James Y Zou, Venkatesh Saligrama, and Adam T Kalai.
Man is to computer programmer as woman is to homemaker? debiasing word embed-
dings. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Ad-
vances in Neural Information Processing Systems, volume 29. Curran Associates, Inc.,
2016.

[2] Francisco Vargas and Ryan Cotterell. Exploring the linear subspace hypothesis in gender bias mit-
igation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language
Processing (EMNLP), pp. 2902–2913, Online, November 2020. Association for Computational
Linguistics.

[3] Shauli Ravfogel, Yoav Goldberg, and Ryan Cotterell. Log-linear guardedness and its implications.
In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Vol-
ume 1: Long Papers), pp. 9413–9431, Toronto, Canada, July 2023. Association for Computa-
tional Linguistics

[4] Park, K., Choe, Y. J., & Veitch, V. (2023). The Linear Representation Hypothesis and the Geometry of Large Language Models.

We are deeply grateful to reviewer YBBU for the careful consideration of the manuscript. The feedback has been invaluable for the revision of the manuscript.

The organization and clarity of this paper needs improvement to enhance its readability. For ... the focus of this paper.

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.

I would greatly appreciate more discussions on the ... a superior performance than the presented baselines.

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.

The selection of datasets—Waterbirds, CelebA, and MultiNLI—are frequently ... discussion is needed on this line of methods and why they are not included in comparison.

See the general comment (point 1) and the revised manuscript (in particular Section 4.3).

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, ... of the exact compared baseline method.

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.

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?

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.

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.

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.

[1] Abhinav Kumar, Chenhao Tan, and Amit Sharma. Probing classifiers are unreliable for concept removal and detection. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 17994–18008. Curran Associates, Inc., 2022.

[2] Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In First Conference on Causal Learning and Reasoning, 2022.

We are grateful to reviewer AJaV for examining the paper and for the detailed feedback. Below, we address the feedback point-by-point.

The method depends on the availability of group labels (i.e., main task and spurious concept label), which is usually unavailable during training time.

See the general comment (point 3) and the revised manuscript (appendix E).

The method assumes the pixel corresponding to the main concept and the spurious concept doesn't overlap. This may not always hold true --- for e.g., if the main concept is the shape and the spurious concept is colour, the pixels can overlap.

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.

Definition of spurious concept in ** eqn 1** is not correct.

$$\text{label and the spurious features are independent}$$

They are independent but correlated in the training data (spurious correlations)

$$\text{while the conditional and marginal distributions are same}$$

I think this is incorrect. If both the conditional and marginal distributions are the same, the joint >distribution will be the same too.

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.

Another major is the use of pre-trained Resnet50. Since it is trained on a large ImageNet data, it can extract features related to both main and spurious concepts.

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

The method is benchmarked against enough baselines. Baselines should also include methods that do not use linear subspace projections/estimations, such as [1,2].

See the general comment (point 1) and the revised manuscript (in particular Section 4.3).

How do you define main concept and spurious concept? For e.g., in CelebA, gender classification is a much harder problem than blond vs. non-blond hair.

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.

[1] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is suffi- cient for robustness to spurious correlations. In The Eleventh International Conference on Learn- ing Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.

[2] Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On fea- ture learning in the presence of spurious correlations. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 38516–38532. Curran Associates, Inc., 2022.

[3] Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. Domain-adjusted regression or: ERM may already learn features sufficient for out-of-distribution generalization. In NeurIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications, 2022

We thank reviewer SKcn for carefully reading the manuscript and the critical feedback, which has greatly helped us to improve the manuscript.

One of my concerns is regarding the experiment settings and baselines. ... from debiased learning algorithms.

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

The suggested algorithm contains double for loops, which looks costly.

We understand that the reviewer raises this point. The double for-loop increases the computational cost for JSE, compared to methods such as INLP. However, we are only estimating a limited number of linear classifiers. JSE involves estimating $d\_{\mathrm{mt}} \times d\_{\mathrm{sp}}$ logistic regressions. In our experiments on realistic datasets, the dimensions of the respective subspaces were never found to be larger than 10. We address the issue of computational costs in the revised manuscript.

The author leveraged PCA to reduce the computational cost, but it would bring out the information loss.

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.

Assuming the known $y\_{\mathrm{sp}}$ is not practical. Recently published debiased learning algorithms do not require spurious labels.

See the general comment (point 3) and the revised manuscript (appendix E).

Is it possible to train ... the original waterbirds dataset?

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.

Does the orthogonality assumption always hold for every dataset and trained model?

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.

Could you compare the computational cost of JSE with INLP?

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.

Why $V\_{\mathrm{sp}}^\perp$ is used to train a last layer instead of $V\_{\mathrm{mt}}$? Is there an ablation study?

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.

[1] Fahim Dalvi, Abdul Rafae Khan, Firoj Alam, Nadir Durrani, Jia Xu, and Hassan Sajjad. Dis- covering latent concepts learned in BERT. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022

Thank you for providing comprehensive responses and addressing the reviewer's concerns in the manuscript. While many of my major concerns have been resolved through the rebuttal and revisions, I have decided to maintain my initial score. My remained concerns are regarding the baselines. Notably, the performance of newly added baselines, particularly GDRO, appears lower than indicated in the existing literature. Additionally, the DFR results are still missing even though they are cited in the related works. I believe that addressing these issues in the revision could strengthen the manuscript. However, the current version may not meet the criteria for acceptance.