How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs
Authors: Samet Demir
Zafer Dogan
Venue: NeurIPS 2025 poster ()
Keywords: Transformer
in-context learning
nonlinear MLP
data mixing
Gaussian universality
Primary area: theory

Abstract: Pretrained Transformers demonstrate remarkable in-context learning (ICL) capabilities, enabling them to adapt to new tasks from demonstrations without parameter updates. However, theoretical studies often rely on simplified architectures (e.g., omitting MLPs), plain data models (e.g., linear regression with isotropic inputs), and single-source training—limiting their relevance to realistic settings. In this work, we study ICL in pretrained Transformers with nonlinear MLP heads on nonlinear tasks drawn from multiple data sources with heterogeneous input, task, and noise distributions. We analyze a model where the MLP comprises two layers, with the first layer trained via a single gradient step and the second layer fully optimized. Under high-dimensional asymptotics, we prove that such models are equivalent in ICL error to structured polynomial predictors, leveraging results from the theory of Gaussian universality and orthogonal polynomials. This equivalence reveals that nonlinear MLPs meaningfully enhance ICL performance—particularly on nonlinear tasks—compared to linear baselines. It also enables a precise analysis of data mixing effects: we identify key properties of high-quality data sources (low noise, structured covariances) and show that feature learning emerges only when the task covariance exhibits sufficient structure. These results are validated empirically across various activation functions, model sizes, and data distributions. Finally, we experiment with a real-world scenario involving multilingual sentiment analysis where each language is treated as a different source. Our experimental results for this case exemplify how our findings extend to real-world cases. Overall, our work advances the theoretical foundations of ICL in Transformers and provides actionable insight into the role of architecture and data in ICL.

How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs
Samet Demir', Zafer Do
MLIP Research Group, KUIS Al Center, Koç University *Department of EEE, Koç University (sdemir20, zdogan) Ckuu 0ku. edu.tz tr
Abstract
Pretrained Transformers demonstrate remarkable in-context learning (ICL) capabil- ities, enabling them to adapt to new tasks from demonstrations without parameter updates. However, theoretical studies often rely on simplified architectures (e.g.. omitting MLPs), plain data models (e.g., linear regression with isotropic inputs), and single-source limiting-lmiitnn their relevance to realistic settings. this work, we study ICL in pretrained Transformers with nonlinear MLP heads on nonlinear tasks drawn from multiple data sources with heterogeneous input, task, and noise distributions. We analyze model where the MLP comprises two layers, with the first layer trained via single gradient step and the second layer fully optimized. Under digendimensional asymptoties, equivalent in ICL error to structured polynomial predictors, leveraging results from the theory of Gaussian universality and orthogonal polynomials. This equivalence reveals that nonlinear MLPs meaningfully enhance ICL performance- -atemaacce--attuulaarry particularly nonlinear tasks compared linear baselines. It also enables precise analysis of data mixing effects: we identify key properties of high-quality data sources (low noise, structured covariances) and show that feature learning emerges only when the task covariance exhibits sufficient structure. These results are validated empirically across various activation functions, model sizes, and data distributions. Finally, we experiment with a real-world scenario involving multilingual sentiment analysis where each language is treated as different source. Our experimental results for this case exemplify how our findings extend to real-world cases. Overall, our work advances the theoretical foundations of ICL in Transformers and provides actionable insight into the role of architecture and data in ICL.
prove that such models
Introduction
Transformers [38] have emerged as dominant architecture across a broad range of machine learning applications. A particularly striking capability of these models in-context learning (ICL)-the the ability to adapt to new tasks from a few examples presented as input, without parameter updates [ As this phenomenon becomes central to modern AI systems, deeper theoretical understanding of when and why Transformers succeed at ICL is increasingly critical.
To make this challenge tractable, recent work has studied simplified settings, often focusing linear
tasks and attention-only models that omit multilayer perceptrons (MIPs) 3 42 45 47). However, such abstractions overlook the empirical observation that nonlinear MIPs integral to Transformer performance in practice. While a few recent studies have begun to investigate MLPs in the ICL context [23 32). their analyses are restricted either nonstandard architectures, specific activation assuming single, homogeneous data source.
types, or overly simplistic data
Corresponding author Source code:
Transformers
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In contrast, real-world settings involve Transformers with nonlinear MLPs trained with data from multiple heterogeneous sources, where data quality vary. This raises two open questions:
(i) How do nonlinear MLPs shape ICI. behavior under realistic data conditions? (ü) What role does mixture of training data play in shaping the Transformer'' capabilitiess for ICI. and feature learning?
To address these questions, we analyze the ICL performance of Transformers with a nonlinear MLP head trained under multiple data sources with different distributions, focusing on nonlinear regression tasks in high-dimensional regimes, We adopt an asymptotic lens where sample size, context length, and hidden dimension jointly diverge, and draw on Gaussian universality theory [3 to analytically characterize model behavior. Specifically, we study Transformer with linear attention and two-layer nonlinear MLP head, where the first layer trained with a single gradient step and the second is fully trained. We show that this model is asymptotically equivalent to finite-degree polynomial model in terms of ICL error. This equivalence allows us to (i) explain the performance gains due to nonlinear MLPs and (ii) study how different training data distributions affect ICL. outcomes.
Our theoretical findings are corroborated by simulations, which demonstrate that nonlinear MLPs significantly improve ICL performance over linear counterparts. We further reveal that data sources with structured covariances and low noise are critical for effective learning. Next, we show that feature learning in Transformers depends crucially on the structure of task distributions, and that certain mixtures can either enhance or suppress the model's 's ability to learn useful features. Finally, we provide an experimental result on real-world scenario (multilingual sentiment analysis), displaying how our results can apply in scenario where different languages are treated as different sources.
Overall, our contributions
follows:
1. We show that a Transformer with a nonlinear MLP head is asymptotically equivalent to a polynomial model, and outperforms its linear counterpart on nonlinear ICL tasks.
2. We analyze the role of data mixing in ICL and identify key coperrtes--structuree covariances for input and task vectors and low target noise that define
data sources.
3. We uncover the interaction between data mixing and feature learning, showing that structure in task distributions is essential for meaningful feature learning to occur.
2 Related Work
In-context learning with Transformers The discovery of in-context learning (ICL) in large-scale Transformers (7] has prompted a surge of empirical and theoretical investigations. Empirically, studies such as 44 33 37 highlight that ICL capabilities grow with model size, suggesting an emergent behavior critical to modern foundation models. On the theoretical side, controlled synthetic especalyyy linear regression-har become a standard benchmark for analyzing ICL mecha- nisms in simplified settings [47 [6] These studies often adopt linearized attention mechanisms for tractability and argue that Transformers learn to simulate algorithmic procedures during pre- training [5 2 (25 42,
26 35 Yet, the precise nature of these learned procedures remains 28 offer generalization analysis for Transformers with a linear attention (without MLPs) in linear ICL tasks, laying the groundwork for extensions to more realistic settings.
29 15
unclear. Notably.
Nonlinear MLPs in Transformers While much of the theoretical literature abstracts away the MLP components of Transformers, recent efforts have started to examine their impact on ICL. For example, 24 studied ReI_U-based MIPs in classification tasks, while 2 32 analyzed architectures where MI. layers precede attentioondevvatin from the standard Transformer design in which MLPs follow attention blocks. Furthermore, (1] studied a simplified Transformer architecture with a fixed (non-trainable) nonlinear attention mechanism -omittn weights-and an optional MLP head. Overall, these studies either restrict task types, limit activation functions, adopt nonstandard architectures, or omit training of the attention weights, leaving open questions about the role of nonlinear MLPs in Transformers for nonlinear ICL in realistic settings. Our work addresses this gap by analyzing a Transformer with trained nonlinear MLP head consistent with the canonical architecture and showing asymptotic equivalence of a polynomial surrogate model to the Transformer model in terms of the ICL error. Moreover, prior works focus primarily on single-source training with simplified data distributions, whereas we consider multiple heterogeneous data sources and explicitly characterize the impact of data mixture ratios on the ICL performance.
omitting the training of key, query, and value

---

Gaussian universality The concept of Gaussian universality from random matrix theory has proven valuable for the asymptotic analysis of neural networks. Under isotropic Gaussian inputs, random feature models (two-layer neural network with random first layer and trained second layer) can be
shown to be equivalent to Gaussian models with matched first and second moments [7 21 g These results have been extended to broader settings, including empirical risk minimization [3]] mixture-distributeed inputs [ and structured covariances [[2]. More recently, the equivalence has been rigorously established for two-layer networks trained with one gradient step [30] I0 [3]. However, these universality results have not yet been extended to Transformer architectures or in- context learning scenarios. Our work bridges this gap by showing that a Transformer with nonlinear MLP is asymptotically equivalent to a finite-degre polynomial model thus connecting Gaussian universality theory to Transformer-bar based ICL for complex data scenarios involving data mixing.
3 Problem formulation
3.1 In-context learning of nonlinear regression
We investigate the in-context learning (ICL) capabilitiie of pretrained Transformer architectures nonlinear regression tasks. Given a sequence of input-output pairs pairs-referred
referred to as the context
(1 31), (222 12), ..... (a.,, ve), (#z+1, ?),
where each input 2; Rd and corresponding response ys E R are sampled i.i.d. from an unknown joint distribution, the objective is to predict ye+1 for a new input E++1 using the context. The context length denoted by f. In the nonlinear setting, we assume the input-output relationship cannot be captured by any linear function. Thus, the model has to learn to estimate the nonlinear input-output relationship from a set of training contexts.
3.2 Data model
While most of the theoretical works assume a single data source, in practice, the models are pretrained with datasets consisting of mixture of data sources [46). To reflect this, consider S different data sources. Thus, when constructing dataset, we first sample the source index s from categorical distribution with the following probabilities: (s i) pa for {0, I....,& S- 1}. Then, for each data source, we posit a nonlinear relationship between x and y governed by a context-specific parameter vector & E IRd, subject additive Gaussian noise. Although & remains constant within given context, it is ra-saed across different contexts, imposing the requirement that the model estimate & from the observed pairs before generalizing to the new input.
Formally, for a given source index we first sample the task vector Els ~ N(M,, Σg.,) and using the sampled task vector, data is sampled
(Els)=(xc(6 + 6//1,
sii8 ~ (u Er,s), and vils = &. ||E| E,,, 11
where o. : + IR is a nonlinear function and G/ss ~ (O, A?) is noise. This setup encapsulates structured nonlinear regression and classification tasks with mixed data sources.
Following 47 we construct an embedding matrix Z E [R(++) by stacking the feature vectors and their corresponding labels, using zero placeholder for the unknown label:
T
Ø2|s
xt/s
xs41|s]
((d+1)((-11)
Y1l5
y2|s
yels
0
where we omit the source index for Z to streamline the presentation.
3.3 Transformer model
In this work, we consider a Transformer model involving single block / layer of linear attention and an MIP head (two-layer neural network). While it is possible to extend the setting to those involving multiple blocks of attention and MI.P layers, following many theoretical studies [47 27 in the literature, we focus on the single block case to streamline our theoretical analysis. Similarly, while

---

nonlinear attention mechanisms (such as softmax attention) are popular in the empirical literature, we use linear attention since it is analytically tractable [47 27 and variants of linear attention can lead to performance similar to softmax attention [20] Another benefit of focusing on linear attention is that it allows us to isolate the role of the nonlinear MI.P head in capturing the nonlinear nature of our ICL tasks since nonlinear attention can also play role in capturing such nonlinear relations. In the rest of this section, we first describe the linear attention cechanismm-inrroducee linear Transformer without MLP head our baseline- and then, we explain our Transformer model with MLP head together with the training procedure.
3.3.1 Linear attention (linear Transformer)
The Transformer must leverage this representation to infer the underlying mapping and predict Ve+1- Then, the output of linear attention in the Transformer can be computed as:
1 A = Z ====+VZ(K2"((02) (ZZ)
where K, Q, V appropriately sized key, query, and value matrices, respectively. When predicting 11+1, the relevant output the linear Transformer (without MLPs) will be 1841.8+1 i.e., the element of corresponding the 0 element the embedding matrix Z. Using the reparameterization in 27 (see Appendix A), this prediction can be simplified to:
Jtinear vec(T) vec(()"vec(HH)
vec(T))
Hz))
(4)
where vec(-) denotes the vectorization operation, r E [Rdx(d+1) is the parameter matrix formed using the entries of K, Q, V matrices while Hz is defined
Hz == Øe+1 It Eige y.
Eige v E (dx(d+))
Note that the parameter matrix I' is trained for the linear Transformer case as
vec(T)) Σ V -vec-) vec vec(ez))) H2 + |||T|
arg min F
where À the regularization constant. Here, ((Z), ()))) denotes the n samples (contexts) used for the training where Z is formed by 2
3.3.2 Transformer with nonlinear MIP
ICL of nonlinear tasks requires nonlinear models. As such, Transformers include nonlinear MLP layers for capturing such nonlinearities. Therefore, we consider the Transformer with a nonlinear MLP head, for which the prediction of the model can be written
Inonlincar == w oFFeeFeeeHH))
(5)
Vi
where o IR is the activation of MLP, and the parameters are F Rkxd(d+1) and e R Here, F encapsulates the parameters of the attention (denoted by I in the description of the linear attention above) and the parameters of the first layer in the MI.P head while ur is the weights corresponding to the second layer of the MLP.
Training procedure To simplify our analysis while keeping feature learning for the MLP, consider the following training procedure with two stages 9
i) Gradient descent on the first layer: Given {(2,, o/ 1)+1)) ))= a set of training samples drawn using the procedure described above, we first fix at the initialization and perform single gradient descent step on F with respect to squared loss. The gradient update with step size >> 0 is given as
F==F+t
(6)
where gradient matrix G is defined
ww TwT (( o oF(FTT) A,
G 71
VI
with H [[veo[ H 21), vec( H Hz)),
and ÿ 1i V2.1.... vec )1²

---

ii) Ridge regression for the second layer: A fresh sample set {(2), v2411)33=11 is used to train w:
n Σ
1 = (Fvec(HH))) 2 + *1/2012
i == arg min
Vet1
n 2=1
vie
where À 2 0 is the regularization constant.
Here, by avoiding data reuse, we ensure F remains independent of the second training set, which simplifies theoretical analysis [4] Furthermore, while our training procedure differs from standard end-to-end training techniques, this two-phase scheme -gradient descent on the first layer followed by ridge regression on the second- is well-motivated by prior theoretical work on MLPs 2 This separation enables tractable analysis while preserving meaningful feature learning dynamics.
3.4 Evaluating ICL Performance
To evaluate ICL performance, we consider a uniform mixture of data sources, i.e., IP(s = 1/S for all i E (0,1,..., 1}. Under this setting, we define the ICL error
8-1 M E (Jet1 -) $]. 8=0
(8)
where ý denotes the model prediction, given by (4) for the linear Transformer or by (5) for the Transformer with a nonlinear MLP. This formulation enables controlled experimentation with training- time data mixtures while evaluating performance as average across all sources. Such approach aligns with practical multi-source training scenarios, where models are trained on heterogeneous datasets and validated by averaging performance over the constituent sources [46]. Note that, we also provide ICL errors corresponding to each source (for cases when it might be relevant) in Appendix G
4
Main results
In this section, present both theoretical and empirical results. We begin by outlining the assumptions underlying our asymptotic analysis. Next, we introduce our main theoretical result establishing asymptotic equivalence. We conclude with simulation studies and real-world experiment that validate this equivalence and further illustrate how performance influenced by factors such as sample size, context length, hidden dimension, feature learning, and data mixing.
4.1 Assumptions
We formalize theoretical analysis under set of assumptions designed to capture the high- dimensional nature of in-context learning (ICL) and the architectural properties of Transformers.
Assumption 4.1 (Proportional limit). The input dimension d, context length e, number of training samples n, and hidden dimension & diverge jointly d, n,n,k + oc, while maintaining constant ratios f/d, n/ /d, and k/n in +
This assumption aligns with regimes studied in prior work: l/d and n/d are critical for ICL
performance in linear Transformers (27). while k/n ensures the model capacity scales with data [2]]
Assumption 4.2 (Input statistics). The covariance matrices of input vectors satisfy || 112 = (d)
and Tr (E...) > d while the mean vectors satisfy ||,,,, (dd/22)
The assumption on the input statistics naturally constrains the input distribution. To streamline the exposition, restrict our analysis to the zero-mean case (Mu.s 0), with the generalization to non-zero mean inputs discussed Appendix F
Assumption 4.3 (Scaling of the step size). The step size y scales as = o(d2).
By permitting step size to scale with the input dimension while the dimension tends to infinity, we introduce nontrivial feature learning, allowing the model to learn nonlinear features 13.
5

---

Assumption 4.4 (Covariance of the attention output). For all task indices i. the attention output satisfies
([Coovvee(Hz))))) = :) Tr(Cov(vec(H =]));
and admits a low-rank decomposition:
Cov(vec(Hz)) = La(a+1) +
Σ 0. Yo.4 Yo.4 Ys.g
(9)
where V, E Zt, 0s.9 > 0, and (%a.g) are orthonormal vectors Rd(d+1),
Assumption 4.5 (Joint scaling). The step size and covariance of attention outputs satisfy
||Coov (vec Hz))lle - o(d²).
Assumptions 4.3 4.5 ensure tractability in asymptotic analysis while allowing the feature learning.
Although these assumptions might look unnatural at first glance, they enable us to connect our setting involving ICL by Transformer with nonlinear MLP to prior results [13] for MLPs, facilitating our theoretical analysis. We would like to note that Assumptions 4.3 4.5 represent limitations of our theoretical results. Nevertheless, our experimental results on both synthetic and real-world datasets suggest that these assumptions can be partially relaxed in practice.
Assumption 4.6 (Initialization). The vector NO, + and the matrix F = [fi..... fu] has
La(a)) rows f: N (0, TT(Cn)(ve<(Hz)))
This initialization is consistent with standard practices for MIP initializations 18
Assumption 4.7 (Target function). Each task-specific function o, R > R is Lipschitz.
Assumption 4.8 (Activation function), The activation R R has bounded derivatives and
satisfies E2- (2) < co, admitting a Hermite expansion 34 Chapter 11.2]:
(x) Σ Hy hy = (===)
(10)
where Hy is the j-th probabilist's Hermite polynomial.
Assumptions 4.7.4.8 4.8 accommodate a broad class of nonlinearities for both the target and activation functions, while maintaining analytical tractability via Hermite expansions of the activation. Although certain functions, such as polynomials, may not strictly satisfy the bounded-derivaaive condition, our theoretical results are expected to remain valid when the derivatives of the activation are bounded with high probability for inputs ~ (0,,1) 1).
4.2 Asymptotic equivalence
We present theoretical results in terms of asymptotic equivalence between models. The main technical challenge in our analysis arises from characterizing the distribution of the attention outputs, denoted by Hz)). To address this, we begin with an equivalent statistical representation of the random feature mapping Fvec Hz), as formalized in the following lemma. This representation later serves the foundation for our analysis of the Transformer with nonlinear MLP (5)
Lemma 4.9 (Asymptotic distribution of Fvec Hz)). Under Assumption 4.6 the entries of are i.i.d. N 1/TO.///e H 1)), ensuring that the components Fvec H 2) have unit variance. Let f, denote the i-th row of F. Given the assumptions
stated in Section 4.1 we obtain
f.f vec Hz)) N'O, 1) converges in distribution.
(11)
for all i kk).
Proof. Set Tr(Cov(vec( (H))) Given Hz, we have
vec(()))) Hz ~ N v/ve (vve(Hz))l) ll2tt)
(12)
Since (H2))l/tt concentrates around 1, the desired result follows. See Appendix B

---

