Mitigating Spurious Bias with Last-Layer Selective Activation Retraining

1. Regarding learned latent representations cannot be factored into source, spurious and noise components, and theoretical backing of this method

   Thanks for raising this concern. We apologize for the confusion. Actually, we do not assume that learned representations can be factored into individual components, and our method does not aim to disentangle spurious dimensions from core dimensions. In a nutshell, we aim to identify dimensions that reflect a model's spurious prediction behavior and then block the contributions from these dimensions for retraining. To better explain our method, we have provided a theoretical analysis in Section 3.3. Specifically, for dimension (neuron) selection, our principal is to select those that reflect spurious prediction behaviors of the model, as formalized in Proposition 1. We call these dimensions as spurious dimensions. We have refined our definition on the spurious and core dimensions in the revised paper (Line 233-236). To further justify our algorithmic designs, we provide two theorems in Section 3.3. Theorem 1 proves that the spuriousness score defined in Eq. (5) indeed selects dimensions that meet our selection principal. Theorem 2 proves that our method, LaSAR, can effectively mitigate spurious bias in the model. Please kindly refer to our revised paper for details.

2. Regarding experiments on 9-Class ImageNet / ImageNet-A

   Thanks for the suggestion. We have provided the experimental results on the ImageNet-9 and ImageNet-A datasets in Table 7 in the revised paper. Our method still works in this challenging setting.

3. The related work section should be updated a bit with relevant works in the area

   Thanks for the suggestion. We have added discussions on the given works in the related work section. Please kindly refer to our revised paper for details.

4. Could also provide results in terms of balanced accuracy (other than accuracy and WGA)?

   Thanks for the suggestion. We have provided balanced accuracy for our methods. We hope to emphasize that while balanced accuracy may provide an overall measure of the performance of a method, WGA is the most commonly used metric in this area and the majority of works reported performance under this metric. Other than WGA, we also provied Acc. gap in our paper, which measures the performance discrepancy across different data groups. We hope these metrics help reinforce your confidence in the effectiveness of our methods.

   Method	Waterbirds	CelebA	MultiNLI	CivilComments
   LaSAR	$94.0_{\pm 0.2}$	$89.1_{\pm 1.1}$	$80.8_{\pm 0.5}$	$85.9_{\pm 0.0}$
   LaSAR$^\dagger$	$93.7_{\pm 0.2}$	$89.7_{\pm 0.3}$	$80.0_{\pm 0.4}$	$86.2_{\pm 0.1}$

5. I think that retraining on a part of the validation set may lead to "unfair" comparison with baseline methods that are not trained also on validation data

   Thanks for raising this concern. We have specifically marked methods (with $^\dagger$) that use validation set for training in Table 1 and Table 2, allowing for fair comparison.

6. Reweighting the classification layer essentially does not "remove" the bias from the whole model, but it is just a correction. Do you think this might be an issue in certain cases?

   Thanks for your question. Our Theorem 2 proves that this in generally would not be an issue, and our method can bring the retrained model very closer to the unbiased model in the parametric space than its original version. If the retrained model does not meet certain robustness criteria in practice, then retraining the whole model would be necessary. However, as shown in the Corollary 1 in Appendix, the unbiased model can only be achieved if the whole model can be retrained with unbiased data. Nevertheless, it is worth to mention that our method is an efficient and post-hoc bias mitigation method which may be of independent interest.

1. Regarding the two very important assumptions of the method: (1) some features in the latent embedding that are responsible for encoding the spurious correlation are confined to single neurons, and can be masked away, and (2) spurious features can be distinguished from core features by looking at the error density.

   Thank you for raising this concern. We would like to clarify that our method does not assume that a neuron exclusively represents a spurious feature that can simply be masked. In reality, a neuron may encode a mixture of spurious and core features. Our method does not attempt to distinguish spurious features from core features directly. Instead, it focuses on selecting neurons that produce high activations for an undesired class, highlighting the model's strong reliance on spurious correlations between the undesired class and certain features represented by those neurons.

   To support the motivation behind our proposed method, we provide a comprehensive theoretical analysis in the paper: (1) outlining the principle for selecting neurons (Proposition 1), (2) proving that the spuriousness score based on the distribution of neuron activations can effectively identify the desired neurons (Theorem 1), and (3) demonstrating that our selective activation retraining method can effectively mitigate spurious bias (Theorem 2). For further details, please refer to the revised paper.

