Navigate Beyond Shortcuts:
Debiased Learning through the Lens of Neural Collapse

Abstract

Report issue for preceding element

Recent studies have noted an intriguing phenomenon termed Neural Collapse, that is, when the neural networks establish the right correlation between feature spaces and the training targets, their last-layer features, together with the classifier weights, will collapse into a stable and symmetric structure. In this paper, we extend the investigation of Neural Collapse to the biased datasets with imbalanced attributes. We observe that models will easily fall into the pitfall of shortcut learning and form a biased, non-collapsed feature space at the early period of training, which is hard to reverse and limits the generalization capability. To tackle the root cause of biased classification, we follow the recent inspiration of prime training, and propose an avoid-shortcut learning framework without additional training complexity. With well-designed shortcut primes based on Neural Collapse structure, the models are encouraged to skip the pursuit of simple shortcuts and naturally capture the intrinsic correlations. Experimental results demonstrate that our method induces better convergence properties during training, and achieves state-of-the-art generalization performance on both synthetic and real-world biased datasets.

Report issue for preceding element

1 Introduction

Report issue for preceding element

Figure 1: Illustration of (a) Neural Collapse phenomenon on balanced datasets, where the simplex ETF structure maximizes the class-wise angles, and (b) Biased classification on datasets with imbalanced attributes, where the model takes the shortcut of attributes to make predictions and fails to collapse into the simplex ETF. The color of points represents different class labels and the shape of points represents different attributes.

Report issue for preceding element

When the input-output correlation learned by a neural network is consistent with its training target, the last-layer features and classifier weights will attract and reinforce each other, forming a stable, symmetric and robust structure. Just as the Neural Collapse phenomenon discovered by Papyan et al. [24], at the terminal phase of training on balanced datasets, a model will witness its last-layer features of the same class converge towards the class centers, and the classifier weights align to these class centers correspondingly. The convergence will ultimately lead to the collapse of feature space into a simplex equiangular tight frame (ETF) structure, as illustrated in Fig. 1(a). The elegant structure has demonstrated its efficacy in enhancing the generalization, robustness, and interpretability of the trained models [24, 5]. Therefore, a wave of empirical and theoretical analysis of Neural Collapse has been proposed [2, 8, 6, 25, 26, 38, 40], and a series of studies have adopted the simplex ETF as the optimal geometric structure of the classifier, to guide the maximized class-wise separation in class-imbalanced training [20, 43, 37, 35, 36].

Report issue for preceding element

However, in practical visual recognition tasks, besides the challenge of inter-class imbalance, we also encounter intra-class imbalance, where the majority of samples are dominated by the bias attributes (e.g., some misleading contents such as background, color, texture, etc.). For example, the widely used LFW dataset [12] for facial recognition has been demonstrated severely imbalanced in gender, age and ethnicity [3]. A biased dataset often contains a majority of bias-aligned samples and a minority of bias-conflicting ones. The prevalent bias-aligned samples exhibit a strong correlation between the ground-truth labels and bias attributes, while the scarce bias-conflicting samples have no such correlation. Once a model relies on the simple but spurious shortcut of bias attributes for prediction, it will ignore the intrinsic relations and struggle to generalize on out-of-distribution test samples. The potential impact of biased classification may range from political and economic disparities to social inequalities within AI systems, as emphasized in EDRi’s latest report [1].

Report issue for preceding element

Therefore, the fundamental solution to biased classification lies in deferring, or ideally, preventing the learning of shortcut correlations. However, previous debiased learning methods rely heavily on additional training expenses. For example, a bias-amplified auxiliary model is often adopted to identify and up-weight the bias-conflicting samples [23, 19, 16], or employed to guide the input-level and feature-level augmentations [13, 18, 21]. Some disentangle-based debiasing methods, from the perspective of causal intervention [32, 42] or Information Bottleneck theory [30], also require large amounts of contrastive samples or pre-training process to disentangle the biased features, significantly increasing the burden of debiased learning.

Report issue for preceding element

In this paper, we extend the investigation of Neural Collapse to the biased visual datasets with imbalanced attributes. Through the lens of Neural Collapse, we observe that models prioritize the period of shortcut learning, and quickly form the biased feature space based on misleading attributes at the early stage of training. After the bias-aligned samples reach zero training error, the intrinsic correlation within bias-conflicting samples will then be discovered. However, due to i) the scarcity of bias-conflicting samples and ii) the stability of the established feature space, the learned shortcut correlation is challenging to reverse and eliminate. The mismatch between bias feature space and the training target induces inferior generalizability, and hinders the convergence of Neural Collapse, as shown in Fig. 1(b).

Report issue for preceding element

To achieve efficient model debiasing, we follow the inspiration of prime training, and encourage the model to skip the active learning of shortcut correlations. The primes are often provided as additional supervisory signals to redirect the model’s reliance on shortcuts, which helps improve generalization in image classification and CARLA autonomous driving [33]. To rectify models’ attention on the intrinsic correlations, we define the primes with a training-free simplex ETF structure, which approximates the “optimal" shortcut features and guides the model to pursue unbiased classification from the beginning of training. Our method is free of auxiliary models or additional optimization of prime features. Experimental results also substantiate its state-of-the-art debiasing performance on both synthetic and real-world biased datasets.

Report issue for preceding element

Our contributions are summarized as follows:

Report issue for preceding element

• For the first time, we investigate the Neural Collapse phenomenon on biased datasets with imbalanced attributes. Through the empirical results of feature convergence, we analyze the shortcut learning stage of training, as well as the fundamental issues of biased classification.

  Report issue for preceding element

• We propose an efficient avoid-shortcut training paradigm, which introduces the simplex ETF structure as prime features, to rectify models’ attention on the intrinsic correlations.

  Report issue for preceding element

• We demonstrate the state-of-the-art debiasing performance of our method on 2 synthetic and 3 real-world biased datasets, as well as the better convergence properties of debiased models.

  Report issue for preceding element

2 Related Works

Report issue for preceding element

Debiased Learning. Extensive efforts have been dedicated to model debiasing, but they are significantly limited by additional training costs. Recent advances can be divided into three categories: reweight-based, augmentation-based, and disentangle-based. Based on the easy-to-learn heuristic of biased features, reweight-based approaches require pre-trained bias-amplified models to identify and emphasize the bias-conflicting training samples [23, 19, 16]. Augmentation-based approaches, with the guidance of explicit bias annotations, conduct image-level and feature-level augmentations to enhance the diversity of training datasets [13, 18, 21]. Other disentangle-based approaches attempt to remove the bias-related part of features, from the perspective of Information Bottleneck theory [30] or causal intervention [32, 42], but at the cost of substantial contrastive samples. Additionally, model debiasing is also well studied in graph neural networks [4, 41], language models [7, 22] and multi-modal tasks [34, 11].

Report issue for preceding element

