Title: Don’t blame Dataset Shift! Shortcut Learning due to Gradients and Cross Entropy
Venue: NeurIPS 2023
Accepted: yes
Keywords: shortcut learning

spurious correlations

perfect stable feature

perception tasks

implicit bias in optimization

improving inductive biases
Abstract: Common explanations for shortcut learning assume that the shortcut improves prediction only under the training distribution. Thus, models trained in the typical way by minimizing log-loss using gradient descent, which we call default-ERM, should utilize the shortcut. However, even when the stable feature determines the label in the training distribution and the shortcut does not provide any additional information, like in perception tasks, default-ERM exhibits shortcut learning. Why are such solutions preferred when the loss can be driven to zero when using the stable feature alone? By studying a linear perception task, we show that default-ERM’s preference for maximizing the margin, even without overparameterization, leads to models that depend more on the shortcut than the stable feature. This insight suggests that default-ERM’s implicit inductive bias towards max-margin may be unsuitable for perception tasks. Instead, we consider inductive biases toward uniform margins. We show that uniform margins guarantee sole dependence on the perfect stable feature in the linear perception task and suggest alternative loss functions, termed margin control (MARG-CTRL), that encourage uniform-margin solutions. MARG-CTRL techniques mitigate shortcut learning on a variety of vision and language tasks, showing that changing inductive biases can remove the need for complicated shortcut-mitigating methods in perception tasks.

Don't blame Dataset Shift! Shortcut Learning due to Gradients and Cross Entropy

Abstract

Commonexplanations for shortcut learning assume that the shortcut improves prediction under the training distribution but not in the test distribution. Thus, models trained via the typical gradient-based optimization of cross-entropy, which we call defaultERM , utilize the shortcut. However, even when the stable feature determines the label in the training distribution and the shortcut does not provide any additional information, like in perception tasks, defaultERM still exhibits shortcut learning. Why are such solutions preferred when the loss for defaultERM can be driven to zero using the stable feature alone? By studying a linear perception task, we show that defaultERM 's preference for maximizing the margin leads to models that depend more on the shortcut than the stable feature, even without overparameterization. This insight suggests that defaultERM 's implicit inductive bias towards max-margin is unsuitable for perception tasks. Instead, we develop an inductive bias toward uniform margins and show that this bias guarantees dependence only on the perfect stable feature in the linear perception task. We develop loss functions that encourage uniform-margin solutions, called margin control ( MARG -CTRL ). M ARG -CTRL mitigates shortcut learning on a variety of vision and language tasks, showing that better inductive biases can remove the need for expensive two-stage shortcut-mitigating methods in perception tasks.

1 Introduction

Shortcut learning is a phenomenon where a model learns to base its predictions on an unstable correlation, or shortcut , that does not hold across data distributions collected at different times and/or places [Geirhos et al., 2020]. A model that learns shortcuts can perform worse than random guessing in settings where the label's relationship with the shortcut feature changes [Koh et al., 2021, Puli et al., 2022]. Such drops in performance do not occur if the model depends on features whose relationship with the label does not change across settings; these are stable features.

Shortcut learning is well studied in cases where models that use both shortcut and stable features achieve lower loss than models that only use the stable feature [Arjovsky et al., 2019, Puli et al., 2022, Geirhos et al., 2020]. These works consider cases where the Bayes-optimal classifier - the training conditional distribution of the label given the covariates - depends on both stable and shortcut features. In such cases, shortcut learning occurs as the Bayes-optimal predictor is the target of standard supervised learning algorithms such as the one that minimizes the log-loss via gradient descent ( GD ), which we call defaultERM .

However, in many machine learning tasks, the stable feature perfectly predicts the label, i.e. a perfect stable feature . For example, in task of predicting hair color from images of celebrity faces in the

CelebA dataset [Sagawa et al., 2020a], the color of the hair in the image determines the label. This task is a perception task. In such classification tasks, the label is independent of the shortcut feature given the stable feature, and the Bayes-optimal predictor under the training distribution only depends on the stable feature. DefaultERM can learn this Bayes-optimal classifier which, by depending solely on the stable feature, also generalizes outside the training distribution. But in practice, defaultERM run on finite data yields models that depend on the shortcut and thus perform worse than chance outside the training distribution [Sagawa et al., 2020a, Liu et al., 2021, Zhang et al., 2022]. The question is, why does defaultERM prefer models that exploit the shortcut even when a model can achieve zero loss using the stable feature alone?

To understand preferences toward shortcuts, we study defaultERM on a linear perception task with a stable feature that determines the label and a shortcut feature that does not. The perfect linear stable feature means that data is linearly separable. This separability means that defaultERM -trained linear models classify in the same way as the minimum ℓ 2 -norm solution that has all margins greater than 1 ; the latter is commonly called max-margin classification [Soudry et al., 2018]. We prove that defaultERM 's implicit inductive bias toward the max-margin solution is harmful in that defaultERM -trained linear models depend more on the shortcut than the stable feature. In fact, such dependence on the shortcut occurs even in the setting with fewer parameters in the linear model than data points, i.e. without overparameterization. These observations suggest that a max-margin inductive bias is unsuitable for perception tasks.

Next, we study inductive biases more suitable for perception tasks with perfect stable features. We first observe that predicting with the perfect stable feature alone achieves uniform margins on all samples. Formally, if the stable feature s ( x ) determines the label y via a function d , y = d ◦ s ( x ) , one can achieve any positive b as the margin on all samples simultaneously by predicting with b · d ◦ s ( x ) . We show that in the same setting without overparameterization where max-margin classification leads to shortcut learning, models that classify with uniform margins depend only on the stable feature.

Building on these observations, we identify alternative loss functions that are inductively biased toward uniform margins, which we call margin control ( MARG -CTRL ). We empirically demonstrate that MARG -CTRL mitigates shortcut learning on multiple vision and language tasks without the use of annotations of the shortcut feature in training. Further, M ARG -CTRL performs on par or better than the more expensive two-stage shortcut-mitigating methods [Liu et al., 2021, Zhang et al., 2022]. Wethen introduce a more challenging setting where both training and validation shortcut annotations are unavailable, called the nuisance-free setting. In the nuisance-free setting, MARG -CTRL always outperforms defaultERM and the two-stage shortcut-mitigating methods. These empirical results suggest that simply incorporating inductive biases more suitable for perception tasks is sufficient to mitigate shortcuts.

2 Shortcut learning in perception tasks due to maximizing margins

Setup. We use y , z , x to denote the label, the shortcut feature, and the covariates respectively. We let the training and test distributions ( p tr , p te ) be members of a family of distributions indexed by ρ , F = { p ρ ( y , z , x ) } ρ , such that the shortcut-label relationship p ρ ( z , y ) changes over the family. Many common tasks in the spurious correlations literature have stable features s ( x ) that are perfect, meaning that the label is a deterministic function d of the stable feature: y = d ◦ s ( x ) . For example, in the Waterbirds task the bird's body determines the label and in the CelebA task, hair color determines the label [Sagawa et al., 2020a]. As s ( x ) determines the label, it holds that y | = p ρ ( x , z ) | s ( x ) . Then, the optimal predictor on the training distribution is optimal on all distributions in the family F , regardless of the shortcut because

The most common procedure to train predictive models to approximate p tr ( y | x ) is gradient-based optimization of cross-entropy (also called log-loss); we call this defaultERM . DefaultERM targets the Bayes-optimal predictor of the training distribution which, in tasks with perfect stable features, also performs optimally under the test distribution. However, despite targeting the predictor that does not depend on the shortcut, models built with defaultERM still rely on shortcut features that are often less predictive of the label and are unstable, i.e. vary across distributions [Geirhos et al., 2020, Puli et al., 2022]. We study defaultERM 's preference for shortcuts in a data generating process ( DGP ) where both the shortcut and the perfect stable feature are linear functions of the covariates.

2.1 Shortcut learning in linear perception tasks

Let Rad be the uniform distribution over { 1 , -1 } , N be the normal distribution, d be the dimension of x , and ρ ∈ (0 , 1) , B > 1 be scalar constants. The DGP for p ρ ( y , z , x ) is:

This DGP is set up to mirror the empirical evidence in the literature showing that shortcut features are typically learned first [Sagawa et al., 2020a]. The first dimension of x , i.e. x 1 , is a shortcut that is correlated with y according to ρ . The factor B in x 1 scales up the gradients for parameters that interact with x 1 in predictions. For large enough B , model dependence on the shortcut feature during defaultERM goes up faster than the stable feature [Idrissi et al., 2022].

The training distribution is p tr = p 0 . 9 and the test distribution is one where the shortcut's relationship with the label is flipped p te = p 0 . 1 . Models achieve worse than random test accuracy ( 50% ) if they exploit the training shortcut relationship and the predicted class flips when the shortcut feature flips. We train with defaultERM which uses log-loss: on a data point ( x , y ) the log-loss is

With d = 300 and B = 10 , we train a linear model on 1000 samples from the training distribution p ρ =0 . 9 , and evaluate on 1000 samples from p ρ =0 . 1 .

Observations. Figure 1a shows that when trained with defaultERM , the linear model does not do better than chance ( < 50% ) on the test data even after 50 , 000 epochs. So, even in the presence of the perfect feature x 2 , the model relies on other features like the shortcut x 1 . Since the final training loss is very small, on the order of 10 -9 , this result is not due to optimization being stuck in a local minima with high loss. These observations indicate that, in the linear setting, gradientbased optimization with log-loss prefers

(a)

Average accuracy and loss curves.

LoSS

100

100

80

train

10-3

60

test

10-6

40

10-9

0

1

2

3

4

5

0

1

2

3

4

5

1e4

1e4

Loss

100

3

Tr, Y = Z

80

Te,Y=Z

Tr, Y≠Z

60

Te,Y≠Z

40

20

10

0

5000

10000

0

5000

10000

epochs

epochs

models that depend more on the shortcut than the perfect stable feature.

To better understand this preference we focus on the errors in specific groups in the data. Consider the classifier that only uses the shortcut z and makes the Bayes-optimal prediction w.r.t p tr : arg max y p tr ( y = y | z ) . We call instances that are classified correctly by this model the shortcut group, and the rest the leftover group. We use these terms for instances in the training set as well as the test set. In this experiment y is positively correlated with z , hence the shortcut group consists of all instances with y i = z i and the leftover group of those with y i = z i .

Figure 1b gives accuracy and loss curves on the shortcut and leftover groups for the first 10000 epochs. The test accuracy for the shortcut group hits 90% while the leftover group test accuracy is < 40% , meaning that the model exploits the shortcuts. Even though a model that relies solely on the shortcut misclassifies the leftover group, we see that the training loss of the learned model on this group approaches 0 . The model drives down training loss in the leftover group by depending on noise, which results in larger test loss in the leftover group than the shortcut group. Thus, fig. 1b demonstrates that the defaultERM -trained model classifies the training shortcut group by using the shortcut feature while overfitting to the training leftover group.

∕

Shortcut dependence like in fig. 1 occurs even with ℓ 2 -regularization and when training neural networks; see appendix B.1 and appendix B.4 respectively. Next, we analyze the failure mode in fig. 1, showing that the shortcut dependence is due to defaultERM 's implicit bias to learn the max-margin classifier. Next, we study the failure mode in fig. 1 theoretically, showing that the shortcut dependence is due to defaultERM 's inductive bias toward learning the max-margin classifier.

