SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\ell_ {\sigma \mathrm {- d a m p}} (\mathbf {y}, f _ {\theta}) = \ell_ {\log} \left(\mathbf {y} \left(1 - \sigma \left(\frac {\mathbf {y} f _ {\theta}}{T}\right)\right) f _ {\theta}\right)
$$
For large margin predictions  $\mathbf{y}f_{\theta} > 0$ , the term  $1 - \sigma (\mathbf{y}f_{\theta}(\mathbf{x}) / T)$  damps down the input to log-loss. The largest the input to  $\ell_{\log}$  can get is  $0.278T$ , 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  $\ell_{\sigma \text{-damp}}$  occurs when the margin is  $0.278T$  on all samples. To demonstrate empirical advantage, we compare
Figure 2: Using  $\sigma$ -damped log-loss yields linear models that depend on the perfect stable feature to achieve near perfect test accuracy. The middle panel shows that  $\sigma$ -damping maintains similar margins in the training shortcut and leftover groups unlike unconstrained log-loss, and the right panel shows  $\sigma$ -damp achieves better leftover test-loss.
standard log-loss to  $\sigma$ -damped loss on eq. (1); see fig. 2. The left panel of figure fig. 2 shows that test accuracy is better for  $\sigma$ -damp. The middle and right panels show the effect of controlling margins in training, where losses on shortcut and leftover groups hover at the same value.
Second, we design the  $\sigma$ -stitch loss, which imitates log-loss when  $\mathbf{y}f_{\theta}(\mathbf{x}) < u$  and penalizes larger margins  $(\mathbf{y}f_{\theta} > u)$  by negating the sign of  $\mathbf{y}f_{\theta}(\mathbf{x})$ :
$$
\ell_ {\sigma \text {- s t i c t h}} = \ell_ {l o g} (\mathbf {1} [ \mathbf {y} f _ {\theta} (\mathbf {x}) \leq u ] \mathbf {y} f _ {\theta} (\mathbf {x}) + \mathbf {1} [ \mathbf {y} f _ {\theta} (\mathbf {x}) > u ] (2 u - \mathbf {y} f _ {\theta} (\mathbf {x})) \tag {3}
$$
As the third MARG-CTRL loss, we directly penalize large margins via a log-penalty:
$$
\ell_ {\operatorname {m a r g} - \log} = \ell_ {l o g} \left(\mathbf {y} f _ {\theta} (\mathbf {x})\right) + \lambda \log (1 + | f _ {\theta} (\mathbf {x}) | ^ {2}) \tag {4}
$$
The fourth MARG-CTRL loss controls margins by penalizing  $|f_{\theta}(\mathbf{x})|^2$
$$
\ell_ {\mathrm {S D}} = \ell_ {\text {l o g}} (\mathbf {y} f _ {\theta} (\mathbf {x})) + \lambda | f _ {\theta} (\mathbf {x}) | ^ {2} \tag {5}
$$
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 default-ERM on the linear perception task; see figs. 6 to 8. We also run MARG-CTRL on a neural network and show that while default-ERM achieves test accuracy worse than random chance, MARG-CTRL achieves  $100\%$  test accuracy; see figs. 10 to 13 in appendix B.4.
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 default-ERM 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  $(\mathbf{w}_z > 0)$ . In their set up, the linear models are overparameterized and it is unclear when  $\mathbf{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  $\mathbf{w}_z > 0$  but test accuracy is  $100\%$ . In contrast to Nagarajan et al. [2021], theorem 1 gives DGPs where the leftover group accuracy is worse than random, even without overparameterization. Ahuja et al. [2021] also consider linear classification with default-ERM with a perfect stable feature and conclude that default-ERM 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 default-ERM 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 default-ERM 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.
MARG-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 default-ERM produces models that depend more on the shortcut and select hyperparameters of the two stage process using validation group annotations. In the nuisance-free setting where there are no validation group annotations, the performance of these methods can degrade below that of default-ERM. 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 losses; 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  $\mathbf{y}f_{\theta}(\mathbf{x})$ , gradient descent has an inductive bias to follow the  $\ell_2$ -regularization path. All these inductive biases prefer shortcuts if using them leads to lower loss within an  $\ell_2$ -norm-budget. MARG-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(\mathbf{x}) = \mathbf{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 default-ERM. 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 nuisance-free 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 Civil-Comments 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  $\sigma$ -damp,  $\sigma$ -stitch, MARGLOG, 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
CelebA	WB	Civil
ERM	72.8 ± 9.4	70.8 ± 2.4	60.1 ± 0.4
CNC	81.1 ± 0.6	68.0 ± 1.8	68.8 ± 0.2
JTT	75.2 ± 4.6	71.7 ± 4.0	69.9 ± 0.4
marg-log	82.8 ± 1.1	78.2 ± 1.9	68.4 ± 1.8
σ-damp	79.4 ± 0.6	78.6 ± 1.1	69.6 ± 0.4
SD	81.4 ± 2.5	80.5 ± 1.4	69.9 ± 1.1
σ-stitch	81.1 ± 2.2	75.9 ± 3.4	67.8 ± 2.8
Table 1: Mean and standard deviation of test worstgroup accuracies over two seeds for default-ERM, JTT, CNC,  $\sigma$  -damp,  $\sigma$  -stitch, SD, and marg-log. Every MARGCTRL method outperforms default-ERM on every dataset. On Waterbirds, MARG-CTRL outperforms JTT and CNC. On CelebA, SD, marg-log, and  $\sigma$  -stitch beat JTT and achieve similar or better performance than CNC. On CivilComments,  $\sigma$  -damp and SD beat CNC and achieve similar performance to JTT.
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 MARG-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.
MARG-CTRL improves over default-ERM. We compare MARG-CTRL to default-ERM on CelebA, Waterbirds, and Civilcomments. Table 1 shows that every MARG-CTRL method achieves higher test worst-group accuracy than default-ERM on all datasets. Default-ERM achieves a mean test worst-group accuracy of  $70.8\%$ ,  $72.8\%$  and  $60.1\%$  on Waterbirds, CelebA, and Civilcomments respectively. Compared to default-ERM, 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. Default-ERM'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 optima
Figure 3: Loss curves of default-ERM on CelebA for two combinations of LR and WD. The combination with the larger learning rate (blue) achieves  $72.8\%$  test worst-group accuracy, beating the other combination by  $20\%$ . The model that achieves the best validation (and test) worst-group accuracy is the one at epoch 13 from the blue run. This model achieves similar loss in both groups and the full data model suggesting that large LR and WD controls margins from exploding (higher training loss in all panels) and avoids systematically smaller margins in the leftover group compared to the shortcut group.
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 default-ERM'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 default-ERM, including those not in the plot, is the one at epoch 13 on the blue curves. This model achieves similar train and test losses  $(\approx 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 default-ERM with large LR and WD mitigates shortcuts when controlling margins like MARG-CTRL.
MARG-CTRL performs as well or better than two-stage shortcut-mitigating methods. Two-stage shortcut mitigating methods like Correct-n-Contrast (CNC) and Just Train Twice (JTT) aim to mitigate shortcuts by using a model trained with default-ERM to approximate group annotations. They rely on the assumption that a model trained via default-ERM 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 default-ERM-trained model and use them to mitigate shortcut learning in a second predictive model. JTT 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  $\sigma$ -stitch perform similar to CNC while all MARG-CTRL techniques outperform JTT. On Waterbirds, all MARG-CTRL methods outperform JTT and CNC. On CivilComments,  $\sigma$ -damp and SD perform similar to JTT and outperform CNC. CNC'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.
MARG-CTRL is faster than JTT and CNC. MARG-CTRL takes the same time as default-ERM, 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 MARG-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 default-ERM, JTT, and CNC. On average, the MARG-CTRL methods close at least  $61\%$  of the gap between default-ERM in the nuisance-free setting and the best performance in table 1 on every dataset. In contrast, CNC and JTT sometimes perform worse than default-ERM.
We study why default-ERM — gradient-based optimization of log-loss — yields models that
CelebA	WB	Civil
ERM	57.5 ± 5.8	69.1 ± 2.1	60.7 ± 1.5
CNC	67.8 ± 0.6	60.0 ± 8.0	61.4 ± 1.9
JTT	53.3 ± 3.3	71.7 ± 4.0	53.4 ± 2.1
marg-log	74.2 ± 1.4	77.9 ± 0.3	66.8 ± 0.2
σ-damp	70.8 ± 0.3	74.8 ± 1.6	65.6 ± 0.2
SD	70.3 ± 0.3	78.7 ± 1.4	67.8 ± 1.3
σ-stitch	76.7 ± 0.6	74.5 ± 1.2	66.0 ± 1.0
Table 2: Average and standard deviation of test worst-group accuracy over two seeds of MARG-CTRL, default-ERM, JTT, and CNC in the nuisance-free setting. Hyperparameter selection and early stopping use label-balanced average accuracy. All MARG-CTRL methods outperform default-ERM, JTT, and CNC on all datasets.
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 default-ERM'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. MARG-CTRL improves over default-ERM 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 MARG-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, max-margin 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  $\ell_2$  improves over max-margin classifiers with  $\ell_2$ : all experiments use overparameterized models trained with weight decay and MARG-CTRL outperforms default-ERM. 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 default-ERM [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,  $\ell(a) = \frac{1}{1 + a^K}$  for some  $K > 0$  — do not 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  $\mathbf{y}$  is not determined by  $s(\mathbf{x})$ :  $\mathbf{y} \perp (\mathbf{x}, \mathbf{z}) \mid s(\mathbf{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.