Strong Correlations Induce Cause Only Predictions in Transformer Training
Authors: Haihan Zhang
Yimu Zhang
Cong Fang
Venue: ICLR 2026 Poster (2026)
Keywords: Implicit bias
Transformers
Optimization
Causal robustness prediction
Primary area: optimization

Abstract: We revisit when Transformers can prioritize causes over spurious effects by viewing the problem through data correlation strength and the implicit regularization of gradient descent. We identify a phenomenon called Correlation Crowding-Out (CCO) arising from the training dynamics of Transformers. Specifically, under strongly correlated causal features, gradient descent filters out spurious cues and converges to a predictor that relies almost exclusively on the causes. Theoretically, using a simplified Transformer model trained on data from a minimal causal chain, we introduce a Dominant-coordinate condition that characterizes when CCO arises and explain its mechanism as a coupling of ''occupation'' and ''crowding-out''.  ''Occupation'' denotes to rapid growth of weights aligned with the dominant causal direction while non-dominant directions remain small. ''Crowding-out'' denotes to attention logits align with separation directions favoring the causal branch, suppressing descendants. We provide convergence guarantees for both the optimization trajectory and generalization. Our empirical results on simulated and real examples across various tasks including vision and natural language demonstrate the procedure. Together, these results show that, under suitable conditions, standard training alone can induce cause only prediction.

Published as a conference paper at ICLR 2026
STRONG CORRELATIONS INDUCE CAUSE ONLY PRE- DICTIONS IN TRANSFORMER TRAINING
Haihan Zhang', Yimu Zhang', Cong Fang!.2.
1 State Key Lab of General AI, School of Intelligence Science and Technology, Peking University 2 Institute for Artificial Intelligence, Peking University (zhanghaihan] @stu. pku. edu.cr.
ABSTRACT
We revisit when Transformers can prioritize causes over spurious effects by view- ing the problem through data correlation strength and the implicit regularization of gradient descent. We identify a phenomenon called Correlation Crowding- Out (CCO) arising from the training dynamics of Transformers. Specifically, un- der strongly correlated causal features, gradient descent filters out spurious cues and converges to a predictor that relies almost exclusively on the causes. The- oretically, using a simplified Transformer model trained on data from a mini- condition that character-
mal causal chain, introduce Dominant-coording izes when CCO arises and explain its mechanism as a coupling of "occupation" and "crowding-out"" "Occupation" denotes the rapid growth of weights aligned with the dominant causal direction while non-dominant directions remain small. "Crowding-out" denotes the attention logits align with separation directions favor- ing the causal branch, suppressing descendants. We provide convergence guaran- tees for both the optimization trajectory and generalization. Our empirical results on simulated and real examples across various tasks including vision and natu- ral language demonstrate the procedure. Together, these results show that, under suitable conditions, standard training alone can induce cause only prediction.
INTRODUCTION
Whether data-driven models can extract causal invariances from observational data and thereby de- liver robust predictions has long been central hope in Al (Pearl 2009 Peters et al. 2016 Arjovsky al. 2019; Schölkopf et al., 2021 Fan et al.] 2024). Yet models trained by empirical
risk min-
2020 Ye et al. 2024). indiscriminately exploiting any correlation, including spurious cues unrelated to the true causal mechanisms (Sagawa al. 2020 Qiu et al., 2023). This pattern is widely documented across modalities and tasks (Geirhos et al.
imization are often prone to shorteut learning (Geirhos et al. 2020: Shah et al.
2019: McCoy et al. 2019; [i et al. 2025). The rise of
Transformers and LLMS sharpens this tension: these systems can sometimes rely on shallow arti- facts (Bender et al. 2021 Tang et al.
2023 Du et al.] 2023 Varma et al. 2024 Jinn et al. [024]] Gui Ji. 2025), yet they also produce answers that appear strikingly logical and robust in certain scenarios (Brown et al. 2020; WWii et al. (2022] Kojima et al. 2022) Yuan et al. 2024). Recent theory offers partial clues for why Transformers sometimes appear causal, but does not yet answer the cause only generalization question. On stylized in-context tasks, Transformers trained on Markov Chain sequences can recover parent sets and estimate transition probabilities in-context (Edelman et al.) Nichani et al., 20244) D' Angelo et al.) (2025). These results suggest how attention might recon- struct graph edges from observational sequences, but they rely on designed ICL setups rather than generic pipelines with spurious features and do not show when spurious information is suppressed at both train and test time. In parallel, large margin analyses show that gradient descent (GD) pushes query- key parameters toward max- -argin separators (Tarzanagh et al. 2023 Ataee Tarzanagh et al.) 2023 Vasudeva et al. 2025). While this suggests separation can emerge during training, it does not characterize how such separation filters out spurious features or yields cause only risk guarantees. This landscape motivates basic question:
"Corresponding author

---

Published as a conference paper at ICLR 2026
When and through what mechanism can Transformer training produce predictors that rely on causes while ignoring spurious effects?
We answer this by uncovering and analyzing Correlation Crowding- (CCO). CCO is a training phenomenon in which, under a uniform dominance gap where causal feature is more strongly associated with the target than any competing spurious feature, GD drives Transformer to progres- sively suppress spurious features and converge to a predictor that relies almost exclusively on the causal feature. Crucially, the dominance condition does not require spurious features to be weak: many remain highly correlated with the target and may even surpass non-dominant causal coor- dinates. What matters a persistent margin favoring the dominant causal direction.
Remarkably, strong causal correlation in the data alone does not guarantee cause only prediction for generic estimators. In Example 25] even under a dominance gap, population least squares retains a constant fraction of spurious feature. Thus, CCO is not a corollary of data dominance; it hinges
on optimization induced implicit regularization that actively crowds out spurious features. This
occurs without explicit invariance penalties (Arjovsky al.
2019:
Shapiro
2017 Fan et al. 2024)
or celli--enviromment training (Peters et al., 2016 Fan et al.
2024
Xu et al.
2024
Our perspective complements existing analyses of correlation driven learning dynamics. Prior work has shown that neural networks exhibit a simplicity bias, often preferring features that are highly correlated with the label or easier fit (Belkin al., 2019] Moayeri et al. [022] Morwani et al.)
2023 Qiu et al., 2023; Xue et al. [023] Yang et al. (024)). When spurious features are more
predictive or less complex, they tend dominate early training, delaying or even entirely inhibiting the learning of causal features (Shah et al., 2020 Yang et al., 2024). These studies underscore that
correlation strength and feature complexity critically shape learning trajectories, and they reinforce the notion that deep models
vulnerable superficial shorteuts. In contrast, we focus the opposite regime: when the causal features themselves dominate in predictiveness. We formalize CCO as the mirror image of shortcut learning. Intuitively, if causal feature explains the target with overwhelming strength, the model has little incentive to rely on weaker spurious cues.
Building this premise, we demonstrate CCO empirically and provide a theoretical account of its mechanism. To theoretically understand this behavior, we analyze a simplified two-layer Trans- former trained data from a causal chain (x + y + z) generative process. Our theory provides Dominant-Cooddnnaee Condition on the data, which quantifies how strong the x-y correlation must be for CCO to occur. Under this condition, the training dynamics unfold in two coupled phases. In the first "occupation" phase, within the Transformer's 's that aligns with the dominant causal coordinate in x grows rapidly to a stable magnitude, while weights in other directions remain small. This expansion makes the causal direction salient and es- tablishes it as the primary signal driving the predictions. Next comes the "crowding-our" phase: the Transformer's 's attention mechanism gradually shifts its query-key aligmment toward the max-margin separator between the transformed causal and spurious features (roughly, x - Z). Consequently, the attention weights concentrate almost entirely the causal x branch, effectively gating out the spu- rious z branch. Through this two-phase process, GD steers the model toward a cause only solution without any specialized regularization for invariance.
sublayer, the weight vector
By elucidating the mechanism behind CCO, we contributes a more nuanced perspective on Trans- former's generalization: while spurious shorteuts are serious and pervasive concern, there exist regimes in which strong causal signals can turn GD into an ally for causal learning. In such regimes, the implicit regularization of GD yield cause only generalization, even in the absence of multiple training environments or explicit causal objectives.
1.1 OuR CONTRIBUTION
We introduce and formalize the new phenomenon CCO.
We elucidate CCO's mechanism with both theory and experiments.
2
RELATED WORK
Spurious Correlations and Invariance Learning. Across vision, language, and ERM-trained deep models including modern Transformens and Readilyy latch onto shorteut cues and spu-
2

---

Published as a conference paper at ICLR 2026
rious correlations, leading to brittle generalization under shift (Geirhos et al. 2019 2020 Zhou
eet al. 2021 Du et al.] 2021 Tang et al.] [023] Du et al.] 2023] Yuan et al. 2024). A major theo- seeks mechanisms correlations,
retical response is invariacee learning: instead of trusting
stable across environments. Two canonical frameworks are Invariant Causal Prediction (Peters et al.)
2016 Meinshausen al. 2016), which tests for subsets of covariates that render the conditional law of the target invariant across interventions or environments, and Invariant Risk Minimization (Ar- jovsky et al.; 2019), which encourages representations that admit a single optimal classifier across environments. Both lines have spurred extensive follow-ups and critiques clarifying assumptions, identifiability, and practical limitations (Ghassami et al. 2017; Heinze- Deml et al. 2018 Phister
et al. 2019. Rothenhäusler et al.) 2019 2021 Rosenfeld et al. 2021: Lin et al. 2022b Kamath
ct al. 2021 Lu et al. 2021 aa Zhou et al. 2022 Lin Li al. 2022a). In parallel, distributionally robust
optimization offers a complementar lens by minimizing worst-case (group) risk under distributional
shifts (Shapiro 2017 Sagawa et al. 2020 Duchi & Namkoong 2202]] Gao et al., 2024). More re-
cently, Environment- crwpammennt--nvrrine Invariant Linear Least Squares and variants show that, when cross-environment heterogeneity is sufficiently strong, a regularized least-squares estimator can recover invariant fea- tures with generalization guarantees while quantifying heterogeneity (Fan et al.] 2024 Gu et al.] 2025a Xu et al., 2024] Gu et al., (025b). Most invariance methods posit either explicit regularizer or environment partitions; comparatively less is known about when standard GD a Transformer will, by its own dynamics, yield an cause only predictor. Our work targets precisely this gap.
Implicit Bias. Implicit bias refers to the tendency of (S) GD, even without explicit regularization, to select solutions with special structure and generalization properties, widely regarded as a key to the success of over-parameterized models. Such as for logistic regression, (S) GD converges in
direction to the max- maax-margin classifier (Soudry et al.) 2018 Ji & Telgarsky, 2019 Wu et al. 2025
Cai et al. 2025); in over- parameterized linear models. (S) GD can display benign overfitting (Zou
al., 2021; Wu al. (022), double descent (Lu al. 2023 Zhang et al. 2025), ) and scaling
laws (Bordelon et al. 2024: Lin et al., 2024);; and for quadratically parameterized models, (S) GD
implicitly favors low- complexixy solutions and exhibits incremental learning (Li et al. 2018 Vaske- Vicius et al.) 2019 Woodworth et al., 2020; HaoChen et al. 2021 Li et al., 2021 Jin et al.) 2023 et et
2024). Turning to Transformers, a growing theory literature dissects how attention evolves under GD. For single-head ViTs, GD is shown concentrate attention label-relevant tokens, yielding progressively sparse maps (Jelassi et al., 2022). Li et al. 2023a). These results clarify which inputs receive mass under training induced anisotropy, but they are agnostic causal structure. On stylized in-context tasks, Transformers trained on Markoy chain sequences learn the set of parent tokens and estimate transition probabilities in-context; related mechanistic work induction heads explains how attention circuits implement dependency tracking and copying behaviors (Lu et al.)
20216, Olsson al. 2022 Li et al., 2023b Edelman al. 2024 Nichani et al. 2024; D Angelo
et al. 2025). These analyses are posed in designed ICL setups and do not address under generic training with spurious descendants when and why GD yields a cause only predictor. complemen- tary line shows that GD on attention pushes query-key parameters toward max-margin separators, establishing that separation can emerge during training; yet this does not identify which side of the margin corresponds to causal versus spurious directions, nor when separation suffices for cause only generalization (Tazzanagh et al. 2023 Atace Tarzanagh et al. 2023 Vasudeva et al. 2025).
3 CCO: A PHENOMENOLOGY IN TRANSFORMER TRAINING
CCO refers to training phenomenon in Transformers whereby, if there exists dominant causal fea- ture whose association with the target exceeds that of any competing spurious feature by uniform gap. GD learns a predictor that progressively suppresses spurious features and relies almost exclu- sively on the causal one. Crucially, this dominance condition does not require spurious features to be weak: many can remain highly correlated with the target and some may surpass non-dominant causal features. What matters is a persistent gap favoring the dominant causal direction.
CCO unfolds through two coupled effects:
(I) Occupation (early rise): within representation and prediction layers (e.g., embeddings, feed- forward blocks, attention heads), weights aligned with dominant, highly predictive causal feature grow rapidly to a stable, large scale, while spurious features aligned directions remain small, ren- dering the causal signal salient to the optimizer.
3

---

Published as a conference paper at ICLR 2026
(a) Input image
(b) Iter=0
(c) Iter=50
(d) Iter=500
Figure I: This figure shows how attention shifts during VIT training on fogrrouundforeground causal disentanglement task. (a) is the input image. Early in training (b, iter O), attention is diffuse across left bird (causal feature) and right bird (spurious feature). As training proceeds (c, iter 50), attention weight rises the left bird illustrating the occupation phase. By (d, iter 500), attention is concentrated almost entirely on the left bird, with the right bird and background receiving near zero weight, illustrating the crowding- phase.
(II) Crowding-out (attention selection): multi- attention progressively aligns its logits with sep- aration directions that prefer causal over spurious features (e.g.. larger query-key margins for causal tokens), concentrates attention mass causal features, and suppresses spurious features.
We verify the above phenomenon Fig During ViT training, the attention map first show a rapid growth in left bird (causal feature) in the occupation stage, while the attention right bird (spurious feature) remain small growth. Then in the crowding-out stage, the attention allocated left bird significantly surpasses the attention to the right bird, with the right bird's attention appearing very faint in the attention map.
3.1 WHY CCO ARISES
(I) Intrinsic strong causal correlation: CCO emerges when the data exhibits a property whereby the correlation between causal feature and the target is consistently stronger than that of any spurious feature. This is common rather than contrived: many real datasets can be viewed as mixtures of latent environments in which causal-target relationships remain relatively stable, whereas spurious features oscillate across environments. When pooled, these oscillations destructively interfere, re- ducing trrret correlation relative causal-target correlation. Equivalently, causal features concentrate stable signal, while spurious features disperse unstable variance, making the dominant causal direction statistically more salient.
(II) Implicit regularization GD in Transformers: early strong-signal directions (Occupation) steer gradients toward causal features, inducing a directional bias in the learned representation. Attention then transduces this bias into selection (Crowding-out, assigning higher weight to causal features and down-weighting spurious ones, thereby approaching invariant, cause only solution without explicit invariance penalties.
3.2 STRONG CAUSAL CORRELATION ALONE DOESN'T T ENSURE CAUSE ONLY
important to stress that strong causal correlation in the data does not by itself alone guarantee cause only prediction for generic estimators. In Example 25), we show that even under dominant causal correlation, population linear regression retains a constant fraction of the spurious features, remaining using spurious features to predict noise. This demonstrates two points: (a) strong causal alignment alone does not ensure spurious suppression; and (b) CCO not a trivial corollary of data dominance but instead relies the implicit regularization induced by Transformers and GD.
4
THEORETICALLY ANALYSIS OF CORRELATION CROWDING-OUT
4.1 PROBLEM SETUP
We provide a theoretical explanation for CCO by analyzing a specialized Transformer module trained on data generated by the causal chain y z. When the dominant feature of exhibits

