When Data Imbalance Helps: Robust Generalization Through Shortcut Saturation

arXiv:2607.10116 · cs.LG, cs.AI · Submitted 2026-07-11 · Read on arXiv

cs.LG, cs.AI

Submitted: 2026-07-11

Updated: 2026-09-17

License: http://creativecommons.org/licenses/by/4.0/

The gist: We study robust generalization under spurious correlations: tasks where a shortcut feature is correlated with the true label in training but anti-correlated in an adversarial held-out split.

Terminology

Abstract

We study robust generalization under spurious correlations: tasks where a shortcut feature is correlated with the true label in training but anti-correlated in an adversarial held-out split. Varying the spurious ratio r (the fraction of training examples where shortcut = true label) and model capacity, we find a counterintuitive result: data imbalance promotes generalization in sufficiently capable models. On a synthetic task where the true label is sum parity of an integer sequence and the shortcut is the parity of the maximum-valued element, a 2-layer, 2-head transformer generalized (reached 100% adversarial accuracy) in 0% of seeds at r = 0.50 but 77% of seeds at r = 0.90. The effect is absent in 1-layer models, where imbalance instead traps the model on the shortcut. Through mechanistic analysis -- gradient conflict dynamics, circuit evolution, and QK/OV circuit ablations -- we characterize a mechanistic pathway consistent with imbalance promoting generalization.

Sources

Related papers