The lemma simplifies the term Fvec Hz), enabling a more tractable analysis for our study of in-context learning in Transformers with nonlinear MLPs. As result, we concentrate on the joint behavior of (Fvec Hz), E æx+1) rather than z De++), 1. given that ye+1 == "x æe-1)) + ce+1- The subsequent corollary demonstrates that (Fvec( H E m241) converges in distribution a jointly Gaussian vector in the high-dimensional limit that we consider.
Corollary 4.10 (Joint distribution of (Fvec Hz),, E" x++1) conditioned on & and s). Suppose that the task vector and data source
given. Under our assumptions (given in Section 4.1
(Fvec Hz), E" 2x41) becomes jointyy Gaussian with a certain mean and covariance.
Lemma 4.9 and Corollary 4.10 collectively characterize the distributions of both the Fvec(Hz) and the label Me+1, which are utilized for our subsequent study of the asymptotic equivalence.
Next, we need to consider the impact of updating F with one gradient descent step as described in ( F = := F r/ where G' is the gradient matrix and 1j is the step size. For this purpose, we look into the gradient matrix G the following lemma.
Lemma 4.11 (Decomposition of the gradient matrix). The gradient G defined in 7 admits the following decomposition: G uv + A. where == aw for some R and v == H // (m/k). Here, we is the dominant rank-one term while 1|2 is the negligible residual term.
tail bounds, this can be shown similarly to the decompositions of the gradient matrix for simpler neural network settings in the literature 44 See Appendix C C
Proof. Using
With these components established, existing asymptotic analyses of two-layer neural networks [33 to study the in-context learning error defined in (8) for Transformers equipped with nonlinear MLPs. The following theorem identifies polynomial model that asymptotically equivalent to such a Transformer architecture.
now leverage the preceding results in conjunction with
Theorem 4.12 (Equivalent polynomial model). Suppose our assumptions described in Section 4.1 Consider we train the first layer F of the MLP according to (6). Then, the Transformer with nonlinear MLP in (5) performs asymptotically equivalent to the model
- Jo'o,(Fveec Hz))
(13)
Vk
with respect the ICL error S where the second layers in the models trained separately. Here, op IR - IR is degree-p polynomial function with a residual term, which is defined as
P Ōp(x) = Σ c (x)
+ cG2 for N),,
(14)
where H; R. denotes the i-th (probabilist's) Hermite polynomial, c; are the corresponding Hermite coefficients and e₂ is the residual term such that Ez~N1o. "tt.n)@@(/)))==mnppp (x) The smallest degree that is necessary for the equivalence depends on the joint distribution of (Fvec( Hz), c" 11+1) and the activation function a, but a finite p suffices under our assumptions.
regarding the distribution of (Fvec( H2). xx+11) - in combination with the orthogonality property of Hermite polynomials under the Gaussian measure in order prove this theorem. See Appendix
Proof. We employ our established theoretical results 4.9
Theorem 4.12 establishes an asymptotic equivalence between the ICL. behavior of a Transformer with an MLP and that of polynomial model in terms of the ICL errors. This result is significant for several reasons. Firstly, [13 admits a much simpler analysis, allowing precise characterization of ICL error behavior of the Transformer under data mixing and feature learning. Second, sheds that it effectively learns
light on the function class learned by the Transformer-specine low-degree polynomial approximation of the task function. This equivalence opens the door to optimizing nonlinearities in MLPs (inside the Transformer) through polynomial surrogate analysis. These insights extend beyond classical intuition and align with recent trends in Gaussian equivalence literature [1 30 13 Note that although similar equivalence results exist for supervised learning by
two-layer neural networks 30 13, this work represents a cornerstone in adapting those results to
novel setting involving ICL by Transformers with MLPs.
7

---

10
tush
Tranfbennrr
Equivalest
10
Lince Transfomer
6x
ICI Pi.
6x.30
408
6x30"
4x
4x300
3x:10
2+
110
03.5
1.0
1.5
20
1.3
1.5
2.0
2+
(a) Number of samples
(b) Context length
(c) Hidden dimension
(L (E= and k 0.5d%)
(n k 0.5d²)
(n 0.5d and f=
Figure 1: Effects of sample size, context length, and hidden dimension on the ICL errors for linear Transformer 4 Transformer with a nonlinear MLP 5 and the equivalent model (13). The number of data sources is S = 2 with equal probability: (s 0) P(s = 1) 1/2. For the input vectors, Mx.s 0 and Σ,, Ia for all s. For the task vectors, Ng.e 0 for all while Σg,0 Ia and ΣS,1 + 0'y'" for some 0 x de and y E IRd with / = The target function o, is ReLU, and the target noise is A, 0.01 for all The Transformer is used with two different activation functions: ReLU and tanh. Here, d = 80, n = d², and À 5 x 10-5 The degree of the equivalent polynomial model is set to 4. The average over 20 Monte Carlo runs is plotted.
4.3 Experimental results
Leveraging our equivalence results, we systematically compare the performance of the Transformer with nonlinear MLP that of the linear Transformer. We also empirically validate the equivalence between the Transformer and its polynomial surrogate, highlighting their predictive alignment under various conditions. We illustrate the data mixing effects in ICL exhibited by Transformers with nonlinear MIP. Next, we provide results depicting the interplay between feature learning, data mixing, and properties of the data sources. Finally, we consider an example real-world scenario involving multilingual sentiment analysis and display how our findings extend to this case.
4.3.1 Data setup
In experiments, consider setting with data sources. For the synthetic (simulation) scenarios, the first source is an isotropic Gaussian in both input and task space, representing low-quality (noise-like) dataset, while the second source features non-isotropic covariance in either the input or task distributions, introducing additional structure and serving as higher-quality signal. This configuration enables us to analyze how ICL performance varies under mixtures differing data quality. The rest of the experimental details and the setting of our real-world (multilingual sentiment analysis) experiment are described in the corresponding captions.
4.3.2 Effect model dimensions
We first assess the role of sample size (n)) context length (f), and hidden dimension (k) on ICL. error. We highlight two consistent observations across all cases:
Results are shown in Figure !
(i) The nonlinear MIPP Transformer (5) outperforms the linear one in terms of ICI. error ( Although outperforming the linear Transformer may sound trivial without the training of the first layer F in (6), there exist many cases where the Transformer with the nonlinear MLP head fails to outpertorm the linear one 14]. Therefore, outperforming the linear one a reasonable baseline and indicator of effective feature learning.
(ii) The nonlinear MLP Transformer closely aligns with its polynomial surrogate model (3). even at moderate dimensionalities, empirically validating asymptotic equivalence results.
These findings underscore both the benefits of nonlinear MLPs and the theoretical significance of our equivalence framework. Additionally, we observe a characteristic double descent phenomenon G in the ICL error: with respect to sample size in Figure I a and hidden dimension in Figure [[). The effect of increasing context length shown Figure (bb). revealing a consistent reduction ICI. error, reflecting improved task estimation with longer contexts.

---

0.40
Tarefortee
Malrg Ratko
648
065
0.0
0.3
0372
0.55
0.00
0.1
00
0.2
0.4
0.5
10
03
1.0
10
10"
50"
Aegal No 110
(a) Input covariance
(b) Task covariance
(c) Noise variance
Figure 2: Data mixing effects the ICL emor. performance of Transformers with ReLU activation (denoted with diamonds (5). and the equivalent model (denoted with lines (13)) are illustrated for different mixing ratios p. For this figure, P controls the data source mixture for the training,
and ((s = 1) = p. while the ICL error is calculated as an average over data sources B Here, d 80, = d, n A 0.5d², À 5 x 10-5 1j x de and the target function 6e is ReLU for all s. We initially consider the following setting: for the input vectors, Mr.s = 0 and and Σg, Ia for all s; and the target noise is A. 0.01 for all For each subfigure, we modify one data property while keeping the rest same and show its effect on the ICL error: in (a), we focus on the input covariance of the second data source by using nx = l + Yz7>> such that | ll2 is changed by varying 0gi in (b), we concentrate on the task covariance of the second data source by using Σg, = Ia + Y such that || la is changed by varying 0g: and in (c), we modify the noise of the second data source A1 while the noise of the first source Ag set to 0.2 and fixed. We set the degree of the equivalent polynomial model to = 5. The average of 20 Monte Carlo runs is plotted.
define (s 0) ==
Eg,s La for all s; for the task vectors, M4g.*
4.3.3 Impact of data mixing
We next explore how mixing structured and unstructured data sources affects ICL performance, shown in Figure 2 In these experiments, we vary the second data source by modifying its input
covariance (a), task covariance (b), or noise variance (c), while keeping the first source fixed.
ii) In Figure a), increasing the proportion of samples from data source with structured inputs significantly improves ICL performance.
(ii) Figure b) shows similar trend when the task vectors have structured covariance, with ICL етror decreasing as their prevalence increases.
(iii) In Figure c), lowering the noise variance in the second source leads performance gains, underscoring the sensitivity of ICL to data quality.
When studying the impact of the data mixing, the ICI. errors corresponding to each source can be of interest, so we provide them in Appendix G where we observe behaviors similar to those in Figure Overall, these results demonstrate that Transformers effectively exploit structured data, and that the relative quality of data sources-deting by structure and noise noise-directoe shapes ICL.
4.3.4 Feature learning with data mixing
While prior results in Figures I 2 used fixed step size 7 x d², we now study the interaction between feature learning and data mixing. Specifically, we assess how varying influences ICL error under two distinct structural conditions: (a) structured input covariance and (b) structured task covariance.
Results in Figure 3 (b) show a notable asymmetry:
(i) In (a), increasing does not improve ICL error, consistent with our intuition that the first-layer updates do not capture a useful signal when task vectors are isotropic. Here, the intuition comes from the fact that the feature learning enables MLPs to align with task vectors [4] Section 3.2]. (ii) In contrast, (b) reveals significant performance gains with increasing as the model can extract informative related features when the task vectors have structured (non-isotropic) covariance.
This contrast highlights that feature learning via gradient updates primarily enhances representations of task-related directions. The extent of benefit depends on both the structure present in the data and the relative mixing ratios of the sources. These findings emphasize a rich interplay between feature learning dynamics and data heterogeneity in ICL.

---

0.500
0423
0.10
0.4
Micing Ratlo
0.301
Froet
0,575
62
02755
0,300
0.323
0.255
0.30
0.38 Tranderrze wath MLP
Equivalect Medel
10
0.200
0.
1.5
2.0
Seguin)
(a) Structured inputs
(b) Structured tasks (| d)
(c) Real-world problem (multilingual sentiment analysis)
2| [B x /
Figure 3: Feature learning with data mixing synthetic and real-world data: the effect of different step sizes 7/ is illustrated. For the synthetic data scenarios in (a) and (b), we start with the same initial setting as Figure In (a), we just modify the covariance of the input vectors from the second data source, while we only change the covariance of the task vectors from the second data source in (b). In each case, we add an additional rank-one structure to the covariance, which is the same as Figure 2 We plot the average of 20 Monte Carlo trials for (a) and (b). For the real-world data scenario (c), we foeus the effect of feature learning ICL errors for multilingual sentiment analysis using the Multilingual Amazon Reviews Corpus [22]. This dataset contains customer reviews (with text and star ratings) in multiple languages, such as English and German. By treating English and German reviews as two distinct data sources, we can vary the mixing ratio across languages, allowing us evaluate our framework in a more realistic setting. We consider English reviews as source 1 and German reviews as source O, so mixing ratio of = 1 corresponds to entirely English data, and p p= entirely German data. For labels {3:}, the review star ratings are demeaned and scaled to lie in the range [-1,1] 1], making the task regression-likes For inputs {æ:}. each review text is embedded using the multilingual text embedding model called "multilingullee--smae generated 384-dimensional embeddings are reduced to 64 dimensions via PCA (principal component analysis), and then normalized. We group I input-label pairs (of the same language) together to form context so that we make the problem compatible with our ICL setting. The rest of the details for 0.25d and X == 5 x 10-5, The degree of the equivalent polynomial model is set to 5. The mean of 100 Monte Carlo trials is illustrated in (c).
[3). the
(c) are
64, n
4.3.5 Data mixing in real-world scenario: multilingual sentiment analysis
For a real-world example, we consider multilingual sentiment analysis using the Multilingual Amazon Reviews Corpus 22 which contains customer reviews in multiple languages. Treating each language as a distinct source, we examine feature learning under data mixing, as illustrated in Figure (3.c). Consistent with our theoretical predictions, the ICL errors of the Transformer model closely align with those of the equivalent model, even in this multilingual real-world setting. We observe clear trend (similar to Figure (b)) whereby increasing the step size 1) reduces the ICL error, indicating effective feature learning. Finally, a higher proportion of English reviews can lead to increased performance, consistent with the known strength of the embedding model (for text) on English data.
5 Conclusion
This work advances the theoretical and empirical understanding of in-context learning (ICL) in pretrained Transformers with nonlinear MLP heads, especially under high-dimensional asymptotics. By leveraging techniques from Gaussian universality and orthogonal polynomials, we establish that such Transformers are effectively equivalent structured polynomial predictors in terms of ICL error. This equivalence not only clarifies the functional role of nonlinear ML.Ps in Transformers but also enables precise analysis of their learning behavior. Our findings reveal that nonlinear MLPs substantially enhance ICL on nonlinear tasks, outperforming their linear counterparts. Furthermore, we explored the impact of having mixture of different data sources during the training, characterizing the properties of a good data source and providing insight about choosing the mixture ratio. Namely, we found that good (igh-quuaitt) data sources are those having low target noise and structured covariances for input and task. We showed that feature learning depends on the data mixture, and specifically, having a structured covariance for the task vectors needed for feature learning. Finally, we illustrated an example showing our findings extend to real-world scenario involving multilingual sentiment analysis. Overall, this work contributes to the theoretical understanding of the impact of data distribution and feature learning ICL by Transformers.
10

---

Acknowledgments and Disclosure of Funding
We acknowledge that this work is supported partially by TÜBITAK under project 124E063 in ARDEB 1001 program. We extend our gratitude to TÜBITAK for its support. S.D. is supported by an AI Fellowship provided by KUIS AI Center and PhD Scholarship (BİDEB 2211) from TÜBITAK. Z.D. is also partially supported by the BIDEB 2224-A program from TÜBITAK.
References
[1] Amirhesam Abedsoltan, Adityanarayanan Radhakrishnan, Jingfeng Wu, and Mikhail Belkin. Context-scaling versus task-scaling in in-context learning. arXiv preprint urXiv:2410.12788 2024.
[2] 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.
[3] Ekin Akyürek, Dale Schuurmans. Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm in-context learning? investigations with linear models. In International Conference on Learning Representations, 2023.
[4] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems, 2022.
[5] Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, 2023.
[6] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical varrince trade-off. Proceedings of the National Academy 2019.
of Sciences, 116(3))155899111
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry. Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[8] Romain Couillet and Zhenyu Liao. Random matrix methods for machine learning. Cambridge University Press, 2022.
[9] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. In Conference on Learning Theory, 2022.
[10] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at time. arXiv preprint arXiv:2305.18270 2023.
[11] Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová. Universality laws for gaussian mixtures in generalized linear models. In Advances in Neural Information Processing Systems, 2023.
[12] Samet Demir and Zafer Dogan, Random features outperform linear models: Effect of strong input-label correlation in spiked covariance data. arXiv preprint arXiv:2409.20250, 2024.
[13] Samet Demir and Zafer Dogan. Asymptotic analysis of two-layer neural networks after one gradient step under gaussian mixtures data with structure. In International Conference on Learning Representations, 2025.
[14] Samet Demir and Zafer Dogan. Asymptotic study of in-context learning with random trans- formers through equivalent models. arXiv preprint arXiv:2509.15/33 15/52, 2025.
[15] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers lear to achieve second- order convergence rates for in-context linear regression. In Advances Neural Information Processing Systems, 2024.
11

---

[16] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems, 2022.
[17] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. The gaussian equivalence of generative models for learning with shallow neural networks. Mathematical and Scientific Machine Learning, pages 226477. PMLR, 2022.
[18] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. attp://www.ddeepeerriigbboo org
[19] Friedrich Götze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials in a-sub-exponential random variables. Electronic Journal of Probability, 26:1 - 22, 2021.
[20] Dongchen Han, Yifan Pu, Zhuofan Xia, Yizeng Han, Xuran Pan, Xiu Li, Jiwen Lu, Shiji Song. and Gao Huang. Bridging the divide: Reconsidering softmax and linear attention. In Advances in Neural Information Processing Systems, 2024.
[21] Hong Hu and Yue M Lu. Universality laws for high- sigh-dimensional learning with random features. IEEE Transactions on Information Theory, 69(3): 1932-196229444 1964, Mar. 2023.
[22] Phillip Keung, Yichao Lu, György Szarvas, and Noah A Smith. The multilingual amazon reviews corpus. In Conference Empirical Methods in Natural Language Processing, 2020.
[23] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In International Conference Machine Learning, 2024.
[24] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear transformers learn and generalize in-context learning? In International Conference on Machine Learning, 2024.
[25] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, 2023.
[26] Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. In Advances in Neural Information Processing Systems, 2024.
[27] Yue Lu, Mary Letey, Jacob A Zavatone-Veth. Anindita Maiti, and Cengiz Pehlevan. In-context learning by linear attention: Exact asymptoties and experiments. In NeurIPS 2024 Workshop on Mathematics of Modern Machine Learning, 2024.
[28] Yue M Lu, Mary Letey, Jacob Zavatone-Vent Anindita Maiti, and Cengiz Pehlevan. Asymp- totic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, )22(29))e2002299122 2025.
[29] Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In International Conference on Learning Representations, 2024.
[30] Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. theory of non-linear feature learning with one gradient step in two-layer neural networks. In International Conference on Machine Learning, 2024.
[31] Andrea Montanari and Basil N Saeed. Universality of empirical risk minimization. In Confer- ence on Learning Theory, 2022.
[32] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained transformer efficiently learns low-dimensional target functions in-context. In Advances Neural Information Processing Systems, 2024.
12

---

[33] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
[34] Ryan O'Donnell. Analysis of Boolean Functions. Cambridge University Press, 2014.
[35] Core Francisco Park, Ekdeep Singh Lubana, Itamar Pres, and Hidenori Tanaka. Competition dynamics shape algorithmic phases of in-context learning. arXiv preprini arXiv:2412.01003, 2024.
[36] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. In Advances in Neural Information Processing Systems, 2024.
[37] Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models a mirage? In Advances in Neural Information Processing Systems, 2023,
[38] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems, 2017.
[39] Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv: 1011.3.0 2010.
[40] Roman Vershynin. High-dimensional probability: An introduction with applications data science, volume 47. Cambridge University Press, 2018
[41] Mariaa Vladimirova, Stéphane Girard, Hien Nguyen, and Julyan Arbel. Sub-weibull distribu- tions: Generalizing sub-gaussian and sub-exponential properties to heavier tailed distributions. Star, 91)xe318, 2020.
[42] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, 2023.
[43] Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, and Furu Wei. Multilingual e5 text embeddings: A technical report. arXiv preprint arXiv:2402.05667. 2024.
[44] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022.
[45] 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 International Conference on Learning Representations, 2024.
[46] Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. In International Conference on Learning Representations, 2025.
[47] Ruiqi Zhang, Spencer Frei, and Peter L. Bartlet. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
13