Neural Collapse. Discovered by Papyan et al. [24], the Neural Collapse phenomenon reveals the convergence of the last-layer feature space to an elegant geometry. At the terminal phase of training on balanced datasets, the feature centers and classifier weights will collapse together into the structure of a simplex ETF, which is illustrated in Section 3.1. Recent works have dug deeper into the phenomenon and provided theoretical supports under different constraints or regularizations [2, 40, 8], as well as empirical studies of intermediate features and transfer learning [26, 38, 25, 6]. Considering the class-imbalanced datasets, Fang et al. [5] point out the Minority Collapse phenomenon, where features of long-tailed classes will merge together and be hard to classify. As a remedy, they fix the classifier as an ETF structure during training, which guarantees the optimal geometric property in imbalanced learning [36], semantic segmentation [43], and federated learning [37]. To take a step further, our work fills the gap of Neural Collapse analysis on biased datasets with shortcut correlations.

Report issue for preceding element

Avoid-shortcut Learning. The recent inspiration of avoid-shortcut learning aims to postpone, or even prevent the learning of shortcut relations in model training. With well-crafted contrastive samples [27, 29, 28] or artificial shortcut signals [33, 42], avoid-shortcut learning has demonstrated its efficacy in image classification, autonomous driving and question answering models. One of the representative methods is named prime training, which provides richer supervisory signals of key input features (i.e., primes) to guide the establishment of correct correlations, therefore improving generalization on OOD samples [33]. In this work, we leverage the approximated “optimal" shortcuts as primes to encourage the models to bypass shortcut learning.

Report issue for preceding element

3 Preliminaries

Report issue for preceding element

3.1 Neural Collapse Phenomenon

Report issue for preceding element

Consider a biased dataset $\mathcal{D}$ with $K$ classes of training samples, we denote $\mathbf{x}_{k,i}$ as the $i$-th sample of the $k$-th class and $\mathbf{z}_{k,i}\in\mathbb{R}^{d}$ as its corresponding last-layer feature. A linear classifier with weights $\mathbf{W}=[\mathbf{w}_{1},...,\mathbf{w}_{K}]\in\mathbb{R}^{d\times K}$ is trained upon the last-layer features to make predictions.

Report issue for preceding element

The Neural Collapse (NC) phenomenon discovered that, when neural networks are trained on balanced datasets, the correctly learned correlations will naturally lead to the convergence of feature spaces. Given enough training steps after the zero classification error, the last-layer features and classifier weights will collapse to the vertices of a simplex equiangular tight frame (ETF), which is defined as below.

Report issue for preceding element

Definition 1 (Simplex Equiangular Tight Frame) A collection of vectors $\mathbf{m}_{k}\in\mathbb{R}^{d}$, $k=1,2,...,K,d\geq K-1$ is said to be a $k$-simplex equiangular tight frame if:

Report issue for preceding element

$$\mathbf{M}=\sqrt{\frac{K}{K-1}}\mathbf{P}(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\mathrm{T}})$$ (1)

where $\mathbf{M}=[\mathbf{m}_{1},...,\mathbf{m}_{K}]\in\mathbb{R}^{d\times K}$, and $\mathbf{P}\in\mathbb{R}^{d\times K}$ is an orthogonal matrix which satisfies $\mathbf{P}^{\mathrm{T}}\mathbf{P}=\mathbf{I}_{K}$, with $\mathbf{I}_{K}$ denotes the identity matrix and $\mathbf{1}_{K}$ denotes the all-ones vector. Within the ETF structure, all vectors have the maximal pair-wise angle of $-\frac{1}{K-1}$, namely the maximal equiangular separation.

Report issue for preceding element

Besides the convergence to simplex ETF structure, the Neural Collapse phenomenon could be concluded as the following properties during the terminal phase of training:

Report issue for preceding element

NC1: Variability collapse. The last-layer features $\mathbf{z}_{k,i}$ of the same class $k$ will collapse to their class means $\mathbf{\overline{z}}_{k}={\rm Avg}_{i}\{\mathbf{z}_{k,i}\}$, and the within-class variation of the last-layer features will approach 0.

Report issue for preceding element

NC2: Convergence to simplex ETF. The normalized class means will collapse to the vertices of a simplex ETF. We denote the global mean of all last-layer features as $\mathbf{z}_{G}={\rm Avg}_{i,k}\{\mathbf{z}_{k,i}\},k\in[1,...,K]$ and the normalized class means as $\mathbf{{\tilde{z}}}_{k}=(\mathbf{z}_{k}-\mathbf{z}_{G})/\lVert\mathbf{z}_{k}-\mathbf{z}_{G}\rVert$, which satisfies Eq.1.

Report issue for preceding element

NC3: Self duality. The classifier weights $\mathbf{w}_{k}$ will align with the corresponding normalized class means $\mathbf{{\tilde{z}}}_{k}$, which satisfies $\mathbf{{\tilde{z}}}_{k}=\mathbf{w}_{k}/||\mathbf{w}_{k}||$.

Report issue for preceding element

NC4: Simplification to nearest class center. After convergence, the model’s prediction will collapse to simply choosing the nearest class mean to the input feature (in standard Euclidean distance). The prediction of $\mathbf{z}$ could be denoted as $\arg\max_{k}\langle\mathbf{z},\mathbf{w}_{k}\rangle=\arg\min_{k}||\mathbf{z}-\mathbf{\overline{z}}_{k}||$.

Report issue for preceding element

Figure 2: Comparison of (a) testset accuracy and (b-d) Neural Collapse metrics on unbiased (CIFAR-10) and synthetic biased (Corrupted CIFAR-10 with the bias ratio of 5.0%) datasets. All vanilla models are trained with standard cross-entropy loss for 500 epochs. The postfix -Aligned and -Conflicting indicate the results of bias-aligned and bias-conflicting samples respectively. The NC1 metric evaluates the convergence of same-class features, NC2 evaluates the difference between the feature space and a simplex ETF, and NC3 measures the duality between feature centers and classifier weights. The vertical dashed line at the epoch of 60 divides two stages of training.

Report issue for preceding element

3.2 Neural Collapse Observation on Biased Dataset

Report issue for preceding element

Besides the findings on balanced datasets, some studies have explored Neural Collapse under the class-imbalanced situation [36, 5]. Taking a step further, we investigate the phenomenon on biased datasets with imbalanced attributes, to advance the understanding of biased classification. To examine the convergence of last-layer features and classifier weights, we compare the metrics of Neural Collapse on both unbiased and synthetic biased datasets. As shown in Fig. 2, we report the result of NC1-NC3, which corresponds to the first three convergence properties in Section 3.1 and respectively evaluates the convergence of same-class features, the structure of feature space and self-duality. The details of NC metrics are concluded in Tab. 1.

Report issue for preceding element

