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

Authors: zhikang  chen
min  zhang
sen  cui
haoxuan  li
gang  niu
mingming  gong
changshui  zhang
kun  zhang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:51:15

Abstract: The spurious correlation between the background features of the image and its label arises due to that the samples labeled with the same class in the training set often co-occurs with a specific background, which will cause the encoder to extract non-semantic features for classification, resulting in poor out-of-distribution generalization performance. Although many studies have been proposed to address this challenge, the semantic and spurious features are still difficult to accurately decouple from the original image and fail to achieve high performance with deep learning models. This paper proposes a novel perspective inspired by neural collapse to solve the spurious correlation problem through the alternate execution of environment partitioning and learning semantic masks. Specifically, we propose to assign an environment to each sample by learning a local model for each environment and using maximum likelihood probability. At the same time, we require that the learned semantic mask neurally collapses to the same simplex equiangular tight frame (ETF) in each environment after being applied to the original input. We conduct extensive experiments on four datasets, and the results demonstrate that our method significantly improves out-of-distribution performance.

Main paper: 
 introduction: 
 The out-of-distribution (OOD) problem refers to the fact that the training dataset and the test dataset have different distributions in different environments, after the training process is completed, what we want is to have correlation between semantic features and label, however, the presence of spurious feature(environmental information) can make false correlation between environmental information and labels. The semantic feature in an image refers to the object of the class, while the spurious feature in an image refers to the background or the environment. As shown in Figure 1 , the background color for digits 0 and 2 is predominantly orange, while the background color for digit 1 is predominantly green. In this figure, the semantic feature includes digits 0, 1 and 2, whereas the spurious feature includes orange and green backgrounds. environment variations from the training dataset, on the one hand, they require a well-partitioned environment in advance, but real datasets usually do not have environment labels; on the other hand, under cross-environment conditions, these methods do not align feature prototypes well between the test environments and the training environments. In this paper, we equate semantic features with invariant features and spurious features with variable features.We first explored the limitations of previous OOD learning methods from a new perspective and conducted a toy example experiment, as shown in Figure 1 . The digits 0 and 2 are associated with orange background, and the digit 1 is associated with green background. The phenomenon of neural collapse, which means that after sufficient training, the categories will collapse to a simplex ETF such that the angles between the feature prototypes of two neighboring categories are equidistant, Thus, in Figure 1 , after the training is completed, the feature prototypes of the digits 0, 1 and 2 occur in three equal parts. Here, to measure the extent of neural collapse, we propose to compute the Frobenius norm (F-norm) of the difference between the feature prototypes and the standard simplex ETF. A smaller F-norm indicates a closer proximity to the standard simplex ETF. The ERM approach [Vapnik et al., 1998 ] that leads to the failure of OOD generalization is due to the endogenous nature of the class features not being able to align the class prototypes after training is complete in different environments. Although IRM-based methods [Arjovsky et al., 2019] theoretically learn a feature representation such that the last layer of the feature extractor is similar across environments, we empirically found that IRM-based methods are not aligned very well, which motivates using neural collapse to align feature for OOD generation.In this paper, in order to bridge this gap, we believe that a feature extractor for OOD generalization across environments should ensure that semantic features are aligned to the same simplex ETF, and we have verified from our experiments that better alignment can significantly improve OOD performance. Specifically, our method can be applied both with and without environment labels. In the absence of environment labels, we can automatically partition the environments and assign local models to different environments. By predicting the variable components of the input, we take the logits corresponding to the label (one-hot encoding) from the predictions of the local models in different environments, form a vector, and select the maximum value of the vector. The corresponding environment of this maximum value is assigned as the new environment for the input. When the new environment shows minimal changes compared to the old environment, we consider the environment partitioning to be complete. In Figure 3 , we provide a detailed example to illustrate our environment partitioning method. When environment labels are available, we learn masks to extract semantic components, we firstly fix a simplex ETF classifier, and in different environments, for semantic features, we pull them all to the corresponding position of the same class prototypes, thus realizing the alignment operation.The main contributions of this paper are summarized as follows:• We explore the OOD problem from a new perspective, namely the use of neural collapse to guide feature alignment for OOD generalization. • We explore the separation of semantic features and spurious features under conditions without environment labels, as well as with given environment labels. • We conduct extensive experiments on four publicly available datasets to validate the effectiveness of the proposed methods.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We have done this work.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Out-of-Distribution Generalization. Out-of-distribution (OOD) generalization has been a topic of significant interest in the field of machine learning and computer vision. Researchers have explored various approaches and techniques to improve the robustness and generalization capabilities of models on unseen target domains. These OOD methods use different strategies to update the model, mainly including metalearning [Li et al., 2018a , Zhang et al., 2023] , domain alignment [Ajakan et al., 2014 , Li et al., 2018b] , regularized training [Zhang et al., 2021 , Krueger et al., 2021 , Shi et al., 2021 , Xu and Jaakkola, 2021] , causal learning [Ahuja et al., 2021 , Krueger et al., 2021 , Koyama and Yamaguchi, 2020 ], etc. For example, IRM [Arjovsky et al., 2019] , a representative OOD method, learns invariant representations with a classifier optimal to domain changes as a regularization term. SANDMask [Shahtalebi et al., 2021] The comparison of our method, the ERM method, and the IRM method in terms of OOD accuracy (%).OOD methods have theoretically and experimentally demonstrated their effectiveness in learning domain-invariant features, most of them require the use of environment labels, which does not meet the requirements of real-world applications. Furthermore, we empirically found that the information extracted by the last layer of the feature extractor based on the IRM-based method in different environments cannot be well-aligned. In this paper, taking a step forward, we propose a novel method that aligns the semantic features extracted by the OOD feature extractor to the same simplex ETF under neural collapse. Note that our method can be applied both with and without environment labels.Neural Collapse. The phenomenon known as neural collapse, first identified by [Papyan et al., 2020] , refers to the observation that when the number of training samples across different classes is balanced, both the feature vectors in the final layer and the classifier vectors tend to converge to the simplex ETF upon the completion of training. Recently, several studies [Ji et al., 2021 , Zhu et al., 2021 , Tirer and Bruna, 2022 , Zhu et al., 2023b , Li et al., 2023 , Xie et al., 2023 , Yang et al., 2023 , Beaglehole et al., 2024 , Fisher et al., 2024 , Guo et al., 2024 , Kothapalli et al., 2024 , Súkeník et al., 2024] have utilized this phenomenon to guide the training process in imbalanced data sets. Among them, Yang et al. [2023] involves pre-allocating a fixed number of classes in a simplex ETF for continual learning, guiding the learning of minority classes in subsequent incremental steps, thereby ensuring convergence of intra-class features to specified positions and maximizing and uniformly separating inter-class features. Xie et al. [2023] addresses class imbalance by designing a novel loss function, namely the attraction-rejection balanced loss. Li et al. [2023] apples neural collapse into federated learning scenarios. Under distributed conditions, neural collapse is used to guide the alignment direction of each client, and the personality of each client model is maintained through fine-tuning. The aforementioned methods only consider collapsing the same class onto a single point in a fixed simplex ETF without accounting for the impact of intra-class spurious correlations. We are the first to utilize the concept of neural collapse to address spurious correlations within classes, thereby enhancing out-of-distribution (OOD) performance.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We have discussed the limitations of our work in conclusion part.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
