Implicit Regularization & Shortcut Learning

In safety-critical domains like healthcare, biomedical data often contains complex, unknown, or hard-to-eliminate biases. When machine learning models train on this data, they frequently rely on “shortcuts”—spurious features that are strongly correlated with the target label in the training distribution but remain unstable across different environments.

Learning these shortcuts severely impedes the model’s ability to identify invariant, core features. As a result, a model might achieve high overall accuracy on average, but fail catastrophically on out-of-distribution or minority samples where the spurious correlation does not hold.

My PhD thesis focuses on modifying the inductive biases of models to prioritize robust, fair performance across all subpopulations. Specifically, I explore how optimization hyperparameters, like batch size and learning rate, govern a model’s reliance on these shortcuts.

The Optimizer’s Role: SGD vs. Gradient Descent

Gradient-based optimizers are often treated as black boxes that converge to a max-margin solution, but they play a fundamental role in what a model actually chooses to learn. Even when fully informative invariant features exist, models often favor spurious features because they artificially increase the margin, accelerating the reduction of empirical loss during optimization.

Our research uncovers a striking phenomenon: Stochastic Gradient Descent (SGD) systematically suppresses dependence on spurious features, while full-batch Gradient Descent (GD) does not confer this benefit and may actually exacerbate shortcut reliance.

This happens due to the implicit regularization inherent to the optimizers. While GD penalizes the squared norm of the full-batch gradient (favoring flatter minima), SGD additionally penalizes the variance of gradients across mini-batches.

Figure 1. Implicit regularization of GD and SGD on a four-point dataset with FIIF. (a) Four-point dataset with a fully informative invariant feature (FIIF). The data lie in two dimensions: an invariant feature ($y$) and a spurious feature ($z$). Majority samples are shown in saturated colors. The invariant solution $w_{\mathrm{INV}}$ achieves perfect classification. Similarly, the solutions minimizing $C(w)$, $C_{\mathrm{GD}}(w)$, and $C_{\mathrm{SGD}}(w)$, denoted by $w^\star$, $w^\star_{\mathrm{GD}}$, and $w^\star_{\mathrm{SGD}}$, respectively, also achieve perfect accuracy, but with larger margins with respect to majority samples. (b) Comparison of the minima of $C(w)$, $C_{\mathrm{GD}}(w)$, and $C_{\mathrm{SGD}}(w)$, together with schematic trajectories illustrating the optimization flows of SGD (blue) and GD (red). (c,d) Implicit regularization landscapes of GD and SGD, where darker regions indicate lower values. Notably, SGD imposes a weaker penalty on solutions with smaller $w_z$, thereby favoring parameters that rely less on the spurious feature.

The Power of Small Batch Sizes & Large Learning Rates

We established both theoretically (using a four-point data generation model) and empirically across deep neural networks that the robustness gains from SGD are directly tied to hyperparameters: the effect intensifies with smaller batch sizes and larger learning rates.

Figure 2. Joint effect of learning rate and batch size on WGA and ACC. Across all datasets, smaller batch sizes consistently achieve higher worst-group accuracy (WGA) when in-distribution generalization is maintained, indicating improved robustness to spurious correlations. Extremely small or large learning rates can lead to optimization failure and poor in-distribution generalization, in which case robustness comparisons become unreliable. These results highlight the importance of the learning-rate-to-batch-size ratio in shaping SGD’s implicit regularization and its effect on robustness.

By viewing each mini-batch as a sampled “domain” from the training distribution, smaller batch sizes naturally increase the variability in subpopulation composition. Suppressing gradient variance under these highly fluctuating conditions forces the model to abandon shortcuts and rely on stable, invariant features, significantly improving Worst-Group Accuracy (WGA).

Furthermore, we found that training with small batch sizes provides a much stronger inductive bias when paired with explicit debiasing methods (like DFR, AFR, and EVaLS), yielding WGA improvements of 25-37% on complex datasets with multi-level spurious correlations.

This perspective reveals that robustness can emerge naturally from the inductive biases of standard optimizers, reducing the need for exhaustive hyperparameter searches in explicit shortcut-mitigation frameworks (Mirzaie et al., 2026).

Training with stochastic gradient descent (SGD) at moderately large learning rates has been observed to improve robustness against spurious correlations, strong cor- relation between non-predictive features and target labels. Yet, the mechanism underlying this effect remains unclear. In this work, we identify batch size as an additional critical factor and show that robustness gains arise from the implicit regularization of SGD, which intensifies with larger learning rates and smaller batch sizes. This implicit regularization reduces reliance on spurious or short- cut features, thereby enhancing robustness while preserving accuracy. Importantly, this effect appears unique to SGD: gradient descent (GD) does not con- fer the same benefit and may even exacerbate shortcut reliance. Theoretically, we establish this phenomenon in linear models by leveraging statistical formulations of spurious correlations, proving that SGD systematically suppresses spurious feature dependence. Empirically, we demonstrate that the effect extends to deep neural networks across multiple benchmarks.

The Optimizer’s Role: SGD vs. Gradient Descent

The Power of Small Batch Sizes & Large Learning Rates

References

2026