Max-margin classifiers depend more on the the shortcut than the stable feature. We consider training a linear model f θ ( x ) = w ⊤ x where w = [ w z , w y , w e ] with defaultERM . Data from eq. (1) is always linearly separable due to the perfect stable feature, but many hyperplanes that separate the two classes exist. When a linear model is trained with defaultERM on linearly separable data, it achieves zero training loss and converges to the direction of a minimum ℓ 2 -norm solution that achieves a margin of at least 1 on all samples [Soudry et al., 2018, Wang et al., 2021, 2022]; this is called the max-margin solution. We now show that for a small enough leftover group, large enough scaling factor B and dimension d of the covariates, max-margin solutions depend more on the shortcut feature than the stable feature:

Theorem 1. Let w ∗ be the max-margin predictor on n training samples from eq. (1) with a leftover group of size k . There exist constants C 1 , C 2 , N 0 > 0 such that with probability at least 1 -1 / 3 n over draws of the training data, it holds that B w ∗ z > w ∗ y .

The size of the leftover group k concentrates around (1 -ρ ) n because each sample falls in the leftover group with probability (1 -ρ ) . Thus, for ρ > 0 . 9 , that is for a strong enough shortcut, the condition in theorem 1 that k < n / 10 will hold with probability close to 1 ; see appendix A.5 for more details.

The proof is in appendix A. The first bit of intuition is that using the shortcut can have lower norm because of the scaling factor B . Using the shortcut only, however, misclassifies the leftover group. The next bit of intuition is that using noise from the leftover group increases margins in one group at a rate that scales with the dimension d , while the cost in the margin for the other group only grows as √ d . This trade-off in margins means the leftover group can be correctly classified using noise without incorrectly classifying the shortcut group. The theorem then leverages convex duality to show that this type of classifier that uses the shortcut and noise has smaller ℓ 2 -norm than any linear classifier that uses the stable feature more.

The way the margin trade-off in the proof works is by constructing a linear classifier whose weights on the noise features are a scaled sum of the product of the label and the noise vector in the leftover group: for a scalar γ , the weights w e = γ 󰁓 i ∈ S leftover y i δ i . The margin change on the j th training sample from using these weights is y j w ⊤ e δ j . For samples in the shortcut group, the margin change looks like a sum of mean zero independent and identically distributed variables; the standard deviation of this sum grows as √ d . For samples in the leftover group, the margin change is the sum of mean one random variables; this sum grows as d and its standard deviation grows as √ d . The difference in mean relative to the standard deviation is what provides the trade-off in margins.

We now discuss three implications of the theorem.

First, theorem 1 implies that the leftover group sees worse than random accuracy ( 0 . 5 ). To see this, note that for samples in the leftover group the margin y ( w ∗ ) ⊤ x = w ∗ y -B w ∗ z +( w ∗ e ) ⊤ y δ is a Gaussian random variable centered at a negative number w ∗ y -B w ∗ z . Then, with Φ e as the CDF of the zero-mean Gaussian random variable ( w ∗ e ) ⊤ δ , accuracy in the test leftover group is

∕

Second, the leftover group in the training data is overfit in that the contribution of noise in prediction ( | ( w ∗ e ) ⊤ δ | ) is greater than the contribution from the stable and shortcut features. Formally, in the training leftover group, w ∗ y -B w ∗ z < 0 . Then, due to max-margin property,

Third, many works point to overparameterization as one of the causes behind shortcut learning [Sagawa et al., 2020a, Nagarajan et al., 2021, Wald et al., 2023], but in the setup in fig. 1, the

linear model has fewer parameters than samples in the training data. In such cases with nonoverparameterized linear models, the choice of defaultERM is typically not questioned, especially when a feature exists that linearly separates the data. Corollary 1 formally shows shortcut learning for non-overparameterized linear models. In words, defaultERM -that is vanilla logistic regression trained with gradient-based optimization - can yield models that rely more on the shortcut feature even without overparameterization .

Corollary 1. For all n > N 0 -where the constant N 0 is from theorem 1 - with scalar τ ∈ (0 , 1) such that the dimension of x is d = τ n < n , for all integers k < n × min 󰁱 1 10 , τ C 1 log 3 n 󰁲 , a linear model trained via defaultERM yields a predictor w ∗ such that B w ∗ z > w ∗ y .

If defaultERM produces models that suffer from shortcut learning even without overparameterization, its implicit inductive bias toward max-margin classification is inappropriate for perception tasks in the presence of shortcuts. Next, we study inductive biases more suited to perception tasks.

3 Toward inductive biases for perception tasks with shortcuts

The previous section formalized how defaultERM solutions, due to the max-margin inductive bias, rely on the shortcut and noise to minimize loss on training data even in the presence of a different zero-population-risk solution. Are there inductive biases more suitable for perception tasks?

Given a perfect stable feature s ( x ) for a perception task, in that for a function d when y = d ◦ s ( x ) , one can achieve margin b ∈ (0 , ∞ ) uniformly on all samples by predicting with the stable b · d ◦ s ( x ) . In contrast, max-margin classifiers allow for disparate margins as long as the smallest margin crosses 1 , meaning that it does not impose uniform margins. The cost of allowing disparate margins is the preference for shortcuts even without overparamterization (corollary 1). In the same setting however, any uniform-margin classifier for the linear perception task (eq. (1)) relies only on the stable feature:

Theorem 2. Consider n samples of training data from DGP in eq. (1) with d < n . Consider a linear classifier f θ ( x ) = w ⊤ x such that for all samples in the training data y i w ⊤ x i = b for any b ∈ (0 , ∞ ) . With probability 1 over draws of samples, w = [0 , b, 0 d -2 ] .

Theorem 2 shows that uniform-margin classifiers only depend on the stable feature, standing in contrast with max-margin classifiers which can depend on the shortcut feature (theorem 1). The proof is in appendix A.6. Thus, inductive biases toward uniform margins are better suited for perception tasks. Next, we identify several ways to encourage uniform margins.

Margin control ( MARG -CTRL ). To produce uniform margins with gradient-based optimization, we want the loss to be minimized at uniform-margin solutions and be gradient-optimizable. We identify a variety of losses that satisfy these properties, and we call them MARG -CTRL losses. MARG -CTRL losses have the property that per-sample loss monotonically decreases for margins until a threshold then increases for margins beyond it. In turn, minimizing loss then encourages all margins to move to the threshold.

Mechanically, when models depend more on shortcuts than the stable feature during training, margins on samples in the shortcut group will be larger than those in the leftover group; see the right panel in fig. 1b where the train loss in the shortcut group is lower than the leftover group indicating that the margins are smaller in the leftover group. This difference is margins is a consequence of the shortcut matching the label in one group and not the other, thus, encouraging the model to have similar margins across all samples pushes the model to depend less on the shortcut. In contrast, vanilla log-loss can be driven to zero in a direction with disparate margins across the groups as long as the margins on all samples go to ∞ . We define MARG -CTRL losses for a model f θ with the margin on a sample ( x , y ) defined as y f θ ( x ) .

As the first MARG -CTRL loss, we develop the σ -damped log-loss: we evaluate log-loss on a margin multiplied by a monotonically decreasing function of the margin. In turn, the input to the loss increases with the margin till a point and then decreases. For a temperature T and sigmoid function

σ , the σ -damped loss modifies the model output f θ and plugs it into log-loss:

For large margin predictions y f θ > 0 , the term 1 -σ ( y f θ ( x ) / T ) damps down the input to log-loss. The largest the input to ℓ log can get is 0 . 278 T , found by setting the derivative to zero, thus lower bounding the loss. As log-loss is a decreasing function of its input, the minimum of ℓ σ -damp occurs when the margin is 0 . 278 T on all samples. To demonstrate empirical advantage, we compare

standard log-loss to σ -damped loss on eq. (1); see fig. 2. The left panel of figure fig. 2 shows that test accuracy is better for σ -damp. The middle and right panels shows the effect of controlling margins in training, where losses on shortcut and leftover groups hover at the same value.

Second, we design the σ -stitch loss, which imitates log-loss when y f θ ( x ) < u and penalizes larger margins ( y f θ > u ) by negating the sign of y f θ ( x ) :

As the third MARG -CTRL loss, we directly penalize large margins via a log -penalty:

The fourth MARG -CTRL loss controls margins by penalizing | f θ ( x ) | 2 :

This last penalty was called spectral decoupling ( SD ) by Pezeshki et al. [2021], who use it as a way to decouple learning dynamics in the neural tangent kernel ( NTK ) regime. Instead, from the lens of MARG -CTRL , SD mitigates shortcuts in eq. (1) because it encourages uniform margins, even though SD was originally derived from different principles, as we discuss in section 4. In appendix B.2, we plot all MARG -CTRL losses and show that MARG -CTRL improves over defaultERM on the linear perception task; see figs. 6 to 8. We also run MARG -CTRL on a neural network and show that while defaultERM achieves test accuracy worse than random chance, MARG -CTRL achieves 100% test accuracy; see figs. 10 to 13 in appendix B.4.

4 Related work

A large body of work tackles shortcut learning under different assumptions [Arjovsky et al., 2019, Wald et al., 2021, Krueger et al., 2020, Creager et al., 2021, Veitch et al., 2021, Puli et al., 2022, Heinze-Deml and Meinshausen, 2021, Belinkov and Bisk, 2017]. A different line of work focuses on learning in neural networks in idealized settings [Yang and Salman, 2019, Ronen et al., 2019, Jo and Bengio, 2017, Baker et al., 2018, Saxe et al., 2013, Gidel et al., 2019, Advani et al., 2020].

Shah et al. [2020] study simplicity bias [Valle-Perez et al., 2018] and show that neural networks provably learn the linear function over a non-linear one, in the first epoch of training. In a similar vein, Hermann and Lampinen [2020] show that neural networks can prefer a linearly-decodable feature over a non-linear but more predictive feature, and Scimeca et al. [2021] make similar observations and use loss landscapes to empirically study which features are easier to learn. Simplicity bias alone only describes neural biases early in training and does not explain why more predictive stable features are not learned later. Unlike simplicity bias which focuses on linear versus nonlinear features, max-margin bias is the reason defaultERM prefers one linear feature, the shortcut, over another, the stable feature, like in the synthetic experiment in section 2.

While Pezeshki et al. [2021] allow for perfect features, they hypothesize that shortcut learning occurs because when one feature is learned first, other features are gradient-starved and are not learned as

well. They focus on a special setting where feature representations for different samples have inner product equal to a small constant to show that models can depend more on the imperfect feature than the perfect feature. In this special setting, they show that penalizing the magnitudes of what we call the margin mitigates shortcuts; this method is called spectral decoupling ( SD ). However, as we show in appendix B.5, the assumption in Lemma 1 [Pezeshki et al., 2021] is violated when using a linear model to classify in the simple linear DGP in eq. (1). However, SD on a linear model mitigates shortcuts in the DGP in eq. (1); see B.5. Thus, the theory in Pezeshki et al. [2021] fails to not explain why SD works for eq. (1), but the uniform-margin property explains why all the MARG -CTRL losses, including SD , mitigate shortcuts.