out-of-distribution generalization: 
 In the given dataset D := (X, Y) = {(x i , y i )} N i=1 , typically, during the training process, the data set D is divided into a training set D tr and a test set D te , where D tr is sampled from the training distribution P tr (X, Y), and D te is sampled from the test distribution P te (X, Y). There is a model that can be divided into a feature extractor f (•; ω f ) and a classifier g(•; ω g ). For out-of-distribution (OOD) scenarios, the training distribution is not observable, and the training distribution differs from the test distribution, i.e., P tr (X, Y) ̸ = P te (X, Y). Specifically, according to previous work, we have multiple environments in the training set, that is, E = {e 1 , e 2 , ..., e E }. In different environments, for X in the training set, we can divide it into semantic parts and spurious parts, the correlation between these two is unstable, thus it is prone to generating spurious correlations.Neural collapse refers to a phenomenon observed during the final stages of training on balanced data (post-zero training error) [Papyan et al., 2020] . It reveals a simplex ETF structure formed by the last-layer features and the classifier, which can be defined as:Definition 1 (Simplex Equiangular Tight Frame). A simplex Equiangular Tight Frame (ETF) refers to a matrix that is composed of K vectors v i ∈R d , d ≥ K -1 and satisfies:EQUATIONwhere V = [v 1 , v 2 , ..., v K ] ∈ R d×K , U ∈ R d×K allows a rotation and satisfies U T U = I K , I K represents identity matrix, and 1 K stands for all one matrix.For any v i ∈ V, i ∈ [1, K] and satisfies:EQUATIONwhere δ i,j equals to 1 when i = j, otherwise it is equal to 0, for all v i ∈ V, i = [1, K], there is the same L2 norm, and the inner product between any two vectors is -1 K-1 .Neural Collapse. The phenomenon of neural collapse is guaranteed by the following four findings: variability collapse (NC1), convergence to the simplex equiangular tight frame (NC2), convergence to self-duality (NC3), and simplification to the nearest class center (NC4).NC1: During the final stages of model training, the final layer feature vectors of the same class will collapse to the class mean, i.e., for all k,k W → 0. Σ k W = 1 n k n k i=1 (f k,i -f k ) (f k,i -f k ) T ,where n k represents the number of samples in class k, f k,i denotes the feature obtained from the i-th sample of class k, f k represents the mean feature of class k.The class means of all classes will converge to the vertices of a simplex ETF centered on the global mean, as defined in Definition 1, fk= (f k -f G ) / ∥f k -f G ∥, where f G = K k=1 N k i=1 f k,irepresents the global mean of the feature.The feature prototypes centered on the global mean will align with the corresponding classifier weights, which means that the classifier weights converge to the same simplex ETF, i.e, fk = v k / ∥v k ∥, where v k represents the classifier weight for class k.NC4: Through the above points, it can be ensured that the network classifier converges to the nearest class center, i.e, arg max k ⟨f , v k ⟩ = arg min k ∥f -f k ∥, where f represents the features of the sample.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