---

Deferred proofs
In the following appendices, we provide deferred proofs. We first consider the case of zero-mean inputs, i.e., Mz,s O for all 8, and address the extension to non-zero mean inputs, i.e., Mz.s # 0, in the final section. This organization streamlines the initial presentation while preserving generality.
Appendix A iintroduces a reparameterization of the attention mechanism that simplifies the form of the attention outputs. Appendix B analyzes the distribution of Fvec( zz., key quantity in our analysis, and proves Lemma Appendix presents a decomposition of the gradient matrix corresponding single gradient step on the first layer ( establishing Lemma 4.11 Appendix builds on prior results to establish a conditional Gaussian equivalence. Appendix contains the proof of our main theoretical result (Theorem (4.2). Finally, Appendix [F considers an extension of the analysis the setting with non-zero mean inputs.
Notation
We adopt standard notation following 18]. For a random vector v, its covariance is denoted by Cov(v). The spectral norm and Frobenius norm of matrix A are denoted by ||A||2 /|/ and ||A||
respectively, while Tr(A) denotes the trace. Matrix entries and slices are indicated by Ar.j A.,. and Ai.: where and j denote the row and column indices, respectively. vec(.) denotes the vectorization operation. We write f(-) x g()) to denote that f and 9 are of the same asymptotic order with respect to the diverging dimensions. We use standard asymptotic notation: f(d) (gg(d) if (d)//((d + C < do for some C >> 0, and ((d = o(g(d) if ((d)/g(d) + 0 as d + co. Furthermore, we define Ö(f(-)) as shorthand for (ff(-) polylog d) to suppress polylogarithmic factors for brevity. Element-wise multiplication denoted by O and the Kronecker product by @. Conditional random variables are written as X C, representing the distribution of X given condition C. Finally, we use > NO, 1) to denote that the random variable converges in distribution the standard normal.
A Reparameterization of the linear attention
We begin by reparameterizing the linear attention mechanism in accordance with prior work 44 29 23 32 27 to simplify subsequent analysis. Specifically, we consider the following form of linear
attention:
A
(Z
(S1)
where K, Q, and denote the key, query, and value matrices, respectively. Due to the structure of the embedding matrix introduced in (2). the output of linear attention, denoted Mtinears corresponds to the (d + 1, 1)-th entry A, i.e.,
Utinear A4t1.411 Va+1, (K"Q)Z,,et)
(S2)
This formulation highlights that the parameters relevant to Mtinear are V241,; and K KTQ. Q. To isolate these, express the matrices follows:
Mi T
and = M
(S3)
U21
U22
where * denotes components not involved in predicting Ve+11 while the submatrices Mn E dxd 22 m21 E Rdd and U22 E R capture the relevant contributions. Using these components, the prediction Yunear can be rewritten in terms of the input features and previous outputs as:
xe+1, M Σ ya + 022
Σ v +
MT
Σ iSF+1
x U21 + +m21 Σ yo V2L)
ist
(S4)
where (.--) is the standard inner product and omit the source index for brevity.
14

---

In most of the theoretical literature 47 to since the leading term
2
29
23
32
the parameters U21 and m221 are commonly
E1+1, 122 M
in (S4) alone is often sufficient to accurately predict ye+1- More recently, [27] relaxed this assumption by allowing m21 to be trainable while keeping v21 = O, as this modification retains analytical
tractability. Following this approach, we adopt setting in which 121 0 and the remaining parameters are trainable.
Under this formulation, the output of the linear attention mechanism can be compactly expressed as:
Vnmear = vee(T)" (Hc
(S5)
where the trainable parameters are consolidated into the matrix F, and the processed version of the embedding matrix Z denoted by Hz. These are defined as:
U22
Mi m211 and Hz == xe+1 & Eige 1 Eigev
(S6)
The advantage of this reformulation lies the fact that the linear attention output is now a linear function of both the trainable parameters I' and the data-derived matrix Hz. As a result, the matrix can be efficiently optimized via ridge regression.
Similarly, in the case of a Transformer architecture equipped with a nonlinear MLP block, the above formulation enables us to express the prediction as:
Inonlincar == Vk
"T a vee Hz)),
(S7)
where e [xxd((+11 serves a role analogous to that of vec(T). with each row F representing the weights associated with one of the k nonlinear neurons. The function denotes nonlincar activation function, allowing the model to capture complex input-output relationships, and w R trainable vector that linearly combines the outputs of these nonlinear units. Overall, this formulation captures the structure a Transformer with linear attention followed by a nonlinear (two-layer) MLP block.
B Asymptotic distribution of F Fvec (vec(Hz))
In this section, we analyze the asymptotic distribution of the term Fvec (oo( which plays a central role in the Transformer with nonlinear MLP layer. After training the first-layer parameters as described in (6), the prediction of the model (5) takes the form
(Ev
1 V w (Evec(Hz)) = Vi w" (Fvec(Hz) + (Gvec(Hz)))
where G denotes the gradient matrix defined in (7). Since (vec(Hz) is the primary input to the nonlinear activation function understanding its distribution is natural and informative starting point.
Let f. denote the i-th row of the matrix F, and define t : Tr (Ce( H ))). For the purpose of this analysis, we initially assume that Hz is fixed. Under Assumption 4.6 each row f. is drawn independently from Gaussian distribution N Ta(a+1)/t). Consequently, we have
ff veel Hz A (o, (2ccH)))1//t))
Thus, the task of proving Lemma 4.9 reduces to showing that ||vec (z)l concentrates around 1,
which would imply that f. vec H z + (0,
To establish this concentration result, we begin with the following lemma, which characterizes Hz as an outer product of Gaussian vector and sub-exponential vector, conditional on the data source index and the task vector Els. The lemma also provides concentration bounds for the norms of these constituent vectors. For brevity, we omit explicit conditioning in some expressions where the context makes it clear.
15

---

1 Lise you. Lemma B.1. Let b = L Eigge
which implies that H2 = D1+1b'. Suppose that the
data source index s and the task vector Els are given. Then, conditioned on 8, the vector x++1 is Gaussian by definition, and the conditional vector (E(E. exhibits sub-expon exponential tails. Furthermore, conditioned on (s, Els), the following concentration bounds hold with high probability:
||
(C(Co(((2)))) =01),
(S8)
| Tr((E[bb')] /T (T/bbbb"")) o(1).
(S9)
(x.))"
Proof. Given and s we have za/s (Ma,s, En,s), and y: Φ where is Lipschitz by Assumption 4.7 and 6//s is Gaussian by definition.
c/ls,
|
y:a:
We can express b as: b &iige bi, where b, = y
Since Φ, is Lipschitz and æ; is sub-
Gaussian, the composition 1/4 is sub- Gaussian with constant due Gaussian concentration of Lipschitz functions [40] Theorem 5.2.2].
By Assumption 4.2 we have |E (d) and Tr(E...) x d, æ; is sub-Gaussian with norm Odd//4. Then, since the product of sub-Gaussian random variables is sub-exponential 40 Lemma 2.7.7], the term yee; is sub-exponential with norm (dd'/2,, and yf is sub-exponential with constant norm. Thus, b; is sub-exponential vector with norm (dd
Applying the (vector version of) Bernstein's 's inequality 440 Corollary 2.8.3] to the vector (an average over samples of sub-exponential) implies that the deviation (b - E[(]) has sub-exponential tails with norm (d (t)). Under Assumption 4.1 t/d IR, and hence the norm becomes (dd-3/)) 3/1).
Finally, applying the Hanson- -rrinnt- inequality 40 Theorem 6.2.1] yields the bound (S8) For SS9 we rely on an extension of the Haanonn-rright inequality to sub-exponential vectors 19
Now that we have bounds on the norms of the vectors used to construct Hz, we can analyze the concentration of the norm of vec Hz) in the following corollary, which builds Lemma B.1 Corollary B.2. With high probability, the following concentration result holds:
[||vec( Hz) [ Els)))
||vec (Hz) (s, {(s{ o(1),
Hz) (l,{))
t
where = (T(Cov(yec(Hzz))
Proof. We first note that:
vec (ec(H)) vec(#(+10") we+1
Using this, we compute the covariance:
Cov(vee Hz)) B[[(b 10+1) (b ar(+1)
(S10)
TB[( me+1) (b² @ 1)
(S11)
TE((b")) 8 (x21
(S12)
Ebbb") @ Cov(
(S13)
where we used properties of the Kronecker product (transpose and mixed-product), the independence of b and 2441, and the linearity of expectation.
Similarly, we have:
|ive H2)) lbb xet1l ||b] ||
(S14)
and
Tr(COv(vec(H,)) Tr(E(b"']) Tr(Cov(zz+1)).
(S15)
16

---

Define the relative errors:
0 ||| Tr(Cov((e+1)) and 5, = |161 Tr(E(66'])
(S16)
Tr(Cov(ze+1))
Tr(E(bb²1)
Then, conditioned on (s, Els), we get:
H2) Tr(Cov(vee((H)) (H)c)) | Tr(Cor(vec(H2)))
|| 1221b1122 - (r(Cov((e++)) Tr(B(bb²))
(S17)
Tr(Cov(ze+1)) Tr(E(bb"])
S 82 + 8 + 0zb0
(S18)
=())
(S19)
where the final step follows from Lemma B.1 which showed that both 6g and 6, are o(1).
So far, our analysis has focused on the case where Hz is conditioned on the data source and the task vector E|s. However, by Assumption 4.4 we have
L -= (Tr(Cov(vec) Hz)
i) - Tr(Cov(vec( Hz) - ())
for any data source indices i, Furthermore, Corollary B.2 showed that ||vec (H)) around 1 for all and Els. This leads to the following unconditioned concentration result.
[|// concentrates
Corollary B.3. Without conditioning the data source task vector, the following holds with high probability:
Ive Hz)l1
||vec
Hz) (8, Els)
S max *
= o(1),
(S20)
t
t
as a consequence of Lemmal B.1 and Corollary B.2
Proof. We apply union bound over the (finite) data sources and integrate over the distribution of Els. Since the conditional deviation is uniformly bounded by o(1), the same bound extends to the unconditioned case.
Remark B.4. (Implications of task vector randomness) Corollary B.3 implies that the variability introduced by the task vector E|s does not significantly affect the norm of the attention output vec Hz), with deviations bounded by o(1). This concentration result justifies simplification in our analysis: we may treat {/s effectively fixed and analyze the system under worst-cass E s rather than integrating over its full distribution. This reduces analytical complexity while preserving rigorous control over variation across tasks. effect, randomness of task vectors can absorbed into a uniform bound, allowing our results to hold uniformly over task diversity.
This completes our analysis of the concentration of ||vec Hz)l//tt around 1, and thereby concludes proof Lemma 4.9
C Decomposition of the gradient matrix
Having characterized the distribution of Fvec( (Hz), we now turn our attention the effect of single gradient step during training, as defined in Equation ( Specifically, we analyze the structure of the gradient matrix G, following the approach of
Recall the definition of the gradient matrix 77)
1 Vi
www (FFH)) o (F("") Å.
G =
wÿ"
(S21)
complication arises from the presence of the derivative of the activation function o,, which appears as multiplicative factor. To simplify this expression, we apply an orthogonal decomposition to the derivative:
o' a
17

---

where a := (0, () (z)) is the average slope, and o'1 (2) := d')) captures the nean deviation. The decomposition justified by the fact that the random variable entering o namely. f!h, converges in distribution to (0, 1) for all i, by Lemma 4.9
Using this decomposition, the gradient matrix G can be rewritten as:
G -u
nv wÿ"" o H
nk o
H, (S22)
spike
A
where we define == aae and H (/nVk).
This decomposition separates the gradient into two components: A spike term uo"" which expected dominate, and residual term A.
Our next goal is to establish that ||uo"||2 > the spectral structure of G.
A ||2. thereby confirming that the spike term governs
A new technical challenge arises in this context: unlike prior work
where H is composed of i..d. Gaussian vectors, in our setting each row of H is realization of vec (ec(Hz). As shown in Lemma B.1 (Appendix B Hz is the outer product of Gaussian vector and sub-exponential vector, making vec( H heavy-tailed [4] random vector. This deviates from the sub sub-Gaussian assumptions in prior analyses and introduces additional complexity in bounding the norms of H.
Nevertheless, in the following lemma, we characterize a spectral norm bound for H that accommo- dates this heavy-tailed structure.
Lemma C.1. Under our assumptions, the spectral norm of the matrix H satisfies the following bound with high probability:
H 12 [[Cor(vec(H2)))) O(d).
(S23)
Proof. Corollaries B.2 and B.3 imply that
vee H Hz)ll
1
Tr(Cov(vec H = O(P)
with high probability, given that II((Co(vee(HH2)) x dd by Assumption 4.4 It follows that ||vec H 222ll2 (d)
holds with high probability. We now apply result from (39 Theorem 5.44] concerning the spectral norm of a random matrix with independent heavy-tailed rows. This yields the stated bound on H 2- Note that the assumption n/d E (from Assumption 4.1 ensures the validity of this application.
We now establish high-probability bounds for other key random quantities involved in the analysis. Specifically, the following lemma bounds the norms of w, F, and ÿ.
Lemma C.2. Under assumptions, the following bounds hold with high probability:
(i) ||l|| Õ(1) and |//l -(k-1/2)
(ii) [[F|] Ö(1)
(iii) l|0ll2 = (22 and |0 = Ö(1)
where k, n (dd) as specified by Assumption 4.1
Proof. (i) Follows from standard sub- Gaussaansi norm bounds, specifically 40
Proposition 2.5.2 and
Theorem 3.1.1], in conjunction with Assumption 4.6
(ii) The spectral norm bound on F under Assumption 4.6 is direct consequence of the concentration of the spectral norm for nub-Gaussian random matrices; see [40] Theorem 4.4.5).
(iii) The bounds for follow from the Gaussian concentration of Lipschitz functions 40 Theorem
5.2.2], noting that each element of ğ defined via a Lipschitz transformation of sub-Gaussian variables.
18

---

Next, we derive high-pnn bygh-probbbility norm bounds for key quantities appearing in the residual term A of the decomposed gradient expression in (S22). Specifically, we consider the terms a' (FHH) o' (F (F and w² of (
Lemma C.3. Under our assumptions, the following bounds hold with high probability:
(i) |d H 01().
(ii) [o Okk).
(iii) oFFHoFFFF H| = 112-1/1)
where k= (d.), specified by Assumption 4.1
Proof. Condition on H. By construction, the rows of F are independent centered Gaus- sian vectors (with generally anisotropic covariance inherited from ( The func- tion (+) is bounded and Lipschitz by Assumption 4.8 so we apply Gaussian concentration for Lipschitz functions [40] Theorem 5.2.2]. This gives us the rows of o' (FH") as independent sub-Gaussian vectors with sub-Gaussian norm proportional to || [ || Tr(Cov (vec(Hz)) -1/2 with Tr(Cov(vee (e))) k by Assumption 4.4 Finally, employing the spectral norm bound for matri-
ces with independent sub-Gaussian rows 39 Theorem 5.39 and Eq. (5.26)), and then removing the conditioning, yields the stated bound.
(ii) We bound the full derivative term as
| )ll & (I + a: 1x|]ixx
where a := on (=)] and 1xx denotes a matrix of all ones. Since | || 1l2 = /kn = O(k)
under Assumption 4.1 the result follows from part (i).
(iii) Similar to (i), using Gaussian concentration of Lipschitz functions 40 Theorem 5.2.2] along
with [//||2 = (1) from Lemma C.2 gives the claim.
With all necessary tigh-probability bounds in place, we derive bound on the norm of the residual term appearing in the gradient decomposition. Recall:
| 2
1
wÿy
o (FH
H
(S24)
nvk
nk
2
1
wy T
H
2
www" a (FF
|
nvk
(S25)
- nk
o oF 11
1l [
1l
nvb
(S26)
-(k
(S27)
which holds with high probability, where B E (0, 1) satisfies [[Cov(vec(Hz)) (2))) Ok- olk2) under Assumption 4.5
and
To derive this bound, we first apply the triangle inequality, and then use the identity for Hadamard products:
ab C diag(a) diag(b).
which allows to factor out the vectors and simplify norm computations. The final result follows by applying high-probability bounds established Lemmas C.T
and C3
Similarly, we can bound the norm of the spiked term uv::
I ||2 = || nVb
H
1110
O(k-1/2),
(S28)
which also holds with high probability.
19

---

Taken together, the bounds in (S27 and (S28) confirm that the leading contribution to the gradient matrix arises from the rank-one term u whereas A constitutes negligible residual. This completes the proof of Lemma
D Conditional Gaussian equivalence
With the results from the preceding appendices, we now establish new result: a conditional Gaussian equivalence, which will play key role in the proof of Theorem (4.12 To formulate this result, we first define subspace relevant for the conditioning in the equivalence argument.
Lemma D.1 (Decomposition of Fvee Hz) conditioned on data source s). Let be fixed data source. Define the subspace
S = (0, Ys.1. Ys,2 .... Ys.r.):
where v is the spiked direction from Lemma 4.11 and Ys.1, YYr, are the spiked directions of Cov(vec( H)) as specified in Assumption 4.4 Let P and P1 denote the orthogonal projection matrices onto SSan(S) and its orthogonal complement, respectively. Then,
Fvec( Hz) (F
n (
+ FSK..
(S29)
where K. := () vec Hz).
Proof. This decomposition follows directly from the definition of orthogonal projections and the linearity of matrix multiplication.
Lemma D.1 yields decomposition of the transformed hidden state Fvec Hz) into components: a "bulk" term that conditionally Gaussian, and a structured term aligned with the low-rank subspace S induced by spiked directions. Crucially, this implies that, conditional the coefficient vector Kg, the transformation Fvec H 2) is sub- Gaussian. This motivates the following conditional Gaussian equivalence theorem, which approximates the nonlinear feature map a Fvec (v) with Gaussian counterpart conditioned on k.).
Theorem D.2 (Conditional Gaussian equivalence), Under the assumptions in Section 4.1 define the nonlinear feature mup (vvec(Hz)) = a ((Fvec(Hz)) and lei projection onto the orthogonal complement of the subspuce S defined in Lemma D. Then, the following conditional Gaussian feature map is equivalent to up Hz)) in terms of TCL. error, when conditioned on data source s and alignment vector tigi
P1 vec( H be the
(vec( (z))8, = k.)) + V. 10, (s8, <s,, 1/2
(S30)
where g ~ N(0, I). and
v(s, k..) = E of (vvec (Hz)) 8, kks
V (s, K..)) = E (fvecc(H)))) 200"² 8, K/
(s,, R.)) = Cov (o(Fvec( (H))e(((H))))))))))))))))))))(
8, K.s)
T (s K. (S,
Proof. The proof strategy follows the standard approach to Gaussian equivalence random features 10
developed in B1 [|]00
[3]. In particular, it is sufficient to establish central limit theorem (CLT)
for the bulk component (F n A oo conditional on s, Ka. Once this CLT is in place, the remainder of the proof mirrors that of [3 and is omitted here for brevity.
Conditional CLT K. E R
For any Lipschitz function Ç R22 - IR, for all s E (O....,, S- 1} and for all
lim sup E [c (w7 (vvce(eeeee (vvec(Hz))) (² 8, K5. E [s(( (15 ( (Hz)). {2x) 8, K/s
0,
d.k of
(S31)
20

---

where the supremum is taken over it E fw E R |wlle - 0 [|w)l O for some >> and & Rd satisfies || = C/v] s)ll Here, k d - oo such that k/ k/d R. as specified in Assumption 4.1
This conditional CLT establishes the equivalence of the original and Gaussian feature maps, (Hve)(H)) and vvvec ( z); K²), in terms of their behavior under any test function Ç. condi- tional on and K., Notably, the supremum ensures robustness to variation in task vectors, accounting for worst-case alignment.
We begin proving the conditional CLT by recalling the decomposition from Lemma D.I
Evee Hz) = (F + +) Pec Hz) + FSKs.
Define r == + (1) P1 veel (z and = FSx, so that Fvee Hz) = r + The random vector r is approximately Gaussian because:
7 ||| - asymptotically (see Appendix C
Pi an orthogonal projection (|) ||2 = 1),
Lemma shows that Fvec(Hz) converges Gaussian distribution,
Assumption 4.4 ensures P1 eliminates spiked directions, leading to r - NO Ix).
Consequently, we may write r & Fq. where _ denotes equivalence in distribution, FF is a random feature matrix satisfying the conditions in [21 and q (O, T). Meanwhile, is deterministic conditional on (F, K.). Define the neuron-wite conditional activation:
)gll)...) (f) = o, (, a oy),
j=
This representation allows us to view the feature map as a collection neuron-speen speroon activations with fixed shifts cj. Importantly, the CLT for random features in [2 applies even when activations vary across neurons, a point further supported by related results in 1113.
Utilizing Corollary 10 and applying the central limit theorem for heterogeneous neuron activations [1] Theorem 2]. we obtain convergence in distribution of ((veec(H2) to Gaussian approximation (vcc s, r..) Furthermore, the odd activation function assumption required be omitted here, since both feature maps share the same conditional covariance structure. See for related proofs and further technical details.
21 can
31
10 13
So far, we have established an equivalence with respect to the training error. While not stated explicitly earlier, this training error corresponds to that ridge regression applied to the second-layer weights in our two-stage training setup. To extend this asymptotic equivalence to the ICL error defined in (8) which is our main object of interest-we require an additional technical condition, formalized below.
Assumption D.3. Consider a perturbed optimization objective:
Ta(c) = min Σ WEC,
Veli
TH ((vecc(Hz)) + 1 cEww,, Vi
1
(S32)
j=1
where R is scalar perturbation parameter, and E(w) denotes the ICL error from ( associated with second-layer weight vector w. The constraint set C. is the constraint set in the conditional CLT })} and it is defined as Cs == (w E R* |w||| O(1), |w]] I oo = OOk -(k-()) for some > 0. Then, there exists a constant c* > O such that, for all E [-c", the function Tn(c) converges pointwise to a limiting function (cc), which is differentiable at 0.
Although this assumption may appear somewhat artificial at first glance, enables the use of convexity-based arguments in establishing generalization error equivalence (ICL errors in our case), as seen in prior work on Gaussian universality [I]] Assumption 5]. Similar assumptions are also employed in related studies asymptotic equivalence 221 31 13] Importantly, this condition arises primarily as technical requirement of the proof method, rather than as limitation on the broader applicability of the Gaussian equivalence results established in this work.
21

---

With Assumption D.3 in place, we are now able to extend the asymptotic equivalence result to the ICL error, thereby completing the proof of the conditional Gaussian equivalence, as established in 13
Theorem D.2 establishes an asymptotic equivalence between two feature maps with respect to the ICL error, under the condition that their first two conditional moments match. This result implies the existence of an equivalent (conditional) Gaussian replace the original feature map without affecting ICL performance. Leveraging this equivalence, now prove that the Transformer model with nonlinear MLP is asymplotically equivalent to polynomial model, as formalized in Theorem 4.12
w² (vvec(H2)) ); 8, K. that can
The full proof is presented below.
E Equivalent polynomial model
We aim to establish the asymptotic equivalence between Transformer model with nonlinear MLP and the polynomial model described in Theorem Our strategy relies on the conditional Gaussian equivalence result stated and proved in Appendix D (Theorem D.2). This result asserts that two models with matching conditional means and covariances yield the same ICL error (8). Thus, to prove equivalence, it suffices to demonstrate that the first two conditional moments of the two models coincide.
Fix data source index s, and recall the orthogonal decomposition with respect to the subspace defined by the matrix S:
Fvee Hz) (F + nA)PP vec(Hz) FSKs:
(S33)
where K. == (S" vec( (zz), as established in Lemma D.1 The first term, (F -
n P vec (zc(H)), behaves asymptotically like Gaussian random variable (see the proof of Theo- rem D.2), while the second term, FSK., is deterministic conditional on 8, and Kis.
According to [13 Theorem 4], if conditional Gaussian equivalence holds and the deterministic
component FS,, vanishes at rate of Ö(k-1 for some 8 > O, then there exists a finite polynomial degree p such that the polynomial activation 6-(-) (as defined in Theorem 4.12) yields equivalent generalization performance to that of the original nonlinear activation ø o.). In their proof, the vanishing nature of the deterministic term is used to show that (+)) and a induce the same first
two conditional moments, thereby ensuring asymptotic equivalence in generalization (ICL.) error.
Similarly, in our setting, the vanishing nature of the term F FSK, together with Assumption 4.8 im- plies the following bounds:
E
(Fveo (2)))
8, K.
(Fve((((2))) 8, K2x
= =o(1,
(S34)
12
E
(Fvec H ))o""
8, E 2[ Vev(Heee))) (H2))
K/g
=(1),
(S35)
Cov (a( (o(Fvect H2))
8, 8, Cov (F Hz)
8, Kg)
o(1),
(S36)
where o := vee Hz). These bounds suffice to prove the equivalence between the original activation (.) and the polynomial approximation âp(+). (.), explained by (13). Therefore, it remains to verify that
Sk,1 Ō(K-)) for all k))
(S37)
for some > O, where f. denotes the i-th row of F.
To analyze 1s. Sk,,, begin by expanding the expression:
lf. Sk, = f. S(S vec Hz)
(S38)
<
vec H 2) + 9
(S39)
where we used the decomposition f. f: 79:, and g: is the i-th row of the gradient matrix G defined in (7). The first term corresponds to the contribution from the randomly initialized feature matrix F, while the second accounts for the effect of the gradient update.
22

---

Finally, invoking the gradient decomposition and the corresponding bounds from Appendix C together with the fact that P S(S (onss))--30 S has rank (11) and Assumptions 4.4 4.5 we conclude that both terms in (S39) vanish at the desired rate. In particular, the scaling in Assumption 4.5 the low-rank structure P. and the spiked-covariance in Assumption en4ees that the contributions of remain controlled, which
the 17-weggttd gradient term and the spiked directions in Cov(vec completes the proof.
In summary, our results indicate that although the output of the attention layer, vec H, exhibits a heavy-tailed distribution 41 the application of the random matrix F attenuates the heavy tails. Moreover, training the first-layer weights F with single gradient step has negligible effect on the analysis. Together, these insights allow us to transfer known results from supervised learning with two-layer neural networks to the in-context learning setting of Transformers with nonlinear MLPs. This connection opens promising avenues for analyzing complex, realistic models using Gaussian equivalence techniques.
F Extension to inputs with non-zero mean
In the main body of our proofs, we have assumed zero-meat inputs, i.c., Mr.. = 0 for all data sources to streamline the exposition and focus the core aspects of the Gaussian equivalence framework. However, in practical scenarios, input data often has non-zero mean, necessitating generalization of theoretical results handle such cases.
The key observation enabling this extension is that the mean vector Ma.. of each Gaussian data source introduces structured, low-rank perturbation in the representation of the attention output. Formally, we have:
Hz (#++1 5- p22s) ,l,,bb',
(S40)
which separates the stochastic (zero-mean) and deterministic (mean) components. Consequently, the vectorized attention output vee H2)) becomes
but its non-central second moment still captures the combined statistical behavior of both components.
To rigorously extend our results this more general setting, the following changes are required.
1. Substitution of covariances with second moments:
In all analytical steps where the covariances of x, is and vec(Hz) appear, need to replace with the corresponding non-central second moment:
Cov(vec( Hz))
E vec H vec Hz)
Σx,, Cov((rs(s)
E [a...: [s].
This replacement maintains the validity of our results since the spike due the mean treated analogously to spikes in the spiked covariance model Assumption 4.4
2. Bounding mean- -nduued terms in gradient decomposition:
The primary technical adjustment required in the proof of the gradient decomposition
(A (Appendix (
The decomposition must now include additional cross-terms involving Hrr,,, which must be carefully bounded using concentration inequalities and structural assumptions on the data distribution (e.g., bounded mean norm, low-rank behavior).
These bounds ensure that the mean-induced components do not asymptotically domi- nate or distort the gradient behavior established under the zero-mean assumption.
By implementing the above modifications miaiic using non-central second moments and bounding mean-induced gradient terms-we can extend our theoretical guarantees to the case of non-zero- mean Gaussian inputs. This generalization not only reinforces the robustness of our analysis but also broadens its relevance to real-world applications, where input means rarely zero in practice.
23

---

G ICL errors per-source in the setting of Figure 2
For the sake of completeness, we illustrate the ICL errors per-source (corresponding each source) in the setting of Figure 2 Mathematically, the per-source error defined as [[ - 8) = 8)
where source indicator E {0. 1] since consider settings with different data sources. For all of the figures below, on the left, the ICL error corresponding to source 0, i.e., E[ 0) 0).
is plotted while the ICL enror for source 1, i.e., E[(ye+1 2 s 11. is shown the right. Figures 4 5 6 iillsstrae the per-source ICL errors in the settings of Figure (a), (b), and (c), respectively. Namely, Figure shows the changes in the per-source ICL errors when varying the input covariance and Figure 5 displays the effects of changing the task covariance on the per-source ICL errors, while Figure b depicts the per-source ICL. errors for different noise levels. In all of these cases, the first data source (source 0) is kept fixed while the aforementioned properties of the second data source are modified, as detailly explained in the caption of Figure 2 The results indicate two primary points. First, per-source ICL eπrоrs for the Transformer (approximately) match those of the equivalent model, indicating that the equivalence specified by Theorem 4.12 is useful for studying per-source ICL errors as well. Second, the trends for the ICL errors (which is the average of the per-source errors) in Figure 2 are also observed for each of the per-source errors, providing further evidence for our conclusions.
Source 0
Source
Innndemeertt MLP
Equivalont Medel
04
0.4
Miking Rallo
1C1 Lover
0C3. Limer
0.3
0.3
D-4
022
0.2
0.1
0.1
legell
(
Figure 4: Per-source ICL errors in the of Figure ((a): impact of varying input covariance of source on the per-source ICL errors.
Source
Source
0.40
0.40
Miking Ratio
0.0
0.30
0
KCL
0:25
0.25
Rigutcalens Madid
015
0.
1/2
0.15
0.
1.0
Figure 5: Per-source ICL errors in the case of Figure 2 b): effect of altering the task covariance of source on the per-source ICL errors.
Source
Source
0.
0.9
Tranfoonner with MLP
Equivalent
103. 1300
0.7
D.2
0.6
04
0.4
10
Figure 6: Per-source ICL errors in the case of Figure 2 c): result of changing noise level of source on the per-source ICL emrors.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
They correctly reflect the paper's content and contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are explained when describing the problem setting.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "L "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:[Yes)
25

---

Justification: Full set of assumptions are given in the main text while the proofs are provided in the appendix (supplementary).
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Experimental (simulation) setting and all of the details are provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification: The code for the experimental results can be found at: https: com/KU-MLIP/DeEE
by- TognTrranforrrss
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
ssicnpolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All experimental (simulation) details are explained in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Error bars are not needed due to the nature of our results. Instead, we plot the average over a reasonable number of Monte Carlo runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
27

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: We have run our simulations on a modern laptop. Thus, there is no significant resource requirement report.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
Our research conform with the NeurIPS Code Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
As theoretical work, this paper does not have direct impact on the society,
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: As theoretical work, this paper do not poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: As an licensed asset, our work only use the Transformer model, which is properly cited/credited to the owners.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
29

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy for what should or should not be described.
31


Reviews:
[{"review_id": "5UhzfVR9L2", "reviewer": "Reviewer_JzvP", "summary": "This research paper investigates in-context learning (ICL) in pre-trained transformer models with non-linear multilayer perceptron (MLP) layers, focusing on non-linear regression tasks derived from diverse datasets with heterogeneous input, task, and noise distributions. It builds upon previous efforts that analyzed ICL dynamics in the limited context of linear regression using attention-only models. Leveraging prior studies on feature learning in two-layer neural networks following a single gradient step—particularly in structured Gaussian mixture settings—this work extends those insights to transformer architectures incorporating non-linear MLP layers. Additionally, it examines the interplay between data mixing and feature learning, demonstrating the significance of structured task distributions in facilitating effective feature representation. The efficacy of ICL is assessed using a uniform mixture of data sources. Empirical simulations reveal that the performance of non-linear MLP-based transformers closely approximates the Hermite polynomial surrogate, even in moderate-dimensional settings. Further analysis explores the impact of mixing structured and unstructured data on overall model performance.", "questions": "1) Prior work has shown that linear transformers exhibit gradient descent-like behavior during in-context learning. Does the equivalence established in this study imply a similar mechanism in the linear setting? This will help connecting the results prior efforts to explain ICL.\n2) For multi-layer transformers with non-linear MLPs, does the theoretical connection to Gaussian Universality continue to hold?  Any simplifying assumptions to explain ICL for multi-layer transformer models with MLP? \n3)  Can the analytical techniques used here be extended to explain ICL behavior in models with multiple layers of bi-directional LSTM or xLSTM architectures?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1) Provides a detailed analysis of the learning dynamics of in-context learning (ICL) in transformer models with non-linear MLP layers, applied to non-linear regression tasks.\n2) Establishes the asymptotic equivalence between transformers with non-linear MLP layers and Hermite polynomial representations, in terms of ICL error—offering a tractable analytical model for studying such architectures.\n3) Supports the theoretical framework with empirical evaluations, demonstrating close alignment between the analytical model and simulation results even at moderate dimensions.\n4) Empirical results further indicate that the quality of data sources—specifically, those with structured covariance—plays a significant role in reducing ICL error\n5) The appendix also includes a useful extension to the data sources having non-zero mean.\n\nWeaknesses:\n1) Prior work has shown that linear transformers exhibit gradient descent-like behavior during in-context learning. Does the equivalence established in this study imply a similar mechanism in the linear setting?\n2) For multi-layer transformers with non-linear MLPs, does the theoretical connection to Gaussian Universality continue to hold?  Any simplifying assumptions to explain ICL for multi-layer transformer models with MLP?\n3)  Can the analytical techniques used here be extended to explain ICL behavior in models with multiple layers of bi-directional LSTM or xLSTM architectures?", "was_revised": true, "final_justification": "I carefully reviewed the authors’ responses to my comments. My observations regarding their rebuttal are as follows:\n- The authors acknowledged the validity of the concerns raised and responded in general terms. However, the responses did not offer any substantive analysis beyond what was already presented in the original manuscript.\nRegarding the authors’ responses to comments from other reviewers, I note the following:\n- On real-world applicability: The authors present a comparison using a simple sentiment analysis task. While the results demonstrate a reasonable alignment between their model and a standard Transformer, it is important to recognize that this constitutes a toy example and may not reflect broader applicability.\n- On the discussion of linear Transformer models and two-layer MLPs with frozen layers: The authors have made an effort to incorporate and contrast relevant literature that was previously omitted. It is worth emphasizing that this paper represents one of the early attempts to explain in-context learning (ICL) using simplified model architectures.\nOverall, the paper introduces tools and methodologies that contribute to understanding the learning dynamics of ICL. These techniques may serve as a foundation for future work. Based on these observations, I will maintain my current score.", "rebuttal": "Thank you for your thoughtful and constructive review. We are pleased that you recognized key strengths of our work, including the detailed analysis of in-context learning (ICL) with Transformers involving nonlinear MLPs, the asymptotic equivalence to Hermite polynomial representations, and the strong alignment between theoretical and empirical findings. Below, we address each of your insightful questions:\n\n> 1. Prior work has shown that linear transformers exhibit gradient descent-like behavior during in-context learning. Does the equivalence established in this study imply a similar mechanism in the linear setting?\n\nYes, our results are fully compatible with prior work on linear Transformers. Specifically, in the linear case, the analysis reduces to the known gradient descent-like behavior for ICL. The novel contribution here lies in extending this understanding to Transformers with nonlinear MLP heads. We show that such architectures yield predictors equivalent to finite-degree Hermite polynomial surrogates—thereby enriching the class of functions learnable in-context. This extension provides a principled explanation for improved ICL performance on nonlinear tasks, particularly in mixed or structured data scenarios.\n\n> 2. For multi-layer transformers with non-linear MLPs, does the theoretical connection to Gaussian Universality continue to hold? Any simplifying assumptions to explain ICL for multi-layer transformer models with MLP?\n\nThis is an excellent question and a natural direction for future research. While the extension to multi-layer Transformers is conceptually plausible, our current analysis relies on the simplification afforded by a single gradient-like step (for the first layer of the MLP). This enables precise derivation of asymptotic equivalence results under Gaussian universality. Multi-layer architectures introduce additional complexity due to inter-layer interactions and compositional nonlinearities. Extending our analysis to this setting would likely require additional assumptions—e.g., layer-wise independence or structure—to preserve analytical tractability.\n\n> 3. Can the analytical techniques used here be extended to explain ICL behavior in models with multiple layers of bi-directional LSTM or xLSTM architectures?\n\nExtending our techniques to recurrent architectures such as bi-directional LSTMs or xLSTMs is non-trivial, primarily due to their architectural differences—especially temporal recurrence and gating mechanisms. Nonetheless, the underlying analytical principles in our work, including Gaussian universality and polynomial surrogate modeling, offer a potential foundation for exploring such extensions. Achieving rigorous theoretical insights in these settings would, however, require new developments to handle the temporal and stateful dynamics unique to recurrent models.\n\n---\n\nOverall, we sincerely appreciate your positive assessment and the thought-provoking questions, which we believe point toward exciting future research directions."}, {"review_id": "RUxRyyvkDW", "reviewer": "Reviewer_G3y6", "summary": "This paper presents a theoretical analysis on how Transformers with nonlinear MLP heads perform in-context learning when trained on heterogeneous data. The authors establish a tractable analytical framework by proving an asymptotic equivalence between complex models and simpler polynomial predictors. This result is achieved through an application of Gaussian universality theory and Hermite expansions under high-dimensional, proportional-limit asymptotics. The function surrogate enable derive closed-form expressions for ICL error, offering novel insights into how data mixing weights, noise levels, and covariance structures collectively determine learning outcomes.", "questions": "1. How feasible to extent the analysis to nonlinear attention, perhaps via kernelized approximation?\n2. How do the dynamics of the polynomial surrogate change empirically after multiple gradient steps?\n3. Could the data mixing heuristics be validated on a real-world benchmark, such as a multilingual corpus where mixing weights control language fractions?", "limitations": "The analysis is limited to regression tasks, leaving its applicability to classification, or language modeling tasks unclear", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n- This work fills a long-standing gap in the literature by incorporating both nonlinear MLP components and multi-source data mixing—two ubiquitous features of real-world models that prior linear-only analyses could not address. \n- The analysis is rigorous and elegant. The authors skillfully coupling Gaussian universality with a rank-one decomposition of the gradient to analyze the feature learning dynamics after a single update step in first MLP layer.\n- The results are insightful. By quantifying how factors like mixing weights, noise levels, and covariance structure shape ICL performance, the work provides clear data curation heuristics—for instance, prioritizing low-noise, structured covariances data sources. \n- These theoretical predictions are validated through simulations that match the theory.\n\n**Weaknesses**\n- The linear-attention simplification limits the practical applications. The nonlinear attention mechanism is fundamental to the success of modern Transformer. There is no evidence or discussion on how feasible the analysis can be transfer. \n- Real feature learning requires multiple optimization steps, and the dynamics likely change as optimization progress. The current analysis rely on one-step gradient on the first MLP layer is too restrictive.\n- Requiring attention output covariances to have identical traces across all tasks and admit low-rank decompositions seems unlikely to hold in practice for natural language or vision contexts. The authors note this but do not quantify robustness.\n- The motivation of assumption 4.5 is unclear and not argued. \n- The experiments are exclusively rely on synthetic Gaussian data without testing on real benchmarks.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thoughtful and thorough evaluation. We are particularly grateful that you recognized the novelty and rigor of our approach, the integration of nonlinear MLPs and data mixing, and the empirical alignment with our theoretical predictions. Below, we address each of the concerns and questions raised:\n\n> The linear-attention simplification limits the practical applications. The nonlinear attention mechanism is fundamental to the success of modern Transformer. There is no evidence or discussion on how feasible the analysis can be transfer.\n\n>  How feasible to extent the analysis to nonlinear attention, perhaps via kernelized approximation?\n\nWe acknowledge that the use of linear attention is a simplifying assumption and agree that extending the analysis to nonlinear (e.g., softmax) attention is a valuable direction. We chose linear attention to preserve analytical tractability, which enabled our novel integration of nonlinear MLP heads and data mixing analysis. That said, recent empirical work (Han et al., NeurIPS 2024) suggests that carefully designed linear attention variants can approach softmax performance—supporting the continued relevance of the linear setting for both analysis and applications.\n\nYour suggestion to explore kernelized approximations for nonlinear attention is especially insightful. We agree this is a promising path for extending theoretical tools to richer architectures, though it introduces additional technical challenges and remains open for future work.\n\n*(Han et al., NeurIPS 2024): Bridging the Divide: Reconsidering Softmax and Linear Attention*\n\n> Real feature learning requires multiple optimization steps, and the dynamics likely change as optimization progress. The current analysis rely on one-step gradient on the first MLP layer is too restrictive.\n\n> How do the dynamics of the polynomial surrogate change empirically after multiple gradient steps?\n\nThis is an important point. Our analysis focuses on a single (but large) gradient step to allow rigorous characterization of the feature dynamics, similar in spirit to prior work analyzing supervised learning with MLP settings (e.g., references [3, 9]). Empirically, we observe that this one-step update captures key aspects of feature learning and already yields strong alignment with theory. Extensions to multiple steps are indeed plausible—in fact, under appropriately scaled step sizes, our framework can accommodate a small number of steps without substantial analytical disruption. Nonetheless, a full multi-step theory, especially with nonlinearities, remains an exciting and non-trivial direction for future work—even for simpler MLP-only settings (see, e.g., reference [13]).\n\n> Requiring attention output covariances to have identical traces across all tasks and admit low-rank decompositions seems unlikely to hold in practice for natural language or vision contexts. The authors note this but do not quantify robustness.\n\nWe appreciate this concern. The assumption of identical trace and low-rank structure is indeed primarily for analytical convenience, helping us isolate the core dynamics of interest. However, our experiments suggest that the predictions remain qualitatively robust even when this assumption is mildly violated. We will explicitly clarify this point in the revised manuscript and highlight the robustness behavior under partial relaxation.\n\n> The motivation of assumption 4.5 is unclear and not argued.\n\nThank you for pointing this out. Assumption 4.5 controls the relationship between the step size and the covariance norm, ensuring that the feature update is nontrivial but remains within the regime where our asymptotic analysis is valid. While we briefly mention this in the current version, we agree that the rationale can be better justified. We will revise the manuscript to clearly articulate the role, motivation, and practical relevance of this assumption.\n\n> The experiments are exclusively rely on synthetic Gaussian data without testing on real benchmarks.\n\n> Could the data mixing heuristics be validated on a real-world benchmark, such as a multilingual corpus where mixing weights control language fractions? \n\n> The analysis is limited to regression tasks, leaving its applicability to classification, or language modeling tasks unclear\n\nThese are all excellent suggestions. We focused on synthetic Gaussian settings to enable precise theoretical modeling and tight control over key variables like mixing weights, noise, and covariance. That said, we agree that empirical validation on real-world problems—such as language modeling tasks with multilingual corpora where mixing is naturally interpretable—is an important next step. We will include a discussion of this in the revised manuscript as a clear future direction. Finally, we would like to note that while our presentation focused on regression tasks for simplicity, our theoretical results trivially apply to classification (if the target/label function is chosen appropriately) as well. We will mention this point in the revised paper.\n\n---\n\nOverall, we are grateful for your positive review and constructive suggestions. Your feedback has already helped improve the clarity and potential impact of our work, and we look forward to incorporating these insights into the final version."}, {"review_id": "dN5NsKNPlJ", "reviewer": "Reviewer_Z4dn", "summary": "The paper consists of theoretical and empirical parts. In the theoretical part, they consider 1 layer transformer (linear self-attention+2 layer MLP with non-linear activation) (equivalently written as in Equation (5)) where (equivalent representation of) all self-attention + first MLP layer parameters are trained with a single gradient descent step, and after that, second layer in MLP is taken to be ridge regression loss minimizer on separate data. The main theoretical result is that this transformer is asymptotically (Assumption 4.1) equivalent to structured polynomial predictor (Equation 13), when observing ICL error on ICL non-linear regression data (from mixture of sources). Empirically, they confirm this equivalence by simulations, and experiment with the impact of model size and data characteristics to successfully learned ICL.", "questions": "1. Can you elaborate on meaning of equation 13? Is it not the same as Equation (5) but with a polynomial activation (and updated F)? What is the benefit (apart maybe theoretical convenience) of showing that we can replace activation function with a polynomial?\n2. In the sentence on lines 130-131, there is no mention of transformers being able to capture non-linearities in data with a non-linearity on self-attention (such as usually used softmax). Can you say why, or mention it in this sentence?\n3. Can you justify more the focus on the average task error across all sources in the evaluation (i.e. the introduced ICL Error)? This makes the train and test distribution different. Instead, I would perhaps consider per source error (which is also tractable in the setting of the paper – 2 sources). Have the authors looked into this already?\n4. In line 168 (potentially elsewhere as well), it is not specified which matrix norm is being used?\n5. On line 282, where does this intuition come from?\n6. Have you tried having more than 2 data sources in the simulations?", "limitations": "There is a limitation introduced with the training procedure of the model. It isn’t clear to me how this limitation would  be solved if we turn to training a transformer with GD or Adam.", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "## Strengths: \n\nThe theoretical paper has nicely set notation, well explained data and model setup. The theoretical assumptions and results are clearly stated. They follow some nicely cited previous work. The experiments are also clear.\n\n## Weaknesses:\n\nThe goal of the paper would be to bridge the theoretical setup closer to the practical use cases of transformers, by incorporating nonlinear MLP after linear self-attention. Apart from the theoretical assumptions in section 4.1, there is also the somewhat weird training procedure of the MLP. Namely, training one linear layer with only one step of GD and fitting the other with ridge regression is also very far from practice, but this divergence is not discussed.\n\nI'm not sure of the significance of the main theoretical result (Theorem 4.12) is. This may be my misunderstanding, but it would be good for the authors to comment on this more/discuss this in the paper. Namely, how this understanding may go beyond existing/classical intuition (which is what the current Conclusion reads as).\n\nSimilarly, the empirical part seems to lack a significant/original conclusion in my opinion:\n- Figure (1) isn’t so surprising to me, namely conclusion on line 256 (nonlinear > linear)\n- In Figure 2a), I don’t agree that the conclusion is that a more structured covariance matrix leads to more ICL (line 269). It can happen that ICL on the first (easier) source is easy to learn and the model has acquired that ability. In addition, ICL on structured data (second source) is harder, and the model learns it only with more examples from that source. Then the decrease in the ICL error would be justified (when increasing $\\rho$), but it wouldn’t be the case that more structured data leads to better ICL. Another reason to believe this is the case is to observe the $\\rho=0.6, 0.8$ cases: The ICL error decreases and then increases again as structure is increasing. I suspect that plotting the error per source may reveal that the error on the first source is slightly increasing, while on the second source is decreasing a lot—together making the U-shaped curve.\n- Conclusion on line 273 is also not surprising to me (lower noise leading to performance gains).\n- I’d plot some more, higher values than 2 in Figure 3b) to check that the claimed trend of decreasing is true.\n\nI would explain Equation (5) in more detail, e.g. It is not stated that $\\Gamma$ and the first linear layer in MLP are incorporated in $F$.", "was_revised": true, "final_justification": "I appreciate the authors' engagement, but my concerns still persist, as noted in my response to the rebuttal. I will retain my current rating, though I understand that it is lower than the other reviewers' scores. I will not \"fight\" to reject this paper, but would like to keep my opinion expressed (hence not changing the score).", "rebuttal": "Thank you for your detailed review and the thoughtful feedback. We appreciate your recognition of the clarity in our theoretical exposition, assumptions, and experimental design. Below, we respond to your concerns and questions point-by-point:\n\n> The goal of the paper would be to bridge the theoretical setup closer to the practical use cases of transformers, by incorporating nonlinear MLP after linear self-attention. Apart from the theoretical assumptions in section 4.1, there is also the somewhat weird training procedure of the MLP. Namely, training one linear layer with only one step of GD and fitting the other with ridge regression is also very far from practice, but this divergence is not discussed.\n\nWe agree that our training procedure differs from standard end-to-end training pipelines. However, this two-phase scheme—gradient descent on the first layer followed by ridge regression on the second—is well-motivated by prior theoretical work on MLPs [3, 9]. It can be interpreted as an idealized form of a two-phase learning schedule with layer-wise learning rate separation: a large learning rate applied briefly to early layers, followed by fine-tuning of later layers. This separation enables tractable analysis while preserving meaningful feature learning dynamics. We will expand on this interpretation in the revised paper.\n\n> I'm not sure of the significance of the main theoretical result (Theorem 4.12) is. This may be my misunderstanding, but it would be good for the authors to comment on this more/discuss this in the paper. Namely, how this understanding may go beyond existing/classical intuition (which is what the current Conclusion reads as).\n\n> Can you elaborate on meaning of equation 13? Is it not the same as Equation (5) but with a polynomial activation (and updated F)? What is the benefit (apart maybe theoretical convenience) of showing that we can replace activation function with a polynomial?\n \nThank you for raising this important point. Theorem 4.12 establishes an asymptotic equivalence between the ICL behavior of a Transformer with an MLP and that of a polynomial model. This result is significant for several reasons:\n\n1. **Analytical tractability**: Equation (13) admits a much simpler analysis, allowing precise characterization of ICL error behavior of the Transformer under data mixing and feature learning.\n2. **Interpretability**: It sheds light on the function class learned by the Transformer—specifically, that it effectively learns a low-degree polynomial approximation of the task function.\n3. **Nonlinearity design**: This equivalence opens the door to optimizing nonlinearities in MLPs through polynomial surrogate analysis.\n\nThese insights extend beyond classical intuition and align with recent trends in Gaussian equivalence literature [e.g., 13, 18]. We appreciate the suggestion and will clarify the conceptual impact of this result in the camera-ready version.\n\n\n> Similarly, the empirical part seems to lack a significant/original conclusion in my opinion: \n\nWe respectfully disagree with you in this regard. While some findings may align with prior intuitions, we believe our contributions go beyond merely confirming them. For example, the relationship among in-context learning, task covariance, and feature learning (illustrated in Figures 2b and 3b) is quite interesting and original.\n\n> - Figure (1) isn’t so surprising to me, namely conclusion on line 256 (nonlinear $>$ linear)\n\nThe observation in Figure 1 is meaningful because nonlinearity alone does not guarantee better ICL performance. As shown in prior work (e.g., [18]), nonlinear models can underperform linear models if feature learning is ineffective. Our result demonstrates that even one-step feature learning enables the Transformer to outperform its linear counterpart—highlighting the utility of the MLP layer in ICL.\n\n> - In Figure 2a), I don’t agree that the conclusion is that a more structured covariance matrix leads to more ICL (line 269). $\\cdots$ Then the decrease in the ICL error would be justified (when increasing $\\rho$), but it wouldn’t be the case that more structured data leads to better ICL. $\\cdots$\n\nRegarding Figure 2a and line 269: we clarify that our conclusion is **not** that more structured covariance *always* leads to better ICL. Rather, increasing the **proportion** of structured data (denoted with $\\rho$) improves ICL. This aligns with your interpretation, and we will reword the corresponding sentence to avoid misreading.\n  \n> - Conclusion on line 273 is also not surprising to me (lower noise leading to performance gains).\n\nWhile the conclusion that lower noise leads to better performance may seem unsurprising in a standard single-source setting, our focus is on how noise interacts with data mixing in multi-source scenarios. In this context, analyzing the role of noise becomes nontrivial and offers valuable insights into how noise levels influence the effectiveness of different mixing strategies—an aspect that, to our knowledge, has not been explicitly studied before.\n\n> - I’d plot some more, higher values than 2 in Figure 3b) to check that the claimed trend of decreasing is true.\n\nFor Figure 3b: we have indeed run simulations for larger values of $\\log_d(\\eta) > 2$ and found that the decreasing trend in ICL error continues, confirming the conclusion presented.\n\n> I would explain Equation (5) in more detail, e.g. It is not stated that $\\Gamma$ and the first linear layer in MLP are incorporated in $F$.\n\nThank you for pointing this out. In response to your feedback, we will revise the explanation around Equation (5) to explicitly clarify that both $\\Gamma$ and the first linear layer of the MLP are incorporated into the definition of $F$.\n\n> In the sentence on lines 130-131, there is no mention of transformers being able to capture non-linearities in data with a non-linearity on self-attention (such as usually used softmax). Can you say why, or mention it in this sentence?\n\nAs explained in the original Transformer paper (Vaswani et al., 2017), softmax (the nonlinearity in the attention) is primarily used to compute normalized weights for value aggregation. While this also makes the model nonlinear, we isolate and study the role of nonlinear MLPs in Transformers for capturing task-related nonlinearity by using linear attention. We will add a clarification in the relevant sentence.\n\n\n> Can you justify more the focus on the average task error across all sources in the evaluation (i.e. the introduced ICL Error)? This makes the train and test distribution different. Instead, I would perhaps consider per source error (which is also tractable in the setting of the paper – 2 sources). Have the authors looked into this already?\n\nGreat question. We have examined per-source ICL error and found that it also behaves in line with our theoretical predictions. However, we focus on the average ICL error to:\n\n1. Have a single evaluation criterion, capturing the overall ICL generalization capacity across mixed-task distributions.\n2. Use a fixed evaluation criterion independent of mixing weights, making it easier to interpret performance across different scenarios.\n\nWe will clarify this rationale in the revised version and may include additional plots in the appendix showing per-source behavior.\n\n> In line 168 (potentially elsewhere as well), it is not specified which matrix norm is being used?\n \nAs noted in the appendix (Notation section), $\\|\\cdot\\|$ denotes the spectral norm and $\\|\\cdot\\|_F$ the Frobenius norm. We will move the notation section to the main text if space permits.\n\n\n> On line 282, where does this intuition come from?\n\nThe intuition is derived from the fact that feature learning enables the model to align with task vectors $\\boldsymbol{\\xi}$, which can be shown via an analysis of one‑step gradient (e.g., see Section 3.2 of reference [3] for a detailed explanation of this behavior in the case of supervised learning with MLPs). Therefore, when task vectors are isotropic (as in Figure 3a), there is little structure to exploit, and feature learning becomes ineffective—resulting in flat ICL error curves. \n\n\n> Have you tried having more than 2 data sources in the simulations?\n\nYes, we have experimented with 3 or more sources and observed similar trends. We chose the 2-source setup for clearer visualization and more direct interpretation of mixing effects. If needed, we are happy to share additional results.\n\n\n> Line 136, fix indexing i and j. Equation 8, remove commas.\n\nThank you for these suggestions. We will revise the text and correct the noted typos accordingly.\n\n---\n\nOverall, we believe our responses address all the concerns raised in your review. We hope that these clarifications, along with your acknowledgment of the paper’s strengths, lead to a reconsideration of your evaluation and ultimately, an increase in your rating. We appreciate your time and thoughtful evaluation."}, {"review_id": "gUu47AGOPZ", "reviewer": "Reviewer_CMBv", "summary": "The authors study in-context learning (ICL) for a *linear-attention* Transformer whose MLP head has two layers: the first is updated by **one gradient step**, the second is fitted.  \nUnder joint high-dimensional scaling, they show (Theorem 4.12) that the model’s ICL error matches that of a suitably chosen degree-p polynomial predictor.\n\nEmpirically (Figures 1–3), they confirm that:\n\n- The nonlinear MLP head beats the purely linear baseline.\n- The polynomial surrogate tracks the MLP’s performance.\n\nTheir framework also enables analysis of “data-mixing” effects, showing that mixtures with **structured covariance** and **low noise** are most effective for ICL.", "questions": "See weaknesses.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written and clear.\n2. The paper is addressing an important problem of the role of MLP in in context learning.\n3. The assumptions are stated clearly.\n\nWeaknesses\n\n1. The paper only considers *linear attention*, which limits its generality. Prior work has shown that for any statistical ICL task, there exists a fixed **non-linear attention** mechanism that achieves Bayes error as the context length tends to infinity, an assumption used in this paper. (Section 5.1 of this paper [^1] )\n\n2. The same work[^1] provides shows evidence that how the MLP enables ICL in one-layer Transformers with any nonlinear attention head. The current paper does not cite this or contrast its contributions, weakening its novelty and context.\n\n3. Both the current paper and prior work[^1] restrict their analysis to one-layer Transformers. However, empirical studies show a significant performance leap from one to two layers, with diminishing returns improvement  beyond two layer. Understanding how ICL emerges across **multiple layers** is crucial, and ignoring this makes the theoretical claims less applicable to real-world models.\n\n---\n\n[^1]: Abedsoltan, A., Radhakrishnan, A., Wu, J., & Belkin, M. (2024). *Context-Scaling versus Task-Scaling in In-Context Learning*. [arXiv:2410.12783](https://arxiv.org/abs/2410.12783)", "was_revised": true, "final_justification": "The authors have convinced me of the distinction between their work and prior studies. Although this work does not fully address the mechanism of in-context learning (ICL), it provides valuable insights in simplified settings.", "rebuttal": "Thank you for your detailed and constructive feedback. We appreciate your recognition of the paper’s clarity, relevance to ICL, and well-stated assumptions. Below we address your concerns and outline how the revised version will incorporate the suggested clarifications and citations.\n\n> 1. The paper only considers linear attention, which limits its generality. Prior work has shown that for any statistical ICL task, there exists a fixed non-linear attention mechanism that achieves Bayes error as the context length tends to infinity, an assumption used in this paper. (Section 5.1 of this paper [^1])\n\nWe agree that nonlinear attention plays a critical role in real-world Transformers. Our choice to focus on trainable linear attention is intentional, as it enables rigorous high-dimensional asymptotic analysis, which is currently out of reach for nonlinear attention mechanisms. This analytical tractability allows us to isolate the role of the nonlinear MLP head and precisely quantify its impact on ICL performance. That said, recent empirical work (Han et al., NeurIPS 2024) suggests that carefully designed linear attention variants can approach softmax performance—supporting the continued relevance of the linear setting for both analysis and applications.\n\n*(Han et al., NeurIPS 2024): Bridging the Divide: Reconsidering Softmax and Linear Attention*\n\n> 2. The same work [^1] provides shows evidence that how the MLP enables ICL in one-layer Transformers with any nonlinear attention head. The current paper does not cite this or contrast its contributions, weakening its novelty and context.\n\nWe thank the reviewer for pointing out this relevant work (Abedsoltan et al., 2024). We will cite it in the revised version and highlight the following contrasts. Specifically, Abedsoltan et al. (2024) study a simplified Transformer architecture with a fixed (non-trainable) nonlinear attention mechanism—omitting key, query, and value weights—and an optional MLP head. They show that this setup converges to a optimal (consistent) estimator in the infinite-context limit via a Hilbert estimate. In contrast, our work investigates the role of a nonlinear MLP head in a Transformer with trainable linear attention, under a data mixing setting. We rigorously show that the ICL error of this model asymptotically matches that of a finite-degree polynomial predictor—thus providing a new theoretical understanding of how nonlinear MLPs enable feature learning and improve performance over linear baselines (as shown in Figure 1). Importantly, we also analyze how these gains depend on the structure of data mixing and training of the MLP.\nWhile both papers explore MLPs in the context of ICL, our contributions are complementary and novel in several key ways:\n\n- Our analysis explicitly models heterogeneous task mixtures with structured covariances, revealing how data mixing governs feature learning—an aspect not studied in Abedsoltan et al. (2024).\n\n- We derive an asymptotic equivalence between Transformer predictors and finite-degree polynomial models under high-dimensional scaling, providing quantitative performance predictions.\n\n- Our empirical results (Figures 1–3) validate this theoretical correspondence, demonstrating that nonlinear MLP heads consistently outperform linear ones in alignment with the predicted trends.\n\nTo further clarify the distinction, we present a detailed comparison between our work and the cited study in the following table:\n\n|                  | Abedsoltan et al. (2024)                                               | *Ours*                                                                                    |\n|------------------|------------------------------------------------------------------------|-------------------------------------------------------------------------------------------|\n| **Attention**    | **Fixed (non-trainable)** *non-linear* attention  feature map (Q=K=V=I; row‑wise $\\ell_1$‑normalized kernel smoother) with an optional MLP head after it      | **Trainable** *linear* attention with an nonlinear MLP after the attention                        |\n| **Data model**   | **Homogeneous** tasks (linear regression, teacher-student, decision trees) | **Heterogeneous** task mixtures with *multiple* data sources (nonlinear regression sources, with structured data, task, and noise) |\n| **Theory**       | **Linking attention to kernel smoother**: Consistency (optimal in the infinite‑context limit) via the Hilbert estimate   | **Rigorous high-dimensional analysis**: asymptotic equivalence to a polynomial model   |\n| **New insights** | Suitably chosen non‑linear attention can achieve optimal ICL in the infinite‑context limit          | Precisely how MLPs in Transformers drive ICL and when feature learning emerges from *data-mixing*         |\n\n\n> 3. Both the current paper and prior work[^1] restrict their analysis to one-layer Transformers. However, empirical studies show a significant performance leap from one to two layers, with diminishing returns improvement beyond two layer. Understanding how ICL emerges across multiple layers is crucial, and ignoring this makes the theoretical claims less applicable to real-world models.\n\nWe agree that multi-layer Transformers are highly relevant, and we appreciate the suggestion. Our choice to analyze a one-layer model (involving linear attention and a two-layer neural network) is consistent with much of the foundational ICL theory (e.g., [23, 41, 37]) and enables precise analytical results that are not yet tractable in deeper models. The one-layer case is not only analytically clean but also sheds light on where and how ICL behavior emerges—in this case, via MLP-driven feature learning under data mixing.\nWe will clearly state this as a limitation and a natural direction for future work. We believe our theoretical framework forms a rigorous starting point for eventual generalizations to deeper architectures.\n\n---\n\nIn summary, we thank the reviewer for highlighting important directions and related work. We will revise the paper to:\n* Cite and clearly differentiate from Abedsoltan et al. (2024),\n* - Clarifying the novelty of our polynomial surrogate theory under high-dimensional scaling,\n* - Highlighting our unique analysis of data mixing and the role of nonlinear MLPs in Transformers,\n* Explicitly discuss the one-layer limitation and suggest paths to extend our work.\n\nWe hope these clarifications will positively impact your evaluation. We believe our contributions offer new theoretical tools and insights into ICL with MLPs—complementary to prior work and of relevance to both theory and practice.\n\n\n*[^1] (Abedsoltan et al., 2024): Context-Scaling versus Task-Scaling in In-Context Learning, arXiv:2410.12783.*"}]

Author rebuttal:

Meta-review: This work develops an asymptotic analysis of in-context learning (ICL) in pre-trained transformer models with non-linear MLP heads. The authors present insights into the ICL behaviors of such transformers as well as how the mixture ratio of the heterogeneous pretraining data source will impact the behavior.

While concerns were raised regarding the assumptions on the linear attention, the setup of the gradient descent steps, and the difference from the existing works, the authors did an excellent job in rebuttal in addressing the reviewers' concerns. The new experiments on sentiment classification with transformers provide evidential justification for the alignment between the analyzed setup and the real-world case. By incorporating all the promised revisions in the rebuttal and the discussion period, despite the remaining concerns, the results can be valuable for the community.
Decision: Accept (poster)