Generalization Bound for a Shallow Transformer Trained Using Gradient Descent
Authors: Brian Mwigo
Anirban Dasgupta
Venue: Accepted by TMLR ()

Abstract: In this work, we establish a norm-based generalization bound for a shallow Transformer model trained via gradient descent under the bounded-drift (lazy training) regime, where model parameters remain close to their initialization throughout training. Our analysis proceeds in three stages: (a) we formally define a hypothesis class of Transformer models constrained to remain within a small neighborhood of their initialization; (b) we derive an upper bound on the Rademacher complexity of this class, quantifying its effective capacity; and (c) we establish an upper bound on the empirical loss achieved by gradient descent under suitable assumptions on model width, learning rate, and data structure. Combining these results, we obtain a high-probability bound on the true loss that decays sublinearly with the number of training samples $N$ and depends explicitly on model and data parameters. The resulting bound demonstrates that, in the lazy regime, wide and shallow Transformers generalize similarly to their linearized (NTK) counterparts. Empirical evaluations on both text and image datasets support the theoretical findings.

Published in Transactions on Machine Learning Research (01/2026)
Generalization Bound for a Shallow Transformer Trained Using Gradient Descent
Brian Mwigo ITT Gandhinagar
brian. metan.mwigooiiie wiigingggittgtt........
Anirban Dasgupta IIT Gandhinagar
anérbandg@hitgnn.ccm
Reviewed on OpenReview: https: // openreview. net/ forum? id= t3qUeMOTS2
Abstract
In this work, we establish a norm-based generalization bound for a shallow Transformer model trained via gradient descent under the bounded-drift (lazy training) regime, where model parameters remain close to their initialization throughout training. Our analysis proceeds in three stages: (a) formally define hypothesis class of Transformer models constrained to remain within small neighborhood of their initialization; (b) derive upper bound on the Rademacher complexity of this class, quantifying its effective capacity; and (c) we establish an upper bound on the empirical loss achieved by gradient descent under suitable assumptions on model width, learning rate, and data structure. Combining these results, we obtain a high-probability bound on the true loss that decays sublinearly with the number of training samples N and depends explicitly on model and data parameters. The resulting bound demonstrates that, in the lazy regime, wide and shallow Transformers generalize similarly to their linearized (NTK) counterparts. Empirical evaluations on both text and image datasets support the theoretical findings.
1 Introduction
Deep learning models have achieved remarkable success on language and vision lasks that were once con- sidered intractable for neural networks. Transformer architectures, in particular, have played a central role in this progress, underpinning largo-scalc language models such as GPT-4 (AAcham cf all, 2023), LLaMA [Louvron el all 2023), and Gemini (Team all 2023). Beyond language, Vision Transformers (Dosovilskiy allll1 2020) have demonstrated competitive or superior performance on image classification and generation tasks. These empirical successes have fueled optimism about the capabilities of modern neural networks and, in some cases, speculation about the emergence of artificial general intelligence. Despite this progress, the theoretical understanding of Transformer models particularly their generalization behavior remains limited.
A central challenge in the theory of deep learning is to explain why modern neural network models gener- alize well. For Transformers, several recent works have developed generalization bounds by controlling the generalization gap, that is, the difference between the true loss and the empirical loss, LD( - Ls(s), over suitable hypothesis class F (Edelman all, 2021; Trauger & Tewarr 2022: En Cu 2022 ). Such bounds are valuable for understanding capacity control and urchitectural inductive biases, but they require training to complete in order to evaluate the empirical loss Ls(f) before any conclusions about the true loss can be drawn.
An alternative line of work directly upper-bounds the true loss Lr( by combining statistical complexity control with explicit analysis of optimization dynamics. This approach yields guarantees on generalization without requiring post hoc evaluation of the empirical loss. For fully connected networks, Arora of all (2019) established such a bound for overparameterized two-layer ReLU networks trained by gradient descent, lever- aging convergence results in the neural tangent kernel (NTK) regime. This framework also applied to

---

Published in Transactions on Machine Learning Research (01/2026)
deep fully connected networks using Neural Tangent Random Features in Cao &Z Gni (2 However, this optimization-aware true-loss perspective is less explored for Transformer architectures.
In this work, we take a step toward closing this gap by deriving a generalization bound for a class of Trans- former models whose parameters remain close to their initialization throughout training. This bounded-drit assumption corresponds to the lazy training regime, which is commonly associated with highly overparame- terized networks trained using small learning rates. Under this regime, training dynamics are approximately linearized around initialization, enabling tractable analysis while still capturing important aspects of over- parameterized learning.
Our analysis proceeds in three steps. First, we formally define a hypothesis class of shallow Transformer models whose parameters remain within a fixed-radius neighborhood of their initialization. Second, we derive an upper bound on the Rademacher complexity of this class using covering- number arguments, thereby controlling its statistical capacity. Third, we leverage recent global convergence theorem for shallow Transformers (WWn elll 202a1) obtain an explicit upper bound on the empirical loss for all models in this class. Combining these ingredients yields high- probbabibtt upper bound on the true loss.
Specifically, our main contribution is an optimization-aune aware generalization bound that directly upper-bounds the true loss for a class of shallow Transformer models trained in the bounded-drey drift (lazy) regime. The resulting bound decreases sublinearly with the number of training samples and holds uniformly over model dimension dm, under suitable overpavametertattion and data assumptions.
2 Related Work
Generalization bounds for neural networks. Generalization in neural networks has been extensively studied using classical complexity measures such as VC dimension, Rademacher complexity, and covering numbers. Early norm- and margin-based bounds for fully connected net were developed in BBrrlee al (2017); Neyshabur et al. (2015 20IZ), with later refinements incorporating spectral norms, path norms, and data-dependent quantities (Neyshabur et all 2018 PPras of all, 2018 Golowich et all, 2017 Li of all, 201;; Arora ef. al, 22T8)). PAC-Bayes and compression- approaches further improved the tightness of these bounds and, in some cases, yielded non-vacuous guarantees for large networks (Zhou el all, 201;; a 201;; Long & Sedgh, 201)). Despite this progress, classical complexity-based bounds are widely rocognized to be quantitatively loose for modern overparameterized architectures, motivating alternative perspectives that incorporate oplimization dynamics data-dependent structure.
Generalization bounds for Transformers. Several works have specialized norm-based generalization analyses to Transformer architectures. Edelman et all (2024) derive a generalization gap bound for Trans- formers that scales logarithmically with sequence length by exploiting bounded self-attention norms and sparse variable creation. This line was strengthened by [Trauger & Tewari (2022)), who obtain sequence- length-independeent generalization gap bound. Beyond norm-based bounds, PAC-Bayes and algorithmic stability frameworks have been applied to Transformers and large language models. Compression brasdeessoo-basedd based PAC-Bayes analyses yield non-vacuous generalization bounds for large protrained models by relating gener- alization to compressibility and posterior complexity (Loth all 2123) Stability-based approaches analyze sensitivity to data task perturbations; for example, IL et all (20281T study in-context learning through multitask stability, while related stability analyses for attention and fine-tuning appear in (Ya all (2122
Yang et al. (202)5) study generalization of Transformers for multi-step reasoning tasks.
Optimization-awarre bounds and NTK-based analyses. A different line of work derives generalization guarantees by explicitly coupling optimization dynamics with statistical complexity control in overparame- terized regimes. For fully connected networks, Arora ct all (2m)²) establish a truc-loss bound for two-layer ReLU networks trained by gradient descent using neural langent kernel (NTK) theory, and & GT (20л) apply this approach to deep networks via neural tangent random features. More broadly, NTK- based analyses characterize convergence and genoralization behavior of wide networks under lazy training dynamics (Aaoot
all 2018; Arora et all 20ha). These works motivate optimization aware analyses for more structured architectures such Transformers.
2

---

Published in Transactions on Machine Learning Research (01/2026)
Training dynamics of Transformers. Understanding the optimization and convergence behavior of Transformers has recently attracted significant. attention. Wi el all (2022) establish global convergence guarantees for shallow Transformers under suitable overparameterization and data assumptions, providing
a Transformer analogue of NTK-style results for fully connccted networks. Related works analyze conver-
gence rates, representational dynamics, and optimization stability allention- architectures (Kohler
& Krzyzak
2023; Huang et all. 2024a: Shen of all 2023: Gurevych et alg 2022; Ahn et all, 2023, Tarzanagh
al
2123
Li el all 20238a Tian all 2023; Song all 2024 Huang all 2IPK; Chen
2024: Yang
2024:
Deora et 2023 Gao et all 2024) Huang et al. 2024 lelassi et all 2022 Wang et al 2024
Yang et al. 202bat Li all
2025 Yang et al. 2022b).
Positioning of the present work. Our work complements the above literature by providing, optimization-awary true-loss bound for Transformer architecture. Specifically, combine (i) Rademacher complexity bound for bounded-drift Transformer hypothesis class with (ii) an explicit empirical-loss bound derived from global convergence theorem for shallow Transformers (lWi et all, 2022). Unlike norm-based, PAC-Bayes, and stability-basee based approaches that bound Lx(f) (s(f), our analysis directly upper-bounds CD(l) under lazy-regime training dynamicc. We emphasize that our results apply to single-layer Transform- ers under bounded drift and are therefore complementary i.e., not directly comparable to existing analyses of deep non-bounded drift Transformer models.
3 Preliminaries
3.1
Problem Setup
3.1.1
Training Examples
n=1 E xd are the instances and y A IN e R are the labels. d, is the sequence length of the inputs and d is the input dimension.
We are given N training examples ((x,, vn)l where {X,)
3.1.2 Model
The model used in this work is popular transformer encoder which is also used by Wu et al 2022 Given input X E Idd.xd, we define each of the transformer layers.
Self-attention layer The self-attention layer is defined as follows;
A - 0, (XW (XWO) (XW;))
Vdm
where a, is the row-wise softmax, Wo, WK, Wy Rdmxd the query, key and value matrices in the self-attention layer. dm is the model dimension. We shall be interested in the effect of the self-attention layer on each row (K.)) of the input X where E [da]. We therefore define 8 as the i-th row of the softmax output;
W WKX T
XW W
B; =0
ddm
=a.
ddm
We also define z; as the final output of the self-attention layer for each row ((1);
(XWT Wo(
= (XW XT
ddm
3

---

Published in Transactions on Machine Learning Research (01/2026)
Feed -forward ReLU layer The layer with ReLU activation function is defined as follows;
A _ (A,))
where or the ReLU activation function. For ease of calculations, Wnr is set Wif I RRdm Xdm Once again, define k; as the final output of the Feed -forward ReLU layer for each row X (f:);
k; or(z) - -o (WiXTo,, 0.
XW Wo(
Vdm
Average Pooling layer The pooling is applied column-wise to reduce sequence length dimension from d, to 1. This is done to ensure a scalar output from our transformer.
0g (A2)
where oo represents the column-wise average pooling, We can also define in terms of each ky
1 d.
d
K WW
fpre
Σ k;
d Σ or (WrXTo, os
dg
dmm
Output layer
The final output layer is defined follows;
f(x) 2 worpr
where wo E Rd is the weight vector in the output layer. We can as well define the final model output ((X) in terms of each row ( of the input. X;
XW Wo (o( (19) (1
((x) - - Σ or (WYXTo,
ddm
Define as a vector representing the union of all parameters of the transformer model as shown below;
= ( WK. Wv, wom
When we pass single input X E Rdd xd to the model, the output is given f(X) E R. When we give all inputs to the model as batch (X, J.N E Nxd,xd the output of the model will be (f(X,) e RN and output of the last hidden layer will be Fore & (N E Nxdm
3.1.3 Initialization
Similar to Wi et al (022
the LeCun initialization described below. The parameters Wq. Wx. Wy are initialized as (()) NN 2), w(y) ~ ~N(O.) 2). w(() ~ ~N(0, 2) for E [dm]] and j E [d] while w is initialized as wo NO, + for E [dm].
3.1.4 Empirical Loss
We consider any loss function e(J(xa), wn) which is 1-Lipschitz in the first argument;
1 N Cs(s) = N
Σ
n=1
4

---

Published in Transactions on Machine Learning Research (01/2026)
This empirical loss is to be optimized using Gradient Descent algorithm shown below;
Input: data (X Un Hr))==1 step size
Initialize weights as follows: 0 == {Wo &
WK: WP, V. w
for to 1 do
W W - 7 Vw )
W W W 7. w. E
W W WV -7- -.V E(0')
O = w' 7 Vwo e(0")
end for
Output: the model based on 0"
3.1.5
True Loss
We interested in upper bounding the true loss defined follows;
CD(f) (xxx
3.2 Rademacher complexity
The theorem of Rademacher complexity is widely used to compute generalization bounds for machine learning models. As per Mohri ct all (2012) theorem 3.1 and Arora cf all (201н)) thcorem B.1, suppose that the loss function E(,-) is bounded in [0, c] and is p-Lipschitz in the first argument. Then with probability at least 8 over the sample $ {(X,, Mn)l of size N:
log(2/6)
(up(C Ls()) S 2pRs(F) + 3eV JEF
2N
where CD(f) is the true loss, Ls(f) is the empirical loss and (F) is the empirical Rademacher complexity of a function class F for samples S = {(X,, (n)) ) of size N defined as follows;
RF(F)
N E Be
sup Σ on (Xx)) JEF
In order construct our generalization bound, we shall upper bound both the Rademacher complexity (F)F)) and the training loss Ls(f) for all F.
3.3 Covering number bound
For given class F. the covering number Noc (Fic (X,) (122) is the smallest size of collection (a cover) C F such that Vf F. C satisfying faxx]f[(XX f(X,,)ll Se. n The Rademacher complexity of the class with respect to samples S {(Xn, yn) can be upper bounded using the covering number F (Edelman ef. 2021);
Rs(F) < c- inf 8+ 1, 1//
log (X: (N M 112) de
620
for some constant e> and |f] < for all f F.
5

---

Published in Transactions on Machine Learning Research (01/2026)
4 Results
In this section, we develop a theoretical framework to analyze the generalization properties of Transformer models whose parameters remain close to their initialization during training. We begin by formally defining a class of models that satisfy this bounded-drift property, which corresponds to the lazy training regime. We then derive an upper bound on the Rademacher complexity of this class, followed by an upper bound on the empirical loss using convergence guarantces under gradient descent. Combining these results, we present our main theorem that establishes high-probability bound the true loss. Finally, we discuss the scope and limitations of our findings in light of existing results, and conclude with key insights and directions for future work.
For ease of proof, and without loss of generality, let us set the input feature dimension d to be equal to the model dimension dm i.e. d = dm-
4.1 Defining a class of Transformer models whose weights stay close to their initialization
To rigorously analyze the generalization behavior of Transformers, we first need to formalize the notion of models whose parameters remain close to their initialization throughout training. This assumption (often referred to as the bounded-drn drift assumption) characterizes the lazy training regime, where model updates are small, and the network operates in a nearly linear regime around initialization. In this subsection, we define the parameter space and construct a hypothesis class of Transformer models confined within a ball of radius R centered at the initialization point. This setup enables the derivation of subsequent complexity and loss bounds under analytically tractable conditions.
Recall that we defined 0 a vector representing the union of all parameters of the transformer model shown below;
0 (WQ. WK. Wv, wom
The squared f2-norm of the parameter vector can be expressed as the sum of the squared Frobenius norms (for matrices) and squared l f--ormm (for vectors);
e = Well + W + We 1| + |
We can therefore say that for all training steps >> 0;
1 0 = W Wt++ W + [ w + | w + 1| w0ll
1 Ra- R²x + Rs + R5
wgll: ; < R2, W l Rk, || -w9lf s Rv, 1/ -wylt Ro for some positive constants Ro, Rv, Ro. Rx
where || -WW+t--oll
Setting R == R3 + R2 Rg gives |0 0 < R. We then define our hypothesis class FR
comprised of the transformer models whose parameters 0 stay in a ball close to 0 for all training steps >>00 0;
J.. = (fo(X.)) :t 0, g+ 0%12 s R}
4.2 Upper bounding the Rademacher complexity
To establish a generalization bound, we must first control the capacity of the hypothesis class of models under consideration. The Rademacher complexity provides data-dependent measure of this capacity,

---

Published in Transactions on Machine Learning Research (01/2026)
quantifying how well the model class can fit random noise. In this subsection, derive upper bound the Rademacher complexity of the bounded-drift Transformer class defined above. Our result shows that under reasonable assumptions on the input features and parameter norms, the Rademacher complexity scales as (V log(Av/ ())), }) indicating that generalization improves with an increasing number of samples and controlled parameter magnitudes. This bound parallels similar results for shallow transformer and provides the foundation for our overall generalization analysis.
The following lemma gives an upper bound on the Rademacher complexity of our class of transformer models I.e., upper bound on Rs(FR
Lemma 1. Suppose that we have Πy Wle Ry. 11wlll2 + Ro, пK
l| IF
Rk.
F RQ where Ro, Rv, RK, RQ remain as defined above. Also assume that the inputs have full rank and are bounded as ||X,||r I Vx,Rx for all n E [N] where Rx some positive constont. The empirical Rademacher complerity of the class of Transformer models FR (fo(X,) Vt > 0, 1. - 001 - R given f = (W WK. Wv.wo) can be upper bounded as follows
Rs( Rs(F 0
1 NVN
N 1 log AV
Av²)))
where S hides logarithmscc dependencies on quantities besides N and da; nonv(Vd,Rx) and P (Vd,Rx)² ((Vamniv) (Vmm)nnon))) log(Nd,)
Proof. Define the following quantities for simplicity F + no = + лK | RR, 1iQ = [ + RQ where Ro, Rv, Rk: RQ remain as defined above in section B.1l
Our class of interest in section was F {X)(X,) ||0r+| 0 - R} and we want compute upper bound the empirical Rademacher complexity Rx("" which is given follows:
N
X, W (o(X((f)))
R=(') Ee-unif((-1.11)
1
WVVX
sup W Wo.
En d
w Gr
v Vdm
Wo.WV: WV:
WwO 2 Eno Wv STV
"x Vim
1
N
de
X X, Wq Wo (x
Σ or
W.X,
sup
E W
Nd,
Vdm
WO. W
Wy.WV: Wy:
R=1
1=1
1100
22 Sno Wy
SMV
Wa
Vim

---

Published in Transactions on Machine Learning Research (01/2026)
Applying subadditivity of the supremum:
1 d Nd.
X W WiWWWo(X Wx((x ()) 7
Rs(Fhh) 5
B Bewnif((-1.1))
sup wo. W Wo. Wy: I So Wv PSVVVV
Σ En wo or Wv 0s
dm
K "K" Vdm
vdm
XnW K W (X (1.))
d,
E, Be-umi(((11.1))
sup wo K
En wo or Wv
Nd,
Vdm
W..We Wv:
mm
100 2250 [W We
| STV
4m
4m
N
X,W, W (x (1.1
E (ewunif((-L.L)
sup
Σ
Cn w or
W. WvX, 0s
v dm
W K Wy.Wy: Wy:
wwlll STO Wv
"K"0 vdm
A Rs(
for any fixed E [da]. Hence,
Rs(Fh)) s Rs(
where Goe is defined as follows
0% = (x(li))) + woa as( x. w (x(( Vam
w Wo We vim
|lol| no, Wyyle S nv.
The following lemma gives an upper bound on Rs(g Its proof can be found in the appendix section;
Lemma 2. For any fixed e > 0 and X XN Rd.xe xd such that |Xalle S vTllx for all n E [N], the Rademacher complerity of go, satisfies the bound given below;
P N
N A
Rs(gh) S o
1 + log
P
where S hides logarithmic dependencies quantities besides N and d,, A nonv(Vo,Rxx) and P (VT,Rx)² ((vimin)) + (vT log(Nd.).
8

---

Published in Transactions on Machine Learning Research (01/2026)
Finally, the upper bound on the Rademacher complexity Rs(FFi)) can be given as;
Rs(FR)
S R9
N )
P
2
1+ log
Av AVP
N
P
N P
S 0
1 + log
AV
N
This completes the proof.
4.3 Upper bounding the empirical loss
Having bounded the complexity of our hypothesis class, we next analyze the empirical loss achieved by gra- dient descent under the bounded-drift condition. This subsection establishes that, given suitable conditions on model width, learning rate, and data structure, the empirical loss decays exponentially during training. Using results from convergence analyses of Transformers in the lazy regime, we derive an explicit upper bound on the empirical loss function of key quantities such as a, p. and no- This provides quantitative connection between network conditioning, data complexity, and training behavior, ensuring that even under restricted parameter updates, the model achieves low empirical loss with high probability.
Define a as the minimum singular value of FO pret i.e., a - amin (
(Fpre) and also define ((0) as follows;
( 7 ff( -
We now state the following assumption about the input data matrix X;
Assumption 3. Assume that the input data has full ronk and is bounded as |X|le Vd, Rx with some positive constont Rx. Furthermore, For any data pair (Xns Xw), with F n' and n' E [N), then we assume thal;
P([ X X X 5
with some constant è > 0
The lemma below gives an upper bound on the empirical loss for all training steps > 0.
Lemma 4. Suppose that we have пソ | l|r Rv, no l + Ro. пK
1| || RK nQ
[ F Ro, &Q [ RRo Ex [ ll2 + RK, &v A Rv where Ro, Rv Rx. Ro remain defined earlier. Under assumption G. if dm {(NN), a² 2 8pM (20)),) a (32p² (2))))))))))))))))) and ((0) is any loss function which is 1-Lipschitz in the first argumend, then with probabilily al least 1 Se -m_/// 8 8 exp(-S2((N - 1)- for proper 8, when training using GD with small step size S 1/k where l is conslant depending (Eq,EK, En, (00)), ) dm p, the empirical loss can be bounded follows for all t 0;
a²
Cs(fe)) S min 8pM
32
where δ omits the logarithmic factor and the other quantities defined as follows; p N Rx n1 + (4 +),,,, & n8(1 + (4/dm) RR3 df (nă + +
M max(fv (L . 2/ de \ max(mv 11
R R (2/
(2/v o

---

Published in Transactions on Machine Learning Research (01/2026)
Proof. For the purpose of simplification, define the following quantities at initialization;
Se _ IW + Ro - [ + Ro _ nQ
EK _
[WA 2| + Rx - [W*
Rx 1
5v A W la + Rv - W || + Rv - Tγ
1o 1lll2 Ro
where Rq, RK, Rv, Ro are as defined before. As mentioned carlier, a is the minimum singular value of F pre: According to Wn ct al (212 theorem 1, under assumption G. if dm 2 ((N*), a² 2 8pM (2) and a 2 (32p²z /2)())))) no, then with probability at least Sem - 8 - cxp(-S((N -))-dd")) for proper 6, GD converges to global minimum follows for sufficiently small step size y $ 1/k with k as constant depending on (εQ,EK, Ek, Ex, no.
T.e., a - min P pre) and (0) is given (0)) 2 llf(0) y
no, 200)), P. d
8(0") = (1 11--2)
a²
NoRy' (2/ dm, RydEEKEvOOO (R and p N1/2 2/2 Rx f r1 + (4/dm) (E3 (x)).
where M max(fv R
We can observe that ((') decays exponentially as training proceeds. This implies the following bound;
4(0') - (00).
V 2
From the first condition i.c.,a² 2 8pM (2)) we can say that (000) s a M²) We therefore end up with the bound below;
(00')= < 128 M
Vt 2 0
From the second condition i.e.,a 2 (32p²x, / 24 08) therefore end up with the bound below:
//oO, we can say that 4 s (2048) We
(00)) < a no
AA 2 0
2048p*2
Combining the two bounds on (0'), we oblain the final bound as;
a
(00)) S min
V1 0
128 2048pt
Our empirical loss i.e,, Ls(sor) N Ei {(Sor (Xn), yn) for all t 6> 0 can be bounded as follows;
N
Ls(for <
Σ
{(for (X.),
-E(yn vn))
N a=l
Σ
for (X,,) - Val
because ((.,)) is 1-Lipschitz in the first argument
N
1
f. y112
VN
24D (
N
a²
a no
< min
SpMVN
32
10

---

Published in Transactions on Machine Learning Research (01/2026)
where M = max(fv Ry (2 (2 and 222 22 p _ d Rx a n1(1 + (4/dm) R de (5Q + (
Upper bounding Eo using no, l, K,, nQ, the upper bound on the empirical loss for all training steps can therefore be written as;
a²
a'no
[s(fer) I min
S 32p22vN
where M max Ro Ro (2/v Bxd (2/ R3 and & a nō(11 + ] R2 )k)) This completes the proof.
4.4 Main result
With both the Rademacher complexity and empirical loss bounds in place, we now combine these results to obtain our main theorem i.e., a high-probability upper bound on the true loss (expected generalization error) of shallow Transformer models under bounded parameter drift. The theorem demonstrates that the true loss decreases with the number of training samples N and depends explicitly on the initialization scale, data structure, and model dimension dn.. The bound captures the essence of lazy training: when model parameters remain near initialization, generalization ion behavior aligns with that of linearized models governed by the neural tangent kernel (NTK). While this result establishes a rigorous theoretical foundation for shallow Transformers, it also highlights the need for further analysis to extend such guarantees to deeper and more expressive models.
Theorem 5. Suppose that we have yy |Welle + Ry,no no lw l2 + Ro. πK [ + RK nQQ 1 RQ, εQ =
I [Qall2+ RQ; ÊK || + RK. εv I + Ry where Ro. Rv. Rg. RQ remain defined earlier. Under assumption B. if dm 2 ((N*), a² 8pM 24 a² 2 (32p²zv/40/ (20(()))mm)))) and (()) is any loss function which 1-lipschitz in the first argument, then with probability at least Se dma - 1)=dd;))) if the transformer model brained using Gradient Descend with small slep sice 7 S 1/k where k is constant depending EK, no,,(o,00000 {(0%), bounded as follows;
25
p, dm ) the true loss LD(f) can be
Lo(J) S min
a²
a'no no
+0
P (
1 log
log
(
A
N
log It
8pM 22
where Z omits the logarithmic factor, S hides logorithmic dependencies on quantities besides N, d, and 5 and the other quantities are defined as follows; p è N /22 2 (1 (4/dm. R dee2 (52 () & e 18 (4/dm) Rx Ln (ng nzk)),
M = max Ro no R R8 (2/ dm deE
M max (n R R (2. R3 d. Ro (2/ R3 d R
A = пoтL (vd,Rx) and P (x), ((Vdmiv) + (Vdm))xnom))) log(Nd.).
Proof. Recall that we defined our hypothesis class as follows;
Fdº (fo(X,) ||0 - 001 - R}
Let us set R; = for i E 21,2.... .... R}. This means that we can define class of models whose parameter
norm is bounded as ||0 001 1/ < R; for i (1, R} as follows;
FR. [So
1 | 00 R.}
11

---

Published in Transactions on Machine Learning Research (01/2026)
From Rademacher complexity and union bound over finite set of R,'s, for any random initialization (0))
with probability at least 1 - 8 over the sample S = {(X Un) Wn)) of size N, we have that;
log 28 2N
[LD(J) - Ls(I)) 2Rs(FR) +
sup
feFR;
for all (1,2,3....,,R) Note that R, S R for all i .1. R} which implies that Rs(Fi, S Rs(Fh)) for any E (1,2...., RJ. This gives the following bound on (s(,) for all E 21. R};
Rs(FR) (V~ 1+ log A (1/)))
where P= (Vd,Rx)² ((VamTv)& (Vdmminnan)) 8) log(NNd, and A = non From lemma B!. 1) d)-dd the training loss for our transformer
with probability at least S0-a1 - model can be bounded as follows for all t> 0;
a² 8pMvN /
a ao
Cs(fer) S min
where p A N1/2d Rx 8 18 (4/2m)) (n + + and
M max(ny Ro¹, Mo Putting everything together, with probability atleast 1 8c dm! that;
( (2/ R 1 (2/ R² d, 1Q Rg)
25 expl (1)(N-))-dd))) we have
a²
a a
P VV
log R N
Ip(()) S min
+0
+ log A
AV
+
SPM 32
where < hides logarithmic dependencies quantities besides N, d, and 6. This completes the proof.
4.5
Discussion
Our main theorem provides a generalization bound for a class of Transformer models whose weights remain close to their initialization during training. This bounded-drift assumption effectively constrains the training dynamics what commonly referred to as the lazy training regime. In this setting, the model behaves similarly to its linearized form around initialization, which has been extensively studied in the context. of neural tangent kernels (NTK).
The implication of this assumption is that the results derived here apply most accurately to wide and shallow Transformers trained with sufficiently small learning rates that prevent the parameters from deviating significantly from their initial values. Such a setting captures the carly or ncar-lincar training phase of overparameterized models, where the NTK remains nearly constant and generalization can be controlled through classical complexity measures such as the Rademacher complexity.
Compared to prior results generalization bounds for neural networks under the lazy regime (c.g., for two-layer networks and linearized models), our bound maintains a similar dependence on the sample size N and model width dm. Specifically, the o( Log(A/ /$) term scales analogously to existing NTK-based results, while the dependence on a and in the empirical loss component rellects the conditioning of the pre-activation features and the interaction between model width and sequence length. However, unlike the more general results that extend to deep networks or non-lazy regimes through complex stability analy- ses or overparameterization assumptions, our result explicitly applies to single-layer (shallow) Transformer architectures only.
12

---

Published in Transactions on Machine Learning Research (01/2026)
We also note that relaxing the bounded-dritt assumption to capture non-lazy or feature-learning regimes remains an open problem. In such regimes, weights undergo significant changes during training, resulting in evolving representations and coupling across layers. Extending our analysis to this setting would require novel techniques to handle the dynamic cvolution of the NTK or an equivalent representation matrix. Sim- a
ilarly, extending the bound to deeper multi-layer architectures would require careful control of layerwise dependencies, possibly through hierarchical complexity bounds or layerwise Lipschitz control arguments.
Overall, our results contribute to the growing thcoretical understanding of Transformers under simplified but analytically tractable conditions. They highlight the relationship between network width, data complexity, and generalization under bounded parameter drift, reinforcing the intuition that in the lazy regime, wide and shallow Transformers behave as near-linear models governed by their initialization structure.
Comparison with other norm-based Transformer bounds. Our result differs substantively from norm-based bounds tailored to Transformers. Edelman al (202T) obtain a gap bound that scales only log- arilhmically with sequence length by analyzing bounded-norm self-attention and the sparse variable creation inductive bias, while Trauger Tewari (2112)) strengthen this line by proving eqquencelengtth wewpnneeeeennt---deeeeeeeeee independent gap bounds via covering-nuntn numbbrr analysis that upper-bounds the Rademacher complexity of Transformer classes and applies also to masked-token training objectives. In contrast, our theorem provides a. true-loss bound that explicitly couples data-dependent optimization ion guarantee (lazy-regime convergence for shal- low Transformer) with capacity term. Practically, this means (i) worm-based results (Edellan et all, 2021 Trauger & Tewari, 20224 can be evaluated after training completes by plugging in norms to bound [fD()) - (f),,, (f)),, and they cleanly characterize how sequence length enters (logarithmically or not at all), whereas (ii) our bound ties generalization to optimization-side quantities such as a, p, and (no. 7K, TK,, "T, under bounded drift and thus characterizes when low true risk is guaranteed without first computing (f()). Conceptually, the results agnostic control of the generalization gap (with refined sequence-length dependence), while our analysis isolates the lazy, shallow regime and provides optimization-aware control of the level of the true loss.
complementary: norm-based bounds offer architecture-widde training-regime-
Relation to PAC-Bayes and stability-based bounds. Compared to PAC-Bayes or compression bounds for LLMs (Lotli et all, 2023: Zhou et all, 21TS) our result targets a different regime and objective: we op- erate under a bounded-drin (lazy) assumption and obtain a true-loss bound for a single-layer Transformer trained by gradient descent, with explicit dependence on quantitics like a, p, and (nq. Mk. 7 no).))). PAC- Bayes bounds, in contrast, typically control [Cp(f) ((s))) via a data-informed posterior and compression, and have been instantiated for large, deep, pretrained language models withont requiring lazy dynamics. Stability-based results for Transformers (e.g., in-context learning) quantify sensitivity to sample perturba- tions through stability coefficients (Li 2122E), again bounding the generalization gap rather than the true loss. Thus, our contribution is complementary: it connects optimization driven lazy training (bounded drift + shallow width assumptions) to generalization, while PAC-Bayes and stability provide depth-agnostic, regining--egimm-ggnien
controls on the gap. Extending our approach to remove the lazy assumption or to handle deeper stacks could help bridge these viewpoints, potentially yielding hybrid bounds that combine optimization-awaree truc-loss control with posterior- or stability-basod gap terms (Yao ef all 2025)
5
Limitations
While our analysis provides a rigorous generalization bound for shallow Transformer models trained under bounded parameter drift, it is important to clarify the scope and limitations of the results.
Loose nature of Rademacher-based bounds. Our gencralization analysis relies on upper bounds on the empirical Rademacher complexity of a constrained hypothesis class. It is well known that classical Rademacher complexity bounds for neural networks are often quantitatively loose, especially for large, over- parameterized models such as Transformers. In particular, these bounds are typically worst-case and may not tightly reflect the effective capacity of models encountered in practice. Consequently, while our results establish principled scaling behavior with respect to sample size, width, and norm constraints, the re-
13

---

Published in Transactions on Machine Learning Research (01/2026)
sulting bounds should be interpreted as qualitative guarantees rather than sharp predictors of empirical performance.
Restrictive nature of the lazy (bounded-drift) regime. A central assumption in our analysis is that the Transformer parameters remain close to their initialization throughout training, corresponding to the lazy or NTK regime. Although this assumption enables tractable theoretical analysis and connects our results to classical kernel-based generalization theory, it does not fully capture the training dynamics observed in practical Transformer training, where substantial feature learning and representation drift typically occur. As result, our guarantees most relevant for wide, shallow Transformers trained with sufficiently small learning rates, for early training phases before significant departure from initialization.
Comparison with existing analyses of Transformer training dynamics. There cxists a growing body of prior work that studies Transformer training dynamics using alternative analytical tools, including mean-field limits, dynamical systems analyses, signal propagation, and task-specific convergence character- izations. These approaches often provide finer-grained insights into representation learning, optimization trajectories, or task-dependent inductive biases beyond what is captured by Rademacher-baid based complexity measures. Our work is not intended to subsume these analyses; rather, it complements them by offering unified generalization bound that explicitly couples optimization guarantees (via lazy-regime convergence) with statistical capacity control for well-defined hypothesis class.
Architectural and depth limitations. Our results apply only to single-layer (shallow) Transformer ar- chitectures. Extending the analysis to deeper stacks introduces significant technical challenges, including the control of layerwise interactions, accumulation of approximation errors, and evolving attention representa- tions. Addressing these challenges would likely require new techniques beyond the bounded-drift framework employed here, such hierarchical complexity bounds, stability arguments, or depth-dependent norm con- trol.
Outlook. Despite these limitations, we view our analysis as useful step toward understanding how opti- mization dynamics and generalization interact in Transformer models under analytically tractable conditions. An important direction for future work is to relax the bounded-drif assumption and to integrate alternative toolssuch as stability, PAC-Bayes, or feature-lcarning analyses with optimization-aware bounds, in order to more faithfully capture the behavior modern deep Transformers trained beyond the lazy regime.
6 Conclusion
In summary, we established a generalization bound for shallow Transformer models trained in the bounded- drift (lazy) regime, where the model parameters remain close to their initialization throughout training. By combining Rademacher complexity analysis with an upper bound on the empirical loss, we obtained a probabilistic bound on the true loss that decreases with the number of samples N and depends explicitly on model and data. parameters.
Our theoretical results align with existing findings for wide, overparameterized models analyzed under the NTK framework, but they are specific to single-layer Transformers. This limitation ensures that our claims remain within the thcoretical scope supported by the bounded-drift assumption. Extending the analysis to deeper architectures or non-lazy training regimes (where substantial feature learning occurs) remains an important direction for future research.
Through this work, we provide rigorous foundation for understanding how bounded-drift dynamics influence generalization in Transformer models and set the stage for future extensions that aim to capture the richer behavior of modern, deeper architectures trained beyond the lazy regime.
Acknowledgments
Anirban Dasgupta acknowledges the support by SERB MATRICS and SERB CRG grants and the support from N Rama Rao Chair Professorship.
14

---

Published in Transactions on Machine Learning Research (01/2026)
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ngge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arv preprint rrXie:2308.08774, 2023.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement proconditioned gradient descent for in-context learning. ArXiu, abs/2306.002m) 2023. URL https:
seapp.sematticsccoolaa wtw.:eawwa...ii......ooooooorraa
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. CoRR, abs/ abs/1802.05296,, 2018. URL http://araxv.orr/aa//B/222
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurIPS, 2019a.
Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fino-grained analysis of optimization and generalization for overparameterized two-layer neural networks. CoRR, abs/1901.08584 2019b. URL
Peter L. Bartlett, Dylan J. Foster, and Matus Telgarsky. Specttrall--norrmalizeed margin bounds for neural networks. CoRR, abss1700.0843) 2017. URL htpp://aaxxv.orr//////77/...88
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. CoRR, abs/1990 abs/1905.13210, 2019. URL http://ar..v.oor/////1115..
Minshuo Chen, Xingguo Li, and Tuo Zhao. On sout of famill of recurrent neural networks. CoRR, abs/ 19t/19990.1294,, 2019. URL http://ariiv.org//////9111111110
Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. Proceedings of the 57th Annual org/CorpusII:
ACM Symposium on Theory of Computing, 2024. URL
267500364
Puncesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimiza- tion and generalization of multi-head attention. ArXiv, abs/2310.12680, 2023. URL https://api. semanticscholar. wwwpoo/oros.::2243300.3.333
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 words: Trannformers for image recognition at scale. CoRR, abs/2010.11929, 2020. URL. htaps://arxiv http://axvvvvooooaaaa ottp:/bb///00..111/221
Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. CoRR, abs/2110.10090, 2021. URL https://arxiv..rg/ / 2110. 10090.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What single attention layer learn? study through the random features lens. Advonces in Neural Information Processing Systems, 36, 2024.
Cheng Gao, Yuan Cao, Zihao Li, Yihan He, Mengdi Wang, Han Liu, Jason M. Klusowski, and Jianqing Fan. Global convergence in training large-scale transformers. ArXiu, abs/2410.23610,, 2024. URL https:
wtt.oollaw..rrroosss:::::... 273707777
Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size-independent sample complexity of neural net-
works. CoRR, abs/ abs//112.006541, 2017. URL http://arxiv. wttp://orgbbbb///...oo
Iryna Gurevych, Michael Kohler, and Gözde Gill Sahin. On the rate of convergence of classifier based on transformer encoder. IEFE Transactions Information Theory, 68(12):8139 8155, 2022.
15

---

Published in Transactions on Machine Learning Research (01/2026)
Ruiquan Huang, Yingbin Liang, and Jing Yang. Non-asymptotic convergence of training transformers for next-token prediction. ArXiv, abs/2409. 17335, 2024a. URT. https://api swmanticscoooar..cmmmm org/
Corpu81D:212210094
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. ArXiu, abs/ 231000022. 2023. URL https://api. semanticacholar. wtt.e//io.....ooooooooooaaaaaaa.
Yu Huang, Zixin Wen, Yuejic Chi, and Yingbin Liang. How transformers learn diverse attention correlations in masked vision pretraining. In ICML 2024 Workshop on Theorelical Poundations of Foundalion Models, 2024b.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In NewrIPS, 2018
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822 37836, 2022.
Michael Kohler and Adam Krzyzak. On the rate of convergence of over- parametrizeddd transformer classifier learned by gradient descent. arXiv preprinl XrXi:23312 22 17007, 2023.
Hongbo Li, Lingjie Duan, and Yingbin Liang. Provable conteeet learning of nonlinear regression with trans- formers. ArXin, abs/2507. 20443, 2025. URL https://appls http:s///pp...emmttiiccccc ser.icccooemmmtiiccccccccolrrrrr o0tp
Hongkang Li, M. Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. ArXiv, abs/2302.00015, 2023a. URL
Arg/CorpusID: :020pp711:2268271011 2568271011
Xingguo Li, Junwei Lu, Zhaoran Wang, Jarvis D. Haupt, and Tuo Zhao. On tighter generalization bound for deep neural networks: Cnns, resnets, and beyond. CoRR, abs/ 1509..559,,,, 2018. URL http://arxiv. org/abs/ htrgabs/1806..o 05159
Yingeong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, 2023b. URL ettp:////pp...mmmticccccoolll htt.aacc.............roraaaaaaa.
Philip M. Long and Hanie Sedghi. Size-free genernlizaiinnn bounds abs/1905.1 12600, 2019. URL.
bounds for convolutional neural networks. CoRR, wwtp:///wwwwraabb/..b.....
Sanae Lotfi, Marc Finzi, Yilun Kuang, Tim G. J. Rudner, Micah Goldblum, and Andrew Gordon Wilson. Non-vacuous gencralization bounds for large language models. ArXiu, abs/2312.17173, 2023. URL https:
//appi.emmaticscholar.cco httppp/..aaaa.....c..a..oo.ooooa 2666773226
Mchryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar, Foundations of Machine Learning. The MIT Press, 2012. ISBN 026201825X.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Nase-based capacity control in neural networks. 00036.
CoRR, abs/ abs/1503.00036, 2015. URL
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. A pac- bayebaaesian approach to specttraly--normaaizeedd margin bounds for noural networks. CoRR, abs/1707.09566 2017. URL attp:
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards under- standing the role of over-parametrizaatoon in generalization of neural networks. CoRR, abs/ hbs/1805.1200 12076. 2018. URL httpp://aaxxv..og/bb///////s
Konstantinos Pitas, Mike Davies, and Pierre Vandorgheynst. Pay-bay margin bounds for convolu- tional neural networks technical report. CoRR, abs/1801.00171, 2018. URL http://arxiv.org/aa
00111.
16

---

Published in Transactions on Machine Learning Research (01/2026)
Wei Shen, Ruida Zhou, Jing Yang, and Cong Shen. On the training convergence of transformers for in-context classification. arxiv preprint or arXiv:24 10. 11778, 2024.
Bingqing Song, Boran Han, Shuai Zhang, Jie Ding, and Mingyi Hong. Unraveling the gradient descent
dynamics of transformers. ArXiu, abs/2411.07538, 2024. URL https://apl www.pp....aaanccccco....oooooo.a
Lorpus1D:2739699724
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. ArXiv, abs/23088 260922,0.16688 16898, 2023. URL htt.:pppp..etaaiicccccolollla htt.:/aattttcccoooooooooooooopaa
2613952061
Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean Cean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: family of highly capable multimodal models. arXiv preprint ar Xrx::23312 11805, 2023.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. ArXio, abs/2305.1633 16380, 2023. URL https:
//api. wwwanticscholatttttttttetcaeaaae httpoooaaar.rr/orpuusssIS. I0:
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. orXiv preprint arXiv:2302.1999 18971, 2023.
Jacob Tranger and Ambuj Tewari. Sequence length independent norm-based generalization bounds for transformers. In International Conference on Artificial Intelligence and Statistics, pp. 1405 1413. PMLR, 2024.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Loc. Transformers provably learn sparse token selection while fully-connected nets cannot. arXie preprint or arxiv:2466. 06803,,6633.... 2024.
Yongtao Wu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher. On the convergence encoder-only shallow transformers. Advances in Neural Information Processing Systems, 36, 2024.
Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, and Yingbin Liang. Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. ArXiu, abs/24100096055, 2024. URL hitps:
//api. semammitttiicchllrrr org//oorus11:22733663
Hongru Yang, Zhangyang Wang, Jason D. Lee, and Yingbin Liang, Transformers provably learn two- mixture of linear classification via gradient flow. In The Thirteenth International Conference on Learning Represenlations, 2025a. URL https://openreviier http::///pppeeeeeeeeee.e/roooooo wtt.:///ooommm////A/AA//...
Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. Multi-head transformers provably learn sym- bolic multi-step reasoning via gradient descent. ArXiv, htss/5008.0222,, 2025b. URL https://api. Bemanticscholar. wtg//orrpusD::280565333
Xinhao Yao, Hongjin Qian, Xiaolin Hu, Gengze Xu, Wei Liu, Jian Luan, Bin Wang, and Yong Liu. The- oretical insights into fine-tuning allention mechanism: Generalization and optimization. In Proceed- ings of the Thirty-Fout Fourth International Joint Conference on Artificial Intelligence, IJCAI 2025, Mon- treal, Canada, August 16-22, 2025, pp. 6838. ijcai.org, 2025. doi: 10.24963 2224966/ICCAL.22255//// URL https://doi.org/ 1tt..rr0///..229/11aaaaa 202557776..
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz. Non-vacuous generaliza- tion bounds at the imagenet scale: a pac-bayesian compression approach. In International Conference on Learning Representations, 2018. URL https://api. semanticscholar otrp:/aaaa....r///ppppaaa///572
17

---

Published in Transactions on Machine Learning Research (01/2026)
A Proof of lemma 2
We want. to obtain an upper bound on Rs( where OR is defined as follows;
X W K Wo Wo( (x))))))) W
))
I la Eno
O = (x((a))] - wa Wv 0s
Wy S
ddm
W R We
2222
Vdm
Vdm
Let's begin by defining the following bounds on the matrioes;
[Wv 1l v WWll2 la by S nM
Wv 22.1 &
Wil12,1 + Bv S ddm
W
) 111.2 + Bx) (11W8ll2.1 + BQ)
dmaknl Vim
Vnankn
Vdm
Vdm
2,1
| 12.00 S Bx & (xlll Vd.R Vn [N]
where bv, By, BR, Bq, Bx are positive constants and Ro, Rv, Rg. RQ, Rx remain defined earlier. The norm |- 12,, interpreted as first taking the 4-nor for each column of a matrix and then summing these column norms. Define another class 9 as shown below;
w S
WW (W thy
G3 = (x((1)) -
- wwo (WvX, (x, W (o(
[ S Ila. By
w Wo
(nww2) 11 #x (lw 12, -
ddn
ddm
2.1
The following lemmna gives an upper bound on the log covering number of the class 9
Lemma (lEdelman el all 2021) Corollary 4.5). For any fioed and X XN E Rd.xd such that | 12.000 S Bx for all E [N], the covering number of got satisfies the bound given below;
log No (x)) (12)
8)
(1IWllla + Br) + (( 11.2 (1atta W ( (IIwella + bv)
Vdn
$ B²-
log Nd.)
2
where < hides logarithmic dependencies quantities besides N and dy.
Upper bounding the norms
l|2,1 and
|| 1|2,00 using the Frobenius norm, || lle, we end up with;
((Vdm))) (VImIKIQQN)) $)
No : (X l12) S (Vd,Rx)2.
(Nd.)
This can also be written as;
P ?
log (9R' se; (Xa) 1N ll (12) S
where P= (VT,Rxx ((VTmin)) + +(TTunnnon)) log(Nd.).
18

---

Published in Transactions on Machine Learning Research (01/2026)
We can now write the bound the Rademacher complexity Rs(9 as follows for some constant 0 and lfl & A for all f e 9h:
A
No ( (X,, b Inaiill l12)) de
Ry (9h)
c. inf
6+
820
N
A
P
c. inf
6
de
820
22
P
c. inf
6 +
de
N
5
P
A
e inf
log
620
N
cV N
1 +
log
A
P
Note that [f| A for all
9% A be obtained as follows;
X. R
Wo (1
woo,
WV.
ddm
X W W (1:3)
S/wolle l
0r
Wv WvX
ddm
X K Wo (x(X (2.3) (
||woll
Or
Wv WvX 0s
V dam
X, N.WWWOWN GiA
WvX, lwolll
(because [|ar(2)]l2 ||2||2)
ddm
WW W (X (1.1)
[wrl22 wolll
X,To,
ddm
X, W.W. K W (X, (1:3)
|woll2 Wv
122
| 60
ddm
|woll2 Wv 12 | 11
(because (|a,(2)ll2 z (2).())]1 = 1)
|woll2 Wr |X,llle
(1ww) 12 (o) W || Ry),v
nonv (vd
This means that A = nonv Rx).
B
Experiments
B.1 Image Classification
We use the transformer model defined in section k to perform classification of images. From MNIST dataset, we extract the images belonging to classes and and create our new dataset. Each image of size x 28 is broken into tokens each of dimension d 64. The main goal of the experiments to demonstrate that the test loss of the trained transformer model decreases with increasing number of of samples i.c., = 400, N 1200 and N 10000. This trend holds for all the values of model dimension which tested
19

---

Published in Transactions on Machine Learning Research (01/2026)
i.e. dm 64, dm 1024 and dm. 4096. The learning rate used is 0.1, the optimization algorithm is batch gradient descent and the loss function is the cross-entropy loss. The results for the experiments are presented below. Each figure MH- shows the training loss and test loss of the transformer model as training proceeds.
0.7
0.6
N-1200
N-10000
0.3
0.2
0.1
200
400
800
800
1000
Epachs
0.7
N=400
0.6
N-1200
0.5
N-10000
0.3
0.1
1000
Epochs
Figure 1: Evolution of training loss (top) and test loss (oss(botoo)) for each epoch of training for model dimension
dm 64.
0.6
0.5
N-10000
0.3
0.2
0.1
200
400
600
800
1000
Epachs
N=400
0.6
N-1200
0.5
N=10000
0.4
0.3
0.1
Epochs
Figure 2: Evolution of training loss (top) and test loss (bostt))) for each epoch of training for model dimension
d, 1024.
20

---

Published in Transactions on Machine Learning Research (01/2026)
0.7
N=400
0.6
N-1200
N-10000
0.5
0.4
0.0
200
400
600
800
1000
Epachs
0.7
N-1200
N-10000
0.5
0.3
0.1
0.0
200
400
600
800
1000
Epachs
Figure 3: Evolution of training loss (top) and test loss (bottom) for each epoch of training for model dimension
dm 4096.
Table Lowest test loss and the epoch at which was achieved for different values of dm and N.
dm
N
Lowest Test Loss
Epoch
64
400
0.1048
147
64
1200
0.0668
165
64
10000
0.0263
546
1024
400
0.1839
95
1024
1200
0.0760
271
1024
10000
0.0045
251
4096
400
0.1269
139
4096
1200
0.0899
169
4096
10000
0.0123
312
B.2 Text Classification
We also perform similar experiments for text classification using the 20 Newsgroups dataset restricted to the categories sei and sci.space. Each text document is represented where each token corresponds to 50-dimensional Glo Ve embedding (d 50). This forms the input sequence for our shallow Transformer model. The main goal of these experiments is to demonstrate that the test loss of the trained Transformer model decreases as the number of training samples increases, i.e., for N = 400, N 1200, and N = 10000. This decreasing trend in test loss is observed consistently across all model dimensions tested, namely d, 64 and d,, = 1024. The learning rate used in the experiments 0.01, the optimization algorithm is batch gradient descent, and the loss function employed is the mean squared error (MSF) loss. The results presented below, where each figure R shows the training loss and test. loss trajectories of the Transformer model as training progresses over 2000 epochs.
sequence of 40 tokens (d, 40),
21

---

Published in Transactions on Machine Learning Research (01/2026)
0.45
N0000
0.40
N-1000
0.35
0.30
Traning
0.25
0.20
0.15
250
500
750
1000 Epodhs
1250
1500
1750
2000
400
N- 1200
0.6
0.4
0.2
250
500
1000
1250
1750
2000
Epochs
Figure 4: Evolution of training loss (top) and test loss losstto)) for each epoch of training for model dimension
dm 64.
4-400
N- N-10000
0.35
loss
0.30
0.2
0.2
0.15
250
500
750
1000
1250
1500
1750
2000
Epochs
N= 600
N- 1200
0.30
100000
0.20
0.10
250
500
1000 Epochs
1250
1750
2000
Figure 5: Evolution of training loss (top) and test loss losstbottom) for each epoch of training for model dimension d_,. 1024.
22


Reviews:
[{"review_id": "A9oJxXdKuS", "reviewer": "Reviewer_aNGY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The main contribution of this paper is establishing the first direct generalization bound for shallow Transformers by integrating training dynamics analysis with Rademacher complexity, achieving an $O(\\sqrt{P/N})$ convergence rate.", "claims_and_evidence_yn": "No", "claims_explanation": "The paper contains two technical errors, which compromise the rigor and numerical accuracy of the results:\n- In Lemma 1, the reduction from $\\mathcal{F}$ to $\\mathcal{G}$ incorrectly equates matrix \"full rank\" with \"statistical independence\" of rows, deriving the equality $\\mathcal{R}_S(\\mathcal{F})=\\mathcal{R}_S(\\mathcal{G})$. However, full rank only implies linear independence; subadditivity of supremum should yield an inequality instead. \n- In Lemma 4, the second term should be $\\frac{\\alpha^3 \\eta_O}{32\\rho^2\\hat{z}\\sqrt{N}}$, but the paper incorrectly writes $\\rho^4$.", "audience_yn": "Yes", "audience_explanation": "Transformer generalization theory is a timely topic, and the methodology of combining training dynamics with complexity analysis is novel. Despite the technical errors, the corrected work would provide valuable theoretical tools for understanding Transformer generalization.", "requested_changes": "1. Remove the independence claim; apply subadditivity to obtain $\\mathcal{R}_S(\\mathcal{F}) \\leq d_s \\cdot \\mathcal{R}_S(\\mathcal{G})$; adjust sequence length dependence accordingly.\n\n2. Correct the exponent from $\\rho^4$ to $\\rho^2$ in the second term; verify the derivation chain; update Theorem 5 and all dependent results.", "broader_impact_concerns": "NA", "additional_comments": "NA", "extra_scores": {}, "extra_text": {}}, {"review_id": "G2qUuopORE", "reviewer": "Reviewer_3jxZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new norm-based generalization bound for a shallow Transformer trained with gradient descent. The approach builds on three interconnected components. First, the authors define a restricted hypothesis class of Transformer models whose parameters remain within a bounded neighborhood of their initialization at every training step. This assumption is motivated by the over-parameterized regime, where parameter drift is limited in practice. Second, they derive an upper bound on the Rademacher complexity of this class, employing covering number estimates to characterize the function space. Third, they upper bound the empirical loss for this class by invoking convergence guarantees for shallow Transformers trained with gradient descent, specifically relying on recent results about global convergence. Together, these steps yield a final theorem that provides an explicit upper bound on the true loss, which tightens sublinearly with the number of training samples for any model dimension. To provide empirical evidence, the authors conduct small-scale experiments on a binary subset of MNIST, showing that the test loss decreases as the number of training examples grows, and that this trend is consistent across model dimensions. The principal contribution lies in extending the methodology of directly bounding the true loss—previously demonstrated for fully connected networks—to the Transformer architecture.\n\nThe strengths of the paper are threefold. First, the work addresses an important gap in the theory of deep learning by providing a direct true-loss bound for Transformers, which had not previously been studied in this form. Second, the proofs are thorough and mathematically careful, with detailed derivations of the Rademacher complexity bound and a clear connection to convergence results. Third, the structure of the argument is well-motivated: the authors justify their assumptions, systematically develop the lemmas, and integrate them into a coherent theorem that contributes meaningfully to our understanding of Transformer generalization.\n\nDespite these strengths, the paper also suffers from weaknesses. The bounded-drift assumption on parameters is very restrictive, and while the authors argue that this reflects over-parameterized training, the conditions are not carefully validated in practice. The empirical section is underdeveloped: using only two-class MNIST, without comparisons to alternative bounds or to real-world Transformer training, makes the evidence weak relative to the theoretical ambition. Finally, the contribution is somewhat narrow. The focus is restricted to shallow, single-head Transformers, and while the authors suggest extensions to deeper and more general models, no concrete results are provided, which limits the impact.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes, the claims are supported in the sense that the theoretical derivations are correct and internally consistent. The use of covering numbers and Rademacher complexity is rigorous, and the application of the shallow Transformer convergence theorem is appropriate. However, the evidence is not fully convincing with respect to practical relevance. The experiments are too limited to demonstrate that the bound meaningfully predicts generalization in real settings. The authors overstate the empirical side of their contribution, since the experiments only verify a trivial monotonic relationship between sample size and test loss. Thus, while the theoretical claims are well supported, the practical claims are weak.", "audience_yn": "Yes", "audience_explanation": "Yes, the theoretical community studying generalization in deep learning, and in particular those focusing on Transformers, will find the results of interest. The extension of direct true-loss bounds to this architecture fills a gap in the literature and contributes to the theoretical toolbox available for analyzing large models. However, the readership interested in applied or empirical aspects of Transformers may be less engaged, since the results are highly stylized and the experiments are too modest to carry practical significance.", "requested_changes": "A critical change is the need for stronger empirical validation. At minimum, the experiments should move beyond two-class MNIST and include synthetic or toy sequence modeling tasks that more closely resemble the intended domain of Transformers. Without this, the experimental section adds little value. A second critical change is to clarify the scope and limitations of the main theorem, especially regarding the bounded-drift assumption and the shallow architecture. The current text suggests broader implications than are warranted. A third critical improvement would be to situate the contribution more carefully against existing generalization bounds for Transformers, beyond the short discussion of Edelman et al. (2021) and Trauger & Tewari (2024). Comparisons to other frameworks, such as PAC-Bayes or stability-based approaches, would help contextualize the novelty. Non-critical but useful improvements include polishing the presentation of proofs for readability, especially in the long lemma sections, and expanding the discussion section to include more explicit insights about what the bound implies in practice.", "broader_impact_concerns": "There are no major ethical concerns raised by this work. The paper is entirely theoretical and does not involve sensitive data, applications, or downstream deployment. The broader impact statement should, however, acknowledge explicitly that the results are limited to shallow and constrained Transformer models, and thus their implications for the safety or reliability of large-scale Transformer deployment are speculative.", "additional_comments": "Overall, this paper makes a meaningful but narrow contribution to the theoretical understanding of Transformers. Its strength lies in extending direct true-loss bounds to a new architecture with rigorous mathematical treatment. Its main weakness lies in the limited empirical evidence and the restrictive assumptions, which limit practical significance. With stronger experiments and clearer positioning within the broader literature, the paper would merit acceptance. As it stands, it is a borderline case: strong on theory but weak on applicability.", "extra_scores": {}, "extra_text": {}}, {"review_id": "xch7XFPYfW", "reviewer": "Reviewer_UayJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the convergence of gradient on a restricted class of shallow transformers. The transformer architecture consists of the standard self-attention layer with a drastically simplified MLP layer with no learnable linear projections, and further, the class of models have the property that the model parameters do not deviate too much from initialization. The main result is an upper bound on the Rademacher process of this restricted class of transformers, and an upper bound on the empirical loss of 1-Lipschitz loss functions. \n\nIn my opinion, the results obtained in this paper are NOT very interesting. First, the assumptions on the transformer class are restrictive, in particular, 1. while the weights do not deviate from initialization is a key property in deep learning theory, but one should **prove** that gradient descent with proper step size indeed ensures the weights do not deviate too much, rather than assuming it directly, 2. in practice, MLP layer is quite important, this paper assumes the MLP layer with only ReLU activation, which largely restricts the complexity of the function class and the expressiveness of the model. Due to these simplifications, the proof is rather straightforward: one could simply assume initial bounds on the weights, which would directly lead to bounds on the weights at time $t$ due to the assumption, and the Rademacher process can then be upper bounded via Dudley's entropy integral. It adds little to no technical insights into proving the convergence of GD on transformers in general.", "claims_and_evidence_yn": "No", "claims_explanation": "Proofs are provided, yet the justifications for the assumptions on the transformer class include the weights are close to the initialization and simplified MLP layer are quite weak. While it is known in deep learning theory that by tuning down the step size, the parameters do not deviate too much from initialization, which was the key to prove convergence for classical ReLU networks, authors should formally prove this is indeed the case when running GD with proper step size, rather than simply assuming it.", "audience_yn": "Yes", "audience_explanation": "I think generalization for transformers is a very interesting problem, yet the restrictions of this paper are too limited to my opinion, this also limits its audience.", "requested_changes": "I believe in order for this paper to be accepted, it would significant strengthening of the results:\n\n1. Instead of simplifying the MLP layer, authors should include the learnable linear projection of MLP;\n\n2. Instead of assuming the weights do not deviate from initialization, authors should try to prove that gradient descent with proper step size indeed gives such condition. This would no doubt be much more difficult, authors should at least provide stronger and more convincing argument / evidence on **why** it makes sense to assume the weights will stay close to the initialization. As proven in the paper, in order for the generalization bound to hold one already requires the intermediate dimension $d_m\\geq \\Omega(N^3)$, such an overparametrization phenomenon is key in proving the convergence.", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)