Metrics Computational details NC1 ${\rm\mathcal{NC}_{1}}=\frac{1}{K}{\rm Tr}(\Sigma_{W}\Sigma_{B}^{\dagger})$, where ${\rm Tr}$ is the trace of matrix and $\Sigma_{B}^{\dagger}$ denotes the pseudo-inverse of $\Sigma_{B}$ $\Sigma_{B}=\frac{1}{K}\sum\limits_{k\in[K]}(\mathbf{\overline{z}}_{k}-\mathbf{z}_{G})(\mathbf{\overline{z}}_{k}-\mathbf{z}_{G})^{\mathrm{T}}$, $\Sigma_{W}=\frac{1}{K}\sum\limits_{k\in[K]}\frac{1}{n_{k}}\sum\limits_{i=1}^{n_{k}}(\mathbf{z}_{k,i}-\mathbf{\overline{z}}_{k})(\mathbf{z}_{k,i}-\mathbf{\overline{z}}_{k})^{\mathrm{T}}$ NC2 ${\rm\mathcal{NC}_{2}}=\Big{\lVert}\frac{\mathbf{W}\mathbf{W}^{\mathrm{T}}}{||\mathbf{W}\mathbf{W}^{\mathrm{T}}||_{\mathcal{F}}}-\frac{1}{\sqrt{K-1}}(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\mathrm{T}})\Big{\rVert}_{\mathcal{F}}$ NC3 ${\rm\mathcal{NC}_{3}}=\Big{\lVert}\frac{\mathbf{W}\mathbf{\overline{Z}}}{||\mathbf{W}\mathbf{\overline{Z}}||_{\mathcal{F}}}-\frac{1}{\sqrt{K-1}}(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\mathrm{T}})\Big{\rVert}_{\mathcal{F}}$, where $\mathbf{\overline{Z}}=[\mathbf{\overline{z}}_{1}-\mathbf{z}_{G},...,\mathbf{\overline{z}}_{K}-\mathbf{z}_{G}]$

Report issue for preceding element

Table 1: The metrics of evaluating the Neural Collapse phenomenon, which are generally adopted in previous studies [46, 38, 24]. $\|\cdot\|_{\mathcal{F}}$ denotes the Frobenius norm, $\mathbf{I}_{K}$ is the identity matrix and $\mathbf{1}_{K}$ is the all-ones vector.

Report issue for preceding element

When trained on unbiased datasets (black lines in Fig. 2), the model displays the expected convergence properties, with metrics NC1-NC3 all converge to zero. We owe the elegant collapse phenomenon to the right correlation between the feature space and training objective, which is also supported by the analysis of benign global landscapes [46, 45].

Report issue for preceding element

However, when trained on biased datasets, the training process exhibits two stages: first the shortcut learning period and then the intrinsic learning period, as divided by the vertical dashed line. During the shortcut learning period, the accuracy of bias-aligned samples increases quickly, and the NC1-NC3 metrics show a rapid decline (green lines with $\blacktriangle$ in Fig. 2). It indicates that when simple shortcuts exist in the training distribution, the model will quickly establish its feature space based on the bias attributes, and exhibit a converging trend towards the simplex ETF structure.

Report issue for preceding element

After the bias-aligned samples approach zero error, the model turns to the period of intrinsic learning, which focuses on the intrinsic correlations within bias-conflicting samples to further reduce the empirical loss. However, although their final loss reduces to zero, the bias-conflicting samples still display low accuracy and poor convergence results (green lines with $\blacktriangledown$). It implies that the intrinsic learning period merely induces the over-fitting of bias-conflicting samples and does not benefit in generalization. We attribute the failure of collapse to the early establishment of shortcut correlations. Once the biased feature space is established based on misleading attributes, rectifying it becomes challenging, particularly with scarce bias-conflicting samples. In the subsequent training steps, the misled features of bias-conflicting samples will hinder the convergence of same-class features, thereby halting the converging trend towards the simplex ETF structure and leading to a non-collapsed, sub-optimal feature space.

Report issue for preceding element

To break the curse of shortcut learning, we turn the tricky shortcut into a training prime, which effectively guides the models to focus on intrinsic correlations and form a naturally collapsed feature space (blue lines in Fig. 2). Our method is presented in the following sections.

Report issue for preceding element

4 Methodology

Report issue for preceding element

Figure 3: The illustration of our method. We take the class climbing from BAR [23] as an example, which contains samples of human climbing but with the bias attribute of different backgrounds (as indicated with the color of image frames). The framework contains: 1) Prime Construction: Before training, a randomly initialized ETF structure is constructed as the shortcut primes, and 2) During Prime Training, the prime features $\mathbf{m}_{b}$ are retrieved based on the bias attribute $b$ of the input samples, to guide the optimization of learnable features $\mathbf{z}$ towards the intrinsic correlations. The classifier $F_{\theta}$ will take both the learnable features and fixed prime features to make predictions. 3) In Unbiased Classification, the prime features are assigned as null vectors to evaluate the debiased model on test distributions.

Report issue for preceding element

4.1 Motivation

Report issue for preceding element

Following the previous analysis, we highlight the importance of redirecting the model’s emphasis from simple shortcuts to intrinsic relations. Since the models can be easily misled by shortcuts in the training distribution, is it feasible to supply a “perfectly learned" shortcut feature, to deceive the models into skipping the active learning of shortcuts, and directly focusing on the intrinsic correlations?

Report issue for preceding element

We observe in Fig. 2 that the NC1-NC3 metrics show a rapid decrease during shortcut learning, but remain stable in the subsequent training epochs. However, if the training distribution does follow the shortcut correlation (with no obstacle from bias-conflicting samples), the convergence will end up with the optimal structure of simplex ETF, just as the results on unbiased datasets. This inspires us to approximate the “perfectly learned" shortcut features with a simplex ETF structure, which requires no additional training and represents the optimal geometry of feature space.

Report issue for preceding element

Therefore, following the outstanding performance of prime training in OOD generalization [33], we introduce the approximated “perfect" shortcuts as the primes for debiased learning. The provided shortcut primes are constructed with a training-free simplex ETF structure, which encourages the models to directly capture the intrinsic correlations, therefore exhibit superior generalizability and convergence properties in our experiments.

Report issue for preceding element

4.2 Avoid-shortcut Learning with Neural Collapse

Report issue for preceding element

Building upon our motivation of avoid-shortcut learning, the illustration of the proposed ETF-Debias is shown in Fig. 3. The debiased learning framework can be divided into three stages: prime construction, prime training, and unbiased classification. Firstly, a prime ETF will be constructed to approximate the “perfect" shortcut features. Then during the prime training, the model will be guided to directly capture the intrinsic correlations with the prime training and the prime reinforcement regularization. In evaluation, we rely on the intrinsic correlations to perform unbiased classification. The details are as follows.

Report issue for preceding element

Prime construction. When constructing the prime ETF, we first randomly initialize a simplex ETF as $\mathbf{M}\in\mathbb{R}^{d\times B}$, which satisfies the definition in Eq. 1. The dimension $d$ is the same as the learnable features, and the number of vectors in $\mathbf{M}$ is determined by the categories of bias attributes $\mathbf{b}_{i}\in\{1,...,B\}$, which are pre-defined in the training distribution. After initialization, the vertices of prime ETF $[\mathbf{m}_{1},...,\mathbf{m}_{B}]$ are considered as the approximation of the “perfect" shortcut features for each attribute, which serve as the prime features for avoid-shortcut training. During training, the prime features will be retrieved based on the bias attribute $b$ of each input sample.

Report issue for preceding element

Prime training. During the prime training, we take the end-to-end model architecture with a backbone $E_{\phi}$ and a classifier $F_{\theta}$. For the $i$-th input $\mathbf{x}_{i,b}$ with the bias attribute of $b$, we first extract its learnable feature $\mathbf{z}_{i,b}=E_{\phi}(\mathbf{x}_{i,b})$ with the backbone model, and retrieve its prime feature $\mathbf{m}_{b}$ based on the bias attribute $b$. The classifier $F_{\theta}$ will take both the learnable feature $\mathbf{z}_{i,b}$ and the prime feature $\mathbf{m}_{b}$ to make softmaxed predictions $\hat{\mathbf{y}}=F_{\theta}(\mathbf{z}_{i,b},\mathbf{m}_{b})$. The standard classification objective is defined as:

Report issue for preceding element

$$\min_{\phi,\theta}\mathcal{L}_{\rm CE}(\mathbf{x},\mathbf{y})=\sum_{i=1}^{N}\mathcal{L}(F_{\theta}(\mathbf{z}_{i,b},\mathbf{m}_{b}),\mathbf{y}_{i,b})$$ (2)

where $\mathbf{y}$ is the ground-truth label. In our implementation, we use the standard cross-entropy loss as $\mathcal{L}$, and concatenate the prime features after the learnable features to perform predictions.

Report issue for preceding element

In essence, we provide a pre-defined prime feature for each training sample based on its bias attribute. The prime features, with a strong correlation with the bias attributes, can be viewed as the optimal solution to shortcut learning. By leveraging the already “perfect" representation of shortcut correlations, the model will be forced to explore the intrinsic correlations within the training distribution. The prime-guided mechanism targets at the fundamental issue of biased classification, without inducing extra training costs.

Report issue for preceding element

Prime reinforcement regularization. Given the prime features, the model is encouraged to grasp the intrinsic correlation of the training distributions. However, we raise another potential risk that, despite the provided “perfectly learned" shortcut features, the model may still pursue the easy-to-follow shortcuts, leading to the redundancy between the learnable feature $\mathbf{z}_{i,b}$ and the fixed $\mathbf{m}_{b}$. We point out that the model may not establish a strong correlation between the prime features and the bias attributes, and continues to optimize the learnable features for the missing connections.

Report issue for preceding element

Therefore, we introduce a prime reinforcement regularization mechanism to enhance the model’s dependency on prime features. We encourage the model to classify the bias attributes with only the prime features, and the regularization loss is defined as:

Report issue for preceding element

where $\mathbf{m}_{\rm null}$ is implemented as all-zero vectors with the same dimension as $\mathbf{m}_{b}$, and $\mathcal{L}$ is the standard cross-entropy loss. In the ablation studies in Section 5.3, we observe an improved generalization capability across test distributions, as the result of the strengthened reliance on prime features. Regarding the entire framework, we define the overall training objective as:

Report issue for preceding element

$$\min_{\phi,\theta}\mathcal{L}_{\rm CE}(\mathbf{x},\mathbf{y})+\alpha\mathcal{L}_{\rm RE}(\mathbf{x},\mathbf{b})$$ (4)

where $\alpha$ is the hyper-parameter to adjust the regularization.

Report issue for preceding element

Unbiased classification. In evaluation, we rely on the intrinsic correlations to perform unbiased classification. Given a test sample $\mathbf{x}_{i,b}$, we extract its learnable feature $\mathbf{z}_{i,b}$ and set its prime feature as $\mathbf{m}_{\rm null}$ to obtain the final output $\hat{\mathbf{y}}=F_{\theta}(\mathbf{z}_{i,b},\mathbf{m}_{\rm null})$.

Report issue for preceding element

4.3 Theoretical Justification

Report issue for preceding element

Based on the analysis of Neural Collapse from the perspective of gradients [43, 36], we provide a brief theoretical justification for our method.

Report issue for preceding element

With the priming mechanism, we denote the $i$-th feature of the $k$-th class as $\widetilde{\mathbf{z}}_{k,i}=[\mathbf{z}_{k,i},\mathbf{m}_{i,b}]\in\mathbb{R}^{2\times d}$, which represents the concatenation of learnable feature $\mathbf{z}_{k,i}$ and prime feature $\mathbf{m}_{i,b}$ based on its bias attribute $b$. To keep the same form, we also denote the classifier weights as $\widetilde{\mathbf{w}}_{k}=[\mathbf{w}_{k},\mathbf{a}_{k}]\in\mathbb{R}^{2\times d}$, where $\mathbf{w}_{k}$ represents the weight for intrinsic correlations and $\mathbf{a}_{k}$ represents the one for shortcut correlations. We observe that, due to the fixed prime features during training, $\mathbf{a}_{k}$ will quickly collapse to the bias-correlated prime features of class $k$, and can be viewed as constant after just a few steps of training. With the definition, the cross-entropy (CE) loss can be written as:

Report issue for preceding element

We follow the analysis of previous works and compute the gradients of $\mathcal{L}_{\rm CE}$ w.r.t both classifier weights and features.

Report issue for preceding element

Gradient w.r.t classifier weights. We first compute the gradient of $\mathcal{L}_{\rm CE}$ w.r.t classifier weights $\mathbf{\widetilde{W}}=[\widetilde{\mathbf{w}}_{1}...,\widetilde{\mathbf{w}}_{K}]$:

Report issue for preceding element

$$\frac{\partial\mathcal{L}_{\rm CE}}{\partial\widetilde{\mathbf{w}}_{k}} =\sum_{i=1}^{n_{k}}-(1-p_{k}(\widetilde{\mathbf{z}}_{k,i}))\widetilde{\mathbf{z}}_{k,i}+\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}p_{k}(\widetilde{\mathbf{z}}_{k^{\prime},j})\widetilde{\mathbf{z}}_{k^{\prime},j}$$

$$\leq\underbrace{\sum_{i=1}^{n_{k}}-(1-p_{k}^{(b)}(\mathbf{m}_{i,b})-p_{k}^{(l)}(\mathbf{z}_{k,i}))\widetilde{\mathbf{z}}_{k,i}}_{\rm pulling\ part}$$

$$+\underbrace{\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}(p_{k}^{(b)}(\mathbf{m}_{j,b^{\prime}})+p_{k}^{(l)}(\mathbf{z}_{k^{\prime},j}))\widetilde{\mathbf{z}}_{k^{\prime},j}}_{\rm forcing\ part}$$ (6)

where $p_{k}^{(l)}$ and $p_{k}^{(b)}$ are the predicted probabilities for class labels and bias attributes, calculated with softmax:

Report issue for preceding element

$$p_{k}^{(l)}(\mathbf{z}_{k,i})=\frac{\exp(\mathbf{z}_{k,i}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime}=1}^{K}\exp([\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k^{\prime}},\mathbf{a}_{k^{\prime}}])}$$ (7)

$$p_{k}^{(b)}(\mathbf{m}_{i,b})=\frac{\exp(\mathbf{m}_{i,b}^{\mathrm{T}}\mathbf{a}_{k})}{\sum_{k^{\prime}=1}^{K}\exp([\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k^{\prime}},\mathbf{a}_{k^{\prime}}])}$$ (8)

In Eq. 4.3, the gradient w.r.t classifier weights are divided into two parts. The pulling part is composed of features from the same class that pulls $\mathbf{w}_{k}$ towards the direction of the $k$-th feature cluster, while the forcing part contains the features of other classes and pushes $\mathbf{w}_{k}$ away from their clusters. The weight factor of each feature $\widetilde{\mathbf{z}}_{k,i}$ represents its influence on the optimization of $\widetilde{\mathbf{w}}_{k}$, which implicitly plays the role of re-weighting in our method.

Report issue for preceding element

We assume that class $k$ is strongly correlated with bias attribute $b$. As the weight $\mathbf{a}_{k}$ is observed to collapse quickly to the bias-correlated prime feature $\mathbf{m}_{b}$, the probability $p_{k}^{(b)}\propto\exp(\mathbf{m}_{b}^{\mathrm{T}}\mathbf{a}_{k})$ of bias-aligned samples (with prime features $\mathbf{m}_{b}$) are much greater than that of bias-conflicting samples (with prime features $\mathbf{m}_{b^{\prime}}$). Thus, with the weight factors in Eq. 4.3, the pulling and forcing effects of bias-aligned samples will be relatively down-weighted, and the impact of bias-conflicting samples will be up-weighted. The re-weighting mechanism of gradient mitigates the tendency of pulling $\widetilde{\mathbf{w}}_{k}$ towards the center of bias-aligned samples, which alleviates the misdirection of bias attributes.

Report issue for preceding element

Gradient w.r.t features. Similarly, we compute the gradient of $\mathcal{L}_{\rm CE}$ w.r.t the feature $\widetilde{\mathbf{z}}_{k,i}$:

Report issue for preceding element

$$\frac{\partial\mathcal{L}_{\rm CE}}{\partial\widetilde{\mathbf{z}}_{k,i}} =-(1-p_{k}(\widetilde{\mathbf{z}}_{k,i}))\mathbf{w}_{k}+\sum_{k^{\prime}\neq k}^{K}p_{k^{\prime}}(\widetilde{\mathbf{z}}_{k,i})\mathbf{w}_{k^{\prime}}$$

$$\leq\underbrace{-(1-p_{k}^{(b)}(\mathbf{m}_{i,b})-p_{k}^{(l)}(\mathbf{z}_{k,i}))\mathbf{w}_{k}}_{\rm pulling\ part}$$

$$+\underbrace{\sum_{k^{\prime}\neq k}^{K}(p_{k^{\prime}}^{(b)}(\mathbf{m}_{i,b})+p_{k^{\prime}}^{(l)}(\mathbf{z}_{k,i}))\mathbf{w}_{k^{\prime}}}_{\rm forcing\ part}$$ (9)

In the gradient w.r.t features, the pulling part directs the feature $\widetilde{\mathbf{z}}_{k,i}$ towards the weight of its class $\mathbf{w}_{c}$, and the forcing part repels it from wrong classes. Regarding the weight factors, the probability of bias attribute $p_{k}^{(b)}$ also re-weights the influence of classifier weights. Bias-aligned samples, with high $p_{k}^{(b)}$ probability, will have smaller pulling effects towards the classifier weight $\mathbf{w}_{k}$, which avoids the dominance of bias-aligned features around the weight centers and hinders the tendency of shortcut learning. In comparison, the bias-conflicting samples are granted stronger pulling and pushing effects, which strengthens their convergence toward the right class. The detailed theoretical justification of our method, along with the comparison with vanilla training, are available in Appendix A.

Report issue for preceding element

5 Experiments

Report issue for preceding element

Table 2: Comparison of debiasing performance on synthetic datasets. We report the accuracy on the unbiased test sets of Colored MNIST and Corrupted CIFAR-10. Best performances are marked in bold, and the number in brackets indicates the improvement compared to the best result in baselines. $(^{\ast})$ and $(^{\diamond})$ denote methods with/without bias supervision respectively.

Dataset Ratio(%) Vanilla LfF^⋄[23] LfF+BE^⋄[19] EnD^∗[30] SD^∗[42] DisEnt^∗[18] Selecmix^⋄[13] ETF-Debias Colored MNIST 0.5 32.22$\pm$0.13 57.78$\pm$0.81 69.69$\pm$1.99 35.93$\pm$0.40 56.96$\pm$0.37 68.83$\pm$1.62 70.53$\pm$0.46 71.63$\pm$0.28 (+1.10) 1.0 48.45$\pm$0.06 72.29$\pm$1.69 80.90$\pm$1.40 49.32$\pm$0.58 72.46$\pm$0.18 79.49$\pm$1.44 83.34$\pm$0.37 81.97$\pm$0.26 (-1.37) 2.0 58.90$\pm$0.12 79.51$\pm$1.82 84.90$\pm$1.14 65.58$\pm$0.46 79.37$\pm$0.46 84.56$\pm$1.19 85.90$\pm$0.23 86.00$\pm$0.03 (+0.10) 5.0 74.19$\pm$0.04 83.96$\pm$1.44 90.28$\pm$0.18 80.70$\pm$0.17 88.89$\pm$0.21 88.83$\pm$0.15 91.27$\pm$0.31 91.36$\pm$0.21 (+0.09) Corrupted CIFAR-10 0.5 17.06$\pm$0.12 31.00$\pm$2.67 23.68$\pm$0.50 14.30$\pm$0.10 36.66$\pm$0.74 30.12$\pm$1.60 33.30$\pm$0.26 40.06$\pm$0.03 (+3.40) 1.0 21.48$\pm$0.55 34.33$\pm$1.76 30.72$\pm$0.12 20.17$\pm$0.19 45.66$\pm$1.05 35.28$\pm$1.39 38.72$\pm$0.27 47.52$\pm$0.26 (+1.86) 2.0 27.15$\pm$0.46 39.68$\pm$1.15 42.22$\pm$0.60 30.10$\pm$0.54 50.11$\pm$0.69 40.34$\pm$1.41 47.09$\pm$0.17 54.64$\pm$0.42 (+4.53) 5.0 39.46$\pm$0.58 53.04$\pm$0.76 57.93$\pm$0.58 45.85$\pm$0.21 62.43$\pm$0.57 49.99$\pm$0.84 54.69$\pm$0.29 65.34$\pm$0.60 (+2.91)

Report issue for preceding element

Report issue for preceding element

Table 3: Comparison of debiasing performance on real-world datasets. We report the accuracy on the unbiased test sets of Biased FFHQ, Dogs & Cats, and BAR. The class-wise accuracy on BAR is reported in Appendix D. Best performances are marked in bold, and the number in brackets indicates the improvement compared to the best result in baselines. $(^{\ast})$ and $(^{\diamond})$ denote methods with/without bias supervision respectively.