methodology: 
 In this section, we delve into a comprehensive discussion of our proposed method, NCFAL, along with its specific implementation. Section 4.1 outlines the framework of the proposed method for out-of-distribution based on neural collapse. We use a fixed ETF classifier to guide the alignment of invariant features across different environments, facilitating better learning of masks to separate invariant and variable features. In Section 4.2, we provide a detailed explanation of how environments are partitioned. In Section 4.3, we elaborate on how a fixed ETF classifier is used to guide mask learning. In Section 4.4, after obtaining invariant features using the trained mask, we input them into the neural network for subsequent training until convergence.Neural collapse reveals the optimal geometric structure of the classifier and feature prototypes after sufficient training (i.e., the simplex ETF). It inspires us to use a simplex ETF as a fixed classifier from the beginning, guiding the alignment of invariant components across different environments. Therefore, we propose a novel algorithm for the OOD generation inspired by neural collapse, NCFAL. Specifically, as described in Section 4.2, in real-world scenarios, environment information is often unknown. Hence, we need to automatically partition environments, and when certain conditions are satisfied, we consider the environment to be sufficiently well-partitioned. As described in Section 4.3, to improve model generalization, after obtaining environment information, we guide the alignment of invariant components across environments using a fixed ETF classifier. Given that class imbalance is likely across different environments, we employ a loss function to mitigate the impact of this imbalance. We alternate between Section 4.2 and Section 4.3, enabling the model to learn improved masks for partitioning invariant components. In Section 4.4, after learning the mask, we obtain semantic features for subsequent training of the predictive model. Specifically, as shown in Figure 2 , our methodological framework is illustrated. Initially, when x is input, it passes through 1 -m for the separation of variable components, allowing for prediction using corresponding models in different environments. Subsequently, an environment selection process assigns the input to a designated environment, yielding a new environment label. When the difference between the new environment label and the previous one is less than a given threshold or a specified number of partitioning iterations is reached, the mask learning process begins. In this phase, x from different environments is input. Since invariant features are extracted, they should collapse onto the same simplex ETF in any environment. This approach guides the alignment of invariant features across variable environments.This chapter focuses on environment partitioning in scenarios where environment labels are unknown. The partitioning process relies on the intermediate module depicted in Figure 2 , which takes the input information and outputs a new set of environments E, where each environment reflects a type of spurious correlation present in the input. We divide the input x into semantic and spurious features, using the invariant mask m to distinguish between them, resulting in semantic features Φ(x) and spurious featuresΨ(x): Φ(x) = m ⊙ x, Ψ(x) = (1 -m) ⊙ x.Subsequently, the variable features Ψ(x) are input into local models, which are associated with the environments, for identification. We designed a two-stage partitioning method to iterative partition environments until convergence.Environment Local Model Learning Stage: Let X tr e be the interaction set in environment e. We aim to use local model to represent the environment e. Intuitively, the primary distinction between different environments lies in their interpretation of spurious correlations. Therefore, we model environments based on these spurious correlations. Specifically, for the interactions in environment e, we predict using a variable representation learning model Γ (e) . In other words, to describe environment e, we learn the predictive model as follows:EQUATIONwhere X tr e represents the training data X corresponding to environment e, ω e represents model parameters corresponding to environment e. Environment Partitioning Stage: In this stage, we have already developed environment models that can be used to assess the spurious correlations in different environments. Consequently, interactions should be separated based on these spurious correlations. To distinguish environments and the interactions within the input, we employ the following formula, which determines that an interaction belongs to the environment with the highest probability of identifying it. Note that predictions are based on the variable component Ψ(x), which represents spurious information. The maximum value of this tensor is selected, indicating that the variable features of the input are more distinguishable in the environment corresponding to the maximum value. Thus, the input is assigned to that environment.EQUATIONwhere for any environment e ∈ E, we select the one with the maximum logit as the corresponding environment for input x.We illustrate this environment selection process with an example in Figure 3 . Considering a logit 0 with orange background, when it is inputted into different environment models, the corresponding predictions are [0.5, 0.3, 0.2] in environment 1 and [0.4, 0.5, 0.1] in environment 2. By extracting the logits at the positions corresponding to the labels(one hot encoding), we obtain vectors [0.5, 0.4].Taking the maximum value, we can determine the corresponding environment as environment 1.This chapter primarily addresses OOD generalization by learning masks to separate invariant representations when environment labels are known. Our approach is guided by the perspective of neural collapse to align invariant representations across environments to a pre-fixed simplex ETF. As shown in Figure 3 , once a suitable environment partitioning mask is learned, we apply it to the input x to obtain the invariant components. Thus, in different environments, for the semantic features Φ(x) in all e i , e j ∈ supp(E), we have: P ei (Y|Φ(x)) = P ej (Y|Φ(x)).In different environments, we use corresponding models to align the invariant features. First, we fix a standard simplex ETF as the alignment direction. Given the potential issue of the number of different classes imbalance across different environments, we use a balanced loss to assign weights proportional to the quantity of each class for the training process:EQUATIONf = v/∥v∥, v = P (h, ω p ), h = f (Φ(x), ω f ),where n e,k represents the number of samples of class k in environment e, v k ∈ V represents the feature vector corresponding to class k on the fixed simplex ETF V, f represents the normalized feature v obtained after projection mapping the original feature h, β represents the learnable temperature, γ represents the sample balancing parameter, ω f represents the parameters of the feature extraction module in the model, ω p represents the parameters of the mapping layer in the model. When learning the mask, we add some random noise to initialize the mask:m i = max{0, min{1, m i + ϵ}}, ϵ ∼ N (0, σ 2 ), m i ∈ m.After completing the above training steps Eq.( 5), we perform a clipping operation on the mask to fix the range of m i :m i = max{0, min{1, m i + ϵ}}, m i ∈ m.In this chapter, we apply the environment partitioning and mask learning in Sections 4.2 and 4.3 respectively to the input data to extract semantic features. These features are then fed into the predictive neural network for training until the network converges:EQUATIONwhere ω * represents the parameters of the final predictive model.  
