Neural Collapse Inspired Feature Alignment for Out-of-Distribution Generalization

Respond to W1: Thanks for the comments. We would like to explain as follows:

The OOD methods encompass techniques for addressing spurious correlations. Our comparison methods are comprehensive, covering both approaches that handle spurious correlations, such as IRM, VREx, GroupDRO, and other OOD methods like DG, including MLDG, MMD, and CDANN. Additionally, we have compared our work with those proposed by the reviewer. Our setting is distinct in that [a-e] address bias issues, while we focus on spurious correlations and follow [1] to benchmark our approach against state-of-the-art baselines.

Comparison with [a]: [a] involves utilizing an auxiliary network to identify and eliminate spurious features in the feature space early on. While, our method leverages the phenomenon of neural collapse to achieve feature alignment, thereby eliminating spurious features.

Comparison with [b]: [b] utilizes HSIC to achieve independence and learns invariant representations on biased data. In contrast, our method employs a two-stage interactive process, utilizing neural collapse to ultimately achieve learning of invariant features.

Comparison with [c]: [c] achieves invariant representation learning during network training by partitioning the hidden layer space into shortcut and complementary subspaces, guided by loss functions. In contrast, our method automatically partitions across different environments, leveraging the properties of class-invariant features to align class-specific features in various environments.

Comparison with [d]: [d] also follows a two-stage process: the first stage utilizes ERM for training, and the second stage handles instances where ERM identification errors occur. In contrast, our approach first partitions environments. Then, we employ neural collapse to align class-specific features.

Comparison with [e]: [e] involves constructing two networks: one amplifies bias information using GCE, representing variable information, while the other learns debiased information, representing semantic information. In contrast, we employ a two-stage interactive approach to learn invariant networks.

[1] Map: Towards balanced generalization of iid and ood through model-agnostic adapters, 2023.

Respond to W2: Thanks for your questions. We would like to explain this question as follows: We are currently following existing work where neither the baseline models in the literature nor the existing benchmarks incorporate stacked features.

Respond to W3: Thanks for the comments. We would like to clarify it as follows: We follow the approach in [1,2], where we introduce some noise during the initialization of masks to enhance their randomness. This enables the masks to better capture invariant features.

[1] Invariant Representation Learning for Multimedia Recommendation [2] Kernelized Heterogeneous Risk Minimization

Respond to W4: Thanks for the comments. We would like to clarify it as follows: In the introduction, we employ the F-norm to assess the feature conditions resulting from different methods. According to the theory of neural collapse, after network training, features collapse onto a simplex. To visually illustrate this phenomenon, we use the F-norm to measure the discrepancy between the features post-training and a standard simplex, as depicted in our first figure. In the main text, we utilize Eq. 5 to align the post-trained features to the standard simplex using a push-pull mechanism. In different environments, invariant features are inherently more similar, thus they consistently align to the same ETH.

Respond to W5: We will incorporate significant improvements in the subsequent revisions of the paper.

Respond to Q1: Thanks for the very constructive comments. We would like to clarify it as follows: In Table 1, we directly apply the IRM algorithm to restrict parameter variations, thereby achieving invariant learning without utilizing a masking mechanism. In Table 2, to demonstrate that our approach with masking can also yield improved results, we substitute the process of learning masks with the IRM loss as shown bellow. $\mathcal{L}\_{mask}=\mathbb{E}\_{e\in\mathcal{E}}\mathcal{L}^{e}+\alpha\\|\mathrm{Var}\_{e\in\mathcal{E}}(\nabla\_{\Theta^{mask}}\mathcal{L}^{e})\odot\mu\\|^{2}+\lambda\\|\mathbf{m}\\|^{2},$ where the $\mathcal{L}^{e}$ is the classification loss, the second term is the constraint across environments, and the third term is a regularization term. $\mathcal{L}^e$ is the average loss value inside the environment $e$. Consequently, the accuracies differ between the two tables.

Respond to Q2: Thanks for the comment. We would like to clarify it as follows: In our paper, we have clarified that when the environment is unknown, we first randomly partition the data into environments. Subsequently, through iterative interaction between environment partitioning as described in Algorithm 1 and mask learning, we achieve convergence in both environment partitioning and mask learning processes. We learn an invariant mask to distinguish between spurious and semantic features in the input data. The learning phase of the invariant mask minimizes the loss function defined in Eq. 5. When this process converges, we assert that the features align well with the standard simplex post-training. Importantly, since we minimize the loss on the masked features, which represent semantic features, there is no issue of conflicting learning between two types of features.

Respond to Q3: Thanks for the comments. We would like to clarify it as follows: When the number of environments is unknown, we treat it as a hyperparameter in accordance with our previous method. By adjusting the size of this hyperparameter, we aim to achieve optimal out-of-distribution (OOD) generalization performance. In our paper, we also conducted ablation experiments to investigate scenarios where the number of environments is unknown.