Nagarajan et al. [2021] consider tasks with perfect stable features and formalize geometric properties of the data that make max-margin classifiers have non-zero weight on the shortcut feature ( w z > 0 ). In their set up, the linear models are overparameterized and it is unclear when w z > 0 leads to worse-than-random accuracy in the leftover group because they do not separate the model's dependence on the stable feature from the dependence on noise. See fig. 14 for an example where w z > 0 but test accuracy is 100% . In contrast to Nagarajan et al. [2021], theorem 1 gives DGP s where the leftover group accuracy is worse than random, even without overparameterization. Ahuja et al. [2021] also consider linear classification with defaultERM with a perfect stable feature and conclude that defaultERM learns only the stable feature because they assume no additional dimensions of noise in the covariates. We develop the necessary nuance by including noise in the problem and showing defaultERM depends on the shortcut feature even without overparameterization.

Sagawa et al. [2020b] and Wald et al. [2023] both consider overparameterized settings where the shortcut feature is informative of the label even after conditioning on the stable feature. In both cases, the Bayes-optimal predictor also depends on the shortcut feature, which means their settings do not allow for an explanation of shortcut dependence in examples like fig. 1. In contrast, we show shortcut dependence occurs even in the presence of a perfect stable feature and without overparameterization. Li et al. [2019], Pezeshki et al. [2022] focus on relative feature complexity and discuss the effects of large learning rate ( LR ) on which features are learned first during training, but do not allow for perfect features. Idrissi et al. [2022] empirically find that tuning LR and weight decay ( WD ) gets defaultERM to perform similar to two-stage shortcut-mitigating methods like Just Train Twice ( JTT ) [Liu et al., 2021]. We view the findings of [Idrissi et al., 2022] through the lens of MARG -CTRL and explain how large LR and WD approximate MARG -CTRL to mitigate shortcuts; see section 5.

M ARG -CTRL is related to but different from methods proposed in Liu et al. [2017], Cao et al. [2019], Kini et al. [2021]. These works normalize representations or the last linear layers and linearly transform the logits to learn models with better margins under label imbalance. Next, methods like Learning from Failure ( LFF ) [Nam et al., 2020], JTT [Liu et al., 2021], and Correct-n-Contrast ( CNC ) [Zhang et al., 2022] build two-stage procedures to avoid shortcut learning without group annotations in training. They assume that defaultERM produces models that depend more on the shortcut and select hyperparamters of the two stage process using validation group annotations. In the nuisancefree setting where there are no validation group annotations, the performance of these methods can degrade below that of defaultERM . In contrast, better characterizing the source of shortcut learning in perceptual problems leads to MARG -CTRL methods that are not as reliant on validation group annotations (see nuisance-free results in Section 5). Without any group annotations, encouraging uniform margins via MARG -CTRL mitigates shortcuts better than JTT and CNC .

Soudry et al. [2018] characterize the inductive bias of gradient descent to converge in direction to max-margin solutions when using exponentially tailed loses; Wang et al. [2021, 2022] then prove similar biases toward max-margin solutions for Adam and RMSProp. Ji et al. [2020] show that for general losses that decrease in y f θ ( x ) , gradient descent has an inductive bias to follow the ℓ 2 -regularization path. All these inductive biases prefer shortcuts if using them leads to lower loss within an ℓ 2 -norm-budget. M ARG -CTRL provides a different inductive bias toward producing the same margin on all samples, which means gradient descent veers models away from imperfect shortcuts that lead to disparity in network outputs. Such inductive biases are suitable for tasks where a feature determines the label ( h ( x ) = y ).

5 Vision and language experiments

We evaluate MARG -CTRL on common datasets with shortcuts: Waterbirds, CelebA [Sagawa et al., 2020a], and Civilcomments [Koh et al., 2021]. First, MARG -CTRL always improves over defaultERM . Then, we show that MARG -CTRL performs similar to or better than two-stage shortcut-

mitigating methods like JTT [Liu et al., 2021] and CNC [Zhang et al., 2022] in traditional evaluation settings where group annotations are available in the validation data. Finally, we introduce a more challenging setting that only provides class labels in training and validation, called the nuisancefree setting . In contrast to the traditional setting that always assumes validation group annotations, the nuisance-free setting does not provide group annotations in either training or in validation. In the nuisance-free setting, MARG -CTRL outperforms JTT and CNC , even though the latter are supposed to mitigate shortcuts without knowledge of the groups.

Datasets. We use the Waterbirds and CelebA datasets from Sagawa et al. [2020a] and the CivilComments dataset from Borkan et al. [2019], Koh et al. [2021]. In Waterbirds, the task is to classify images of a waterbird or landbird, and the label is spuriously correlated with the image background consisting of land or water. There are two types of birds and two types of background, leading to a total of 4 groups defined by values of y, z . In CelebA [Liu et al., 2015, Sagawa et al., 2020a], the task is to classify hair color of celebrities as blond or not. The gender of the celebrity is a shortcut for hair color. There are two types of hair color and two genders in this dataset, leading to a total of 4 groups defined by values of y, z . In CivilComments-WILDS [Borkan et al., 2019, Koh et al., 2021], the task is to classify whether an online comment is toxic or non-toxic, and the label is spuriously correlated with mentions of certain demographic identities. There are 2 labels and 8 types of the shortcut features, leading to 16 groups.

Metrics, model selection, and hyperparameters. We report the worst-group test accuracy for each method. The groups are defined based on the labels and shortcut features. The more a model depends on the shortcut, the worse the worst-group error. Due to the label imbalance in all the datasets, we use variants of σ -damp, σ -stitch, MARG -LOG , and SD with class-dependent hyperparameters; see appendix B.6.2. For all methods, we use the standard Adam optimizer [Kingma and Ba, 2015] and let the learning rate and weight decay hyperparameters be tuned along with the method's hyperparameters. We first report results for all methods using validation worst-group accuracy to select method and optimization hyperparameters and early stop. For both JTT and CNC , this is the evaluation setting that is used in

existing work [Liu et al., 2021, Idrissi et al., 2022, Zhang et al., 2022]. Finally, in the nuisance-free setting where no group annotations are available, we select hyperparameters using label-balanced average accuracy. Appendix B.6 gives further details about the training, hyperparameters, and experimental results.

5.1 M ARG -CTRL mitigates shortcuts in the default setting

Here, we experiment in the standard setting from Liu et al. [2021], Idrissi et al. [2022], Zhang et al. [2022] and use validation group annotations to tune hyperparameters and early-stopping.

M ARG -CTRL improves over defaultERM . We compare MARG -CTRL to defaultERM on CelebA, Waterbirds, and Civilcomments. Table 1 shows that every MARG -CTRL method achieves higher test worst-group accuracy than defaultERM on all datasets. DefaultERM achieves a mean test worst-group accuracy of 70 . 8% , 72 . 8% and 60 . 1% on Waterbirds, CelebA, and Civilcomments respectively. Compared to defaultERM , MARG -CTRL methods provide a 5 -10% improvement on Waterbirds, 7 -10% improvement on CelebA, 7 -10% improvement on Civilcomments. These improvements show the value of inductive biases more suitable for perception tasks.

Large LR and WD may imitate MARG -CTRL in ERM . DefaultERM 's performance varies greatly across different values of LR and WD on, for instance, CelebA: the test worst-group accuracy improves by more than 20 points over different LR and WD combinations. Why does tuning LR and WD yield such improvements? We explain this phenomenon as a consequence of instability in optimiza-

tion induced by large LR and WD which prevents the model from maximizing margins and in turn can control margins. Figure 3 provides evidence for this explanation by comparing defaultERM 's loss curves for two LR and WD combinations.

The blue curves in fig. 3 correspond to the run with the larger LR and WD combination. The model that achieves the best validation (and test) worst-group accuracy over all combinations of hyperparameters for defaultERM , including those not in the plot, is the one at epoch 13 on the blue curves. This model achieves similar train and test losses ( ≈ 0 . 4 ) and thus similar margins in the shortcut group, the leftover group, and the whole dataset. The red curves stand in contrast where the lower LR results in the leftover group having higher training and test losses, and therefore smaller margins, compared to the shortcut group. These observations together support the explanation that defaultERM with large LR and WD mitigates shortcuts when controlling margins like MARG -CTRL .

M ARG -CTRL performs as well or better than two-stage shortcut-mitigating methods. Twostage shortcut mitigating methods like Correct-n-Contrast ( CNC ) and Just Train Twice ( JTT ) aim to mitigate shortcuts by using a model trained with defaultERM to approximate group annotations. They rely on the assumption that a model trained via defaultERM either predicts with the shortcut feature (like background in Waterbirds) or that the model's representations separate into clusters based on the shortcut feature. The methods then approximate group annotations using this defaultERM -trained model and use them to mitigate shortcut learning in a second predictive model. J TT upweights the loss on the approximate leftover group and CNC uses a contrastive loss to enforce the model's representations to be similar across samples that have the same label but different approximate group annotations. Appendix B.6.1 gives details.

Table 1 compares MARG -CTRL to JTT and CNC on Waterbirds, Celeba, and CivilComments. On CelebA, SD , marg-log , and σ -stitch perform similar to CNC while all MARG -CTRL techniques outperform JTT . On Waterbirds, all MARG -CTRL methods outperform JTT and CNC . On CivilComments, σ -damp and SD perform similar to JTT and outperform CNC . C NC 's performance on Waterbirds differs from Zhang et al. [2022] because their reported performance requires unique large WD choices (like WD set to 1 ) to build a first-stage model that relies most on the shortcut feature without overfitting to the training data.

M ARG -CTRL is faster than JTT and CNC . M ARG -CTRL takes the same time as defaultERM , taking around 1 , 20 and 60 minutes per epoch for Waterbirds, CelebA, and CivilComments respectively on an RTX8000 GPU. In contrast, on average over runs, JTT takes around 6 , 80 , 120 minutes per epoch and CNC takes around 8 , 180 , 360 minutes per epoch. Thus, MARG -CTRL performs as well or better than JTT and CNC while being simpler to implement and computationally cheaper.

5.2 M ARG -CTRL mitigates shortcuts in the nuisance-free setting

Work like [Liu et al., 2021, Zhang et al., 2022] crucially require validation group annotations because these methods push the work of selecting models for mitigating shortcuts to validation. Determining shortcuts itself is a laborious manual process, which means group annotations will often be unavailable. Further, given a perfect stable feature that determines the label and a shortcut that does not, only models that rely on the stable feature more than the shortcut achieve the highest validation ac-

curacy. Thus, we introduce a more challenging setting that only provides class labels in training and validation, called the nuisance-free setting . In the nuisance-free setting, models are selected based on label-balanced average accuracy: the average of the accuracies over samples of each class.

Table 2 reports test worst-group (WG) accuracy in the nuisance-free setting. On all the datasets, every MARG -CTRL outperforms defaultERM , JTT , and CNC . On average, the MARG -CTRL methods close at least 61% of the gap between defaultERM in the nuisance-free setting and the best performance in table 1 on every dataset. In contrast, CNC and JTT sometimes perform worse than defaultERM .

6 Discussion

Westudy why defaultERM -gradient-based optimization of log-loss - yields models that