experiment setup: 
 Datasets. Following the work [Gulrajani and Lopez-Paz, 2020] , we evaluate our method with baselines on benchmark datasets, using four datasets, namely ColoredMNIST, ColoredCOCO, CO-COPlaces and NICO. ColoredMNIST is colorized on the MNIST dataset. The image dimensions were set to [2, 28, 28], the digits [0, 1, 2, 3, 4] are designated as category 0, while the digits [5, 6, 7, 8, 9] are designated as category 1. ColoredCOCO dataset is derived from the COCO dataset, which includes a selection of ten categories. Background color alterations were applied using ten different colors. All images are configured with dimensions of (3, 64, 64) . COCOPlaces employs the same classes and settings as ColoredCOCO, with the distinction that we sample images from Places as spurious information. NICO dataset is a real-world dataset, including 10 subclasses for animals and 9 subclasses for vehicles. In total, our split consists of 4,080 samples of dimension (3, 224, 224) and 2 classes of the classification task. More details can be found in Appendix A.1. For each dataset, we partition it into two subsets, d 1 and d 2 , based on the environment, with the ratio of sample quantities between d 1 and d 2 being 9:1. Subset d 1 from the training environment is used for model training, while d 2 is employed for testing the models within the training environment. In the testing environment, d 1 is utilized to assess the out-of-distribution (OOD) performance of the models, while d 2 is utilized according to DomainBed standards for selecting the best model.Architecture. On ColorMNIST, training is conducted using a 4-layer convolutional neural network. For the ColoredCOCO and COCOPlaces datasets, we adhere to the setup outlined in Ahmed et al. [2020] , Gulrajani and Lopez-Paz [2020] , employing ResNet8 for training. On the NICO dataset, training is performed using ResNet18.Baselines. In order to demonstrate the advantages and effectiveness of our approach, we conduct comparative tests against different OOD learning methods. including (1) IID learning: ERM [Vapnik et al., 1998] (2) OOD learning ( sixteen methods): IRM [Arjovsky et al., 2019] , VREx [Krueger et al., 2021] , ARM [Zhang et al., 2021] , GroupDRO [Sagawa et al., 2020] , MLDG [Li et al., 2018a] , MMD [Li et al., 2018b] , IGA [Koyama and Yamaguchi, 2020] , SANDMask [Shahtalebi et al., 2021] , Fish [Shi et al., 2021] , CDANN [Li et al., 2018c] , TRM [Xu and Jaakkola, 2021] , IB ERM [Ahuja et al., 2021] , IB IRM [Ahuja et al., 2021] , CondCAD [Ruan et al., 2021] , CausIRL CORAL [Chevalley et al., 2022] , MAP [Zhang et al., 2023] .Table 1 presents a comparison between our method, the ERM algorithm, and 16 other OOD algorithms. It is evident that our approach outperforms all methods under the given environmental conditions. Moreover, in situations where the environment is unknown, our method demonstrates adaptive capabilities in partitioning the environment, achieving superior OOD performance. The reason behind this phenomenon, we believe, lies that sometimes manually partitioning the environment may not be optimal for neural networks to identify and discern spurious correlations, however, Utilizing neural networks to partition the environment based on predictions might be more adapt at separating semantic and spurious features. Particularly for COCOPlaces, due to the complexity of backgrounds, manual partitioning may not always achieve appropriate segmentation. Therefore, from the results, when the environment is known, the OOD accuracy is 33.3%; however, using our method for environment partitioning yields a result of 36.7%.Table 1 : Average accuracy (%) of OOD on three toy and one real datasets using different methods.ColoredMNIST ColoredCOCO COCOPlaces NICO ERM [Vapnik et al., 1998] 51.5 ± 0.1 45.4 ± 0.9 20.1 ± 0.7 73.6 ± 1.9 IRM [Arjovsky et al., 2019] 60.3 ± 2.8 49.2 ± 0.3 27.1 ± 0.9 75.8 ± 2.0 VREx [Krueger et al., 2021] 52.9 ± 1.2 48.8 ± 0.7 26.2 ± 0.7 76.9 ± 0.7 GroupDRO [Sagawa et al., 2020] 38.5 ± 1.5 49.1 ± 0.6 26.9 ± 0.6 74.6 ± 2.4 MLDG [Li et al., 2018a] 29.4 ± 0.6 11.9 ± 0.8 14.6 ± 0.5 68.4 ± 2.7 MMD [Li et al., 2018b] 50.6 ± 0.1 50.4 ± 0.8 26.3 ± 1.7 78.2 ± 1.2 IGA [Koyama and Yamaguchi, 2020] 50.5 ± 0.1 11.0 ± 0.6 10.8 ± 0.3 48.1 ± 1.3 SANDMask [Shahtalebi et al., 2021] 58.6 ± 6.5 49.2 ± 1.2 25.9 ± 1.4 72.8 ± 1.5 Fish [Shi et al., 2021] 28.0 ± 1.5 41.7 ± 0.5 19.3 ± 2.1 77.0 ± 1.2 CDANN [Li et al., 2018c] 41.7 ± 3.5 38.4 ± 1.5 19.4 ± 1.0 72.8 ± 1.8 TRM [Xu and Jaakkola, 2021] 44.2 ± 5.0 47.5 ± 0.6 24.8 ± 1.1 73.0 ± 0.9 IB ERM [Ahuja et al., 2021] 50.2 ± 0.2 45.4 ± 1.1 20.2 ± 1.0 77.7 ± 1.9 CausIRL CORAL [Chevalley et al., 2022] 28.7 ± 1.3 51.5 ± 1.1 26.1 ± 1.1 75.7 ± 0.9 CondCAD [Ruan et al., 2021] 49.2 ± 0.5 41.2 ± 0.7 20.8 ± 0.3 73.9 ± 1.4 IB IRM [Ahuja et al., 2021] 53.8 ± 1.8 33.9 ± 0.6 14.8 ± 2.3 70.2 ± 2.2 ARM [Zhang et al., 2021] 28.1 ± 0.0 33.0 ± 0.6 25.1 ± 0.2 76.4 ± 1.6 MAP [Zhang et al., 2023] 52.6 ± 0.5 50.9 ± 1.3 26.9 ± 1.0 76.8 ± 1.4 Ours 66.4 ± 0.2 58.0 ± 0.2 33.3 ± 1.7 85.4 ± 0.6 Ours (w/o env) 66.9 ± 2.4 56.9 ± 1.1 36.7 ± 0.9 85.9 ± 0.2The comparison between other methods employing masking techniques and our method. Table 2 presents the effectiveness of our proposed method for learning masks, we conducted comparisons with the IRM and REx methods, which incorporate regularization to facilitate learning of the mask through the gradient or variance. In contrast, our approach leverages the phenomenon of neural collapse to guide the optimization direction for semantic feature enhancement instead of regularization. Simultaneously, we compared our method with the ERM approach, confirming the superiority of using masks to acquire the invariant parts. The comparison between mask is applied at the pixel level and the feature level. To validate the generality of our method, we conducted comparative experiments at both the pixel level and the feature level in Table 3 . It was observed that sometimes extracting invariant features at the feature level could yield better results. In particular, improved performance was observed on the ColoredCOCO and COCOPlaces datasets. This phenomenon primarily stems from the fact that feature-level information is high-dimensional, capturing relationships that are difficult to discern. The comparison between randomly and using maximum likelihood probability method to split environments. Compared with random environment splits in Figure 4 (a), our method, which involves assigning different models to distinct environments and selecting maximum likelihood probabilities outputted by the models, was found to effectively partition environments.(a) The method for selecting environments &RORUHG01,67 &2&23ODFHV(b) The number of environments to partition The comparison between different number of environment divisions. Compared with different number of environment splits in Figure 4 (b), we can find that when the number of divided environments is closer to the actual number of given environments, better performance is achieved, thereby validating the effectiveness of our environment division method. 