We compared the papers provided by the reviewer. Due to the unavailability of code for papers [a, c, e], we evaluated the methods from [b, d] on ColoredCOCO and COCOPlaces. We also compared our approach with Rubi [f] and LearnedMixin [g]. Our method demonstrates superior performance compared to these approaches.

[a] Overcoming simplicity bias in deep networks using a feature sieve. ICML2023.
[b] Learning de-biased representations with biased representations. ICML2020.
[c] Chroma-vae: Mitigating shortcut learning with generative classifiers. NeurIPS2022.
[d] Just train twice: Improving group robustness without training group information. ICML2021.
[e] Learning from failure: De-biasing classifier from biased classifier. NeurIPS2020. [f] Rubi: Reducing unimodal biases for visual question answering. Neurips2019. [g] Don’t take the easy way out: Ensemble based methods for avoiding known dataset biases. EMNLP2019.

We sincerely appreciate the reviewer for raising the score. Regarding the concerns you mentioned, we would like to offer further clarification. Although we are currently unable to use CAM visualization to demonstrate our mask performance, we plan to include a mask visualization in the final version to provide a more intuitive presentation. Below, we present more granular experimental results and detailed analysis to clarify the mask learning process for the two types of features.

Experiments

We designed an experiment using our proposed method for a two-stage interactive learning process to obtain partitioned environments and invariant masks $m$. To verify that $m$ can effectively distinguish between spurious and invariant (or semantic) features, we conducted additional experiments. Specifically, we used the feature information from $m \odot x$ and $(1-m) \odot x$ to predict during the test phase and compared their performance. The results show that using $m \odot x$ significantly improves the model's out-of-distribution (OOD) generalization performance compared to using $(1-m) \odot x$. This demonstrates that our proposed mask accurately captures both invariant (or semantic) and spurious features.

Analysis

We initialize the masks, so that $m \odot x$ may not accurately reflect the input's semantic features. To obtain a semantically representative mask $m$, we align features of different categories from various environments to a standard simplex equiangular tight frame (ETF), leveraging the property of neural collapse (where features from different classes form a canonical simplex ETF upon balanced training). In response to the reviewer's concerns, we can address them using two metrics in our paper.

• The variation in environment partitioning: It is significant at the beginning because the initial mask $m$ does not accurately capture the invariant features of the input. Consequently, $1 - m$ also fails to effectively represent spurious features, leading to suboptimal environment partitioning. However, as training progresses with the neural collapse algorithm, the semantic mask $m$ improves, leading to better separation of semantic features and spurious features. Ultimately, environment partitioning stabilizes, showing minimal changes when it is complete. The table below shows changes in environment partitioning between the new environment and old environment over multiple interactions using our method in ColoredCOCO dataset. It indicates that the variation in environments decreases over interactive time, suggesting that $1 - m$ increasingly captures spurious information for better environment partitioning.

where $\Delta N = N^{(e)}\_{new} - N^{(e)}_{old}$. $\Delta N$ represents the number of sample changes between the previous iteration (i.e., the old environment) and this iteration (i.e., the new environment). As the number of iterations increases, the change in the number of samples between environments gradually converges to 0, representing the environment partitioning fit.

• The loss incurred from using the neural collapse method for alignment in the models in different environments: Initially, the loss is high because $m$ is initialized. However, as the number of interactions increases, the loss gradually decreases to a low level. At this point, $m$ effectively captures the semantic features across different environments. The table below shows the initial loss at different stages of interaction using our method in ColoredCOCO dataset. It reveals a successive decrease in loss, indicating that $m$ increasingly captures the invariant (or semantic) features of the input more effectively.

where $\mathcal{L}\_{mask}=-\log\frac{n\_{e,y}^\gamma\exp(\beta\cdot\mathbf{v}\_y^T\mathbf{f})}{\sum\_{k} n\_{e,k}^\gamma\exp(\beta\cdot\mathbf{v}\_k^T\mathbf{f})}, k\in K, e \in \mathcal{E}$.

We sincerely hope the reviewer will reevaluate our paper based on the additional experiments and explanations provided above. We look forward to your feedback if you have any further questions!

Respond to Q1: Thank you for the reviewer's suggestions. We will include citations to these papers in a subsequent version.

We sincerely appreciate the reviewer's positive feedback on our manuscript. Due to the time constraints of the rebuttal period, training on datasets like NICO, which require extensive computational resources, took longer than anticipated. Therefore, we have to present all the supplementary experimental results now. We sincerely hope that the reviewer will re-evaluate our work and decide the final rating based on the following discussions.

The following papers seem missing from the neural collapse literature that may be helpful.

We thank the reviewer for providing three important neural collapse literature, and would like to incorporate the following discussion in our final version.

• [1] employs the neural collapse algorithm in transfer learning to facilitate few-shot learning, enabling the training of a linear classifier on top of the learned penultimate layer and achieving good results. In contrast, our approach utilizes the neural collapse phenomenon in the context of imbalanced categories across different environments to learn invariant masks and perform environment partitioning, ultimately enhancing OOD generalization.
• [2] primarily leverages the phenomenon of neural collapse to guide semantic clustering of categories in self-supervised learning. This is achieved by using regularization to promote clustering within the same category and to separate clusters of different categories. In contrast, our method aligns features of the same category across different environments to learn invariant feature masks. These masks are then used to enable the network to better learn invariant features, thereby enhancing OOD generalization.
• [3] represents a highly significant contribution that can greatly enhance the quality of our work. It primarily focuses on using the neural collapse phenomenon to introduce a novel method for measuring the generalization bounds of neural networks. This method, being empirically non-vacuous and largely independent of depth, offers superior generalization bounds compared to traditional measurement approaches. In future versions, our method can draw on the approach presented in this paper to conduct a theoretical analysis of the generalization bounds.

We also investigate recent advances in neural collapse algorithm for other applications as below.

• [4] investigates the use of the neural collapse phenomenon within a continual learning setting. The scenario involves introducing new classes in a few-shot manner after the model has been trained on a large dataset, leading to an imbalance between old and new classes and resulting in severe catastrophic forgetting. By employing the neural collapse method, the model can align old class features and separate new class features, thereby preserving knowledge of old classes while learning new class knowledge effectively.
• [5] addresses the problem of sample imbalance in large models using the neural collapse method. By leveraging the property of neural collapse where features between classes exhibit a standard ETF distribution after balanced training, the method guides the collapse of features across different categories. This approach enhances the performance of large models in the context of sample imbalance.

The experiments are a bit lacking, but are convincing.

To make the experiment section in our paper more convincing, we have added extensive experiment as below.

• We conduct additional ablation experiments to explore the effect of mask position (i.e., different neural network layers) on the OOD performance, by applying masking to the output features of different layers in the network..

We found that putting the mask in the output layer (i.e., the previous layer of the linear classifier) gave the optimal results in general, which is also consistent with the findings of neural collapse.

• In comparison with the two-stage HRM method, we found that our approach stably outperforms HRM.

• We also compared our method with more competing baselines and found that our method outperforms them.

We will definitely put the above additional experimental results into our revised manuscript -- thank you so much!

[1] "On the role of neural collapse in transfer learning." arXiv, 2021.

[2] "Reverse engineering self-supervised learning." NeurIPS, 2023.

[3] "Comparative generalization bounds for deep neural networks." TMLR, 2023.

[4] "Neural collapse inspired feature-classifier alignment for few-shot class incremental learning." ICLR, 2023.

[5] "Bridging the gap: neural collapse inspired prompt tuning for generalization under class imbalance." KDD, 2024.

[6] "Heterogeneous risk minimization." ICML, 2021.

[7] "Learning de-biased representations with biased representations." ICML, 2020.

[8] "Just train twice: Improving group robustness without training group information." ICML, 2021.

[9] "Rubi: Reducing unimodal biases for visual question answering." NeurIPS, 2019.

[10] "Don't take the easy way out: Ensemble based methods for avoiding known dataset biases." EMNLP, 2019.

Respond to W1: Thank you for pointing out the issue. We indeed omitted the comparisons in our manuscript. We have actually demonstrated this in Figure 1, where we used the F-norm to measure the degree of alignment. A smaller F-norm indicates that, after training, the feature prototypes are closer to the standard ETF.

Respond to W2: Thank you for the reviewer's reminder. We sincerely apologize for the misunderstanding caused by our negligence. A precise description will be provided in future revisions.

As shown in Figure 3, our method represents the variable features of the input as $(1-m) \odot x$. By using $(1-m)$, we obtain the variable features (equivalent to the environmental information) from the previously partitioned data. The neural networks corresponding to different environments are then trained on these variable features. After training, the neural networks from different environments predict the respective logits. For each input’s true label, the index with the maximum value determines the environment of the input. The structure of the local model is the same as the local model used for alignment with the neural collapse algorithm later.

Respond to W3: Thanks for the valuable comments. We would like to explain as follows: The primary difference between the two methods is the location of the masking. Pixel-level masking is applied to the input image, offering stronger interpretability. Feature-level masking, on the other hand, applies masking to the output features of the neural network, achieving better results and reducing the dimensionality of the masking. To address the reviewers' concerns, we performed experiments by applying masking to the output features of different layers in the network.

Respond to W4: Thanks for the valuable comments. We would like to explain as follows: Our method aligns invariant features across different environments by leveraging the phenomenon of neural collapse. For the same category, invariant features should be similar; thus, we use the same ETF for alignment, facilitating mask learning. Although the HRM algorithm also employs a two-stage approach, it partitions environments through clustering and then learns the mask through regularization. We conducted comparative experiments with the HRM algorithm to evaluate our method.