2. Why not provide visualizations using linear dimensionality reduction for the motivating example?

   Thank you for raising this concern. We chose T-SNE because we used the multi-dimensional input data for visualization and there is no linear relationship between input dimensions.

3. How were the baselines implemented? are they openly available (it could make sense to provide the links) or did you re-implemented them

   The baselines we compared against are openly available, and their respective links are provided in the original papers. To ensure a fair and reliable comparison, we used the reported performance metrics from these methods as presented in their papers, as they were evaluated on the same datasets and under the same settings. Leveraging the authors' official results avoids potential discrepancies that could arise from re-implementations, ensuring the integrity of our evaluation. We believe this approach provides the most accurate and convincing comparison between our method and the baselines.

4. How many spurious features were masked away in each of the examples?

   Thank you for the question. Indeed, our method masks the same last-layer neurons for all examples. The table below presents the number of masked neurons and their corresponding percentages among all neurons across the four datasets.

1. Regarding the reasoning of the proposed method to identify the spurious features and theoretical results

   Thanks for raising this concern. We have added a theoretical analysis in the paper to offer insights into our detection method. In summary, we first gave in Proposition 1 the principal for selective activation, which is the core of LaSAR. Proposition 1 specifies the property of a neuron that is affected by spurious bias. Specifically, our selection method select neurons that have high activation values but not beneficial for predicting the target. Then, we theoretically showed in Theorem 1 that our proposed spuriousness score indeed selects neurons in the last layer in a way that follows the proposed principal in Proposition 1. In other words, the score in general can effectively identify neurons affected by spurious bias. Finally, we proved in Theorem 2 that LaSAR can mitigate spurious bias in a model by bringing the model closer to the unbiased one in the parameter space. Please kindly refer to Section 3.3 in the revised paper. All changes are marked in blue color.

2. Why did you define the spuriousness score as (5)?

   Thanks for the suggestion. The score is defined such that a large difference between $\mu_{\text{mis}}$ and $\mu_{\text{cor}}$, i.e., a large $\delta_{i}^y$, indicates a high likelihood of the $i$'th dimension being affected by the spurious bias in the model. In other words, the model incorrectly amplifies a spurious feature in the neuron activation when it should not. In contrast, a negative $\delta_{i}^y$ shows the importance of the $i$'th dimension for predictions as most correctly predicted samples tend to have high activation values on this dimension, while most incorrectly predicted samples have low activation values. We have added this explanation in the revised paper.

3. Regarding the definition of WGA in line 157

   Thanks for pointing this out. We have corrected this in the revised paper.

4. In Section 3.2, they used a synthetic motivating example. It may be better to use a real motivating example.

   Thanks for the suggestion! To better highlight our motivation and contribution, we have replaced the synthetic experiment with a theoretical analysis section and moved the synthetic experiment to the experiment section. Please kindly refer to the revised paper for more details.

5. In line 321, the authors may want to say "equation (6) and equation (7)" rather than equation 6 and equation 7.

   Thanks for pointing this out. We have corrected this in the revised paper.

6. Regarding relation to variable selection in statistics

   Thanks for providing this new perspective. Variable selection methods, such as L1 regularization, does not take into account whether those variables represent a spurious feature or not, and they can be easily biased by the imbalanced distribution of the data. In fact, using traditional variable selection methods may actually amplify spurious bias as they are based on statistics of the data. In contrast, our method not only exploits the distributions of neuron activation values but also utilizes the label information to jointly characterize the properties of spurious features and incorporate them into our algorithmic design.

1. Regarding the contribution

   We hope to emphasize that our contribution does not focus on "misclassified examples of validation data and retraining all layers or the linear head". Indeed, we focus on mitigating spurious bias from the perspective of neuron masking, which has the potential to train a robust model without any group information. Our contributions include: (1) an automatic detection framework that identifies neurons affected by spurious bias, (2) a spurious bias mitigation method that blocks spurious neurons for retraining, and (3) we demonstrate the effectiveness of our proposed method on benchmark datasets and provide theoretical analyses (in the revised paper) to justify our design choices.