---

Published as a conference paper at ICLR 2026
sufficiently strong association with y, the implicit regularization of GD leads the learned predictor to filter out z and rely almost exclusively on X.
4.1.1 DATA GENERATIVE PROCESS
We consider the causal chain x + + 2, where x, E Rd are vector covariates and y E IR is scalar response. The response is sparse quadratie signal in x:
= x (w (w*)0²+ +
with noise 1 x, Ele] O, and Var(c) = o². The descendant z depends on via an L-Lipschitz function f IR- Fd and additive noise Rd
f(y) + fLy.
We assume the moment and boundedness conditions:
E xx
E[x+2|= Var(x + E,
0
1a-1
and almost surely supi xzyl S Bx, sup2s)sd |zy Bx, ld| S Be, supisisa |Ey1 S Be, supisiga | I |If (0)ll L (r + Be) Bg == By. The ground truth is sparse and binary: w E [0, w = 1, and [supp(w")] S We observe i.i.d. samples { (x* v 2t) }
from (x..
The chain x - + a minimal DAG that that captures the key trade-off behind CCO: causal parent that determines y. versus spurious descendant z is induced by This reduction is pur- poseful and representative. For example, in sentiment analysis, content features - sentiment label rating y - label derived auxiliary fields generated downstream z (Gururangan al. 2018). So that z a descendant induced spurious correlate of y while x carries the causal signal.
In this pattern, descendants furnish alluring but non invariant shortcuts, phenomenon widely doc- umented across deep learning (Geirhos et al. 2020). By positing one dominant, highly y y-predictive direction in x while allowing to be strongly, yet causally correlated with y. Thus, the x + y + z pattern offers a principled, portable abstraction: is simple enough for precise anal- ysis yet representative of broader scenarios where CCO is expected to emerge.
4.1.2 MODEL ARCHITECTURE
We adopt two-key attention architecture and augment inputs with fixed positional encodings 81, S2 E RM:
S
z
82
E R-M+d.
We parameterize the query the gating vector q² k! = z' and the values as v = x and v = z
v E I M+d take the keys as k == x and
Two-key Attention.
Define the logits
- = (99
C, (q')
and weights
a = -a a
Che + oh
By softmax translation invariance, a = d((')) (k: k-)) =
- -z')) =: Pi-
The attention output (per sample) is
h 1. o v + a v oa x 21
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 GD on the two-key attention model
1: Input: {(x', y², )')}²= f encodings $1, $2, stepsizes {m.) (Be). initialization scale a, iterations T.
2: Positional Encoding:
2
=
x
%
0
3: Init: w
OM+dd
alı
4: for t = 0,1, 0..,.... T 1 do
5:
for i to do
pj s o( 1) (x 2 21)). .),)), g/. +- (v/x* + (1 -1- -1-)))))) (wr r gz m5²
6:
7: W 1+1 w' ltt Σ rl (p/*** +(1 (1-11pp)) (1-pp)z²) O
n
8: F + vt & Ei= r p/(1 pt) (x-- -2-)" ((x+))=2)(*
(x (x+ 22
n
9: Return: (w v v(+1).
Squared-paranen Head and Loss.
We predict with quadratic parameterization feed-forwer
layer:
M+d
-'1= (h ( 02
102 Σ
w h
(w, v)
1 y²)²
This quadratic parameterization feed-forward layer can be seen as a special diagonal neural network, essentially position wise FFN that provides anisotropic multiplicative gains and thus retains feature learning capacity through the attention mixed representation. This parameterization can be further
generalized by ÿ = (h++)" ((*+.t) @2 - (w-))) 62)
GD the two-key attention model is summarized in Algorithm
Our module is exactly a single-head dot-product attention applied per sample with two keys/values, one for the cause path and one for the descendant path. It the special case of Transformer atten- tion block where Wo, WK. Wv are identity projections, so the query is the learned gating direction v, and the tokens X and 2. This reduction keeps the softmax competition geometry and the value mixing mechanism intact while stripping away projection layers that would obscure the optimization dynamics. The quadratic parameterization head a diagonal, position wise FFN that provides nonnegative per-coordinate gains. Studying this minimal attention FFN pair is theoreti- cally meaningful: it isolates the allocation dynamics behind the implicit bias we analyze, preserving the key nonlinearities (softmax and multiplicative gains) that produce CCO.
The distinct fixed encodings $1 * 82 attach branch identity to keys and values and inject sample- independent margin (V (s) 82) into the logit difference. When x; and 2; are weakly separated
early in training. the offset (v TT() [s] $2) prevents the gate from collapsing to 1/2 and ensures a non-degenerate gradient, thereby guaranteeing identifiability of branches and stable training dynam- ics. This mirrors the role positional embeddings in Transformers.
4.1.3 DOMINANT-COORDINANN CONDITION.
We characterize which patterns of strong correlation are sufficient for CCO to emerge. The two conditions below formalize (i) population-level dominance one causal coordinate and (ii) per-sample margin along that coordinate.
Define 8, = [x(x (w)) (x, + =;)] measures the cross-moment between response y and the combined coordinate x, + zy. The adjustment μj == Ble(x, + 2,)] accounts for noise leakage. ofth s + μ the effective signal which governs the drift of gradient updates. We also write my = E x + 2 = Ejj + 23 and maj == E[(x +x) + 2;)] = Σnj + Suly which capture the second-monen momeent scales of the combined features.
Condition 1. The effective signal satisfies that sgll > 2m1 + maxx>1 (4 |16 m21).

---

Published as a conference paper at ICLR 2026
Condition requires effective signal the dominant feature is sufficiently strong to exceed that of other competitor by a uniform gap. The assumption is mild, it allows strong descendant induced corre- lations other coordinates but prevents the dominant causal direction from being overwhelmed. Under Condition | the GD dynamics preferentially amplify the squared weight on the dominant coordinate, creating the occupaney that initiates CCO.
Condition 2. There exist constant Ti, > O such that for every sample = =1,..., n: (i) Nontrivial gap: X zil 2 71- (ii) Sign stability: sgn(x) (x} zi) sgn(xi)- (ūi) Dominant-cordinee
margin lower bound: x 2 (r 1) Bx B, + 72-
In combination with Condition 1 Condition 2 guarantees that GD on the gate parameter V towards the max-margin solution on {x, } drives p + and thereby squeezes out the descendant branch. In short, Condition ensures occupancy, whereas Condition 2 ensures crowding out, com- pleting the CCO mechanism.
These conditions are satisfiable in bounded, Lipschitz settings. Importantly, as detalied in Ex- ample 26] they do not exclude the empirically relevant regime where some non-dominant causal coordinates less correlated with y than descendant coordinates: it can happen that for some j> 1 with wj = 1, Cov(x;,)) v) < Cov(2j.u).
4.2 MAIN RESULT
We next formalize when and how CCO emerges in our two-key attention model. Under the Dominant-coordinate condition, the first theorem provides a mechanistic account of CCO during training. The second theorem provides a generalization guarantee: with high probability, the learned predictor filters out the descendant 2, relies almost exclusively on the causal x, and attains test risk near the cause only level.
Theorem 1 (CCO's Mechanism). Under Condition 7 and Condition consider GD with initial-
ization scale a Yo loglogd/i d/n and the following stepsize schedule: (i) For S S Ti ==
43 minst E N : wi 2 11. set = and Be = 0. (ii) For Ti < S T + TT, with
T x exp (VIISI + (Bx + 3²) set Y) = 0 and Br = B. (iii) For T + T2 < s
T + T2 + T3 =: T" set The = and B = O with Tž
x log 7t Tointar)
Then, with
probability at least S. the squared-paramee head satisfies
w a Vogd for i supp(w"), W
w. S for & supp(w").
vn
Meanwhile, the query (gating) iterate q' = 5² obeys A = á logt + p', where ú is the max- margin solution on {x, -2:) and p¹ a bounded residual. Consequently, pi 2 1 for all isism n.
This theorem explains the mechanism by which CCO arises during optimization. Under the dominant-coordina condition, the dominant causal direction becomes visible GD: the gate's gradient aligns with the separation direction (x+ 2 and tracks max-margin ray with logarith- mically diverging norm, so the attention weight concentrates on the x-branch. As the gate filters out the descendant branch, the squared-paramette head fits the ground-truth weights w" up to the error on active coordinates and a 1/ tail on inactive ones.
Role of Positional Encodings. Distinct fixed encodings $ + $2 attach branch identity and in- troduce a sample-independent margin in the gate logit, (V (V') (si -2). This symmetry breaking enables the two-key attention identify the dominant feature and drive the attention weights to select the branch associated with it, thereby catalyzing CCO.
Theorem (Generalization CCO). For an independent test triple (x,3, (,,,)), 2), there exists an event =V
Ω with Pr(s)) 2 1 - 8 VIl=|B3+ ( (Bx+
2ln(2d)) such that conditioned on Ω,
|
o²
< r ro? log d
p = = - 2T)) 2 d
-
and
L
2
7

---

Published as a conference paper at ICLR 2026
With high probability (strengthened when ||s|β x d), the learned gate continues to prefer the causal branch test distribution, i.c. p² is bounded away from 0 and close 1. Moreover, the test loss approaches the cause only noise floor 0²/2 at rate O(r o² log d/n), indicating that the predictor essentially relies
while filtering
the test distribution.
Theorem 2 controls generalization when train and test share the same data distribution. We next show that the same CCO predictor remains robust under test time shifts that perturb y -
Corollary (Robust generalization under y y-> shifts). At test time, change the - mechanism so that z' = (V) + & and assume sup, |zy| I By There exists an event Q with Pr([2) 2
8
1 ( B
221(28)) such that conditioned on S
|
o²
ro roolog logd
D²' d(( -1 2
E d2
((xxxx') (w
(wT T (
Ω
S
5 FURTHER DISCUSSION
Positioning of CCO. CCO arises under purely correlational training with single environment, no environment labels, and no explicit invariance regularizers. Yet when a dominant causal correlation is present and GD's implicit bias takes hold, the learned predictor moves beyond correlation toward causation: it increasingly relies on causal features while largely discounting spurious correlates. Meanwhile, multi-environment invariance methods also seek causally aligned predictors, but they pursue this goal by explicitly leveraging cross environment heterogeneity.
When Can Transformers Learn Causation? CCO offers a concrete path to cause only behav- ior under standard Transformer training, but it is not unique, and its assumptions need not always hold. In practice, Transformers/LLMM. frequently exploit shortcuts and spurious cues (Bender et al. Du al. 2023] Tang al.) 2023] Jin et CCO also has limits: it benefits from a
Invariance
Causation
Approx
Apprex
Mulii-Ennironment
CCO
Invariance Learning!
2021
2024
Environmental lggterogeneity
strong causal correlation; when spurious cues are comparably strong or plentiful, single environ- ment ERM may still lean on them. In this regime, multi-environment invariance learning that plicitly leverages heterogeneity remains essential for causal generalization.
Strong Causal
Envir-Eavironment
Currelation
Environment
Correlation
Figure 2: Positioning of CCO.
Practical Insights. CCO suggests actionable insightss for training: (i) amplify causal alignment in data to widen the dominant causal gap; (ii) employ mild attention sparsity or large step schedules to accentuate strong features. These steps do not enforce invariance, but they increase the likelihood that standard training will self select a cause only solution when the data permit.
6 EXPERIMENTS
6.1 SIMULATED EXPERIMENTS
We realize the GD on the two-key attention model in Algorithm and present the simulation result in this section. We consider the case where the data are generated from the same causal chain - + z. The structural assignment for each variable is defined x ~ ((
= x" (w*) + c, z Cy + &, where c, E are independent standard normal distributed and we set w as an all-ones vector. The results are shown in Fig 3) We calculate the weight pL., and display E,ri, We then run GD for 5000 iterations with batchsize 64, and the dimension of data d {5. 10). We in the first 100 iterations corresponding to the occupation phase, while in the crowding out stage p remains 1, while w slowly decreases to the minimum value.
its average across the batch F
that p", increases rapidly to 1 in all cases
8

---

Published as a conference paper at ICLR 2026
2.0
||w 11.4-5
1.0
[|w Il.d=10
0.8
1.0
p,d=10
0.6
Average
wegght
occuration
crowding out
0.2
0.4
0.0
0.2
200
400
Iteration
600
800
1000
200
400
600 500
1000
(a)
(b)
Figure 3: Simulation results for the GD on the two-key attention model. (a): the curve of l|w w w*1le and the average of with d (5, 10). (b): the first component of w quickly reaches its
optimum during occupation phase, while the other components slowly approach their optima during the crowding-ond phase.
95
Pre
Pra
VI
Rea
Dechat
Prec -
ElHcannnee b4
--- Pruc -.0
Oracle
Dewc e 0.55
70
83
0.6
07
0%
0' 10
02
0.4 0.6
0.8
Strength
Text Strengsh
Figure 4: Experiments waterbirds dataset. (a): The test accuracy with bias strength Piest 0.02 bias strengths on DeiT-Small, ResNet34 and EfficientNet-BB trained across a full sweep of training bias strengths from 0.5 0.99. Oracle is the accuracy no-biased test data using meil-Small trained without bias. (b): The test accuracy with bias strength Ptest sweeping from 0.02 to 0.99 on DeiT-Small trained across full sweep of training bias strengths from 0.5 to 0.95.
6.2 EXPERIMENTS ON REAL DATA
Experiments on Vision Task. We consider an image object classification task on the birds. The
target is to classify water birds (Y = 1) and land birds (Y = 0 in the CUB dataset (Wah et al,] (011). To eliminate confounding due to foreground-backgroumm setting where one bird species species on the right side acts as the spurious bias 2, both appearing in the foreground. We set the bias strength in the train dataset to 0.9, i.e. Prain P( yly) 0.9. This ensures that any observed attention shift cannot be attributed to low-level feature quality differences (e.g., texture richness semantic complexity) between foreground and background.
asymmetry altogether, we introduced the left side serves as the true target label and another bird
The results in Fig consistently show that the causal features progressively occupy and crowds out the spurious features (whether background or another bird). We find that the attention map on the left bird raise rapidly in the first 50 iterations, while the attention map on the right side seldom changes, illustrating the occupation phase. By iter 500, attention is concentrated almost entirely on the left side, with the bird on the right side receiving near zero weight, marking crowding-owd These findings confirm that the observed behavior reflects genuine optimization-drive cause preference not artifacts of feature disparity.
We conducted fair experiments on Waterbirds using DeiT-Small (from timm with ImageNet pretrain- ing) alongside ResNet34 and EfficientNet-BA (from torchvision, also pretrained, with comparable about 20M parameter counts), training all models for 1,000 epochs at a learning rate of le-4 across a full sweep of bias strengths from 0.5 to 0.99. As shown in the Fig (a), DeiT-Small maintains signif- icantly higher accuracy at strong bias levels (e.g., 0.9), demonstrating that Transformers can better capture the underlying causal signal-lett left side bird type- deee-desptt overwhelming spurious correla- tions with right bird, suggesting an advantage over CNNs in leveraging stronger semantic features when spurious cues dominate.

---

Published as a conference paper at ICLR 2026
0.30
Train Steps
orgin oue
50% Train Steps
421-1778
99% Train Steps
0.20
1.3
0:15
6:15
0.10
(rowding vx
0.00
DVD
(h)
Figure 5: Experiment results on natural language task. (a): the test loss when mask the noun, adj, verb or their combination the text. (b): the saliency scores of each token when input hate this DVD, it's awful." to the model at 1%, 50%, 99% of the training steps.
We also added out-of distrribbiit (OOD) test experiments in Fig 4 (b). We constructed water- bird dataset with base spurious correlation of varying training bias strengths Prain. measuring test accuracy on OOD data where the test bias strengths Pies E [0. 1]. The curve reveals that when Prrain 2 0.9, test accuracy drops as bias increases, indicating that the model fails to learn the invari- ant causal feature (bird type) and instead relies heavily on the spurious background cue. However, once Pirin S 0.85, test accuracy rises significantly and remains high (above 95%), which is the hallmark of CCO: the model effectively crowd out the spurious features and learn the cause only prediction. Therefore, when the spurious correlation is under the threshold, transformer can obtain cause only predictor which exhibits robust generalization at test time.
Experiments on Natural Language Task. We conduct the sentiment classification task on the
Amazon reviews dataset (He & McAuley 2016) which consists of reviews from amazon. Here Y e {1. 2, 4, 5) represents the reviewer's rating, X denotes the associated adjectives and verbs, and Z indicates the nouns related to the product itself. We finetune the
model Devlin et al. (019)) for 50k steps, employing the Adam optimizer Kinggaa & Ba (2014) with a learning rate of Ic-5. When constructing the test data, we mask the noun, adj, verb or their combination in the text. As shown in 5 test loss with masked NOUN+VERS decay rapidly corresponding the occupation phase. We also observe a final upward trend the test loss with masked ADJ+VERB, indicating that the attention allocated to NOUNs is being crowded out by cause features. Fig (b) display the saliency scores computed by the gradients target class score relative to input embeddings, which show which tokens most influence the model's decision. The result indicates that the cause features (hate, awful) crowds out the spurious features during the training process.
7 CONCLUSION
In this paper, we identify a new training phenomenon for Transformers training dynamics called CCO, showing that strong causal alignment in the data, coupled with the implicit regularization of GD, can drive the model toward cause only prediction. We demonstrate CCO empirically and de- velop a theoretical account of its two phase mechanism (occupation and crowding-out) While not the only route to causal generalization, CCO offers a concrete answer to when and through what dynamics standard Transformer training can suppress spurious features and rely almost exclusively on causal ones. The results spark that: amplifying causal alignment in data and designing train- ing procedures that accentuate causal signals can make Transformers more likely to learn causally grounded predictors.
ACKNOWLEDGMENTS
This work is supported by the NSF China (No.s 92470117 and 62376008).
10

---

Published as a conference paper at ICLR 2026
9
ETHICS STATEMENT
Our paper complies with the ICLR Code of Ethics.
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 2019.
Davoud Ataee Tarzanagh, Yingcong Li. Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Neural Information Processing Systems, 2023.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical bias-variame trade-off. Proceedings of the National Academy of Sciences, 116(32): 15849)31584911585 2019.
Emily M Bender, Timnit Gebru, Angelina McMillan-Major and Shmargaret Shmitchell. On the dangers of stochastic parrots: Can language models be too big? Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, pp. 610-623, 2021.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In International Conference on Machine Learning, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Neural Information Processing Systems, 2020.
Yuhang Cai, Kangjie Zhou, Jingfeng Wu, Song Mei, Michael Lindsey, and Peter L Bartlett. Implicit bias of gradient descent for non-homogeneous deep networks. arXiv preprint arXiv:2502. 16075, 2025.
Francesco D' Angelo, Francesco Croce, and Nicolas Flammarion. Selective induction heads: How Transformers select causal structures in context. In International Conference on Learning Repre- sentations, 2025.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Mengnan Du, Varun Manjunatha, Rajiv Jain, Ruchi Deshpande, Franck Dernoncourt, Jiuxiang Gu, Tong Sun, and Xia Hu. Towards interpreting and mitigating shortcut learning behavior of NLU models. In Proceedings of the 202 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 915 915-929, 2021.
Mengnan Du, Fengxiang He, Na Zou, Dacheng Tao, and Xia Hu. Shortcut learning of large language models in natural language understanding. Communications of the ACM, 67(1):100-120, 2023.
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. The Annals of Statistics, 4998): 2021.
Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evo- lution of statistical induction heads: In-context learning markov chains. In Neural Information Processing Systems, 2024.
Jianqing Fan, Cong Fang, Yihong Gu, and Tong Zhang. Environment invariant linear least squares. The Annals of Statistics, 52(5):2268-22) 22922, 2024.

---

Published as a conference paper at ICLR 2026
Ragnar Frisch and Frederick V Waugh. Partial time regressions as compared with individual trends. Econometrica: Journal of the Econometric Sociery, pp. 387-401, 1933.
Rui Gao, Xi Chen, and Anton Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Research, 72(3)111771119,, 2024.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained CNNs are biased towards texture; increasing shape bias im- proves accuracy and robustness. In International Conference on Learning Representations, 2019.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 ((1))665-663) 2020.
AmirEmad Ghassami, Saber Salchkaleybar, Negar Kiyavash, and Kun Zhang. Learning causal structures using regression invariance. In Neural Information Processing Systems, 2017.
Yihong Gu, Cong Fang, Peter Bühlmann, and Jianqing Fan. Causality pursuit from heterogeneous environments via neural adversarial invariance learning. The Annals of Statistics, 53(5):2230- 2257, 2025a.
Yihong Gu, Cong Fang, Yang Xu, Zijian Guo, and Jianqing Fan. Fundamental computational lim- its in pursuing invariant causal prediction and invariance-guided regularization. arXiv preprint arXiv:2501.17354 2025b.
Shurui Gui and Shuiwang Ji. Mitigating spurious correlations in LLMs via causality awwrr post- training. arXiv preprint arXiv:2506.09433, 2025.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A Smith. Annotation artifacts in natural language inference data. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguis- ties: Human Language Technologies, Volume (Short Papers), 2018.
Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. Conference on Learning Theory, 2021.
Ruining He and Julian McAuley. Ups and downs: Modeling the visual evolution of fashion trends with one-class collaborative filtering. In Proceedings of the 25th International Conference World Wide Web, pp. 507-517, 2016.
Christina Heinze-Demi. Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models. Journal of Causal Inference, 602):20170016, 2018
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision Transformers provably learn spatial structure. In Neural Information Processing Systems, 2022.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference Learning Theory, 2019.
Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason Lee. Understanding incre- mental learning of gradient descent: A fine-grained analysis of matrix sensing. In International Conference on Machine Learning, 2023.
Zhijing Jin, Jiarui Liu, Zhiheng Lyu, Spencer Poff, Mrinmaya Sachan, Rada Mihalcea, Mona T. Diab, and Bernhard Schölkopf. Can large language models infer causation from correlation? In International Conference Learning Representations, 2024.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk mini- mization capture invariance? In International Conference on Artificial Intelligence and Statistics, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 44126698, 2014.
12

---

Published as a conference paper at ICLR 2026
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are sero-shot reasoners. In Neural Information Processing Systems, 2022.
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vi- sion Transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations, 2023a.
Jiangyuan Li, Thanh Nguyen, Chinmay Hegde, and Ka Wai Wong. Implicit sparse regularization: The impact of depth and early stopping, Neural Information Processing Systems, 2021.
Yicong Li, Xiang Wang, Junbin Xiao, Wei Ji, and Tat-Seng Chua. Transformer-empowered invariant grounding for video question answering, IEEE Transactions on Pattern Analysis and Machine Intelligence, 47(11):5510-952) 2025.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-varameee wer-parameterized matrix sensing and neural networks with quadratic activations. In Conference on Learning Theory, 2018.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do Transformers learn topic structure: Towards mechanistic understanding. In International Conference Muchine Learning, 2023b.
Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter Bartlett, and Jason D. Lee. Scaling laws in linear regression: Compute, parameters, and data. In Neural Information Processing Systems, 2024.
Yong Lin, Hanze Dong, Hao Wang, and Tong Zhang. Bayesian invariant risk minimization. In Computer Vision and Pattern Recognition, 2022a.
Yong Lin, Shengyu Zhu, Lu Tan, and Peng Cui. Zin: When and how to learn invariance without environment partition? In Neural Information Processing Systems, 2022b.
Chaochao Lu, Yuhuai Wu, Jose Miguel Hermandez-Lobato and Bernhard Schölkopf. Nonlinear invariant risk minimization: A causal approach. arXiv preprint arXiv:2102.12355. 2021a.
Haoye Lu, Yongyi Mao, and Amiya Nayak. On the dynamics of training attention models. In International Conference on Learning Representations, 2021b.
Weihao Lu, Haobo Zhang, Yicheng Li, Manyun Xu, and Qian Lin. Optimal rate of kernel regression in large dimensions. arXiv preprint arXiv:2309.04222 2023.
Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. In Anmual Meeting of the Association for Computational Linguistics, 2019.
Nicolai Meinshausen, Alain Hauser, Joris M Mooij, Jonas Peters, Philip Versteeg, and Peter Bühlmann. Methods for causal inference from gene perturbation experiments and validation. Proceedings of the National Academy of Sciences, 113(27)7361-736) 2016.
Mazda Moayeri, Sahil Singla, and Soheil Feizi. Hard imagenet: Segmentations for objects with strong spurious cues. In Neural Information Processing Systems, 2022.
Depen Morwani, Jatin Batra, Prateek Jain, and Praneeth Netrapalli. Simplicity bias in 1-hidden layer neural networks. Neural Information Processing Systems, 2023.
Eshaan Nichani, Alex Damian, and Jason D. Lee. How Transformers learn causal structure with gradient descent. In International Conference on Machine Learning, 2024.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint 2209. 1895, 2022.
Judea Pearl. Causal inference in statistics: An overview. 2009.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant pre- diction: Identification and confidence intervals. Journal of the Royal Statistical Society Series B: Statistical Methodology, 78(5):94--1012, 2016.
13