Dataset Ratio(%) Vanilla LfF^⋄[23] LfF+BE^⋄[19] EnD^∗[30] SD^∗[42] DisEnt^∗[18] Selecmix^⋄[13] ETF-Debias Biased FFHQ 0.5 53.27$\pm$0.61 65.60$\pm$2.27 67.07$\pm$2.37 55.93$\pm$1.62 65.60$\pm$0.20 63.07$\pm$1.14 65.00$\pm$0.82 73.60$\pm$1.22 (+6.53) 1.0 57.13$\pm$0.64 72.33$\pm$2.19 73.53$\pm$1.62 61.13$\pm$0.50 69.20$\pm$0.20 68.53$\pm$2.32 67.50$\pm$0.30 76.53$\pm$1.10 (+3.00) 2.0 67.67$\pm$0.81 74.80$\pm$2.03 80.20$\pm$2.78 66.87$\pm$0.64 78.40$\pm$0.20 72.00$\pm$2.51 69.80$\pm$0.87 85.20$\pm$0.61 (+5.00) 5.0 78.87$\pm$0.83 80.27$\pm$2.02 87.40$\pm$2.00 80.87$\pm$0.42 84.80$\pm$0.20 80.60$\pm$0.53 83.47$\pm$0.61 94.00$\pm$0.72 (+6.60) Dogs & Cats 1.0 51.96$\pm$0.90 71.17$\pm$5.24 78.87$\pm$2.40 51.91$\pm$0.24 78.13$\pm$1.06 65.13$\pm$2.07 54.19$\pm$1.61 80.07$\pm$0.90 (+1.20) 5.0 76.59$\pm$1.27 85.83$\pm$1.62 88.60$\pm$1.21 79.07$\pm$0.28 89.12$\pm$0.18 82.47$\pm$2.86 81.50$\pm$1.06 92.18$\pm$0.62 (+3.06) BAR 1.0 68.00$\pm$0.43 68.30$\pm$0.97 71.70$\pm$1.33 68.25$\pm$0.19 67.33$\pm$0.35 69.30$\pm$1.27 69.83$\pm$1.02 72.79$\pm$0.21 (+1.09) 5.0 79.34$\pm$0.19 80.25$\pm$1.27 82.00$\pm$1.24 78.86$\pm$0.36 79.10$\pm$0.42 81.19$\pm$0.70 78.79$\pm$0.52 83.66$\pm$0.21 (+1.66)

Report issue for preceding element

Report issue for preceding element

5.1 Experimental Settings

Report issue for preceding element

Datasets and models. We validate the effectiveness of ETF-Debias on general debiasing benchmarks, which cover various types of bias attributes including color, corruption, gender, and background. We adopt 2 synthetic biased datasets, Colored MNIST [15] and Corrupted CIFAR-10 [10] with the ratio of bias-conflicting training samples {0.5%, 1.0%, 2.0%, 5.0%}, and 3 real-world biased datasets, Biased FFHQ (BFFHQ) [18] with bias ratio {0.5%, 1.0%, 2.0%, 5.0%}, BAR [23], and Dogs & Cats [15] with bias ratio {1.0%, 5.0%}.

Report issue for preceding element

As for the model architecture, we adopt a three-layer MLP for Colored MNIST and ResNet-20 [9] for other datasets. Since BAR has a tiny training set, we follow the previous work [19] and initialize the parameters with pre-trained models on corresponding datasets. All results are averaged over three independent trials. More details about datasets and implementation are available in Appendix B.

Report issue for preceding element

Baselines. According to the three categories of debiased learning in Section 2, we compare the performance of ETF-Debias with six recent methods. For reweight-based debiasing, we consider LfF [23] with auxiliary bias models, and its improved version LfF+BE [19]. For disentangle-based debiasing, we consider EnD [30] and SD [42], which stem from the Information Bottleneck theory and causal intervention respectively. For augmentation-based debiasing, we consider DisEnt [18] and Selecmix [13], to include both the feature-level and image-level augmentations.

Report issue for preceding element

5.2 Main Results

Report issue for preceding element

Comparison on synthetic datasets. To display the debiasing performance, we report the accuracy on the unbiased test set of 2 synthetic datasets in Tab. 2. It’s notable that ETF-Debias consistently outperforms baselines in the generalization capability towards test samples, on almost all levels of bias ratio. We observe that some baseline methods (e.g., EnD) do not display a satisfactory debiasing effect on synthetic datasets, as they rely heavily on diverse contrastive samples to identify and mitigate the bias features. In contrast, our approach directly provides the approximated shortcut features as training primes, which achieves superior performance on synthetic bias attributes.

Report issue for preceding element

Comparison on real-world datasets. To verify the scalability of ETF-Debias in real-world scenarios with more diverse bias attributes, we test our method on 3 real-world biased datasets in Tab. 3. We observe that ETF-Debias shows an even greater performance gain on real-world datasets than on synthetic ones, which may be attributed to the semantically meaningful prime features constructed with the simplex ETF structure. On the large-scale BFFHQ dataset, our method achieves up to 6.6% accuracy improvements compared to baseline methods, demonstrating its potential in real-world applications.

Report issue for preceding element

Convergence of Neural Collapse. In Fig. 2, we display the trajectory of NC metrics during training on the Corrupted CIFAR-10 dataset. Guided by the prime features, the model establishes a right correlation and shows a much better convergence property on biased datasets, contributing to the superior generalization capability. More convergence results are available in Appendix C.

Report issue for preceding element

5.3 Ablation Study

Report issue for preceding element

Ablation on the influence of regularization. To measure the sensitivity of our method to different levels of prime reinforcement regularization, we compare the accuracy on the unbiased test set with $\alpha$ range from 0.0 to 1.0 in Fig. 4(a). It’s been shown that the debiasing performance remains significant with different strengths of regularization, and achieves extra performance gain with the proper level of prime reinforcement.

Report issue for preceding element

Figure 4: Ablation studies on hyper-parameter and prime features. We report (a) test set accuracy on different datasets, with hyper-parameter $\alpha$ ranging from 0.0 to 1.0, and (b) test set accuracy on 5 datasets with different prime features. The shaded areas represent the standard deviation, and the bias ratio of all datasets is 5%.

Report issue for preceding element

Ablation on the influence of ETF prime features. As illustrated before, we choose the vertices of ETF as the “perfectly learned" shortcut features, thus redirecting the model’s attention to intrinsic correlations. To demonstrate the efficacy of ETF prime features, we compare the results of randomly initialized prime features with the same dimension as the ETF-based ones. As shown in Fig. 4(b), the randomly initialized primes suffer a severe performance degradation, underscoring the advantages of ETF-based prime features in approximating the optimal structure.

Report issue for preceding element

5.4 Visualization

Report issue for preceding element

