SGD Provably Prioritizes a Shortcut Spurious Feature in the XOR Model

Tyler LaBonte; Vidya Muthukumar

stat.MLJun 29, 2026

SGD Provably Prioritizes a Shortcut Spurious Feature in the XOR Model

Authors: Tyler LaBonte, Vidya Muthukumar

Organizations: H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology · School of Electrical and Computer Engineering, Georgia Institute of Technology

Abstract

Neural networks are known to be susceptible to over-reliance on spurious correlations. However, the precise mechanism by which models exploit shortcut features is not fully understood, and algorithms to mitigate this behavior rely on as yet unjustified assumptions about the learned representations. In this work, we provide the first end-to-end theoretical characterization of spurious feature learning for two-layer ReLU neural networks trained by online minibatch SGD on the logistic loss. We consider data drawn from the high-dimensional Boolean hypercube with a quadratic signal function (namely XOR) and a linear spurious correlation. We show that SGD learns the spurious feature first, and exponentially fast. Moreover, the optimization dynamics couple the spurious and signal features, with a stronger spurious component inhibiting signal feature learning. Our analysis reveals precise phase transitions in the learning dynamics. In the first phase, alignment between the signs of the spurious feature and second-layer weight drives rapid growth of the spurious feature. In the second phase, large majority group margin slows learning and the signal feature remains suppressed. When the spurious correlation is maximally strong, we show theoretically that the spurious feature dominates even at the sample complexity threshold where XOR would be learned in isolation (i.e., if the spurious feature was absent). In contrast, when the correlation strength is constant, we provide preliminary empirical evidence that the model can eventually learn the XOR signal, although the spurious feature is not forgotten.

Explore similar work

Apr 28, 2026cs.LG

Feature Repulsion and Spectral Lock-in: An Empirical Study of Two-Layer Network Grokking

Tian (2025) proves a repulsion theorem (Theorem 6) for the matrix

B = (\widetilde{F}^\top \widetilde{F} + ηI)^{-1}

during the interactive feature-learning stage of grokking: similar features have negative off-diagonal entries

B_{j\ell}

, producing an effective repulsive force that drives them apart. However, the theorem does not specify when this mechanism becomes empirically observable, nor whether it leaves a measurable spectral signature in the parameter updates. We test this directly on Tian's modular addition setup (

M = 71

K = 2048

, MSE loss) and observe a clear structure-mechanism dissociation. The predicted sign rule holds robustly on the top-200 most-similar feature pairs across activations (empirical sign-match rising from 0.865 to 0.985 on

σ= x^2

across 5 seeds, and saturating at 1.000 on

σ= \operatorname{ReLU}

). However, the spectral signature in the parameter updates is strongly activation-dependent. With

σ= x^2

, a simple slope detector on the rolling eigengap

σ_2 / σ_3

ΔW

fires in 15/15 grokking seeds at epoch 174 (IQR [173,174]) and in 0/15 non-grokking controls, with 229

\times

late-stage magnitude separation; the spectrum is rank-2. In contrast, with

σ= \operatorname{ReLU}

, the detector never fires and the spectrum remains effectively rank-1. This dissociation aligns with Tian's Theorem 5 distinction between focused (power-law) and spreading (ReLU) memorization: while the sign structure of

B

depends only on

\widetilde{F}^\top \widetilde{F}

, how feature repulsion translates into weight updates critically depends on the activation derivative

σ'

Yongzhong Xu

Feature Learning Grokking Two-Layer Neural Networks Rotational Repulsion Mechanism Empirical Study

Ask PDF

May 11, 2026cs.LG

The Benefits of Temporal Correlations: SGD Learns k-Juntas from Random Walks Efficiently

We study how temporal correlations in the data can make certain sparse learning problems efficiently learnable by gradient-based methods. Our focus is on Boolean k-juntas, a canonical sparse learning problem known to pose barriers for gradient-based methods under independent uniform samples. We show that this picture changes when the samples are generated by a lazy random walk on the hypercube. In this setting, the temporal dependencies can be exploited by a two-layer ReLU network trained using stylized-SGD with a temporal-difference loss, which compares target and predicted increments across consecutive samples. For every fixed k, the resulting sample complexity is essentially linear in the ambient dimension d. By contrast, we show that for large-batch gradient methods using standard convex pointwise losses, temporal correlations do not provide the same advantage.

Elisabetta Cornacchia, Dan Mikulincer, Elchanan Mossel

Spatio-Temporal Correlations Stochastic Gradient Descent Optimal Sample Complexity Temporal Difference

Ask PDF

Jul 3, 2026cs.LG

Implicit Bias of SGD in Multivariate ReLU Networks: Effective Width Collapse

We study the implicit bias of noisy stochastic gradient descent in training wide two-layer ReLU networks for multivariate regression. In a mean-field regime, the training dynamics are approximated by a Wasserstein gradient flow that converges to a unique stationary measure. We characterize the structure of this stationary measure and the predictor it represents. We show that, despite the network being infinitely overparameterized, the learned predictor admits an effectively finite representation: the input weights and biases align along finitely many directions, leading to an effective width collapse. In particular, the solution function is continuous piecewise affine, with affine regions determined by the cells of a finite hyperplane arrangement. The number of learned directions, and hence hyperplanes, is bounded above by

2\mathcal{P}-1

, where

\mathcal{P}

denotes the number of linear dichotomies realizable on the training inputs. We further establish a non-redundancy property of the learned representation by proving that each learned direction induces a unique ternary activation pattern on the training data. Consequently, the complexity of the learned predictor is governed by the combinatorial geometry of the training data.

Shuang Liang, Tom Jacobs, Guido Montúfar

Rectified Linear Unit Networks Stochastic Gradient Descent Gradient Flow Mean-Field Limit Representation Collapse Implicit Bias

Ask PDF