---

Published as a conference paper at ICLR 2026
Niklas Pfister, Peter Bühlmann, and Jonas Peters. Invariant causal prediction for sequential data. Journal of the American Statistical Association, (14 114(527):1264 1276, 2019.
GuanWer Wen Qiu, Da Kuang, and Surbhi Goel. Complexity matters: Dynamics of feature learning in the presence spurious correlations. NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk mini- mization. In International Conference Learning Reprexentations, 2021.
Dominik Rothenhäusler, Peter Bühlmann, and Nicolai Meinshausen. Causal dantzig. The Annals of Statistics, 47(3):1688-1722 2019.
Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, and Jonas Peters. Anchor regres- sion: Heterogeneous data meet causality. Journal of the Royal Statistical Society Series B: Statis- tical Methodology, 83(2):215-246 2021.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. In International Conference on Learning Representations, 2020.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 10965)x612-33, 2021.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In Neural Information Processing Systems, 2020.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algo- rithms. Cambridge University Press, 2014.
Alexander Shapiro. Distributionally robust stochastic programming. SIAM Journal on Optimization, 27(4):2258 2275,
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70): 1-57, 2018.
Ruixiang Tang, Dehan Kong, Longtao Huang, al. Large language models can be lazy learn- ers: Analyze shortcuts in in-context learning. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 4645-4657, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023.
Maya Varma,
Delbrouck, Zhihong Chen, Akshay Chaudhari, and Curtis Langlotz. Ravl: Discovering and mitigating spurious correlations in fine-tuned vision-language models. In Neural Information Processing Systems, 2024.
Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. In Neural Information Processing Systems, 2019.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. Transactions on Machine Learning Research, 2025.
Roman Vershynin. Digh--iimssional Probability: An Introduction with Applications in Data Sci- ence. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, Cambridge, edition, 2026.
C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The caltech-uesd 2tls-2200-201 dataset. Technical Report CNS-TR-201 1-001, Californa Institute of Technology, 2011.
14

---

