Towards a Theoretical Understanding of In-context Learning: Stability and Non-I.I.D Generalisation
Authors: Yingjie Wang
Yutian Zhou
Shi Fu
Yuzhu Chen
Yongcheng Jing
Leszek Rutkowski
Dacheng Tao
Venue: ICLR 2026 Poster (2026)
Keywords: In-context Learning
Generalisation Error
Primary area: learning theory
TL;DR: This paper establishes generalisation bounds for Transformer-based models in in-context learning under non-i.i.d. scenarios.

Abstract: In-context learning (ICL) has demonstrated significant performance improvements in transformer-based large models. This study identifies two key factors influencing ICL generalisation under complex non-i.i.d. scenario: algorithmic stability and distributional discrepancy. First, we establish a stability bound for transformer-based models trained with mini-batch gradient descent, revealing how specific optimization configurations interact with the smoothness of the loss landscape to ensure the stability of non-linear Transformers. Next, we introduce a distribution-level discrepancy measure that highlights the importance of aligning the ICL prompt distribution with the training data distribution to achieve effective generalisation. Building on these  insights, we derive a generalisation error bound for ICL with asymptotic convergence guarantees, which further reveals that token-wise prediction errors accumulate over time and even lead to generalisation collapse if the prediction length is not properly constrained. Finally,  empirical evaluations are provided to validate our theoretical findings.

Published as a conference paper at ICLR 2026
TOWARDS A THEORETICAL UNDERSTANDING OF IN-
CONTEXT LEARNING: STABILITY AND NON-I.I. GEN-
ERALISATION
Yingjie Wang
Yutian Zhou
Shi Fu
Yuzhu Chen
Yongcheng Jing'
Leszek Rutkowski
Dacheng Tao
Generative AI Lab, College of Computing and Data Science, Nanyang Technological University
FFudan University
Uniieeriity of Science and Technology of China
"The AGH University of Krakow
dacheng. tao Hammalggmm.......mm com, apollo. gpollo.yt. yt......... zhou@gma 1. com shi 01 1 @e. ntu. edu. sg, cyzkrau@mail www.uustce.. en yongcheng. janeeenggiiiggie jing@ntu.edu. sg, rutkowski@agh.edu.
ABSTRACT
In-context learning (ICL) has demonstrated significant performance improvements in transformer-basee large models. This study identifies two key factors influencing ICL generalisation under complex non-i.idd, scenario: algorithmic stability and dis- tributional discrepancy. First, we establish a stability bound for transformer-batt based models trained with mini-batch gradient descent, revealing how specific optimiza- tion configurations interact with the smoothness of the loss landscape to ensure the stability of non-linear Transformers. Next, we introduce a distribution-levt discrepancy measure that highlights the importance of aligning the ICL prompt distribution with the training data distribution achieve effective generalisation. Building on these insights, derive generalisation error bound for ICL with asymptotic convergence guarantees, which further reveals that token-wint predic- tion errors accumulate prediction length not properly constrained. Finally, empirical evaluations are provided validate theoretical findings.
time and even lead generalisation collapse if the
INTRODUCTION
In recent years, the Al community has witnessed the emergence influential Large Models (LMs)
such as Generative Pretrained Transformers (GPTs) (Brown et al. 2020 Achiam al., 2023
Radford et al.] [2188] 2019), LLaMa (Touvron et al. 2023), and Pathways Language Model (PaLM)
(Chowdhery al. 2023). A particularly attractive characteristic of L.Ms is their in-context learning (ICL) capability, which enables effective predictions on downstream tasks using only a short context, without requiring any parameter fine-tuning (Black et al.) 2022).
Recently, the empirical success of ICL has attracted growing interest in theoretically analyzing its generalisation capability. [i al. (2023) establish optimization-indeeendenn generalisation bounds for ICL under i.i.d. inputs or trajectories derived from dynamical systems. Other works incorporate training dynamies and prompt structure into the analysis, examining how architectures and optimization strategies influence ICL performance (Huang et al.) 2024 Li et al.) (2024a] Chen al. 2024). Notably, Wu al.) 2024) establish statistical task complexity bound for the attention model pretraining and indicates pretrained model closely matches the optimally tuned ridge regression by achieving nearly Bayes optimal risk on unseen tasks. However, these studies rely on simplifying data assumptions that limit their applicability to real-world settings, such as the pairwise orthogonal token pattern imposed by (Huang al. 2024 [i et al. 2024a) and the independent token sampling assumption in (Chen et al.
2024
Wu et al. 2024).
"Corresponding author.

---

Published as a conference paper at ICLR 2026
Table I: Theoretical Contributions (V-has the given information, X-hasn't the given information)
Multi-Head Generalisation Optimization Distribution
No Special
Orthogonality
Multi-Layer Analysis
Dependent
Shift
Input Structure
Free
Li ct al. 2024b
Feng al. 2023
X
x
J
Chen ct al. 2024
V
X
Bai et al. 2024
X
X
X
Yag ct al. 20224b
X
Li et al.
22224a
X
X
Ours
Special input structure refers to prompts structured a specific format satisfy theoretical constraints. Orthogonality-Fee refers to data that is not constrained by orthogonal patterns in its generation
This paper relaxes these ideal assumptions and theoretically analyzes the generalization of nonlinear Transformers for next-token prediction in ICL, leveraging algorithmic stability (Bousquet & Elisseeff
2002, Charles & Papatliopoulos 2017]] Liu et al., (007)) and discrepancy measure (Kuznetsov &
Mohri 2015 2020; Wang et al. 2022). Our main theoretical contributions are:
Algorithmic Stability and Discrepancy Measure: Algorithmic stability ensures that small changes in training data do not cause large inference variations. We theoretically identify conditions under which Transformers achieve stability under mini-batch gradient descent and quantify discrepancy across different scenarios. Theorem reveals three key insights: 1) for sufficiently smooth loss landscape, algorithmic stability is well-controlleds and allows iteration number to scale polynomially with the training sample size; 2) in non smooth scenarios, stability deteriorates rapidly as iterations number increase, especially with small learning rate, making it advisable to limit iterations number to a logarithmic scale relative to the sample size; 3) regardless of whether the landscape is sufficiently smooth, an appropriately chosen step size can ensure that the convergence rate of algorithmic stability achieves O(N-1), where N denotes the sample size. The discrepancy measure captures distribution shift between training and target data. To quantify this discrepancy, Theorems 23 establish a on sequential Rademacher complexity for the non-i.i.d. setting.
asymptotically vanishing bound for the i.id. case, and bound based
Theorem 4 establishes the generalisation error of Transformer-basee models under ICL scheme by leveraging algorithmic stability and the discrepancy measure, revealing: 1) In the ideal i.i.d. data scenario, the ICL generalisation error achieves a convergence rate of ( with appropriately chosen iteration number and batch size, regardless of the loss landscape's smoothness; 2) In the non-i.i.d. data scenario, effective generalisation requires properly weighting training samples and suitable ICL prompting, particularly when the loss landscape exhibits insufficient smoothness; 3) The generalization сттor accumulates across the intermediate tokens generated by the model. Theorem 5 suggests that, to ensure effective generalisation, the length of next-token predictions should be constrained to grow at most logarithmically with the sample size.
Generalisation Bounds:
2 RELATED WORK
A major line of work investigates the approximation capabilities of ICL in solving diverse tasks, while another focuses their generalization and dynamic training behavior, aiming to establish theoretical guarantees for adaptation to unseen tasks under i.i.d. and distribution shift settings. In the research line of approximation analysis, [Aküüre et al.) (2023); Bai et al.) (224)) demonstrate that Transformers are expressive to conduct many machine learning algorithms in context, such as ridge regression and Lasso regression. Moreover, a series of studies prove the existence of Transformer architectures capable of implementing gradient-based methods and their variants when given appropriate prompts (Von Oswald et al.) 2023: Ahn et al., 2023 Ding et al. 2024). Within a self-training framework, Fu et (2025) establish the generalization error of Transformers with ICL. A particularly influential subclass of ICL prompts, Chain-of- Thought (CoT), has been extensively studied as structured form of in-context reasoning. Several works show that CoT-enhanced Transformers expressive than their standard counterparts (Feng et al.) 2023 Li et al., 2024c) |Merrll & Sabharwal 2023). Specifically, Malach (2024) prove that next-token predictors trained on CoT data can efficiently
strictly more
2

---

al. (2024a) provide theoretical evidence for the superiority of multi-head attention and standard Transformers over single-head and recurrent baselines in various reasoning settings. More recently, Gong et al. (2025) examine the emergence of ICL capabilities autoregressive next-token prediction models through PAC-Bayee theory. In addition, L et al. (2025) provide sample complexity and bounds for training Transformers to acquire CoT capabilities under a token orthogonality assumption. Recent theoretical studies have provided elegant geometric and optimization- -based explanations of in-context learning through structured concept representations. These works substantially deepen the mechanistic understanding of how semantic geometry and task-vector behavior emerge in transformer models (Bu et al. 2024; 2025). Our work tackles a complementary question. Without assuming generalization framework via Published as a conference paper at ICLR 2026
simulate any computabbe function, while Li et al. (2024b) show that Transformers learn multi-layer perceptrons in context.
In another research line, Huang al. (2024) explore the training dynamics and generalisation of
ICL on single-attention Transtormers. Huang et al. (2024) analyze the generalization properties of single-head attention Transformers, while Chen et al. (2024) study the gradient flow dynamics in multi-head architectures for multi-task linear regression. Further, (Ci et al.] (2024) and Yang
any latent concept geometry, develop distribution-shin algorithmic stability and discrepancy.
To clearly highlight our contributions, Table provides comparative analysis of existing theoretical works, emphasizing key differences in assumptions and results. Unlike prior works that rely on restrictive assumptions, such as orthogonal token patterns (Li et al.) 2024b Feng et al.) (023) or i.i.d. sampling (Chen et al. (024)), our analysis does not require idealized input structures and explicitly handles non-i.idd settings with distribution shift. This makes our generalization bounds applicable to a wider range of realistic ICL scenarios, including those where training and inference environments differ significantly.
PROBLEM SETUP
Suppose we have sample of size N. where the i-th sample variable is denoted as (X', C'), with X' representing the query variable and C' = Importantly, our theoretical results allow for these sample variables to follow distinct distributions.
representing the length-N, output sequence.
A typical length-Np ICL prompt consists of an example set D {(X*, C)) which is contex- tually associated with the pair (X*, C²), followed by a query input X We formally represent the prompt as P = !D' X'), where [D', X denotes the concatenation of the example set D' and the query input X' into single flattened input vector. In practice, we typically predict each intermediate token j = 1, ..... Ne. in an autoregressive manner, where the prompt for predicting the j-th token incorporates the token from the previous j 1 steps. Accordingly, we denote the integrated prompt for j-th token prediction as pid = [P', , Ci, ,i,,..,C-_1).
In practice, instead of relying on the correct intermediate tokens, the estimated intermediate tokens are more commonly used to predict the next-token. Under this more general scenario, we define pioj (pw-11 T(P))-1) with pioo P where is a Transformer-based model. However, this approach inevitably results in error accumulation. Appendix G establish the gap between the generalisation performance with p and pi.,, highlighting the impact of these accumulated errors. For convenient reference, Appendix A provides a summary of the notations used in this paper.
3.1 TRANSFORMERS ARCHITECTURE
This section introduces the widely adopted non- lon-linear Transformer architecture, which comprises self-attention mechanisms and a multi-layer perceptron (MLP) module.
Definition 1. (Multi-head Self-Attention Module) For any given length-N, prompt
z
2/
E N₂D
P
3

---

Published as a conference paper at ICLR 2026
Algorithm I Mini-batch Gradient Descent Optimizer for Transformer
Input: Observations S = {('', c') Initialization 0%, Max-lter Q, 0. Batch Size [B].
Output: 0 0= 00
For: S Q: -- q 1;
Stochastically Sampling C 09 = 00
To B EiEB VoL(T);
suppose that there are N₂ attention module A(:) RNNDD R,D. with parameters On DDxD and {(Vm, Qm, Km)} E DRD for each attention module 1,.... Na. The attention
score associated with i-th token (A(.))., NX XD R1xD is given by
N
A(P).: Σ
Σ
softmax 2m Vm
Om
where the softmax mapping is defined by
ext QmkKms
softmax (2. (mmmax(?m@mmmm))
N₂ ex Rm
/j=
The vector-based form can be derived easily:
A(P)
A(P)
E NRXD
softmax (PwaxxpPQmmma PQ Km) Pm
A(P)
Definition 2. (MLP Module) For any given matrix Z E N,D (token-wise) MLP layer with hidden dimension D is denoted as M(Z) ReLU(ZW,)W, E xDD where Wi, W2 RR DxD are parameters matrices.
Given any prompt P, we have the following inference process of l-layer Transformer
H' = TH = M² (A'(
1,... L,
where H' the output [-layer block of Transformer and H P. Consequently, the Transformer architecture with I. layers can be expressed as T(P) = Th TL-1 ... ... o T'(P). It is important to highlight that typical usage, only the last token from the final layer, denoted as T(P)... is utilized as the output corresponding to the queried response.
3.2 TRAINING WITH STOCHASTIC GRADIENT DESCENT
This paper considers a training process where each training example is aligned with the test setup. This learning scheme ensures that the model learns to mirror the inference process at test time. Furthermore, the empirical risk formulation employed
this work is also widely used in both theoretical analyses (Li al. 2024a Yang et al. [2024b), empirical studies from practical applications Miin et al. 2022), and dataset development Longpre et al. (2023).
Given N-size sample set S c'),, the training objective is formulated as:
N
N
L(T)
Σ
N.
where q;, i = 1,.., N, represent the weights for the training data, reflecting its relative importance in the overall optimization process.
4

---

Published as a conference paper at ICLR 2026
Ousput
Input
lpput Prompt
Input
Toaingat
Curput
Denue:
Training Dats
Training Data
Touniing Dott
Promot
Figure 1: Algorithmic stability quantifies the sensitivity of an algorithm to perturbations in the training data, where higher stability (i.e., lower sensitivity) typically indicates better generalisation. Discrepancy measures the divergence between the target distribution and the training data distribution, assessing how well the training data represents the target data.
Our goal predict unknown sequence C based on the given ICL prompt P The correspond- ing population version is
1 N. No
C(T)
Σ
E [(T(PN-IS-))
Additionally, the expected risk, which takes error accumulation into account, is expressed
1 No
C
E
CN+1 (T((PNNNNS)))...... (es,
N
For notational simplicity, we use = fO m. m. Ω Ki my WT. L.N.
2)1=1 to represent all trainable
parameters. Moreover, the training details using mini-batch GD is summarized in Algorithm| [
4 THEORETICAL ANALYSIS
We assume that Ts()) denotes a Transformer model trained using mini-batch GD on the dataset {p', e') The main concern is how bound the generalisation error in terms of the difference between the population risk and the empirical risk evaluated at Ts. We address this question using tools from algorithmic stability and distributional discrepancy, illustrated in Figure ]
4.1 MINI-BATCH GD- DE-PEEENDENT ALGORITHMIC STABILITY
This paper utilizes variant of commonly-usen uniform stability in statistical learning theory (Bous. quet & Elisseeff 2002). Let S denote the dataset S with i-th sample replaced by an independent sample drawn from the same distribution. The algorithmic stability is defined as below.
Definition 3. A randomized algorithm G that maps N -isize dataset S to estimator Ts has uniform stability the following inequality holds
N Eg (Tts C², -(Tss (phss1) C2) S3, LI., VS
To establish the upper bound on uniform algorithmic stability, we introduce the following assumptions. Assumption 1. [Boundedness] The norm of each row of the input prompt P and the norm of each response vector Cor j 1, 1,..., Ne, ....,, N,, uniformly bounded by constants Bp and

---

Published as a conference paper at ICLR 2026
Bc, respectively. Additionally. for any attention head m I..... Na and any layer1 = L, ...... L, the parameter norms satisfy the following conditions [|W|]l2 S Bw,, [| I Bw 11Q 12 S
-l.....
Bq, || S BQ. ||Vm S Bv. 110 < Bo-
This mild boundedness assumption is widely utilized in various theoretical studies (Bai et al. 2024 Zhang et all 2023). Indeed, the boundedness assumptions in our theoretical analysis can be further relaxed to unbounded settings, with the theoretical results still holding. For example, one can replace the assumption of a hard bound on inputs with light-tailed distribution assumption (e.g.. inputs or features have sub- Gaussian (ails)AAttia & Koren (2024). This means extremely large input values are exponentially unlikely, effectively limiting the influence of outliers without requiring an absolute bound. Under this assumption, we thus denote the maximum value of the loss function as Mr- = sup (.). Additionally, to establish the bound on algorithmic stability, we consider its Lipschitz constant with respect trainable parameters (Definition 4 and the Lipschitz smoothness constant ? (Definition 5. Detailed calculations for both are provided in Appendix H
Definition 4. (Lipschitz constunt) For a Lipschitz function f defined over domain X, the Lipschitz constant T.y is defined as the smallest value such that [ - [(x)l12 S T.yl1y - =112, Vx, E X. Definition 5. (Lipschitz smooth constant) A function f defined over domain X' is said to be Lipschitz smooth if there exists a constant y > such that [V ff(() ll2 S >/| y yll2 for all x, y X.
We then give the bound on the algorithmic stability (See Appendix forr for the detailed proof). TΙk- k a > 0. The algorithmic
Theorem 1. Let Assumption 7 be true and the learning rate be stability satisfies
BM BU-T BU
QITy Nya 21241
1+ 1-40
if S
BS
BMLL m7 Q Nya
if 1 1-1a(11-4
20
where Mf. Le, and 7 for Transformer are given in Equations 7- There constants are related to Transformer architecture, e.g., depth L and the number of attention head. For example, since quantities such as Mt grow exponentially in L (see Equation 7 sufficient condition for stability
is that the depth grows at most logarithmically with N.
Remark 1. The algorithmic stability bound depends on the Lipschitz smoothness constant ", batch size B, number of iterations Q, dataset size N, and learning rate decay a. For small 7, stability is better controlled, while for large 7, stability degrades rapidly with Q, especially when o is small. A larger datasel N improves stability, but increasing B the maximum loss M, worsens it. This aligns with existing studies indicating that small- batch SGD tends to yield superior generalisation performance compared to large-batch SGD or full-bateh GD (Keskar al. 2017 Masters & Luschi 2018 LeCun et al.) 2012 Wilson & Martinez 2003). To maintain stability, it is beneficial to use smaller batch sizes, moderate a, and smooth the loss function to keep y small.
The following corollaries examine its asymptotic behavior under two distinct scenarios, characterized by the smoothness of the loss landscape.
Corollary [Well-conditioned Smoothness] Let the conditions in Theorem be true, and Çi and Ça be arbitrary non-negative real numbers that control the growth rates of the batch size and iteration count,
respectively. If the loss landscape is sufficiently smooth, i.e., S (2 + v1 4a(1 a)), and
the upper bound Mt. Lipschitz, (smooth) constants Le and ry are bounded. By putting [B] = ((N5)
and Q = O(N5)) into Eq. 3), the upper bound on algorithmic stability is 3 O 147 -1).
Some techniques such as regularization can be used to ensure that the loss landscape is smooth. Corollary [ captures fundamental trade-off between optimization and stability. Increasing the number of iterations Q (and/or the batch size) generally improves optimization and reduces the empirical risk (which is observable). At the same time, stability analysis shows that larger Q amplifies the accumulated perturbations along the optimization path, thereby worsening the stability coefficient 3 and enlarging the generalization gap.
Corollary 2. [Insufficient Smoothness] Let the conditions in Theorem 1 be true. If > (1
Vi 4a(1 1 a)), by putting [B] O(N5:), Q OIn N) into Eq. (3) we get B = Ö(N5--1).

---

Published as a conference paper at ICLR 2026
Corollary 2 indicates that when the Lipschitz smoothness constant is overly large, constraining iteration growth to a logarithmic scale effectively mitigates instability.
4.2 DISCREPANCY MEASURE
Given the potential distribution shift between training and target data, a suitable metric that does not impose distributional assumptions is essential for quantifying their divergence. This paper extends discrepancy metric inspired by Kuznetsov & Mohri (2015) to make hypothesis-space independent.
Definition 6. (Discrepancy Measure) For the estimator Ts, the discrepancy measure is defined as
N Σ
N
disc(q) = Ne 1=1
EN+I5
Eis
1=1
where Eg, = [(((((((P (,) """l((""""""""""" )
The disc(q) measures the degree of misalignment between the target task distribution and the training distribution. We then show how this discrepancy can quantified under different scenarios.
L..d. Scenario: In the ideal i.i.d. case, where the training and target distributions match, the discrepancy admits the following asymptotic property (see Appendix B for proof).
Theorem 2. Let Ts be a learning algorithm that uniformly 3-stable. Suppose the training data and test sample are i.i.d. Then, with confidence at least 5, vō (). the discrepancy satisfies
disc(q) I 2/ //og(2/8) where B is defined in Eq equation 3 and thus disc(q) - as N + 0o provided that 3| N + 0.
The condition B - 0 is easy to satisfy under standard choices of the training weights. For example, if we take uniform weights 9; I/N for all samples, then ||q]] N N/1 and thus the requirement becomes simply 3
1/2). Theorem [ shows that such a decay rate for B is achievable under multiple concrete regimes. For instance, Corollary implies that o(N NN2-1/2 holds whenever ( + (1 + 7) < 1/2, where Çi and &2 characterize the growth rates of the batch size and the iteration count Q, and y is the Lipschitz-smootthee
parameter of the loss.
Non-i..d Scenario: If the target domain is entirely unrelated to the training domains, achieving accurate predictions becomes nearly impossible. Therefore, we consider a scenario where at least some training domains share meaningful relationship with the target domain. Formally, suppose that there exists an effective prompt such that the example distribution set drawn from distribution related to the training distributions, ensuring N. Ne [EN+IJ - E.er U E S c, > 0, where I C {1....,N N} is the index set that refers to the related training data, and 1, is the corresponding weight. Techniques such as incorporating more diverse training data and designing more effective ICL prompts can help reduce by better aligning the training and test environments. For this non-i.i.d scenario, Thcorem provides an sequential Rademacher complexity based upper bound on disc(q) (See Appendix F for detailed proof).
Theorem 3. Under the above situation and Assumption 7 with confidence at least - 8, there holds
N
N - ( + 3.M T log (Rx((
disc(q)
< + sup TEH
N. i= 1=1
+ Mella- - vial log 8
where the sequential Rademacher complexity Rx (fo T}) over measurable hypothesis space H. (see Definition for more details) satisfies RN ([e T}) ART. + N.Brlla vII. Loj is the Lipschitz constant given in Eq. 12), and R max Bc. (Bw, Bw, By Br
In the non-i.i.d. setting, Theorem reveals how the complexity of the hypothesis space involved in the second and third terms affects disc(q). For instance. a more complex hypothesis space, characterized by higher sequential Rademacher complexity, allows the model to fi arbitrary
Remark 2.
7

---

Ç1. Ço Published as a conference paper at ICLR 2026
patterns in the training prompts, increasing its sensitivity to distribution shift and thereby amplifying the discrepancy. It suggests that regularization techniques, such as weight norm constraints. may help control this complexity and thus improve alignment between training and testing distributions. In addition, the weight discrepancy ||q vll offers a theoretical explanation for the effectiveness of finetuning, which reweights training samples toward those relevant the target.
4.3 GENERALISATION ERROR ANALYSIS
Building on the above analysis, this section derives an upper bound on the generalisation errors L(Ts) ((Ts). The detailed proof is provided in Appendix D
Theorem 4. Under Assumption [ let Ts be a 8-stable learning algorithm and (q).... .N.. be any weight vector used in training objective. For any 8 O, each of the the following bounds holds
with confulence at least 1 - 6:
N
N₂
((Ts)
S
1
ΣΣ
q. (p")s( + disc(a) + ||||] 1 22llg|22MM 2 log
No
1=1
where B is defined in Eq. (3)
The following corollaries provide a more detailed characterization of the asymptotic behavior under both i.i.d. and non-i.i.d. data settings.
Corollary 3 (Asymptotic Behavior under i.i.d Scenarios). Let the conditions in Theorem 4 Aand i.i.d. assumption be true. Let qi N. |B] = O(NS:), and Q = ((N a) When the loss function scape
well-conditioned smoothness, with confidence least 6, 0 < 8 < 1, there holds
N N.
C(Ts) S N.N
((Ts(p )eay ct) + N 2 ==
when Çi + Gor 147 2- b) When Lipschitz smoothness constant is large such that 7 > (2xx) +
VI 4a(1 a)), by setting [B] O(N(:)( 1.2 and Q OOIn there holds
4
((Ts) S NON ==
ΣΣ ((Ts( ).. d)) N 2, 2222lgg
with at least confidence 1- 8. Remark 3. In the ideal i.i.d. setting, the corollary above establishes how the generalisation error bound achieves the fastest convergence rale of O(N-b) under different levels of loss landscape smoothness. Specifically, when the loss function is sufficiently smooth, the are tuned such that 2Ç1 + 2527
147 = 1. However, when the smoothness constant is large, exceeding the
threshold (2a)- / 4a(1 -a)), a)), to achieve the convergence rate O(N- the number of iterations recommended scale logarithmically with the sample size.
Scenario
Parameter Settings
Convergence Rate
Lidd & Smooth
B O(NG), Q O(N42), 26 m21=1
O(N-5)
Lidd & Non-smooth
B O(NG), Q = O(I N), S
O(N
Non-i.id & Smooth
q - + |1و O(N+).
O(NG), [ = O(NN.), O(N4),
B O(N()) Q -
Goz <1
O 0(ymar(6ss11 147 -1))
145
Non-iidd & Non-smooth
(R
B O(NG), O(In N₂ O(N9)
Table 2: Summary of Generalisation Error Bounds under Different Scenarios.
Corollary 4. [Asymptotic Behavior under Non-i.i.d Scenarios] Let the conditions in Theorem Hand
3 be true. a) If the loss landscape is sufficiently smooth and if lq ----va vll + ||q|| = O(N7s), then by setting [ (NN)) and O(N52), ), for any 6 > O, with confidence at least 8, there holds:
N
N q: (p
N
C(Ts) <
)a... ct) + sup
Σ (v (p
TER
mc 520 -1ca)
2 log
+
8

---

Published as a conference paper at ICLR 2026
b) llq 11-vii ( and ||q]] = (NN)), then by setting [B] O(N Q OIn N), and the ICI prompt length as N, = (22 for any 6 > O, with probability at least 6, there holds
((Ts)
No q: (T )a..s cb)
sup TEH
Iu - (.)))))((( )e.s
A 1.5.) 2 +e.
log
Remark 4. Corollary 4acaaryyyyo characterizes ICL generalization under non-i..d. settings by establishing two upper bounds under distinct smoothness conditions. The results show that smoother loss landscapes and better alignment between training and test prompt distributions (i.e., small l|q -vl) yield improved generalization.
Remark 5. From Corollary 4
to achieve better cross-domain generalization (i.e., minimizing L(Ts)). shall minimize the following optimization problem:
No
min- EE
(a + sup 2 (T
Σ (o (T.((())) + Aulis 9l12 + + 22119112
TER
where x1 and X2 are regularization parameters. The entire optimization procedure can be decom- posed into two stages. In the first stage, we solve for the optimal sample-weige vector q by optimizing the latter three terms. This subproblem can be computed via DC programming (Tao & An (998) or mradient-based methods. Once the optimal sample weights have been obtained, we then optimize the first term accordingly to learn the final model parameters.
In practical scenarios, the model typically uses its own estimated token to predict subsequent tokens. This approach, by its nature, leads to cumulative errors as inaccuracies in earlier steps propagate forward. The corresponding theoretical result and proof are provided in Appendix G
NUMERICAL EVALUATION
Our experimental setup follows Li et al. 2023). where all evaluations conducted using the same GPT-2 architecture implemented via the HuggingFace Transformers library (Wolf et al. (2020). consisting 12 layers and 8 attention heads. All empirical evaluations are conducted using NVIDIA H20 GPUs with 80GB of memory.
Evaluation on i.i.d data scenario: the ideal i.i.d. setting, we focus on validating the asymp- totic behavior predicted by Corollary | and the error accumulation characterized in Theorem 5 We consider a = 10-dimensional linear regression task, where each in-context example is of the form (p, c). For each sample i, given parameter vector 3 e Rd, generate length-L sequence using the recurrence relation of = -i +e, for | = I...., L, where the initial query of (0, 0.1 and the noise term ~ (0, 0. The prompt is constructed by concatenating two such examples along with the query input cf into a single flattened vector. Each parameter vector 3 E Rd is independently sampled from (V(0.,, 0.1La). We set the sample size N (50, 100, 200, 400, 800, 1600), use uniform training weights di = 1/N, and set the batch size to |B| // to ensure sufficient training. For evaluation, we independently generate 1000 i.i.d. test samples. We fix the number of optimization iterations to Q 200, set the learning rate decay exponent to a = 1, and systematically vary the sequence length N. E 21, 4, 5, 6, 7,8,9) 9}.
The Generalisation Error Convergence Analysis: We evaluate the generalization enror as the number of training samples N increases. Figure 2 demonstrates that the error decreases and asymptotically vanishes as N - Do. consistent with the theoretical prediction in Corollary for the i.i.d. setting. Results are shown for sequence lengths and 2; other lengths follow similar trends but are omitted due to large differences in error magnitude, which would obscure the overall pattern if plotted together.
The Error Accumulation Analysis: Figure ()b shows that the generalization error increases with sequence length, following approximately polynomial trend. In particular, once the sequence length exceeds a threshold near N, the error rises sharply. Moreover, this threshold shifts to larger values the sample size increases. These empirical findings support Theorem

---

Published as a conference paper at ICLR 2026
(1000 hort:
(a) Asymptotic behavior
bb) Error accumulation
Figure 2: The generalisation error under i.i.d scenario. For extended results under non-i.i.d. distribu- tions, additional sequence lengths and overfitting risk, please refer to Figures 4 and 5 in Appendix !
In addition to the i.i.d. scenario presented above, we also conduct evaluations under non-i.idd settings, which are detailed in Appendix [] These experiments are designed to assess the robustness of our theoretical claims, particularly under distribution shift conditions where training and test domains exhibit structural divergence. The results demonstrate consistent alignment with our theoretical bounds, especially regarding the influence of distributional discrepancy and prompt reweighting.
CONCLUSION
This study derives ICL generalisation eтror bound with asymptotic convergence analysis by exam- ining algorithmic stability under mini-batch GD and a distribution-leww discrepancy measure. Our results reveal how optimization settings interact with the smoothness of the loss landscape to ensure algorithmic stability, and how, when combined with high-quality prompts, they enable effective ICL. generalization. On the theoretical side, future work should develop tighter generalization bounds using techniques such as gradient stability. On the practical side, our findings inform algorithm design, including strategies like weighted training samples.
ETHICS STATEMENT
This paper is theoretical and does not involve human subjects, personally identifiable information, sensitive data. All proofs and theoretical analyses conducted under standard mathematical assumptions and are intended to advance the understanding of large models' generalization ability. We foresee ethical concerns with the content or potential applications of this work.
REPRODUCIBILITY STATEMENT
Our work is primarily theoretical, and all theoretical results are presented with formal statements, clearly defined assumptions, and complete proofs provided in Section ( and Appendix[] To support the practical relevance of our findings, we include empirical experiments conducted using publicly available open-source models, specifically, the GPT-2 architecture implemented via the HuggingFace Transformers library (Wolf et al. 2220)). All implementation details, including model configurations and data processing steps, are described in the Numerical Evaluations section. No proprictary data or code is used, and we are committed to making our results fully reproducible.
ACKNOWLEDGMENTS
This project is supported by the National Research Foundation, Singapore, under its NRF Professor- ship Award No. NRF-P2024-001.
10

---

Published as a conference paper at ICLR 2026
REFERENCE
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, al. Gpt-4 technical report. arXiv:2303.08774 2023.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. In Proceedings of the International Conference on Learning Representations, 2023.
Amit Attia and Tomer Koren. general reduction for high-probability analysis with general light- tailed distributions. arXiv preprint arXiv:2403.02287. 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, 2024.
Fan Bao, Guoqiang Wu, Chongxuan Li, Jun Zhu, and Zhang. Stability and generalization of bilevel programming in hyperparameter optimization. In Advances in Neural Information Processing Systems, 2021.
Sid Black,
Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang. et al. Gpt-neox-200 20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745 2022.
Olivier Bousquet and Andre Elisseeff. Stability and generalization. Journal of Machine Learning Research, 2(3):499 -9666226,,,6--- 2002.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, et al. Language models are few-shot learners. In Advances Neural Information Processing Systems, 2020.
Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Taiji Suzuki, Qingfu Zhang, and Hau-San Wong. Provably transformers harness multi-concept word semantics for efficient in-context learning. In Advances in Neural Information Processing Systems (NIPS). 2024.
Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Qingfu Zhang, Hau-San Wong, and Taiji Suzuki. Provable in-context vector arithmetic via retrieving task concepts. In International Conference on Machine Learning, 2025.
Zachary Charles and Dimitris S Papailiopoulos. Stability and generalization of learning algorithms that converge to global optima. In Proceedings of the International Conference on Machine learning. 2017.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. In The Thirty Seventh Annual Conference on Learning Theory, 2024.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Yingqian Cui, Jie Ren, Pengfei He, Jiliang Tang, and Yue Xing. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401. 2024.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning. In Proceedings of the International Conference on Learning Representations, 2024.
11

---

Published as a conference paper at ICLR 2026
Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. In Advances in Neural Information Processing Systems, 2023.
Shi Fu, Yingjie Wang. Yuzhu Chen, Xinmei Tian, and Dacheng Tao. A theoretical perspective:
How to prevent model collapse in self-consuming training loops. In International Conference Learning Reprexentations (ICLR), 2025.
Zixuan Gong, Xiaolin Hu, Huayi Tang, and Yong Liu. Towards auto-rea regressive next-token prediction: In-context learning emerges from generalization. In Proceedings of the International Conference on Learning Representations, 2025.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Proceeding of the International Conference on Learning Representations, 2024.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In Proceedings of the International Conference Learning Representations, 2017.
Vitaly Kuznetsov and Mehryar Mohri. Learning theory and algorithms for forecasting non- non- stationary time series. In Advances in neural information processing systems, 2015.
Vitaly Kuznetsov and Mehryar Mohri. Discrepancy-based theory and algorithms for forecasting non-stationary time series. Annals of Mathematics and Artificial Intelligence, 88:367-399, 2020.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus Robert Müller. Efficient backprop. In Neural Networks: Tricks of the Trade, pp. 9-48. Springer Verlag, 2012.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Chen. How do nonlinear transformers learn and generalize in in-context learning? In Proceedings of the International Conference Machine Learning, 2024a.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlincar transformers for chain-of-thought inference: A theoretical generalization analysis. In Proceeding of the International Conference on Learning Representations, 2025.
Yingcong Li, Muhammed Emrullah lidiz, Dimitris Papailiopoulos, and Samet Oymak. Transform- as algorithms: Generalization and stability in in-context learning. In Proceedings of the International Conference on Machine Learning, 2023.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting ohain-ofthhuught: Compositionality through in-context filtering and learning. In Advances in Neural Information Processing Systems, 2024b.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In Proceedings of the International Conference on Learning Representations, 2024c.
Tongliang Liu, Gábor Lugosi, Gergely Neu, and Dacheng Tao. Algorithmic stability and hypothesis complexity. In Proceedings of the International Conference on Machine Learning, 2017.
Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Yi Tay, Denny Zhou, Quoe Le, Barret Zoph, Jason Wei, and Adam Roberts. The flan collection: Designing data and methods for effective instruction tuning. In International Conference on Machine Learning, 2023.
Eran Malach. Auto-regressive next-token predictors are universal learners. In Proceedings of the International Conference on Learning Representations, 2024.
Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv:1804.0766 001.1100400662,. 2018.
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint 100.922310.07923, 2023.
12

---

Published as a conference paper at ICLR 2026
Sewon Min, M. Lewis, Hannaneh Hajishirzi, et al. Metaicl: Learning to learn in context. In Proceedings of the North American Chapter of the Association for Computational Linguistics, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving Language Under- standing by Generative Pre-Training. OpenAI, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alexander Rakhlin, Karthik Sridharan, and Ambuj Tewari. Sequential complexities and uniform martingale laws of large numbers. Probability Theory and Related Fields, 161(1):111-133, 2015.
Pham Dinh Tao and Le Thi Hoai An. A de optimization algorithm for solving the trust-region subproblem. SIAM Journal on Optimization, 8(2):476-505, 1998.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv:2302. 13971, 2023.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the International Conference on Machine Learning, 2023.
Yingjie Wang, Xianrui Zhong, Fengxiang He, Hong Chen, and Dacheng Tao. Huber additive models for non -tationti time series analysis. In Proceedings of the International Conference on Learning Representations, 2022.
D Randall Wilson and Tony R Martinez. The general inefficiency of batch training for gradient descent learning. Neural networks, 16(10):1429) 425-1441, 2003.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural System Demonstrations, pp. 38 28-45,77 2020.
Language Processing:
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In Proceedings of the International Conference on Learning Representations, 2024.
Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, and Liwei Wang. Do efficient transformers really save computation? In Proceedings of the International Conference on Machine Learning, 2024a.
Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408. /0147, 2024b.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. / 19420, 2023.
13

---

Published as a conference paper at ICLR 2026
Appendix
CONTENTS
Introduction
Related Work
3
Problem Setup
4
Theoretical Analysis
5
Numerical Evaluation
6
Conclusion
10
Reference
11
Contents
14
Notations
15
B Proof Sketch
15
Algorithmic Stability (Proof of Theorem 1)
16
Generalisation Error Bound (Proof of Theorem
19
E The Upper Bound Discrepancy Measure Under Li.d Scenario
20
The Upper Bound on Discrepancy Measure Under Non-i.i. Scenario
21
Error Accumulation Analysis (Proof of Theorem 5]
23
H Gradient, Hessian Matrix and Lipschitz (Smooth) Constant
25
Numerical Evaluations
31
The Use of Large Language Models (LLMs)
33
14

---

Published as a conference paper at ICLR 2026
A NOTATIONS
For clarity and ease of reference, Table presents comprehensive summary of the notations used throughout this paper. The input and ICL variable spaces for the i-th sample are denoted by X and C", respectively, while X and C" represent the corresponding input and ICL random variables. Their specific realizations are given by x² and c'. The dataset consists of N samples, with N₂ denoting the length of the prompt P. Ne the number of demonstration examples, and Ne the number of steps in the ICL inference process. The Transformer model employs N, self-attention heads, and the batch size in the mini-batch GD optimization scheme is denoted as BB. The ICL prompt for the i-th sample, containing Ne examples followed by a query, is represented by P while extends this by incorporating j additional reasoning steps. The estimated version of this prompt is given by P The parameters associated with the m-th attention module in the (-th layer of the Transformer are represented as O Vay. Q and K corresponding to the output, value. query, and key matrices, respectively, while Wi and W22 denote the parameters the MI.P in the 1-th layer. Finally, the empirical risk denoted by (()), while [(0) represents the expected risk associated with the ICL prompt P. and [ (0) denotes the expected risk when using the estimated ICL prompt P. accounting for potential deviations due to reasoning inaccuracies.
Table 3: Notations
Notations
Descriptions
X, C, X
the input and output variable space for i-th sample, respectively the input and output random variables for i-th sample, respectively the realizations of X and C for i-th sample, respectively the sample size the length-N,, prompt the size of demonstrations the length of inference the number of self-attention heads the batch size in Mini-Bate GD optimization scheme
r'.c
N
Np
N
N.
N
B
P
the i-th ICL prompt variable with N. examples followed by query
p&j
the i-th ICL prompt variable with N, examples followed by query and j tokens the i-th ICL prompt variable with N, examples followed by query and estimated tokens
pi,
O
represents the parameter associated with the m-th attention module in the I-th layer represents the parameter associated with the m-th attention module in the I-th layer represents the parameter assuciated with the m-th attention module in the I-th layer represents the parameter associated with the m-th attention module in the I-th layer represents the parameter associated with MLP in the /-th layer represents the parameter associated with MLP in the 1-th layer the loss function
Vi
Q
Kt
W
W
((-)
((0)
the empirical risk
(())
the expected risk associated with P the expected risk associated with prompt P
(EA(O)
B PROOF SKETCH
Our main results follow a structured sequence of steps. Below we summarize the logical chain of the analysis and involved technical tools.
Step 1: Formalizing stability and discrepancy. We begin by defining two key notions: (i) the algorithmic stability which measures how sensitive the mini-batch SGD-trained Transformer is to replacing a single training example; and (ii) the discrepancy measure disc(q), which quantifies the distributional mismatch between the seportance-wiightee training distribution and the target prompt distribution. These quantities jointly determine the generalization behavior of in-context prediction.
Step 2: Decomposing generalisation error. We establish the general decomposition (see the proof of Theorem 4 in Section D):
((Ts) L(Ts) disc(q) + , (vanishing statistical term),
15

---

Published as a conference paper at ICLR 2026
which shows that the generalization error consists of the training loss term, the distribution-sie term disc(q), the algorithmic stability 3 and vanishing statistical term. Hence, to obtain explicit bounds on ((Ts), must control both 3 and disc(q).
Step 3: Bounding the stability of Transformers under mini-batch SGD. To control 3. we analyze how perturbing a single training sample influences the multi-head, multi-layer Transformer during Q mini-batch SGD updates. Using the layer-wise Lipschitz and smoothness constants derived in Appendix H we derive recurrence relating the perturbed and unperturbed updates. Solving this recurrence yields the stability bounds in Theorem T
Step 4: Bounding discrepancy. We next characterize disc(() under both i.i.d (Theorem 2 and non-i.i.d scenarios (Theorem 3) by employing concentration inequalities together with the notion of Sequential Rademacher Complexity (see Definition in Appendix
Step 5: Combining the bounds. Substituting the stability bound (Step 3) and discrepancy bounds (Step 4) into the decomposition (Step 2) yields our final generalization results.
Finally, Figure
3 (see Appendix outlines the technical tool used for our theoretical analysis.
Appendix C
Appendix E&F
Inductive Probability Bounding (B.1 (B.))
Sequemial Rademacher Complesity (De 7)
Lipschitz Emoothrese Constranns (Section
Concentration Inqquiitios (eg.. Arim's Hoeflling Inoqulity)
Theorem
Theorems 2 &
Algorithmic Stability Bound
Discrepancy Measure Bound
Appendix
Martingale Difference Sequences Azuma's 's Inequality (1 (Lemma 1)
Theorem
Generalisation Bound
Appendix G
Inductive Bounding
Theorem
Error Accumulation Analysis
Figure 3: Proof sketch: the logical dependencies among stability, discrepancy, generalisation, and стror accumulation analyses.
C ALGORITHMIC STABILITY (PROOF OF THEOREM I
Building on the insights from the algorithmic stability bound for SGD under bilevel optimization (Bao et al. (2021), this section derives upper bound algorithmic stability of the Transformer model when trained with the mini-batch GD optimizer.
For any sample pairs (P&j, CJ), we give a general version of expected risk with weights vy, 1,... No, which is defined by
L(T) = Σ o ((TPP'--))
(6)
Let ( represent the optimization process after q iterations, following Algorithm | Define 6, ||0(a]) 8 (4))) [|r, where g(4) and 6' (q) are the respective outputs of T and T where S and S! differ by single data point.
16

---

Published as a conference paper at ICLR 2026
Under Assumption |
we establish the following bound:
")) - =I
Prob(5a0 = 0)T [ - C(T | 52 = =0
Probb 0) L(T > 0
Prob(500 = 0) Σ u
| ()) (pi-- ,, 03) - -((T ( (p- 1,)1, =0
Prob(6
0) Σ
B |E(T
(0) (P<S-))) c2) E(T ( (p--- 33)l =
1=1
By L (0 = =0] + Prob(óg,, > 0) MyMs
where Lgge is the Lipschitz constant of the loss function E with respect to 0 associated with (P*) Cý), and Mi is the upper bound of the loss function & associated with (Pij C2).
C.I BOUNDING PROBABILITY TERMS
If the optimization algorithm Tg (10) does not select the i-th sample within the first qu iterations, then
5yo 0. By induction, we obtain:
B 1--1- N-1
B
Bqo
Prob(6 = =0) =
2 N
C
N
Thus, we also have:
Bgo N
Prob(ō > >0)
Correspondingly, we have Prob(6gs > 0) S
Bgo N
As result,
N-c
N-c Bιo
" -
Σ
Us Le E 16 = =0
Σ M.M.
N
C.2 RECURSIVE BOUND ON 150 = 0]
Let v(1) Ny EiEn, VoL(T), and let o (4) be its counterpart using perturbed data. Denote
7 y= Min, U:7i by the Lipschitz smooth constant. The update rule in Algorithm giiess:
El 0] = Prob(1 E 16, 0, E B₂] + Prob(1 & B,) 60 0, 1 $ By]
= [[ (((9-1) 1) + nl 1 19-1) - -11-_1)111600 =0]
N
N B E |1000-1) 1) -g'(1- 1) + mT 7 (1' (1- - v/ 111 = o]
+
< Cg- E 16g0 =0] + Dg-1,
where
B (N B1( + n4 14-17)
2 L
C
D
N
N
By induction, we obtain:
16 0 S Σ
Dj
II Ck-
j=
k=
17

---

Published as a conference paper at ICLR 2026
C.3 BOUNDING II C
Since
C (1
using the inequality + S et, we obtain
II l=
C2 I exp
(1 (1-B/NY
Σ ?
Amj+1
Thus,
9-1
66 =0 S
Σ j=40
Dy exp
(1 B/Nyr
Σ
M
Aen++1
C.4 FINAL BOUND ON B
Combining the above results, we get
No m
R
-
Σ exexp 90
N B N
Σ
TKk
Σ
Ta. +
Bgo N
Mi,-
k=j+1
Denote by L E == Ui and Mj = EE ei M,i, where Li,e, Mi,e and y are obtained in Finally, optimizing to leads to the stability bound:
Section H.9
Mr Bgo
8 x
min T0 ()...QQ
L
D3
exp
?
Σ
j=
k=j-1
We denote the original objective function by
Q Σ
2n L
Mr
H(4)) = Le
exp
N
N
+=
k=j+1
B
Σ
Q
1
Qa
)7
Mr
as x
N
ju
jo
N
j=
B0
Q
M,
)
+
N
ja
N
11
BQ
Q1
11a---a
90
1-07-a
Mr Bgo
N
1- 07 a
N
2L²BQL-a
2L3BQar
1-0r--
Mgo
N(1 -0y 0.)
N(1 a2
(.)
N
2LBB
(Q'-0 - (Q--a-Qu0) gu -a7-aa
Mg
N(1 07 aa)
N
The goal is to minimize (go), cnsuring that:
I N(1 -a)aa))) -Q'-a-Quu0l 1- +
2L3B
M, Bqo
8 min
N
dH Setting Ago
0, obtain:
2(1 - - -
M,B
90
-ay & +
N = 0.
N(1 - 07 a)
18

---

Published as a conference paper at ICLR 2026
For optimal selection of q², using 7)s ba approximate:
22
aii
Mt
Finally, by setting w; = 1 No the upper bound on stability is
B - H(4)) N(1 07 a7-a)
BLBB
Q'-a-00
big
2ul1 7) al+1)
M2
all
Mj बोंन कों Q=+=N-1
BL
a(1+))-1 2 22x(2,022+100 a0+t) Lg
x
L Q
117
We finally derive the desirable result
a(1-)--1 RM, a(ll)) rulll アলन
if y < 1+ 1+ 1a(1 1-1(1
>> 0,
B B
al1 Nya BM, a(1tT)
L, t2
if 1-14(1--00
10> 0.
Nyaa
2a
D GENERALISATION ERROR BOUND (PROOF OF THEOREM
This proof closely follows the approach of Theorem (Kuznetsov Mohri 2015). with the key distinction that we extend the analysis to weighted average version. For the sake of completeness, we present it here, beginning with an essential concentration inequality.
Lemma 1. (Azuma's 's Inequality) Suppose (Yu..., ,Y,) is a martingale difference sequence with respect to the filtration Fo C F C ... C CFN. If
0x S [[(++/ S b44,, VO
then the following probability bound holds:
2e
Prob Σ
M 26
5 2exp
Emm(b (b: a;)²
For notational simplicity, we denote, for the i-th sample, the random variables as Z' = (Pi, C*), Zj (pij-1 Cj), and the sequence Zj:m (Z) ...,, 2
..m)....).............. We define S(i) as the sequence set
(Z) Z Z +1 ,,,,.
where ž' is independently drawn from the same distribution of Z Now, consider the following quantities:
A3
E EøjiIN 2311
((Tss (1 Ezits
[((Ts(): 1
and
BJ Enj [e(T ((T)())))> Zz 11=1
[(Ts, Zz
where 2,+1 ~ ((-1z) independent of 7 and 7 By construction, we observe that:
Byt g 2.01 [A}] 0,
and
E Bztt 2 [Bj] 0.
These equations indicate that both sequences A 1,... No and B I.... .N. form martingale difference sequences. By applying Azuma's 's Inequality (Lemma I). for any > with probability least 1 5/2, we obtain:
--1
4 8' 9 = 119 2
Σ
19

---

Published as a conference paper at ICLR 2026
and
N-1
log 4 8
Σ
9 S |
where M,,t is the upper bound of the loss function associated with input P.. Summing both inequalities, we obtain:
(A) + Bj) s 2 2 8
Next, we define the weighted sequences:
N A' Σ GA4.
B' = Σ cj
where Σラ= C 1. Since these sequences also form martingale difference sequences, we apply the definition of uniform stability, which states that:
N₂ Σ ==1
G E3
Sn
where
Eg 12 [(T 25t²))
-Byz-1 [((TSO) )(2b
Thus, with probability at least 8, we obtain:
4 Σ 9 + B*) s 2 2 2log 8'
where M₁ Em w; M, By the definition of algorithmic stability, it follows that:
N
4 (A' + B*) < 2 2
Σ
and
N 9: cy B2g,, N,aaN,22(( 1T(( 2,)1,
2 1991118 + 2119ll2 2 8" Cye(Ts, ZJ)
i=1
Finally, using the definition of discrepancy, we arrive at the final bound:
N G 2z) + disc (q) + +| 2/1 2 log 8
2
((Ts) S
By taking wa and combining the upper bound on , obtain the desirable result.
E THE UPPER BOUND ON DISCREPANCY MEASURE UNDER I.L.D SCENARIO
Lemma (Asymptotic Vanishing of Discrepancy). Lei Ts be a learning algorithm that is uniformly 3-stable. Suppose the training data and test sample are ii.dd, and ||q]] 1. Then, the discrepancy
term satisfies [disc(q)] < 2ll (log(2/6) and thus disc(q) - 0 as N - oo provided that
8x 0.
Proof. Fix token index j E {l,..., N.) and define
N
Dy Dy:=E(((T, ZN+ 1
Σ q Eie(Ts, 23) 12=-1
20

---

Published as a conference paper at ICLR 2026
We analyze each summand
A = Bl[Tss 2M 1-B(Ts, 2j) Zbi-1].
Introduce an intermediate model (( Ts(s) trained on S S [Z Decompose A, as:
A B(((Ts, 1 Eie(Ts: 2j) Z [ 2j) | B[e(Ts, 2j) z
(A)
(B)
Since Z3 ~ 2 are i.i.d. and independent of Ts once 71:N is fixed, we have:
(C(TT,,2 12 B[ 2j) ZI:N1.
Hence, =A)=0. 0.
To control the second term, define a filtration Fi == o ,..., Z and define Doob martingale:
X₁ == [Ts 23) Z ..... Z²) = ...... N.
Since Ts is 8-uniformly stable, replacing one sample Zt changes the expected loss by at most B. So:
X, X S B.
Then, Azuma-Hoeffding implies that with probability at least 1 - 8:
B[((Ts, Zc,) Z INN E[e(ss Zi,j) 1 s 3 2( - i 1) 12g(2))).
Using |///22 < and Jensen's inequality:
N
N 822(N - i - log((2)) < 23]/a]| /log(2/0)
|D < q A S ||al2
Finally, averaging over j:
1 No
(disc(a))
S
Σ
Ne 1=1
We complete the proof.
F THE UPPER BOUND ON DISCREPANCY MEASURE UNDER NON-L.LD SCENARIO
We firstly introduce the definition of Sequential Rademacher utilized in (Rakhlin et al., 2015 Kuznetsov & Mohri 2015 2020).
Definition 7. [Sequential Rademacher Complexity] Let (o1.....o7) be sequence of Rademacher random variables (each o; independently taking values =1 with equal probability), and let . Qr) E IR" be a given weight vector. For function class defined sequential data 21, .., 2T, the sequential Rademacher complexity is
......,,1"
Rea (G)
sup EE sup Σ oe q g(((())) GEG t=1
where the supremum is over all complete (depth-N) binary trees adversarial sequences 2((1i....,0t. 0t-1)
In simpler terms, R20 (G) measures how well & can fit random signs {o,} in an online or sequential manner.
21

---

Published as a conference paper at ICLR 2026
F.1 FUNCTION CLASS: TRANSFORMER HYPOTHESIS SPACE
We fix a Transformer architecture (with Na heads per layer, hidden dimension D, and L layers), and let 0 collect all parameters (Qm, Km, Vn. Om: Wi, W across L layers. Denote the overall parameter space by θ, and suppose we have a norm constraint |0| I A, bounding all weight matrices operator norm (or some suitable layerwise norm). Let FTrans be the function class:
Frrans
{ P Jo:P-T(P) T(P)
0 E |0| s
For sequential inputs P ot), this means the Transformer is invoked on each partial prompt Pi.--
F.2
REWRITING THE LOSS A COMPOSITE FUNCTION.
Let us set
g(P) T(P)n. RD,
and define function
((x,y) = -113-
where y E P. Then
Then
((T) dg(P), Y) [IT(P) Y12.
Hence the loss class G = {{(T)) is precisely 6o G where
G = {P( T(P) T E H
We check how o(x,y) = |x - 2 depends on x. Let R = max Bc, max H1
}1}}} be a constant.
Suppose ||x|| S R and ||y|| S R for all feasible (x, We can show that ((x, ) -ip-hiit in with
Ls S R,
because
x -5-yy y11 - 1x 12| S 4R - xx||.
Thus. if the model outputs x = P and targets Y remain within ball of radius RL then v(:y) is -ips-Lipshhiz in the first coordinate.
F.3 SEQUENTIAL RADEMACHER COMPLEXITY OF TRANSFORMERS
Inside the expectation [.]. the random variables {ar} are independent Rademacher signs. Let us write:
N
N
E sup JEF 1=1
Σ
[())]]
< 4RL; E,
Σ
Oe
Thus, remains to bound E. [II ΣN of 2 1l] A typical assumption in bounding Rademacher-
based complexities that each 2; has finite norm || S VnBp. Then according to the fact that < t then
E Ex(0:00000 E. -o (o.4::0222)) = O,
have the following:
E
Oe
qt
24
E
1 t=1
oe q 24
17]
Σ of |1 2 S VnBr
Putting all these pieces together:
Ry (1(T))) s 4RL
This shows that under constraints and bounded inputs/targets. the sequential Rademacher complexity of the squared-2z loss class is finite and depends primarily on the Lipschitz constant of the loss w.r.t. the model's ''s output, as well as the base complexity of the Transformer itself.
22

---

Published as a conference paper at ICLR 2026
While the above bound may appear loose (e.g. exponential in the number of layers L), nonetheless demonstrates qualitatively that the capacity of squared-lz losses controlled by parameter norms, data magnitude R, sequence length T, and any submultiplicative structure in the Transformer layers.
Remark.
1 No For given hypothesis space and define by disc(q) == SuPTEH N. Eド= EN- Eis where E E ((TT(P)))))
By further combining the fact disc(q) dísc(q) with the following lemma Kuznetsov & Mohri 2020), we obtain the final result.
Lemma 3. For any with probability at least 1 for all f E F and all > 0, we have
N
N
Σ (l) 1
Σ q: f(z.) + ||ql22 + 6 M.v log T RT + Me Mellalley log
t=1
Combining this lemma with above results, we complete the proof.
G ERROR ACCUMULATION ANALYSIS (PROOF OF THEOREM 5)
In practical scenarios, the model typically uses its estimated token to predict subsequent tokens. This approach leads to cumulative errors as inaccuracies steps propagate forward. Denote (T) = ((T(P C?)) the population risk at i-th step prediction such that L(T) = N. Σm (T) We have the following relation between (T) and L(T).
Theorem 5. (Error Accumulation Analysis) Let L(T) and (T) be defined in Eqs Assume the conditions in Theorem 4 hold. For any 0 < 1, we have
N
N N. (1+
EA (BA( S No
LN.
+ 8-=1
II
(1
Σ
4 E 1
N
k=
where 81) is the Kronecker delta.
It is evident that L (defined in Eq. 12 increases with the number of layers L, and polynomially with the prompt length, as well linearly with the model size parameters N, and D. Based on this observation, we focus on analyzing the impact of inference length on the generalization error bound.
Corollary 5. |Generalisation under i.i Scenario] Let [B| ((NO) ((NG)).G 1 Q ((n and
No O(I ())a). With at least confidence 1 - ô, there holds
N
4
C (AA(Ts) (Ts)
am, an, (Ts((TTs(')))))) +) tby 2 2oo
where 1b = (log nı) Sa(log 225222 N-1/2 and the weights Tj I,..., No equal to
N.
N N
L N
Σ (1+ ( 4
8-=1 Lir
II (1+
=
This result suggests, in scenarios where error accumulation occurs, the length of the inference process should be constrained to a logarithmic scale relative to the sample size to ensure effective generalisation. Notably, this finding aligns with (Merrill & Sabharwal 2023), which states that transformers with logarithmic number intermediate tokens may exhibit enhanced computational power. Similar results can be easily extended the non-i.i.d. setting.
23

---

Published as a conference paper at ICLR 2026
Proof. According to the definitions of [(0) and (R), we have the follows
1
EA 1
(0) ((0))
No (()
No
No (P )..: -
1),,, 13
BIIT(PNN+J-1) - -(P.N+1--1)..11
N
According to the Lipschitz property of Transformer, we have
N
-
N
L
No
L
(0))
No
i=1 j=1
For notational simplicity, we denote ARA
-E CEA(() and A = Ej=1 ()) Then we have
L Σ A
A < Am
t=1
We denote by S, 3m=111 A.A.A Since
Li SS + An
Sn = =5= + EA
S (1 +
n
+ Σ II (1+ 4
1=2
k=
<
Σ
j== II
5 Σ
I1+ 4
c..
cz.
1=1
1=2
1
Thus, we have
L1 -
(FA(O) < (()) +
Σ
A (0)
Ne
i=1
1 No
L1 N--1
N II
N
<
C, +
81=1
(1
)
N
N₂ j=l
2=
L No-1
No-1
N N
1
CN.
+
1
+
8-=1
II
(1
N
Ne
L+
j=1
This completes the proof.
We then simplify these weights. Using the logarithmic approximation, there holds
N--1
N
II
1
Li ~
exp
(
Σ
< N
i
Similarly,
24

---

Published as a conference paper at ICLR 2026
N--1
Li
N₂
1/7
II
i
i+1
Ami+1
Approximating the summation, we have
No
Ne
47
No
N
44
Σ 1=/+1
dx.
i+1
j
Evaluating the integral yields
NLT
(NN - - 1-tز
NeT
1+ j
T
Lig
Thus, the dominant term in the simplified bound is:
O (
By combining above results with the generalisation bound established above, we obtain Corollary
5
H GRADIENT, HESSIAN MATRIX AND LIPSCHITZ (SMO0TH) CONSTANT
This section gives the gradient and hessian of Transformer models. Note that, n, N, S n S No No in this section refers to the length of input prompt.
H.1 GRADIENT W.R.T. HA, AND ITS NORM UPPER BOUND
It is easy to obtain the gradient of loss w.r.t. final output H i.e.
OC
2(H y).
dH.
To establish the upped bound on |l lae ll2, we need to bound the upper bound on the output |
The definition of Transformer models yields
H M (A"((((
and
(A'
1):11 ()).))
Bw, Bwg sup |A H 112.
We have
N₂ = Σ
(IA mm.
softmax (H- Qu KL (HL-
H4-1 0 ll2
m
m=l
n
< Bo
H 1-11 112
Jji
s By BoNa
sup j=1....
H
1|2-
Combined with above result, we obtain
sup | < Bw, Bw, By BoN
H S (Bw, Bw, By N. sup H 22-
sup i=
Under Assumption D we have
sup |2 y)lla & (2 Bw, By Bo LBp Bp +
/2Bc =: Ct
Similarly, the maximum of loss function can be bounded by
Mf 2( (Bw, Bw, By Bo N. 24 Br + 2B2
25

---

Published as a conference paper at ICLR 2026
H.2 GRADIENT W.R.T. W/ AND ITS NORM UPPER BOUND For any / = =1, -1,... L, the gradient w.r.t. W2, we get
dc =
0C OH
OH 0H
HH aw
aw
OH
OH
According to the definition of M and A, we have
OH ReLU((((H"
aw?
Recalling the Lipschitz property of Transformer, we have
HHe R. 0H
öH le I L 4
CLe
We then can bound the gradient w.r.t W by
dC
|r & CrCIA Bw, (H
Ow
where
||4 (H''
1 Vn sup
(AA(H'
i
I VoBy
sup
| S VTB (Bw, Bp Cw
H.3 GRADIENT W.R.T. W AND ITS NORM UPPER BOUND
We next to give the upper bound on the norm of the gradient w.r.t. W Similarly, for any / = 1,.... L. we have
Oc
OC
HHL
0H
H ow
awe
OH
OH
and
0H
0. (H'
R
W R A (
ow!
where
R (A4 (H' WW)
0, otherwise
Ríj The upper bound is
ac
aw ||f S CCCCL Bw. RR' A' Bw (
H.4 GRADIENT W.R.T. Q AND ITS NORM UPPER BOUND
For any 1= 1= L, the gradient w.r.t. Q is
OL
dL OHL
HH
0H
0H
IQ
0H
aQ
and
diag(R))
diag(
0H
0H' DA'(H
(H'
aQ_ m
9AA (H (H)
aQu
aQu
diag(R' W
where
(HA(HH
softmax' (H' 1 kt H l
aQ
and
softmax's (l)tmax(Z)([- -sottaa )
26

---

Published as a conference paper at ICLR 2026
Note that Z = H '
Kn "
(H" and
Vn
||oftmax' 1I S
2
Then the corresponding upper bound is
OL
0A'(H''-1)
||r
s C.CLAVIBw, Bwg
||5
aQ m
iQ
s CrCrLLBW, Bwg Bx By Bo |
|
s CLCCLLn Bwi Bw BR B 1-1
3-1 2-11
21
341
Bo
N₂
2 Bp.
H.5 GRADIENT W.R.T. Ky AND ITS NORM UPPER BOUND
Similarly, the corresponding upper bound of the nomn of the gradient w.r.t. K is
OH
0A'
x C.CLLVnBW, Bwa
OK?
OK!
S
CLOIEnDBWr
91-1 2-1
Bw
2
BQ By
36-1
Bo = No 2 B
36-1
3:--1 11
H.6 GRADIENT W.R.T. VL AND ITS NORM UPPER BOUND
For any 1 = 1,.... L, the gradient w.r.t.
VL is
Oc
dC HHL
OH
0H 0H
0H 0V?
OV
L
R,
m/
and
(iag(R)))/wy)
OH
0H
(H'HH
(A'
TV!
-H
diag R!W)
where
0 (A' = softmax (H' Qt (H'- H 1100!
OV!
Then we can bound 1l
|| by
OC
(H'-1
||
< CrC.I.VTBB Bw,
OVT
QvL m
T
s CrCLLBWW, Bw. Bo -
< CrCLA Bw, Bw, Bo
H |
C, Bly, Blv By 1
H.7 GRADIENT W.R.T. O AND ITS NORM UPPER BOUND
Similarly, there also holds
ac
0L
HHL
OH
8H
n.:
80
TH
-
OH'
80!
and
diag(R\)wi/g
0H
(A'H-'1)
(A
10 m
0.4'
(H
00
00
diag(R
27

---

Published as a conference paper at ICLR 2026
where
0.A' 0 (H' softmax (H'
(H' (H'-1K)))) H H
80
Then we can bound
||
0H' 8oo
|| by
0c ||
(A
||r
& CLCLLVBBW, Bw 80
a0 mn
CL Bw, By H ( 11p
Cr.CEi Bw, Bw, By H H 1|
CrCran& Blv, Blv Bl, B/ N Bp.
H.8 HESSIAN MATRIX
We firstly calculate the Hessian of Transformer T' w.r.t H'-1 and its upper bound.
For 1-th layer, the Hessian matrix is
Hr = - ")w/) = (A'
where
N
V2 Hi
1
V2 2m
softmax (7m) (VH->1 7m) V.l Or + V z.,, (2 Vi 7m V.l O
== 1
and
Zn = H 2' (H' (H''))
The norm of each component in Lipschitz smooth constant for T w.r. H is
[cbng(ReLU((((((())))))))) vnD,
- (H' (H Cooftmaz N₂ By Bo Bo H 11 H +
where a conservative bound Csoftmar D2/8. Then we have
[Halle s
VnDi N.D Bo Bw, [| 2).
For W and W3,, their Lipschitz smooth constants are O. For On we have the Hessian matrix
Ho:
Vg ReLU(A(((I--)w///wy
A(H''-)W(WW
where
No
A'
Σ m=1
V',OO, V2 softmax(2m) (Vo., Zm) VL 0 + Vz, softmax V2 Zm Vil O
and
|| Φr A'l By | 1|-
For Ki we similarly have
V A < 2 By
By Bo B3 H 1
For Vl and O the Lipschitz smooth constants 0.
28

---

Published as a conference paper at ICLR 2026
H.9 LIPSCHITZ CONSTANT Le, LIPSCHITZ SMOOTH CONSTANT Y AND MAXIMUM OF LOSS FUNCTION Mr
According to the results in Appendix H the upper bound on Le, Mf and y are
Mr 2(Bw, By By BoN.) (ap)))BBBp. 2Bg.
(7)
(N
Le
Σ
Un
E(On, + Cw. - No + Ck CV + C))
(8)
N,+11
1=1
where
Gvv = CLCLLB VnBy By No (Bw, Bwa) 111 1 Bp
Gh, C2CrLBW VnBy Bo (Bw, Bw,)' 1oB-B-BBBB Bp
11-1 #-1
#-1 x-1 9(1-1)
0/ CrCuin Bwi Bws Bx By
2
Bo No Bp.
8-1 2111
sh1
2-1 30-1)
C₂ = Cr.Ciin Bw;
Bw 2
BoBy²
Bo 2
Na B₁
Cy = Cn.Can Blv Bl By BB
Bp.
C CrCann Bl, B By B No-1 Bp.
The Lipschitz smooth constant is
(N++1)N.
Σ
Na(1l +7k)
7
On
N
where
(VTDDN.By Bo Bg BR Bw, Bw, | [[[ (|HHH 4-1
N. D2 N2D²By By Bo By B 1|
and
nk (VnDDN..V BR Bo Bw, Bw, H || [ -1
H N₂D² Bo | 113
where
- S n (Bw, Bw, By BoNa)' Bp.
Putting the above results into the upper bound of stability will obtain the desirable results.
H. LIPSCHITZ CONSTANT W.R.T INPUT LT
In a Transformer, the ReLU activation function is piecewise linear and thus non-differentiable at certain points. In particular, the concept of a Jacobian, defined in terms of the network's outputs relative to its inputs, indicates how those outputs vary with small changes in the inputs. The Jacobian at a point x computed via the chain rule during backpropagation. However, it is only well-defined all ReLU nodes are differentiable at that point, meaning their inputs must be strictly positive or strictly negative. Consequently, if an input equals zero, one must assume the existence of aub-gradient within [0, 1].
According to the chain rule, the Jacobian at a point p (namely H if defined, can be compactly represented as:
Jp [T] [p T4-1 ............ T4 JH (T4] JHH [
29

---

Published as a conference paper at ICLR 2026
To obtain J/r-- 1,.., L, need to calculate [1M" and Jn--i[4),, respectivell Since both M and A' map from IR xDD D to HDDxD, their Jacobian matrices have the same form
J
J
JiN, ERNDDDN,,DDDD
J/
J2
M2N,
(10)
JN
JNz2
JNN,NN
We firstly give the Jacobian matrix for JA (M We denote by A A'HH-11 E RNxDD, Recall the definition of mapping
MF(A)
(A
MM = RcLU(A'W/)
ERXxD
(
where M.(AA) = ReLU(A,/ By taking partial derivatives, for any i,j E [Npp] we have
aM,(A)) (Re.U((,WW,WZ = 6, WTTT
where G encodes the activation pattern of a layer caused by the input x, and 6,j is the Kronecker delta. The matrix G is diagonal matrix, having as elements if the corresponding neuron is active, otherwise Os for inactive neurons. The Jacobian is the same for all the points strictly inside linear region with the same activation pattern. Since ReLU networks are piece-wist wise linear in nature, the Lipschitz constant is exactly equal to the p-norm of the Jacobian one such linear region in the input domain. Thus, the jacobian matrix for M a diagonal block matrix, having W W G as elements if j. otherwise Dxx for i + j. then have
sup |A [M']l2 S N2 DB/, Bw2- A
In fact, there have been some studies analyzed JEp-1/ [A4). l,.... L and the Lipschitz constant of attention. Since transformer T is map from TR N2 to R the element of Jacobian is Jl
(H' The Jacobian of the softmax also well-known. Suppose that softmax(u)
8H
R N.x1 Then we have
diag(v) - TV
Ou
Recall the definition of mapping
Af (HH-))
N.
AH
(HHH-11
softmax
(H'-'Q', Ky (H (
E NRND
(H
where
N
A.
Σ
(
and
Mj = softmax QQ' K (H"-))
For any 1 = 1,... L, by taking partial derivatives we obtain that
No
J/
Σ
My (VL))
vd o
0M, OH!
No
N
H. Vi' o M
HE. (K'h) (8.)))
30

---

Published as a conference paper at ICLR 2026
where M' := d - M, Mr. with sup ||||| || s 2. and Eig R RNNNN is binary matrix with zeros everywhere except the (i.j) entry.
Under assumption for any = ..... L, we then have
sup |لأين|| || & BospN H
and
sup || [A''ll s sup | Co HH'
where Lip is a positive constant.
Then the Lipschitz constant of layer Transformer is
L Bq Bp II
Ly = BB
N Biv Bw By Byy
(11)
Specifically, if the output Transformer model is assumed to take the last token at L-layer, its Lipschitz constant is
2 LT.
Li =N=
(12)
This completes the proof.
I NUMERICAL EVALUATIONS
Evaluation on non-i.id data scenario: In the non-i.id scenario, besides the asymptotic behavior and error accumulation, we additionally validate the impact of distributional alignment (quantified by l|q - vlI) in Corollary 4 We consider a complex scenario where the training and test tasks are drawn from related but non con-identical distributions. Specifically, for each sample i, given parameter vector 9r, we similarly generate length-L sequence via the recurrence relation of - + e for 1= d,....., L. where cj ~O, La) is the query and (O, 0.1 Ia). Different from i.idd scenario, in this scenario, the training parameter 3 is drawn from a mixture of two Gaussian distributions Puzin (x) p1 (1, Lạ) P2 (0.2, 0.1 Ia), where P1 and P2 the weights such that + P2 = 1. The test data is drawn from a different distribution: (r) = (V0.1, 0.1 Ia). We consider the prompting format that incorporates single in-context example. For training samples, the in-context example is drawn from the same distribution as the query task, i.e., both query and support samples share the same 8 For test samples. the in-context example instead drawn from (0.2, 0.1 Ia). regardless of the test query's s distribution. This design aims to reduce the distributional divergence between the support and query examples the test setting.
Under this setting, we assign group-wise importance weights based on the product of pairwise overlaps between the training and test component distributions:
((k) Overlap(-V(uu,, (0..1, 0.1 Ia))
12),, -=) 1,2
where k- 1, 2 means the two class distributions of prompts in training process, and the total overlap between two distributions p(z) and q(x) is defined as Overlap(p, q) = - min(p(z), q(x)} } dz.
This value lies the interval [o, 1). where indicates complete distributional alignment and 0 denotes no overlap. We then can approximate the optimal weights R. as:
G()
GO)
(()
(2)
ERN
==
PIN
PIN
P2N
P2N
N
P2N
where ((k) is re-nomalized constant such that Ee ((k) 1, pı and p2 are the proportions of the two training components, and N is the total number of training samples. These weights are then uniformly assigned to all training prompts according to their source component. For training weights, we set = (0, 0.2, 0.4, 0.6, 0.8, }} and let the training weights be
W
w)
PIN
P1
P2N
P2N
N
31

---

Published as a conference paper at ICLR 2026
(a) Distributional
Alignment
(b) Error accumulation
(c) Error accumulation
Analysis
Figure 4: The generalisation error under non-i.i.d scenario.
The impact of distributional alignment on non-i.id generalisation: We evaluate the role of distribution mismatch by explicitly controlling the norm || vlI, which quantifies the divergence between the empirical training distribution q and the ideal importance-weightee distribution v. Figure 4a shows that the le distance between q and v steadily increases, and the non-i.i.. generalisation ability correspondingly deteriorates, manifesting Corollary 4 which asserts that tighter alignment between the training and test prompt distributions yields better generalisation under distribution shift. Moreover, it underscores the importance of high-quality prompts for non-i.i.d. settings, since they reduce the gap between training and test distributions and thus improve generalisation.
larger generalisation error. This observation validates
The Generalisation Error Convergence Analysis: We evaluate the generalization error as the number of training samples N increases. Figure Hb) demonstrates that the error decreases and asymptotically vanishes as N + oc, consistent with the theoretical prediction in Corollary 4 for the non-i.i.d. setting.
The Error Accumulation Analysis: Figure 4c) shows that the generalization
increases with sequence length, following an approximately logarithmic trend. In particular, once the sequence length exceeds threshold near in N, the error rises sharply. Moreover, this threshold shifts to larger values as the sample size increases. These empirical findings support Theorem under non-i.i.d
(a) The Length of Sequence is 6
(b) The Length of Sequence is 7
(c) The Length of Sequence is 8
Figure 5: Generalisation error progression over optimization steps in the non-i.i.d. setting.
Overfitting Risk: Figure how the generalisation error evolves with the number of optimization steps. As the sequence length increases, the task becomes more complex and the loss landscape grows more non-smooth, resulting heightened risk of overfitting. These observations align with the conclusions of Theorem 4
Empirical evaluation on realistic data:We conduct an additional NLP experiment on sentiment- classification task. The training set consists of labeled movie reviews, and the test-time prompts contain several review-label demonstration pairs. We collected approximately 600 movie reviews from Douban, segmented them into sentences, and fine-tuned base GPT-2 model. Another 100 reviews were prepared for in-domain testing. Using the same procedure, we also constructed literary-text test set from online literature platforms create distinct out-of-domain distribution.
After fine-tuning on movie reviews, we examined how the discrepancy measure disc(q) relates to generalization behavior. To this end, formed trrget-prompt mixtures spanning both movie-review and literary domains, with mixture ratios ranging from 0:7 to 7:0. The results are reported in Table 4 Mixtures containing higher proportion of literary prompts correspond to larger disc(g). as literary
32

---

Published as a conference paper at ICLR 2026
texts difter more substantially from movie reviews (empirically. their bidirectional KL divergences around 12). These nigher-discrepange mixtures exhibit moderately increased predictive loss, whereas mixtures more aligned with the training distribution (smaller disc(q)) show lower loss and stronger in-context performance. Overall, the observed trend is consistent with the qualitative dependence predicted by theoretical analysis.
Prompt Config
Loss
Top-1 Acc.
Prompt Config
Loss
Top-1 Acc.
7:0
0.9319
90.78%
3:4
0,9416
90.64%
6:1
0.9363
90.71%
2:5
0.9434
90.61 %
5:2
0.9374
90.66%
1:6
0,9474
90.57%
4:3
0.9398
90.64%
0:7
0.9477
90.60%
Table 4: Prompt configuration vs. performance on sentiment classification.
The Validation of Assumption on Lipschitz Constant: Although the theoretical and empirical constants need coincide numerically, observing that the empirical estimates follow the same scaling laws across model sizes and datasets confirms the asymptotic tightness of our Lipschitz and smoothness bounds. To validate we approximate the constant by sampling multiple imputs, computing gradient and Hessian norms, and taking the maximum observed value. This approach effectively captures the dominant scaling behavior and serves reliable empirical proxy. As shown in Eqs. 7 g the Lipschitz (smoothness) upper bound depends on factors such as QKV matrix size, model depth, and other architectural parameters. We varied these factors to examine their influence, with the results summarized in Table 5 The empirical results reveal clear scaling patterns of the Lipschitz constant with respect to key architectural parameters such as model depth, embedding dimensions and attention head. The consistent asymptotic behavior provides empirical evidence supporting the effectiveness of our theoretical Lipschitz (smoothness) bound.
Number of Layers
# Lipschitz (Layers)
Attention Heads
Lipschitz (Heads)
Embedding Dim
Lipschitz (Embedding)
12
19.38
4
19.38
1218
19.38
43.46
36
59.52
32
38.71
5712
108.31
70.22
11044
908.
60
820.52
938,
1 1168.58
1961.24
Table Lipschitz-related quantities across Transformer configurations.
J THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used LLMs (c.g., ChatGPT) only as a general-purpose writing assistant. Its roles were limited to polishing language (grammar and clarity), and concise rephrasing or shortening of paragraphs without adding technical content. The LLMs did not generate research ideas, problem formulations, proofs, theorems, algorithms, experiments, results, figures, or evaluations. All technical content (definitions, lemmas/theoremss, proofs, algorithms), experimental designs, and conclusions are solely by the authors and were fully verified by us. The authors take full responsibility for all text this paper. The LLM is not an author.
33


Reviews:
[{"review_id": "E8EZNHs5TH", "reviewer": "Reviewer_J7sq", "summary": "In this paper, generalization of in-context learning is studied using the tools of algorithmic stability. First a stability bound for transformer based architecture is derived, and then such a stability bound is coupled with discrepancy measure and provides the generalization guarantees for i.i.d sequences and non i.i.d sequences. Empirical investigation on synthetic tasks are provided to support the theory.", "questions": "i) In the context of Theorem 2, do you give an example of scenario when $\\beta \\| q\\|_2 N \\to 0$ ?\n\nii) In table 2, the paper presents the convergence rate, does it not take into account the convergence of training loss ?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper tackles an important question of generalization of language models and proposes a general framework and approach for generalization bound for incontext learning using algorithmic stability and discrepancy measure. The flexible approach can also handel non i.i.d data scenario's with discrepancy measure and is interesting.", "weaknesses": "A following things limit the applicability or significance of the result. \n\ni) A key limitation of the paper is that its results lack clear interpretability, and the novelty of the proposed approach is not effectively communicated to the reader. The work would be significantly strengthened by presenting a concrete example, such as one involving in-context learning. This would provide a practical scenario where the derived generalization bounds are tangible and make intuitive sense, helping to ground the paper's theoretical contribution \n\nii) What is the specificity of the result to incontext learning or transformer architecture ?\n\niii) The exteme dependence on  the iteration number, the stability becomes worse and worse with the  number of iterations (Q) and it is sometimes logarithmic in the Q meaning the result is not applicable for a single pass over the data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qOK6PgZUDZ", "reviewer": "Reviewer_gCrj", "summary": "This paper develops a theoretical framework for understanding the generalisation behavior of in-context learning (ICL) in Transformers under non-i.i.d. settings. The authors derive a generalisation error bound for the algorithmic stability and distributional discrepancy measure and conduct empirical evaluations to validate their theoretical findings.", "questions": "1. Could the authors include more realistic ICL tasks in the experiments to better support and validate their theoretical claims?\n2. Could the authors provide a more detailed discussion on the practicality and justification of the boundedness and Lipschitz smoothness assumptions in real-world settings?\n3. These assumptions (boundedness and Lipschitz smoothness) appear to be quite general and could potentially apply to many neural network architectures or tasks. Do the theoretical results in this paper offer insights that are specific to the Transformer architecture or the ICL problem in particular?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper studies an interesting and important question regarding the stability and non-i.i.d. generalization of in-context learning (ICL), and it may offer valuable practical insights.\n\n2. According to Table 1, this work considers more general and realistic settings compared to prior theoretical studies.\n\n3. The paper is well-presented, featuring clear illustrative figures, concise proof sketches, and well-organized comparisons with related works (e.g., Table 1).", "weaknesses": "1. The experiments are purely synthetic and do not include realistic NLP or multimodal datasets. This limits practical impact.\n\n2. Boundedness and Lipschitz smoothness may not hold for real Transformer loss landscapes; discussion of how these assumptions approximate practice would help.\n\n3. It seems that these assumptions (boundedness and Lipschitz smoothness) are very general and could apply to any neural network architectures or tasks. Therefore, it is unclear whether the theoretical results in this paper truly provide any insights that are specific to the Transformer architecture or the in-context learning (ICL) problem.", "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": "PMAvKcXfJk", "reviewer": "Reviewer_bLpF", "summary": "This paper derives generalization guarantees for non-linear multi-layer/multi-head Transformers under ICL, by coupling mini-batch-GD-dependent uniform stability with a hypothesis-space-independent discrepancy measure; the bounds highlight (i) optimization- and smoothness-aware choices of step-size/batch/iterations, (ii) the need to align prompt distributions between training and inference, and (iii) error accumulation across generated tokens implying an at-most logarithmic growth of prediction length for reliable generalization, corroborated by experiments.", "questions": "It would be good to include the discussions of arXiv:2508.09820 and arXiv:2411.02199, which consider generalization analysis over non-orthogonal data.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The theoretical results are seemingly sound and cover both smooth and non-smooth regimes.", "weaknesses": "* The definitions of $\\zeta_1$ and $ \\zeta_2 $ are missing or unclear around lines 295–305. Please specify them explicitly for completeness.\n\n* It is counter-intuitive that the non-smooth counterpart allows ( Q ) to be exponentially smaller than that in the smooth case (Corollary 2 vs. Corollary 1). The theoretical intuition behind this discrepancy should be clarified.\n\n* The claim in lines 300–302 remains vague without quantitative support. A more formal analysis is needed to characterize the continuous Pareto frontier of the purported trade-off.\n\n* The manuscript does not clearly explain why small-batch SGD achieves better generalization than its large-batch counterpart. A brief theoretical sketch in the proposed framework would strengthen the argument. Similarly, the Remarks section could more clearly outline how the asymptotic behaviors arise from the assumed settings, rather than only restating theorems.\n\n* (Relatively minor point) Experiments on realistic datasets would make the findings more convincing and demonstrate the applicability of the theory.", "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": "kY849nY0WJ", "reviewer": "Reviewer_rUZ4", "summary": "This paper studies the ICL generalization of Transformers by characterizing the algorithmic stability and distributional discrepancy. The authors especially establish comprehensive discussion in different scenarios, including smooth or non-smooth loss functions, and i.i.d. or non-i.i.d. data. Some experiments are conducted for supporting the theory.", "questions": "1. I don't quite get the discussion of Figure 2 (b). I cannot see why the generalization error increases following a \"logarithmic\" trend with sequence length. \n\n2. It seems experiments in Section I are more interesting to justify the theoretical results. Why not put them in the main body?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The theoretical analysis is impressive and looks solid. \n\n2. The discussion is comprehensive enough to cover many cases.", "weaknesses": "1. The practical insight of the analysis is unclear. I am not sure how the proposed results can be used to explain any phenomenon or improve performance. This makes this work less interesting to me. \n\n2. This paper mainly focuses on Transformers with multiple heads and multiple layers. However, it is not clear how the number of heads and layers affects the theoretical results. I don't know whether the derived bounds are tight and can be quantitatively verified by experiments. It is also not clear how and why the results and the analysis of Transformers differ from non-Transformer models. Therefore, I cannot specifically evaluate the novelty of the analysis.\n\n3. The message from Section 4.1 is quite awkward. It combines both a brief introduction of the theoretical results and the main proof technique. However, a proof sketch should introduce the logic chain of establishing the proof rather than only mentioning the theoretical tools used in the proof.", "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)