conclusion: 
 We explore a new perspective to understand the inherent limitations of ERM and IRM-based methods, which fail due to the inability to align semantic features across environments, resulting in reduced generalization performance. By leveraging the phenomenon of neural collapse to guide the alignment of semantic features across environments. Compared to other OOD methods, we have significantly improved OOD performance. Moreover, in real-world scenarios where environment labels are unknown, our method addresses this by training local models to different environments, automatically achieving environment partitioning. This method greatly reducing the cost of manual annotation and expanding the applicability of our method. Additionally, we will investigate and improve more effective and efficient environment partitioning techniques in future work. We can also explore applying this method in other domains, such as segmentation or few-shot spurious correlation issues. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We have provided computer resources in appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We are make sure to preserve anonymity.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We have discussed in boarder impact in appendix.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This work focuses on addressing the OOD generation, which is not related to the misuse risk.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: We use public datasets in our experiments.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: A.1 Datasets ColoredMNIST, introduced by IRM for assessing spurious correlations in out-of-distribution (OOD) problems, follows the configuration outlined in [Zhang et al., 2023] . In accordance with this setup, we colorized the MNIST dataset with red and green hues to establish a strong correlation between digits and colors. The image dimensions were set to [2, 28, 28], and the correlation coefficients for two training sets and one test set were specified as (0.9, 0.8, 0.1), indicating the proportion of red and green within category 0. The digits [0, 1, 2, 3, 4] are designated as category 0, while the digits [5, 6, 7, 8, 9] are designated as category 1.ColoredCOCO dataset is derived from the COCO dataset, which includes a selection of ten categories: airplane, bird, boat, bus, dog, horse, motorcycle, train, truck, and zebra. Background color alterations were applied using ten different colors. Their RGB values are [0, 100, 0], [188, 143, 143] , [255, 0, 0], [255, 215, 0], [0, 255, 0], [65, 105, 225] , [0, 225, 225] , [0, 0, 255], [255, 20, 147] and [160, 160, 160] . The number of samples for each training environment is 400 for each class but the testing environment is 200 for each class. All images are configured with dimensions of (3, 64, 64).COCOPlaces employs the same classes and settings as ColoredCOCO, with the distinction that we sample images from Places as spurious information [Liu et al., 2021] , such as b/beach, c/canyon, b/building facade, s/staircase, d/desert/sand, c/crevasse, b/bamboo forest, f/forest/broadleaf, b/ball pit and o/oast house. Moreover, some random places are also used, i.e., k/kasbah, l/lighthouse, p/pagoda, r/rock arch, w/water tower, w/waterfall, z/zen garden.NICO dataset is a real-world dataset including photos of animals and vehicles captured in a wide range of contexts (or backgrounds). There are 10 subclasses for animals and 9 subclasses for vehicles, with each subclass having 9 or 10 different contexts. Following [Lin et al., 2022] , we select a subset of this dataset to simulate the spurious correlation of different contexts and classes (animal or vehicle), which is similar to the setting of ColoredMNIST. More specifically, we make use of both classes that appear in four overlapped contexts: "on snow", "in forest", "on beach" and "on grass" to construct two training environments and one testing environment. In total, our split consists of 4,080 samples of dimension (3, 224, 224) and 2 classes of the classiffcation task. We propose a new perspective for addressing the OOD generation problem by using neural collapse to align semantic features across different environments. Our method can be applied both when environments are given and when they are unknown, as it can autonomously partition environments. This approach enables the model to independently learn semantic features, thereby achieving better generalization in future practical applications.In our experiments, we conduct all methods on a local Linux server that has two physical CPU chips (Intel(R) Xeon(R) CPU E5-2640 v4 @ 2.40GHz) and 32 logical kernels. All methods are implemented using Pytorch framework and all models are trained on GeForce RTX 2080 Ti GPUs.NeurIPS Paper Checklist 


Decision: This paper received 3 reviews from experts in the field. The paper received the following reviews:  1 Accept, 1 Borderline Accept and 1 Borderline Reject.

The main concerns presented by the reviewers included experimentation and clarity. The author rebuttal with additional experiments persuaded the reviewers to increase their ratings. The AC agrees with the points presented by the reviewers and the decision for this paper is to accept. Please include the additional experiments in your paper and address the reviewers’ concerns about clarity in the camera ready version.

Reviews: review: paper_summary: paper_summary: The spurious correlation between image background features and their labels is a significant research problem, and the existing research suffers from the issue of difficult decoupling. In this paper, we propose a new approach to solve the spurious association problem by alternately performing environment segmentation and learning semantic masks from the perspective of neural collapse. Extensive experiments are conducted on four datasets and the results show that the proposed method significantly improves the out-of-distribution performance.
strength_weakness: strength_weakness: W1:  In lines 48-50, it is mentioned that IRM-based methods learn similar representations from different environments, indicating a lack of proper alignment. Could you provide a corresponding experiment to demonstrate this phenomenon?

W2: In Figure 3, the explanation of the middle module that uses logits to judge the environment is unclear. Could you please clarify the structure of the local models, the number of local models used, and the specific meaning of the logit values?

W3: Could you explain the differences between masks based on pixel-level and feature-level approaches? If using feature-level masks, what is the impact of different network-layer features on model performance?

This work addresses an important and interesting question by introducing neural collapse from an invariant perspective, which I believe can provide valuable insights to the community. However, my main concern is that the same mask is used to learn both invariant and variable feature information. What are the advantages of the mask learning mechanism proposed in this paper compared to HRM's [1] mask mechanism?

[1] Heterogeneous Risk Minimization
questions: questions: For more information see Weaknesses.
limitations: limitations: Yes, the authors have adequately described the limitations in their submission.
paper_summary: paper_summary: The paper leverages the neural collapse inspired ETF behavior to simulate different environments in datasets, and uses it for OOD classification.
strength_weakness: strength_weakness: The paper seems generally consistent and well merited. The experiments are a bit lacking, but are convincing.
questions: questions: The following papers seem missing from the neural collapse literature that may be helpful: 

1. https://arxiv.org/abs/2112.15121
2. https://proceedings.neurips.cc/paper_files/paper/2023/hash/b63ad8c24354b0e5bcb7aea16490beab-Abstract-Conference.html
3. https://openreview.net/pdf?id=162TqkUNPO
limitations: limitations: N/A
paper_summary: paper_summary: This paper addresses the problem of spurious correlations caused by environments from where data are collected.
The proposed method applies a mask to input data to separate spurious and semantic features.
The masked input data are fed into a local model specialized to each environment.
Each local model is trained to induce neural collapse for OOD generalization.
strength_weakness: strength_weakness: - W1: Comparison with not only OOD generalization methods but also spurious correlation (sometimes called bias or shortcut) methods is necessary. Methods that can automatically detect and split spurious and semantic features have been developed [a-e].
- W2: Types of spurious features that the proposed method can handle need to be clarified. Can the proposed method handle spurious features in superposition, e.g., objects and textures?
- W3: The rationale behind the proposed method needs to be clarified. For instance, it is unclear why the method adds the noise to the mask when learning it.
- W4: Deeper analyses in the experiments would make the paper more interesting. For example,      
  - Whether the neural collapse is achieved by the proposed method should be confirmed in the experiment.    
  - Visualizing learned masks would produce more valuable insights.