depend on the shortcut even when the population minimum of log-loss is achieved by models that depend only on the stable feature. By studying a linear task with perfect stable features, we show that defaultERM 's preference toward shortcuts sprouts from an inductive bias toward maximizing margins. Instead, inductive biases toward uniform margins improve dependence on the stable feature and can be implemented via MARG -CTRL . M ARG -CTRL improves over defaultERM on a variety of perception tasks in vision and language without group annotations in training, and is competitive with more expensive two-stage shortcut-mitigating methods. In the nuisance-free setting, where even validation group annotations are unavailable, MARG -CTRL outperforms all the baselines. The performance that M ARG -CTRL yields demonstrates that changing inductive biases can remove the need for expensive shortcut-mitigating methods in perception tasks.

Without overparameterization, uniform-margin classifiers are unique and learn stable features only, while max-margin classifiers can depend more on shortcuts. With overparameterization, maxmargin classifiers are still unique but uniform-margin solutions are not which necessitates choosing between solutions. The experiments in section 5 suggest that choosing between uniform-margin classifiers with penalties like ℓ 2 improves over max-margin classifiers with ℓ 2 : all experiments use overparameterized models trained with weight decay and MARG -CTRL outperforms defaultERM . Further, our experiments suggest that uniform-margin classifiers are insensitive to the WD and LR choices, unlike max-margin classifiers; appendix B.8 shows that MARG -CTRL achieves high performance for all LR and WD choices but ERM requires tuning.

Theorem 1 also explains how balancing may or may not improve dependence on the stable features. For example, a weighting-based approach produces the same max-margin solution as defaultERM [Sagawa et al., 2020b, Rosset et al., 2003], but subsampling leads to a different solution that could depend less on the shortcut. For the latter however, models are more prone to overfitting on the smaller subsampled dataset. Similar observations were made in Sagawa et al. [2020b] but this work extends the insight to tasks with perfect stable features. Comparing ERM and MARG -CTRL on subsampled data would be fruitful.

Any exponentially tailed loss when minimized via gradient descent converges to the max-margin solution in direction [Soudry et al., 2018]. Thus, theorem 1 characterizes shortcut learning for any exponentially-tailed loss. However, losses with decreasing polynomial tails - for example, ℓ ( a ) = 1 1+ a K for some K > 0 -donot converge to the max-margin classifier. One future direction is to show shortcut-dependence results like theorem 1 for polynomial-tailed losses, which in turn would mean that all common classification losses with a decreasing tail impose inductive biases unsuitable for perception tasks.

In the tasks we consider with perfect stable features, Bayes-optimal predictors rely only on the stable feature. A weaker independence condition implies the same property of Bayes-optimal predictors even when y is not determined by s ( x ) : y | = ( x , z ) | s ( x ) . For example, in the CivilComments dataset a few instances have ambiguous labels [Xenos et al., 2022] meaning that there may not be a perfect stable feature. Studying uniform margins and other inductive biases under this independence would be fruitful.

Acknowledgements. This work was supported by the NIH/NHLBI Award R01HL148248, NSF Award 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science, NSF CAREER Award 2145542, ONR N00014-23-1-2634, Google. Aahlad Puli is partly supported by the Apple Scholars in AI/ML PhD fellowship.

References

Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, 2020.

Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran Haque, Sara M Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. Wilds: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 5637-5664. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/koh21a.html .

Aahlad Manas Puli, Lily H Zhang, Eric Karl Oermann, and Rajesh Ranganath. Out-of-distribution generalization in the presence of nuisance-induced spurious correlations. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id= 12RoR2o32T .

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations , 2020a. URL https: //openreview.net/forum?id=ryxGuJrFvS .

Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning , pages 6781-6792. PMLR, 2021.

Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Re. Correct-ncontrast: a contrastive approach for improving robustness to spurious correlations. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 26484-26516. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/zhang22z.html .

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research , 19 (1):2822-2878, 2018.

Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Bernhard Schölkopf, Caroline Uhler, and Kun Zhang, editors, Proceedings of the First Conference on Causal Learning and Reasoning , volume 177 of Proceedings of Machine Learning Research , pages 336-351. PMLR, 11-13 Apr 2022. URL https://proceedings.mlr.press/v177/idrissi22a.html .

Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning , pages 10849-10858. PMLR, 2021.

Bohan Wang, Qi Meng, Huishuai Zhang, Ruoyu Sun, Wei Chen, Zhi-Ming Ma, and Tie-Yan Liu. Does momentum change the implicit regularization on separable data? Advances in Neural Information Processing Systems , 35:26764-26776, 2022.

Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=fSTD6NFIW_b .

Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overfitting: Interpolation and invariance are fundamentally at odds. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=dQNL7Zsta3 .

Mohammad Pezeshki, Sékou-Oumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks. In Thirty-Fifth Conference on Neural Information Processing Systems , 2021.

Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. arXiv preprint arXiv:2102.10395 , 2021.

David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). arXiv preprint arXiv:2003.00688 , 2020.

Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning , pages 2189-2200. PMLR, 2021.

Victor Veitch, Alexander D'Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance to spurious correlations: Why and how to pass stress tests. arXiv preprint arXiv:2106.00545 , 2021.

Christina Heinze-Deml and Nicolai Meinshausen. Conditional variance penalties and domain shift robustness. Machine Learning , 110(2):303-348, 2021.

Yonatan Belinkov and Yonatan Bisk. Synthetic and natural noise both break neural machine translation. arXiv preprint arXiv:1711.02173 , 2017.

Greg Yang and Hadi Salman. A fine-grained spectral perspective on neural networks. arXiv preprint arXiv:1907.10599 , 2019.

Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate of neural networks for learned functions of different frequencies. Advances in Neural Information Processing Systems , 32, 2019.

Jason Jo and Yoshua Bengio. Measuring the tendency of cnns to learn surface statistical regularities. arXiv preprint arXiv:1711.11561 , 2017.

Nicholas Baker, Hongjing Lu, Gennady Erlikhman, and Philip J Kellman. Deep convolutional networks do not classify based on global object shape. PLoS computational biology , 14(12): e1006613, 2018.

Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 , 2013.

Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Madhu S Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks , 132:428-446, 2020.

Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. arXiv preprint arXiv:2006.07710 , 2020.

Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. arXiv preprint arXiv:1805.08522 , 2018.

Katherine L Hermann and Andrew K Lampinen. What shapes feature representations? exploring datasets, architectures, and training. arXiv preprint arXiv:2006.12433 , 2020.

Luca Scimeca, Seong Joon Oh, Sanghyuk Chun, Michael Poli, and Sangdoo Yun. Which shortcut cues will dnns choose? a study from the parameter-space perspective. arXiv preprint arXiv:2110.03095 , 2021.

Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-ofdistribution generalization. Advances in Neural Information Processing Systems , 34:3438-3450, 2021.

Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning , pages 8346-8356. PMLR, 2020b.

Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Mohammad Pezeshki, Amartya Mitra, Yoshua Bengio, and Guillaume Lajoie. Multi-scale feature learning dynamics: Insights for double descent. In International Conference on Machine Learning , pages 17669-17690. PMLR, 2022.

Weiyang Liu, Yandong Wen, Zhiding Yu, Ming Li, Bhiksha Raj, and Le Song. Sphereface: Deep hypersphere embedding for face recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 212-220, 2017.

Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-distribution-aware margin loss. Advances in neural information processing systems , 32, 2019.

Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis. Labelimbalanced and group-sensitive classification under overparameterization. Advances in Neural Information Processing Systems , 34:18970-18983, 2021.

Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: De-biasing classifier from biased classifier. Advances in Neural Information Processing Systems , 33:20673-20684, 2020.

Ziwei Ji, Miroslav Dudík, Robert E Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Conference on Learning Theory , pages 2109-2136. PMLR, 2020.

Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion proceedings of the 2019 world wide web conference , pages 491-500, 2019.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision , pages 3730-3738, 2015.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015. URL http://arxiv.org/abs/1412.6980 .

Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. Advances in neural information processing systems , 16, 2003.

Alexandros Xenos, John Pavlopoulos, Ion Androutsopoulos, Lucas Dixon, Jeffrey Sorensen, and Léo Laugier. Toxicity detection sensitive to conversational context. First Monday , 2022.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Xinlong Feng and Zhinan Zhang. The rank of a random matrix. Applied mathematics and computation , 185(1):689-694, 2007.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id= Bkg6RiCqY7 .

Nimit Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. Advances in Neural Information Processing Systems , 33:19339-19352, 2020.

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id= lQdXeXDoWtI .

A Appendix: Proof of Theorem 1, Corollary 1, and Theorem 2

A.1 Helper Lemmas

A.1.1 Bounding norms and inner products of isotropic random vectors.

The main lemmas of this section are lemma 3 and lemma 4. We will then use these two to bound norms of sums of random vectors and inner products between the sum and a single random vector in lemma 5. We first list some facts from [Vershynin, 2018] that we will use to bound the probability with which norms and inner products of Gaussian random vectors deviate far from their mean.

Definition 1. (Sub-Gaussian norm) For an r.v. x , the sub-Gaussian norm, or ψ 2 -norm, is

An r.v. is called sub-Gaussian if its ψ 2 -norm is finite and for some fixed constant c

A Gaussian r.v. x ∼ N (0 , σ 2 ) has an ψ 2 -norm of G σ for a constant G = 󰁴 8 3 . 2

Definition 2. (Sub-exponential norm) For an r.v. x , the sub-exponential norm, or ψ 1 -norm, is

A sub-exponential r.v. is one that has finite ψ 1 -norm.

Lemma 1. (Lemma 2.7.7 from [Vershynin, 2018]) Products of sub-Gaussian random variables x , y is a sub-exponential random variable with it's ψ 1 -norm bounded by the product of the ψ 2 -norm

Lemma 1 implies that the product of two mean-zero standard normal vectors is a sub-exponential random variable with ψ 1 -norm less than G 2 .

Lemma 2. (Bernstein inequality, Theorem 2.8.2 [Vershynin, 2018]) For i.i.d sub-exponential random variables x 1 , · · · , x d , for a fixed constant c = 1 (2 e ) 2 and K = 󰀂 x 1 󰀂 ψ 1

Next, we apply these facts to bound the sizes of inner products between two unit-variance Gaussian vectors.

Lemma 3. (Bounds on inner products of Gaussian vectors) Let u , v be d -dimensional random vectors where each coordinate is an i.i.d standard normal r.v. Then, for any scalar 󰂃 > 0 such that 󰂃 ≤ G 2 √ d , for a fixed constant c = 1 (2 e ) 2

Proof. First, the inner product is u ⊤ v = 󰁓 d i u i v i ; it is the sum of products of i.i.d. standard normal r.v. ( σ = 1 ). Then, by lemma 1, each term in the sum is a sub-exponential r.v. with ψ 1 -norm bounded as follows:

We can apply Bernstein inequality lemma 2 to sub-exponential r.v. to the inner product and then upper bound the probability by replacing K with the larger G 2 in eq. (6)

Substituting t = 󰂃 √ d in the above gives us:

Using the fact that 󰂃 ≤ G 2 √ d to achieve the minimum concludes the proof:

Lemma 4. Let x be a Gaussian vector of size d where each element is a standard normal, meaning that 󰀂 x i 󰀂 ψ 2 = G . Then, for any t > 0 and a fixed constant c = 1 (2 e ) 2 , the norm of the vector concentrates around √ d according to

Proof. Equation 3.3 from the proof of theorem 3.1.1 in [Vershynin, 2018] shows that

As x has i.i.d standard normal entries, max i 󰀂 x i 󰀂 ψ 2 = G , concluding the proof.

A.1.2 Concentration of norms of sums of random vectors and their inner products

This is the main lemma that we will use in proving theorem 1.

Lemma 5. Consider a set of vectors V = { δ i } where δ i ∈ R d of size T V ≥ 1 where each element of each vector is drawn independently from the standard normal distribution N (0 , 1) . Then, for a fixed constant c = 1 (2 e ) 2 and any 󰂃 ∈ (0 , G 2 √ d ) with probability ≥ 1 -2 exp( -󰂃 2 c G 4 )

and with probability ≥ 1 -4 T V exp( -󰂃 2 c G 4 )

Further, consider any set U of vectors U = { δ i } of size T U , where each vector also has coordinates drawn i.i.d from the standard normal distribution N (0 , 1) . Then, with probability ≥ 1 -2 T u exp( -󰂃 2 c G 4 )

By union bound, the three events above hold at once with a probability at least 1 -2(2 T V + T u + 1) exp( -󰂃 2 c G 4 ) .

Proof. We split the proof into three parts one each for eqs. (7) to (9).

Proof of eq. (7). As δ is a vector of random i.i.d standard normal random variables, note that 1 √ T V 󰁓 i δ i is also a vector of i.i.d standard normal random variables. This follows from the fact that the sum of T V standard normal random variables is a mean-zero Gaussian random variable with standard deviation √ T V . Thus dividing by the standard deviation makes the variance 1 , making it standard normal.

Then, applying lemma 4 with t = 󰂃 gives us the following bound:

Proof of eq. (8) We split the inner product into two cases: T V = 1 and T V ≥ 2 .

Case T V = 1 . First note that due to lemma 4,

∀ j ∈ V, p 󰀓 󰀂 δ j 󰀂 < √ d -󰂃 󰀔 ≤ p 󰀓󰀏 󰀏 󰀏 󰀂 δ j 󰀂 -√ d 󰀏 󰀏 󰀏 > 󰂃 󰀔 ≤ 2 exp( -c 󰂃 2 / G 4 ) . Then, the following lower bound holds with probability at least 1 -2 exp( -c 󰂃 2 / G 4 )

To summarize this case, with the fact that 1 -2 exp( -c 󰂃 2 / G 4 ) ≥ 1 -4 T V exp( -c 󰂃 2 / G 4 ) , we have that

with probability at least 1 -4 T V exp( -c 󰂃 2 / G 4 ) .

Case T V ≥ 2 . First note that,

∕

For each of the T V different δ j 's, using lemma 4 bounds the probability of the norm 󰀂 δ j 󰀂 being larger than √ d -󰂃 :

In the case where T V ≥ 2 , we express the inner product of a vector and a sum of vectors as follows

∕

∕

and noting that like above, 1 √ T V -1 󰁓 i ∈ V,i = j δ i is a vector of standard normal random variables, we apply lemma 3 to get

∕

Putting these together, by union bound over V

∕

∕

∕

Thus, with probability at least 1 -4 T V exp 󰀓 -c 󰂃 2 G 4 󰀔 , none of the events happen and

∕

Thus, putting the analysis in the two cases together, as long as T V ≥ 1

Proof of eq. (9) Next, we apply lemma 3 again to the inner product of two vectors of i.i.d standard normal random variables:

with probability at least 1 -4 T V exp 󰀓 -c 󰂃 2 G 4 󰀔 .

By union bound over U

Thus, with probability at least 1 -2 T u exp 󰀓 -c 󰂃 2 G 4 󰀔 , the following holds, concluding the proof

Lemma 6. Let { x i , y i } i ≤ n be a collection of d dimensional covariates x i and label y i sampled according to p ρ in eq. (1) . The covariates x i = [ ± B y i , y i δ i ] , where + B in the middle coordinate for i ∈ S shortcut and -B for i ∈ S leftover . The dual formulation of the following norm-minimization problem

is the following with ζ ⊤ = [ B, 1 , 0 d -2 ]

where X is a matrix with y i x i as its rows.

Proof. We use Lagrange multipliers λ ∈ R n , ν ∈ R to absorb the constraints and then use strong duality. Letting ζ ⊤ = [ -B, 1 , 0 d -2 ] , X be a matrix where the i th row is x i y i ,

Now, we solve the inner minimization to write the dual problem only in terms of λ , ν . Solving the inner minimization involves solving a quadratic program, which is done by setting its gradient to zero,

Substituting w = 1 ( ζν + X ⊤ λ ) in eq. (11)

A.2 Shortcut learning in max-margin classification

We repeat the DGP from the linear perception task in eq. (1) here.

Theorem 1. Let w ∗ be the max-margin predictor on n training samples from eq. (12) with a leftover group of size k . There exist constants C 1 , C 2 , N 0 > 0 such that

with probability at least 1 -1 3 n over draws of the training data, it holds that B w ∗ z > w ∗ y .

Before giving the proof of theorem 1, we first give the corollary showing overparameterization is not necessary for theorem 1 to hold.

Corollary 1. For all n > N 0 -where the constant N 0 is from theorem 1 - with scalar τ ∈ (0 , 1) such that the dimension d = τ n < n , theorem 1 holds.

a linear model trained via defaultERM yields a predictor w ∗ such that B w ∗ z > w ∗ y .

Proof. We show that for a range of k , for all n ≥ N 0 theorem 1 holds for some d < n . Note that theorem 1 holds for n ≥ N 0 , d = C 1 k log(3 n ) and

Setting d ≤ τ n for some τ ∈ (0 , 1) such that d < n means that theorem 1 holds if

Absorbing this new upper bound into the requirements on k for theorem 1 to hold, we get that for any scalar n > N 0 , τ ∈ (0 , 1) , d = τ n , theorem 1 holds for

In turn, even though d < n , a linear model trained via defaultERM converges in direction to a max-margin classifier such that w ∗ with B w ∗ z > w ∗ y .

Proof. (of theorem 1) We consider two norm-minimization problems over w , one under constraint w y ≥ B w z and another under w y < B w z . We show that the latter achieves lower norm and therefore, max-margin will achieve solutions w y < B w z . The two minimization problems are as follows:

From eq. (16), any w that satisfy the constraints of the dual maximization problem will lower bound the value of the optimum of the primal, 󰀂 w stable 󰀂 2 ≥ W stable . From the eq. (17), substituting a guess in w shortcut that satisfies the constraints yields an upper bound, 󰀂 w shortcut 󰀂 2 ≤ W shortcut . The actual computation of the bounds W shortcut , W stable is in lemmas 7 and 8 which are proved in appendix A.3 and appendix A.4 respectively. We reproduce the lemmas here for convenience.

Lemma. (7) Consider the following optimization problem from eq. (16) where n samples of x i , y i come from eq. (1) where x i ∈ R d :

Let k = | S leftover | > 1 . Then, for a fixed constant c = 1 (2 e ) 2 , with any scalar 󰂃 < √ d , with probability at least 1 -2 exp( -c 󰂃 2 / G 4 ) and ∀ integers M ∈ 󰀅 1 , ⌊ n 2 k ⌋ 󰀆 ,

Lemma. (8) Consider the following optimization problem from eq. (16) where n samples of x i , y i come from eq. (1) where x i ∈ R d :

Let k = | S leftover | ≥ 1 . Then, for a fixed constant c = 1 (2 e ) 2 , with any scalar 󰂃 < 1 3 󰁴 d k < √ d , with probability at least 1 -2(2 k +( n -k ) + 1) exp( -c 󰂃 2 G 4 ) , for γ = 2 d -4 󰂃 √ kd ,

Together, the lemmas say that for any ∀ integers M ∈ 󰀅 1 , ⌊ n 2 k ⌋ 󰀆 and 󰂃 < 1 3 󰁴 d k , with probability ≥ 1 -2 exp( -c 󰂃 2 / G 4 )

and with probability at least 1 -2(2 k +( n -k ) + 1) exp( -c 󰂃 2 G 4 ) , for γ = 2 d -4 󰂃 √ kd > 0 ,

First, we choose 󰂃 2 = 2 G 4 c log(3 n ) . This gives us the probability with which these bounds hold: as k < 0 . 1 n we have k +2 < n 2 and

Next, we will instantiate the parameter M and set the constants C 1 , C 2 and the upper bound on k in theorem 1 to guarantee the following eq. (separation inequality):

which then implies that 󰀂 w shortcut 󰀂 2 < 󰀂 w stable 󰀂 2 , concluding the proof.

Invoking the conditions in theorem 1 and setting the upper bound on k . We will keep the 󰂃 as is for simplicity of reading but invoke the inequalities satisfied by log(3 n ) from theorem 1:

Now we let C 1 = 2 G 4 cC 2 for a constant C ∈ 󰀃 0 , 1 3 󰀄 3 , such that

We next find a C ∈ 󰀃 0 , 1 3 󰀄 such that eq. (separation inequality) holds with M = 5 , which upper bounds k :

Simplifying W shortcut and W stable . To actually show W shortcut < W stable in eq. (separation inequality), we compare a simplified strict upper bound on the LHS W shortcut and a simplified strict lower bound on the RHS W stable

For the simplification of the RHS W stable of eq. (separation inequality), we will use the fact that d ≥ 2 G 4 cC 2 log(3 n ) k . Given the assumption n > N 0 , choosing N 0 to be an integer such that log(3 N 0 ) ≥ 40 cC 2 G 4 means that log(3 n ) > 40 cC 2 G 4 and we have

which gives us, for M = 5 ,

Now, we produce a simpler upper bound on the first part of the LHS of eq. (separation inequality): recalling that γ = 2 d -4 󰂃 √ kd , and substituting in the upper bounds on 󰂃 ,

3 The 1 3 comes from requiring that 󰂃 < 1 3 󰁴 d k from lemma 8.

Next is a simpler upper bound on the second part of the LHS of eq. (separation inequality). Again with γ = 2 d -4 󰂃 √ kd ,

Now setting gives the lower bound on B from theorem 1:

Formally,

By combining the upper bound from eq. (30) and the upper bound from eq. (29), we get an upper bound on the whole of the LHS of eq. (separation inequality), which in turn provides an upper bound on W shortcut :

because k ≥ 1 . Note the upper bound is strict.

Concluding the proof. Now, we show that a C exists such that the following holds, which implies W shortcut < W stable , which in turn implies eq. (separation inequality) and the proof concludes:

The above inequality holds when

The right hand side holds when the quadratic equation (1 + C ) 2 -󰁴 10 9 (1 -4 C ) is non-positive, which holds between the roots of the equation. The equation's positive solution is

Setting C to this quantity satisfies the requirement that C ∈ (0 , 1 3 ) .

Thus, a C exists such that eq. (separation inequality) holds which concludes the proof of theorem 1 for the following constants and constraints implied by C and M = 5 :

where G is the ψ 2 -norm of a standard normal r.v. and c is the absolute constant from the Bernstein inequality in lemma 2.