Published as a conference paper at ICLR 2026
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Neural Information Processing Systems, 2022.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, 2020.
Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last iterate risk bounds of SGD with decaying stepsize for overparameterizee linear regression. In International Conference on Machine Learning, 2022.
Jingfeng Wu, Peter Bartlett, Matus Telgarsky, and Bin Yu. Benefits early stopping in gradi- ent descent for overparameterized logistic regression. In International Conference Machine Learning, 2025.
Yang Xu, Yihong Gu, and Cong Fang. The implicit bias of heterogeneity towards invariance: A study of multi-environment matrix sensing. In Neural Information Processing Systems, 2024.
Yihao Xue, Siddharth Joshi, Eric Gan, Pin-Y Yu Chen, and Baharan Mirzasoleiman. Which features are learnt by contrastive learning? On the role of simplicity bias in class collapse and feature suppression. In International Conference on Machine Learning. 2023.
Yu Yang. Fric Gan, Gintare Karolina Dziugaite, and Baharan Mirzasoleiman. Identifying spuri- biases early in training through the lens of simplicity bias. In International Conference Artificial Intelligence and Statistics, 2024.
Wenqian Ye, Guangtao Zheng, Xu Cao, Yunsheng Ma, Xia Hu, and Aidong Zhang. Spurious corre- lations in machine learning: survey. arXiv preprint arXiv:2402. /2715, 2024.
Yu Yuan, Lili Zhao, Kai Zhang, Guangting Zheng, and Qi Liu. Do LLMs overcome shortcut learn- ing? An evaluation of shorteut challenges in large language models. In Conference on Empirical Methods in Natural Language Processing, 2024.
Haihan Zhang, Weicheng Lin, Yuanshi Liu, and Cong Fang. Learning curves of stochastic gradient descent in kernel regression. In International Conference on Machine Learning, 2025.
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 mil- lion image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(6); 146,1144521144, 2018.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift. In International Conference Machine Learning. 2021.
Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang- Sparse invariant risk minimization. In International Conference on Machine Learning, 2022.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Benign overfitting of constant-stepsize SGD for linear regression. Conference on Learning Theory, 2021.
15

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 USE OF LLMS
We used LLMs for language polishing.
A.2 ADDITIONAL EXPERIMENT
(a) Input image
(b) Iter=O
(c) Iter=500
(d) Iter=1000
Figure 6: This figure shows how attention shifts during ViT training on a background-labe (target)
task. (a) is the input image. Early in training (b, iter 0), attention is diffuse across background (causal feature) and bird (spurious feature). As training proceeds (c, iter 500), attention weight rises on the background illustrating the occupation phase. By (d, iter 1000), attention is concentrated almost entirely on the background, with the bird receiving near zero weight, marking crowding-out.
Background-label (target) task. We consider image object classification task the background with birds. The target is to classify water environment (Y = 1) and land environment (Y 0). We generate datasets by combining the bird images in the CUB dataset (Wah et al. 201T) and the background images in the Places dataset (Zhou et al. 2018) using specific probabilities, which is similar to the waterbird setting in (Sagawa al. (2000) with different target. We set the pixels related to birds as and place 70% of all water birds against a water background and 70% of all land birds against a land background, generating a dataset with 30k images. We then train the vision Transformer model (Dosovitskiy et al.) (202]) using the dataset, fixing the input image size to 224, with patch size set to 16, learning rate set Ie-4, and batch size set 16. The results are displayed in Fig(o Initially, attention grows rapidly in the background with only slight increase on the bird. Later, during crowding-out, the map is rapidly dominated by background attention, while the bird's attention becomes very faint.
Sensitive of sign-stability. In order to empirically verify how sensitive our mechanism is to mild violations of sign-stability, we consider the image object classification task the background with birds. The target is to classify water bird (Y = 1) and land bird (Y = 0). We generate datasets by combining the bird images in the CUB dataset (aah et al. 2011) and the background images in the Places dataset (Zhou et al. 2018) using specific probabilities. But here we flip the label with probability Piip:
[ ---,, with probability phip:
Y. with probability 1 Ptip
We also place Piran of all water birds against water background and Pusin of all land birds against land background, generating dataset with 30k images. We then train the vision Transformer model using the dataset, fixing the input image size to 224, with patch size set 16, learning rate set to le-4, and batch size set to 16. As shown in Fig 7 we scan the Peip from 0 to 0.5, and find that when Phip 0.2, the crowding-out behavior can still be observed in the model with the accuracy reaches over 90%, proving that the mechanism is sign-stable When Pflp < 0.15, the test accuracy remains robust across different bias strengths, indicating that CCO is still effective and the model is able to leam invariant cause only prediction.
CCO boundary conditions. To empirically characterize the boundary conditions of CCO, we intro- duce controlled cause- predictvee correlation by flipping the true bird class label Y with probability Pnip in a Waterbirds- setup, where water birds and land birds are composited onto matching backgrounds with strength Punin (measuring spurious-predictivs correlation). As shown in Fig 7
16

---

Published as a conference paper at ICLR 2026
103
40
Rgy
6m 65
10
Pg-00 02
0.3
0.0
Pip 0.02
Pap 0.0
05
0.2
24
0.5 100 155
02
0.A
0.5
15
(a) Puin
0.7
(b) Prain 0.9
Figure 7: Test Accuracy on Deit-Small trained across full sweep of label flip probability Ptip from O to 0.5 with Puain fixed to 0.7 and 0.9. We sweep the test bias strength Piest from 0.02 0.99.
0.8
Attention on
Attention
Text Accuccc
80-
0.7.accground 0.7.oralee Prain
backkgrundd Porin 0.9.orade
0.5
0.4
0.00
0.05
0.10
0.15
0.20
0.25
0.30
Pree
(a)
(b)
Figure 8: (a): Test accuracy of Deit-Small trained dataset with various Ptti and Puain the bird classification task. The background curves show test accuracy on the background dataset without bird. The oracle curves show test accuracy the dataset with Piest = 0.5. (b): The dynamic of attention weight for x, z during training in simulation on standard transformer.
when Pnip < 0.3 in (a) or Pfip S 0.05 in (b), the Vision Transformer maintains high and stable test accuracy across varying bias strengths, and its attention concentrates on the bird rather than the eevidre
background evidence that CCO is active and the model learns an invariant predictor based on the
causal features x. Even at Pnip = 0.2 in Fig 7 (a,, the model still achieves over 90% accuracy and exhibits crowding-out behavior, demonstrating robustness to mild violations of sign-stability. How- ever, as Ptip increases further toward 0.4 or 0.5, the performance deteriorates sharply, indicating that CCO collapses once the correlation gap falls below critical threshold. These results establish that CCO operates effectively when the pruse-predictive correlation remains sufficiently stronger than the spurious one, defining practical boundary beyond which the mechanism longer reliably emerges.
Generalization efffecs of spurious-dominant correlation. We test the accuracy of ViT on the background dataset without bird in specular setting with dominant correlations between and z, and weaker ones between x and y. The results are shown in (a). When the correlation between and z in the dataset is strong, we observe specular result: the model achieves high accuracy on the background-onle test set, indicating that primarily relies on features associated with z for prediction. In contrast, when the correlation between y and is relatively weaker compared to the correlation between x and y, the model's s accuracy on the background-only test set becomes very low. In this regime, the CCO mechanism emerges: the model's ''s attention focuses predominantly causal features x, effectively crowding out those spurious features z.
Image classification on CelebA. We conduct our experiment on the classification task on CelebA dataset. This classification task aims to predict the presence of beard from CelebA images, where the target label is spuriously correlated with gender. We trained ResNet-34, EfficientNet-B4 and DeiT-Small with comparable parameter counts on this dataset under standard settings, using the Adam W optimizer with learning rate of 1e-4.
17

---

Published as a conference paper at ICLR 2026
Table 1: The accuracy of ResNet-33, EfficientNet-BA. and DeiT-Small on train dataset and test dataset including (1) (Test Set 1) masking out facial regions, and (2) (Test Set 2) masking out every- thing except the facial regions.
Model Deit-small
Train Accuracy
Test Set 1 Accuracy (L)
Test Set 2 Accuracyy (1)
0.987
0.552
0.893
ResNet-33
0.992
0.577
0.861
EfficientNet-B4
0.979
0.573
0.802
Table 2: The loss of BERT trained on datasets with various P( y). The table shows the loss on train dataset and test dataset including (1) (Oracle Test): and are independent, P(z v) = 0.5, and (2) (Biased Test): P(= y) = 0.02.
P(z y)
Final Train Loss
Oracle Test Loss (L) Biased Test Loss (4)
0.5
0.64
0.61
0.65
0.9
0.62
0.65
0.68
0.95
0.67
0.77
0.87
In the test set, we evaluated two masking conditions based on the bounding box (bbox) annotations provided by the dataset: (1) (Test Set 1) masking out facial regions, and (2) (Test Set masking out everything except the facial regions and the result are shown in Tab I On Test Set 2, DeiT-Small outperformed both ResNet-34 and EfficientNet-Bi indicating the CCO mechanism crowding out spurious features for accurate beard prediction. The performance gap observed on Test Set 2, where only facial regions are visible, underscores that when the dataset contains strong but misleading associations (like gender bias), DeiT-Small leverages its capacity to attend to all parts of the image equally and identify the most predictive elements- the beard itself-thus achieving higher accuracy. This supports the hypothesis that under certain conditions, partieularly those involving complex spurious correlations, Transformers exhibit robustness and adaptability that enables them to focus on invariant cause only prediction, enhancing their generalization capabilities on unseen data.
Controllable Spuriousness evaluation in NLP task. We construct NLP evaluation settings where the degree of spurious correlation is known and controllable. When ground-truth labels y are avail- able and the data- metta-enerating process allows intervention, we can deliberately manipulate the associ- ation between a potentially spurious variable (c.g., the name of item) and the label y. By sampling instances according to a fixed conditional distribution P(2 y) = we can break or calibrate the spurious link between and y.
The Amazon reviews dataset provide the label of scores which is the target y and the name of item, which is a measurement of Varying p across experimental conditions allows systematic study of how model behavior changes with the strength of the z-y association. The table shows the final test loss of BERT under various p. where BERT remains lower test loss when p = 0.9, demonstrating that transformers pick up the stronger causal signal in NLP data.
Simulation on standard transformer. We conduct simulation experiments on standard multi-token transformer. We take two-token X = [x, z] as the input, and the causal chain is + y + where x, 1 E Rdd are vector covariates. We set
y x= x W2
z = Way +
Here and { are both Gaussian random vectors, with variances of 0.1 and 1, respectively. We set W, = Id, w 0.1 1a. We then train a 2-layer standard multi-token Transformer with a learning rate of le-3 and the dynamic of attention weight for x, z during training shown in Fig g The attention weight curve demonstrates that the model initially assigns comparable attention to both the cause X and its effect Z, but shows a two-stage shifts focus toward X while sharply suppressing attention to Z. This "occupation" and "crowding out" behavior aligns with the CCO mechanism. Consequently, the model learns to rely on direct evidence rather than attending to indirect, spurious predictive pathway.
18

---

Published as a conference paper at ICLR 2026
A.3 PROOF OF THEOREM 1
The proof of Theorem in three stages. In Stage I, we show that the squared-parameter FFN rapidly amplifies the weight on the dominant causal coordinate while keeping all other coordinates small. In Stage 2. GD on the gate parameter V towards the max-margin solution on {x; ) drives pi and thereby squeezes out the descendant branch. Finally, in Stage 3, after the descendant is nearly excluded, the parammeee FFN recovers the sparse ground truth.
A.3.1 STAGE
Theorem 3. Let TÝ min {tE N: w 2 4).
Under Condition 1 suppose the step sizes satisfy
3 21 y and Br O, and the initialization scale is a o logd/n
Ve = 1) <
Then with
probability at least 1 - 22. the iterate w
1
0. w
satisfies
S w/ S 2 and [w Ti s
Lollond/n oo
for j >
d
Proof of Theorem 3 Throughout Stage we set Br = 0, hence V' = - for all 1 T7. With the
two-key attention, this implies p (0) 2/2 for every sample S i S n and iteration 1 S 1 T7.
By the structure of the ganared-paramett head, we keeps the form w [!
The update of w' satisfies
w We analyze non-dominant coordinates (j > 1) and the dominant coordinate (j 1) in turn.
?
[ (x' + +z) T (w') 02 - (2')022(66))
(0-))w===-000 - e/ es) (x' + +2") o ww. (1)
By the defination of Ti, we have w < 1/4, for & < log( (dd /aly o log d/n))
TT. We now prove for j > and 1 S 5
Ti f log d/n holds by induction.
AT. K 16 11
mi
log(1 +1))
32
The update for coordinate j >>> 1 is
T2
n
w/
17
1
xxf + (x} + z (w))
17 > (x') (w" (x) 2 w3j
+ (x) + 2
Σ
w/
1
(x ) w (x(x + w/
2n
i=1
k=
Separating population terms from sampling deviations gives the multiplicative form
w (
mij
(wi (w + 4j.
2
where O is expressed as the following:
1
1
Σ
(xf + (x) + +z))= - mıj
(wi) (wi)²(wg (w²)
2
n
+
x
T
(w*) 62 (x) + 21)) - 8) w/ +
Σ
(x) + z;) - Hg wl
n
1
(x} + 21{ (wi) (xx) + +z) w/.
2n
By Lemma 6
(concentration) and boundedness, using w 1 and the inductive hypothesis /w/[ s
eaz log d/n. 12
obtain
Φ1 |w]]
1
A < (02 + 0t)
[w]l +
+
|wil = e, w5|w3h
(3)
32
2d
19

---

Published as a conference paper at ICLR 2026
where eg equation
3) and, for large d, leyl - 1/2 maxxx>1 (15/th + ). Combining equation 2
and
3 w 1+1
(1595
Mij
2 1 11 17
max
+
+ Ej
[w]
jil
32
(4)
(1o9tt
mi1
2 1 +
max j>l
+
32
By the above inequality, /w)"
++1
I can be bounded by
I 1
(1)
+ 17
5
max j>l
[199t]
+
m1,
[w]
4
32
loxid /ary /alvoollu log d/m)) log(1 +7p)
(2)
1 + 17
5
max 4 f31
eff
15
V o log d/n
32
d3
(5)
logid /aty (ond/m)) (oR(1---1))
5
[155
myj
02 log d/n
1 + 17
max
+
32
d3
v 02 log d/n
d2
0 log d/n dB
where (1) uses equation 4 (2) uses
and (3) uses Condition 1
Therefore, by induction, for all j > 1 and all 1 St t
loggdd /A1 Toy 1-77)))))))))))))))))))))))
ATT, we have /w/] S
o log d/n
de
We then prove for j = 1 and S 1 <
log(dd /1) log( 11777
A Tr. w/+1 2 (1 + 77:)w/ where x
15 cff mi 16
32
1+1 is given by The update for w/+1
n
w
7
+ 2))
(w/)² + [ Σ (x*) (xx)T(W>)) (w )02 + +z)) wi
2
1 d
+ L/ Σ
(x} + 21) wi
(x)
+ +z/wi
2n
1=1 1=
Separate expectations and deviations, we have
w w
w + 77 (s wi
mi
(wi)²)
+ 17 A|.
n
n
A = 2 11
x + 2))
m1 (wi)3
3
(w") 32 (xf + 21) 81 w
1=1
+
1
Σ
+
21 11
w
1
(x) + 2,) ( (w/)) (xi + 21) wi.
2n
1=1 j=2
By Lemma and boundedness,
B3 d
Φ1
A < (Φz + de) (wil +
/ 3 +
(w²)2
/wil-
(6)
2
2
1=2
Since for j > 1, /w// < v o log d/n d²
and 0 < wi
< 1/2, we have _ < 1/d. Hence by
equation 6
B3 x=
m1 + 01 )
2 1 + 17
02 0.
2 + 178
2d
32
20

---