- W5: What is described in the introduction and what is done in the proposed method seems to be different. Although L42 states that `we propose to compute the Frobenius norm (F-norm) of the difference between the feature prototypes and the standard simplex ETF`, the F-norm does not appear in the proposed method.
- W6: Writing and formatting can be improved. There are many inconsistent spellings. For example,    
  - Is "variable features" in L153 the same as spurious features?  
  - The meaning of "interaction" in L189, 192, and so on is unclear. Maybe "training?"  
  - Such inconsistent spellings occur from Section 4.

[a] Tiwari, Rishabh, and Pradeep Shenoy. "Overcoming simplicity bias in deep networks using a feature sieve." ICML2023.    
[b] Bahng, Hyojin, et al. "Learning de-biased representations with biased representations." ICML2020.    
[c] Yang, Wanqian, et al. "Chroma-vae: Mitigating shortcut learning with generative classifiers." NeurIPS2022.    
[d] Liu, Evan Z., et al. "Just train twice: Improving group robustness without training group information." ICML2021.    
[e] Nam, Junhyun, et al. "Learning from failure: De-biasing classifier from biased classifier." NeurIPS2020.
questions: questions: - Q1: Why do the accuracies of IRM (w/ env) in Tables 1 and 2 differ?
- Q2: When environment labels are not available, how are spurious and semantic features learned? Is there a possibility that the two types of features are learned conversely?
- Q3: How do we determine the number of local models when the total number of the environments is unknown?
limitations: limitations: Discussed in Section 6.

Comments: title: Thank you for your rebuttal
A kind request for your feedback about the summary of our clarifications.
Many thanks for raising the score and we also provide further analysis to address the reviewer's concerns.
Please kindly find our supplemental references as follows.
We would like to supplement detailed discussion on neural collapse literature and more convincing experiments.
Thank you for the reviewer's further discussion.
The Additional Experiments according to the reviewer's suggestion
comment: Dear reviewer F1YK,

Since the discussion period will end in a few hours, we will be online waiting for your feedback on our rebuttal, which we believe has fully addressed your concerns.

We would highly appreciate it if you could take into account our response when updating the rating and having discussions with AC and other reviewers.

Authors of # 10491
Thanks for the further experiments.
Most of my concerns about experiments were addressed, but I think writing should be improved.
I would like to retain the current score.
Thank you for your rebuttal!

I want to thank the authors for their tremendous efforts during the rebuttal process. The additional comparative experiments have significantly improved the quality of the paper. The authors conducted relevant experiments to demonstrate the impact of mask positioning, thoroughly analyzed the differences between the proposed method and the HRM approach, and provided extensive experimental evidence to validate the effectiveness of their method, addressing my concerns. Based on these points, I have decided to raise my score to 7.
Many thanks for taking time on reviewing our paper. We would like to provide a summary of the key issues raised by the reviewer. We addressed the reviewer's questions through both **rigorous analysis** and **comprehensive experiments**.

1. **Experiments:** To verify whether $m$ can effectively distinguish between spurious and invariant features, We trained the invariant network using $m \odot x$ and $(1-m)\odot x$, and compared their performance during the final testing phase.

Methods|ColoredMNIST|ColoredCOCO|COCOPlaces| 
|---|---|---|---|
**Ours (m)**|**66.9±2.4**|**56.9±1.1**|**36.7±0.9**
Ours (1-m)|21.8±0.9|16.3±1.8|8.4±0.6

2. **Analysis:** we can address them using two metrics in our paper.

- **The variation in environment partitioning:** 

Num |0|1|2| 3 |4 |5|6|7|8|9|10|
|---|---|---|---|---|---|---|---|---|---|---|---|
ColoredCOCO|3200|2011|1099 | 696 | 442| 509 | 405| 325 | 280|235| 174

- **The loss incurred from using the neural collapse method for alignment in the models in different environments:**

Step |0|1|2| 3 |4 |5|6|7|8|9|10|
|---|---|---|---|---|---|---|---|---|---|---|---|
ColoredCOCO|10.69|5.48|2.25 | 1.69 | 1.55| 1.47 | 1.21| 1.02 | 0.82|0.65| 0.50

We sincerely hope the reviewer will re-evaluate our paper based on the additional experiments and explanations provided above. We look forward to your feedback if you have any further questions!
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.

Methods|ColoredMNIST|ColoredCOCO|COCOPlaces| 
|---|---|---|---|
**Ours (m)**|**66.9±2.4**|**56.9±1.1**|**36.7±0.9**
Ours (1-m)|21.8±0.9|16.3±1.8|8.4±0.6

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

Number of interactions|0|1|2| 3 |4 |5|6|7|8|9|10|11|12|
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
$\Delta N$|3200|2011|1099 | 696 | 442| 509 | 405| 325 | 280|235| 174|75|23

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.