A.3 Lower bounding the norm of solutions that rely more on the stable feature

Lemma 7. Consider the following optimization problem from eq. (16) where n samples of x i , y i come from eq. (1) where x i ∈ R d :

Let k = | S leftover | > 1 . Then, for a fixed constant c = 1 (2 e ) 2 , with any scalar 󰂃 < √ d , with probability at least 1 -2 exp( -c 󰂃 2 / G 4 ) and ∀ integers M ∈ 󰀅 1 , ⌊ n 2 k ⌋ 󰀆 ,

Proof. By lemma 6, the dual of eq. (16) is the following for ζ = [ -B, 1 , 0 d -2 ] and X is an n × d matrix with rows y i x i :

Now by duality which means any feasible candidate to eq. (32) gives a lower bound on 󰀂 w stable 󰀂 2 .

Feasible Candidates for λ , ν . We now define a set U ⊂ [ n ] , and let λ i = α | U | > 0 for i ∈ U and 0 otherwise. For M ∈ (1 , ⌊ n 2 k ⌋ ] , we take 2 Mk samples from the training data to be included in U . Formally,

U = S leftover ∪ (2 M -1) k a random samples from S shortcut , which gives the size | U | = 2 Mk . Then, we let ν = α 2( M -1) 2 M > 0 .

Note that for the above choice of λ , X ⊤ λ is a sum of the rows from U scaled by α | U | . Adding up k rows from S leftover and k rows from S shortcut cancels out the B s and, so in the B is accumulated | U | -2 k = 2( M -1) k times, and so

As λ has α | U | on | U | elements and 0 otherwise, λ ⊤ 1 = α

As we set ν = α 2( M -1) 2 M ,

For the chosen values of ν , λ the value of the objective in eq. (32) is

Letting

the objective is of the form α -α 2 Γ 4 . To maximize with respect to α , setting the derivative of the objective w.r.t α to 0 gives:

This immediately gives us

and we lower bound this quantity by upper bounding Γ .

By concentration of gaussian norm as in lemma 4, with probability at least 1 -2 exp( -c 󰂃 2 G 4 )

In turn, recalling that | U | = 2 Mk

The upper bound on Γ gives the following lower bound on 󰀂 w stable 󰀂 2 :

A.4 Upper bounding the norm of solutions that rely more on the shortcut.

Lemma 8. Consider the following optimization problem from eq. (16) where n samples of x i , y i come from eq. (1) where x i ∈ R d :

Let k = | S leftover | ≥ 1 . Then, for a fixed constant c = 1 (2 e ) 2 , with any scalar 󰂃 < 1 3 󰁴 d k < √ d , with probability at least 1 -2(2 k +( n -k ) + 1) exp( -c 󰂃 2 G 4 ) , for γ = 2 d -4 󰂃 √ kd ,

Proof. Let k = | S leftover | . The candidate we will evaluate the objective for is

High-probability bounds on the margin achieved by the candidate and norm of w The margins on the shortcut group and the leftover group along with the constraints are as follows:

Due to the standard normal distribution being isotropic, and y j ∈ { -1 , 1 } , y j δ j has the same distribution as δ j . Then, we apply lemma 5 with V = S leftover , U = S shortcut -which means T v = k and T u = ( n -k ) -to bound the margin terms in eq. (40) and 󰀂 w 󰀂 2 with probability at least

Applying the bound in eq. (9) in lemma 5 between a sum of vectors and a different i.i.d vector,

Applying the bound in eq. (8) from lemma 5

The margin constraints on the shortcut and leftover from eq. (40) respectively imply

We choose β = 1 + γ󰂃 √ dk , which implies an inequality that γ has to satisfy the following, which is due to d -3 󰂃 √ kd > 0 ,

Now, we choose

Computing the upper bound on the value of the objective in the primal problem in eq. (17) The feasible candidate's norm 󰀂 w 󰀂 2 is an upper bound on the solution's norm 󰀂 w shortcut 󰀂 2 and so

By lemma 5 which we invoked,

To conclude the proof, substitute β = 1 + γ󰂃 √ dk and get the following upper bound with γ = 2 d -3 󰂃 √ kd :

A.5 Concentration of k and intuition behind theorem 1

Concentration of k around (1 -ρ ) n . Denote the event that the i th sample lies in the leftover group as I i : then E [ I i ] = 1 -ρ and the leftover group size is k = 󰁓 i I i . Hoeffding's inequality (Theorem 2.2.6 in [Vershynin, 2018]) shows that for any t > 0 , k is at most (1 -ρ ) n + t √ n with probability at least 1 -exp( -2 t 2 ) :

Letting ρ = 0 . 9 + 󰁴 log 3 n n and t = √ log 3 n , gives us

To connect ρ to shortcut learning due to max-margin classification, we take a union bound of the event that k < 0 . 1 n , which occurs with probability at least 1 -1 3 n and theorem 1 which occurs with probability at least 1 -1 3 n . This union bound guarantees that with probability at least 1 -2 3 n over sampling the training data, max-margin classification on n training samples from eq. (1) relies more on the shortcut feature if ρ is above a threshold; and this threshold converges to 0 . 9 at the rate of 󰁳 log 3 n / n .

A.6 Bumpy losses improve ERM in the under-parameterized setting

Theorem 2. Consider n samples of training data from DGP in eq. (1) with d < n . Consider a linear classifier f θ ( x ) = w ⊤ x such that for all samples in the training data y i w ⊤ x i = b for any b ∈ (0 , ∞ ) . With probability 1 over draws of samples, w = [0 , b, 0 d -2 ] .

Proof. Letting X be the matrix where each row is y i x i , the theorem statement says the solution w ∗

First, split w ∗ = [ w ∗ z , w ∗ y , w ∗ -y ] . Equation (43) says that the margin of the model on any sample satisfies

We collect these equations for the whole training data by splitting X into columns: denoting Y, Z as vectors of y i and z i and using · to denote element wise operation, split X into columns that correspond to y , z and δ respectively as X = [ Y · Y | Y · Z | X δ ] . Rearranging terms gives us

The elements of Y · Z lie in { -1 , 1 } and, as the shortcut feature does not always equal the label, the elements of Y · Z are not all the same sign.

Solutions do not exist when one non-zero element exists in ( b -w ∗ y ) 1 -w ∗ z Y · Z By definition of w ∗

Denote r = ( b -w ∗ y ) 1 -w ∗ z Y · Z. and A = X δ . Now we show that w.p. 1 solutions do not exist for the following system of linear equations:

First, note that A = X δ has y i δ i for rows and as y i | = δ i and y i ∈ { -1 , 1 } , each vector y i δ i is distributed identically to a vector of independent standard Gaussian random variables. Thus, A is a matrix of IID standard Gaussian random variables.

Let U denote D -2 indices such that the corresponding rows of A form a matrix D -2 × D -2 matrix and r U has at least one non-zero element; let A U denote the resulting matrix. Now A U is a D -2 × D -2 sized matrix where each element is a standard Gaussian random variable. Such matrices have rank D -2 with probability 1 because square singular matrices form a measure zero set under the Lebesgue measure over R D -2 × D -2 [Feng and Zhang, 2007].

We use subscript · -U to denote all but the indices in U . The equation Aw = r implies the following two equations:

As A U is has full rank ( D -2 ), A U w = r U admits a unique solution w ∗ U = 0 -because r U has at least one non-zero element by construction. Then, it must hold that

∕

For any row v ⊤ ∈ A -U , eq. (44) implies that v ⊤ w ∗ equals a fixed constant. As v is a vector of i.i.d standard normal random variables, v ⊤ w ∗ is a gaussian random variable with mean 󰁓 ( w ∗ i ) and variance 󰀂 w ∗ 󰀂 2 . Then with probability 1 , v ⊤ w ∗ will not equal a constant. Thus, w.p.1 A -U w ∗ U = r -U is not satisfied, which means w.p.1 there are no solutions to A w = r .

Case where ( b -w ∗ y ) 1 -w ∗ z Y · Z is zero element-wise As X has rank D -2 , X δ w ∗ δ = 0 only when w ∗ δ = 0 .

Each element in ( b -w ∗ y ) 1 -w ∗ z Y · Z is either b -w ∗ y + w ∗ z or b -w ∗ y -w ∗ z . Thus,

Adding and subtracting the two equations on the right gives

2(

-

B Appendix: further experimental details and results

B.1 DefaultERM with ℓ 2 -regularization.

In section 3, we show defaultERM achieves zero training loss by using the shortcut to classify the shortcut group and noise to classify the leftover group, meaning the leftover group is overfit. The usual way to mitigate overfitting is via ℓ 2 -regularization, which, one can posit, may encourage models to rely on the perfect stable feature instead of the imperfect shortcut and noise.

We train the linear model from section 3 with defaultERM and ℓ 2 -regularization - implemented as weight decay in the AdamW optimizer [Loshchilov and Hutter, 2019] - on data from eq. (1) with

b

w

∗

y

) = 0

and

2

w

∗

z

= 0

.

d = 800 , B = 10 , n = 1000 . Figure 4 plots accuracy and losses for the ℓ 2 -regularized defaultERM with the penalty coefficient set to 10 -8 ; it shows that ℓ 2 -regularization leads defaultERM to build models that only achieve ≈ 50% test accuracy.

For smaller penalty coefficients, defaultERM performs similar to how it does without regularization, and for larger ones, the test accuracy gets worse than defaultERM without regularization. We give an intuitive reason for why larger ℓ 2 penalties may lead to larger reliance on the shortcut feature. Due to the scaling factor B = 10 in the synthetic experiment, for a fixed norm budget, the model achieves lower loss when using the shortcut and noise compared to using the stable feature. In turn, heavy ℓ 2 -regularization forces the model to rely more on the shortcut to avoid the cost of larger weight needed by the model to rely on the stable feature and the noise.

B.2 Margin control ( MARG -CTRL )

In fig. 5, we plot the different MARG -CTRL losses along with log-loss. Each MARG -CTRL loss has a "bump" which characterizes the loss function's transition from a decreasing function of the margin to an increasing one. These bumps push models to have uniform margins because the loss function's derivative after the bump is negative which discourages large margins. The hyperparameters - like temperature in σ -damp or function output target in MARG -LOG -affect the location of the bump and the slopes of the function on either side of the bump.

B.3 M ARG -CTRL on a linear model

In fig. 6, we compare defaultERM to σ -stitch. In fig. 7 and fig. 8, compare SD and MARG -LOG respectively to defaultERM . The left panel of all figures shows that MARG -CTRL achieves better test accuracy than defaultERM , while the right most panel shows that the test loss is better on the leftover group using MARG -CTRL . Finally, the middle panel shows the effect of controlling margins

in training; namely, the margins on the training data do not go to ∞ , evidenced by the training loss being bounded away from 0 . Depending on the shortcut feature leads to different margins and therefore test losses between the shortcut and leftover groups; the right panel in each plot shows that the the test losses on both groups reach similar values, meaning MARG -CTRL mitigates dependence on the shortcut. While defaultERM fails to perform better than chance ( 50% ) even after 100 , 000 epochs (see fig. 1), MARG -CTRL mitigates shortcut learning within 5000 epochs and achieves 100% test accuracy.

B.4 MARG -CTRL vs. defaultERM with a neural network