2. How is LaSAR fundamentally different from AFR?

   Thanks for the question. AFR is a sample-level method in the sense that it mitigates spurious bias by up-weighting samples that have high prediction losses. In contrast, LaSAR works in the latent embedding space and mitigates spurious bias by selectively blocking neurons that are identified as spurious. Moreover, AFR identifies which samples are important for retraining using the prediction loss from an ERM trained model. LaSAR blocks certain neurons for all the training samples and does not prefer one sample over another during retraining. LaSAR utilizes the distributions of neuron activations to identify spurious bias. Moreover, LaSAR can work in the completely unsupervised spurious bias mitigation setting.

3. Regarding comparison with JTT and AFR when tuned on worst-class accuracy.

   Thank you for bringing this up. In the table below, we provide a comparison between AFR, JTT, and LaSAR, where JTT and AFR were tuned using worst-class accuracy. Typically, models tuned with worst-class accuracy tend to perform worse than those tuned with worst-group accuracy, as the table demonstrates. We want to emphasize that, even without access to group information, LaSAR is already competitive with state-of-the-art methods tuned using worst-group accuracy.

   Method	Waterbirds	CelebA
   JTT	84.2$\_{\pm 0.5}$	52.3$\_{\pm 1.8}$
   AFR	89.0$\_{\pm 2.6}$	68.7$\_{\pm 1.7}$
   LaSAR	91.8$\_{\pm 0.8}$	83.0$\_{\pm 2.8}$

4. Regarding theoretical guarantees about the convergence and stability of the selective activation retraining process

   Thanks for pointing this out. In the revised paper, we have provided a theoretical analysis on a linear model with synthetic data. The main results we obtained are: (1) the spuriousness score indeed selects neurons that are affected by spurious bias (Theorem 1), and (2) the selective activation retraining process will converge to a point in the parametric space of the model closer to the unbiased model (Theorem 2). Please kindly refer the Section 3.3 in the revised paper for details.

5. How does the method ensure that it doesn't accidentally block neurons representing valid but complex feature combinations rather than truly spurious correlations?

   Thank you for raising this important point. We apologize for any confusion and would like to clarify that our method does not explicitly aim to block purely spurious features or activate purely core features. As shown in Eq. (11), even for a simple model, neuron activation represents a mixture of spurious and core components from the input. The principle behind our selective activation method is explained in Proposition 1. Fundamentally, we aim to identify spurious behaviors in the model by analyzing neuron activations. In our theoretical framework, it is $\gamma^T\mathbf{w}_{\text{spu},i}<0$. Blocking the identified neurons effectively removes contributions arising from these spurious behaviors. While this approach cannot fully ensure that it won’t unintentionally block neurons representing valid but complex feature combinations, both the theoretical result (Theorem 2) and empirical evidence (performance on benchmark datasets) indicate that the tradeoff is worthwhile.

6. How does the method handle cases where features might be spurious in some contexts but valid in others?

   Thank you for raising this question. As mentioned in Lines 229–232, we argue that in a well-defined classification task, a spurious dimension identified for one class cannot serve as a key contributor to predicting another class. For instance, in a task that involves classifying between "rectangle" and "blue color", a dimension strongly associated with "blue color" for the "rectangle" class cannot reliably predict the "blue color" class when given a blue rectangle, as this would result in ambiguity. Therefore, we block the identified spurious dimensions across all classes.

Thank you for your thoughtful feedback and for reviewing our responses. We sincerely appreciate your encouragement to incorporate the suggested changes and discussions into future iterations of the paper.

We would like to highlight that many of these clarifications and improvements have already been incorporated into the current revision:

• Section 3.3 now provides formal guarantees on LaSAR’s spurious bias mitigation.
• We clarified the definitions of spurious and core dimensions (L237–238) and explicitly discussed the assumption of spurious and core features (L238–241).
• Differences between LaSAR and retraining methods like AFR and JTT are elaborated (L133–134, L344–346, Section A.3), emphasizing LaSAR’s novelty in dimension-level spurious bias control (L87–88, L91–92), particularly in unsupervised settings (L59–60, L92–95).

All major changes are highlighted in blue for ease of review. These updates aim to further clarify our contributions and underscore LaSAR’s theoretical and practical strengths.

We remain committed to refining the manuscript to ensure clarity and impact. In light of the substantial improvements already made, we kindly request that you reconsider your score to reflect the progress in the latest version.