Published as a conference paper at ICLR 2026
log(d"/44 /4ly/4////// / log d/nl)
ATi. wf 2 (1+ + -))) y Voologd logd/n d/n Therefore,
This implies that for S -
log( 1-7))
log(d" /4( o0 log d/n))
V 1o d/n for j >
we obtain that
1 = Then we have |w S
log( 177)
and W TT-1
S 1. Then,
w
w S 1 + 17 59 + 02 + de +
B
+
mi + Φ1
(7)
2d
32
Since satisfies y) 81 + 22 + th +
x+*
ma t:db
S 1, we have w
- 2
2d
32
Consequently, we obtain $ S w Ti
and
Ti W
o² log d/n for j >
The lemmas required for the Theorem 3 are listed below.
Lemma 4, For any i n and 1 S j d, (llf(()) (0) + (r B.) Be = B2.
Proof of Lemma 4
Because w E {0, 13dd and (supp(w") ST,
(xx') (w*)
Σ
x)
Σ
Bxx
(8)
jEsupp(w*)
jEsupp(w*)
Together with |e|| s B.,
|n'l = (xx') (w") +
$ rB + B..
(9)
For any coordinate i d,
(o((') < (f, + (0) - f3(0)) &
f(0)lo & lf,(0) L
Taking supy and using equation 9 yields
Ilf()")l>> $ [l + Lr + B.).
(10)
By the triangle inequality,
zz)l
15 + +E s 1/300') +1 S 11 + Be-
Combining with equation 10 gives
|z)| s [lf(0)]l + L(BBx + -B.) B
and thus ||z'|oo S B2 almost surely.
Lemma (Bernstein inequality for bounded distributions (Theorem 2.9.5, Vershynin (026))). Let (E:) be independent, EE; l&il < M a.s., and Var(&) S Then for any t
nf?
Pr
E
2
< 2exp
20 2/2 M1
1=1
hence with probability 2 6,
2o log(2/8)
2M log(2/6)
3 Σ
+
(B1)
n
3n
i=l
21

---

Published as a conference paper at ICLR 2026
Lemma 6. Define Bx+2 Bx + Bg. For any ò (0,1), with probability least 1 6, the following hold simultaneously for all $ j d:
(x} + (x) + 25) mss)
<Φ = 2B (V 21og(6d/0)
2 log (6d/6)
max 1SRJ Sd
n
3n
1=1
2 (v 221og(6//0)
log (od//)/))
max
((x))
(w )02) (x) + 23}
S 02 = r Bxx Bxt2
ISiSd
π
3n
1=1
23) - Hj I Φ = Be Bxtx (V
(21og(6u/b)
log(6d/6)
max 151Sd
n2
3n
where Bx-z:= Bx + By and the factor log log(6d/6) accounts for union bound over d coordinates and the three families.
Proof of Lemma 0 For fixed 1 S k,j S d, set set = (xl xx(x) + 25) maj- Then |=(1 S
M2 == 2B x+z and Var(= < == M2. Applying Bernstein and union-bounding over j with probability 2 8/3.
(x) + 22) - -ms) s s0I
max
i=1
where
2 logg(6d/0)
2 2log(6d/5)
Φ1 2B2 x+x
3n
Let S; = (x'))(x''))) (w*) 02 and E.20 = Si(x) + =2,) - 83- Since (w~)02 E {0, 1]d with support size
S r. [St] S rxx, hence E 2) S MM == rB Bxtx and Var(=. (( S UM = MM.. Bernstein
union bound gives with probability at least 2 - 8/3
(xx)) (w )2) (x) + zz)) 8j S 62,
max
11
with
2log(6d/0)
2 (6g(/)/))
Φ = r Bx Bxtx
n
3n
Ler = (x² -2))-1433 Then (
< M. = B, Bx and Var(=())) S == M². Bernstein
plus union bound yields with probability
8/3
(x) 21) 14j S de:
where
(v
22og(6d/88) log(6d//))
de = B. Bxtz
+
3n
A.3.2 STAGE
Notation and max-margin solution.
Set = 81 S2 and define
E Md
||u|| S v |>s13 + d + +B)2 a.s.
11
x 2
22

---

Published as a conference paper at ICLR 2026
Let {u') Y== be the samples form and consider the empirical &g max- max-margin separator
û E arg min
WERM 2 l s.t. (u')Tw 2 1, i ===1,..... n..
Its empirical margin is "emp == // Let S S [n] be the support set with (u') û 1 for i E S.
By the KKT conditions there exist multipliers o; 0, nonzero only on S, such that
ù Σ a.u"
iES
|
Lemma 7. Define Usep ==
For every i, (u') usep = 1, hence Usep satisfies all margin
constraints and separates the sample from the origin with margin 1. Since minimizes ||u||| over the feasible set, ||اا|2 $ ||seep 2 1/||s|12. and thus
2 |||||2-
Vemp
Proof of Lemmal 7 Since u² 1 (s; x² x² and usep = [s/|||| 0). (u²)T Usep TT (s/1[s|3)
Thus Usep is feasible and the norm bound follows.
By the properties established in Lemma each sample gradient can be written as (V,(v') Vol(v T u'). Each Ø; is monotonically decreasing to zero, Car-smooth, and has a (C,, n.) tight exponential tail. Hence, Phase II the implicit bias of GD drives the direction toward the & max- margin solution while the norm diverges, which forces the gate weight P to converge to one. We formalize this as the following Theorem [ The proof follows [Soudry al.] (2018) Thm. 9), with minor adaptations to handle sample dependent Φ; and the dominant coordinate condition.
Theorem 8. Let Th de Vo2 log d/n. Under Condition 2 suppose the step sizes satisfy 7pe O. 3t = B < Ogn (U)2 where Con is a constant defined in Lemma 9 there exists bounded residual p¹ such that
4' á logt + p',
and, in particular,
ú
lim
min(v (T logt - +00.
V 2
|
Consequently, T Ty, P
Ti +
(V T uut)= 2 V d/n for all
Proof of Theurem We first prove that can be expressed by V p'. Since p' =r and I' V á log t ú, then
logt l + p² with bounded
(ist1-r9) | |10+1 - r' -r( -
+' 11=2112.
(11)
By Lemma (A), we have o ((V4) un < O. By the definition ú, we have û un 2 1. This implies
1 1+ 1 1)
1++1 - r'1 = (V - á log
= 32 VL(V-)| + ||ا1 log (1 + (1++)) ++))) 2
+ 23n VV(VV) log 1 (1+ 1) (12)
1
< 8 VC(V*)
+ - ||]l
By Lemma 9 (B) and Lemma 10 we have
VL((V) = Co < 00,
lim | 0.
(13)
23

---

Published as a conference paper at ICLR 2026
Therefore,
20
|| - 1² =0 (1),
Σ
||r - r'1 = Co < 20.
(14)
[=
11 fort t > ti, (r'+1 - r²) _ r² v Ct min 1--50mmmmmmm))
By Lemma
Thus,
||-1| |
|| p - r011 + 2 (rrtt - r²) nh
k=
(15)
I Co +2
Ck
10.
k=
00.
This implies ||=*]| is bounded and further ||||| is bounded.
In Stage 2. w' = whi is fixed and the gate updates e follow:
ri+1 v n (V'),,) रh O.
Ve((T'=== = (fixx + (1 - p$) z²) T (w²) 02 - -(x') T (w-)02=
(047)2-(xx)
(16)
(x' 2 (wTi)
@2
p (1 -ri) (x' - -2)).
where p = a ((v*) (x' - z')), and att) lte-e.
1
(V(("") can be further expressed
VC((")
19 (xf zi)
Ti Ci x + a + C2
(xi -zi) ( (w.""
+ Gi
p (1 - p|) (x' 2²)
(17)
pl (xxi zi)² (wit)) + C3 (xi di) (xi zi) (w
+ t
pi (1 - pi) (x* -))-
where Çi, C/. S3. Ci are small quantities, and ej is constant.
Gi Σ (x) 21)
- = (x)) (w])) 2
)
1=2
(w"") 5i + (c))², d (xf - -zi) (w TT)
C3 = 2 (x- z1)
32 -x-
(xi -f
of z (w))
1i
Σ (x) (w;) 2
(w;)²-()
Let 0 u²) L((vt) such that
6 (u) (a()) (x² 21)² (w/tt)
+ + - -xx (xi-c) 21) (w/+)² + c/
D (w) (1 σ (u))
24

---

Published as a conference paper at ICLR 2026
Define
[(V) = 1
o v n'), u' x 2²
(18)
Then VL(()- _
n Σ (V') which is equal to the gradient of Z(V) :- ((V) - inff L(V) 2 0.
Without loss of generality, we assume infy L(V) = O.
Since W1 Ti
[1/4, 1/2] and |
S
o log d/n
for all > 1. For each sample i, define and
d²
|x]] S Bx [z)] S B2 |e"| S B. a.s. Then
K v Bx By
B₂ 1 S
C3
G S C4 d
I -
d
d
d
for explicit constants C3, C4 depending only on (Bx, B2, B.).
Let
A = (x] - (w Bo, = (x) - (G) (x] - (w,
and absorb the small remainders into
C3 C4 d
(3(w) = (w) G + Ci K((WW] S
for all u (since 0 S s((u) S 1)).
Then the scalar driving term becomes
o(u) Ite Lemma 9 (Properties of (u())) (u) has the following properties:
(u( = (o(u) 4 - Bo, + (ō(u)) ((u) -u)),
(A) Monotonicity. For all and u, (1) < O.
(B) Second-derivative control. For all u, (1)] S Car, hence
|V>L(V)] S Car
max (U)2
U:= [u u "]
where
Car
1 5[ + B2)²
+ 7 (Bx
1
1
1Bxx
Bz +
+r Bx + B.) (Bx + Bz) +1).
64
$
4
(C) Exponential tails. Let
fi(u) (b((u) = (Bu, o(u)A
(w)
there exist C; > 0, H4; > O and uo, such that. all u > uo.i.
C2( (1 - e eMiue S fi(u) S C₁ t keu)e-
Proof of Lemma 9 (A) Note that
Bo,i A (xi at) (xi zi) (w 12 (xf - 2/)² (w Tiyd
=xx 2i (w7 (1- x
(xx) (w3)² +
j=2
Since w/ E 1//4, we have 1 - (w,"i) 2 1 - (1/2)² 3/4.. By Condition (iii).
(1- (wFF
xi + Σ (x))
(w3)² 2 + 6
3 |x1|| - Σ (x)) (w:)² 1 |e'1
3=2
J=2
3 4
>
| - B + Ba) 2 T2.
25

---

Published as a conference paper at ICLR 2026
Moreover, Condition 2 (iii) also implies
sgn ( (1 - (w/ (i) xi
Σ 5=2
(x)) (wz)² +e') - +c-)= sgn(xx)).
+
By Condition 2 (i) and Condition 2 (ii).
(11-(W)))) x
d +
[xx - zi)
Σ j=2
(xj) (wr)² +o² (vx))++e+)
27112
sgu((x} zi)
(1-(1)7))) xi
+
Σ
(xx) (w,)² + +0))- +1.
j=2
Therefore,
Bu,s - 100) 2
1
T1 T2.
16
When d is large enough,
(u) A!00 Ba. + (5(n) <
11 12
(19)
32
Therefore, for all u E R,
d((u) (o(u) - Bu, (((1)) (1u)(1-((u) 0.
(B) Let s=t o()), g()) s(1 - s) E 1/4). Define h(u) =A!") Bo. + (b()). Noting
o'))) g C.. we have
(u(u) (() + h(u)g = (4 + (3) g())² + +h(u) (1 - 2s) g(u).
Hence
[w)( 5 14
+ Cál +1 +14,0) + cšl + |Bo + [K((W)) s Car,
4
where, using [x] - 2 S Bx
Bg and w 1/221
(Bx
B₂)²
C3
1
Bz
14,0 + c3l <
Bo,ll < 4
Bx +
+ r Bx + Be) (Bx + Be).
16
d
4
and thus one can take the explicit bound
Car: := 5 B2)
(
64 + 1 (Bx+
1
Bz
+r Bx + B. (Bx + B2) + C3 + C1
4 4 (Bx + 4
Bx
BB)
d
Consequently, for the Phase-ll scalar objective
[(v)
1 Σ 4()) - n n
-
x'==-'Z
T1
Let U := [u', u
na E R(M+d)xn
VC((T) 1
( u')
T' = 1
UD
n
n
D(V) = diag
)....,,'({ un))
D s C." mmax (U)2
v'c(
< 1
26

---

Published as a conference paper at ICLR 2026
(C) Let
:- ==
(Bo.s o - (i(w))
(1 d(u)),
a(u) Ite
1
Set
:= Bo - C Ci Dj = 4!0 + S3.
(0)
Assuming the dominance condition ensures C, > 0 and, for d large, D, 2 0 (since 1 > 0 and [Cš] = ((/dd).
Expanding e em - 0, fi(u) = CGe (D; -02s)
2u + (30; - 3D;;) e 3u Oie
=10;
Factor C;
fe(u) Cee " (1 gie + c.e. O 3u
small large
Choosing M0, and μr > 0 such that lael lafl-- + the single factor (1 + (en). Then, we have for all u0,b
bbe e- S e Smm absorbs all higher orders into
C,
C - e euueee S fi(u) S Ci(1 (1++ e
Denote C, exp (-u ú T u.) az Let v -loot ú, then p² -x + ú.
Lemma 10 (ouudyy et (2018 Lemma 10)). Consider GD updates v 4t1 =V - BVL(V)).
3 < Cgn amax (U)2, mm)
then the GD sequence satisfies
80
VV(("))2
< co,
lim 1-000
VL(V')1 = O.
t=0
Lemma 11. Define Mmin misi ,4 There exists constant C and t, such that for t > t, (r'+1 - r²) r² Ct min (0.1000..0..mmmmm...... 00.mmnn)) 0 = arg minggs u Tu' > 1.
Proof of Lemma 77] Since under our setting each sample's 's scalar loss (u) satisfies that -d((u) has an exponential tail, [Soudyy et (2018 Lemma 11) yields this lemma.
T 1+ - 1)) r'
(x - 14) -T 10---(8886 - ú log
3 Σ di (49) (r
Tre
log
1 +
t
t
n
idS
(20)
(A1)
A2
1 C. cxpl (ecppplüle u') + ((') u²) (ne)
iES
Denote by Xs the matrix whose columns are the support vectors, and let the orthogonal projection onto the subspace spanned by these support vectors. Then Pú ú.
For A1. firstly, the following shows that T lim; //VL(V/))))))))) = 0.
o(t).
Lemma
10 shows that
n T |Ti BB Σ VL(V*) - -û -úlogt
s=
(21)
$ T
(vT) i logg - ii) 31 min ñ VL(V*) = o(1)
T Sast
27

---

Published as a conference paper at ICLR 2026
Then A1 can be bounded from above by:
1 b
1 + ;) max [
1
log (1+
r'
log
S max [á Pr"Pr',
,d" _]]²²
(22)
|II| à ||| G,t. if |r 1
(tt)), if |r > 4
By Lemma 9 we have (u) < 0. Then,
1
V
n v n < 0.
(23)
This implies that L(V) does not have finite critical points v. With lim || = 0, we
have lit Tu'
u²= = +00. By Lemma there exists 4 > 0 such that for L > 41 and i E [m].
(V")"mn" S = S (i (v0) (24)
For Ag, when > 11, (0:((V') [[((()))TT\)] s 2exp(- (u))
B Σ (F') u') S 8
Σ
oi (v')) T. (T') r'
n
S
(nF: r''0
8
Σ
2 2exp(- (xp((-(u'] (r)))
:gS:(u')TT
2,9
Σ
a exp ( (-u un-nn''n (n'')T Tx
n
(4S:(u)
'20
23
Σ
(nn) exp ( (-ū
n
igS:(u')
< 2,3 exp
(-
min ū n 1
Isisn
(25)
For As, the proof is divided into two cases (u') and (u') rr < 0.
If ( 2 0, for > t1. by equation 24
Tx²
3 l C exp( -upp--uie u²) + o (vt) T ut) Tx'
(26)
BC 1-1 exp (-ū (11 + the exp (-paūT exp (-( -1] (T.)
n
(1 1f0 < Tr' 0.50,
C exp( - á uu)) + o' (V')T u') Tre
3C;
exp
-(1 + mi) min u un (- 1
(27)
Isisn
3C;
exp
- + (4)
min û Tn'
1155ppmmia
n
Isisn
28

---

Published as a conference paper at ICLR 2026
(I) If (u') r > t -055, we have
8 1 Cxxpp(- ú u') + oi (v') T ut) Tu') T'
80 1 exp (-u u²) [[1+t tM exp (- u exp (-t-u.sHe) - -1] (u') Tx²
8C: 1- exp (-uTu') [[1 t exp (-w:ū (1 t-.. + +1/
1 (r') Tr²
30, 1 exp (-ūTu') It exp (-p;ū" (1 - t 0.35p + mhi - t 0.s + thi (
0. (t t(2)
(28)
(II)) If (u')Tr²
e2, consider to > to such that tg > exp (mini ū w') (c0.50-11) Ft.
-
Then we have
1
C. exee(p(i n -n")) + o5 (v') T ut) (r
(29)
BC:
exp(-
max ú u') (1 - - 0.5.e)) c2t isisn
n
If (u')
Tr < 0, still consider three case. (I) If 0:5... S ( p² < 0, since - 0f > 0, we have
B 1 C exp(( ú uTu') + =o((()) T u') (u') r" <
SC
exp(
(u') Tr'|
n
n
3C₁
exp(
max 18/5n
T "u")t 1- 0.5p
n
3C;
exp(
max
n
15/5n
(30)
(f < 0.33p
for t >> t1
C exp(( ( ū u') + oil (v') ut) (u')Tr²
(31)
BC: - u¹ [1-e r' (1- (1-(1-1 (1-le-uTu ( """))""))) (u))
11
We then show that there exists t/ > tj such that for > ti the right hand of the above inequality is negative. Since (t'le ú u'u'e (u')
exp (-p (Vt) u' + 0. there exists t > to such
that for t > tts, (t'le uu un (un)) M4 < J. lfe (n') 3, then
Tt (1- at-le-aTun a (*)
2 1.5 > 1.
(32)
If e (n') Tr' < 3, then
(ur)
1 - - - uu e- -(1))T my
>e- (w') r'
(1 (3t-le- -u ut )tt)
(33)
2 (1 + -..5m.)
(1 (1- th (3e (3e-ūi ())))) mn
21 0.5pp 0.5pu
14 (3e in n Mi
1.5g
3e T P4
By taking t3 > 1/ such that for t > t. equation 32 and equation 33 larger than 1, we can obtain that
3
C expp( - ú nn) + o' (v!) T u²)
(r')T>' <0. 0.
(34)
n t
29

---

Published as a conference paper at ICLR 2026
(II)) If (u') T -.50, there exists t > tj such that for t t/3.
e-(n')'r' (1- u'n ) 2 1.5.
(35)
Then there exists constant ej such that
3 1 C, exp( u') + (+t) ut) (u') (u') < -att
(36)
1
Flnally, consider > l2 maxisign (s: If ||r|| > then
max iES
(u') (u') T
(u') Pr"
1 S Pr'l
1 2 8 iin (Xs)
(37)
iES
Let 62 = /IS| min (Xs) then there exists E S such that [(u')Tr'] 2 62. By equation
29
equation 36 and equation 22
we can obtain that there exists > 0, such that for t >> t2.
(r*t1-r)) - -r²) r') r' S -ct- + o(t
(38)
Then there exists ts > tg such that for & t3.
(r++1 r¹) r' 0 Ct (in(1,,10.55mmmi
(39)
If ||rr|| > <1, then by equation 27 equation such that for & > t2, we have
28
equation
30 and equation 31 there exists C3
(r (1) T re cg t- min (0.100110..5mmm..
(40)
Taking C 63 and t1 t,, we obtain that for t > tı, (1tt1 -xr) T xt sCt mp.55
A.3.3 STAGE
0
After Stage 2, we have with probablity at least - - W
satisfies 1 5 w < 2
w Th
for j S (1),
Yoologd/n d/n
ww T5
<
o lollog//n and for j E S
T
o log d/n d²
Moreover,
d3
de
the gate iterate satisfies p4
+i = oo(v (T T
v o log d/n
for all We now show
that, in Stage with T5.
log 7 Tog(ar))
and T = Ti + Ty + T3 GD yields wT w
Since Stage 3 freezes the gate, V remains ú log T3 + p² where ú the max max-margin solution on {x. 2, 1 T bounded residual, and T3 d2 02 log d/n. This completes the proof of Theorem The Stage 3 bounds follow by combining Lemma 12 and Lemma 13
w3 11l00 at its Stage 2 form: 4 T"
0 oolood logd
lllog//n
As a preparatory step, we provide a recursion expression of w in Stage 3, which will be used repeatedly in the proof. The update of w in Stage 3 takes the form:
++ = wt 1
17
(x') T (wt)02-(11--- - (1 (1-p,,) T2 (x* - TT)'
1=1
(41)
-(x') (ee))(w=)Bm) (w2 (00))02-s) -ss) (x' (1 p1 (xx - 2
30

---

Published as a conference paper at ICLR 2026
Collecting the small factors involving (1 (1-p2 p.²) yields:
n x (xx) ((W((( -(w))))
"
w = w' ② la 1
n Σ x'6
71
i=1
To (1 Ps
x''x' -z') = (w+))2
(1 Ps (x' xx)(x') T ((w0)02= (w²)
(42)
(1 Pj
T5 ) (x' -z')c; c
i=1
]
(1
To 2 (xt - x()(x² - 2t) T ((w (w=')02 P. )
We then decompose w' over the support of w Let supp (w*) = S. with |S| = 7. Define
w's = w'
② 1s, W's- w C 1ge. Then we define,
(x")
H
((w3) 02
(w2))2)
71
(x')
(w'..
1
x'c
1 1
Σ
(1
Pi T5
x² (x'x =z')((')00 (ww))2
n 1=1
n1
i=1
1
(1
Pi
(x' - zz) xx)) T ((w')
- (w2)
+ Σ (1 p, (x' -z') -z')es Ei
1=1
1
Pi )² (x- z z')(x² - -z') (w=))22
1=1
(43)
Then the dynamic of w can be expressed as
w r+1 w' 3 Ia r/H ((w3) 02 - 02
02) - r¹ + +ee
(44)
With H
H = didi((aa,,, 1, 1....,1,,, 1), let c == a for j = and c = 1 for j > 1. Since (w;)² = w] E (0.
w = w/t1 (1
11 (l-n(((ww))) (w 2 -c-r,
r ( e.)
88 w/th = w (1 17 (T+ + +e)))
Lemma 12. Define Bi C1 rr log(dr) and Bg C2 / 0² log(dr) for absolute constants C > 0.
Assume the step size 0 1 < (I+ Set the phase lengths
max (a(w]" T5 )2, 17 )
B
T
log
+
T =
log
T : T + Tg.
4m
5a B1
4m
4n
B2
Then, with probability at least 1 ae, the following statements hold. For all t < T + Tū
02 logd
Vj E S: 0 s1
E w $ d
n
Proof of Lemma 12 Define the inductive property P(t) for t > T
P(I): : Vs E [T,4) 4, Vj E 0 w/ 1 B1, As T 4. Vj S w/ S 5
log Vo d/n.
31

---

Published as a conference paper at ICLR 2026
Base case holds by initialization. Assuming P(t), with concerntration Lemma 14 and 20 we bound
15 16 17 18
19
|x s 2 S Bi, || S BBB,, hence ||r -e' 1oc B1. By Lemma with B = B₁ and the stepsize, we have S w S + B for j S. For j E 5
w w s (1 (1)) w/
Ti wj Tit S B. w
Since T3 = T2 + Tī satisfies
d d log d/n
Tš S nBi
log
Ti "max,S" S
then w Ti < d= Vo" log d/n for all 7 S Tg. Thus P(t + 1) holds. By induction, the claim
holds.
Lemma 13. For all 2 T3 + TT and all j E E.
[w] - 1| S B1.
If n 2 log then for all t 2 T3 + Tī + T = Tī + T3 and all j E S,
[w] 1| s B2.
Proof of Lemmal 13
Fix E S and ot := c[w])) C fa, = 1/4, To = =12 = 1/8. When
w/ < Yes
wft ww (1 n((l1 (w/)))
with lbel S Bj. Since 1 1- (w()))²² 2 15/6 64, assuming B1 S 11/64 yields c(1 (w 2 11/64.
hence w/t1 2 (1 + 1 11/64) w' From w, E [1/4, 1/2]. in at most
64
1 1 1-""" 33
T(B:) <
log
11
T
7
steps we have w/ 2 1 - Yw, i.e., [t cl S 7. Then by Lemma
22 (with B B1),
|t+1 cl S (1 Kn) | cl B1nBI,
K 3. 8 = 2.1,
and whenever [xg - c| 2 AcB with 2 23/k (A = 7),
| cl s (1 n)) 20x cl.
Hence after
T < log I1z c
5
+
41
Bc
4n
iterations, |xy, - cl S 5cB,, which implies /w/ 1 S 5B.
This implies ((w3)) 102 (w*) 102) I C. Therefore,
2og22og² (dr)
| & L, r C. Bi S CT
B2,
100 S sCoV
o² logd < B3.
n
By Lemma 21 with B B2, for
Tī T + +T5.
5
B1
T5 4n
log
+ |w/ - 1| S 5B2
B2
Lemma 14. Given vector u supported on for any 8 E (0, 1), with probability at least
[1: ( (xx)
H)ul -H)
< B2
2
log 2dr) .r -rllulocc
V n
32

---

Published as a conference paper at ICLR 2026
Proof of Lemma
14
Fix j E [d]. Then
[ n. Σ (x*)
H)ul, = Σ
Σ x [xxxxx1)
ug.
KES
For each k; x
[- B2 B Hoeffding implies
2nt²
Pr Σ x E[x,xa] 2t) eexp
(2
union bound over (j,k) E [d] gives the uniform deviation. Hence,
[[(: [x(x) H H)mp B B log( 2dr) 6
Σ
Jus S B² / log( 2dr) ..rrll -r|[ullsos
kES
and taking the maximum over yields the claim.
Theorem 15, Given E IRd, for
for any å E (0, 1), with probability at least 1 5,
x (x I
$
maxfa, 13 + B d 2 (2d (6)
|ulloo.
Proof of Lemma 15
Decompose
x (x T u== Hu 1 (x+)) H H)u.
i=1
Since H diag(a, 1, 1,...,1) 1), Hulloo S max 13 ||u|bo..
For the deviation, define
Yı == x (x u E[, Ju],
( { x xx -H)u
= - Yia
n 1=1
2=1
The Y2 are i...d, mean zero, and almost surely
Yesl S
x uul
S Bx (d ||u|ox
B& d ||u||oo-
By Hoeffding's 's inequality,
nt f dB 1 120
Pr Σ Yis 21)
< 2exp
A union bound over 1, 1,.... d gives, with probability at least 8. --8. 2 $ogg(222/ /8)
max 153
Yis
B d ||//oo
n
Combining with the population bound yields the result.
Lemma 16. For any 8 E (0, 1), with probability at least
$ Bx B.
log
Proof of Lemma
16
Apply Hoeffding coordinatewise and union bound over coordinates.
33

---

Published as a conference paper at ICLR 2026
Lemma 17. Define By rBx B. and Bx Bx + lIf(0))loo + LB + Bg. Then, given R for any 8 (0,1), 1), with probability at least - 8,
2¹)
us
x (x' z²) uge
BBx 11+
2lor(4d/6)
(rlluslloo + (d r) l ll00
Consequently, with Pmin = min; p.: T2
Di x x'(x'= z (w (w=)
S(1 - Pmin) Bx (1
2log(dd/8)
(ril(w's)0)) 22 + (d (- 1021100)
Proof of Lemmaa 77] By the model assumptions, |xx| I Bx almost surely. With y x (w" y02 e and |xx 1 _xx. we have yy| S rB + B. =: By Since z = f(y) + with f being L-Lisschit
and 10 I Bg. follows that
|2| - (1) Dc + L/yl + oo & (1f(0))loo + LB, + Be.
Hence ||x - 20 - Bx + [0))l - LB, + Be =: Bxz.
Fix any j E [d]. For the j-th coordinate of E xx* - 2²) uss we have
1
n
-
((x - zt) us)
xz)| |X 1 | S Bx ||s||ı
n
i=1
Taking expectation shows the same bound for the population mean. For the centered fluctuations, define
y(i.s) x) ((x' z us) E 2x'' uz))
Then |
< 2 Bx Bxx ||us|| almost surely, and by Hoeffding's inequality,
Pr (1S) 2 BxBxalll ||ssllu V 2log(2d/6)
8
S
n1
2d
A union bound over j [d] yields, with probability least 1 5/2,
x" 2')J us S Bx
1
2log(2d/5)
||us|l..
Using ||us||| < r| us oo completes the bound for the us term. An identical argument with use (and |uge 11 5 (d - r) r>||uusllaa (loo) gives, with probability at least 1 0/2.
< BxBxa (1+ / /210g(20/8))
1/4 Σ x (xf z²)
use
||use I|1-
union bound over the two events replaces (2g(2d/8) 8) by log(4d/6) and yields the first display.
For the consequence, note 0 < 1 Ti < (1 Pumi)). Factoring this out and applying the first bound with u (w (so uş = (w's) and Ug- = (w's)) 002) proves the second display.
Lemma 18. Under the same assumptions Lemma 17 for any E (0, 1), with probability at least 8,
Σ xx' - z xx)(x²) T us
+ 1
Σ
(x" z²) uss=
1=1
1=1
2 loog 4d
(rl) (rluusloo + (d ---r)) r)]llss=loo))
<BxBxe
1+
34

---

Published as a conference paper at ICLR 2026
Consequently, with Pmin min, Pi
25
Σ (1 ) (x' 2x)(x') -(w))22)
log 2 (rl((w)) (w3)) 1loo + (d - (B--))((wss.. 02 l10))
<(1 Pmin) Bxx 1+
a
Proof of Lemmal 18
Same proof Lemma 17
exchanging the two factors.
Lemma 19. Let Sn
E Ei=((x' (x' z²) ci
R.d. Then for any E (0, 1), with probability at least
1 1, 6,
n
Σ p.²² )x² 2')
2d 200log( logl
20
2
log(
2d 8
$1 Pmmin)
Bxx Bxx
Bxa Be
3
n
Proof of Lemma 19 Fix E [d] and define x(6) = (x) - 23)(s, 6.. Since || xx'-2
-2,)
x 2211 11x0 S Bxx a.s., we
[ [[(x; - 2,)c]. By Cauchy- have (x(² S Bxx B. and Var(X.!)) S B 0². Let 11
Schwarz and the a.s. bound, 1a51 S Bxx Let y,U) = x μjj then |4,0) S xx 1nyl S
Bx (Be + σ) and Var(Y,(1)) = Var(X(!)) I 0² . By Bernstein's 's inequality, for any > 0,
nt²
Pr
2 t
< 2 2xxp
2B 02 Bxx (B. + o
Choose = Bxxe 22og(2d/8)
+ Bxe (B. + to) (2g(2)/)) and take union bound over j = 1....,d ,d
to obtain, with probability at least - 8,
X
(1)
$ 111ll Bxx
2lo((2//8)
+ Bxe + log(2d/8)
2
n2
n
n
DD
Since ||o1 I Bxz and Be to S BBa we have
n
x
< oxx + Bxa
2log(2d/5)
+ 3
2
B.
log(2d/6)
n
im1
Finally, 0 1-1-2 p;
(1 Pamin) implics
(1 po (xx z¹) z()
1=1
2 (2og(2d/8)
3 Bxx B, log 12g(2)/))
2
<(1 Pmiin
Bxx + Bxao
which completes the proof.
Lemma 20. Let Pmin min; p.a For any 8 E (0, 1), with probability at least 1 6,
L/ 1=1
Σ (1 T3 (x - -z'z 2x² - (
< (1 - pmin)² #B2 1+
2log(4d/6)
(wl((w)) 02 + (d (w:))) )09
35

---

Published as a conference paper at ICLR 2026
Proof of Lemma 20
Apply the same argument as in Lemma 17
Consider the scalar update
w+ Go)(w) == w - n(ww? - cc-b)), c e21, 16| S B, 17 1>0, 0
and let x: == cw?, Tb(x) == sx1 n/ - +c+ b)))2 (so x+ -T(x))).
Lemma 21 (Invariance and monotonicity). If0 1) S 8c] then for every |b| S B the map Ge.b is nondecreasing on [0, 1 + B] and Ge,b (10, 1 - B) [ 00.1 B).
Proof of Lemma 21 We have Gr,, = 11 cc-c
b). On [0. 1 1+ B}. max - c b) S
3c(1 + B) hence G² 0 if n S 1/[3c(1 + B) which is implied by 1/[8c(1 + B) For invariance, note that (0))(0)
and by monotonicity in both wo and b,
Gablw)= - Ge.B(L (1 + B) = (1 (1+ (1- - 1) B) S 11+B,
since c((1 + 1)) - B 2 B 0 and (c((( BB) s 1. Nonnegativity of the bracket
is ensured because
(1 n--n - c- ci bi) -1 B)
min we|0,1+B),.bb |b
c 20
under $ 1//[8c[1 + 9)
Lemma 22 (Local contraction with bounded noise). Fix y E (0,1/4), assume |b| s B, B S 1/8 and 1) S 1/8. If cl S % then
x cl I (1 kn)) |- cl + B B, Moreover, is [x cl 2 AcB with A 2 23/ K (e.g. A = 7), then
K k= 3/5 B = 2.1.
1x+ s (1 |n) |- dl s (1- }n) |r cl.
Consequently,
2
To
1
log
wwo - 11 SAB.
Kq
AcB
Proof of Lemma 22 Let y : x (c+ b). Then (x) = (1 (1- my) - my - 2nz).. On |æ - c) S Y
-
|b| S B, we have |y| S 7 + B 3/8 and 2x |a| 2 2(c - 7) (7+ B) 2 9/8. Hence
[Tg(,)] $ (1 + nyy) - (22 (b)) < - ((2r - (b]) - 1n))++>>(3 - 1y))
Using (2r- |y] v| 2c 2B 3/4, |2r - |y/) < (3/8) 27/64, and < 1/8,
we get (x) < 1- n. By the maluevalue theorem, [Tb(z) Tb(c) S - - c with
K = 2/3. Also [T(cc) cl = 1e(1+mp) nb)² c| $ 2en + cT B² < 2.1 B. 1/8. Combining
gives the first claim. If [ cl AcB, then
x c - (1 xn + ( x - cl s (1- mn) r cl
whenever A 2 23/n. This yields the second claim and the exponential -eoonential-time bound.
A.4 PROOF OF THBOREM 2
Recall ú be the max-margin solution on J and Tomp := 1/1l
Lemma 23, Fix e E (0, and 8 E (0, 1). With probability at least - 8 over the training sample,
2/ s 2 + d( Bx + B2Y
2 2lm(2/8)
Pr((ū) T u > E) 2 1-
|s]|211 - V)
36

---

Published as a conference paper at ICLR 2026
Proof. Let U {u', ,..., ,u) be the training sample and U' {u" ...... um*) be another i..d. samples.
Let // iill Vemp and u/// Define the ramp function at threshold and width s >
1, t
((E,)) =
1 to 8 1 < <
0,
12 8.
Then d [0, 1] and is L/ishiishitt. Set 7emp - E Yemp (1 =) 0. Since min, 6 o'u'
Temp 2 E + s, the empirical ramp loss is zero:
Σ (TE u²) =
(45)
and 1 S E) I (), ( Now apply the high-pr probpbbobbbitty Rademacher uniform deviation
bound (Shalev-Shwartz & Ben-David 2014), Theorem 26.5) to the loss L viE.s) (bounded by 1),
over the linear class F {u + 1 I 1}: with probability least 1 - 8,
B[e(o w)] $
E u') 2E
Bu [R(EoFoU')] +
2 ln(2/5)
Using the contraction lemma SSallv-Shward & Ben-David 2014), Lemma 26,9), Hi FoU') < R(F U'). By the linear class Rademacher bound (Shaley-Shwartz & Ben-David 2014))
Lemma 26.10), R(F U') R/vn since ||u|| s R a.s. Because the empirical ramp loss zero, we get
2R svn +
2lm(2/8)
E ((VEE,)) u)] I
Finally, Pr(oT T u () < E (e.s) T u)] yields
2R
2ln(2/0)
Pr ((á) u ( -
Yemp (1 E) +V
which implies the complementary lower bound for Pr((ū) u > E). Combining with Lemma gives the explicit version.
Define
s x 2
Rad := 1|s|11 (Bx + Bz)².
8
S
$2
Since with probability at least - O(1/d2) the output "
= Ti + T5 + +T satisfies
d
ilog T2 + t
p 11= = 0(1),
T
02 log d/n
Moreover,
w 1l,o S log d/n
W
Iww3 -1s|loo S S
log
Hence
a² log d
[w"""002 (w*) 02 r
I[("" 20 1 +r()+ "02
n
Let a ((v ((VT)TTu) and, for E (0,1/2))
0 (p= > 1 7 }
:= log 1-r
37

---

Published as a conference paper at ICLR 2026
Choose E, (K, Cp Rad)/ R with (( ul Cp Rad a.s.; then
{u" u Ex} C Q, By
Lemma 23 with probability at least 1 - over the training sample,
2 Rad
2 2m/02)
Pr((),) 2 1 |s||a11 8r)
(46)
In particular, taking T = + 0² log (/ (n and log T2 2 4Cp Rad + 2 log >> nd Ing a) gives - E, 2
1 1/4 x11 and thus Pr(S2)) close to one; moreover p 2 - r 2 1 7=-
Lemma 24. Let F := Amax B x)(2 - x)" [2,) - Crd. Cr = (Bx + BB) Then
E [(w"" : x²²) 0²
S2
2
S Smmax (H) [l(w" T" 02 (w") ) 0 1/ + 72 rr ||( T" 102 11 + B. T vT, [[(w" )02 12-
Proof. Let p1z (x) EEly x = (w*)02 and (p> + (1 p)) (WT )02 Then we have
9 - Pz = x T ((w") y02 - (w") 02 + (1 - p) g,
g = (z - x) (w T" )02
(47)
Since c 1 x, E|c] = O. we have
L - 0²/2 12/22/ = 1/2 (ū }1 - [( )=)c].
(48)
Conditioning on S2, (so 0 $ - p 5 T), we get
E [ 3 [2,] 2 (w T" )02 )02) H(( (w T" )02 - (w2) + 2 7 Elg² [tr]
2 Amax (H) ([w))((w))) T") y02 (w*) 002) 122 + 22 r, l1(w Th 10
For the cross term, [iÚ pz)e | [2]] = [B[(1 - p)ge [4 S B T E[gl [2,) S BeT T v Tr |l(ww T 102 1|2
Combining yields the claim.
Proof of Thhoremm 2 Let Ω = Ω2 The probability lower bound follows from equation 46 with
E, and the stated scale of Tž. Since p" 2 r 2 3. Apply Lemma
24 and substitute
[I((w"))o -w | o² log d [[ T" y0 1 S r. 1' S Ord, and 7 S o² logd obtain the nd
stated rate.
n
Proof of Corollary 7 Define u' =
x =z' and Rad' :=
| [ - B²))² Replicating
the high-probability calibration of the gate (as in the proof of Theorem 2 with u replaced by u' and
Rad by Rad' yields the event & with
8 Rad'
2 22(2))
Pr([)] 2
|s|12v/
n
on which nf = =0((v [ un) 2 1 - 1/d2.
For the risk, write 9 = (pT x + (1 - p" )z') T (w T" 02 in Lemma 24
002 and Ho xT (w*)02, Decompose (as
Mr = ((w T" )02 - (w*))2)
82) + (1 - x- )g', g'= (z' - x
(w (w
Conditioned on Ω we have 0 < 1 - -ptT' S $1/02. Hence,
[ - p.) 2 [2] < 2 X Amax (H) [[ 002 (w))0211 + 2- 12 ||| T' 00212
where 7 = 1/d and 17 ;= Auax (El(z' () - x) T [)) I Ctd with Cf = (Bx + B²) by
the envelope bounds. The cross term satisfies [E(( - 1x [ < B vT: (w 7" 10202122
Finally, substitute the parameter accuracies from Theorem and [| T 002 2/ S r(1 + o(1)), to conclude
(I T" 022 - (w 1021 13 S r 2og n
r logd
T.
Lix,nnz) (w T" )TT))
Ω
S
2
n
38

---

Published as a conference paper at ICLR 2026
A.5 EXAMPLES
Example 25 (Under Dominant-Coordinat Condition, population linear regression retains a constant
fraction of z). Lel
X1 +c,
21 = cy + 2, ( 2).
with mutually independent coordinates and
Var(x1) a 00 Var(x,) (j 2 2), Var(c) o², Var(E) o
and 1 E EL x. Consider the population linear regression problem
12) E arg min Bl - Bx x - Bg 2) 8x.32
o O, the unique solution is
co²
82,
Pai -0
Bk =0 2). 8x1 oe + c'o2'
o + co22
Equivalently,
202
co² of c
i x1
of c2
so linear regression uses 2j to fit the residual noise e with coefficient /²02/(02 + &²0²), thereby retaining constant fraction of the component.
Proof of Example 25 From y = X1 + e and independence,
Var(y) + o², Cov(x1, = Cov(x,, = = 2).
Moreover,
21 = cy + cxx
zj 0= 2),
(x1 21) = ca,
Cov(x, 21) = -0
Var(z)) = e² (a + 0²) + +
The population projection is
11, = Ex Eiy
Since y X1 + with e 1 x, we have
IT yx O).
y'==y xxx
Likewise, only exj z1 = CX1 ce & projects on x, hence
11x1
Cov(x1, 21)
Var(x;)
ca
e,
z =
112ix X1 CE
and 2/ = f 2. Mrisch-Waugh-Lowwe Theorem (Frisch & Waugh 1933) yields the z-coefficients of the joint regres- sion by regressing y' on z'. Thus,
Cov(zi,8/) Var(z()
Cov(ce+ E) Var(ce E) 2022 + oz
co²
P5.
and z = 0 implies 8 for j 2 2.
The x-oyly regression gives
g B (only-x) 1
(onlly = 0j 2 2).
The joint coefficient equals the above minus what is explained via Z
8 = g (nnly-x) Cov(x1,21) B², = 1 - -c = o + 2 2
o²
Var(x1)
and by symmetry $2.5 = 0 for j 2 2.
39

---

Published as a conference paper at ICLR 2026
Example 26. Let Unif - -Xi (symmetric, X 0). X2 Unif[-Bi B, B). Unif[-FF F, F. E Unif[ 1 & Unif[ 2,, all independent, and xj = 0 for 323. Te1
y
X2 +,,
21 (1 k) L
02 y + E
with < < 1, 02 0. Set 3x == max [X, B), ==
= B. and, B, == E. Then:
(1) Descendant stronger than non-dominant causal coordinate.
Since Var(x2) B²/3 and
Var(y) =: Vu Var(x1) + Var(x2) + Var(c),
Cov(x2, v) = Var(x2) = B²/3, Cov(22, = 02 Vy-
Hence for any 02 > (B2 /3/VV, we have Cov(z2, Z2, y) > Cov y) even though X2 is causal.
(2) Population dominance of coordinate (Condition D. Since
$j = F[y(x; + 2;)], My = c[((x + 2;)].
$9 = 8, + Mg
Then
cff Var(x)) + (1 - Var(y) + (1 Var(e),
$2 = Var(xa) + 02 Var(y) + az Var(e),
and 5 = 0 2 3. Because m1, Maj are bounded under our assumptions and satt is fixed by (B, a2, K), choosing Var(xi) large enough (e.g., increasing X) guarantees
eff
2m11
4 + 1011
max
15
j>1
8
i.e, Condition
(3) Per-sample margins (Condition
2). Since x1 - 21 = KX1 (1 - () + +() E1.
X1 z 2 x/ x|x|| - (1 - ()x +|e] - |5,1 2 RX - (1-) -E)-E -Ei
Thus i X (1- - )) E) + E + TI. then xx 21 2 Ti a.s. Moreover xf (1 x))
E) + E,, then the sign of X1 21 equals the sign of x1 a.s. Finally, (r- + B +
then $ |x1 2 (r Bx B. + a.s. Hence Condition 2 holds.
40


Reviews:
[{"review_id": "Jdcg5ksooK", "reviewer": "Reviewer_1yaN", "summary": "The paper is trying to answer a very specific, surprisingly under-asked question:\nWhen can a vanilla Transformer, trained by plain ERM + gradient descent on one environment, end up using only the causal feature and ignore the spurious ones — without IRM, DRO, group labels, or any of the usual causal tricks?\n\nTheir answer is: it happens when the causal signal is uniformly stronger than every competing spurious signal — and gradient descent in a Transformer has an implicit bias that pushes attention to that dominant causal direction. They call this phenomenon Correlation Crowding-Out (CCO).", "questions": "Please see weaknesses", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Positioning CCO as the “mirror regime” of shortcut learning (spurious wins early vs causal wins early) is neat.\n\n- The occupation → crowding-out story is crisp and gives an intuitive training-dynamics explanation. Theorem 1 cleanly proves the two-phase story—occupation then crowding-out.  The training-time schedule and resulting bounds are explicit.\n\n- The Frisch–Waugh–Lovell construction shows population least squares retains a nonzero spurious coefficient even with a dominant causal feature—so the cause-only outcome is due to optimization + attention geometry, not merely stronger correlations.\n\n- The two-token, two-branch attention setup with a squared-parameter head is cleverly chosen: simple enough for precise proofs,\nstill recognizably “Transformer-like”, and directly exposes the competition between causal vs descendant branches.", "weaknesses": "- Your main theoretical result (Theorem 1) relies on a three-stage schedule where you alternately freeze the head and the gate. Can you clarify whether your informal claims about “standard training alone” are meant to apply to simultaneous GD on all parameters?\nDo you have either (a) a theoretical argument, or (b) empirical evidence, that joint updates (no freezing) still exhibit the same occupation → crowding-out dynamics and cause-only behavior?\nIf not, could you re-scope the claim or discuss when the staged schedule is a realistic approximation to how Transformers are trained in practice?\n\n- All formal results appear to be proved for the squared-parameter head  $\\hat{y} = h^\\top w^{\\odot 2}$, while the difference-of-squares generalization $\\hat{y} = h^\\top \\bigl(w^{\\odot 2} - v^{\\odot 2}\\bigr)$ is only mentioned briefly.\nCan you clarify whether Theorem 1 and Theorem 2 actually extend to the  difference-of-squares parameterization (and hence to a general signed linear head)? If they do extend, could you outline the additional arguments needed to handle the dynamics of both $w$ and $v$ and possible cancellations between $w_j^2$ and $v_j^2$? If they do not, it would be helpful to state explicitly that the current guarantees\nare restricted to the squared-only head.\n\n- Your Dominant-Coordinate Condition combines (i) a uniform population dominance gap and (ii) per-sample margin and sign-stability assumptions on the dominant causal coordinate. How realistic do you believe these assumptions are in typical high-dimensional settings with multiple causal parents, mixed-sign effects, or occasional sign flips?\nAre there weaker or more invariant conditions (e.g., dominance in norm, or dominance on average but not per-sample) under which you conjecture CCO would still hold?\nCould you comment on how sensitive your mechanism is to mild violations of sign-stability or margin—for example, do small fractions of “bad” samples already break the crowding-out behavior?\n\n- In Theorem 2, the lower bound on $p_{T^*}$ is written as $p_{T^\\*} \\ge 1/d^2$, which seems inconsistent with the training result $p_{T^\\* i} \\geq 1 - 1/d^2$ and  with the textual claim that the gate ``prefers'' the causal branch. Should the bound in Theorem 2 instead be $p_{T^\\*} \\geq 1 - 1/d^2$?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PFYORj0B83", "reviewer": "Reviewer_hNry", "summary": "The paper formalizes how transformers learn differently strong signals in data. They define causal signals as the target signals that dominate in data, and verify in a stripped down transformer that the dominating signal is learned over the training duration, edging away the spurious signal. In three experiments, they verify that the transformer architecture indeed ends up learning the stronger, causally relevant signal from the data.", "questions": "- Can the authors provide fair experiments with waterbirds (stronger bias, i.e. 90-95%) and use ResNet and other modern architectures as baselines?\n\n- Can the authors provide even more experiments with realistic data, in which the bias is explicitly known (the NLP dataset is not suitable, as far as I understand, as spuriousness is not labeled or controllable)?\n\n- Do the authors think that transformers are substantially different to CNNs when it comes to picking up the stronger signal in data? Any evidence for this, be it theoretical or empirical if yes?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This reviewer is not deeply familiar with the transformer-theory literature, but the theoretical framing is potentially original and useful. (Perhaps other reviewers/AC with deeper familiarity of the topic could confirm whether similar derivations already exist)\n\n- The paper poses an original and conceptually stimulating hypothesis: that under strong causal alignment, gradient descent in Transformers may implicitly suppress spurious correlations.\n\n- The mathematical formulation is clean and logically structured. The \"Dominant-Coordinate Condition\" and the theoretical two-phase training dynamics are novel in how they connect implicit bias to causal learning.\n\n- The analysis isolates the attention mechanism and offers a mechanistic interpretation (\"occupation\" and \"crowding-out\") that could inspire further study.\n\n- The clarity of exposition is commendable. Even though the setup is theoretical, the derivations are readable and well-motivated.\n\n- The authors attempted to empirically verify their theory, including both simulated and real tasks.\n\n- Conceptually, this paper asks a meaningful question that sits at the intersection of optimization dynamics, causal inference, and representation learning - a very relevant direction.", "weaknesses": "1.) Architectural realism.\nThe experiments and the two-key, gated-query setup differ substantially from classical self-attention. It remains unclear whether the observed CCO mechanism holds in standard multi-token Transformers where Q, K, V are interdependent.\n\n2.) Experimental limitations and questionable causal setup.\nThe empirical evidence is far too weak to substantiate the claims.\n\nThe toy experiments are overly simple and serve mainly as visual confirmation of the theory rather than genuine empirical validation. The NLP experiment has no quantifiable way of measuring spuriousness that is in the data.\n\nIn the Waterbirds experiment, the authors treat the background as the causal feature and the bird type as spurious. This is opposite to the conventional bias/shortcut setup. The background is typically the spurious cue, as it is the easier signal. This inversion undermines the causal argument and makes the results hard to interpret. To be credible, the experiment should make the bird type the causal target and the background the bias.\n\nEven more problematic is the bias strength of 70%. This is far too weak to demonstrate spurious correlation effects. At 70%, even simple architectures easily learn the causal signal.\n\n3.) Replication evidence (my own test).\nTo prove my point that even simple architectures learn the causal signal, I reproduced a small experiment following the waterbirds setup in Sagawa et al. [1]. Using the CUB and Places datasets, I balanced the waterbird/landbird classes (so that accuracy equals balanced accuracy) and applied a 70% background bias. I used the bird type as the label and the background as the spurious attribute, which aligns with standard fairness and bias-mitigation literature. I trained a ResNet-18 under standard ERM and also with the GDRO method [1]. The training set contains the 70% bias, while the test set is perfectly balanced/unbiased. Training on biased data and testing on unbiased data is a reliable method to verify whether methods have learned the causally relevant features [1,2,3,...]. Results, rounded:\n\nOracle (ResNet-18 trained on unbiased data): 96% accuracy\n\nResNet-18 ERM (trained on biased = 70%): 95% accuracy\n\nGDRO (ResNet-18 backbone, trained on biased = 70%): 95% accuracy\n\nThese results show that with such weak correlation strength, even a plain CNN (ResNet) is not misled by the bias and the model simply learns the dominant causal signal. Thus, the authors' experiment cannot reveal anything unique about transformer dynamics. To show a genuine robustness effect, they should raise the bias correlation to ≥ 90% or 95%, as commonly done in prior work [1, 2, 3]. Furthermore, the authors must report quantitative numbers and not only qualitative visual cues.\n\n4.) Lack of baselines.\nThe paper does not compare against CNNs, MLPs, or other baselines. Without this, one cannot tell whether CCO is a transformer-specific property or simply an instance of correlation dominance common to most architectures.\n\n5.) Triviality of the claim.\nThe theoretical result essentially formalizes the intuitive statement that if one signal is much stronger, gradient descent will focus on it. This is not necessarily \"causal learning\". It's correlation bias under strong dominance. Similar statements can likely be derived for CNNs and other architectures. The true open question is, whether transformers produce more causally stable predictions than other networks, or whether their implicit bias provides tighter robustness bounds under stronger spurious dependencies.\n\nIn short, while the theory is elegant, the experiments do not convincingly demonstrate anything beyond the trivial: when one feature is much more predictive, the model will use it. As shown above, even a ResNet-18 attained these results out of the box if the causal signal significantly dominates (as given by the authors experiments). Whether transformers are better for robust learning than other relevant baselines is a very interesting question. As is, this paper is not publication ready, as the claims cannot be proven beyond the trivial. I'm open to discussing my points during the rebuttal.\n\nReferences (as cited in review):\n[1] Sagawa et al. \"Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization.\" arXiv 1911.08731 (2019).\n[2] Makar et al. \"Causally Motivated Shortcut Removal Using Auxiliary Labels.\" AISTATS 2022.\n[3] Liu et al. \"Just Train Twice: Improving Group Robustness Without Training Group Information.\" ICML 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HqXHPowLsl", "reviewer": "Reviewer_DtiM", "summary": "Authors propose to inspect a phenomenon called Correlation Crowding-Out (CCO), which arises during transformers training where, under a precise set of assumptions, transformers learn to exclude spurious features in favor of the true causal factors as training progresses. The authors argue that CCO arises because weights aligned with causal features grow and stabilize during training.\n\nMore generally, a two-phase process is assumed where the models quickly adapts to the true causal features (while still remaining attention to spurious links). The authors provide a detailed theoretical analysis, providing bounds for CCO effects for gradient descent training on transformers. Experiments on synthetic tabular-, visual- and natural language data are provided to support the theory.", "questions": "My questions primarily regard the mentioned weaknesses:\n\n1) Could the authors further elaborate on the insufficiency of anti-causal predictions in the chosen scenario? Under which scenarios would and z->y prediction become invalid, and how does this reflect in the made assumptions and data generation?\n2) Could the authors comment on the possibility of the observed occupation effects being due to varying feature quality or link complexity?\n3) Which assumptions with regard to the model and data generating process do the authors make in their paper? Are these assumptions expected to be observed in real-world scenarios and how do the results generalize beyond the made assumptions?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The research question is interesting as only a limited number of works analyze the conditions under which transformers learn to distinguish between causal and spurious features. The authors present an extensive theoretical analysis on the proposed problem and, to the best of my knowledge, are the first to provide such a formal analysis on the theory of causal transformer learning.\n\nThe paper is generally well motivated and structured. The overall claims and theory seem reasonable and is well laid out. The presented theorems and proofs seem to be generally plausible and consistent, although I did not check them in detail.\n\nTheoretical results are supported by experiments on simulated tabular-, vision- and natural language data and cover a wide range of applications, demonstrating the robustness of the proposed theory under a distinct range domains.", "weaknesses": "1) **Causal Inference.** In unintervened settings (in the Pearlian Causal or Bayesian Network sense) it should be perfectly fine to infer y from z, assuming an equal quality of the links x->y and y->z. Given that the authors assume x->y to be a linear relation and y->z to be a more complex Lipschitz function, x->y should generally be expected to be easier to learn. Particularly the initially described \"occupation\" phase might simply be a byproduct of the different link qualities. I would like to suggest to at least mention this aspect on the possible validity of causal and anti-causal inference in the paper.\n2) **Experimental Alignment.** To test their theory on image data, the authors create a synthetic waterbirds-like dataset. In their setup, the learned foreground and background features, however, differ in quality - \"object-like\" versus \"scenery\". This difference in feature quality could in theory be a cause of the transformer focusing in the foreground feature, artificially introducing the \"occupation\" effect. A more sound setup would be to analyze the evolution of attention on images containing two features of the same type and/or complexity, e.g. two bird of different species in the image. While the crowding-out theory still yields important insights on the evolution of feature and attention, I am unsure whether the initial motivation of shifting attention actually aligns with the described problem setup. \n3) **Intuition on Theory.** The authors employ an extensive theoretical analysis under assumptions that are not further justified. For example, the moment and boundedness conditions (Eq. 3 - it would help referencing equations, if the authors could add numbers beyond the first two equations) and following sup-bounds are no further elaborated. I suspect that the particular setup is a key assumption to why a model should prefer x->y against predicting z->y. Generally, I the authors should discuss more clearly and extensively when the conditions necessary for CCO arise and can actually hold in real-world scenarios.\n4) **Choice of the Model.** The authors do not elaborate why the particular two-key attention architecture is chosen over a standard transformer with concatenated features. The paper could be improved by stating whether this particular architecture is required to obtain the observed results and how it would generalize to larger models. Similarly, it not described whether the value vector $\\tilde{v}$ assigned to $q^t$ is a concatenation of $v_x$ and $v_z$ for both attention products $l_{x/z}$ or whether they are fed separately into $l_x$ and $l_z$. \n\n\n\n**Missing references to the appendix:** The proofs in the appendix should be referenced in the main text. Also Fig.2 is un-referenced.\n\n**Typos:** \n\n- Line 20 and 22: \"to\" -> \"the\"\n- Line 47: GD is abbreviated as GD before the mention in line 57.\n- Line 179: \"feature\"->\"features\"\n- Line 188: \"exhibit\" -> \"exhibits\"\n- Line 205: \"dominance\" -> \"dominant\"\n- Line 206: \"feature\" -> \"features\"\n- Line 225: \"with noise\" is repeated.\n- Line 294: \"a diagonal\" is repeated.\n- Line 306: \"preserves\" -> \"preserving\"\n- Line 478: \"phenomenon CCO for Transformers\" -> \"phenomenon for transformers training dynamics called CCO\" or something similar.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sI7hVR3HBG", "reviewer": "Reviewer_9kRu", "summary": "The authors study the phenomenon that they name CCO (correlations crowding out) in transformers. Under the hypothesis of strongly correlated causal features, the authors show that simplified transformer architectures, trained with GD, can filter out spurious dependencies in the data and synchronize with the causal directions. The analysis of the learning dynamics leads to the observation of two phases: \"occupation,\" in which there is a rapid growth of the weight in the causal direction, and the \"crowding out,\" in which the attention aligns its logits to favor the causal features and suppress the spurious ones.\nThe theoretical results are accompanied by simulations on both synthetic and realistic data.", "questions": "- I wonder how much of the result has to do with _causality_ and how much with _dominant correlation_. Since you do not perform invariance study with interventions, how exactly does the fact that $x$ \"causes\" $y$ impact the theorems? In a specular setting with dominant correlations between $y$ and $z$, and weaker ones between $x$ and $y$, could we see a specular result, or would the inverted causal dependencies forbid it?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- The concept of CCO is interesting and the setting of the theoretical analysis, albeit simplified, correctly captures the main ingredients for the study of CCO in transformers trained with GD.  \n- The two-phases unveiled are both non-trivial and interestingly link causal data structure to architectural properties of self-attention.  \n- The proofs of the theorems are convoluted and contain non-standard ideas.  \n- Causal dominance seems to be a natural requirement satisfied in concrete cases.", "weaknesses": "- The transformer model is rather simplified with respect to practice.  \n- The training algorithm is also greatly simplified, with attention and $w$ learned at different moments. It is not clear from the theorems whether this fundamentally shapes the two-phase behavior described, or whether the two phases would persist also in the case of a standard training protocol. This weakness is partially compensated by Section 6.1, in which we see simulations with standard GD show the same phenomenology.  \n- Figure 4a is hard to see, and it is not clear what point it makes.  \n\n#### Minor remarks\n- There is a mistake in Example 25: in line 1660 the coefficient of $x_1$ is 1.  \n- In line 284 the index should be $t+1$, not $t$.  \n- The notation is sometimes confusing; for instance, why do you introduce $\\tilde v$ if it is equal to $q$?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uKAi3m4N1c", "reviewer": "Reviewer_SPKC", "summary": "This manuscript identifies a phenomenon termed Correlation Crowding-Out (CCO), whereby gradient descent on Transformers under a dominance gap in correlation causes the network to rely solely on causal features.", "questions": "1. Include experiments with mixed-strength spurious features to test CCO boundary conditions.\n2. Discuss links to implicit-bias literature (max-margin GD) quantitatively.\n3. Clarify how “cause” vs “effect” features are labeled in synthetic datasets.\n4. Consider renaming to “dominant-feature selection” to avoid causal over-interpretation.\n5. Although the concept is interesting, it still needs to extend CCO tests to large language models (e.g., BERT for sentiment analysis) to determine whether the two-phase dynamics (occupation -> crowding-out) generalize beyond Vision Transformers.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Conceptually elegant, with positions causal invariance as an implicit-bias phenomenon.\n2. Theoretical development (Dominant-Coordinate Condition) is clear and connects optimization dynamics to causality.", "weaknesses": "1. The “cause-only” notion is defined in correlation space, not interventionally -- terminological overreach.\n2. Limited experimental scope: small-scale ViT toy data only.\n3. Mathematical exposition skips proofs for convergence guarantees.\n4. Does not quantify how large the dominance gap must be in realistic data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)