With d = 100 and B = 10 in eq. (1), we train a two layer neural network on 3000 samples from the training distribution. The two layer neural network has a 200 unit hidden layer that outputs a scalar. Figure 9 shows that a neural network trained via defaultERM fails to cross 50% test accuracy even after 40 , 000 epochs, while achieving less than 10 -10 in training loss.

(a)

Average accuracy and loss curves.

100

100

80

10-3

60

train

test

10-6

40

10-9

20

17

?

0

1

2

3

4

0

1

2

3

4

1e4

1e4

100

2.0

75

Tr, Y = Z

1.5

Tr, Y = Z

Te, Y=Z

Te, Y=Z

50

一

Tr,Y≠Z

1.0

Tr,Y≠Z

25

Te,YZ

0.5

Te,YZ

0

0.0

0

1000

2000

3000

4000

0

1000

2000

3000

4000

epochs

epochs

(b)

Accuracy and loss on shortcut and leftover groups.

In fig. 10, we compare defaultERM to σ -stitch. In fig. 12 and fig. 13, compare SD and MARG -LOG respectively to defaultERM . The left panel of all figures shows that MARG -CTRL achieves better test accuracy than defaultERM , while the right most panel shows that the test loss is better on the leftover group using MARG -CTRL . Finally, the middle panel shows the effect of controlling margins in training; namely, the margins on the training data do not go to ∞ , evidenced by the training loss being bounded away from 0 .

B.5 Spectral decoupling for a linear model on the linear DGP in eq. (1).

We first show that a linear classifier trained with SD achieves 100% test accuracy while defaultERM performs worse than chance on the test data; so, SD builds models with more dependence on the stable perfect feature, compared to Empirical Risk minimization ( ERM ). Next, we outline the assumptions for the gradient starvation (GS) regime from Pezeshki et al. [2021] and then instantiate it for a linear model under the data generating process in eq. (1), showing that the assumptions for the GS-regime are violated.

Figure 7 shows the results of training a linear model with SD on training data of size 1000 sampled as per eq. (1) from p ρ =0 . 9 with d = 300 ; the test data also has a 1000 samples but comes from p ρ =0 . 1 .

Figure 7 shows that SD builds models with improved dependence on the perfect stable feature, as compared to ERM , to achieve 100% test accuracy.

B.5.1 The linear example in Equation (1) violates the gradient starvation regime.

Background on Pezeshki et al. [2021]. With the aim of explaining why ERM -trained neural networks depend more on one feature over a more informative one, Pezeshki et al. [2021] derive solutions to ℓ 2 -regularized logistic regression in the NTK ; they let the regularization coefficient be small enough for the regularized solution to be similar in direction to the unregularized solution. Given n samples y i , x i , let Y be a diagonal matrix with the labels on its diagonal, X be a matrix with x i as its rows, and ˆ y ( X , θ ) = f θ ( X ) be the n -dimensional vector of function outputs where each element is ˆ y i = f θ ( x i ) . In gradient-based training in the NTK regime, the vector of function outputs of the network with parameters θ can be approximated as ˆ y = Φ 0 θ , where Φ 0 is the neural-tangentrandom-feature ( NTRF ) matrix at initialization:

To define the features, the strength (margin) of each feature, and how features appear in each sample, Pezeshki et al. [2021] compute the singular value decomposition ( SVD ) of the NTRF Φ 0 multiplied by the diagonal-label matrix Y :

The rows of V are features, the diagonal elements of S are the strengths of each feature and the i th row of U denotes how each feature appears in the NTRF representation of the i th sample.

To study issues with the solution to ℓ 2 -regularized logistic regression, Pezeshki et al. [2021] define the gradient starvation (GS) regime. Under the GS regime, they assume U is a perturbed identity matrix that is also unitary: for a small constant δ << 1 , such a matrix has all diagonal elements √ 1 -δ 2 and the rest of the elements are of the order δ such that the rows have unit ℓ 2 -norm.

The GS regime is violated in eq. (1). When f θ is linear, f θ ( x ) = θ ⊤ x , the NTRF matrix is

In this case, let us look at an implication of U being a perturbed identity matrix that is also unitary, as Pezeshki et al. [2021] assume. With ( u i ) ⊤ as the i th row of U , the transpose of i th sample can be written as ( x i ) ⊤ = ( u i ) ⊤ SV . Pezeshki et al. [2021] assume that δ << 1 in that the off-diagonal terms of U are small perturbations such that off-diagonal terms of U ( S 2 + λ I ) U ⊤ have magnitude much smaller than 1 , meaning that the terms | ( u i ) ⊤ S 2 ( u j ) + λ | << 1 for i = j and positive and small λ << 1 .

Then,

∕

In words, this means that any two samples x i , x j are nearly orthogonal. Now, for samples from eq. (1), for any i, j such that z j = z i and y i = y j ,

As δ are isotropic Gaussian vectors, around half the pairs i, j will have ( δ i ) ⊤ δ j > 0 meaning 󰀏 󰀏 ( x i ) ⊤ x j 󰀏 󰀏 > 101 . This lower bound implies that U is not a perturbed identity matrix for samples from eq. (1). This violates the setup of the gradient starvation regime from [Pezeshki et al., 2021].

Thus, the linear DGP in eq. (1) does not satisfy the conditions for the GS regime that is proposed in [Pezeshki et al., 2021]. The GS regime blames the coupled learning dynamics for the different features as the cause for defaultERM -trained models depending more on the less informative feature. Pezeshki et al. [2021] derive spectral decoupling ( SD ) to avoid coupling the training dynamics, which in turn can improve a model's dependence on the perfect feature. SD adds a penalty to the function outputs which Pezeshki et al. [2021] show decouples training dynamics for the different features as defined by the NTRF matrix:

As eq. (1) lies outside the GS regime, the success of SD on data from eq. (1) cannot be explained as a consequence of avoiding the coupled training dynamics in the GS regime Pezeshki et al. [2021]. However, looking at SD as MARG -CTRL , the success of SD , as in fig. 7, is explained as a consequence encouraging uniform margins.

An example of perfect test accuracy even with dependence on the shortcut. In fig. 14, we train a linear model with defaultERM on data from eq. (1), showing that even when shortcut dependence

acc

loss

100

100

train

W_Z

80

10-2

test

0.2

60

10-4

0.1

40

train

test

10-6

20

0.0

0

10000

20000

0

10000

20000

0

10000

20000

is non-zero, test leftover group accuracy can be 100% . Nagarajan et al. [2021] consider linearly separable data and formalize geometric properties of the data that make max-margin classifiers give non-zero weight to the shortcut feature ( w z > 0) . In their example, it is unclear when w z > 0 leads to poor accuracy in the leftover group because Nagarajan et al. [2021] do not separate the model's dependence on the stable feature from the dependence on noise. The example in fig. 14 gives an example where w z > 0 but test accuracy is 100% , demonstrating that guarantees on test leftover group error require comparing w y and w z ; the condition w z > 0 alone is insufficient. In contrast, theorem 1 characterizes cases where leftover group accuracy is worse than random even without overparameterization.

B.6 Experimental details

B.6.1 Background on Just Train Twice ( JTT ) and Correct-n-Contrast ( CNC )

JTT Liu et al. [2021] develop JTT with the aim of building models robust to subgroup shift, where the mass of disjoint subgroups of the data changes between training and test times. To work without training group annotations, JTT assumes ERM builds models with high worst-group error. With this assumption, JTT first builds an "identification" model via ERM to pick out samples that are misclassified due to model's dependence on the shortcut. Then, JTT trains a second model again via ERM on the same training data with the loss for the misclassified samples upweighted (by constant λ ). As Liu et al. [2021] point out, the number of epochs to train the identification model and the upweighting constant are hyperparameters that require tuning using group annotations. As Liu et al. [2021], Zhang et al. [2022] show that JTT and CNC outperforms LFF and other two-stage shortcutmitigating methods ([Zhang et al., 2022]), so we do not compare against them.

∕

Correct-n-Contrast ( CNC ) In a fashion similar to JTT , the first stage of CNC is to train a model with regularized ERM to predict based on spurious attributes, i.e. shortcut features. Zhang et al. [2022] develop a contrastive loss to force the model to have similar representations across samples that share a label but come from different groups (approximately inferred by the first-stage ERM model). Formally, the first-stage model is used to approximate the spurious attributes in one of two ways: 1) predict the label with the model, 2) cluster the representations into as many clusters as there are classes, and then use the cluster identity. The latter technique was first proposed in [Sohoni et al., 2020]. For an anchor sample ( y i , x i ) of label y = y , positive samples P i are those than have the same label but have the predicted spurious attribute is a different value: ˆ z = y . Negatives N i are those that have a different label but the spurious attribution is the same: ˆ z = y . For a temperature parameter τ and representation function r θ , the per-sample contrastive loss for CNC is:

The samples i are called anchors . For a scalar λ to trade off between contrastive and predictive loss, the overall per-sample loss in the second-stage in CNC is

C NC uses hyperparameters informed by dataset-specific empirical results from prior work. The original implementation of CNC from Zhang et al. [2022] uses specific values of first-stage hyperparameters like weight decay and early stopping epoch for each dataset by using empirical results from prior work [Sagawa et al., 2020a, Liu et al., 2021]. The prior work finds weight-decay and early stopping epoch which lead defaultERM models to achieve low test worst-group accuracy, implying that the model depends on the spurious attribute. This means the first-stage models built in CNC are pre-selected to pay attention to the spurious attributes. For example, [Zhang et al., 2022] point out that the first-stage model they use for Waterbirds predicts the spurious feature with an accuracy of 94 . 7% .

Without using dataset-specific empirical results from prior work, choosing LR and WD requires validating through the whole CNC procedure. We let CNC use the same LR and WD for both stages and then validate the choice using validation performance of the second-stage model. This choice of hyperparameter validation leads to a similar number of validation queries for all methods that mitigate shortcuts.

B.6.2 Training details

Variants of MARG -CTRL to handle label imbalance. The three datasets that we use in our experiments - Waterbirds, CelebA, and Civilcomments - all have an imbalanced (non-uniform) marginal distribution over the label; for each dataset,

When there is sufficiently large imbalance, restricting the margins on all samples could bias the training to reduce loss on samples in the most-frequent class first and overfit on the rest of the samples. This could force a model to predict the most frequent class for all samples, resulting in high worst-group error.

To prevent such a failure mode, we follow [Pezeshki et al., 2021] and define variants of σ -damp, MARG -LOG , and σ -stitch that have either 1) different maximum margins for different classes or 2) different per-class loss values for the same margin value. Mechanically, these variants encourage uniform margins within each class, thus encouraging the model to rely less on the shortcut feature. We give the variants here for labels taking values in { -1 , 1 } :

With per-class temperatures T -1 , T 1 > 0 the variant of σ -damp is

∕

The 1 . 278 comes in to make sure the maximum input to log-loss occurs at f θ = 1 . However, due to the different temperatures T 1 = T -1 , achieving the same margin on all samples produces lower loss on the class with the larger temperature.

With per-class temperatures T -1 , T 1 > 0 the variant of σ -stitch is

With per-class function output targets γ -1 , γ 1 > 0 the variant of MARG -LOG is

These per-class variants are only for training; at test time, the predicted label is sign ( f θ ) .