Thank you again for your valuable feedback, which has been instrumental in strengthening our work.

1. Including more rigorous analysis or theoretical guarantees would strengthen the paper's claims about the effectiveness of LaSAR.

   Thanks for the suggestion. We have provided a theoretical analysis in the reivsed paper (Section 3.3, highlighted in blue color) with proofs provided in Appendix. In summary, we first gave in Proposition 1 the principal for selective activation, which is the core of LaSAR. Proposition 1 specifies the property of a neuron that is affected by spurious bias. Specifically, our selection method does not select neurons that exclusively represent a spurious or a core feature; instead, we select neurons that have high activation values but not for the target. Then, we theoretically showed in Theorem 1 that our proposed spuriousness score indeed selects neurons in the last layer in a way that follows the proposed principal in Proposition 1. In other words, the score in general can effectively indentify neurons affected by spurious bias. Finally, we proved in Theorem 2 that LaSAR can mitigate spurious bias in a model by brining the model closer to the unbiased one in the parametric space.

2. Provide the distribution of the proposed spuriousness score across neurons in different datasets.

   Thanks for the suggestion. We have added additional visiualizations in Fig. 5 to Fig. 8 in Appendix to further support the finding that distributions of activation values provide useful information for identifying neurons affected by spurious bias (which is also theoretically proved in Theorem 1). The samples used for visualization are selected based on their activation values on the identified neurons, and we selected five samples that have the largest activation values on each identified neuron. In this way, we select samples that are mostly representative of the feature which the neuron represents.

3. Regarding the contribution

   We hope to clarify that our contribution does not focus on demonstrating ``the phenomenon that spurious neurons and core neurons can be separated". Indeed, the prior work offers insights into spurious bias mitigation in the latent representations. Instead of demonstration, our contribution in this paper is to propose an automatic method that can explicitly identify individual neurons affected by spurious bias, which is different from prior works where this can only be achieved after retraining the model via bias mitigation methods, such as balanced sampling. Moreover, our method is fundamentally different from retraining the last layer while up-weighting the incorrect samples. Please kindly refer to the discussion in our answer to your question on this issue (Point 7). Finally, thanks to your suggestion, we have added a theoretical analysis to elucidate the principal of our method and validate our design choices.

4. The neuron masking algorithm assumes that a neuron can represent part of the spurious features, which is a strong assumption that may not always hold true.

   Thanks for raising this concern. We hope to clarify that we do not make the assumption when using the neuron masking algorithm. Indeed, before this step, in Section 3.2.1, we have an additional step to identify neurons that are affected by spurious bias, i.e., neurons that mainly represent spurious features. Therefore, we ensure that the neuron masking algorithm masks out the correct neurons. Additionally, as shown in our theoretical analysis, the step in Section 3.2.1 indeed selects neurons with a strong reliance on spurious features. Empirically, the heatmaps provided in Fig. 5 to Fig. 8 support that a neuron can represent part of spurious features in a biased model.

5. It is unclear why masking the last layer is necessarily better than masking a middle layer.

   Thanks for bringing this up. Indeed, we did not attempt to make the conclusion that masking the last layer is the optimal way of retraining. We made the first attempt to mask out neurons before the last layer for spurious bias mitigation and proved both empirically and theoretically that our approach is effective. We also emphasized in the revised paper that this approach is efficient. Masking neurons in a middle layer brings more freedom in designing methods for spurious bias mitigation, but this would also complicate the analysis and improve computational complexity. We greatly appreciate your insight and recognize it as a promising direction for future work.

6. JTT Algorithm Classification

   Thanks for raising this concern. JTT does not need group information in the training, but it does require group information in the validation data for achieving optimal performance, which satisfies our definition of semi-supervised spurious bias mitigation. We have clarified this point in the related work section.

I appreciate the authors’ response and the additional experiments. I have raised my score by 1 level. However, my main concern remains regarding Comment 7. The authors explain how LaSAR operates differently from JTT, which I acknowledge and agree with. However, my point is that LaSAR’s approach of excluding neurons based on correct and incorrect samples is quite similar to retraining the last layer while up-weighting misclassified samples. While I understand that LaSAR allows for controlling embedding dimensions, this does not seem to be a critical difference to me.