Figure 5: The comparison of visualization results on bias-conflicting samples (first row) between vanilla models (second row) and ETF-Debias models (third row). We display the result of CAM [44] on (a) Corrupted CIFAR-10 dataset, with the bias attribute as different types of corruption on the entire image, (b) Dogs & Cats dataset, with the bias attribute as the color of animals and (c) BFFHQ dataset, with the bias attribute as gender.

Report issue for preceding element

To intuitively reveal the effectiveness of our method, we compare the CAM [44] visualization results on vanilla models and debiased models trained with ETF-Debias. As shown in Fig. 5, the model’s attention is significantly rectified with ETF-Debias. For example, on Corrupted CIFAR-10 dataset, vanilla models are easily misled by the corruptions on the entire images, but with the guide of prime features in our method, the debiased models shift their attention to the objects themselves. It’s also notable that our method circumvents the wrong attention area in classification and encourages the focus on more discriminative and finer-grained regions. On datasets of facial recognition, our method also breaks the spurious correlation on specific visual attributes [19] and considers more facial features, as shown in Fig. 5(c). More visualization results are available in Appendix E.

Report issue for preceding element

6 Conclusion

Report issue for preceding element

In this paper, we propose an avoid-shortcut learning framework with the insights of the Neural Collapse phenomenon. By extending the analysis of Neural Collapse to biased datasets, we introduce the simplex ETF as the prime features to redirect the model’s attention to intrinsic correlations. With the state-of-the-art debiasing performance on various benchmarks, we hope our work may advance the understanding of Neural Collapse and shed light on the fundamental solutions to model debiasing.

Report issue for preceding element

Supplementary Material

Report issue for preceding element

Appendix A Detailed Theoretical Justification

Report issue for preceding element

A.1 Analysis of Vanilla Training

Report issue for preceding element

To illustrate why vanilla models tend to pursue shortcut learning, we follow the analysis of previous works [36, 43] and re-examine the issue of biased classification from the perspective of gradients.

Report issue for preceding element

Following the definition in Section 3.1, we denote $\mathbf{x}_{k,i}$ as the $i$-th sample of the $k$-th class, $\mathbf{z}_{k,i}\in\mathbb{R}^{d}$ as its corresponding last-layer feature, and $\mathbf{W}=[\mathbf{w}_{1},...,\mathbf{w}_{K}]$ as the weights of classifier. In vanilla training, the cross-entropy loss is defined as:

Report issue for preceding element

$$\mathcal{L}_{\rm CE}(\mathbf{z}_{k,i},\mathbf{W})=-\log\left(\frac{\exp(\mathbf{z}_{k,i}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime}=1}^{K}\exp(\mathbf{z}_{k,i}^{\mathrm{T}}\mathbf{w}_{k^{\prime}})}\right)$$ (A.1)

Gradient w.r.t classifier weights. To analyze the learning behavior of the classifier, we first compute the gradient of $\mathcal{L}_{\rm CE}$ w.r.t the classifier weights:

Report issue for preceding element

$$\frac{\partial\mathcal{L}_{\rm CE}}{\partial\mathbf{w}_{k}}=\underbrace{\sum_{i=1}^{n_{k}}-(1-p_{k}(\mathbf{z}_{k,i}))\mathbf{z}_{k,i}}_{\rm pulling\ part}+\underbrace{\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}p_{k}(\mathbf{z}_{k^{\prime},j})\mathbf{z}_{k^{\prime},j}}_{\rm forcing\ part}$$ (A.2)

where $n_{k}$ denotes the number of training samples in the $k$-th class, and $p_{k}(\mathbf{z})$ is the predicted probability of $\mathbf{z}$ belongs to the $k$-th class, which is calculated with the softmax function:

Report issue for preceding element

$$p_{k}(\mathbf{z})=\frac{\exp(\mathbf{z}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime}=1}^{K}\exp(\mathbf{z}^{\mathrm{T}}\mathbf{w}_{k^{\prime}})},1\leq k\leq K$$ (A.3)

In Eq. A.2, we decompose the gradient w.r.t the classifier weight $\mathbf{w}_{k}$ into two parts, the pulling part and the forcing part. The pulling part contains the effects of features from the same class (i.e., $\mathbf{z}_{k,i}$), which pulls the classifier weight towards the $k$-th feature cluster and each feature has an influence of $1-p_{k}(\mathbf{z}_{k,i})$. Meanwhile, the forcing part of the gradient contains the features from other classes to push $\mathbf{w}_{k}$ away from the wrong clusters, and each feature has an influence of $p_{k}(\mathbf{z}_{k^{\prime},j})$. When the vanilla model is trained on biased datasets, the prevalent bias-aligned samples of the $k$-th class will dominate the pulling part of the gradient. The classifier weight $\mathbf{w}_{k}$ will be pulled towards the center of the bias-aligned features, which have a strong correlation between the class label $k$ and a bias attribute $b_{k}$. The biased feature space based on shortcut will thus be formed at the early period of training, as the result of the imbalanced magnitude of gradients across different attributes. Similarly, the forcing part of the gradient is also guided by the bias-aligned samples of other classes, further reinforcing the tendency of shortcut learning. It confirms our observation in Section 3.2 that the model’s pursuit of shortcut correlation leads to a biased, non-collapsed feature space, which is hard to rectify in the subsequent training steps.

Report issue for preceding element

Gradient w.r.t features. Furthermore, we compute the gradient of $\mathcal{L}_{\rm CE}$ w.r.t the last-layer features:

Report issue for preceding element

$$\frac{\partial\mathcal{L}_{\rm CE}}{\partial\mathbf{z}_{k,i}}=\underbrace{-(1-p_{k}(\mathbf{z}_{k,i}))\mathbf{w}_{k}}_{\rm pulling\ part}+\underbrace{\sum_{k^{\prime}\neq k}^{K}p_{k^{\prime}}(\mathbf{z}_{k,i})\mathbf{w}_{k^{\prime}}}_{\rm forcing\ part}$$ (A.4)

In Eq. A.4, the gradient w.r.t features is also considered as the combination of the pulling part and the forcing part. The pulling part represents the pulling effect of the classifier weight from the same class (i.e., $\mathbf{w}_{k}$), which will guide the features to align with the prediction behavior of the classifier. The forcing part, on the contrary, represents the pushing effect of other classifier weights. As we discussed before, the model’s reliance on simple shortcuts is formed at the early stage of training, due to the misled classifier weights toward the centers of bias-aligned features. It results in a biased decision rule of the classifier, which directly affects the formation of the last-layer feature space. Consider the bias-conflicting samples of the $k$-th class, although the pulling part of the gradient supports its convergence towards the right classifier weight $\mathbf{w}_{k}$, the established shortcut correlation is hard to reverse and eliminate, which has been demonstrated with the metrics of Neural Collapse in Section 3.2.

Report issue for preceding element

A.2 Analysis of ETF-Debias

Report issue for preceding element