Details of the vision and language experiments. We use the same datasets from Liu et al. [2021], downloaded via the scripts in the code from [Idrissi et al., 2022]; see [Idrissi et al., 2022] for sample sizes and the group proportions. For the vision datasets, we finetune a resnet50 from Imagenet-pretrained weights and for Civilcomments, we finetune a BERT model.

Optimization details. For all methods and datasets, we tune over the following weight decay ( WD ) parameters: 10 -1 , 10 -2 , 10 -3 , 10 -4 For the vision datasets, we tune learning rate ( LR ) over 10 -4 , 10 -5 and for CivilComments, we tune over 10 -5 , 10 -6 . For CivilComments, we use the AdamW optimizer while for the vision datasets, we use the Adam optimizer; these are the standard optimizers for their respective tasks [Puli et al., 2022, Gulrajani and Lopez-Paz, 2021]. We use a batch size of 128 for both CelebA and Waterbirds, and train for 20 and 100 epochs respectively. For CivilComments we train for 10 epochs with a batch size of 16 .

Per-method Hyperparameters. Like in [Pezeshki et al., 2021], the per-class temperatures T -1 , T 1 for σ -damp and σ -stitch, and the function output targets γ -1 , γ 1 for MARG -LOG are hyperparameters that we tune using the worst-group accuracy or label-balanced average accuracy computed on the validation dataset, averaged over 2 seeds.

For σ -stitch, we select from T -1 ∈ { 1 , 2 } and T 1 ∈ { 2 , 4 , 8 , 12 } such that T 1 > T -1 .

For SD and MARG -LOG , we search over γ -1 ∈ { -1 , 0 , 1 } and γ 1 ∈ { 1 , 2 , 2 . 5 , 3 } for the image datasets and γ 1 ∈ { 1 , 2 } for the text dataset, and the penalty coefficient is set to be λ = 0 . 1

For σ -damp, we search over T -1 ∈ { 1 , 2 } and T 1 ∈ { 2 , 4 } such that T 1 > T -1 .

For JTT , we search over the following parameters: the number of epochs T ∈ { 1 , 2 } for CelebA and Civilcomments and T ∈ { 10 , 20 , 30 } for Waterbirds, and the upweighting constant λ ∈ { 20 , 50 , 100 } for the vision datasets and λ ∈ { 4 , 5 , 6 } for Civilcomments.

For CNC , we search over the same hyperparameter as [Zhang et al., 2022] : the temperature in τ ∈ { 0 . 05 , 0 . 1 } , the contrastive weight λ ∈ { 0 . 5 , 0 . 75 } , and the gradient accumulation steps s ∈ { 32 , 64 } . For the language task in Civilcomments, we also try one additional s = 128 .

B.7 M ARG -CTRL improves over defaultERM on CelebA even without the stable feature being perfect.

CelebA is a perception task in that the stable feature is the color of the hair in the image. But unlike the synthetic experiments, MARG -CTRL does not achieve a 100% test accuracy on CelebA. We investigated this and found that CelebA in fact has some label noise.

We trained a model via the MARG -CTRL method σ -damp on CelebA data with no shortcut; this data is constructed by subsampling the groups to all equal size, ( 5000 samples). This achieves a test worst-group accuracy of 89% . We visualized the images that were misclassified by this model and found that many images with blond-haired people were classified as having non-blonde hair. Figure 15 shows 56 misclassified images where samples with blonde as the true label have a white strip at the bottom while samples with non-blonde as the true label have a black strip at the bottom. The figure shows that images with blonde people can have the non-blonde label, thus demonstrating label noise. Thus, MARG -CTRL improves over ERM even on datasets like CelebA where the stable features do not determine the label.

dataset = celeba

0.9

0.8

0.7

0.6

0.5

0.4

O-damp o-stitch

erm

marg-log

sd

[method]

B.8 Sensitivity of ERM and MARG -CTRL to varying LR and WD

In fig. 16, we compare the test worst-group accuracy of defaultERM and MARG -CTRL on CelebA, for different values of LR and WD . There are 8 combinations of LR and WD for which ERM is run. For each combination of LR and WD , the hyperparameters of the MARG -CTRL method (values of λ , T, v ) are tuned using validation group annotations, and the test worst-group accuracy corresponds to the best method hyperparameters. DefaultERM 's performance changes more with LR and WD than MARG -CTRL , which shows that defaultERM is more sensitive than MARG -CTRL . Only 2 combinations of LR and WD improve ERM beyond a test worst-group accuracy of 60% , while every MARG -CTRL method achieves more than 70% test worst-group accuracy for every combination of LR and WD .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretecally and empirically showed that the inductive bias of default- ERM maximizing the margin causes shortcut learning in  a linear perception task.\nIt proposed uniform margins that leads to models that depend more on the stable than the shortcut feature and suggested loss functions encourage uniform-margin solutions.\n\nSTRENGTHS:\nThis paper analyzes shortcut learning theoretically in terms of margin maximization.\nI have not seen an analysis from this perspective before.\nIt also proposes the concept of uniform margin from theory and suggests a method to prevent shortcut learning.\n\nWEAKNESSES:\nThe theory itself has limited applicability due to the linear model, and we do not know how well it can actually be explained in general terms in actual deep learning models.\n\nQUESTIONS:\nIt seems that this paper forcus on a theoretical analysis in a situation like background shortcuts, where shortcut features make up a large portion of the data as shown in term B. \nLearning  features with a large B preferentially is intuitively reasonable in such a simple situation; however, reading the proof in the Appendix, it seems to require a very tough proof to show that. \nI would like to know where the difficulty of the proof lies, what you devised in the Appendix proof ,and where you derived the key inequality.\n\nShortcut features do not always occupy as much space as stable features as background features. For example, in the following paper, the percentage of partial input is much smaller.\n>Overinterpretation reveals image classification model pathologies Brandon Carter, Siddhartha Jain, Jonas Mueller, David Gifford\nIsn't Bz>y important for the theoretical analysis?\nIf Bz<y, what part of the proof is more difficult? Or what other assumptions are necessary?\n\nLIMITATIONS:\nThe theory itself has limited applicability due to the linear model, and we do not know how well it can actually be explained in general terms in actual deep learning models.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides an in-depth analysis of the phenomenon of \"shortened learning\" in machine learning models, especially in the context of perceptual tasks. The authors confirm that basic empirical risk minimization (ERM) methods tend to prefer models that depend on shortcut features, even when models can achieve zero loss using only stable features. They attribute this to ERM's inductive bias to maximize margins across all samples. To address this, the authors propose an alternative loss function biased inductively with a uniform margin called MARG-CTRL. The paper demonstrates that MARG-CTRL mitigates shortcut learning on multiple vision and language tasks without the use of annotations of the shortcut feature in training or even validation. They also show that MARG-CTRL performs on par or better than the more complex and costly two-step shortcut mitigation method.\n\nSTRENGTHS:\n-This paper presents a thorough analysis of the shortcut learning problem. \n-The authors propose a novel solution, MARG-CTRL, which has been shown to effectively mitigate shortcut learning in several tasks.\n-The paper is well-written and easy to understand.\n\nWEAKNESSES:\n-Although MARG-CTRL is shown to perform well across different tasks, it is not clear how MARG-CTRL perform in scenarios where the shortcut features and the stable features are highly correlated.\n-The paper compares MARG-CTRL with two-stage shortcut-mitigating methods like JTT and CNC, it would be helpful to understand the specific scenarios where MARG-CTRL outperforms these methods and where it does not.\n\nQUESTIONS:\n-How does MARG-CTRL perform in scenarios where the shortcut features and the stable features are highly correlated? Does this affect the effectiveness of MARG-CTRL?\n-What are the trade-offs involved in using MARG-CTRL versus two-stage shortcut-mitigating methods??\n-Are there specific scenarios or types of tasks for which MARG-CTRL may not be as effective?\n\nLIMITATIONS:\nPlease see the weaknesses part.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nHaving had my concerns addressed by the authors I have updated my score.\n-----------------------------------------------------------------------------------------\n\n\n\n* The paper proposes an explanation for why neural networks tend to learn spurious features over stable fratures which lead to a lower loss and better performance.\n* The main argument for this is due to the maxmargin losses induced through cross entropy, to counter this the authors propose several losses which induce a uniform margin.\n* The authors propose a nice example and demonstrate the results on vision and language datasets.\n\nSTRENGTHS:\n* The paper addresses a very important topic of the simplicity bias, and goes a significant way in addressing it. \n* The comments and explanations are clear and concise, and follow a clear narrative.\n* I find the arguments that inductive bias towards uniform margins for perception tasks to be convincing.\n* The experiments are well conducted and clearly support their statements\n\nWEAKNESSES:\n* I find the scatter approach to testing many margin based losses a little un elegant, I would rather know which is the most effective in the paper and then see the results for other attempted losses in the Appendix. But this is just me preference.\n* Is Corollary one explicitly followed up with an experiment? I would like to know if it actually exhibits this behaviour or is there just the potential to?\n* The experimental datasets are rather limited, Waterbirds, CelebA, Wilds etc. Not a negative per se, but it would be nice to see authors trying to improve the testing framework for this probelm.\n\nQUESTIONS:\n* The examples in the paper tend to be restricted to linear models, which is not the case in most perception systems, do you have any intuition if these conclusions carry through to a non-linear case? I'm aware it makes the analysis harder, and I'm not asking you to perform it. But it would be nice to hear any issues that may arrise.\n* Are there any experiments in the non over-parameterised setting? As this is discussed in the paper, I would have expected to see them in the experiments.\n* Do you think these shortcuts are a bit artifificial? For instance in the waterbirds dataset, the background should be an important feature in the decision, but not the primary feature, does this method enable predictions from multiple features? Or do you think it will just rely on one?\n\nLIMITATIONS:\nNot explicitly discussed", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores the phenomenon of models using easy-to-learn spurious features (aka, shortcuts) instead of reliable but harder-to-learn true features. They find in their theoretical that max-margin relies on the spurious feature while controlling for uniforming margin induces learning the true feature. With this insight, they try numerous loss functions that provide a bias towards uniform margins and find improved worst-group accuracy on a variety of tasks.\n\nSTRENGTHS:\n* Experiments show good results for their methods.\n\nWEAKNESSES:\n* The informal statement of Theorem 1 in the main text is too vague in my opinion. In particular, it's not stated what the regime of $n$ is and how these other quantities interact with it. \n* The graphs are tough to read (too small) and the legends overlap with the plots.\n\nQUESTIONS:\n* How does the dimensionality (e.g., choosing $d=100$) come into play? The story setting up the analogy does not in principle rely on the dimensionality all, just the gradient sizes of the noisy and true signals. I'm guessing it has something to do with interpolation, but I think it'd be good to have further discussion of this.\n* What's the connection between the neural network trained in section 2.1 and the linear model theorems? What happens if you train a logistic regression model instead?\n* Similar toy settings as in this paper (such as (Sagawa et al., 2020)) have been examined in the past and seemed to rely on overparameterization for their phenomena. What are the particulars in this paper's setup that allow us to observe phenomena in the underparameterized setting?\n* How should I think of the different MARG-CTRL? What factors cause their difference in performance and how should I make the decision in which to pick?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}