Number of interactions |0|1|2| 3 |4 |5|6|7|8|9|10|
|---|---|---|---|---|---|---|---|---|---|---|---|
$\mathcal{L}\_{mask}$|10.69|5.48|2.25 | 1.69 | 1.55| 1.47 | 1.21| 1.02 | 0.82|0.65| 0.50

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!
[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.
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..**

Methods|ColoredCOCO|COCOPlaces|ColoredMNIST
|-|-|-|-|
Output Layer (suggested)|**63.9 ± 0.5**|43.7 ± 0.7|**58.7 ± 2.8**
Layer 0|62.3 ± 0.6|41.2 ± 0.5|54.2 ± 1.5
Layer 1|63.3 ± 0.4|**44.5 ± 0.3**|53.6 ± 1.7
Layer 2|61.2 ± 0.2|40.2 ± 0.1|50.6 ± 0.3
Layer 3|63.2 ± 0.3|43.2 ± 0.6|51.2 ± 0.9

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

Methods|ColoredMNIST|ColoredCOCO|COCOPlaces|NICO|
|-|-|-|-|-|
Ours|**66.9 ± 2.4**|**56.9 ± 1.1**|**36.5 ± 0.8**|**85.9 ± 0.2**
HRM [6]|61.3 ± 0.6|52.7 ± 0.7|33.4 ± 0.6|79.6 ± 0.9

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

Methods	|ColoredCOCO|	COCOPlaces| NICO
|-|-|-|-|
**Ours**	|**63.9 ± 0.5**	|**43.7 ± 0.7** | **85.4 ± 0.6**
Rebias [7]	|56.0 ± 0.2	|39.2 ± 0.2 | 78.4 ± 1.7
Jtt [8]	|55.3 ± 0.3	|38.0 ± 0.4| 75.5 ± 0.3
RUBi [9]	|53.8 ± 0.5	|32.7 ± 0.7 | 73.4 ± 0.2
LearnedMixin [10]	|52.0 ± 1.4	|30.2 ± 0.5 | 73.7 ± 0.1

We will definitely put the above additional experimental results into our revised manuscript -- thank you so much!
Thank you for your further clarification.
I would like to increase the score.
But I am still curious about how the two types of features are split actually, such as learned masks.
Thank you for the reviewer's further discussion. Below, we provide a detailed analysis of how our method achieves environment partitioning and semantic feature mask optimization through a **two-stage interaction process when environment labels are unknown.** We also explain how the mask is utilized to learn invariant semantic information.

We initialize $m\odot x$ as the semantic or invariant features of $x$. Consequently, $(1-m) \odot x$ represents the spurious or variable features of $x$. Our method iteratively learns and interacts in two stages, gradually achieving environment partitioning and distinguishing between semantic and spurious information.

**Step 1: Environment Partitioning Process:** Since the distinction between different environments arises from variations in spurious features, we use these spurious features to partition environments. Based on the current sample $x$, we obtain spurious feature $\Psi(x) = (1-m) \odot x$, then randomly assign environment labels to the sample. We use the environment partitioning model to make predictions for the samples and reassign environment labels by maximizing the predicted values, as illustrated in Equation 4.

$$e(x) = argmax_{e \in \mathcal{E}} \Gamma^{(e)}(x,\Psi|\omega_e).$$


**Step 2: Mask Optimization and Learning Stage:** For datasets $D^e_{tr}$ under different environments, we apply the invariant feature mask $m$ to obtain the semantic features for each environment. These semantic features $\Phi^e(x)$ are then used for model training across environments. Since the categories are consistent across environments, their semantic features are similar. Therefore, we employ neural collapse methods for feature alignment, aiming to align semantic features from different environments onto a common Simplex Equiangular Tight Frame (ETF), as shown in Equation 5, 

$$\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}.$$

This process results in the updated invariant mask and the newly defined environment partitions after one complete iteration.

After performing multiple rounds of **Step1: environment partitioning** and **Step2: mask optimization**, if the data within the newly defined environments does not significantly differ from that in the environments defined in the previous iteration, we consider the environment partitioning to have converged. At this stage, we use the data from these environments to train the invariant network, aiming to achieve out-of-distribution (OOD) generalization. The final results reflect the effectiveness of our method, demonstrating the successful separation of semantic and spurious features using semantic masks.

We look forward to your feedback if you have any further questions!
Thank you for the rebuttal and additional experiment.

For Q2, my question is whether the mask learns to remove the spurious features from input data.
For instance, in ColoredCOCO, is there any possibility that the mask learns to remove the semantic features, i.e., the objects, and the predictive model learns classification using background colors?
Without prior knowledge such as environment labels, the model and mask cannot distinguish which is the semantic feature.
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.
Methods	|ColoredCOCO|	COCOPlaces|
| -------- | ------------------ | ----------------- |
**Ours**	|**63.9±0.5**	|**43.7±0.7**
Rebias [b]	|56.0±0.2	|39.2±0.2
Jtt [d]	|55.3±0.3	|38.0±0.4
RUBi [f]	|53.8±0.5	|32.7±0.7
LearnedMixin [g]	|52.0±1.4	|30.2±0.5


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