In light of the analysis of vanilla training, our proposed debiasing framework, ETF-Debias, turns the easy-to-follow shortcut into the prime features, which guides the model to skip the active learning of shortcuts and directly focus on the intrinsic correlations. We have provided a brief theoretical justification of our method in Section 4.3, and the detailed illustrations are as follows.

Report issue for preceding element

When trained on biased datasets, we assume each class $[1,...,K]$ is strongly correlated with a bias attribute $[b_{1},...,b_{K}]$. Based on the mechanism of prime training, we denote the $i$-th feature of the $k$-th class as $\widetilde{\mathbf{z}}_{k,i}=[\mathbf{z}_{k,i},\mathbf{m}_{i,b}]\in\mathbb{R}^{2\times d}$, where $\mathbf{z}_{k,i}\in\mathbb{R}^{d}$ represents the learnable features, and $\mathbf{m}_{i,b}\in\mathbb{R}^{d}$ represents the prime features retrieved based on the bias attribute $b$ of the input sample. With the definition, a bias-aligned sample of the $k$-th class $\mathbf{x}_{k,i}$ will have its feature in form of $\widetilde{\mathbf{z}}_{k,i}=[\mathbf{z}_{k,i},\mathbf{m}_{b_{k}}]$, and a bias-conflicting sample will have its feature as $\widetilde{\mathbf{z}}_{k,i}=[\mathbf{z}_{k,i},\mathbf{m}_{b_{k^{\prime}}}],k^{\prime}\neq k$. To keep the same form, we denote the classifier weight as $\widetilde{\mathbf{w}}_{k}=[\mathbf{w}_{k},\mathbf{a}_{k}]\in\mathbb{R}^{2\times d}$, where $\mathbf{w}_{k}\in\mathbb{R}^{d}$ represents the weight for intrinsic correlations and $\mathbf{a}_{k}\in\mathbb{R}^{d}$ is the weight for shortcut features. In the detailed convergence result of Neural Collapse (Section C), we observe that, due to the fixed prime features and their strong correlation with the bias attributes, $\mathbf{a}_{k}$ will quickly collapse into its bias-correlated prime feature $\mathbf{m}_{b_{k}}$, and can be viewed as constant after just a few steps of training. Thus the cross-entropy loss can be re-written as:

Report issue for preceding element

$$\mathcal{L}_{\rm CE}(\widetilde{\mathbf{z}}_{k,i},\widetilde{\mathbf{W}})=-\log\left(\frac{\exp(\widetilde{\mathbf{z}}_{k,i}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k})}{\sum_{k^{\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k,i}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime}})}\right)=-\log\left(\frac{\exp([\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k},\mathbf{a}_{k}])}{\sum_{k^{\prime}=1}^{K}[\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k^{\prime}},\mathbf{a}_{k^{\prime}}]}\right)$$ (A.5)

Gradient w.r.t classifier weights. With the avoid-shortcut learning framework, we justify that the introduced prime mechanism implicitly plays the role of re-weighting, which weakens the mutual convergence between bias-aligned features and the classifier weights, and amplifies the learning of intrinsic correlations. We first analyze the gradient of $\mathcal{L}_{\rm CE}$ w.r.t the classifier weights $\widetilde{\mathbf{W}}$:

Report issue for preceding element

$$\frac{\partial\mathcal{L}_{\rm CE}}{\partial\widetilde{\mathbf{w}}_{k}} =\sum_{i=1}^{n_{k}}-(1-p_{k}(\widetilde{\mathbf{z}}_{k,i}))\widetilde{\mathbf{z}}_{k,i}+\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}p_{k}(\widetilde{\mathbf{z}}_{k^{\prime},j})\widetilde{\mathbf{z}}_{k^{\prime},j}$$

$$=\sum_{i=1}^{n_{k}}-\left(1-\frac{\exp([\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k},\mathbf{a}_{k}])}{\sum_{k^{\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k,i}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime}})}\right)\widetilde{\mathbf{z}}_{k,i}+\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}\frac{\exp([\mathbf{z}_{k^{\prime},j},\mathbf{m}_{j,b^{\prime}}]^{\mathrm{T}}[\mathbf{w}_{k},\mathbf{a}_{k}])}{\sum_{k^{\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k^{\prime},j}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime}})}\widetilde{\mathbf{z}}_{k^{\prime},j}$$ (A.6)

$$\leq\sum_{i=1}^{n_{k}}-\left(1-\frac{\exp(\mathbf{z}_{k,i}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k,i}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime}})}-\frac{\exp(\mathbf{m}_{i,b}^{\mathrm{T}}\mathbf{a}_{k})}{\sum_{k^{\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k,i}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime}})}\right)\widetilde{\mathbf{z}}_{k,i}$$

$$\ \ \ \ +\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}\left(\frac{\exp(\mathbf{z}_{k^{\prime},j}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k^{\prime},j}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime\prime}})}+\frac{\exp(\mathbf{m}_{j,b^{\prime}}^{\mathrm{T}}\mathbf{a}_{k})}{\sum_{k^{\prime\prime}=1}^{K}\exp(\widetilde{\mathbf{z}}_{k^{\prime},j}^{\mathrm{T}}\widetilde{\mathbf{w}}_{k^{\prime\prime}})}\right)\widetilde{\mathbf{z}}_{k^{\prime},j}$$ (A.7)

$$\leq\underbrace{\sum_{i=1}^{n_{k}}-(1-p_{k}^{(b)}(\mathbf{m}_{i,b})-p_{k}^{(l)}(\mathbf{z}_{k,i}))\widetilde{\mathbf{z}}_{k,i}}_{\rm pulling\ part}+\underbrace{\sum_{k^{\prime}\neq k}^{K}\sum_{j=1}^{n_{k^{\prime}}}(p_{k}^{(b)}(\mathbf{m}_{j,b^{\prime}})+p_{k}^{(l)}(\mathbf{z}_{k^{\prime},j}))\widetilde{\mathbf{z}}_{k^{\prime},j}}_{\rm forcing\ part}$$ (A.8)

The predicted probabilities both satisfy $0<(1-p_{k}(\widetilde{\mathbf{z}}_{k,i}))<1$ and $0<p_{k}(\widetilde{\mathbf{z}}_{k^{\prime},j})<1$, and the scaling step from Eq. A.6 to Eq. A.7 is based on the Jensen inequality. The terms $p_{k}^{(l)}$ and $p_{k}^{(b)}$ are respectively the predicted probabilities for class labels and bias attributes, calculated with softmax:

Report issue for preceding element

$$p_{k}^{(l)}(\mathbf{z}_{k,i})=\frac{\exp(\mathbf{z}_{k,i}^{\mathrm{T}}\mathbf{w}_{k})}{\sum_{k^{\prime}=1}^{K}\exp([\mathbf{z}_{k,i},\mathbf{m}_{i,b}]^{\mathrm{T}}[\mathbf{w}_{k^{\prime}},\mathbf{a}_{k^{\prime}}])}$$ (A.9)

$$p_k^{(b)}({𝐦}_{i,b})=\frac{\exp (}{}$$