SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Don’t blame Dataset Shift! Shortcut Learning due to Gradients and Cross Entropy
Aahlad Puli $^{1*}$  Lily Zhang $^{2}$  Yoav Wald $^{2}$  Rajesh Ranganath $^{1,2,3}$
$^{1}$ Department of Computer Science, New York University
$^{2}$ Center for Data Science, New York University
$^{3}$ Department of Population Health, Langone Health, New York University
Common explanations 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 default-ERM, 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 still exhibits shortcut learning. Why are such solutions preferred when the loss for default-ERM can be driven to zero using the stable feature alone? By studying a linear perception task, we show that default-ERM'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 default-ERM'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). MARG-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.
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 default-ERM.
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
*Corresponding email: aahlad@nyu.edu
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. Default-ERM can learn this Bayes-optimal classifier which, by depending solely on the stable feature, also generalizes outside the training distribution. But in practice, default-ERM 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 default-ERM 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 default-ERM 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 default-ERM-trained linear models classify in the same way as the minimum  $\ell_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 default-ERM's implicit inductive bias toward the max-margin solution is harmful in that default-ERM-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(\mathbf{x})$  determines the label  $\mathbf{y}$  via a function  $d$ ,  $\mathbf{y} = d \circ s(\mathbf{x})$ , one can achieve any positive  $b$  as the margin on all samples simultaneously by predicting with  $b \cdot d \circ s(\mathbf{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, MARG-CTRL performs on par or better than the more expensive two-stage shortcut-mitigating methods [Liu et al., 2021, Zhang et al., 2022]. We then 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 default-ERM 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  $\mathbf{y},\mathbf{z},\mathbf{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  $\rho$ ,  $\mathcal{F} = \{p_{\rho}(\mathbf{y},\mathbf{z},\mathbf{x})\}_{\rho}$ , such that the shortcut-label relationship  $p_{\rho}(\mathbf{z},\mathbf{y})$  changes over the family. Many common tasks in the spurious correlations literature have stable features  $s(\mathbf{x})$  that are perfect, meaning that the label is a deterministic function  $d$  of the stable feature:  $\mathbf{y} = d\circ s(\mathbf{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(\mathbf{x})$  determines the label, it holds that  $\mathbf{y}\|_{p_{\rho}}(\mathbf{x},\mathbf{z})\mid s(\mathbf{x})$ . Then, the optimal predictor on the training distribution is optimal on all distributions in the family  $\mathcal{F}$ , regardless of the shortcut because
$$
p _ {t r} (\mathbf {y} \mid \mathbf {x}) = p _ {t r} (\mathbf {y} \mid s (\mathbf {x})) = p _ {t e} (\mathbf {y} \mid s (\mathbf {x})) = p _ {t e} (\mathbf {y} \mid \mathbf {x}).
$$
The most common procedure to train predictive models to approximate  $p_{tr}(\mathbf{y} \mid \mathbf{x})$  is gradient-based optimization of cross-entropy (also called log-loss); we call this default-ERM. Default-ERM 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 default-ERM 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 default-ERM'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\}$ ,  $\mathcal{N}$  be the normal distribution,  $d$  be the dimension of  $\mathbf{x}$ , and  $\rho \in (0, 1)$ ,  $B > 1$  be scalar constants. The DGP for  $p_{\rho}(\mathbf{y}, \mathbf{z}, \mathbf{x})$  is:
$$
\mathbf {y} \sim \operatorname {R a d}, \quad \mathbf {z} \sim \left\{ \begin{array}{l} p _ {\rho} (\mathbf {z} = y \mid \mathbf {y} = y) = \rho \\ p _ {\rho} (\mathbf {z} = - y \mid \mathbf {y} = y) = (1 - \rho) \end{array} , \quad \boldsymbol {\delta} \sim \mathcal {N} (0, \mathbf {I} ^ {d - 2}), \quad \mathbf {x} = [ B * \mathbf {z}, \mathbf {y}, \boldsymbol {\delta} ] . \right. \tag {1}
$$
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  $\mathbf{x}$ , i.e.  $\mathbf{x}_1$ , is a shortcut that is correlated with  $\mathbf{y}$  according to  $\rho$ . The factor  $B$  in  $\mathbf{x}_1$  scales up the gradients for parameters that interact with  $\mathbf{x}_1$  in predictions. For large enough  $B$ , model dependence on the shortcut feature during default-ERM 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 default-ERM which uses log-loss: on a data point  $(\mathbf{x},\mathbf{y})$  the log-loss is
$$
\ell_ {l o g} (\mathbf {y} f _ {\theta} (\mathbf {x})) = \log \left[ 1 + \exp (- \mathbf {y} f _ {\theta} (\mathbf {x})) \right].
$$
With  $d = 300$  and  $B = 10$ , we train a linear model on 1000 samples from the training distribution  $p_{\rho=0.9}$ , and evaluate on 1000 samples from  $p_{\rho=0.1}$ .
Observations. Figure 1a shows that when trained with default-ERM, 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  $\mathbf{x}_2$ , the model relies on other features like the shortcut  $\mathbf{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, gradient-based optimization with log-loss prefers
(a) Average accuracy and loss curves.
(b) Accuracy and loss on shortcut and leftover groups.
Figure 1: Accuracy and loss curves for training a linear model with default-ERM on 1000 training samples from  $p_{0.9}$ , with  $B = 10, d = 300$  (see eq. (1)), and testing on  $p_{0.1}$ . (a) The model achieves  $100\%$  train accuracy but  $< 40\%$  test accuracy. (b) The learned model achieves high test accuracy  $(\approx 90\%)$  on the shortcut group and low test accuracy on the leftover group  $(\approx 30\%)$ . Models that depend more on the stable feature than on the shortcut, achieve at least  $50\%$  accuracy on both the shortcut and leftover groups. Hence the learned model exploits the shortcut to classify the shortcut group and overfits to the leftover group.
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  $\mathbf{z}$  and makes the Bayes-optimal prediction w.r.t  $p_{tr}$ :  $\arg \max_y p_{tr}(\mathbf{y} = y\mid \mathbf{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  $\mathbf{y}$  is positively correlated with  $\mathbf{z}$ , hence the shortcut group consists of all instances with  $\mathbf{y}^i = \mathbf{z}^i$  and the leftover group of those with  $\mathbf{y}^i\neq \mathbf{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 default-ERM-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  $\ell_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 default-ERM'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 default-ERM'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_{\theta}(\mathbf{x}) = \mathbf{w}^{\top}\mathbf{x}$  where  $\mathbf{w} = [\mathbf{w}_z,\mathbf{w}_y,\mathbf{w}_e]$  with default-ERM. 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 default-ERM on linearly separable data, it achieves zero training loss and converges to the direction of a minimum  $\ell_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  $\mathbf{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
$$
\forall n > N _ {0}, \quad \forall \text {i n t e g e r s} k \in \left(0, \frac {n}{1 0}\right), \quad \forall d \geq C _ {1} k \log (3 n), \quad \forall B > C _ {2} \sqrt {\frac {d}{k}}, \tag {2}
$$
with probability at least  $1 - \frac{1}{3n}$  over draws of the training data, it holds that  $B\mathbf{w}_z^* > \mathbf{w}_y^*$ .
The size of the leftover group  $k$  concentrates around  $(1 - \rho)n$  because each sample falls in the leftover group with probability  $(1 - \rho)$ . Thus, for  $\rho > 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  $\sqrt{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  $\ell_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  $\gamma$ , the weights  $\mathbf{w}_e = \gamma \sum_{i \in S_{\text{leftover}}} \mathbf{y}_i \boldsymbol{\delta}_i$ . The margin change on the  $j$ th training sample from using these weights is  $\mathbf{y}_j \mathbf{w}_e^\top \boldsymbol{\delta}_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  $\sqrt{d}$ . For samples in the leftover group, the margin change is the sum of mean one random variable; this sum grows as  $d$  and its standard deviation grows as  $\sqrt{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  $\mathbf{y}(\mathbf{w}^{*})^{\top}\mathbf{x} = \mathbf{w}_{y}^{*} - B\mathbf{w}_{z}^{*} + (\mathbf{w}_{e}^{*})^{\top}\mathbf{y}\delta$  is a Gaussian random variable centered at a negative number  $\mathbf{w}_{y}^{*} - B\mathbf{w}_{z}^{*}$ . Then, with  $\Phi_{e}$  as the CDF of the zero-mean Gaussian random variable  $(\mathbf{w}_{e}^{*})^{\top}\delta$ , accuracy in the test leftover group is
$$
p (\mathbf {y} (\mathbf {w} ^ {*}) ^ {\top} \mathbf {x} \geq 0 \mid \mathbf {y} \neq \mathbf {z}) = p [ (\mathbf {w} _ {e} ^ {*}) ^ {\top} \boldsymbol {\delta} > - (\mathbf {w} _ {y} ^ {*} - B \mathbf {w} _ {z} ^ {*}) ] = 1 - \Phi_ {e} (- (\mathbf {w} _ {y} ^ {*} - B \mathbf {w} _ {z} ^ {*})) \leq 0. 5.
$$
Second, the leftover group in the training data is overfit in that the contribution of noise in prediction  $(|(\mathbf{w}_e^*)^\top \pmb {\delta}|)$  is greater than the contribution from the stable and shortcut features. Formally, in the training leftover group,  $\mathbf{w}_y^* -B\mathbf{w}_z^* < 0$  . Then, due to max-margin property,
$$
\mathbf {w} _ {y} ^ {*} - B \mathbf {w} _ {z} ^ {*} + (\mathbf {w} _ {e} ^ {*}) ^ {\top} \mathbf {y} _ {i} \pmb {\delta} _ {i} > 1 \Longrightarrow (\mathbf {w} _ {e} ^ {*}) ^ {\top} \mathbf {y} _ {i} \pmb {\delta} _ {i} \geq 1 - (\mathbf {w} _ {y} ^ {*} - B \mathbf {w} _ {z} ^ {*}) > | \mathbf {w} _ {y} ^ {*} - B \mathbf {w} _ {z} ^ {*} |.
$$
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 non-overparameterized linear models, the choice of default-ERM 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, default-ERM — 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  $\tau \in (0,1)$  such that the dimension of  $\mathbf{x}$  is  $d = \tau n < n$ , for all integers  $k < n \times \min \left\{\frac{1}{10}, \frac{\tau}{C_1 \log 3n}\right\}$ , a linear model trained via default-ERM yields a predictor  $\mathbf{w}^*$  such that  $B\mathbf{w}_z^* > \mathbf{w}_y^*$ .
If default-ERM 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 default-ERM 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(\mathbf{x})$  for a perception task, in that for a function  $d$  when  $\mathbf{y} = d \circ s(\mathbf{x})$ , one can achieve margin  $b \in (0, \infty)$  uniformly on all samples by predicting with the stable  $b \cdot d \circ s(\mathbf{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 overparameterization (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_{\theta}(\mathbf{x}) = \mathbf{w}^{\top}\mathbf{x}$  such that for all samples in the training data  $\mathbf{y}_i\mathbf{w}^\top \mathbf{x}_i = b$  for any  $b \in (0,\infty)$ . With probability 1 over draws of samples,  $\mathbf{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  $\infty$ . We define MARG-CTRL losses for a model  $f_{\theta}$  with the margin on a sample  $(\mathbf{x},\mathbf{y})$  defined as  $\mathbf{y}f_{\theta}(\mathbf{x})$ .
As the first MARG-CTRL loss, we develop the  $\sigma$ -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
$\sigma$ , the  $\sigma$ -damped loss modifies the model output  $f_{\theta}$  and plugs it into log-loss: