Noise Stability of Transformer Models
Authors: Themistoklis Haris
Zihan Zhang
Yuichi Yoshida
Venue: ICLR 2026 Poster (2026)
Keywords: transformers
simplicity bias
noise stability
regularization methods
spectral concentration
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce noise stability in Transformer models as an alternative proxy for explaining simplicity bias and propose a corresponding regularization method that we observe accelerates grokking.

Abstract: Understanding simplicity biases in deep learning offers a promising path toward developing reliable AI. A common metric for this, inspired by Boolean function analysis, is average sensitivity, which captures a model's robustness to single-token perturbations. We argue that average sensitivity has two key limitations: it lacks a natural generalization to real-valued domains and fails to explain the "junta-like" input dependence we empirically observe in modern LLMs. To address these limitations, we propose *noise stability* as a more comprehensive simplicity metric. Noise stability expresses a model's robustness to correlated noise applied to *all* input coordinates simultaneously. We provide a theoretical analysis of noise stability for single-layer attention and ReLU MLP layers and tackle the multi-layer propagation problem with a covariance interval propagation approach. Building on this theory, we develop a practical *noise stability regularization* method. Experiments on algorithmic and next-token-prediction tasks show that our regularizer consistently catalyzes grokking and accelerates training by approximately $35$\% and $75$\% respectively. Our results establish noise stability as a powerful tool for understanding and improving modern Transformers.

Published as a conference paper at ICLR 2026
NOISE STABILITY OF TRANSFORMER MODELS
Themistoklis Haris"
Zihan Zhang
Department of Computer Science
National Institute of Informatics
Boston University
Tokyo, Japan
Boston, USA tharis@bu. edu
z dihan@nit ac. ip
Yuichi Yoshida
National Institute of Informaties
Tokyo. Japan yyoshida@nii. ac. ip
ABSTRACT
Understanding simplicity biases in deep learning offers a promising path toward developing reliable AL. A common metric for this, inspired by Boolean function analysis, average sensitivity, which captures a model's robustness token perturbations. We argue that average sensitivity has two key limitations: it lacks natural generalization to real-valued domains and fails to explain the "junta-like" input dependence empirically observe in modern LLMs. To ad- dress these limitations, we propose noise stability as a more comprehensive sim- plicity metric. Noise stability expresses model's robustness to correlated noise applied to all input coordinates simultaneously. We provide a theoretical analy- sis of noise stability for single-layer attention and ReLU MLP layers and tackle the multi-layer propagation problem with covariance interval propagation ap- proach. Building on this theory, we develop practical noise stability regularizu- tion method. Experiments on algorithmic and next-token prediction tasks show that our regularizer consistently catalyzes grokking and accelerates training by approximately 35% and 75% respectively. Our results sculpt new connection between signal propagation in neural networks and interpretability, with noise stability emerging as powerful tool for understanding and improving modern Transformers.
single-
next-token-prediction
INTRODUCTION
Simplicity Biases have been a promising direction of study in recent years (Shah al., 2020 Va-
sudeva et al. 2024 Bhattamishra et al., 2022) they provide unifying framework for generaliza- tion, interpretability and robustness. Neural networks, including Large Language Models (LLMs), often converge to the simplest possible functions that explain the training data. Because simpler functions are inherently more interpretable and robust, this bias provides solid theoretical frame- work for improving model reliability.
To quantify simplicity, current research often draws on concepts from Boolean function analysis (O'Donnel] [02T]. In particular, theoretical work on Transformers (Vaswani et al. (017) has high- lighted average sensitivity, the expected change in a model's output given a single-token pertur- bation, a key metric. Bhattamishha et al. (2022) formally studied this metric in Transformers, showing that they learn functions with lower sensitivity than LSTMs. Subsequent work has linked average sensitivity learnabilit:: Hahn (200)) and Hahn & Rofin (2024) demonstrated that Trans- formers struggle to learn functions with high sensitivity, such as PARITY. Further validating its utility, Vasudeva et al. (2024) connected average sensitivity the grokking phenomenon and pro- posed an extension of the metric beyond the Boolean domain.
"This work conducted while the author was an inter NII.

---

Published as a conference paper at ICLR 2026
Despite its usefulness, average sensitivity has notable drawbacks. Theoretically, its origins in Boolean analysis do not readily extend to real-valued domains. Empirically, fails to fully explain the "junta-like" input dependence we observe in models like GPT-2, GEMMA, and ROBERTA, where outputs rely on a small subset of inputs.
To address these shortcomings, propose to instead quantify simplicity bias via noise stability. Un- like average sensitivity's 's one-by-one perturbations, noise stability measures a function's 's resilience noise applied all inputs simultaneously, offering more robust spectral concentration guarantees. This approach naturally extends to real-valued domains via the Seateinn-Uhleneckk operator in the Gaussian measure, preserving formal connection the function spectrum and enabling a more powerful theoretical analysis.
1.1 OUR CONTRIBUTIONS
Our primary contributions are the following:
1. We observe that LLMs like GPT-2 exhibit a "junta-like" input dependence (Figure phenomenon not fully captured by average sensitivity and its extensions (Section 3 To better characterize this behavior, we propose noise stability (Section simplicity metric that naturally extends real-valued domains.
4 as comprehensive
2. We derive theoretical results on noise stability for single-layer Transformers and ReLU FFNs. We also provide novel insights for the multi-layer setting (Section 5 through proxy securrence-based analysis and stability interval propagation technique.
3. We propose noise stability regularization (Section 6 a method that consistently accel- erates grokking across synthetic (modular addition, sparse parity) and non-synthetic (next- token-prediction) benchmarks, reducing the training time required for generalization by xe 35% and 75% respectively.
1.2 RELATED WORK
Simplicity Bias in Deep Learning. The tendency of neural networks converge to simple func- tions has been a subject intense recent study. This simplicity bias (SB) is analyzed from several perspectives. One line of research connects SB to spectral concentration in the Conjugate Neural Tangent Kernel of networks (Yang & Salman 2019 Emami et al. 2021 Vasudeva et al.) 2024). An- other uses Fourier analysis characterize SB bias toward low-frequency functions (Xu al.,
2019 Rahaman et al. 2019). A large body of work investigates SB through the lens of training dynamics, often in shallow linear networks (Morwani et al. 2023 Zhang al. [2019] Yun al.
2020 Chen et al. 2020 Boursier & Flammarion [024]] Chizat & Bach [2200] Gatmiry et al. 2024
Tsoy & Konstantinov 2024). Beyond mechanisms, SB has been linked generalization Valle-
Perez et al. (2018), though it can sometimes lead to degenerate solutions (Shah et al. 2020). and
has been correlated with adversarial robustness (Min & Vidal 20244 Chen et al.) 2020). Specific to our focus, recent work has begun to explore SB in Transformers, particularly through the lens of token-to-token interactions in shallow models (Teney et al. 2025 Rende et al. 2024).
Sensitivity Analysis in Transformers. To develop a computationally tractable proxy for spectral concentration, recent work has adopted average sensitivity from Boolean function analysis. [Bhtt tamishra al. (2022) showed that Transformers are more biased towards low-sensitivity functions than LSTMs, enabling generalization even with noisy labels, Hahn (2020); Hahn & Rofin (2024) established that Transformers struggle to learn high-sensitivity functions like parity, despite having the capacity to represent them. Further, Vasudeva et al. (2024) demonstrated that average sensitivity can also serve as metric for tracking progress in grokking.
Signal Propagation in Neural Architectures Motivated by the challenge of gradient instability during training (e.g., vanishing exploding gradients), a rich body of literature has established
theory of signal propagation in neural networks (Poole et al. 2016 Ji et al. 2023: Lee et al. 2019
Kohan et al. 2023 Lou al. 2022 Wang al. (024)). Leveraging tools from statistical phy sicss, mean-field dynamical systems (Schoenholz et al. 2016). kernel methods (Martens et al. 202T),
and random matrix theory (Pennington et al. 2017 this now-mature theory offers critical insights
2

---

Published as a conference paper at ICLR 2026
into maintaining models the "edge of chaos"-a state of criticality where signal information is preserved without diverging or collapsing. Although theoretical analyses are frequently constrained to models initialization, they provide valuable guidelines regarding the efficacy of optimization techniques such as dropout (Pretorius et al. 2018) and residual connections (Fischer et al. 2023) the mechanics of which were previously under-theorized.
Central to this literature are two key quantities: the correlation between the pre-activations of distinct inputs, and the pre-activation variance The notion of noise stability proposed in our work can be viewed as a more parsimonious analogue to signal correlation and C'-aps. Indeed, our results on propagating noise stability through MLP and Attention layers mirror similar findings the signal propagation literature. However, we instead arrived at noise stability through questions of spectral concentration and simplicity bias, rather than initialization and training stability. This distinction suggests a compelling link between the two fields of study that can be further explored in future work.
Signal propagation has also been examined within Transformer architectures to address pathological issues such as unstable gradients and rank collapse (Saada et al.] 2024 Noci et (022 Huang et al. 2023). These works also generally adopt dynamical systems perspective; in contrast, we utilize discrete viewpoint and explicitly draw connections to simplicity biases. Unifying these theoretical landscapes to encompass both discrete and continuous domains remains promising direction for future research.
Finally, noise-related regularization is not unprecedented. Hua et al. 2023) for instance, propose noise-stability regularization method for fine-tuning Transformers. Our approach differs fundamen- tally in motivation, scope of application, and the definition of correlated noise.
Generalizing Sensitivity to Continuous Domains. The concept of average sensitivity has been generalized to real-valued domains via geometric influences (Keller et al., 2012 2014). This for- mulation is equipped with an analogue of Friedgut's junta theorem for continuous spaces (Bouyrie) 2017). unifying prior results across various discrete and continuous measures (Benjamini et al. Wimmer 2014).
2016
Noise Stability and Sensitivity. Our work is most closely related to that of Li & Mossel 2025)) who study noise sensitivity-a dual measure our noise stability-for hierarchical functions. They use an inductive argument to propagate sensitivity bounds across layers in a simple, non-intersecting neural network. Though they do study Transformers in practice, their layer-wise propagation strategy directly inspired our approach for noise stability intervals in multi-layer Transformers.
SETUP
We define a simplified Transformer as an L-layer model that maps an input sequence X gaxd to distribution over ne classes. Each layer E [Z] contains H attention heads. A head j takes the layer input Y, E Raad and computes an output as,j Rnxd via:
Oig = oY WT a. WK YY Wy,
Here, WQ..j, WK.YJ, Wv.,, E Raxd are weight matrices and is the row-wise softmax The head outputs are concatenated and passed through a Multi-Layer Perceptron (MLP) with a ReLU activation, d(x) = max(0, z): à; Φ((a,,1 o.... (")). The final layer's output â, is then mean-pooled and projected produce class logits Rh using an output matrix Wo Rdxx..
For theoretical simplicity, this definition omits elements like residual connections, layer normaliza- tion, and attention masks, though we include them in our experiments.
Poole et al. (2016) explicitly define iterative maps for these quantities: Q-maps capture variance and C- maps capture correlation.
For convenience, in theoretical results we will ignore the 1/vá factor that appears in these expressions.
3

---

Published as a conference paper at ICLR 2026
2.1 BOOLEAN FUNCTION ANALYSIS
Our work draws on Boolean function analysis, which studies functions f 1111" + R by ex- panding them as multilinear polynomials via the Fourier spectrum: f E Luginl fuxe, where (z) Ieu 22 the basis functions. For full overview, Appendix
A key property is the influence of coordinate E [n), which measures the expected impact of flipping the input zς;
((())-
)}
27
[nf:[[] == (111
2
S3
The total influence across all coordinates is the average sensitivity, 7[ʃ] ビズ measure function's robustness to noise.
Inf: common
MODELS ARE OFTEN "SIMPLER" THAN EXPECTED
The "simplicity" of Boolean function f (=1)² + R is formalized by its dependence few variables. One way to characterize simplicity is spectral concentration, where the most of the Fourier mass is on low-degree coefficients. A function is (e, ks-specrallyy concentrated if its mass on terms of degree i or higher is bounded:
n
Σ w' 11 where W Σ PJ J=k
10
A stricter notion is a k-juunta, a function that depends at most h coordinates. A function with low average sensitivity is simple both senses: a function is always (s, -speccslll concen- trated, and Friedgut's 's Junta Theorem (Kelman et al. 2021 Friedgut 1998 structurally close to a junta:
shows it must also be
Theorem 3.1. For every >> 0, there exists k-junta g 1-1,11" - {-1, such that [f + [(x)]] where the number of variables k on which depends is bounded by k S 20((h/))
While prior work has used average sensitivity analyze model simplicity (Vasudeva et al. 2024
Hahn & Rofin (024), we argue it has significant theoretical and empirical drawbacks as metric for simplicity bias.
Theoretical Drawbacks: Extending to Real- Valued Domains First, the theoretical foundation of average sensitivity in Boolean domains difficult to extend the real-valued functions of deep learning. Approaches using generalized domains that mimic finite field are cumbersome, as sen- sitivity is not naturally defined, and its estimation via sampling is impractical.
A more robust alternative is geometric influence (Keller et al. 2012). defined for a smooth function f: R" + R and measure y as:
19 Σ |0,llo)))))
1=1
Crucially, this measure has strong theoretical backing, including analogue of Friedgut's junta theorem for hypercontractive measures (Bouyrie) 2017). This makes a more suitable tool for analyzing neural networks and use it in empirical study below.
Empirical Drawbacks: Mismatch with LLM Behavior Second, average sensitivity fails to capture the empirical simplicity of LLMs. The exponential bound influential variables from Friedgut's theorem is too loose to explain the behavior of modern Transformers, which exhibit far stronger influence concentration. To demonstrate this, we analyze the geometric influence of each input token for GEMMA-2B, ROBERTA, and GPT-2 on sequences of n = 256 tokens. Friegut's 's theorem predicts < 1024 variables with influence at least 0.1 I[f], yet Figure I only shows 5-10 such variables exist, meaning the bound is loose. Our analysis also reveals three key patterns with further experimental details available in Appendix B
See O' D'nnnlll 2021), Chapter for exposition on Boolean Function Analysis on the hypergrid.

---

Published as a conference paper at ICLR 2026
Total Model Influence Across Token Positions (Last Layer) Multiple Models
GEMMA Total Influence (%)
3.0
ROBERTA influence (%) GPT2 Total Influence (%)
Percentage of Total Influence (%)
2.5
1.0
0.5
0.0
50
100
150
200
250
Token Position
Figure 1: Comparing the per-coordinate geometric influence of three models for 256.
Junta-like Concentration: A small subset of tokens have disproportionally high influence.
Positional Bias:
The first and last tokens are consistently the most influential. This part agrees with observations made in the KV Cache Compression literature about "attention sinks" (Xiao et al. 2023).
Sensitivity:
Every token has a non-zero influence, indicating that the models are sensitive to all inputs, even if asymmetrically.
4 NOISE STABILITY AS MEASURE OF CONCENTRATION
For more holistic characterization of simplicity that offers finer control over spectral concentration and easily generalizes to real- valued domains, we propose noise stability, a concept from Boolean Function Analysis (O' Donnelll 202T) that measures a function's resilience to correlated noise ap- plied to all inputs simultaneously.
This concept extends naturally to real-valued functions in L?(r)) by leveraging the Ornstein- Uhlenbeck (OU) semigroup Tç. This framework uses the basis of Hermite polynomials under the standard Gaussian measure = AO, 1), allowing for a direct spectral interpretation"] The corre- lated pair (X, Y') generated by adding scaled Gaussian noise X:
Definition 4.1. Let L?(r) where y is the Gaussian measure E². For p E (0, let ~ and Y = pX + 2 / p², where Z is independent of X. The noise stability of f is:
Stab,(f) = E [f(X) (X,Y)
Noise stability is useful because it relates directly to the function's spectrum through its Hermite- Fourier coefficients f(a), shown in Appendix
Stab,(f)
Σ (allf((a)
(3)
aEN#
This connection allows us to formally bound a function's spectral concentration. The following lemma shows that if function highly stable (i.e.. its stability is close its total variance), its Fourier mass must be concentrated on low-degree coefficients. For a fixed correlation and spectral tail budget e, the degree concentration T becomes smaller as the ratio 8/2 approaches zero.
*See Appendix
following Andersson Sjögren 2012), for introduction OU Semigroup Theory.

---

Published as a conference paper at ICLR 2026
Lemma (Spectral Concentration via Stability). f f L2(("). If Stab, 2 (1 f some p E 10 and 0 < < 1, then f is -specrspettrllly concentrated for any
T T log: (1-2 8
Proof Sketch. The proof follows from analyzing the action of the Ornteei-UUllenbeck semigroup Tp on the Hermite expansion of f. The full proof is in Appendix D
Spectral Concentration Bounds: Sensitivity vs. Stability We compare the predicted Fourier tail mass (percentage of weight in degrees 15 with 256) for several Transformer models, using bounds derived from average sensitivity versus those from noise stability. The results in Table show that the noise stability bound offers more accurate estimate of spectral concentration.
Model
Tail Mass Bound from [if]
Tail Mass Bound from Sta
GPT-2
0.003
0.0005
BERT
0.04
0.02
ROBERTA
0.19
0.02
GEMMA
0.043
0.0157
Table 1: Predicted Fourier tail mass (percentage of weight in degrees > 15) for Transformer models.
5 ANALYSIS OF NOISE STABILITY IN TRANSFORMER MODELS
We now present our theoretical results the noise stability of Transformer components. We begin by analyzing single ReLU MLP layer and an attention layer, and then use these results analyze the propagation of stability through multi-layer network.
5.1 NOISE STABILITY OF A SINGLE RELU MLP LAYER
We first analyze the stability of MLP layer with ReLU activation, result closely related to the arc-cosine kernel. Consider a pair of p-correlated standard Gaussian inputs (X, Y),, whose joint distribution is:
(x,Y) ~ C0.
P))
The noise stability of the ReLU function is given by the following theorem. Theorem 5.1. The noise stability of the ReLU function under the standard Gaussian measure is:
E (RLU(())) = 2m (V1- + pla arccos p)
(X.Y)
Proof. The proof by direct integration and can found Appendix
While exact, this expression is unwieldy for analyzing layer composition. For practical purposes, wpll-approximmtte by its second-order Taylor expansion around 0:
2T (V1- p² p(á arccos pi)
1 2n
++ p² + O(p)).
(4)
4T
5.2
NOISE STABILITY OF A SINGLE ATTENTION LAYER
We next analyze the noise stability of single attention layer, defined as f(X) = = XWy. The analysis depends critically on the structure of the product == Wo so we consider three representative cases.

---

Published as a conference paper at ICLR 2026
The Identity Case (W = Ia) When W is the identity matrix, the attention mechanism simplifies. In the nigh-dimensional limit, the attention matrix a o(XX") converges to the identity matrix I,s causing the layer to act as linear transformation. This results in linear relationship between input and output stability (see F
Theorem 5.2. Let X NO, luxa) and px + ZV1 p² for Z NN, Luxd) independent of X. Let ((X) oX XWV Then for any [n], [d]:
lim ( p- d0000
(1(v) o1)
Proof Sketch. As + oo, we show that o (XX") + in in probability, The stability calculation then reduces to that of linear layer, with cost of ((1). The full proof is in Appendix F.2
The Low-Rank Case (W = UU?) If W is random low-rank matrix, where U E Rd with k << d, the analysis reduces to the identity case. The matrix acts as projecting the imput vectors into k-dimensional space while approximately preserving their inner products (Matoušek) 2008). Consequently, the projected attention matrix again converges to the identity, and the stability remains linear in p.
transform,
The Unstructured Case (W (O laxa)) When W random Gaussian matrix, modeling a randomly initialized layer, the behavior changes. For large d, we the attention matrix Ax = (XWXX) X concentrates towards random permutation matrix, meaning each output token attends to single, randomly chosen input token.
Correlation Propagstion Attention [Linear Casel
Verlication the Intullive Modeel for Expectation
Etiivatee
Dinect Mensiremen hases Stabily
Unsar Stopet 6.3225. Expected skope 7.58LL
(p- dpl: [[WW])
Extmered EUS00i (FM9)
00 02
Correlation Confiiiinn
Figure 2: Stability of Single Layer Attention (Identity and Unstructured)
The stability now depends on the consistency of this permutation. For given output row i, suppose Ax selects input row k and Ay selects input row K. The stability is non-zero only if the attention pattern is preserved (k k'), Let s(p) P(k k') be the probability that the attention pattern is stable (see Figure 3 for an illustration). The total stability is the product of the linear stability and this probability factor:
p 11WW)
ifk
(I((X)
i
Averaging over the randomness of the patterns we obtain the following result, which we verify empirically in Figure (see Appendix for the proof).
Theorem 5.3. Let X ~ No Inxd) and Y = pX + z for Z ~ N Inxd) independent of X. Let f(X) = x XWWv. Then for any i E [n],j E [d], we have:
lim EL((X) 2 p s()) (I o(1), with s(p) = n J - doo
where Tc. fo are the joint CDF and PDF of bivariate normal distribution with correlation

---

Published as a conference paper at ICLR 2026
Attention Matrix Ax
Attention Matrix Ay with 0.01
Query Index
0.6
0.2
0.0
Key Index
Key Index
Figure 3: Comparing Ax and Ay for d = 128 and p = 0.01.
5.3 STABILITY PROPAGATION IN DEEP TRANSFORMERS
In the single-layer setting, we have shown that ReLU MLPs dampen stability (Theorem while attention layers can preserve it (Theorem 5.2). This raises the question whether such analysis can be performed for the multi-layer setting as well.
FFN Stability Propagation as Recurrence and Weak Dampening Consider ReLU Feed- Forward Network (FFN) and suppose we ignore inter-layer distributional shifts. Such a thought ex- periment not without precedent the moment propagation literature for neural networks (Wright et al. 2024 In this model, the correlation PL after layer L follows the recurrence:
PL 2x (V pL-1 + PL-1 ( arccos(Pt-1)))
(5)
Solving this non-linear recurrence analytically is difficult, so we shall instead use linear ap- proximation from Equation 4 to get the following proxy recurrence, which can be solved more easily:
2
2
(
L-1
PL 2π PL-1 PL
3
3
This suggests that noise stability for multi layer ReLU FFNs exhibits weak dampening, converging to the fixed point 2. This confirmed by a numerical evaluation of Equation (5) Figure 4b Indeed, Figure shows that multi-layer MI. does exhibit weak dampening.
3-
Nakie Sent Hamneer Layers U13 Rell1 FRN
Tiseectry of the Not- Rbccammenze Relation
W
2.3
Murlac (TL.ye5
Denntonn
Stability of a deep MLP network.
(b) Numerical solution Equation
For multi-layer Transformers, recurrence relation analysis does not yield weak dampening. If we let WoW and [I(W)).,ll2 = 7 < 1, recurrence similar to Equation (5) would yield candidate fixed points of チ[--フタ- However, we observe that for ? < the noise stability signal actually dampens fully to zero, suggesting that the attention map alters the distribution enough to preclude the weak dampening behavior. For more details, Appendix
8

---

Published as a conference paper at ICLR 2026
Noise Stability Intervals A more formal approach is to track rigorous upper and lower bounds on the noise stability as they propagate through the network. We derive such bounds for individual MLP and attention layers (Appendix [.] and Appendix L.2). Further empirical work is needed to determine the tightness these bounds in practice.
NOISE STABILITY REGULARIZATION AND ITS BENEFITS
We established in Section 4 and Section that high noise stability is desirable property for cre- ating robust and interpretable models. To this end, we introduce a regularizer designed to orient a model's training towards to away from stable functions. We designed our regularizer be (1) differentiable with respect to the model's s parameters, and (2) data-dependent, meaning that the regularization should be evaluated on the model's outputs for training data, not just its parameters.
Definition 6.1 (Noise Stability Regularization). Let M [U] - A(C) be model, E [U] be an input sequence, S E {0, }} be the orientation parameter, and p [-1, be correlation parameter. The S-oriented noise stability regularizer is defined as:
(X)= (-1)s M(X); M(Y); where:
-x)=
(6)
X
with probability L+e otherwise.
Y, = Z uni))rm((U]),
Setting the orientation parameter S = encourages stability. For a differentiable loss function E, the regularized loss then becomes (M, X) == M, +) R.s.sppX where y a hyperparameter controlling the regularization strength. We test the effect of positively-oriented noise stability (S 1) by training a Transformer from scratch on two tasks known to exhibit "grokking": noisy sparee parity and modular addition.
Note that calculating the noise stability regularizer requires just one additional forward pass through the model per training iteration. It is interesting direction whether the regularization can be applied in a more cost-effective manner.
Experimental Setup We standard decoder-only Transformer (Appendix D For all experi- ments, we compare regularized and non-regularized models. All other hyperparameters, including the random seed, are held constant across (Figure
NOISY k- E-SPARSE PARITY (NSP) We learn the function f(x) θ ic? for an input I E {0, 1}" and a secret sparse index set C [n]] of size Each training label is flipped with fixed probability 77. This problem is well-studied in learning theory (Chen et al.] (2024] Feldman et al. (009). and Transformers can solve it for small values of & (Bhattamishra et al. 2022). We test on inputs of length [10, 100] with & (2, 3). using (7, p) (0.05, 0.05).
MODULAR ADDITION This task is to compute (X + Y) (mod K). We study it as standard benchmark for grokking in Transformers (Nanda al. 2023). For our experiments, we use prime modulus K 113, train for 10, 000 iterations, and set (Y, p) (0.75, 0.25).
Noise Stability Regularization Catalyzes Grokking. Training Transformers on these tasks ex- hibits an "emergence" phenomenon, where validation loss drops suddenly after long period of stagnation. We find that noise stability regularization acts as catalyst for this emergence. For modular addition, regularization reduces the iterations required from 4500 to 3300, 36% acceleration. We observe a similar zs 35% speed-up for the noisy sparse parity task (Figure
Noise Stability Evolution During Training By observing the noisy sparse parity task, find that a Transformer's noise stability naturally decreases during training to match the target function, serving as a leading indicator of generalization (Appendix J.2

---

Published as a conference paper at ICLR 2026
Validation Accuracies: Addition
Training and Validation Losses: Addition
Ri arizition
Regu betcation
hasse Stude By feguuaaraaion Train Loss Noice Stabi By Regulritatio Locs
Aucuresy (46)
Loss
No Reguantzttion Vei Accuraty
Nasse Stade Insuulatiation Vill Accuracy
2200
2000
2200
5000
5000
2000
T030
1300
2000
2300
6000
5000
6000
X0C
Epoch
Epoch
Validation Accuracies: Noisy Sparse Parity
Training and Validation Losses: Noisy Sparse Parity
by Reguutaraatinn Loss
hotse Siati by Regulortztion Loss
Accuracy (50)
LOOS
Stabi Reqularization Accuracy
Epoch
Epoch
Figure 5: Noise Stability Regularization accelerates training.
6.1 NON-SYNTHETIC EXPERIMENTS ON LANGUAGE GENERATION
We also tested noise stability regularization for the task next-token next-token-rrediition on WikiText-2-vl (Appendix (.3). We observed that regularized training reaches high validation accuracy in z 75% fewer iterations (Figure 6). The noise stability of the regularized model notably stays high, while non-regularized models become increasingly unstable.
Validation Aocuracies WicTeed NTK
Noise Stability vs Epoch
Speeh
Cpoct
Figure 6: Noise Stability Regularization for
(NTK) on WikiText-2
7 CONCLUSION
In this work, we introduced noise stability measure of simplicity bias in Transformers, argu- ing theoretically and empirically that better explains the spectral concentration observed in LLMs than average sensitivity. We also proposed noise stability regularizer and found that it unexpectedly catalyzes grokking, offering a potential avenue to accelerate model training. Our findings open sev- eral avenues for future research, including demystifying the mechanism of moment propagation in deep networks and understanding the practical limits and theoretical underpinnings noise stability regularization in LLM settings.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
We thank the anonymous ICLR 2026 Area Chair (QeA8) for highlighting the extensive literature regarding signal propagation. TH additionally thanks the faculty and staff at the National Institute of Informatics for providing a stimulating research environment throughout the summer internship program. YY is supported by JSPS KAKENHI Grant Number 24K02903.

---

Published as a conference paper at ICLR 2026
REFERENCES
Adam Andersson and PETER Sjögren. Ornstein-uhlenbeed theory in finite dimension. Preprint, 12, 2012.
Itai Benjamini, David Ellis, Ehud Friedgut, Nathan Keller, and Arna Sen. Juntas in the & fi-grid and lipschitz maps between discrete tori. Random Structures & Algorithms, 49(2):253-279, 2016.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. arXiv preprint arXiv:2211.12316, 2022.
Etienne Boursier and Nicolas Flammarion. Simplicity bias and optimization threshold in two-layer relu networks. arXiv preprint 2400:244, 2024.
Raphaël Bouyrie. An unified approach the junta theorem for discrete and continuous models.
arXiv preprint arXiv:1702.0000 2017.
Peijie Chen, Chirag Agarwal, and Anh Nguyen. The shape and simplicity biases of adversarially robust imagenet-trainee trained cnns. arXiv preprint arXiv:2006. 09373, 2020.
Xue Chen, Wenxuan Shu, and Zhaienhe Zhou. Algorithms for sparse Ipn and Ispn against low-noise.
arXiv preprint arXiv:2407. htf.::4077/99111.. 19215, 2024.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory, pp. 1305-1338. PMLR, 2020.
Melikasadat Emami, Mojtaba Sahraee-Ardakan, Parthe Pandit, Sundeep Rangan, and Alyson K Fletcher. Implicit bias of linear rnn. In International Conference on Machine Learning, pp. 2982-2992. PMLR, 2021.
Vitaly Feldman, Parikshit Gopalan, Subhash Khot, and Ashok Kumar Ponnuswami. On agnostic learning of parities, monomials, and halfspaces. SIAM Journal on Computing, 39(2):606-645, 2009.
Kirsten Fischer, David Dahmen, and Moritz Helias. Optimal signal propagation in resnets through residual scaling, arXiv preprint arXiv:2305.07715. 2023.
Ehud Friedgut. Boolean functions with low average sensitivity depend on few coordinates. Combi- natorica, 18(1):27-35. 1998.
Lazarus Fuchs and Kurt Hensel. Journal für die reine und angewandte Mathematik. W. de Gruyter, 1866.
Khashayar Gatmiry, Zhiyuan Li, Sashank J Reddi, and Stefanie Jegelka. Simplicity bias via global convergence of sharpness minimization. arXiv preprint arXiv:2410. 10.44001, 2024.
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
Michael Hahn and Mark Rofin. Why sensitive functions hard for transformers? arXiv preprint arXiv:2402.099%. 099663 2024.
Hang Hua, Xingjian Li, Dejing Dou, Cheng-Zhong Xu, and Jiebo Luo. Improving pretrained lan- guage model fine-tuning with noise stability regularization. IEEE Transactions on Neural Nei- works and Learning Systems, 2023.
Zhongzhan Huang, Mingfu Liang, Jinghui Qin, Shanshan Zhong, and Liang Lin. Understanding self-attention mechanism via dynamical system perspective. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 1412-1422, 2023.
Peng Ji, Jiachen Ye, Yu Mu, Wei Lin, Yang Tian, Chittaranjan Hens, Matjaž Pere, Yang Tang, Jic Sun, and Jürgen Kurths. Signal propagation complex networks. Physics reports, 1017:1-96, 2023.
Nathan Keller, Elchanan Mossel, and Arnb Sen. Geometric influences. 2012.
12

---

Published as a conference paper at ICLR 2026
Nathan Keller, Elchanan Mossel, and Arnab Sen. Geometric influences ii: Correlation inequalities and noise sensitivity. In Annales de I'IHP Probabilités et statistiques, volume 50, pp. 1 1121-1139, 2014.
Esty Kelman, Subhash Khot, Guy Kindler, Dor Minzer, and Muli Safra. Theorems of kkl, friedgut, and talagrand via random restrictions and log-sobolev inequality. In /2th Innovations Theo- retical Computer Science Conference (ITCS 2021), pp. 26-16 1. Schloss Dagstubl Leibniz-Zentrum für Informatik, 2021.
Adam Kohan, Edward A Rietman, and Hava T Siegelmann. Signal propagation: The framework for learning and inference in forward pass. IEEE Transactions Neural Networks and Learning Systems, 35(6))8585-8599 2023.
Namhoon Lee, Thalaiyasingam Ajanthan, Stephen Gould, and Philip HS Torr. signal propagation perspective for pruning neural networks at initialization. arXiv preprint arXiv:1906 1e06:06666..777 2019.
Rupert Li and Elchanan Mossel. Noise sensitivity of hierarchical functions and deep learning lower bounds in general product measures. arXiv preprint arXiv:2502.05073 2025.
Yizhang Lou, Chris Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In International Conference on Machine Learning, pp. 14248-14282 PMLR, 2022.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint arXiv:21 10.01765, 2021.
Jiff Matoušek. On variants of the johnson-lindenstraus lemma. Random Structures & Algorithms, 33(2):142-156, 2008.
Hancheng Min and René Vidal. Can implicit bias imply adversarial robustness? arXiv preprint arXiv:2405.15942, 2024.
Depen Morwani, Jatin Batra, Prateek Jain, and Pranceth Netrapalli. Simplicity bias in 1-hidden layer neural networks. Advances in Neural Information Processing Systems, 36:8048 18880088-775,,, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05221, 2023.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198-27211, 2022.
Ryan O' O'Donnell. Analysis of boolean functions. arXiv preprint arXiv:2105. 10386, 2021.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learn- ing through dynamical isometry: theory and practice. Advances in neural information processing systems. 30, 2017.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponen- tial expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
Arnu Pretorius, Elan Van Biljon, Steve Kroon, and Herman Kamper. Critical initialisation for deep signal propagation in noisy rectifier neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International conference on machine learning, pp. 5301-5310. PMLR, 2019.
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. A distributional simplic- ity bias in the learning dynamics of transformers. Advances in Neural Information Processing Systems, 37:96207-96225 2024.
13

---

Published as a conference paper at ICLR 2026
Thiziri Nait Saada, Alireza Naderi, and Jared Tanner. Mind the gap: spectral analysis of rank collapse and signal propagation in attention layers. arXiv preprint arXiv:2410.07799, 2024.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint 1rXiv:1611.0125: 2016.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 999977-9585,, 2020.
Damien Teney, Liangze Jiang, Florin Gogianu, and Ehsan Abbasnejad. Do we always need the simplicity bias? looking for optimal inductive biases in the wild. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 19-90, 2025.
Nikita Tsoy and Nikola Konstantinov. Simplicity bias of two-layer networks beyond linearly sepa- rable data, arXiv preprint arXiv:2405.17229, 2024.
Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louís. Deep learning generalizes because the parameter-function map is biased towards simple functions. arXiv preprint arXiv: 1rXiv:/805.08555 2018.
Bhavya Vasudeva, Deqing Fu, Tiany Zhou, Elliott Kau, Yougi Huang, and Vatsal Sharan. Simplicity bias of transformers learn low sensitivity functions. arXiv preprint arXiv:2403.06925, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Zhichao Wang. Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propa- gation in deep neural networks. In The Thirty Seventh Annual Conference on Learning Theory, pp. 4891-4957. PMLR, 2024.
Karl Wimmer. Low influence functions over slices of the boolean hypercube depend on few coordi- nates. In 2014 IEEE 29th Conference Computational Complexity (CCC), pp. 120-131. IEEE, 2014.
Oren Wright, Yorie Nakahira, and José MF Moura. An analytic solution to covariance propagation in neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 4087- 4095. PMLR, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.1745. 2023.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light deep neural networks. arXiv preprint 1rXi::1901.06523 2019.
Greg Yang and Hadi Salman. A fine ane-grained spectral perspective on neural networks. arXiv preprint 1eXiv:1907.10599 2019.
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. arXiv preprint arXiv:2010.02501, 2020.
Yaoyu Zhang, Zhi-Qin John Xu, Tao Luo, and Zheng Ma. Explicitizing an implicit bias of the frequency principle in two-layer neural networks. arXiv preprint arXiv: 1905. 10264, 2019.
14

---

Published as a conference paper at ICLR 2026
APPENDIX CONTENTS
Essentials of Boolean Function Analysis
16
A.1 Influence and Sensitivity
16
B More Experimental Results with Total Influence
17
Basics of Cernstin-UUlenbere Theory
17
Hermite Polynomial Basis and its Properties
17
The Ornstein- Uhlenbeck Operator and Semigroup
19
The Mehler Kernel
20
From OU Theory to Stability
22
1)
Proof of Lemmal
22
Proof of Theorem
24
Approximating with quadratic
27
Proof of 55.2
28
Preliminaries
28
F2
Proving the Theorem
28
Noise Stability Propagation in the Unstructured Case
30
Full Noise Stability Dampening in Multi-Layer Transformers
31
Noise Stability Interval Propagation
32
MLP Stability Under Bonami- Gaussians
32
1.2
Stability Propagation through single Attention Layer
33
J
Noise Stability Regularization Experiments
35
J.1
Architecture Layout, Hyperparameters and Training Details
35
0.2
Evolution Noise Stability During Training.
36
J.3
Experiments Language Generation
36
15

---

Published as a conference paper at ICLR 2026
A ESSENTIALS OF BOOLEAN FUNCTION ANALYSIS
We now review some basic facts on Boolean Function Analysis. Boolcan function is defined on the hypercube: f {=1} + R. Therefore, can think of the space of boolean functions as I fundamental fact that every boolean function can be represented uniquely as a multilinear polynomial. This is natural outcome of considering the following set of vectors in IR?":
B Xx(x)
II LEU
J3 : U S [m]
For any U, V we have that:
(xxxxv) 2" (Rue(t))
2" II E[r.]
iEVAV
which is equal to 1 if and only U = V. Thus B is an orthonormal basis for the set of Boolean functions and so we can conclude the following:
Theorem A.1 (Fourier Expansion of Boolean Functions). Let f: {+1]*=+ - be boolean function.
Then write f uniquely as a mudtilinear polynomial Fourier Expansion:
f f=
= fo XU UC(m]
where So = XU) the Fourier Coefficients of f. The degree deg(J) of f is defined as the largest cardinality U for which fo +
deg(f) == max{[U]: fu + 0}
A.1 INFLUENCE AND SENSITIVITY
Given f {+1)* R R. define define its sensitivity by considering fluctuations in its output when one single bit pertrubed:
Definition A.2 (Influence). We define the influence of coordinate i as:
f(x) - f(x 2
Inf: [ʃ] = E,,
2
where x is x with the i-th coordinate flipped. Then we define the total influence of f as:
I[11 Σ [nf
We often think of total influence measure of average sensitivity.
The following lemma gives us a Fourier representation of influence:
Lemma 2. If fsxs then:
Inf.[/] = Σ f(S)2
and T[]]
Σ |S| f(S)2 SCIn]
Σ k. W*[]]
534
k=
where W*[f]= Σ f(S)2
1
The influence is related to other important quantities about f: its variance and its degree. Lemma (Poincaré's inequality). The variance of f is:
Varlf] = El 2
Σ f(S)
8
It is true that:
Varlf] < [[]]
The relationship between total influence and degree is very important because it allows us to repre- sent functions with low total influence low degree polynomials.
16

---

Published as a conference paper at ICLR 2026
B MORE EXPERIMENTAL RESULTS WITH TOTAL INFLUENCE
In this section we provide additional experimental details from examining total influence of four widely used models: GPT2, GEMMA-2 --MMAA---2B, RoBERTA. and BERT. As before, run experiments on 7 = 256 tokens and experiment both with being the uniform distribution and the distribution the model learns.
We analyze the geometric influence of these models by using forward hooks in Pytorch and col- lecting the gradients with respect the input embeddings. We analyze positional influence both across layers and different attention heads, finding that in deep Transformers there are often layers with minimal influence towards the output. Such sparsity is interesting phenomenon that could warrant further investigation.
(a) Model proxy Distribution
(b) Uniform distribution
Figure Per Layer Geometric Influence Heatmaps for 4 different models: We observe that GEMMA-2-B has very influential middle layer, while for ROBERTA and GPT2 the few lay- ers are more influential.
C BASICS OF
THEORY
We provide self-contained review of the Omstein-Uhlenbeck (OU) theory, culminating in the derivation of noise stability under the Gaussian measure. We encourage the interested reader to consult the excellent monograph of Andersson & Sjögren (2012) for more details.
C.1 HERMITE POLYNOMIAL BASIS AND ITS PROPERTIES
The Ornstein-Uhlenbed theory starts by considering the Gaussian measure d7(x) # |=|] dr in Rd Let (2) be the space of equare-integrable functions with respect this measure. Then, the Physicist's Hermite Polynomials turn out to be an invaluable way to spectrally decompose functions in this space.
17

---

Published as a conference paper at ICLR 2026
A. Ko Bet
Nottoon
(a) Model proxy Distribution
(b) Uniform distribution
Figure 8: Geometric Influence Plots for two attention heads different models. We again observe a junta-like spectral concentration phenomenon in every model. The first and last token positions are consistently the most influential.
Definition C.1 (Physicist's Hermite Polynomials). We define Ho 1 and:
d"
= (-1)ae dæn e
Some important properties of Hn are captured by the Lemma below: Lemma 4 (Properties of H,). The following identities hold:
d H Hn(x) = 1nH
dx
(a+) = H (x) 21
Also (x) is degree n polynomial with leading coefficien 2n
Proof.
Let
We have by the product rule that:
du
DH,((x) (-1)"e² e 2xDe +-1))² E D(D"
Then the generalized Leibniz formula gives:
12
n1 D (-2x)) D -
DH, (z) (-1)" " 2x a +--1)"
Only the terms where k = 0, non-zero, have:
(x, = (-1)"e" 2rD + (- (-2xD" + ("2) = -2HHn-1
Now, again by expanding (H) we get:
D= e- d dr --H441(1)++ 2xH1((x)
x² + (-11"a
d"
x²
dx
And this gives:
Hx+1(x) 2rH (x)
This allows us, by simple induction, show that the leading coefficient of H (x) is equal 2 1
18

---

Published as a conference paper at ICLR 2026
The most important statement in the Seatein-Uhlenbeckk theory is the following theorem: Theorem C.2 (Hermite Orthogonal Basis). The set [ forms complete orthogonal basis of (7) with ||H,llaac) = 2n/2 2n/2/n/ This motivates the definition and occasional use of the normalized Hermite basis:
1 2
hn
H (x)
which is orthonormal in (7).
Proof. To establish orthogonality, we have by repeated integration by parts that:
H Hm = (-1)"
Hm
(m)ee
(D" je "² dx
VT
1-)
Hm(x). (( (=
1)n
Hm (x)e - dr
If m n then D (x)) = 0 and that gives orthogonality. For m n, because the leading coefficient is 2 we have:
|| = 2n/2
desired.
Finally, we have to establish completeness. Recall from analysis that completeness means for every e (2) and every 0 there must exist other words, we seek show that the linear span of B is dense in (7), or, in even different words, that the closure of the linear span is the whole space (7) It suffices to show that if E 1,2(7) is such that (f. Hn) for all N then 0. To see this, let consider the spectral expansion of f with the following coefficients:
g E span () such that || gllr2(>) c. In
f(n) = / (())))())))((====
By Parseval's 's identity, we have that:
1f(n)12 2nn!
||
n=
And thus / = 0, which concludes the proof.
When considering 1, we define the Hermite basis as tensor product over multi-indices:
H. (x) = II Ha, (x), Nd
The same theorem about orthogonality and completeness holds in high dimensions as well.
C.2 THE ORNSTEIN- UHLENBECK OPERATOR AND SEMIGROUP
Next, we will define the Ornstein- Uhlenbeck operator. This operator is an analog to the Laplacian the (2(r)) space. It uses the adjoint of the partial derivative operator 0 := dm,
0
(0.f.g) = 4//2
(.f(z)g(z)e
=
[xx) (.g(x))e-1u7 de
(Integration by parts)
4/2 20
(J. (2x; 0.)g) Thus arrive the following definition:
RRcalll the adjoint of linear operator is an operator that moves to the other side of the inner product: (Tx,, v) (z, T" T)
19

---

Published as a conference paper at ICLR 2026
Definition C.3 (Omstein-- Constein-Uhlenbecc Operator), The OU Operator is defined as:
L Σ o, -
This operator is symmetric and positive and it has the wonderful property that the Hermite basis actually its eigenfunctions: Theorem C.4 (Bigenfunctions of the OU operator). The sei {Hn are eigenfunctions of
LHa |a|Ha
d
where |a] = Σ
2m1
Proof. Consider d = 1. We have for j # n:
(D" Hj) = nj Hz-1)
And if j we have (D" Hn Hn) D Ha-e: H
2nnil, so D) Hn Hn- Thus we have showed that:
for all d > and i [d]. So we know how the adjoint partial operator acts on the Hermite polynomials. Thus, we can figure out how the OU operator acts on Ho:
d
LH E O,HHH 2m
Σ 2,20, Ha- Σ
a.Ho la|
as claimed.
Now we can define the OU Semigroup. A semigroup a sequence of operators T3 that describe the evolution of process such that To = I and T + s) = Ttt) T(s
Definition C.5 (The OU Semigroup). The OU Semigroup defined as:
(Ti) e -t
where if = (a. T3 acts on f as:
e eh Σ -lll f(a)
aEN#
C.3 THE MEHLER KERNEL
It is easy to verify that (Tt) is indeed a semigroup. A very useful tool for us to analyze properties of certain semigroups is kernels. If semigroup written via a kernel it will be very easy to prove powerful properties on it. A kernel is just the analogs of matrix multiplication. When operator T acts on function f. imagine there being some kind of function K K(x, y) that for each f(y) tells us how much that value "contributes" to (xx):
Ts(x) = f K(x, (w)dd
We want find kernel for the OU semigroup:
Tef(x) [ M7 (z, (2)))))
Mehler already found this kernel Fuchs & Hensel 1866 in 1866!
'Technically there are also some analytic continuity requirements, hut we will not dive into them here.
Please refer to (Andersson & Sjögren 2012) for more details.
20

---

Published as a conference paper at ICLR 2026
Theorem C.6 (The Mehler Kemel). The Kernel for the OU semigroup is:
M7 (x,
Σ aENd
e (z))aa())
Proof. We shall just verify this. Choose some and see how the kernel acts on Hg:
[ /ERR ha Ha()))
e (
E ca
MER
la|<N
la|eN
e-10p (hg. HB) (a)
(only 3 survives)
= c 1|8
H (x))
H,
Now take N - oc and we arrive at the correct result.
The Mehler kernel has a really nice analytical expression. Starting from Hn(3)
(-1) D" D" and considering the Fourier Transform of the Gaussian: F(e- E (x) = vpi.
e- we have:
Hn( = (-1)" 2"²
{ Eiv de
V
And the Mehler kernel can be writter as:
(f v) =
tm N1, (z))
N1, (I)( (=) 2
VT { 2ie2 de
(-ife 1)n 1)" ( de
(Switch Σ and f)
n!
n=0
e
2) (1-e de
Letting g' = [VT c- and taking the inverse Fourier Transform get that:
1 eV² (2, V v1
2i6' x
16'1
Vi-o dg
1 21
00
And this brings to the following important theorem:
Theorem C.7 (Analytical form of Mehler's Kernel). We have that:
1
Mf (z. v) dd/2 - e 212/2/22
Some important consequences of this treatment is the easy, via Hölder's inequality, proof of the non-expansiveness of Tr:
Lemma (Non-expansiveence of OU operator). For any 221 we have that:
ITe S [[T)fllll $ IIffI(
We also easily get Mehler's formula, which will be the starting point in our stability argument. Just perform the change of variables
2
e-2 21
Formally, also need to justify why switching infinite summation and integration is possible by using dominated convergence, but will skip this here.
21

---

Published as a conference paper at ICLR 2026
to get:
Ti((X) / (1x+ dz(2))
where p:= for 2 In other words,
T,S(X)= = Ez-Nro.1) + 1²1
which allows view T₃ as the expected outcome of random process: take X, add some noise it get Y pX Z / p² and the expected value of S(Y).
C.4 FROM OU THEORY TO STABILITY
We are finally ready define stability familiar terms:
Definition C.8 (Gaussian Noise Stability). Let f E ()) and X ~ 7"h where y NV 1). Let E (0, 1) and = pX + ZV-1- p² for ~ ,,, independent from X. We define the stability of f as:
Stab,
E [f())((M) (X2Y)
We note the immediate connection between stability and the OU semigroup:
Stab, f)
Now we know from the spectral expansion of the OU Semigroup action that:
Tof = Σ fll)
aEN
By orthogonality the Hermite polynomials we have:
(f) = Σ aa
f(a) Σ pal pall
Σ
/ at f(a)²
VOEN
actd
aE
If define f(a) v2lala!((a) we finally get:
Stab,(f) Σ fla)
aEN1
D PROOF OF LEMMA
Lemma (Spectral tail bound via stability). Let f E L3((") be square-integraabe function under the standard Gaussian measure 7" Suppose that
Stab, = -EJ(()S(Y)) 2 -)]1[11
for some p E (0,1) and 0 <
1, where (X,Y) are standard Gaussian vectors with correlation p, i.e., Y = px + V1 p 7 for ~ 7 independent of X.
Then f conncettnettt for
log (1 6/E)
log
in the sense that:
ΣΣ f(a)) 2/ < < k2T|a|=k
where f(x) = EnEN f(a)H. (x) the expansion of f in the multivariate physicists' Hermite
polynomial basis.
22

---

Published as a conference paper at ICLR 2026
Proof. Let write the Hermite expansion of as
f(x) Σ (())H.(() aEN
where Hz(z) Ha, (zz), and H₃ is the physicists' Hermite polynomial of degree k. The
Hermite polynomials orthogonal in (y*)) with squared norm |Hall 2| 2!
Parseval's identity in this basis reads: 1 = E(a}} 2l
We define the Hermite level-k weight of f as:
[f((f) = Σ f(a)2. 2
|a|=b
With this notation, Parseval's 's identity is simply = [((())))
The Constein-Uhlenben semigroup Tp acts on the Hermite expansion by
Tpf(2)) -et))>>ple plaa (())
The noise stability can therefore be written
(f) (S,T,f)
Σ paal f(a)² 2l
(w(()
k=
Fix threshold T e N. We can split the sum at degree T::
Sfa
Σ p
(
Σ Eo (
K<T
K2T
Since p E (0, 1), we can upper-bound the terms in the sums by 0² < for T and p² s p² for
k > T. This gives:
St)b S Σ [W())((()++
[f((f) + TT Σ w (((f)
R<T
RST
= (11116
- Σ W(LF
+ + Σ ((((ff(()
k2T
RZT
115112 - p1-
Rearranging the inequality to isolate the tail sum, have:
Σ ((WLF() < 11 Stab,
R2T
Using the assumption that Stab, 2 (1 (1- 1| we can further bound the numerator:
Σ w (TC) < | (1 8) 11
8 11
To ensure this tail sum at most E |f [ suffices to have:
|| 2 S & 11
1
23

---

Published as a conference paper at ICLR 2026
Assuming f is not the zero function, can cancel ll.f and solve for T:
T log p S log (1-9
T log (1 5/E)
log p
This bound is well-defined provided the argument of the logarithm is positive, which holds due to the assumption that d <.
E PROOF OF THEOREM 5.1
Theorem E.1 (Stability of ReLU), Let (X,Y) ~ N 0. | 1) for p (-1,1). Then:
E[ReLU(X))LLU()) 2T (VI p² + p( - arccos p)
Proof. We compute the expectation:
(p) ERRGLU(X))RLLUY) = E|max(0, (0) max(0, Y)].
The joint density function of (X, is
x² 2pry + 112
f(x, v)
exp
2T / p²
2(1 - (2)
Thus, we can compute:
000 Jo xy f(x, v) dr. dr.
T(p)
We perform change to polar coordinates x cos 0, that v. y 2 Then the Jacobian is da dy dr d0 and the exponent becomes:
r sin O r [0, co). 0 E [0, T/2] so
x² - 2pry + v²
2 (cos² 0 - 22p cos A sin 0 + sin" Ø)
(Substitute polar coordinates)
2(1 1²)
2(1 p²)
21 (20))
(sin(28) 2 sin cos 0)
2(1 p2)
So the integral becomes:
1" f
11 - psspinnn)))))) (20))
I(p)
22 cos asn
exp
dr d0
2TT 1
2(1 p²)
1//
cos 1 sin 1
?² psin(20))
3 exp
dr d0.
2T
1²
2(1 - p²)
Now let
a(0) 1 p (2n
2(1 - p2)
Use the substitution - p2 du 2r dr, so:
| 3e a(()r² dr 1 /00 ue -00)j du a(0)2
1
(Standard integral: f0th uc uc-aaui )
2.
24

---

Published as a conference paper at ICLR 2026
So obtain:
4/2
I(p)
2T 2TV11- p² Jo
cos sin 0-
20) d0
n/2
cos sin 0
4T / p² /
a(0) d0.
Now recall that:
1 psin
(1- - p s2nn)
a(0) = 2(1 p²)
= a(0)²
4(1 p²)²
So:
1
4(1 - p
a())
(1 p psin (2)))2
Plugging in:
n/2
4(1 - p²)²
I(p) 4T /T 11 p²
cos f sin 0.
(1 - psspn(2)) d8
(1 p2)3/2
m/2 cos 0 sin A
(1 - psin (2 d8.
n
Now make the substitution u = tan 0, so:
sin()0) 1 + u²
2u
w cos 0 sin 0 d0
du
(1 u
Thus:
cos 0 sin 0
(1 - (1-psin(20) ) d0
2 du
du.
(1 21+
u2)2 (1 p1 Itu
j (u² + 1 2pp))²
This integral can be evaluated by completing the square:
v² 2pu 1 (u (a-p) p)² + +11 (1-1²) 1²)
So, now consider the integral
00
du,
30 ((u p)2 + (1 - 02)
where we define
a:= = v - 22.
We rewrite the numerator as
(u p) + du
1
(u-p)) + a2)
u p-p
du +p x
/
((u p)2 + 0z)2 du
jo
((u p)2 + a²)²
Tn + plz-
25

---

Published as a conference paper at ICLR 2026
Integral 4:
Substitute p. then du dv, and the integration limits become -pto to
11
(02 + a²² dv.
Using the antiderivative
d
2v (0² + 0²²²
du
122 + a²
we have
du
+ +C
(v2 + 22
2(0² + a2)
Evaluating at the limits:
1
11000
1
1
11
lim
2(0² + d²)
+00
2(M² + a²)
2(p² + a²)
v=-p
1
1
0 + 2(p² + a²)
2.1
2'
since 0² + a² 1² -²) =1
Integral 12: Again substitute U = U - P. 5o
100
12
(v² + a2)² dv.
The antiderivative is known:
du
1
U ())
arctan
+C.
(v² + a²)²
(x² 2 + a²)
2a
Evaluating at the limits:
M
M
12 lim
arctan
+00
2a² + a2)
203
a
-p
1
arctan
(p² + a²2)
2a
a
Since
M
M
lim M M
= 0,
and
lim arctan M
(M + a?)
2
get
P
(-²)-
P
I
+
arctan
4
2a²
2a
Using the oddness of arctangent,
arctan( -x) arctan(x))
we rewrite
P
(
I
+
+
aretan
4a
2a
20
Combining Jn and T2:
π
1
P'
I= pla
arctan
4a3
2a
2n
Finally, recall the trigonometric identity:
arctan (
arccos P.
which allows for the final form after multiplying back by the outer constant - 2233/2 1//22/z.
we get:
E (RL = 2 (V1 p² + p(x arccos p))
26

---

Published as a conference paper at ICLR 2026
E. APPROXIMATING WITH QUADRATIC
The expression we gave above for the stability propagation is hard to evaluate. We can approximate with quadratic really well: Lemma (Quadratic Approximation of f())). Define the function
sp)
VI-
p(T arccos p)
2T
for p (-1,1). Then the quadratic approximation of s around p is
1
1
s())
+ 4
0² + o(p²).
2
4x
Proof. First, evaluate the function at zcro:
s(0) 2x
1
(VI- 02 +0. arccos
1
2n
Next, compute the first derivative:
s'(p)
P arccos p P. arccos
d
2n
dp
Using d arccos p
this simplifies to
1-
s'(p)
(π
arccos p).
2T
Evaluating at p = 0:
1
8' 8()))
π
Now compute the second derivative:
1 d (x arccos p) 2x /1 D2
1
(p) 2x dp
and at = O:
1
s²(
Therefore, the quadratic approximation at 1=
1
s(p) ~ x(0) + (0)
p
2m
4T
0.5
-0.5
0.5
Figure 9: Approximating the stability of MLP with quadratic
27

---

Published as a conference paper at ICLR 2026
F PROOF OF THEOREM 5.2
F.1 PRELIMINARIES
We will make use the following mathematical tools in the proofs below:
Lemma 8 (Laurant-Massart Concentration Bounds for Xa Random Variables). Let Z be a chi-
squared random variable with d degrees of freedom. Then:
Pr 2 2 2v + 2u) Se-
Lemma (Concentration of Gaussian Random Variable). Let ~ N(O, o Then:
Pr|z 2 al S =
Proposition (Weighted y Un TR. It is true that:
Inequality). Let a1, On be non-negative constants and
1/2
1/2
Σ
Σ
Σ
03
aj
Proof. First, by the triangle incquality for the fo-norm and the fact that aj 2 0, we have:
2
2
n
n
Σ
زززر
|ay l2
Σ
1l2
Σ
Va
aj
11
j=1
j=1
Next, we apply the Bauchy-Schward inequality to the term on the right-hand side, to obtain:
2
n
n
Σ
Vaz (Vasllusla)) s
Σ
aj
a
11
j=1
which finalizes the proof.
F.2 PROVING THE THEOREM
Theorem F.1 (Stability of Attention Layer, WoW Ta). Let N(0 Inxa) and px + Z 2-1- p² for Z ~ N0, Iaxa) 1 X. Let f Baxd - Raxd where f(X) = oxXx)) XWv, Wv E Rdxx and a is the row-softmax function. Thenfor all E [n] and j E [d] we have that:
lim (f))))))l)))) Pl((WV).a11 d=
Proof. Let Ii...., Fn RRd be the rows of X. We first show that = oXXT) converges to I, in probability as d + with exponential tails in d.
First, fix some [n] and let o² = ||::1|2- that:
We know that 0² ~ xa and so by Lemma g we have
23d
4ld
Pr
22 <
2c -/66 d/61
32
32
Let A be the likely event defined above. Let us condition on happening.
Now, condition on and let V, = (Ii .z)) - o². Since a; and so by Lemma we have:
Ty we have that V, ~ NN0002,00 ar
Pr V 2
σ²
53
= Pr
(1i) 2
xs
Se z
c -02118 <e 23dd/26
2
2
28

---

Published as a conference paper at ICLR 2026
By union bound, we have that:
0²
Pr Ejfi s.t. V2 | A
I (n 1)e 280
Removing the conditioning on x; and A, get that:
Pr Ej V 2 0²
Pr[-A] (n 1)e 250 d S 2e 2e-6/61 + (n
1)e
258
Therefore, with probability at least 1 - E ((d) it holds that V, s o² S
23d
Let call this
event and condition on it.
Now, we focus on the i-th row of W. For j # i, we have:
exp(V,)
(-23) d (Todd)
23
Wij 1 Exg exp(VA) S exp((V) S exp
Conditioning on Ha, we have proven that:
(- (-64 64
41d
0 (ne- (ne- B(a)
Pr
max Wi, > exp
Jfi
In other words, Wij 3 0 as d with an exponential tail.
Similarly, when i = j. we have that:
1
Wa I+ Enji exp(Vi)
I 1 (n 1) exp( 64
23
d)
meaning that W₂
Now that we have shown that W P,, I, as d - let consider the product W -WW. We want to treat this product as e,A', so we show that the error goes to 0:
W. XA 6 XWv Σ
Wiya,WW Ex
We have by the weighted Cauchy-Schwatz inequality that:
> El, Wr
1x <
Σ Wii
< (n 1)
Σ
E [Ma
|| Wv
j
where Ma = maxj Wi [0, 1J. We know that Ma 0 in probability, so by the dominated
convergence theorem we get that E[M,] + 0 as d + oc. Also, E
Σ
|| Wy 113
is a constant
#
with respect to d. Applying Cauchy- Schwartz again, we get:
E
(n 1) E E[M
E
Σ Ir, Wv
d=
Similarly, we can also show that Ey| [|]] + as - Dc.
With that, if we fix j E [d]. let xx,j f(X) )is (x; We have:
E x < B|| Ex 2 0
Elst.gl < By +
29

---

Published as a conference paper at ICLR 2026
We can therefore claim that:
[X)(s f(Y) (l()W)) +w
To see this, observe that:
(x)) (x, (x:Wr) (x:Wv)) exxs (w:Wi)) eyRi + Fexi eyúa Each term has expectation o(1):
E((x;Wv), ey - [l)) VEle,,,) [I(W]))sll o(1) = o(1)
W(( Ex = o(1) symmetrically.
Elxxas Eyasl & V V o(1) o(1) = o(1).
Finally, obtain the final calculation:
lim El5(X), = El(x;Wi)) (v:l o(1) d
(l (Wv):a) (Wi, (Wv) + o(1)
((Wv)a, (Wv)es Bluis v.d1 o(1)
((W))2
(By definition of (X, Y))
p. [(VV):=111++(
as claimed.
G NOISE STABILITY PROPAGATION IN THE UNSTRUCTURED CASE
In this section we prove the following theorem:
Theorem G.1. The noise stability of an attention layer in the unstructured case is:
lim
= = (I(Wv) o(1), with:
d0000
sp) = n [
/
1 (z, y)- fpp drdy
00
where 4 is the joint CDF of a bivariate normal distribution with correlation e and fo is the respec- tive PDF
Proof. We first prove that each row of the softmax matrix is permutation:
Lemma 10. Let Pa. = x W xy Then:
o (P (P... ... P...) 2g ex, where i: arg Pa.j
jE(n)
similar statement also holds for Y.
Proof. The proof of this statement is very similar to argument from Theorem 5.2
We have that:
Σ abbE(d]
Fi.a Wa j.bb
Conditioning on X, Ptş normal with mean and variance ||0:1| ||, 12 ((d²) (Lemma
So write P., = d. Yi where Yig ~ ~N(00 1) conditioned on X. Let he arg max P., JE(ni
30

---

Published as a conference paper at ICLR 2026
arg max y.j. Let A Pa. jE(n]
Pat d(Yik- Yuk). We know that for & * k- the quantity Ak
converges oo in probability as d 00,000
Pu
Ax
30
Σ e 1+ Lofk e e-A,
as d + co. For k = k* however, have:
Pur
1 1 Lefk e 21
Σ. e
This establishes the lemma.
Now, let Lij- v W y and let Pr(arg max p.. arg max Ligh. We know that Ps and Lij have correlation p² as:
jE(n)
sE(m]
]|
E
Σ
Ti,a Wa, WabDs))
Σ
Vi.a Wa Wa.bU
Σ
Wa,d dp²
p,
a.bEld
a.b.cdd(Mm]
The joint conditional distribution of both p.,, and Li,j is standard bivariate normal with correlation p²:
x² - 22-2ppxy xy + y²
(2, v)
exp
2T / pr
2(1 - 14)
It's CDF is:
2 (2,
8 v) du
00 -20
Thus, we can calculate:
20
8 = s(p)
s() =
D (z, (x, (0 (p, y)dady
The remainder of the proof concludes as in the main text by considering the events of the maxima matching not, taking the expectation and de-conditioning.
H FULL NOISE STABILITY DAMPENING IN MULTI-LAYER TRANSFORMERS
Let W K = / and (W (. = I 1. Ignoring distributional shifts, we can combine Theo- rem 5.1 and Theorem 5.2 to get the following recurrence:
(V
(11-rpp_ + P arccos(? PL-1)))
(8)
2m
Substituting the linear approximation of Equation (4) and setting pi = 1. we obtain:
72
2
L-1
2T 4
+ PL.
PL
(4 2)
<T 72)
This suggests that for y < the noise stability propagation through a multi-layer transformer con- verges to 5(1-92T
However, empirical verification suggests this is not the case. Figure 10 shows that for y < we observe full dampening, while for 1 weak dampening remains. This is due to the fact that for * 1 the output of the Transformer layer decreases exponentially with 7, which also affects the noise stability.
31

---

Published as a conference paper at ICLR 2026
Noise Stability vs. Number of Layers (L): Transformer
Staclity IoWo
Simulated Stubiitt for IlWo 0.75
0.5
Simulatca Stablity for BOWo..lla
Estimated Stability EIt (Kh,, -f (.(9)..))
02-
0.1
0.0
10
11
12 33 14
Number Layers
Figure 10: Transformer exhibits full dampening in the multi-layer setting.
I NOISE STABILITY INTERVAL. PROPAGATION
In this section we present some analytical results on the propagation of noise stability intervals. These are lemmata that provide more flexibility towards the task of analyzing noise stability through- out a deep Transformer, though distributional assumptions still need to be made. Such assumptions can be quite detrimental
validity of the analysis, as we have already seen.
1.1 MLP STABILITY UNDER BONAMI-BECKNER GAUSSIANS
Lemma (Propagation of Stability in MLP Layer). Let E R be a random variable where Xij N(ug, 0sz) (not necessarily independent of each other). Consider a random variable Y generated by a scaled Konami-Becking noise process:
aaxi; N(uj, iid, otherwise
with probability Psj E [0.1]
Yú
where ar Let &b IR R be the element-wise ReLU function. Then have that:
El((X.) o(Yi))) = PijaEi + +(1 -Pt))
for all (i,j) E [n]] [d], where:
Mij
Peg
+ pij
Oij
Jis
E2 (ni + o )
Uij
+
Jsyliig
Mij
Oij
and f, Tb is the PDF and CDF of the standard (0, 1) Gaussian distribution.
Proof. Let us fix some j and use the law of total expectation:
Elo(Xx))o((,)) P Più (()))x(XX))
32

---

Published as a conference paper at ICLR 2026
For the second moment, we have:
(ompe)
1
E
[lo(X1))
2
20
2
1
(uT + +1is) (1s)) 6- du
(u 2-my
2T
-uylow
dis
0j
u'e-" du +
20ssUig
2 |mmglas ue ue- 2 du +
poc
é du
2T
-Huy/ow
22.
-uss/os
For the third summand in this expression, we have that:
- du = n/g -1 Φ
Mij
-Mig
mii $
Mig
V2T
-pylos
Fij
asj
For the second summand, is:
20
20igHis
H4
ue du
2
/2TT
Juus/ar
-Miplous
Ozj
For the first summand, we can use integration by parts to get:
00
uee du
(1-me- 1==22=a"22)
+ f -
pc
e=
e
du
22T Huylas
2m
-us//as
o pij
Mij
Pbi
f
+
+
Oij
Oij
Ozj
Combining, we get an expression for the second moment:
E Blo(X.)))1 o²j -Mi
f
Mig
o
+
Mij
2pijoris
Hi
+ p $
Uii
Oij
Oij
Oij
0ij
asj
(o + 22
[ii
+ MigOrif
Hj
Tis
Now for the mean, we have:
8
(e-Pis?
1
[lo(x+))
2e
2
1
(uoij + Hig}e du
(u 3-
V2T
-pis/ay
Splitting up this sum, we have:
Eg = E[o(x))]]
Oij
000
uc 2 du +
Pij
600
e - - du
22T
Jui//or
22T
|lij
Mij
+ pij
T
Osj
Osj
This concludes the proof.
I.2 STABILITY PROPAGATION THRROGHH A SINGLE ATTENTION LAYER
For convenience, the following result defines a stability matrix, which captures the correlation of each input token with the other. Assumming all the entries in that matrix bounded in some interval, we analyze the noise stability propagation through an attention layer. We also assume some structural properties of the weight matrices to allow for stability propagation in our proof:
33

---

Published as a conference paper at ICLR 2026
Lemma 12 (Stability Propagation through an Attention Layer). Let X,Y have stability matrix (Chee = TE[X E Rndxnd with < I Cke S Pr for all k, E [n], E [d], attention layer with matrices
and suppose || loo | S B with probability 1. Consider WK, Wo, Wv E RRddd and denote, for two vectors 3. y E TR. the following quantities:
st (z, = Σ max(0,,2.3) AgE|d
S (z = Σ
min 00,
Suppose that Wv is such that pest (w,, wyy + p. (w,, wy) for all j. E [d, where wy (Wv).s is the j-th column of Wv. Hence, let:
Re inf (peSt (wj, wy) + P. (wy, wy JJ'E(d]
= sup (p. (wij wy
i.'E(d]
Now, let A(X) nxd be the output of the attention layer. Then we have that:
Re
< S ((A( A(Y)))) R. E
(9)
for all i, E [n]. E [d], where E == exp( 4 B [ |0x0 Welloo))))
Proof. Let 5x = -XW and VX = XWv. We have that A(X);s = (S ViS) so:
B(A(X)s A(Y) E (S V.J (8) VAJ
]
St v.
Σ
S
Vilgg
S Spikk V Vidg
k.
Now let q,¹¹ = X1: We R = X4. WK. We have:
explai* kkr)
S
Σ explax ks)
Therefore:
B(A(X) A(Y))))
exp(1x kk.)
explox(i)
E
V
VXVV V
k
exp(gx ks)
Σ kx
exp(X,,, W
exp(Ys,
E
(w
k
Σ
exp(X,, W xX.)
Σ
We will bound the softmax terms using the norms of Wq. Wx x and X. For any 81, 82 we have:
IX Wo < dl B² Wollo |WKlo
And this implies that:
exp(-2dt 22f B² W Wolls) exp(X,, WoW WoWXX
1 exp(2dd exp(2dBB?llw)) B²l [WQll>)
<
n
34

---

Published as a conference paper at ICLR 2026
We now analyze the isolated terms
[[ (e, vW.).)
We can use our prior information on the stability matrix C. Fix some k. and we have that:
Σ
E (e Yee (Wv))))
6,
= [(wr)s(Ww)er E [Kee
6,
E [Re, Rrd
(as LE [XkeYwel E (resT,))
Overall, if we let = exp(Ad B² [K |
[W)lma)) we get that:
Rs
0< S ((())))))((())))) R
J NOISE STABILITY REGULARIZATION EXPERIMENTS
J.1 ARCHITECTURE LAYOUT, HYPERPARAMETERS AND TRAINING DETAILS
We present details of our architecture, training and hyperparameters. These can also be found in our codebase. Each Transformer layer multi-head self-attention followed by a position-wise feed- forward network, with residual connections around both sublayers. We use sinusoidal positional encodings, as well as binary attention masking M. We also use dropout, applied attention weights, attention output, and FFN hidden layer with rate The activation function we use for our FFN is ReLU. To produce a classification label we use mean pooling over the label dimension.
For initialization, linear layers are initialized with (/0 0.02²) and zero biases. Every other learn- able parameter initialized via Xavier initialization. Positional encodings are fixed and not train- able.
For training, we use AdamW with learning rate 7) and l weight decay regularization X. Our loss is the cross entropy loss. We use learning rate scheduler that reduces on validation loss plateau. The patience and factor parameters of the scheduler are hyperparameters we set. Our codebase also provides support for multi-GPU and distributed training, though our models were too small to benefit from such augmentations.
Table 2: Model Hyperparameters and Training Configuration for MODULAR ADDITION
Name
Symbol
Default
Embedding dimension
dmodel
128
Transformer layers
L
2
Attention heads
H
2
Max seq length (PE)
max max_length
512
Vocab size
M
K + 5
Num classes
K 113
Dropout rate
p
0.1
Batch size
256
Epochs
T
7000
Learning rate
n
0.001
Weight decay
X
0.001
Label smoothing
0.0
Scheduler patience
10 epochs
Scheduler factor
0.8
Train samples
2000
Continued on next page
35

---

Published as a conference paper at ICLR 2026
Validation Loss and Noise Stabilities Over Epochs
1.0
Norae SLAdly Boo onl)
03
Velidation Loss
0.4
0.2
Epoch
Figure 11: Evolution of noise stability and validation loss for noisy sparse parity. Stability (green)
begins to decrease before the validation loss (blue) drops, acting as a leading indicator for general- ization.
Table continued from previous page
Name
Symbol
Default
Val/Test samples
200 200
J.2 EVOLUTION OF NOISE STABILITY DURING TRAINING.
We also analyze how noise stability (with p 1/2) evolves during training without regularization, focusing on the noisy sparse parity (NSP) task first because its target function has a known, low noise stability d O'Donnell 2021). We make three key observations:
A randomly initialized Transformer exhibits high noise stability, which decreases through- out training to converge toward the low stability level of the target function (Figure II].
Noise stability is leading indicator for generalization. Figure II shows that stability begins to decrease well before the sharp drop validation loss, signaling that internal model adjustments precede performance improvements.
Stability can serve as a secondary metric for model selection. Among models with similar validation loss, the one whose stability best aligns with the thcoretical properties of the target function may be preferable.
J.3 EXPERIMENTS ON LANGUAGE GENERATION
We also tested noise stability regularization on a language generation task. We trained 4 layer transformer model with dmwseel = 30 and H = 6 the next-token-predice task. The dataset we used was WikiText-2-vl (Small), with sequence length N 20, a vocabulary size of 500, batch size of 200 and 1000 training examples We trained our transformer without noise stability regulariza- tion, with weight decay À = 0.02 and with numerous settings of (p.7). We tracked noise stability, validation loss and validation accuracy throughout training.
We first observe that even in this non-synthetic task noise stability regularization offers deep benefits for training. As shown in Figure 12 the model climbs 70% accuracy within 1000 iterations in the grokking phase. On the other hand, takes the non-regularized model 4000 iterations to do
*ur setup also be found in full in codebase.
36

---

Published as a conference paper at ICLR 2026
so, amounting to 75% speedup. The situation is similar for the validation losses: the regularized models exhibit more stability in training, while the non-regularized training fluctuates validation loss. Noise stability regularization also causes the validation loss to decrease much faster and earlier than without it, while grokking.
Examining the noise stability at 1/2 for the regularized and non-regularized settings (Figure 14) we can see a fundamental difference between the models. The non-regularized model becomes less and less stable, which could explain its instability. Regularized models stay stable as the training dynamics force the model to improve while remaining robust. Understanding the benefits of this regularization better is definitely an interesting direction for future work.
Finally, we can see that noise stability regularization offers benefits for variety of different settings of the hyperparameters (p. 7). Ultimately however, the best performance found via a thorough hyperparameter sweep, Figure 13 shows.
Validation Accuracies: WikiText NTK
Accuracy (96)
y-000pp03
-.4.
440 p-0.
20.15,pa p=0.3
Epoch
Figure 12: Accuracy Comparison on Next-Token-Prediction
Training and Validation Losses: WikiText NTK
yod 0, p=0 Trah
yo0 p=0.5
y-d.O, p-1 Val
1.15 pr0.1 Tiwin
pr02 Tran
pro.
pH0
y-0.A Train
y-0.55
LOSS
Epoch
Figure 13: Training and Validation Losses for Next-Token Prediction
37

---

Published as a conference paper at ICLR 2026
Noise Stability vs Epoch
3.50
3.25
3.00
Noise Stability
2.75
2.50
ye0.15, (pe..)
100
2002
300
400
500
500
700
800
Epoch
Figure 14: Noise Stability Comparison: Regularization vs Non-Regularization. We see that the non- regularized model tends become more unstable, while regularization maintains stability.
38


Reviews:
[{"review_id": "dd8PcpnINA", "reviewer": "Reviewer_vNH8", "summary": "The paper studies noise stability as a quantitative lens on simplicity bias and spectral concentration in Transformers. It formalizes Gaussian noise stability using the Ornstein Uhlenbeck semigroup and Hermite expansions, then proves new results for core components: a closed form for the stability of a ReLU MLP layer, and stability characterizations for attention under identity, low rank, and unstructured regimes. It extends the analysis to depth via a recurrence for MLPs that predicts weak dampening and introduces stability interval propagation to bound stability through deep Transformers. Empirically, the authors propose a noise stability regularizer that is differentiable and data dependent. On grokking style tasks such as modular addition and noisy k sparse parity, the regularizer speeds up emergence by about thirty five percent. They also compare stability based spectral tail bounds with sensitivity based bounds on several pretrained models, finding better agreement with measured concentration when using stability.", "questions": "1. Role of  $s(\\rho)$ beyond initialization. The unstructured attention result introduces a permutation consistency factor $s(\\rho)$ that is derived under random weights. How does $s(\\rho)$ evolve during training when attention becomes more structured. Can the authors estimate or bound $s(\\rho)$ on trained layers.\n\n2. Tightness of interval bounds. The interval propagation method provides upper and lower bounds. On realistic trained networks, how tight are these intervals per layer. A plot of measured stability with interval envelopes would help assess utility.\n\n3. The regularizer allows both orientations. Have you tried S=0 to discourage stability as a stress test on overfitting or memorization. Does decreasing stability hurt grokking or robustness.\n\n4. From synthetic to text. Can you run a small language modeling experiment that logs stability, spectral tail, and validation perplexity every few thousand steps. Even a 100M parameter model on WikiText or FineWeb subsets would help connect theory to practice.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Stability analysis of attention with identity, low rank, and random regimes is new and insightful. The appearance of a permutation consistency factor $s(\\rho)$ in the unstructured case is a neat and interpretable result. Theoretical results are derived from a solid OU semigroup foundation with explicit ReLU formulas and careful discussion of approximations. Interval propagation addresses the realistic case where exact recurrences are brittle. The regularizer is minimal and differentiable, with a clear sampling scheme for $Y$. The experiments control for seeds and hyperparameters and demonstrate a consistent thirty five percent speedup on two canonical grokking tasks. Finally, stability based tail bounds better match measured concentration across common pretrained models than sensitivity based bounds. This supports the relevance of the stability viewpoint for understanding learned spectra in practice.", "weaknesses": "- Assumption gaps for deep Transformers. The multi layer recurrence assumes fixed distributions and idealized conditions. The authors later observe full dampening in practice when $\\gamma < 1$, which highlights a gap between proxy analysis and actual behavior. More systematic experiments that probe when interval bounds are tight would strengthen claims.\n\n- Limited task scope. The regularizer is only evaluated on synthetic algorithmic tasks. Evidence that it helps on language modeling or fine tuning tasks is indirect through tail bound comparisons. A small scale LM experiment that tracks validation perplexity, stability, and tail mass together would improve external validity.\n\n- Sensitivity to $\\rho$ and $\\gamma$. The analysis and the regularizer hinge on these hyperparameters. The paper sets them per task but does not map performance and stability as a function of $\\rho$ and the regularization strength $\\gamma$ across a grid. This makes practical tuning less clear.", "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": "l3NfZjAPnK", "reviewer": "Reviewer_DkBU", "summary": "This paper proposes noise stability as an alternative to average sensitivity for explaining spectral concentration in Transformers. The authors argue that average sensitivity has theoretical limitations in extending to real-valued domains and empirically fails to capture the \"junta-like\" behavior observed in modern LLMs. They provide theoretical analysis of noise stability for single-layer attention and ReLU MLP components, develop a stability interval propagation technique for multi-layer networks, and introduce a noise stability regularization method that accelerates grokking by approximately 35% on algorithmic tasks.", "questions": "1. What is the intended use case for S=0 (encouraging instability)? Are there regimes (e.g., parity-like targets) where anti-stability regularization could help avoid degenerate shortcuts? The paper only evaluates S=1.\n\n2. Appendix J.2 analyzes stability evolution *without* regularization. Showing the same trajectories *with* the regularizer would more directly support the “catalyst” claim and illuminate the mechanism.\n\n3. What is the computational cost of noise stability measure and training with noise stability regularization relative to average sensitivity or other alternatives? (E.g., the load of extra forward pass on a correlated Y and scaling properties) A brief cost analysis would be useful for practice, but I didn’t see this discussed.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "**Substantive theoretical contributions.**\nThe paper gives an exact, closed-form expression for the noise stability of a ReLU layer (Theorem 5.1 / Appendix E) and analyzes single-layer attention under different structural regimes (identity/low-rank and unstructured; Theorems 5.2 and 5.3). It also studies multi-layer propagation, showing weak dampening for deep ReLU MLPs and contrasting behavior for Transformers.\n\n**Thorough and self-contained presentation.**\nThe work includes background on Boolean function analysis and the Ornstein–Uhlenbeck semigroup/Hermite expansion, connecting noise stability to spectral concentration and supporting later results with detailed appendices.", "weaknesses": "**Limited comparison with prior works.**\nSection 3 argues that geometric influence (GI) is a more robust alternative to average sensitivity in continuous domains, but then pivots to advocating noise stability as the primary metric without a direct empirical or theoretical head-to-head comparison (e.g., GI vs. noise stability on the same tasks/models). Related work cites Li & Mossel (2025) for noise sensitivity propagation but does not clearly mention what is new beyond the inspiration or compare results. Adding targeted comparisons would sharpen the contribution.\n\n**The regularizer mechanism requires more intuition and ablation.**\nThe regularizer (Def. 6.1) is clearly specified, but the paper offers limited intuition for why this particular form steers learning toward “simpler” functions beyond the variance-level view. Hyperparameter sweeps for $\\gamma$ and $\\rho$ are also missing.\n\n**Narrow experimental scope.**\nAll experiments are on synthetic algorithmic tasks (noisy k-sparse parity, modular addition), and the grokking speed-up is reported only there. It’s unclear whether the regularizer helps on realistic NLP or larger-scale settings.\n\n**Minor presentation issues.**\nThere’s a typo (“doamins”) in Section 4, and some figures/tables could use more consistent visual styling and larger font sizes for readability.", "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": "nrv2MFIKTh", "reviewer": "Reviewer_QkLJ", "summary": "The paper proposes noise stability as a measure of simplicity of Transformers.\nUnlike average sensitivity, which captures local perturbations, noise stability measures robustness to correlated, global perturbations across the input sequence. Building on this idea, the authors introduce a noise-stability regularizer and show that it improves training and generalization on algorithmic tasks exhibiting grokking dynamics.", "questions": "1. How do the positional-bias and sensitivity patterns in Figure 1 concretely demonstrate the limitation of average sensitivity, and in what way does noise stability resolve it?\n2. Can the authors include a controlled comparison among (i) average-sensitivity regularization, (ii) noise-stability regularization, (iii) weight decay, and (iv) no regularization, to demonstrate that noise stability improves training dynamics beyond heuristic smoothing?\n\n3. What is the additional computational overhead of computing and back-propagating the noise-stability term compared to average sensitivity?\n\n4. In Appendix J.2, the authors show that noise stability correlates with validation loss. Can they further test whether noise stability decreases before validation loss changes, serving as an early indicator of grokking as observed in [3].\n\n[1] Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking Beyond Algorithmic Data. ICLR 2023.\n\n[2] Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S. Du, Jason D. Lee, and Wei Hu. Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking. ICLR 2024.\n\n[3] Bhavya Vasudeva, Deqing Fu, Tianyi Zhou, Elliott Kau, Youqi Huang, and Vatsal Sharan. Transformers Learn Low Sensitivity Functions: Investigations and Implications. ICLR 2025.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The mathematical formulation of noise stability through the Ornstein–Uhlenbeck operator is well presented. The concept of analyzing robustness under global, correlated perturbations rather than local ones is intuitive and sound.\n2. The introduction of a differentiable noise-stability regularizer is a concrete and practical contribution.\n3. The choice of algorithmic tasks is appropriate for studying grokking phenomena. The experiments show correlations between noise stability, validation loss, and generalization behavior.", "weaknesses": "1. The paper motivates the use of noise stability over average sensitivity based on the observed “junta-like” dependence in Transformers, arguing that noise stability captures global correlated perturbations rather than single-coordinate ones. However, the connection between robustness to global perturbations and dependence on a few input variables is not clearly established. It remains unclear why capturing correlated noise should correspond to fewer influential tokens.\n2. The authors claim that noise stability provides tighter spectral tail bounds than average sensitivity, but no theoretical comparison or proof is given. The empirical evidence in Table 1 lacks sufficient methodological detail for verification—for example, the meaning of “percentage of weight in degrees ≥ 15” is not explained.\n3. The claim that “average sensitivity is too loose to explain the behavior of modern Transformers” is not convincingly supported by experiments. The paper would be much stronger with a direct comparison of average sensitivity and noise stability as regularizers.\n4. The paper also does not compare against other regularizers known to affect grokking dynamics (e.g., weight decay, $L_1$ norm [1, 2]). The reported grokking acceleration appears to be based on a limited number of runs and lacks confidence intervals.\n5. The proposed regularizer perturbs discrete token inputs via random replacements, whereas the theoretical formulation of noise stability is defined over Gaussian perturbations. The paper would benefit from clarifying how these two settings relate. \n6. Section 5’s analysis of noise stability in Transformers does not meaningfully yield actionable insights for Transformer training or regularizer design. \n\nMinor: \n- The term “junta-like” appears in the abstract before being defined.\n- In Line 397, the authors mention “three key properties” but list only two.", "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": "jNvzX6if9B", "reviewer": "Reviewer_6UoK", "summary": "The authors propose a method to investigate the simplicity bias in neural networks. Simplicity bias is the tendency of neural networks of learning simple functions to fit the data: as simpler functions are (i) more interpretable and (ii) more robust, assessing simplicity is important to study model reliability. \nA current trend of research frames the problem of simplicity in terms of spectral decomposition. The underlying idea is that simple functions should show spectral concentration around low-order polynomials. Drawing from Boolean functional analysis, it is possible to study the impact of a single token by \"swapping\" its sign and contrasting the average output change with the total variance. The authors criticise the approach. They claim that neural networks employ functions defined on real numbers -and not on binary coordinates- and they criticise the idea of considering one token at the time as this fails to capture their collective behaviour: (i) the \"junta\"-like concentration, i.e., the sparsity of the influential tokens; (ii) the positional bias; (iii) the non-zero influence of all tokens. \nTo address their concern, they propose to extend the concept of noise stability from Boolean analysis - which is the study of the effect of correlated noise applied to all inputs simultaneously. First, they define the concept of stability as the correlation between the original input and a partially perturbed input - corrupted with Gaussian noise. Then, they express this measure employing the spectral decomposition of the function in terms of its Hermite-Fourier coefficients. Finally, they study (i) a simplified RELU MLP and (ii) a simplified Transformer, finding that the first dampens stability, whereas the latter maintains it. \nTo test the impact of penalising with respect to spectral concentration, they propose a stability regularisation: they train transformers on two artificial tasks that exhibit grokking, i.e., a sharp decrease in the loss function value after an initial over-fitting phase, noting that regularisation cuts the number of iterations requires to attain it.", "questions": "1. Line 107/108. Are you sure the dimensions of this equation are correct? \n    2. Table 1. What would happen if we instead considered the geometric influence here? \n    3. Line 279. Maybe a typo: you see \"(see Section 5.2)\" in Section 5.2. \n    4. Figure 2. This is a bit hard to read: the labels are very small. \n    5. Section 5. Can you comment on the results? Why are the different between Transformers and MLPs? Or how does it show in practice?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality:  The authors propose an approach to the study of neural networks that is well-grounded in theory and takes into account the collective behaviour of the tokens in LLMs. They derive analytical solutions for typified examples and introduce a regularisation mechanism. \n Quality:   the derivations are sound; experiments accompany all the claims; a numerical example works as a case study. \n  Clarity: The work is clear.  well-organised and easy to read. The Appendix is organised well and it is easy to follow the derivation of the derivation of the Ornstein-Uhlenbeck Operator. Some relevant entries in the bibliography are incomplete (e.g., Nathan Keller. Elchanan Mossel. Arnab Sen. \"Geometric influences.\" Ann. Probab. 40 (3) 1135 - 1166, May 2012. https://doi.org/10.1214/11-AOP643). The labels of the plots are too tiny and difficult to read. \n Significance The author propose a new method to investigate the simplicity bias that considers all tokens at the same time. Given the importance of collective token dynamics in Transformer, this work goes in the right direction.", "weaknesses": "The model criticizes the use of influence as a proxy of sensitivity and propose noise stability as a new method. One of the reasons of their criticism is that the upper bound of the former method are too loose with respect to the latter when estimating spectral concentration. As evidence, they report the values for four different models in Table 1. However, it is not entirely clear if this is the only advantage of their new method. In other words, it is not entirely clear if using influence or, better still, geometric influence it is not possible to gain the insights they obtain using noise stability. The authors obtain interesting results for MLPs and Transformers, but it is not  clear why the results differ and what might be the reason behind this. Similarly, while the example of Section 6 shows that there seems to a be link between simplicity bias and spectral concentration, it is not immediately clear if noise stability captures it better than influence, if it might be possible to use it too, and if there is a way to experimentally detect the issues they already showcase in Figure 1 in this training setting.", "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)