In-context Learning for Mixture of Linear Regression: Existence, Generalization and Training Dynamics
Authors: Yanhao Jin
Krishna Balasubramanian
Lifeng Lai
Venue: Accepted by TMLR ()

Abstract: We investigate the in-context learning capabilities of transformers for the $d$-dimensional mixture of linear regression model, providing theoretical insights into their existence, generalization bounds, and training dynamics. Specifically, we prove that there exists a transformer capable of achieving a prediction error of order $\mathcal{O}(\sqrt{d/n})$ with high probability, where $n$ represents the training prompt size in the high signal-to-noise ratio (SNR) regime. Moreover, we derive in-context excess risk bounds of order $\mathcal{O}(L/\sqrt{B})$ for the case of two mixtures, where $B$ denotes the number of training prompts, and $L$ represents the number of attention layers. The dependence of $L$ on the SNR is explicitly characterized, differing between low and high SNR settings. We further analyze the training dynamics of transformers with single linear self-attention layers, demonstrating that, with appropriately initialized parameters, gradient flow optimization over the population mean square loss converges to a global optimum. Extensive simulations suggest that transformers perform well on this task, potentially outperforming other baselines, such as the Expectation-Maximization algorithm.

Published in Transactions on Machine Learning Research (11/2025)
In-context Learning for Mixture of Linear Regression: Existence, Generalization and Training Dynamics
Yanhao Jin
yahjin@rudavis.edi edu
Deportment of Statistics University of California, Davis
Krishnakumar Balasubramanian
kbala Abalo@uucdavis edu
Department of Statistics University of California, Davis
Lifeng Lai
(nai@ucdavis. edu
Department of Electrical and Computer Engineering, University of California, Damis
Reviewed on OpenReview: /ttps://openreert opwnreviee.........om net/ forum? id- buZXVuTsti
Abstract
We investigate the in-context learning capabilities of transformers for the d-dimensional mixture of linear regression model, providing theoretical insights into their existence, gen- eralization bounds, and training dynamies. Specifically, we prove that there exists a trans- former capable of achicving a prediction crror order O(yd/n) with high probability, where n represents the training prompt size in the high signal-to- Moreover, we derive in-context excess risk bounds of order O(L/VB) for the case of two mixtures, where B denotes the number of training prompts, and L represents the number of attention layers. The dependence of L on the SNR is explicitly characterized, differing between low and high SNR settings. We further analyze the training dynamics of transform- ers with single linear self-attention layers, demonstrating that, with appropriately initialized paramoters, gradient flow optimization over the population mean square loss converges to a global optimum. Extensive simulat ions suggest that. transformers perform well on this lask, potentially outperforming other baselines, such as the Expectation- Maximization algorithm.
ratio (SNR) regime.
Expeeation-NNaxiiiaatinn
1 Introduction
We investigate the in-context learning (ICL) ability of transformers for the Mixture of Regression (MoR) model (De Veaux 1989 Jordan & Jacobs 1994). The MoR model is widely applied in various domains, including federated learning, collaborative filtering, and healthcare (Deb & Holmes] 2000 Viele & Tong, 2002: Kleinberg & Sandler 20088 Faria & Soromenho, 2010 Ghosh et al., 2020) to address heterogeneity in data, often arising from multiple data sources. In particualr, we consider linear MoR. models where independent and identically distributed samples (Ir, v.) e Rd TR, for i 1,..., n, are assumed to follow the model y: (B,, .r) + U2, where e; ~ No0 (2) represents obscrvation noise, independent of 2;, and 8 e Rd is an unknown regression vector. Specifically, there are K distinct regression vectors { and each B is independently drawn from these vectors according to the distribution { 1-- The goal for new test sample, In+1, is to predict its label 1n+1- Specifically, we are interested in the ICL setup for MoR (Kong
et al. 2020 Pathak et al. 2024).
Classically, the Expectation Maximization (EM) algorithm is a widely used method for estimation and prediction in the MoR models (Balakrishnan et al. 2017 Kwon et al. 2019 Kwon & Caramanis 2020 Wang et al. 2024). A major limitation of the EM algorithm is its tendency to converge local maxima rather than the global maxinum of the likelihood function. This issue arises because the algorithm's performance

---

Published in Transactions on Machine Learning Research (11/2025)
crucially depends on the initialization (Jin et al. 2016
To mitigate this, favorable initialization strategies
based on spectral methods (Chaganty & Liang 2013 Zhang et al. 2016 Chen et al. 2020 are typically employed alongside the EM algorithm.
In an intriguing recent work, through a mixture of theory and experiments, Pathak et al. 2024 examined the performance of transformers for TCL, MoR models. However, their theoretical result suffers from the following major drawback. They only showed that the existence of a transformer architecture that capable of implementing the oracle Bayes optimal predictor for the linear MoR. problem. That is, they assume the availability of { I6-1, which are in practice unknown and are to be estimated. Hence, there remains a gap in the theoretical understanding of how transformers actually perform parameter estimation and prediction in MoR. Furthermore, their theoretical result is rather disconnected from their empirical auservations which focused on ICL. Indeed, they leave open a theoretical characterization of the problem of ICL MoR (Pathak et al. 2024 Section 4).
In this work, we first demonstrate that transformers are capable (in an existence sense) of in-context learning for linear MoRs by effectively implementing the EM algorithm, double-loop algorithm in which each inner loop consists of multiple steps of gradient ascent. We derive in-context excess risk bounds for the global solution of the empirical in-context risk minimizer, precisely quantifying the number of pre-training tasks required to achieve accurate predictions. While the aforementioned existence result and in-context excess risk bounds provide insight into capability of transformers for in-context learning MoR modes, from the practical point of view, it more important to understand the training dynamics of transformers under the MoR. models. Towards that, we also analyze the performance of gradient flow (in the population setting) for ICL MoR models with linear self-attention transformers. Furthermore, through our experiments, we empirically show that trained transformers achieve efficient prediction and estimation in the MoR model while substantially mitigating the initialization challenges typically associated with the EM algorithm. To summarize, we make the following contributions:
We demonstrate the existence of transformer capable of learning MoR. models by implementing the dual-loops of the EM algorithm. This construction involves the transformer performing multiple gradient ascent steps during each M-step of the EM algorithm. In Theorem 3.1 we derive precise bounds on the transformer's ability to make prediction in high signal-to- noise (SNR) regimes. In the special case of two mixtures, Theorem 3.2 also provides the procise high-probabin the estimation of the parameters by the constructed transformer in the high and low-SNR settings.
bound for
we analyze the sample complexity associated with pretraining transformers using a finite number of ICL training instances. Additionally, Theorem 4.2 provides guarantee that the gradient flow of the parameters of single linear self-attention layers will eventually converge to the global optimum under population mean squared loss with appropriate initializations.
In Theorem 4.1
As a byproduct our analysis, we also derive convergence results with statistical guarantees for the gradient EM algorithm applied to a two-component mixture of regression models, where the M-step involves T steps of gradient ascent. We extend this approach to the multi-component case, improving upon previous works, such as [Balakrishnan et al. (2017), which considered only a single step of gradient ascent.
1.1 Related works
Transformers and optimization algorithms: Garg et al. (2022) successfully demonstrated that trans- formers can be trained to perform ICL for linear function classes, achieving results comparable to those of the optimal least squares estimator. Beyond their empirical success, numerous studies have sought to uncover the mechanisms by which transformers facilitate ICL. Recent investigations suggest that transform- ers may internally execute first-order Gradient Descent (GD) to perform ICL, a concept explored in depth by Akyürek et al. (2023), : Bai et al. (2024). Von Oswald et al. (2023a), Von Oswald et al. (002bb), Ahn et al. (2024), Huang et al. (2024) and Zhang et al. (2024). Specifically, Akyürek et al. (2023) identified fundamental operations that transformers can execute, such as multiplication and affine transformations, showing that transformers can implement GD for linear regression using these capabilities. Building on
2

---

Published in Transactions on Machine Learning Research (11/2025)
this, Bai et al. (2024) provided detailed constructions illustrating how transformers can implement convex risk minimization across wide range of standard machine learning problems, including least squares, ridge, lasso, and generalized linear models (GLMs). Further, Ahhn [hhn et al. (2024) demonstrated that a single-layer linear transformer, when optimally parameterized, can effectively perform a single step of preconditioned GD. Zhang et al. (2024) expanded on this by showing that every one-step GD estimator, with a learnable initialization, can be realized by a linear transformer block (LTB) estimator.
Moving beyond first-order optimization methods, Fu et al. 2023 revealed that transformers can achieve convergence rates comparable to those of the iterative Newton ss method, which are exponentially faster than GD, particularly in the context of lincar regression. These insights collectively highlight the sophisticated computational abilities of transformers in ICL, aligning closely with classical optimization techniques. In addition to exploring how transformers implement these mechanisms, recent studies have also focused on their training dynamics the context. of linear regression tasks; see, for example, [hhang et al. (2023) and
Chen et al. (2044). In comparison to the aforementioned works, in the context of MoR, we demonstrate that transformers are capable of implementing ing double-loop algorithms such the EM algorithm.
The analysis of the standard EM algorithm for mixture of Gaussian and linear MoR models has a long-standing history Wu (1983); McLachlan & Krishnan (2007): : Tseng 2004). Recently, Balakrishnan et al. 2017) proved that FM algorithm converges at a geometric rate to a local region close to the maximum likelihood estimator with explicit statistical and computational rates of convergence. Subsequent works (Kwon et al. 2019; 2021) established improved convergence results for mixture of regression under different SNR conditions. Kwon & Caramanis (2020) extended these results to mixture of regression with many components. Gradient EM algorithm was first analyzed by Wang et al. (2015) and [Balakrishnan et al. (2017). It is an immediate variant of the standard EM algorithm where the M-step is achieved by one-step gradient ascent rather than exact maximization. They proved that the gradient EM also can achieve the local convergence with explicit finite sample statistical rate of convergence. Global convergence for the case of two-components mixture of Gaussian model was show by Xu et al. (2016), Daskalakis et al. (2017) and Wu & Zhou (2021). The case of unbalanced mixtures was handled by WWiinbrrerr & Bresler (2022). Penalized EM algorithm for handling high- digh-dimensional mixture models was analyzed by Zhu et al. (2017). Yi & Caramanis
EM Algorithm:
(0155) and WWan et al. (2024). showing that gradient EM achieve lincar convergence to the unknown parameter under mild conditions.
2 Preliminaries
Mixture of regression model: We now formally describe the MoR problem. The underlying true model is described by the equation:
V14 = xT Ui:
(1)
where v, ~ NN, Ta), vs ~ NN(0, T2))) denotes the noise term with variance J2 and 3,'s are i.i.d. random vectors that taking the value 8 with probability oे for k = I...,, K. The vectors St are unknown. For the MoR model equation I we define Rt = - as pairwisc distance between regression vectors, and Rmin mineay R Rmax maxıæ Rty as the smallest and largest distance respectively. The SNR of this problem is defined as the ratio of minimum pairwise distance versus standard deviation of noise
Rmin/ 10
(2)
When the number of the components K = and we represent 8f - 8 the SNR reduces to 7 2//00 In Section 3 we will show that the performance of the constructed transformer solving the MoR problem in general depends on the SNR condition of the problem.
Transformer architecture: We focus on transformers that handle the input sequence H E D by integrating attention layers and multi-layer perceptrons (MLPs). These transformers are structured to process the input by effectively mapping the complex interactions and dependencies between data points in the sequence, utilizing the capabilities of attention mechanisms to dynamically weigh the importance of different features in the context of regression analysis.
3

---

Published in Transactions on Machine Learning Research (11/2025)
Definition 2.1. An attention layer with M heads is denoted as Attno(-) with parameters 0 [(Vm, Qm, Kn Km) mE = DDD, On any input sequence H E DR we have
M
H Attno H) H
1
(V. o ((QmH) (KmH)) e D Σ
where a RR + R is the activation function and D is the hidden dimension. In the vector form,
N
h, h4 + Σ N Σ ((2mm, Kmhr)) Vmhgg.
The prevalent choices for the activation function inchude the softmax function and the ReLU function. Tn our analysis in Section 3] the attention layer (defined in Section 2.1 employs normalized ReLU activation, t t- (t)/N, which is used for technical convenience. This modification does not impact the fundamental nature of the study. There are various works on transformers that use non-softmax function as activation, (e.g. Shen et al. (2023 Bai et al. (2024) for ReLu activation, Luo et. al. (0021) for kernel attention, Qin et al. (2022) for cosFormer).
Remark 2.1.
Definition 2.2 (Attention only transformer). An L-layer transformer, denoted as TFo(-), a composition of L self-attention layers,
(-)(-) Ating Ating:- o... Ating (H)
where H E IR DxN is the input sequence, and the parameter 0 = (Ø)..... (0'....,0") 0%) consists of the attention layers ((0) = ((Vl)) (9) K(l) me(Mrre = DDD
Our theory consists of two parts, (i) the existence of the theoretical transformer that can internally implement the EM algorithun, and (ii) the dynamics of transformers with single linear self-allention layer trained by gradient flow on mixture of regression tasks. In the first part, the input sequence H E R Dx (x+1) has columns
h; [Tt, On a 1,4)"
(3)
hat1 [Ia+1. 10441.44:1 41 Op-d-3 11
where 1/f = yit, and 6, := 1fi < n + 1} is the indicator for the training examples. Then the transformer TFa produces the output H TF 8 (H). The prediction ân-1 is derived from the (d + 1, + 1)-th entry of H, denoted Da+1 (F) = (Im+1)a_1= Our objective is to develop fixed transformer architecture that efficiently conducts ICL for the mixture of regression problem, thereby providing a prediction 2a+1 for yatt under appropriate loss framework. Besides, the constructed transformer in Section 2 can also extract an estimate of the regression components, which is realized by operator (TF(() = extracts the estimate of 8* in the output matrix. In the second part, the embedded input matrix is given by
E
31
D2
In
Futl
R R(d+1)x(x+1)
(4)
31
12
Dn
O
and is fed into a single-layer lincar self-attention layer FLSA R(d-1)x(++1) - ((d+1)x(++1)
fLSA(E: E VE- ETQUKE
(5)
where 0 V). The prediction on the query sample xn+ is given by the bottom-right entry of the matrix by FLSA. i.e. 0n+1 = (E:))(E
Remark 2.2. We highlight the key ideas of our work. In Section 3 Theorems 3.1 and 3.2 provide an explicit construction of a transformer that executes gradient-EM within its forward pass an E-step followed by T gradient updates in the M-step along with the SNR and prompt-length regimes where the error bounds hold. This an existence result. In Section 4 we pre-train transformers by minimizing MSE. Thus, the learning mechanism not tiod to our constructive parameterization and can match or surpass it. Our experiments report end-to-end excess MSE, showing the effectiveness of trained transformers, though this does not by itself certify that they EM internally- even though they capable of doing so.
4

---

Published in Transactions on Machine Learning Research (11/2025)
Remark 2.3. Equation equation 3 is used only for theoretical convenience, not as data- prata-pppprocessing recipe. It specifies a per-token "register" layout that lets the proof show how one attention layer can read/write quantities needed for an EM step, with zero blocks serving as scratch space and t, as position tags. In practice, we feed standard embeddings for (z,, Ye) and a masked query In+1, while a single input projection can create the scratch dimensions automatically. Writing the prompt as E in equation | just a compact reformatting of inputs (matching common implementations) for technical convenience and does not introduce new assumptions.
Notation: Given two functions g(n) and f(z), we say that f(n) Q(g(n)), if there exist constants c > 0 and no >= 0 such that f(n) i c (n) for all ie no- We say that f(n) is (gg(n)) if there exist positive constant C' and no such that 0 & f(n) < Cg(n) for all v no. For vector E Rd its l2 norm is denoted by ||e||2- For a matrix A E Rdxd ||Al|op denotes the operator (spectral) norm of A. We denote the joint distribution of (I, y) in model equation by Pay and the distribution of x by Pa. Besides, we denote the joint distribution of (#1, Vi...., ,.., In, Un, In+1, Mn+1) by P, where {rz, Vili- are the input in the training prompt and In+1 is the query sample. Besides, in Section we use 3/4 E defined as V y:l, and 4 = [(inn11) for I,...,n, n, n + 1 simplify our notation.
Evaluation: Let f H ~ R be any procedure that takes prompt H as input and outputs estimate f on the query 1n+1- We define the mean squared error (MSE) by MSF(J) =- [P[([(H) - yn+1)
3 Existence of transformer for MoR
In this section, we show the existence a transformer that can approximately implement the EM algorithm internally in Thcorem 3.1 and Theorem 3.2. Note that under model Theorem [] the oracle vector that minimizes the mean squared error of the prediction Ep[(xT+1 1n+1) is given by
K
gOR = arg min Ep.,y BER
[( - vn+1) = Σ n B2.
Generally, the transformer constructed in Theorem 3.1 will provide prediction that is close to x goR
Theorem 3.1. Given the input matrix H in the form of equation there exists transformer TF with the number of heads (()) & M = 4 in each attention layers. This transformer TF can make prediction on Un+1 by implementing gradient EM algorithm of MoR problem where T sleps of gradient descent is used in each M-step. When L is sufficiently large and the prompt length n satisfies following condition
d.2
K2 1/3 115 ,
)
d Tmin
h2 8
n e C max dog
log
under the SNR condition
n 2 CKpr log( log(Kop) px); for sufficiently large > O,
(6)
equipped with OTTlog (n/d) attention layers, the transformer has the prediction error Ay read, (TF(H)) - xr T4 30R) upper-bounded by
o (
(1
dk 2 log2
nk² (
dK log( K2 8
v log(dd/0)
nTmin
with probability at least 1 - 96, where Pr = max, THz min, n. is the ratio of marimum mizing weight and minimum mixing weight, Tmin miny nz and ready (H) = ((n+t)am extracts the prediction on query sample.
Theorem demonstrates the feasibility and theoretical guarantees of transformers in solving a general mixture of regression problems under the high SNR condition specified in equation G The error bet
5

---

Published in Transactions on Machine Learning Research (11/2025)
the transformer's prediction and the true response for the query sample is bounded with high probability in the order of / /log(d/0) & (n/8)) This error decreases as the prompt length increases and is affected by factors like the dimension d and the number of components K. The ratio Pr = max, x./minn miny my quantifies the imbalance in mixing proportions. Larger imbalances pr degrade the error bound, indicating that the transformer's performance could worsen when mixture components are highly unbalanced.
In the spocial case of MoR problems with two components B - 3* and "f x = 2 have 30R = O. While predicting zero is not quite meaningful, estimating the true regression coellicient vector B is of interest. Hence, in Theorem 3.2 below, we provide more refined results focusing both on the low and high SNR regimes.
Theorem 3.2. Given imput matrir H whose columns are given by equation S there erists transformer TFs, with the number of heads M < M = 4 in each ottention layers, that can make prediction on Yn+1 by implementing gradient EM algorithm of MoR problem where T steps of gradient descent are used in each M-step. When T is sufficiently large and the prompt length n satisfies
n z Cd (1/6),
(7)
the transformer can approzimates 3 by the second-to- last layer with probability at least 1 - 6:
When n C(dlog n/n)* + reads (TF(H)) - 30 o (dloo"(n/00//)))
When n > C(dlog" n/n)* à then reads (TF(H)) - B* & 0 (Vlloggn//6)//
where reads((F(H))) [TF())))((((())(((e()))))))))))) d extracts the estimate of 8*
In Theorem 3.2 the error depends on the SNR, 7, and exhibits two distinct behaviors: In low SNR set- reflecting the inherent difficulty of recovering 3 in noisy environments. In high SNR settings, the error scales as O(v dllog'dn//// (n/8)/n), showing better performance due to stronger signals dominating the noise. According to Theorem 3.2 the architecture of the constructed transformer varies primarily in the number of layers it includes. In general, with the prompt length and dimension d held constant, the constructed transformer needs more training samples in the prompt in the low SNR settings to achieve the desired precision. The prediction error is order of settings, and is ([(d/n)+) 4) in the low SNR settings. Besides, under the high SNR settings, the constructed transformer needs O( [lg(n//)) attention layers, while needs ((v n/ log(lo((/dd)) attention layers in the low SNR settings.
tings, the error scales as
under the high SNR
The proof of Theorem is provided in Appendix A.7 and details of the proof of Theorem 3.2 can be found in Appendix A.5 In Theorem 3.1 and 3.2 the variance parameter 0 in equation 1 and equation 2 in the noise is assumed to be a fixed known constant, since is used in the construction of the transformer. The SNR condition required Theorem 3. is stricter than that in Theorem 3.2 due technical reasons in the proof. However, in our simulations (presented in Figure I in Section 5 we see that the actual performance of the transformer is still good in the low SNR scenario when the number of components K > 3.
Finally, in Theorem 3.3 we provide the excess risk bound for the transformer constructed in Theorem 3.2
Thcorem 3.3. For any T being sufficiently large and the prompt length satisfics condition equation Define the excess risk R = Ep (un+1 (TF())) inf Ep Ep [x - bn+1)3]. Then the ICL prediction (TF(H) of the constructed transformer in Theorem 3.2
satisfies
olydlog" n//n)
0 < r < (d (n/8)/n)* (n/)))))/ 1/4
R R
(8)
Oddogg n
n > C(dlog (n/8)/)) 1/1
Furthermore, inf Ep [ 22 82 + 1
6

---

Published in Transactions on Machine Learning Research (11/2025)
Theorem 3.1 Theorem 3.2 and Theorem 3.3 provide the first quantitative framework for end-to-end ICL in the mixture of regression problems, achieving desired precision. The excess risk of the constructed trans- former is Odd n//n) under the high SNR settings, and is D(yd/nlon logm) under the low SNR settings. These results represent an advancement over the findings in Pathak et al. (2024), which do not offer explicit error bounds like equation
The transition on the error rate appears in Theorem 3.2 due to the switch of the analysis regime. In high SNR 77 (n/))nn) "4 the population EM map is strongly contractive (curvature x 17² Balancing this contraction with sampling noise ( ( V d (n/6)/n) gives the parametric rate Oly d (n/odd/gg/////////// (n/8)/n). low SNR 7 v (1/8))) 1/4 the map nearly flat; progress is driven by a cubic drift, yielding an is exactly where curvature (xrr²) overtakes sampling error. Plugging this ni into the high-SNR bound recovers the n114 order, so the two regimes meet smoothly.
n scale neighborhood: 0 (1/))))11/0))n) 3/2) The threshold at n ==
4 Understanding Transformer Training on MoR tasks
4.1 Analysis of pre-training
We now analyze the sample complexity needed to pretrain the transformer with limited number of ICL training instances. Prior results from Bai et al. (2024) are only applicable to linear models and are not immediately applicable to the linear MoR models that we focus on in this work. We consider the square loss between the in-context prediction and the ground truth label:
(; 2) = 2 yn+1 clip, ready (TFe
I.....,, L.m =,,,,... 1....,My is the collection of parameters of the transformer and cl)p = (y)) is the standard clipping operator with (a suitably large) radius R > that varies in different problem setups to prevent the transformer from blowing up on tail events, in all our results concerning (statistical) in-context prediction powers. Additionally, the clipping operator can be employed control the Lipschitz constant of the transformer TFa with respect to 0 (see [ai et al.] )024)). In practical applications, it is common to select. sufficiently large clipping any input sequence of interest. Applying clipping operator on the objective functions when training LLM is used in RLHF (sce Ziegler et al. 2019 Ouyang el. al. (2022) ). Denote ||0]| as the norm of transformer given by
where Z : (H, bn+1) is the training prompt, 0 = {(K!h Q!h (1))
radius R. to ensure that it does not alter the behavior of the transformer
M Σ m==1
|0] == max max
(11912)
op K opp)
+
[(00
xE[L]] mc[M]
op
Our pretraining loss is the average ICL loss on B pretraining instances (1: corresponding test ICL loss on a new test instance:
lid
~ n, and we consider the
B Laa(0)== Σ (e; z
( = :10))=Ep(4.)) 0 z)
Our pretraining algorithm is to solve a standard constrained empirical risk minimization problem over transformers with I. layers, M heads, and norm bounded by M"
i = arg min Laa(()) BeG
(9)
MM 0
m Q v(O) max M & M, & M). de[2]

---

Published in Transactions on Machine Learning Research (11/2025)
Theorem 4.1 (Gencralization for pretraining). For MoR problem given by equation 4 suppose that maxigk || & C with some absolule constant C, then with probability at least - 3 (over the pretraining
instonces (Z() Je(B)) the solution 0 to equation satisfies
( & inf (0) + o0 (1 4/ x log g
2nB g
(L)((MA)) + log(1/E)
B
where = log (2 + max { (M', Bgg By, (2B,) Bg Vlg((nd By - log
2n
C log (2 D is
the hidden dimension and M the number of heads.
provided Appendix B.1 Under the low SNR scttings, the constructed transformer generally requires more attention layers than those in the high SNR settings achieve the same level of excess risk. Besides, the theorem highlights fundamental trade-off: while larger models (more layers, heads, and hidden dimensions) increase have the capacity to learn complex patterns, they also require more pretraining datasets (B).
Remark 4.1. The proof of Theorem
4.2 Dynamics of single linear self-attention layer
Next, we investigate the training dynamics of gradient flow for MoR models. For this subsection, consider transformers with linear self-attent attention layers. Given the input matrix in the form of equation appropriately sized key, query and value matrices K, Q. V, the output of a linear attention block is given by fn+1 = (E) a+1 For our technical analysis, following Zhang et al. (002)). we only consider training the model equation 5 over population squared loss between Dn+1 and Un+1, L.c.
1 [P[(0n+1 1o+1)
arg min (F,A(F,O) BEB
And we assume that E Σ 1
3 on the MoR task equation
The gradient descent gives
0 02 LVL 0,)
when minimizing LSA(E, 0) and this could simplified
d0
dt LSAA
Letting 0 and setting ((t) 02 at time t = ke, recognize the left-hand side above the discrete derivative of ((t) at time Hence, we get contimuous-iime ordinary differential equation (i.e. the gradient flow of the paramcters)
d0
dt VSA
It captures the behavior of gradient descent when we running the gradient descent with small step size. In the remaining of this section, we start by rewriting the output of the linear attention module an alternative form. Following Zhang et al. (2023), we define
)
Un u12
V
KTQ=
U1 u-1
therefore, the prodiction on the query sample is given by
1
0n+1 = u221 -x Un u - x y upz
(10)
n
1 _ x Un u- LyTy u
+ u-1
16+1
where X
xx," and y = [y.....,yu] MT We will consider gradient flow with an initialization that satisfies the following assumption
8

---

Published in Transactions on Machine Learning Research (11/2025)
Assumption 4.2.1 (Initialization). Let y be parameter, and let e Raxd be any matrix satisfying θ R,3 0 and tr (00* (R,38T) !00 (R33T))
w-1(0) 7
"1)(0) = u21 = 0,
(11)
(n) = (EB3T) T0T (RBBT)*
Theorem 1.2 below proves that gradient flow will converge to a global optimum under suitable initialization. when the parameter satisfies the condition
Theorem 4.2. Under initialization condition equation 17
2 min (E33T)
9 Vdl n+d+1 max (E833T) ) + Ø
(12)
and EB = Σ - O, we have U21 (t) = (t) = 0 for all t 3 O, and the gradient flow converges to global minimum of the population loss. Moreover, U11 and u-1 converge to Un and w²1 respectively, where
- = (" E33 +
1 + 0²
E
(13)
F
U = (u=,) -1
n2 E33
+
E| 3
02
EB3
(14)
With = and Uh specified in equation and equation 14 the linear self atffattenntonn layer makes prediction En+1
n2 + E33T
E B
v2
In+1 = x 2
+
E89
n
When n is sufficiently large, it holds that Un 2 Ia and 1 -i= Biod 2 Eryyy. Therefore, when is large, the prediction made by linear self-attention layer In+1 & xn Ja (Erx T +3+ Eur) - 1 This shows that the linear self-attention layer effectively learns the optimal predictor in the large-sample limit.
Theorem 1.2 provides crucial insights into the convergence of gradient flow under structured initializations and zero mean assumption for the coefficients. Larger noise variance Ø or smaller sample size necessitates smaller y (scaling of (0)()). The initialization (1) and the trace condition on θ encode prior knowledge of the input distribution E [. acting
preconditioner for efficient learning. In particular, the trace condition in equation [ ensures θ is scaled to interact stably with the data covariance, preventing explod- ing/vann /eg/aiiisii updates. Finally, we remark that if E3 + 0, there are additional terms affecting the dynamics, possibly complicating the convergence. We leave this problem as a possible future direction.
Compared to Theorem 4.1 in Zhang et al. (2023), our results are applicable for the case when label noise is present. Furthermore, generalize the distribution assumption proposed the the coefficient 3. Indeed, in Zhang al. (2023), the sample on the prompt are generated based the noiseless model yı 2 3 with B ~ NO0, T2). Whercas, our analysis only relies on the moment information E, and E the distribution of B. Besides, as mentioned above, our assumption equation I1 precisely characterizes how the initialization condition on the parameters depends on the covariance structure E
5 Simulation study
In this section, we present numeric results training transformers on the prompts described in Section 2 We train our transformers using Adam, with a constant step size of 0.001. For the general sottings in the experiments, the dimension of samples = 32. The number of training prompts are 64 by default (B is other value if heheewise stated). The hidden dimension D 64 by default (D is other value if otherwise

---

Published in Transactions on Machine Learning Research (11/2025)
stated). The training data 2,'s are iid. sampled from standard multivariate Gaussian distribution and the noise 04's are i.idd sampled from normal distribution (O, &2). The regression coefficients are generated from standard multivariate normal and then normalized by its l2 norm. Once the coefficient is generated, it is fixed. The excess MSE is reported. Each experiment is repeated by 20 times and the results averaged over these 20 times. We also point out that assumptions regarding n and in Section is for theoretical constructions. It is not related to the trained transformer practice. The performance of the trained transformer will be better than the constructed transformers in Theorem and 3.2
3.0
test MSE by TF SNR-1
test MSE by TT SNR=1
3.-
by SNR=5
2.5
test MSE by TF SNIt110
test MSE by TF SNR-10
by
test MSE by EM SNItm5
test MSE by EM SNP-55
2.0
SNPm10
2.5
Excess
1.5
Excees
2.0
1.0
1.0
0.5
0.0
0.0
0
10
20
40
50
60
0
10
20
30
50
60
Prompt length
Prompt length
SMR-1
test MSE TF SNR=1
test MSE by TF 5MR=5
test MSE by TF SNR-55
SNR=10
MSE TF 5NR=10
test MSE by EM 5MR=1
test MSE by EM SNR-1
2.5
test MSE by EM 10
SNR-5
test MSE EM 1NR=10
Excess MSE
2.0
Excees
2.5
1.0
0.5
0.0
40
50
Prompt length
Prompt length
Figure 1: Plot of excess testing risk by the transformer (and EM algorithm) v.s. prompt length with different SNRs on MoR lasks with K 2, 5, 20 components.
The initializations of the transformer parameters for all our experiments are random standard Gaussian. As we will see from our results, transformers provide efficient prediction and estimation errors despite this global initialization. A possible explanation for this fact might be the overparamelrization naturally available in the transformer architecture and the related need for overparametrization for estimation in mixture models (Dwivedi et al. 2020 Xu et al. 2024); we leave a theoretical investigation of this fact as intriguing future work.
Performance with different prompt length: In this experiment, we vary the number of components
K = 2, 2,25200. 20. For each case, we run the experiment with different SNR (7 = 1,5, 10). The r-axis is the prompt length, and the y-axis is the test MSE. The number of attention layers is given by L = 4. The performance results of the transformer are presented in Figure 1
From Figure 1 we observe the following trends: (1) With the number of prompt lengths and other parame- ters held constant, the trained transformer exhibits higher excess MSE the low SNR settings. (2) When the prompt length is very small, indicating an insufficient number of samples in the prompt, the resulting
10

---

Published in Transactions on Machine Learning Research (11/2025)
excess test MSE is high. However, with sufficiently large prompt length, the performance of the trans- formers stabilizes and is effective across all SNR settings, leading to relatively small excess test MSE. (3) Additionally, when the prompt length and SNR are fixed, an increase in the number of components tends to result in a larger excess test MSE.
Performance with different number of training prompts: In this experiment, we vary the number of training prompts B from 64 to 512. For each case, we run the experiment with two components (K 2), different SNR (1 1, 10). The r-axis is the number of training prompts, and the y-axis is the test MSE. The length of training prompts is n 64.
test MSE SNP-1
test MSE SNR=10
0.03
0.02
0.01
200
number of training prompts 8
Figure Plot of excess testing risk of the transformer v.s. the number of prompts with different SNRs.
test MSE =NR=5 test MSE SNE=10
0.06
Excess MSE
0.02
0.00
60
80
100
120
Dimension
Figure 3: Plot of excess testing risk of the transformer v.s. the dimension d with different SNRs.
Figure 2 gives the performance of trained transformer with different number of training prompts under three different SNR settings. Based Figure [ we observe that when the number of training prompts is already sufficiently large, the excess MSE is relatively small. Furthermore, as the number of training prompts incrcases, there is general trend of decreasing in the excess MSE.
Performance with different dimension d of samples: In this experiment, we fix the hidden dimension D 256, the number of components K = 2, the number of prompts = 128 and the prompt length is given by n 64. The r-axis is the dimension d of the input sample v; and y-axis the excess test MSE. In this experiment, we evaluate the performance of the trained transformer for various dimensions d 32, 48, 64, 80, 96, 112, 128. The performance of the transformer presented in Figure Observations
11

---

Published in Transactions on Machine Learning Research (11/2025)
from this figure indicate that increasing the dimension significantly raises the excess test MSE. Notably, this increase becomes more pronounced at the lower SNR levels.
An additional experiment on the performance with different number of hidden dimension is provided in Section D.2
6 Conclusions
We have explored the behavior of transformers handling linear MoR problems, demonstrating their in- context learning capabilities through both theoretical analysis and empirical experiments. Specifically, we showed that transformers are capable of implementing the EM algorithm for linear MoR tasks. Additionally, have examined the sample complexity involved in pretraining transformers with finite number of ICL training instances and the training dynamics of gradient flow, offering valuable insights into their practical performance. Our empirical findings also reveal that transformer performance is less susceptible to initial- izations. For future work, understanding the training dynamics of general transformers for MoR problems remains a highly interesting and challenging task. Furthermore, extending our results to non-linear MoR models would be a natural direction.
Prior work (e.g., Raventós al. (2023)) studied ICL for linear regression with coefficients sampled from prior, measuring task difficulty the minimum distance across all regression vectors in the pretraining corpus. In contrast, our paper defines SNR per task, based the separation of the K mixture components within single prompt relative to noise, so having many distinct tasks in pretraining does not automaticall reduce SNR. For technical convenience, we assume fixed ground-truth coefficients S but one could instead modol them with priors (c.g., Gaussian or uniform on the sphere), which would reduce to Raventós et al. (2023) in the K case. Extending our results to this randon-eflects setting requires re-deriving several lemmas and adapting the proofs of Theorems 3.1 which we leave as promising future work.
12

---

Published in Transactions on Machine Learning Research (11/2025)
Acknowledgments
The work of Krishnakumar Balasubramanian was supported in part by National Science Foundation (NSF) under grant DMSS-4134266. The work of Lifeng Tai was supported part by NSF under grants CCF-22-32907 and ECCS-24448268.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement precon- ditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https:://ppeneevie http:////ppaaeeeeee...esss/rrrrp not:/forumm/dd-oggxxeee111
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in- context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
Sivaraman Balakrishnan, Martin J Wainwright, and Bin Yu. Statistical guarantoes for the EM algorithm: From population to sample-based analysis. The Annals of Slalistics, 45(1):777 120, 2017.
Arun Tejasvi Chaganty and Percy Liang. Spectral experts for estimating mixtures of linear regressions. In International Conference on Machine Learning, pp. 1040-1048. PMLR, 2013.
Sitan Chen, Jerry Li, and Zhao Song. Learning mixtures of linear regressions in subexponential time via fourier moments. Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pp. 587- 600, 2020.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of muuli--eaad softmax at- tention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. 19442. 2024.
Constantinos Daskalakis, Christos Tzamos, and Manolis Zampetakis. Ten steps of EM suffice for mixtures of two Gaussians. In Conference on Learning Theory, pp. 704-710. PMLR, 2017.
Richard D De Veaux. Mixtures of linear regressions. Computational Statisties & Data Analysis, 8(3):227-244) 245, 1989.
Partha Deb and Ann M Holmes. Estimates of use and costs of behavioural health care: A comparison of standard and finite mixture models. Health economics, 9(6):475 489, 2000.
Raaz Dwivedi, Nhat Ho, Koulik Khamaru, Martin J Wainwright, Michael I Jordan, and Bin Yu. Singularity, misspecification and the convergence rate of EM. The Annals of Statistics, 48(6):3161-311 2020.
Susana Faria and Gilda Soromenho. Fitting mixtures of linear regressions. Journal of Statistical Computation and Simulation, 80(2):201-223 2010.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. orXie preprinl arXim:2310.17088 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in- context? case study simple function classos. Advances in Newral Information Processing Systems, 35: 30583 30598, 2022.
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 53:19586-19550 2020.
13

---

Published in Transactions on Machine Learning Research (11/2025)
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. Proc. International Conference Machine Learning, 2024.
Chi Jin, Yuchen Zhang, Sivaraman Balakrishnan, Martin J Wainwright, and Michael Jordan. Local maxima in the likelihood of Gaussian mixture models: Structural results and algorithmic consequences. Advances neural in information processing systems, 29, 2016.
Michael Jordan and Robert Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural computation, 6(2):181-21) 1994.
Jon Kleinberg and Mark Sandler. Using mixture models for collaborative filtering, Journal of Computer and System Sciences, :4(1))49-69, 2008.
Jason M. Klusowski, Dana Yang, and WD Brinda. Estimating the coefficients of a mixture of two linear regressions by expectation maximization. IEEE Transactions on Information Theory, 65(6):3515-3555 3524, 2019.
Weihao Kong, Raghav Somani, Zhao Song, Sham Kakade, and Sewoong Oh. Meta-learning for mixed linear regression. In International Conference on Machine Learning, pp. 5394 5404. PMLR, 2020.
Jeongyeol Kwon and Constantine Caramanis. EM converges for mixture of many linear regressions. In International Conference on Artificial Intelligence and Statistios, pp. 1727-1736. PMIRR, 2020.
Jeongyeol Kwon, Wei Qian, Constantine Caramanis, Yudong Chen, and Damek Davis. Global convergence of the EM algorithm for mixtures of two component lincar regression. In Conference on Learning Theory, pp. 2055 2110. PMLR, 2019.
Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the minimax optimality of the EM algorithm for learning two-component mixed linear regression. In International Conference Artificial Intelligence and Statistics, pp. 1405 1413. PMLR, 2021.
Shengjic Luo, Shanda Li, Tianle Cai, Di Hc, Dinglan Peng, Shuxin Zheng, Guolin Kc, Liwei Wang, and Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems, 34:22795- 2222077722207 2021.
Geoffrey J McLachlan and Thriyambakam Krishnan. The EM algorithm and extensions. John Wiley & Sons, 2007.
Long Ouyang, Jefsrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models follow instructions with human feodback. Advances in neural information processing systems, 35:27730-27744, 2022.
Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. In The Twelfth International Conference on Learning Representations, 2024. URL
https://openreview net//forum?iderlL4jj1111
Zhen Qin, Weixuan Sun, Hui Deng. Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint ar Xiw:2200200879 2022.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in neural informalion processing syslems, 36: 14228 14246, 2023.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. arXiv preprint vrXio:2302.06461 2023.
Paul Tseng. An analysis of the EM algorithm and entropy-like proximal point methods. Mathematics of Operations Research, 29(1):27 -4, 2004.
14

---

Published in Transactions on Machine Learning Research (11/2025)
Roman Vershynin. High-dimmensiona probability: An introduction with applications in data science, vol- ume 47. Cambridge university press, 2018.
Kert Viele and Barbara Tong. Modeling with mixtures of linear regressions. Statistics and Computing, 12: 315 330, 2002.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, An- drey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. Tn Interna- tional Conference on Machine Learning, pp. 35151 35174. PMLR, 2023a.
Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zuechet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa- optimization algorithms in transformers. arXiv preprint arXiv:2309.0520 055858, 2023b.
Martin J Wainwright. High-dimensional statistics: A non-asymptotic vieupoint, volume 48. Cambridge University Press, 2019.
Ning Wang, Xin Zhang, and Qing Mai. Statistical analysis for penalized EM algorithm in high-dimensional mixture linear regression model. Journal of Machine Learning Research, 22222)) 85, 2024.
Zhaoran Wang, Quanquan Gu, Yang Ning, and Han Liu. High dimensional EM algorithm: Statistical optimization and asymptotic normality. Advances in neural information processing systems, 28, 2015.
Nir Weinberger and Guy Bresler. The EM algorithm is adaptively-pptimal for unbalanced symmetric Gaus- sian mixtures. Journal of Machine Learning Research, 23(103): 79, 2022.
CF Jeff Wu. On the convergence properties of the EM algorithm. The Annals of statistics, pp. 95- 95-103, 1983.
Yihong Wu and Harrison H Zhou. Randomly initialized EM algorithm for two-component Gaussian mixture achieves near optimality in O(vn) iterations. Mathemalical Statislics and Learning, 4(3), 2021.
Ji Xu, Daniel J Hsu, and Arian Maleki. Global analysis of expectation maximization for mixtures of two Gaussians. Advances in Neural Information Processing Systems, 29, 2016.
Weihang Xu, Maryam Fazel, and Simon S Du. Toward Global Convergence of Gradient EM for Over- Parameterized Gaussian Mixture Models. arXiv preprinl ar Xrx::22407 00490, 2024.
Xinyang Yi and Constantine Caramanis. Regularized FM algorithms: A unified framework and statistical guarantees. Advances in Neural Information Processing Systems, 28, 2015.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. ar preprint xrX1t:2306.09927 2023.
Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In- Learning of Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization. arXiv preprint ar Xrx:224402 14951, 2024.
Yuchen Zhang, Xi Chen, Dengyong Zhou, and Michael I Jordan. Spectral methods meet. EM: A provably optimal algorithm for crowdsoureing. Journal of Machine Learning Research, 17(102):1 44, 2016.
Rongda Zhu, Lingxiao Wang, Chengxiang Zhai, and Quanquan Gu. High-dimensional variance-reduced stochastic gradient expectation-mantt maximization algorithm. In International Conference on Machine Learn- ing, pp. 4180-4188. PMLRR 2017.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Chris- tiano, and Geoffrey Irving. Fine-tuning language models from human preferences. ar aiv preprinl arXiv:1909.0885 2019.
15

---

Published in Transactions on Machine Learning Research (11/2025)
A Proof of Theorems in Section 3
In this section, we illustrate that the transformers constructed in Theorem 3.1 and Theorem Theorem 3.2 solve the MoR problem by implementing the EM algorithm internally while GD is used in each M-step. Previous works (e.g. Balakrishnan et al. (2017) Kwon et al. (2019) and Kwon et al. (021)) focused on the sample-based EM algorithm, typically employing closed-form solutions or one-step gradient approaches in the M-step. We start with detailed explanation on the existence of the transformer in MoR problem with two symmetric components (K = 2, .. = के = and B = - = -B*). For general analysis, we explore the performance of the transformer using gradient descent in steps T within the EM algorithm. To simplify the analysis, we restrict our stepsize a (0, 1) in each gradient descent step in M-step.
Attention layer can implement the one-step gradient descent. We first recall how the attention layer can implement one-step GD for a certain class of loss functions demonstrated by Bai et al. (2024). Let e R²2 - R be loss function. Let (B) 1 function l on dataset (i)ee))) and denote
= E ((3 x., v.) denote the empirical risk with loss
BK+1 == -a
(15)
as the GD trajectory on T. with initialization 30 E TRd and learning rate 0. The foundational concept of the construction presented in Theorem 3.2 derived from Bai ct al. 2024).. It hinges on the condition that bivariate function), can be
the partial derivative of the loss function, d,E (s, 0,((s, () (considered approximated by sum of ReLU functions, which are defined as follows:
Definition A.1 (Approximability by sum of ReLUs). A function g : RR - R is (Eapprox M, C)-
approximable by sum of ReLUs, if there exists a "(M, C')-sum of ReLUs" function
M
[M,
Cma (a 12] Σ m=1
with Σ m==1
cmll & C, max [am]i & 1, am E k+++1 Cm E R me[M]
such that supzel- [g(z) far,o(z)l &
Eapprox.
Suppose that the partial derivative of the loss function, 0, t), is approximable by a sum of ReLUs. Then, T steps of GD, as described in equation [5 can be approximately implemented by employing T attention layers within the transformer. This result formally presented in Proposition C.T
Transformer can implement the gradient-EM algorithm: Proposition illustrates how the trans- former described in Theorem 3.2 is capable of learning from the MoR problem. Using Proposition C.1 we can construct transformer whose architccture consists of attention layers that implement GD for each M-step, followed by additional attention layers responsible for computing the necessary quantities in the E- step. Here a summary of how the transformer implements the FM algorithm for the mixture of regression problem. Following the notation from Balakrishnan et al. (2117). we consider the mixture of regression with two symmetric components and define the weight function:
(y x T99
w w
exp
20
exp
20 (v TT TT) + exp
(b+
20
Denote 300) as the current parameter estimates of 8 in the EM algorithm for the MoR problem. During each M-step, the objective is to maximize the following loss function:
-1 Σ gco (I., ui) (y: - 27 2T + (1 - -w0an (I, v.)) (y: + xx xT))
Qn (800) = 2n
(16)
The update 3(++1) is given by g(++1) arg max Qn 8()). Lemma A.1 below, demonstrates that the function (1)()) minimized in each M-step is approximable by sum of ReLUs.
16

---

Published in Transactions on Machine Learning Research (11/2025)
Lemma A.1. For the function (B) - (xx B, where
(x) B, U,))) wgg (0i, V1 3) (1 wigu Vi Vi x2 3)]
it holds that (1) (s,t) is convex in first argument; and (2) (s, is (0, to, 4, a0papppooie by sum of ReLUs.
Proof of Lemma
A.1 Note that
(s, t) = 3g0 (I,, (t) - s)² (1 (grn s)?.
Taking derivative w.r.t. the first argument yields
0 w (ot, ())( 2) 5) (1- (s,)) = pos (I,, v.) + 2 3gn (2,, v1)) 2.
s),
Hence, t(),t) is convex in the first argument and
(s,00 zug (z,, Vs)(s t) + 2(1 ugt (i))
= Zuggo (I,, (2)((2((s - t)/2) 20(- (s - t)/2)]
+ 21 - gen v1 (20 1)/2) - 20 (-xx + 1)/2)].
Here 01 gu (Ii, Wi), 02 (2; U)), C3 4(1 - Wgte) (I,, vi)) and C4 -4(1 warn (24, v:)). Now, we have |c1 leal + lcsl + Ical < 16 and C,l(s,t) t) is (0, +0, 4, 16) approooioie by sum of ReLUs.
By Lemma A.1 we can design attention layers with T layers that implement the T steps of GD for the empirical loss (1) (B*) as outlined in Proposition C.1 We provide a concise demonstration of the entire process for MoR with two symmetric components below. Starting with an appropriate initialization g(0)
the first M-step minimizes the loss function:
(
1
Σ
22 (F v:)( B)2 + (1 Wator v:) x 8)²).
Following Proposition C.1 given the input sequence formatted as h4 [z, 3/: Oz: 0D-20-3:230000000 1;t;), there exists
transformer with T attention layers that gives the output h, = [x;; vl: 36 0pp -00-2a-3;1;;) 1:4]. Furthermore, the existence of a transformer capable of computing the necessary quantities in the M-step is guaranteed by
Proposition 1 from Pathak et al. 2024h and we restate this proposition in section C in appendix.
It worth mentioning that computing Wgt (I,, Vs) in cach M-step can be casily implemented by affine and be made for the upcoming iterations of the
softmax operation in Proposition C.2 Similar arguments
EM algorithm and we summarize these results in Lemma A.2 and Lemma A.3
Lemma A.2. In each E-step, given the input (T++) = | (T+
(T+1) 2+1
where
11 (T+1) [ani/] 8 Op 2d 3: 1;4; w
o4 u- (a) i= I,....
h (T+1) a+1
| 304 3/: O, 0p-24.
Furthermore, the output sequence takes the there exists transformer TF that can compute w uage (z. form of
(1) (((++1) [ 1/4 .T 000-2-1:: 1; (n vi)l i I,..., n.
(17)
((++1) [ a/0 ((4) 0::20-1:::1;00
(18)
17

---

Published in Transactions on Machine Learning Research (11/2025)
Proof of Lemma
A.2 Note that the output of M-step after t-th iteration is given by H(T-1) where
[1]
(T +1)
A/ ((T+1) = [ (
1DD0022 gg- i= ===..... I....,n
h
(T+1)
[0 : Op -::00-2::::1: 2;
n+1
i...
31
3/2
In
In+1
1/1
1/2
3/1
n+ T
T
ott)
34
3
HT++1)
Op-2d-3
D 2d-3
0 D
D 2d-3
h
In
1
(11. 1/1)
1 3 (T2, V/2
( (I mn
0
After copy down and scale operation, the output is given by
32
In
Intl
3/1
1/2
3/2
11
T
8(1)
3(
(
(()
(1)+
PT
D-33-3
D-3d-3
0 D
D-33-3
1
12
In
1--1) (TI,
w
F2, V2)
W g(tt-))
Fn Vn)
In this step, the negative of 3/ is copied down. After affine operation, the output is given by
21
D2
2n
In+1
y/
1/2
y/
TT
3(0
8(0
3
(8(0)
(2++)
F1
12
I'
Op-3d 34
0 D 3d
O 0p-3d-4 3d-4
Op 3d-3d-4
1
t1
(T1. W1)
g T2, W2)
((tt-1)
(N Vn.)
Here, is computed in the next row. After another affine operation, the output is given by
T2
Fn
In+1
y/1
1/2
y/
g(t)
3
3(
PT
3
DT
H(T+1)(3)
Tn
f1
f2
Fn
D--34-
Op-3d-5 3d-5
Op- 13-3225
Op-3d-5
1
t
t
(I1, 81
F2, U2)
Wott-1) 1) Ena Vn)
18

---

Published in Transactions on Machine Learning Research (11/2025)
Similarly, 81 is computed in the next row. After softmax operation, the output is given by
21
22
Zn
Zn+1
1/2
V'=
(1)
n
T
3
3
3
-97
PT
(4)
T1
Fa
Fi
f2
Fa
0
0 D-3--5
D
0 D
0D-3d-1
41
In
1 (0
(F1, v/)
W 97
(22, 122
#4 (Iny Un)
0
The only change is that the responsibility register is updated: the row/column containing w "ah-") (1i, yk) is overwritten by w pr (Ii; ve) for all tokens i. This update exactly the E-step computed by the forward pass at that stage. After copy over operation, the output is given by
21
T2
Fn 1/
Tu+1 ((1)
1/1
1/2
1+ Pe
alf
3T
BT
(()+
0 D
D-20-3
OD-2d-3 1
Dd-3333
(19)
t
tn
1
000
v1)
st (22 12)
W " (In (Tn Un)
0
Finally, this transformer gives the output matrix H H (T+
as equation 19
Lemma A.3. In each M-step, given the input matrix equation and equation 18 there exists transformer (f( with T + attention layers that can implement T steps of GD the loss function ( : Σ (x] (, vs), where (x) B, yi) = "2 (ax, V:) x. 3) + (1 wa a (Zt (:)) x B)?. Furthermore, the output sequence takes the form
(T+1)
Jn 34 Op-2a-s: 1; 4: (
( (c.,
1,...,n,
(T+1)
84 (1+1)) a((+1), 10020-3:1;1;01
a+1
Proof of Lemma A.3 The conceptual basis of the proof draws from the theorem discussed in Bai et. al.
2024h. By Proposition C.2 in Bai et al. (2024). there exists a function f R² R of form
f(s,t) Σ Cmo (ams + bmt dm) with Σ 1Cml 16, [aml |6ml + |dml & 1, Vm [4],
such that sup(er) (f(,tt t)) & c. Next, in each attention layer, for every m e [4], we define matrices Qm, Km Vm DDD such that
3
0₂
bm
1/
dm
(N 1)mcm
Qmhh+
Kmhg-
R(1 - -t)
Vnh
N
Jj
0 2p
0
0
0
0
0
19

---

Published in Transactions on Machine Learning Research (11/2025)
where D is the hidden dimension which is constant multiple of d. In the last attention layers, the heads ((@+)) km+1l
/ ) m- satisfies
Q k (T) [ 0D-
K Kf 1/ [ OD-d+1].
nyT+D [0 1: Op-4].
2 k a 1p-d+11
(T)
k 1 pl+11 ; :p-al. v |0.: -0:00:0p-
The output of this transformer gives the matrix
T2
Entl T 1n B
V/2
Vn
a(t+1)
g(++1)
(t++1)
g(t+1)
H(++1)
OD 0 2d-3
Op 2d-3
Op 2d
pd-333 1
1
1
t
t2
tn
1
W # (11,81)
w m
(22, 92)
w get (Tas (Tas yn)
0
Combining all the architectures into one transformer, we have that there exists a transformer that can implement gradient descent EM algorithm for To iterations (outer loops) and each M-step (inner loops), it implements T steps of GD for function defined by equation [6] Finally, following similar procodure Theorem of Pathak et al. (2024), the output of the transformer will give 30RR :- ny r1 T2 o(To+))
which an estimate of 30R T1 x23* H R (x(+1),, whose columns are
that minimizes the prediction MSE. The output given by
h. [1i, (l++1) -,
i 1,..., n
h [int1 111 3OR
31 1,1
g(T++1)
In the remaining part of Section 3 we give the proof of the estimation and prediction bound presented in Theorem 3.1 and Theorem 3.2 The transformer described in Theorem A.3 which is equipped with T + layers, implements the M-step of the FM algorithm by performing T steps of gradient descent the empirical loss ( (3). Therefore, it is sufficient to analyze the behavior of the sample-based EM algorithm in which T steps of gradient descent are implemented during each M-step.
To begin, define some notations that utilized in the proof. We denote ((0) any fixed initialization for the EM algorithm. The transformer described in Theorem 3.2 addresses the following optimization problem:
argmin (3)
1 Σ fo (I,, v. x. 3)² + (1 (3- (zt, v1) + x 8)93}
for some weight function Wgrn 10, The transformer generates sequence 3
2 with $(0)
g()) = arg min i (3) as - co. More generally, we denote 3(0)) as the minimizer of the loss function
2. (3) at each M-step. Additionally, a(1-1) ,((-1) represents the sequence generated by applying T + 1 attention layers of the constructed transformer in Lemma A.3 on the loss f.(. (3)).
The approach to analyzing the convergence behavior of the transformer's output, TF(H), involves examining the performance of the sample-basse gradient EM algorithm. This analysis is conducted by coupling the finite sample EM with the population EM, drawing on mothodologies from Balakrishnan al. (2017 and
Kwon et al. (2019).
A.1 Results population gradient EM algorithm for MoR problem
In this section, we present some results regarding the population EM algorithm. Given the current estimator of the parameter 8* to be g(t), The population gradient EM algorithm maximizes (see Balakrishnan et al.
20

---

Published in Transactions on Machine Learning Research (11/2025)
2017
and Kwon et al. (2019))
Q(51 = l g (x,,)) + (1- Wggn) -go
(X
whose gradient is given by E tanh (fz YXT g()) 8]. Rather than using the standard population EM update
g(+1) arg max Q ((0) E tanh
1 YX 310) vx
(20)
02
the output after applying T steps of gradient descent is employed as the subsequent estimator for the parameter 8*, i..
g((+))
g((+1) = a1- (1- (1- tanh (gp yz yx g(0) xx)
)
rx
(21)
where a e (0,1) is the step size of the gradient descent.
In each iteration of the population gradient EM algorithm, the current iterate is denoted by B, the next iterate by 3 and the standard EM update based on equation 20 by Br. We concentrate on single iteration of the population EM, which yields the next iterate 8". Consequently, equation 21 becomes:
+ (1- (1-(1-a)))
(22)
We employ techniques similar to those used in Kwon et al. (2019 for basis transformation. By selecting Uj = B/ in the direction of the current iterate and 02 as the orthogonal complement of Un within the span
of {3,,8)), we extend these vectors to form an orthonormal basis {vi...., va] in Rd. To simplify notation, we define:
61 = (B,1)) = 1312,
6 = (5 v1)
b = (B*,02) 02).
(23)
which represent the coordinates the current estimate 3 and g*. as:
The next iterate Sr can then be expressed
=1- + (1 (1- tanh (a1b) 82
a;
(24)
based spherical symmetry of Gaussian distribution. The expectation is taken over a; NO 1) and Y 0a N(a,h + azbş, 82). Without loss of generality, assuine that b, bf b > 0.
Lemma A.4 is analogous to Lemma 1 from Kwon et al. (2019). It provides an explicit expression for 3 within the establishod basis system, demonstrating among other insights that 8 resides within the 5pan(3,35 Consequently, all estimators of 3 generated by the population gradient EM algorithm remain confined within the span( (*)
Lemma A.4. Suppose that a (0, Define 82 = 02 + +z2. We can write 3' = u1 + b2U2, where &f and 82 satisfy
ij (1 a) Th b + (1 1 (1- (1-))) R),
(25)
8/2 = (1 (1 a) b555.
(26)
Here, S > 0 and R> 0 are given explicitly by
S = E B.o. tanh ( (2ybe (y a1b1)) + w (y + anb") tanh' (sabe (y abb"))]
(27)
and
aft, tanh's (aub) 82 + (y+arbt)) aabb)))
R (02 + | Ba..
(28)
Moreover, S iff b 0 or bf 0.
21

---

Published in Transactions on Machine Learning Research (11/2025)
Proof. The proof of Lemma A.4 is directly adapted from the argument used in Lemma 1 from Kwon et al. (2019), applying equation 24 for our specific context. In equation 24 the inner expectation over is independent of a; for i > 3. Consequently, taking the expectation over a; for i z 3 results in zero, confirming that 3 remains within the plane spanned by 11, U2. This allows us to express S as 3 = B122 ur + 8222 with
(bia) 1/ (1 aa) +(1 (1- (1- Bo Exlan tanh
y) 8 (auy))aaj ya]
(29)
02
8/ (1 (1 ))a)Ba,,aa Bylo1 tanh (baa)
02 y) y 102)
(30)
where the expectation is taken over a, ~N(0, 1), and y az (a + 0zbÿ 2²). The computation from
equation 29 and equation 30 to equation 27 and equation 28 is identical to that in Lemma 1 of Kwon et al.
2019
The findings in Lemma A.4 align with Lemma from Klusowski al. 2019 As the number of ileralions T approaches infinity, the estimator S' converges to the standard population EM update
X, (0-
g Ex-N(oo) (Exxxo/-D(EIxxxxxxx..
tanh
12
For any number of steps T, the angle between B and 8* is consistently smaller than that between and S*. This can be observed by noting that:
b/2 Uj
(1 a1 (1-11a)7)55 T) (1-1) a)) R)
0 o tan ((8',)
tan ((39,
(31)
These relationships demonstrate the geometric convergence propertics of the estimation process. Motivated by equation 31 examine the behavior of the angle between the iterates z00) and _ For clarily, use Oo. O, and 0' to denote the angles formed by 3* with 3(() (the initial iterate), 3 (the current iterate), and 8 (the next iterate), respectively. Using the coordinate representation of 3" equation 25 and equation 26 the cosine and sine 6' can be expressed by
(1-a)* - (1-
R
cos 0'=
1 (111--)) (b*
(1-a) (1-00000000000000000000
sin 0'=
| (1- (RF (2010s6))+2211111111111 +2(1-a) (1- (1-(()))) (b$S+R)
Lemma A.5. There exists non-decrcasing function ((A) on Ae [0, 1] such that
(p() V1 + (S/R)2 2(S/R)b
((1) = 1.
As long as A [ 4) and a E (0. 1), it holds that
sin 0' < (((1 a)²) sin
and
2n²²
(1 v1 (S/R)] 8 RS/ (p())
<
1+
cos
<1
n
Similarly,
cos 0' > $((1 a) cos cos
22

---

Published in Transactions on Machine Learning Research (11/2025)
where
(3 1 + 202) 02) (3) 3 202) > 1,
(00) = V (11
(1)
Proof. We provide the proof for the sine case, and the proof for the cosine case follows a similar approach.
Define À = (1-a) E (0, I], we have
by
Mb1 (1 2)
sin 8'
||3011
V V(Jin + +1 1) + R))² (A-0 0 (1 - X)(35))
Ab1 + (1-
sin 0
(1 RI)² + (1-))(111)))))
Then we define the function (P)) to
Ab1 + (1-A)R
((1)
VAb: + (1 2) R))² + (A.0+ A)(35)))
By symmetry, one can assume that the angles <(3, </3 ,+) 3. The non-decreasing property of (A) can be casily verified by the fact that 3' is located on the line segment between the current iterate and standard population EM updates 3 based on equation 22
In the remainder of this section, we discuss the convergence of the gradient population EM algorithm in terms of distance, as presented Lemma A.T
Theorem A.1. Assume that 0 < t/8, and define 83 02 b If bら or #
bi then we have
18' - 3e1 2 & ((1 - a)² (1 (1-(1 a)²))))))))) )x) |1 - 33112
1/2
(1- (16sin 00) 1?2"
where K = (V1+ min (2),,b))/08) Otherwise, we have
| - 101 $ ((1 - a)T + 0.6(1 (1- 3 3112.
Proof. The proof of this theorem is direct corollary of Theorem from Kwon et al. 2019 by noticing that
1 1 - - (1 (1 - a))) -3*1 131/2
& a1 118- 30112 + (1 (1 -1) 3 181
A.2 Results in sample-based EM algorithm for MoR problem
In this section, we present results concerning the convergence of the sample-based based gradient EM algorithm. We begin by deriving the update rule for the sample-based gradient FM algorithm, which incorporates T steps of gradient descent. Starting from the previous estimate, (((-1), we define E = E N,त, The new estimate, g(0,, obtained by applying T stops of gradient descent to the loss function L.! (4-1(()) (B), specifically:
B(0) -14-11
23

---

Published in Transactions on Machine Learning Research (11/2025)
a n
1
1
+
tanh Ui g(t-11) yo
n
n
=11- - a
(I
aL))
a
tanh
(zzliiz) /(( yo
71
a
tanh
1
U
3((-1)
y:
1=1
=(1- - 3 1) + a (1-1))))1))))
1 Σ tanh g22 Mer; g(t-11) y1... y1...
For the analysis in the remainder of this section, we denote the current iteration as B, the subse- quent itcration resulting from steps of sample-based gradient descent as 3", and the itcration follow- ing T steps of copulation-baser gradient descent as Br. By define fu == - tanh (#roaz) () you, and r :- Etanh (#2TTT)) YX, we have
3 = (I a -- (I- at))) (I--\)))
8 = (I - aI= + (I- (I aI))
In the previous analysis,
3' - 3t = -I-
a -I-
8*) + (I- (1)) - B*).
(BiSern13)
(Br(ar, 3*)
$ 3* = $-1
yeRt tanh
Ey
yr tanh
02
02
= - 1
1
yoaz tanh
Wi(Ti,, (3)
Ey 1
yiuk tanh
(W(aee B)
82
922
:-11
1 n
ys tanh (Uik0t,3))
1 n
+ Đー1
Egg n
Σ
Eg n
Σ
yeo tanh
y2
I=1I
Then [|4llap 1 + +0 by standard concentration result and it requires > Oddog" in the end.
n
Conditioning on the sample covariance matrix has bounded spectral norm, |T/||2 0 (1 Finally, for and its angle with 3*, 0 is less than 7o with n = ( 1777|2 <
each fixed B3 satisfying |18]2 > 10
|
(0.95 + +/Vd)113 - 1000-111 1/..
This can be improved by
8" = (1- (I-(1-a$)")
a))))
tanh 1
(2 3) yuura - 3"
72
=(1- 34) (1 (1- (1-(1- - )2
n
x tanh 1 (92 3)):ii)
yior -3
n
f-1 --
4=
1 Σx, [XX11xx0))1)) Ex,((XAxxn(3]]
- Ex.Y [XYA
A
24

---

Published in Transactions on Machine Learning Research (11/2025)
Σ aays tanh (yer) /02) y, n
Σ
raye tanh (year 3 /02))
+ n
where (3x,()3)= == tanh (yxT3/00) 3/02) - tanh (yxTBM BM /0/02). Then
A 0.9||3 - 1
A2 < (IB 30 1) d dlog² (n|8*12/0)/m,
A3 CV looo (og((//))
with probability at least 1
A.3 Convergence result in Theorem 3.2 under the high SNR setting
We first present the results for parameter estimation under the high SNR regime.
Lemma A.6 (Lemma in Kwon et al. (2021)) For any given r > 0, there exists universal constant > such that with probability least 1 6.
1 11/2 (l
sup 1r
where
11 E XY tanh (YXT B)],
A tanh
Σ 1 T
Lemma A.7 (Lemma in Kwon al. (0021) 6d exp ( nt² 72
For each fixed B, with probability at least 1 - exp(-en)
Σ ya tanh (ui(,,, B))
Ey y tanh (uK(z,, ()]] t
for some absolute constant > 0. Theorem A.2. Suppose that r) 2 Od (n/66))) 1/4 and cos <(3 (()) > 0.95., let {,300} be the iterates of sample-based gradient EM algorithm, then there exists a constant 7/2 E 1) such that
for some absolute constant C and |000| > 3*
(,00)) 8*12 7/ 1- 7/2
o(v (log((//)/8)/1
holds with probability at least 1- 58.
Proof. Without loss of generality, we can assume that 1. Denote 3 as the current iterate, and 3 as the next sample- based iterate. We first consider
3 (I a Σ (I-
Σ
tanh (192 Vi 3)) yeore - 3
n
= 3r) + (1 (1- -1-
$ tanh 1
(Fothee uia B3) Vuue 33
25

---

Published in Transactions on Machine Learning Research (11/2025)
1/4 & y $ for some universal constant
We prove the results two cases, i.e. 7 a and Co Co- When nj > 1, based on the analysis in Kwon et al. 2022Ip, with probability at least 1- 56,
|A|2 < (0.9 + Vdlog" (n1|8(12200) | 8 + Ciy Sil/loog" (n18*12//0//
$ ~1113 - 1 + Civ dlog² 12/0//n
(32)
where: 0.9 + V d logg (n| 3 2 /n). By standard concentration results on E - 1, it holds that with 30 (1/6), [[| a op & (1 a/2)² with probability at least - for appropriately small a. Along with equation 32
- 3d 2 (1 - " 18 18 + (1- (1 )
[11 2) T (1
(1- -1-(1-1 (1-2)))7 1B 3
+ (1- - (1- -
(1-2))) ())) dlog (nl 12/0)
(33)
Define e(n,,) = (1- (1-1)) d (n1,3*122 2 /n and (1 - -)²² + (1- (1-(1-g)) 7)) As long as
y 1, iterate over t based on equation 33 and obtain
[ 3$1 (21.8(t-1) - + c(n, & (111 3/ + (1+ 72)
1/1101000 1212 1- dn,,8) 8).
In the remaining part of the proof, we present an analysis of the convergence behavior of the sample-based gradient EM algorithm when Co (d (n/8)//n/ & < & that
By Lemma from Kwon al. (2021) holds
[[ tanh (YX B) - 3* < (1 11) 8*
To systematically analyze the convergence, we categorize the iterations into several epochs. We define Co = []((0) 3* 12 and assume that during each 1th epoch, the distance ||3 - 3 lies within the interval [1,2--1,C,2 CJ.-1. This stratification is conceptual and does not impact the practical implementation of the EM algorithm. The key idea in this part is the same as Kwon et al. (2021). During the Ith epoch, the improvement in the population gradient EM updates must exceed the statistical error for convergence to occur, formalized as:
(1
(1 a -1--)))) 1 > 2ery (n/8)/n
||*0|| Go2-1 and using triangle inequality || < where e is the constant in Lemmaa A.6 By setting
1,8 12 133 - 18 in 1th epoch when
[ (1- a 13
2062-011 > / d dlog" (m//)// (1-
4c(113*1 + (o2 2) log (m/0)/m
is guaranteed to be true, then it holds that
||A||2 & (1
16 110* 12) 18 111
1
- 1a1 & [1
(1 (1
-
-)
1
|13 -3*11 -3*1/2-
2
16
26

---

Published in Transactions on Machine Learning Research (11/2025)
Recall that 77 2 (I/d)og(n//0)/)) then with appropriately set constants
18812 > (c1 + 1/V d
we can deduce that B moves progressively closer 3 long as C,2 & c2113*112 V d (n/6)/n. This process requires 1(1) iterations per epoch, and after O(log(n/d) epochs, the error bound ||8 3
02 12 -1 (w/6)) is expected to hold. Thus, the convergence rate for (8t) towards 8* is quantified as:
1,300 - 12 < 221 - 881 + 1- 72
V dlog (n/8)/n. dlog2(n/8)/n
A.4 Convergence result in Theorem 3.1 under Low SNR settings
We present several auxiliary lemmas that will be utilized in analyzing the convergence results for sample- based gradient EM iterates.
Lemma A.8 (Lemma in Kwon et al. 2021 There exists some universal constants Cu > 0 such that,
11811 - 41312 - Ca || 1
E Lanh (YX) x)xx]l < [18]2(1 1912 cul13*133)
Theorem A.3. When & Gold (n/)))))) there exist universal constants C3, Cs > such that the sample-based gradient EM updates initialized with ||300)|2 < 0.2 return 309) that satisfies
(8(0) 181 0 co(dlog"n//1)
with probability at least 8 after t Ca(1- (101-(1-app)) (1-a/2)²)
" log(log(n/d)) n/ (nlo2(())/)) iterations.
Proof. The proof argument follows the similar localization argument used in Theorem A.2 Define c(n, = s//log" (n/oo"nn/)))) (n/5)/n with some absolute constant C > 0. We assume that we start from the initialization region
where ||81|2 en (n, (n,6) 8) for some ag [0, 1/2). Suppose that c (n, 6) = |88 en (n, at the 1th epoch for 1> 0. We let C 0 sufficiently large such that
c(n, 8) > 4c, 1c +
sup BeB(
|1- - 10---1p -1l1/1mmm
with = c. During this period, from Lemma A.8 on contraction of population EM, and Lemma A.6 concentration of finite sample EM, we can check that
| All2
11112 0.5|1112 + cu cul18112118*12 + sup BE
1n-S-1m1
1 30t+1 (n, 8) 4 -a+ (n,
|
2
1/811/2 < (1 )
T 11 + (1-(1- a T
(1
9)
||S-1P111 #112
|1112 (1 - (1- (1-1)) 2
1
1 a (n. 8)
30rt3 (n,, 5)
4
Note that this incquality is valid as long as 1 (n, 8) < || $ (n, 6). Now we define sequence aj by
al (1/3))) - 1/2) + 1/2
and 04 + 1/2 as - s0. With this choice of ag, c (d/ (l/n))= 2/4, Hence during the tth epoch, we have
( (1- (1-2)" ") or 1
|/81 < 181| 2 4
27

---

Published in Transactions on Machine Learning Research (11/2025)
Furthermore, the number of iterations required in 1th epoch is
4 (E (n, 8) - c (n1, 8))
(1- (n.
11-1112)
(1 (1- (1-11-7 9) c
Cart1(n (eal-)))a
When it gets into (1+ + 1)th epoch. the behavior can be analyzed in the same way and after going through epochs in total, we have [8]2 & e (n, 6). At this point, the total number of iterations (counted in terms of steps of gradient descent) is bounded by
(1-(1-2)) (m))))))
By taking 1 = - (---CC(1111000 log(1/6) for some universal constant C such that aj is 1/2- for arbitrarily
small 0 > 0, it holds that
1 12 m (12-0((, = (oddlogg(//)))))) 1/4 0/2
with high probability long as t 3 e 8), (d-m(1-(1-a23)) log(1/0) where i some universal con- stant. By taking 0- C/log(d/n) and using triangle inequalities, it holds that || & e(dlog" (n/5)/m) 1/4 and [| 3+1l2 2 & ci (dlog"(n) (n/0)/n)* 1/4 where e₁ is some universal constant under low SNR settings.
To finish the proof, replace 8 by 8/log(1/0) and take the union bound of the concentration of sample gradient EM operators for all I = I,..., C(1 - aa2))))) log(1/0), such that the argument holds for all epochs.
A.5 Proof of Theorem
3.2
Proof of Theorem
3.2
The existence of the transformer follows directly from Lemma A.2 and Lemma A.3
A.6 Proof of Theorem 3.3
For the data generated based on model equation 1 with two components, Bat1 = -3* with probability 2 and 3x+1 = 3* with probability 2 For any choice of B Rd,
Epa.y [[o++ - x 111+11891 = Ep,,, [x 1n41 - x 3 + Un+1)²
0² + Ep,,. (IT+11 Bn+1 - x x }²]
T8) 0² + Epr.y tr(3n41 () - -B)
0² + Ep.., tI T (Bn+1 - (3) - -8)
0² + Ep.,, 18n41 - -B112
1 11 -B12 + 119** 3113-
0² +
Therefore, [[u+11[[1111 1 is minimized at
3OR
2 ge
g
(34)
the optimal risk is given by 82 + 1 12. And same results holds if the estimator 3 depends on previous training instance (x1, (21,91.... V1. ......, Tna vn) and the expoctation taken w.r.t. P. And in general MoR problem, the vector that minimizes the mean squared error of the prediction given by
K
goR= == arg min Ep. [xx B - - = Σ r 37.
BER
28

---

Published in Transactions on Machine Learning Research (11/2025)
Proof of Theorem 3.3
The oracle estimator that minimizes the MSE, i.e. MSE(J) Er[(J(H) -y+t)) ) The output of the transformer is given by
is given by equation 34
Da+1 = ready (TF(H)) = sr, ,3R
where 30R is given by
3OR = T1 -m))
with B = (eada(TF(H)) for L - -1-022)
(
/d) n/ (nd (d1))))m//)) in the low SNR
settings and O log ((lag))) in the high SNR settings. Note that 1 - 3OR la < T1 B|g + (1 T1)1| 2
- E 188 $12.
Under thelow SNR regime, after To 3 O (ogglog(n/dd) (ndl///g/d/o//// outer loops,
- - 30R & 0 (clog(mn00) }}
n
with probability at least 56.
Under the high SNR regime, after To e 0 log (mlag))
outer loops,
dlog (n/8)
1 3OR 1 & 0
with probability least 56.
We would like to bound
Ep (TFF)
inf Ep (x 1n++)² 2²]
3
Note that the Epp (Wh+1 (TF(H))) given by
Ep [x] [x 30R Mu+z)²| Ep (3 gOR + 30R) - yn+z)²
Ep 10T_,300 (OOR)
+2ppp (BRR- T anti(ETT gOR yu+z) + Ep (o 1004, bn+2)
Hence, when कI = π2 = 1,, gOR T1 T² = 0,
Ep
(yn+1 - rcady(TF(
inf Ep (IT -on+)))++)) 27 )²]
3
Ep (r (BOR 180R-190 )30R)) 27 + 2Rp (GOR 3OR) T an 11
-Ep
(30R - (OR) T xn+ x
n+ (3OR - 3ORR
-Ep I ut (3OR_ 3OR) (3OR _OR) I)
=Ep|| 30R12
Under the high SNR settings, it holds that
P 1130R 3OR
< dlog)(n/66//n) 2 21-8 8.
29

---

Published in Transactions on Machine Learning Research (11/2025)
Hence, by integrating the tail probabilities we have
+
B|
P(] -[]30R-B0RRI - -30R122 > vi)ct
to
|[[]3R- 300 * vi)dt
Jol
6
+00
P[]3OR- - 30R e vi)dt
+
P(] 30R| vt)dt. dt.
< 01
Sctting vi (dVllog
and solving for give 8 & nexp { /nt/dd.
By taking
Odlog"n n for some absolute constant C, it holds that
2 n dlog
+00
E 3R
BOR 2
nexp { /t
n
dlog 2
n
(2d + 1) log
0
0
n
dlog 2
n
0
Under the low SNR settings, it holds that
P (1 3R - 3OR < Od' log (og8((///))))) > -.
Hence,
+ P(130R-
30R 1
-30R11 > vi)
J0
0000
P[[1OR 30R > vijdt
1dt
+
P(]BOR 3ORl| vt)dt
x
P(]3OR
3OR|, vt) dt.
Similarly, setting vt 0( log* (n/6)/nt) and solving for give us 8 nexp taking /Vlllo" n/n for some absolute constant C, it holds that
vn/dtt. By
E 1IROR 3OR < o
d'log" n n
+ +x
nexp { ddtiifdt
Je
0
(o(Vd/in //dnloog
=
Combining everything together, it holds that
Ep (Unt1 - read, (TF() inf Ep [x= yn 1
3
30

---

Published in Transactions on Machine Learning Research (11/2025)
O Ing 12 dfnlogn)
> ((dlog"(n/8)) (n/6)/n)* 7 Cd (n/0)//))))))))))
A.7 Proof of Theorem 3.1
In this section, illustrate the existence of transformer that solve MoR problem with K > 3 components in general. Given the input matrix H as equation and initialization of m K there exists transformer that implements E-steps and computes
J M exp (-
1 (ve - x7 xT 202
(1+1) Mi
(++1) 1
(1+1)
(35)
(1) y 5 Ih= exp
1 (ye - 27 )2) 202
since the computation in equation 35 only contains scalar product, linear transformation and softmax op- eration. Next, following same procedure as before, one can construct T attention layers that implement gradient descent of the optimization problem
(1+1) (ve- - (00-ATxz) xe) xe)2 for all j [K]. 7ij
min BERd
as the gradient loss ((xT we) =
((+1) (ye - 3 xT)² is convex first argument and 0gl(s, t) is 2m Bij
(0, +000 4, 16) approximable by sum of ReLUs. Hence, the construction in Lemma A.2 and Lemma A.3 also holds.
at step & 1, the population EM algorithm is defined by the updates Given the estimate - and x, (t-1)
" 1) exp { -
X 1)3
EmjK (1-1) exp (Y XT XT [)))))
8/0 [w] (X,Y)xX
(0 (Fl,, ()x)),
(1) E E[ (X.
In the sample version of the gradient EM algorithm, we define 50 4/4 Σاك w (r., The new estimate, g(t), is obtained by applying L steps of gradient descent to the loss function
( 2n 1-1
1 (0) w (zi, v:) - xx 3)²
Σ
starting from - Specifically,
a = (I -- (i-as(() T 3/1 (I- Σ 5
[ w (0 v:)ai
In the finite sample gradient version of FM, the estimation error the next iteration in this problem is
0 - 83 (1-a210) T "((((-))
-3, ,7) (1 (1 -1-11-200))) 1 " (zi, v:)u., 33
(2.
Define
J.. exp - (Y -x-xx1 x (33 )
Σ N. exp [ XT (3)33
31

---

Published in Transactions on Machine Learning Research (11/2025)
then we have
(X, VX) xT - 0,
since true parameters are fixed point of the EM iteration. Hence,
8 83 (I ( (
88) (1 (I- (1- ( ((()))))))))))))) - "les + B].
eB = Lij (I) 11 - (X, - -x X
B E 3 (x, X) x V
In Kwon & Caramanis 2020 the following results are proved.
Lemma A.9 ([Kwon & Caramanis 2020 ). Under SNR condition
1) > CK Pr log( Kpx)
with sufficiently large C >> 0 and initialization condition
(t-1)
max 1T 0.( & 2
max |((-)) - :32112 & Kp log( px)
for sufficiently small e> 0. Given n > O(mas {dlog (dK2/8), ( (KK21)).(K318) (K2 samples, we get
Kmy² x2
Km
leal2 &
log? (nK2/6) max [1
-3r 19/22 +
log² (nK2/0)
Tmin
Tuin
with probability at least 1 1-8. 8.
Lemma A.10 ( (wwon & Caramanis 2020 ). Under SNR condition
n b CK p. (og(Kx>))
with sufficiently large C 0 and initialization condition
max (t-1)
[T("-1) xz & Tmin
max [[((-1)) -33 & Kpn log( px))
en
for sufficiently small <> 0. There exits some universal constant c3 (0, 1/2)
B & dB max ||st _8 l.
Now, it remains to bound the maximum eigenvalue and minimum eigenvalue of the weighted sample covari- ance matrix (+()) Define the event
Eg = the sample comes from j-th component
and Pze n//nz for j l. Note that
1 n
1 n n
w (t) (r,, v.) 1g,
Σ
w ig
(Ii,
9b
Σ
x
n
32

---

Published in Transactions on Machine Learning Research (11/2025)
By standard concentration results on Σ - I, it holds that with > O(dlog(10)),
3 2
Amax ($ & Amax ($)
with probability at least 1. The concentration of -i= Wij (z,, y1 a 1g, comes from standard con- centration argument for random matrix with sub-exponential norm Vershymin (2018). Since w E (0, 1) and 24 is standard multivariate Gaussian, then by Appendix B.2 in Kwon & Caramanis (2020), it holds that
(0)
+
Σ
(1) a 1e, B[wf (X, 1xxx w
< <0(v
d loog K2/6)
with probability at least 1 - 8. By Lemma A.3 in Kwon & Caramanis 2020 it holds that under the same SNR condition
(X,Y)XXXI) T
Amin
E Juy
Therefore, as long as n > Oddoog(K') Tmin), ) it holds that
Amin
(5() Amin
Σ
2, 18,
m..
Therefore, we have under the same SNR and initialization condition, as long as
n O max (dK* 1/3 dolog( K //// mm)3)
it holds that for appropriately small a,
1(1 - -2- T
12 & max - 33//2), - am == YT,
(36)
Km Tiin
Km
d logg (nK3/6).
leBl2 <
log? (nK2/8) max (| - 2
(37)
n2
Tmin
n
B x
max ||s6
B
(38)
4
|5000
(39)
Tmin
For appropriately small a, we have "rr E (0, 1) Therefore, combining equation 36 cquation
37
cquation
38
and equation 39 together, we have
8 g(() & YT + (1 -rrr (v
Kny²
d
log" (nK2/8) + x/ ) max (1 - -39 2
V V=
Tmin
K Kxya
d
log² (nK2/5)
Tmin
n
with probability at least 1 - 55.
To derive the concentration results for
(0) 1 " wiy (z: - (x, (X, define following events
εg = {10l & TE)
E,2 = {ACKX, Ai)1 (X, & (X,
E,a = (X, 3 - -8t1 > 4V2m
33

---

Published in Transactions on Machine Learning Research (11/2025)
Et, good Ex,1 n Eei E1,
where D = B(1-1) -8, then have the decomposition
K
(2,, (21, v1)
Σ
w (Is, y1) U.) + w (2 yi E goods
+ w (zs. yi)le,
Therefore, we could bound
( W (Ii, Vi) EgnEt
(1) wij (xt, Ui)
Tenee, wood
(t) iy (Iz TE
(t) wis (Irs Uid
E
good
good
wy (Tt Mele,
B wy (FM
respectively. For the first part, note that
Ee n Eg.good 1/p [1 (It, vi)
|lu (ai, y1) Eemny good 1/2
supp 1/2g
p2l
Cpe exp (-
Therefore, with probability least 1- 8/K2,
! 1 Σ wij (ot, Vi) LEr - E3 Wiy (xi, Bi) E
]|
< (p(p exp (-72) op
1 log K 122/))
For the second part, note that
[10]9((0 (xt Ui))" Ee E good 1 1 1
Iuo (0) (Ii, ui) Eenl goodlly llve = sup p 1/2 ED
p=1
P(Ee El, good ) 0 (TTT (K
Therefore,
(1) I E 10 22
(2t, (2) Eeney B[u! (Ii, Vi) gener
*
log(K2/8) n
log (og(K(2/))
Kpr
Similar to the second part, have the following concentration result for the last part:
1 Σ "h 10 (1,,)))),
B[ (a. (.),)))),)] & =0(
log(K2/8)
log( (2(K))8)
n
71
71
Combining three parts together, we have
(Z., v:)
E[ (x, Y)
& x0
log(K2/8)
Σx Pey exp (-7²) +
+
log(
n
K log (K2 numin
K Lits
v Kp
K log (2/6))
(og(K2)6)
34

---

Published in Transactions on Machine Learning Research (11/2025)
with probability at least 1 - 36. Therefore,
| 1= 3OR = 30R| < 1n+1 12/ 3 - -BO0
xn+ l max j - qÿ max 1,3112 l max ( max 113, - (3112)
5
$ log(d/6) (V
K log(K2/0)) (K2
Knzz dd log (nK2/6)
7 +
Mim
Tmin
with probability at least 1 - 95.
B Proof of Theorems in Section 4
B.1 Proof of Theorem 4.1 in Section 4.1
Proposition B.1 (Proposition A.4 Bai 2024 Suppose that (Xaleee is zero-mean random process
given by
N X9
Σ f(z:0) Ex[f(2;0)],
where #1,"" IN are i.i.d samples from distribulion Pz such that the sollowing assumplion holds:
(a) The inder set θ is equipped with dislance p and diameter Further, assume that for some constant A, for any ball Ø' of radius r O the covering number admats upper bound log (6, p) d log(2Ar/8) for all < 2r.
(b) For any fixed 0 E O and sampled from P,, the random variable f(z;0) is (G) Gau--Gassiin random variable.
(c) For any 8, e and & sampled from P,, the random variable f(z; 0) f(2:0) is a SG(B ((0,8))
subGaussian random variable.
Then with probobility at least 1 5, it holds that
Xol & CB. dlog (2Am) + log(1/0)
sup OE
N
where C is universol constant, and denote 1+ B'D/BO.
Furthermore, if we replace the SG in assumption (b) and (c) by SE, then with probability at least 1 - 8, it
holds that
dlog(2Ax) + log(1/8)
dlog(2Am) + log( 1/8)
sup Xe & CB" 6
N
N
1/p
For any p [1, let |H]l,,, R 0, we denote Hn H H 2,00 & R. be the ball of radius R under norm || ||2,00
Em [14a1B)
denote the column-wise (2, p)-mm of H. For any radius
Lemma B.1 (Corollary J.1 Bai et al. (2024)). For single attention layer Oatea {(Vm. Qm, Km)) me[M]
DD and any fixed dimension D, we consider
+ Oattn.B' (0am Oattn < B'}
35

---

Published in Transactions on Machine Learning Research (11/2025)
Then for H E NR. Pattn B. the function (Gatts H) Attme (H) is (B"R")-Lipschitz w.r.t. aatnn and (1+ + B Lii'ssiits w.r.t. H. Furthermore, for the function TF given by
TF (0. clipR (Attnemu (H)).
TF is Bg-Lipschitz w.r.t 0 and Ly-Lipschitz w.r.t. H, where Bg == B R and Bh := 1+ R Proposition B.2 (Proposition J.1 D, we consider
Bai et al. (0244). . For a fixed number of heads M and hidden dimension
OTE,L, {0
(11:1) attn
M M, MDD D. & B).
Then the function TF is (LBH Be) Lipschitz w.r.l 0 E OTF,L. L.B for any feaed H.
Proof of Theorem
4.1 Define events
Ey =
max [e[n+I]aje[P
110,0113 Bu
Eg
max ic(n+1].jE[B]
{10,001b 12] & Br Boyo
and the random process
1 B
Xo
Giol (0 2( - [ic( z)]
where (1:B) are i.dd. copies of Z P, drawn from the distribution P. The next step involves applying Proposition B.1 to the process (Xo) conditioning events e To procoed, we must verify the following preconditions:
(a) By Wainwright (2019). Example 5.8), it holds that log where (r) is any ball of radius r under norm ||
< L(3MD)) + 2r/6),
(b) Ga (0; 2 4B2 and hence 4B2 (B--sub-Gaussian
(c) (m(( Z) Lial (a (0:Z) Z) < 2B, (LB - ह| by Proposition B.2 where Bg = B²² and
Bи + R²
Therefore, by Proposition B.1
conditioning on En o Ey with probability at least 1 - E.
L(MD2) + log(1/E)
sup Xe| < B3
B
where 20L log (2 + max (B', R, (2B,) -1)). Note that
Var(o, B.) Eg, Var 2. Be 3.) + Var [E.[ B. 11..))
Σ RII
K
K
E AcC 8.)
Σ
T day
ox exp { 12
36

---

Published in Transactions on Machine Learning Research (11/2025)
Denote 7i be the sub-Gaussian parameter of x. 3, Then have ki- " &
maxiæK || 22- Besides, yi is sub-Gaussian with parameter at most v T. + 0². Finally, note that
1 - 1 $ 2 ( Summing over and j for all i + J, we have
Σ C43
1 53112 < 4(K 1) Σ
3
which implies (K 1) min:45 10 0 & 4(K 1) - the subGaussian parameter тi as below
13812. Thus, we can derive a lower bound of
18
Min Σ
1|
1=1
Mmin
K
min
33 3
mnin
Rauin-
4
Therefore, + (1+ (0/T)²) T.
102
== 1+ Kn2)C Then by taking
Kmin KR
B, = v /log(BB/EE.
B, 2(11 4/(Tmin Kn log( 2nB/E).
R max m By),
we have P(Ey)= 3 - and P(Ex) 3 1 - by union bound. Hence, with probability at least 35,
(1 4/ Kn)))) 1og) L(MD3)e + log(1/E)
sup Xo <
B
where
log (2+ max {B', R, (2B,)-11)
is log factor.
B.2
Proof of Theorem 4.2
B.2.1
Computation of the risk
Given the prompt in the form
E
32
Fn
Ent1
E I (dd+1)((++1)
(40)
Vi
appropriately sized key, query, and value matrices K, K,Q,V V, the output of linear-allention block is given by
A E
+==E+_VE(K
VE(KE)
We start by rewriting the output of the linear attention module in an alternative form. Following Zhang
et al. 2023) we define
U21 u- ) WKQ
(
U11 U12
)).
Wpv
(41)
therefore, the prediction on the query sample is given by
Dn+1
T - T
Un + U21
1 1x y uip + u-
-y X Un + =ui _y y T Int1
n
(u21 + u -u 11
XT (Un + Bo+
. -nntl
n
+ U2i - XT
n- uip + un
=u T x Un + +u
2 x XB.+1uT + LT u. u Ent1-
37

---

Published in Transactions on Machine Learning Research (11/2025)
B.2.2 Gradient of Loss w.r.t. Parameters
In this section, we compute the gradient of the population loss w.r.t. parameters {u21, u12 (n,,-1).11 simplify the presentation, we define v [ur...,,,a va]" E R", X [ x1] Raxd and following notation
1 - -x-x X (Un + 11)
T (221 + u 1-188++111 11+11)
21
2 0 XB 1 v T Un + u-1
22 = U21
x u 212 + u-1
n
n
-u V- u-1 12
23
n
Since X, u, B are independent, we have
2..A (E, 0) B(in+1 (3n+1; X) vn+1)
(21 + 22 + 23 Bnt1. x²²
+ -2
las E (=1 + + 23 - B2+1) (21 + 22 + 23 - Ba+t)) 0²
0² + [d, E (21 - Bo+1) (21 - Bn+1) + B2222) + (Id, E2ges) + +2 (La, E (2) - Bu+1) 23
We first compute St. Note that X independent of and Bat1, then by making use of Lemma C.1
Ta. B21z1)
(U11 (n-) T. 1 n
(um + u-1 (um1 + u B=+1) 1 x (Un + BntiuTa) la x _
=EE
n
n +
E tr (U11 + But1uT2) (wm + u-18nt)) (u21 + u-1 T. Ia (Un + Pu+iul Ia]
1
Esaer tr tr (u22 + 1-1 (w21 + w=18m+1) Ia) (Un1 + ButiuTe) Ia (Un + Bntiu Ia
n
+
Epn11
(Un1 Be)) (mn) v (u21 + u=1 (Un SutruTe)
1
Bs tr (u21 + u 1 (uz1 + w118++1) La] tr (U11 + ButiuTe)" (U11 Br+ 1o
n
n +
(I)
+ -1 (II),
where
(1) = Egn.s tr UiT U2i Un]] + 2E.... tr (Bn+1uTe) Uy Vunuiii Un1 + 2E tr [UT (u-1 uz Un)
+ 2m, tr (Bn+iute) (u-1) uzl uzi Un| + 2T tr (Bntiuje) u2l (u (1+1) Un]
(Bo+1uT2) T uinna (8n+12Te) + Epnes tr UUi (u-18n+1) (u (B=1)++)) T Unl
+ + tr
+ 2E 2m tr
UIi (u-1 (u (u= (2n+1412) + 2B tr (Batiute) T + ERrer tr (Bn+1u(2) (u11mt)) (u 18 18++1) (Bn+1uTa)
u2 Bu T (Bntzute)
[U] uer Un] + 2E (112 U2n /I Un) Er
+ E..,E
[BT Br+ I u U |1 11 + 2u 2 ER,, tr [ U2 ST Un]
+ "u tr I 1 TT 1
x2 Eg... tr B ST Un
2n² Eg... tr [ Ba+ 8T Bn [T]]
2u 1 12 t 1841]
38

---

Published in Transactions on Machine Learning Research (11/2025)
u², 222 1,ntall2
and
(II) Bones (u21 + uu1 (221 + +u1 tr (Un + )ntiuTs) (Un + 3ntiuTa)
Ega+1 [1z]] U21 2u 3 121 + u² 1 Bn+1] tr [1 Un 2012 Un 1
U21 U1] 21 221 R U1 + 2u 1 - U21 Ir [1]
+ Uhi U21 4u
4u
u21 111111111111111
U21 U21
ui, tr Un BS Ant1
Bn+1 2u-1 21 18+1 2u
1 Un142
م² U1 Bn+zll2
For the cross term in S1, we have
(Id, 121,31+1 B((wa) (w22 u-18++)) - XTX (Un But1u)2) I
C(U11, 412, U21, 1-1
E (u21 + 3=131+1 (U11 Bu+ii(te) Bnt1]
Es 1n Un 1113 1| + Esnes [い An 12
+ 1_ 10 U11 u BBts 10t180-14122n1
Finally,
[as E.n+ 1.3+11441) = E33
S2
E
U21 T X "uuguizo" T Xu21 + 1 XUnUU
n²
2u v. XUnuu12 Xun + 4u v X 3 12 XTv
11 X X 1uTu Xuzi
1u X Bn BaniuT
1 E Unn X "uu X Xuzil
02 EXUn
72
n
2u
4u² 12
+
E [0 XU11 Xuzi] +
E tr X
Brt 1 n+1 x
n
n²
+
4u
E [ X Xw21]
4u
X Bn
n²
n²
02
[u u2 412 + 22, tr (UnUi]) + 2u-1u U21 + 4u² u u tr (13++131+11)
71
4u
E J X 134414221200 Xuzı] + 4u-1 4 E [u'
+
X XBa+
n²
Sg 0 on + 2)
2² U12
and
S₁ = 2 (Jas B(2 Bn+1) 25
2TE {[
(w22 + 4u x x (un + Batiula) - -BT
-
1412
= 200um11 {((22) + u 10+1)) (Un + ButiuTe) - 3
u12
39

---

Published in Transactions on Machine Learning Research (11/2025)
20 u (2 U11 u + 1-18ht1U1 + M-10ht13n+1upp -pu+1| u12]
20 2002m1u2 2 U1 12 20 B Sn+ 12
+ 1 Uni42 + 200 12 n 9n+1 - 10 112-
Now we compute the derivative of the risk R w.r.t. parameter U11, u-1, u12 and U21, where
(1)) (II) 2LsA (E,0) o² p²= - (1) (1)++)++++++---0))))-))))++-
2C(U11, 012 u21, u-1) + E S2 + S3 S4.
n
We first calculate the derivative of all components in risk w.r.t. the scalar u-1
20) 21 ir (U 11 Un)) + 2E 1 t2n U11
du_1
2Br
1122 u2n 2u Un1] + 2u 2u- B
Bn+1
A(1I)
2
21 U21 Un) 4E, n+1 u2n 22n Un 2u tr [Un 1 8
du-1
22
u22 Sn+1
4u E3 8m U U1141 + 2u 2u_ uu 3|
@C(U11 U12 U21, u-1)
E B 11t+1u1BBtti
δu_1
and
0S2
402
202 u-1 tr (UI1UT)
80
+ 112U11221 u21 + 1-14 tr (E,r+1
du_1
n
n
4 n
+
B
[t' X 1 Xu21]
+
n
0S3
20
2n + 2)
u-1 (uj2u12).
Ōu_1
OSa
20
u Un + 200 Euzz Bn 14
+ 18 Unr + 40 u u 202 112
Next, we calculate the derivative all components in risk w.r.t. the matrix parameter Uni
a()) 2ugiuh uhi U11 + 2E U21 uu 8n
unn 1 Un] 2u B U11
2011
2u_ U21 E, n+ "
2u tr Bn+ 1 U 3n+1
+ 2u2 J8 3 U1 + 2u² IE, 1 3 3n 1
a(11)
2 un Un + 2ug U22 E, u12 + 4u E,3 U21 Un
2U11
+ 44uu 10.-1E,3-00 Un 3n 112 + 2u U 1 3n+1
+ E3 BT 1
aC(U11, U12, U21, u-
ETgI + 1,1
8011
and
0S,
202
202 u-1u21
4u E [uu2sT+1X XTox"xvvv Too
2 Un
@U1
n²
40

---

Published in Transactions on Machine Learning Research (11/2025)
0Sg
0,
8111
0S4
1 + 202u² BBn+141
@U11
Then we calculate the derivative of all components in risk w.r.t. the parameter u12:
2(1 = 2EUU U21 22 Bn+1 + 2u 1uuiBB 8n -o- 1
2012
2u 1 Bn+ U22 Bnt1 + 2012 Unn un 1nt1
+ 2 BU 3n _
Bu+1 4u 11 ur [ 121 (3 13 1 -1 20 18u+1l22
2(11)
2U R3 U2 U21 + 2m U21 u12 EB
Bm+1 + 4m U1| RB, 1,3n+1ug18nn1 1218++1
tu12
+ 4u 4 - Bn+1
B
2u 18 B 1 8 +
18
aC(U11,
U12, 212
U21 u-1) E [Ba 3n+ 11 + u-1 E 1
Ou12
and
0S2
202
202 u U +
819 un tr [E B3 u12
U12 U21)
012 +
Ōu12
4u
4u BUn "v XBm+1,
E
112
n+1 X Vud x v u1zu Xuzre X XButi] +
n
n²
6S3
2(n + 2)
_ 11:
Zu12
0S4
421 + 40 u- 1212 U21
du12
Finally, we calculate the derivative of all components in risk w.r.t. the parameter u21!
a(I)
20 1U u21 2E Un 3n+ uu U U21 + 2u 2u_, EUU B
Owz1
Un + U1)) 1 U21 + 2u Mo Bnt1,
a(1I) 24r
[UNUU
U21 (E U1))2) u21 + 2u
[Un U1 R,
Own
2uj
(1.87+18441) u21 4u 13 U + 2u 2u 2u Bn+1,
aC(U11. u12 U21. u-1)
EU11BB+1 + 13.+1uT0041
and
0S2
20²
(uT2u1) u21 u 1-1 + 4u
20t
E [XTuvleeeee vuT To].
Juzi n
n
0Sa 0,
Ju21
aSa
202u u0uu1111112 + 2
When 18,+1 0, then by taking the derivative to zero and solving the equations, have
1 Uij =
41

---

Published in Transactions on Machine Learning Research (11/2025)
and
n n+1 1
1 2u 2 1n11U11 uut 2u 1
202 u
2u En Un
111 Un
0,
which implies
n +
L B 11 1 n
I IT
U1
E EB
Epa+1
B.2.3 Dynamics of parameters
Lemma B.2. Let E E [((+1)x(++)) be embedding matrix corresponding to prompt of length and define
Un (w21)
U
R R((+1)x(++1)
Vec(U) E (d+1))²
(42)
u-1
When E3; = 0, then starting with U12
U21 Oa, the dynamics of parameter U follows
dUn
22 U1 U11 n EBBT + (TETT. 3+ B
12 ) I + -u1E,3
(43)
dt
du-1 - ir U-UUU
n E,3T +
(1 E,3 B (
n ) Un
(44)
dt
and (1) (12 0 for all t O.
Proof. Based on the calculation in Section B.2.2 we see that if we initialize 412 u21 and use the identitios
du12 V u12 LsA(E;)))
dt
du21
V U21 Lsa (E;0),
dl
we can see that du12 dugi = The expression of dtn and dt
du
de can be obtained by plugging u12 = u21 0
dt dt
and IE3 = 0 in the gradient Von LSA (E,0) 0) and du LSA O).
Remark B.1. One can show that with U12 U21 Od,
dUn UT
n 1
E33T + (LET B
02 m) mm
-u2,U11 U1
1 + - UII
dl du_1 dt
n 1 EB8
(1 EB 3
02 =/1 }))
2
UUUHH
- u E33
u
n
n
Hence,
d
di ir [n[U
dt =
(45)
Besides, the loss function could be simplified as
[ n E33 + (-EST B 12) ) UniUi
((nn) -² and thus,
u [Un EE3 3 02
n
n
d
Un((t) = di(U11,
d
(t) = (((U11;011)
dt
0U11
di
du
42

---

Published in Transactions on Machine Learning Research (11/2025)
Lemma B.3. Under initialization condition equation 11
when the parameter y satisfies the condition
2 2 (B38T)
7 6
(46)
vdl ned++1
Max (E33T) + #²
have:
For all t > 0
T 13
n + EØ3T + (-ETT 3 - vdn2
- m) |
u- >
> 0.
(47)
2vd EE8 op
((Un(()) satisfies the Pl-inequality, i.e. there exist p 0 (Iree of () such that
(((Un((),=_110)) (() (((((())))()))) min (U11,
(48)
Then gradient flow converges to global minimum of the population loss. Moreover, WPV and WRQ converge to WPV and w.k respectively, where
PV* W22
1 E33 (
BBT B+
)
E33
(49)
n1
W
(w -1
n (~t
E33T +
EB B
0²
)
E
(50)
n
Proof. We first characterize the behavior of E(Uin(t) (t)( on the gradient flow. Define
r n +
1 -
E39
+
E3 B
I.
(51)
Note that
1
lr u - (u- r 1B33) ']
2
1
tr UnLTT
2u 2 -1 E33 1 E33 TT1E333))
2
1
tr - + E
2u
2
1
22 tr 10
2 tr
|-- (E
2
Hence,
E(U u-1) = う tr [T (u-1Un1 ( - tr
p-1 (IB33T)²
EE3 + 8..
Besides, by chain rule, we have
dē
dē
dUn
dē
du
dt
dU'1"
dt
du-1"
dt
dUn
dUn
du
du-
< 0.
dt
dt
dt
dt
Also, ? is a quadratic function w.r.t, u and (U1,, u_1) = E3 3 + 8² when u- 0. Hence, if we can
show that for all > 0, غ & 3 02, we can claim (t) > by the property of quadratic function. Now
under the initialization condition equation with (0OT (E,33 θ0 (E33T)) 1, at 0 we have
(Un1 (0)
43

---

Published in Transactions on Machine Learning Research (11/2025)
1
tr r (E33T) $ Toe (E38T) Tee (E33T) $ - tr (LBBT)) 2 ree EB 3 0²
2
27 ir ( (E33T) T0e (1888)) L0e - 7² Lr (EB3T) 2 T0e
1
+ 18,3 B 4.00 0²
1
1r 7² o (E381) $ /
2 7 S/vvar op E33 T 6 E39 00
+ EB 3 8+02
2 7 |ee (EB3T)) (va> op 00
P 2) + E,3 B 02
where the second equality follows from the fact that E commute with T. Therefore, as long |00 & 2, i.e. 7
|er then will have u-1(() 0. Now we just show that condition equation 12 is sufficient condition by finding a lower bound for Te8 TT
= tr (R,38 (eeeee0Tee TET3E TeT)
(00 E83 TOT) Amin (EB8T)
(00 3 2 (E33T)
|100| 1 A2 min (E
we have
(03 Hence, under condition > Amin (B33T). Besides, [Pp 2t Amax (E33T) + A
060
17
equation
12
we have
2Amin (I8,,TT
2Amin (E33TT
V a+d+1 Amax (E33T)+ +
22 n
Vd n+1 Amax (E,3TT) + 1,,8T
82
"
2
[ 100
and << TB 3 3+ 0² Next, prove equation
47 by contradiction. Note that can lower bound E(U1,, u-1)
((U11, tr
-
2
I' U u. UT
+ E3 8 0²
2
Lr [u_
E33 UH E, 3 02
du
E33
op |U11| + E 3 + 0²
where the first inequality follows from that and UnU are positive definite. Besides, at time t = 0, we have
1 - E33 E (2 (2- (2-v2r?1 |00T|r) + E3 3 0².
(0) <
Suppose that the opposite of equation 47 holds, i.e.
1 (2-vd,21[[),, 100T|
u-
2vd E33 lop
then we will have
du |IE33T op >
1 2
EB3 TE33 1| (2 vdr2111 [ ||0)Tp))
2
and thus,
?UU w-),,,)))2 =(( (0))
44

---

Published in Transactions on Machine Learning Research (11/2025)
However, also have dt/dt & 0. Therefore, we must have
1 E33 (2- vd>2111/ 1001 > 0.
u- V 2vd B3ST llop
Now we prove PL-inequality. Note that
E(U11, u-1) - min (U11, u-1) = tr T (u-1Un -1 (u-1U11 - - -
2
(u-1Un1 - -
2
Besides
dU'1
u² Unl --1E33"
dt
2
d0
+
aU11 F
>
|n_ Un r
2 |-1Un1 -033
Since E commutes with I, we have
u L-UUT-EE33
uuuUn - TEB33
uu -
-u-1Un-1-EE33 11-11---BB3 r.
Therefore, it holds that
T (u-UUn TB333))) T - EB33)r r-1
and
? (Un1, min J (U1. u_1
2 L -33 12 1-11
1
r 2
Now if we set a such that
2 Iu- T333 2 24
|+
JuiUn E33T 1 1
F
then PL-inoquality holds, i.c.
V u², | 1|
"" r2 - 1-11
> (Un, min (U11,
In particular, using lower bound of 1.,, we can set
T 13
(2 - vdn | 1I)) 0.
(52)
E33T op
Now, by the dynamics of gradient flow,
d d (() = - dUn11
|
du-1 2
$ (p u-)))
min (U11 --1)
di
dt
dt
45

---

Published in Transactions on Machine Learning Research (11/2025)
Hence,
0 (UU11()), min (Un, exp(-ut) [e (e(n u-1(0)) min (
i.e.
lim [e(Un(t), min (U11, 60
lim
r -
0.
1-00
Hence, we have 1-1011 - EB88 Besides, u², tr (UI1UT) = |1 1 we have as - 00
u- | E33
Ft
Un |333 IF - E33 T
C Auxiliary Results
Proposition C.1 (Proposition C.2 in Bai et al. 2024 Let :R. - R be loss function such that dil is (e, R. M,M)) -apprrzzzie by sum of relus with R max { BB Bw, By, 13. Let [.,(3) :- E {(3 wss Mi) denote the empirical risk with loss function on dataset {(1,, y1) se[m]" Then, for any > O, there exists an attention layer {(2ms Km: Vm)) mE[M] with M heads such that, for any input sequence that takes form h [ v/3 B: -2-00ppp---0 2d 3: 1; 1;;]] with [,31| Bud gives output
Ty. [Atnn( [xy; 1/i 3: On 20-3: 1:4]
for all i [N + 1], where
|3 (3- < <E (nB.).
Proposition C.2 (Proposition 1 in Pathak et al. (2024)) ). Given any input matrir e TRP4 that output a matria H' R2xx, following operators can be implemented by a single layer of an autoregressive transformer:
copy copyyddyown down H: k, W, L,I): For columns with index i e I, outputs H' where
Hi = Hk:.,
and the remaining entries are unchanged. Here, l' = + copied 'doun' within columns i E Note, we assume f. z k and that ke & a so that the operator is well-defindd.
(e- k) and W 2 k, so that entries are
ooey_ H, For columns with indez i I, outputs H' with
Hea = H
The remaining entries stay the same. Here entries from column
copied "over" to column i.
mul (H; K,, k" I.T): For columns with index i I, outputs H' where
H = HK4t,HH+H..
for (O...., k}.
Note that &" km where 88 W' e x8"x'' and & k + 8, = + + We assume 8,8'.85 88 2 The remaining entries of H are copied over H', unchanged.
scaled agg( a, KKK,K, Outputs a matrir H' with entries
Hw41 = a Σ Hatss for t (0, I, k}.
JEZ
The set I is causal, that I (1- The remaining entries of H are copied H', unchanged.
46

---

Published in Transactions on Machine Learning Research (11/2025)
test MSE SMI-1
MSE SNIt=10
0.06
hidden dimension
Figure 4: Plot of excess testing risk of the transformer v.s. the hidden dimension D with different SNRs.
soft (,) For the final column 9, outputs matric H' with entries
H
for tE e 10,1,..., kj.
Et-k "kee
The remaining entries of H are copied over H', unchanged.
Lemma C.1 (Lemma D.2 in Zhang et al. (2023]). If x is Gaussian random vector of dimension, mean zero and covariance matrix Ia, and A E Raxd is fixed matrix. Then
E xx Axx J (A AT) +tr(A)
IR Mxd are generated independently from X; N(0, Ia), then it holds If A is symmetric and the rows in X that
E X M. M Ia+ M(M + 1)A.
D Additional Details in Simulation
D.1 Batch EM Algorithm
The procedure of EM algorithm is given by Algorithm In our implementation we stop (or declare the
algorithm converged) if the maximum iteration is attain, or if maxy 1350 ((t- 1) 12 E, where 0.0001 and the maximum iteration Tmax 10000.
D.2 Performance with different number of hidden dimension
In this experiment, we vary the hidden dimension D 34, 64, 128. For each case, run the experiment with two components (K 2), different SNR (1 1,5, 10). The r-axis is the hidden dimension D, and the y-axis is the excess test MSE. The performance of the trained transformer is presented in Figure 4 In the low SNR settings, increasing the hidden dimension helps in improving the transformer's ability to learn the mixture of regression problem. However, excessively large hidden dimensions can lead to sparsity in the parameter matrix, which may not significantly enhance performance further.
47

---

Published in Transactions on Machine Learning Research (11/2025)
Algorithm Batch EM algorithm for Mixture of Regression Problem
Input: Number of prompts B > O, prompts (EO) : = I,..., B} of length n + 1 containing the data [(1,, (s), Zn+1 : i 1,...,n) generated based on Mixture of regression tasks with noise variance 0,
number of components K > O.
Initialize (1)) [0, drawn uniformly the probability simplex.
Initialize 3(0) Rd, drawn uniformly on the sphere radius for j E K] such that cos <(83, > 0.8.
Initialize (0)
% = for all i [B]. [K].
while have not converged do
Update prompt-component assignment probabilitics, for all [B],j e [K].
(vl 1 10T g(
IT- exp
20
(1+1) 76
(1) Σ T3y IT exp {
(0 -0(- (9T 36000
Update the marginal component probabilities for all j E [K]
B
11+1)
(1+1)
B
Update the parameter estimates for all j E [K]
g/(t1) arg min BER
(t 1) (.)) Tig
(
33 x
(1
y/
Update the iteration counter, t + end while
Return: Final set of component centers, (000)
1-=
48


Reviews:
[{"review_id": "V5z1ASO6sK", "reviewer": "Reviewer_ejbj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This work extended previous works into MoR tasks, and explicitly provided the excess risk bounds, which may be of independent interest.\n2. The writing of the work is generally easy to follow.\n\n\nWeakness:\n1. The experiment is a bit disconnected with the theory. This work proved the existence of a Transformer which implements gradient EM, but the experiment only reported excess MSE comparison between Transformer and EM, and it is not clear that whether the Transformer really learned the EM as authors claimed.\n2. The theory is a bit unintuitive to me. Theorem 3.2 indicates an abrupt change in the error, and there seems to be a \"magic number\" of SNR for a phase transition in the error growth, I am not sure whether such results is common in the community, also such pattern is not well-understood in theory and well-supported in your experiments based on my own understanding.\n3. Gradient flow analysis is nice, but all in all it should correspond to a continuous-time algorithm, it would be nice to have more discussion and insights on the discrete-time algorithm design and analysis, otherwise it is unclear how directly the continuous-time results translate to realistic architectures.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work studied ICL for MoR models. It proves the existence of a Transformer which internally implements gradient EM algorithm of MoR problems, and derives the prediction error under low and high SNR regimes; also it provided the excess risk of pretraining with finite number of ICL instances, and further analyzed the training dynamics of gradient flow for single linear self-attention layer. Numerical experiments are provided to demonstrate the effectiveness of the proposed theory.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. I think it would be more convincing if authors can add some experiment to show that the Transformer actually execute gradient-EM, for example, construct a real Transformer following your theory, and plot the errors of prediction of $\\beta$ versus low/high SNR as shown in your Theorem 3.1/3.2.\n2. Regarding Theorem 3.2, whether the change between low SNR and high SNR occurs abruptly, i.e., the error jumps from $O(d\\log^2(n/\\delta)/n)^{1/4})$ to $O(d\\log^2(n/\\delta)/n)^{1/2})$, or is the change gradual as the SNR varies? I may expect the error is a function of SNR in the whole regime so it should change incrementally, otherwise an abrupt change corresponds to a phase transition, I hope to have more explanations on it.\n3. Both Theorem 3.1 and 3.2 require the prompt length $n$ to be much larger than the input dimension $d$, and the model size $D$ is a multiple of $d$ (in experiments $D=64=2d$, also $n=64$), so possibly the theory may not work for the overparameterized case ($D$ is very large), can you comment on it?\n4. Also $n=64$ means the setting $n\\gg d$ may not be satisfied in your experiemnts, so I suggest authors consider adjusting the experimental setup or providing justification to better align with the theoretical assumptions.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZK7w1dsQZL", "reviewer": "Reviewer_3dAV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, this is a strong paper of interest to the community. However, I have some concerns that I would like to see addressed (see weaknesses and requested changes).\n\nStrengths:\n- The paper is split into a logical order, roughly following the order of Existence/representation -> statistical properties -> optimisation. Section 3 looks at the existence/representation abilities of the transformer architecture. Then, section 4 first looks at statistical properties of the estimator (section 4.1), and then shows that such estimators are realisable via gradient flow (section 4.2). The paper ends with some empirical investigations (section 5). \n- It is nice to see two different regimes quantified in theorem 3.3, which describes the excess risk of the MoR model, governed by the value of the SNR. \n- Theorem 3.1 and Theorem 3.2 describe the high probability bounds of the prediction error of the MoR model, with various parameters controlling the bound including dimension $d$ and prompt length $n$. \n\n\n\n\nWeaknesses:\n- Throughout the main paper, there is a lot of heavy lifting done by the phrase \"can implement the EM algorithm\". It is very difficult to understand precisely what the authors mean by this, without consulting Appendix A (and no helpful pointer to Appendix A is given in the main paper). I believe what it means is that the derivative of the function defined in the $E$ step can be approximated by a sum of ReLU functions (internal to the transformer), and this somehow translates to performing iterates of EM algorithm through the layers. I am having difficulty understanding why the EM algorithm in particular is of interest - wouldn't we expect such networks to be capable of representing lots of algorithms? \n    - Appendix A is quite challenging to read. I am having trouble seeing what the difference is between $H^{(T+1)}(3)$ and $H^{(T+1)}(4)$. The authors could direct the reader to what is changing between each line. \n- Equation (3). I am not sure if I am misunderstanding something, so I'm happy to be corrected by the authors here. It seems to me as though a special kind of input preprocessing/representation used, whereby the data is padded with zeros and the indicator function. This seems like it might be different to what people would do in practice, revealing a gap between the theory presented in this paper and the reality of transformers.\n- The first part (existence of a transformer which can implement EM algorithm) seems divorced from the second part (training dynamics under GD). My main concern is that the matrix $E$ is *assumed* in the second part, and this would either not happen in practice or would be inconsistent with the setup in the first part.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors investigate the use of transformer architectures for solving Mixture of linear Regression (MoR) problems. The analysis covers representation, statistical, and optimisation properties. First, they show that the transformer architecture is capable of representing the EM algorithm. Second, they establish that in-context learning estimators with low risk can be established by minimising an empirical risk. Finally, they show that gradient flow can realise such estimators. Hence, transformers for MoR are sufficiently rich, admit desirable statistical properties, and can be optimised (at least under idealised assumptions). They apply their model to some MoR problems, showing empirically faster convergence than EM algorithm, and investigating various hyperparameters such as number of examples and prompt length.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I am keen to upgrade my evaluation of \"Claims and Evidence\". Please add text that either addresses my confusions in the weaknesses above, or appropriately justifies any choices made. Additionally: \n\n- It would be helpful between equations (1) and (2) to state that $\\nu$ is known (I believe, anyway).\n- Remark 2.1. You could possibly further justify not using softmax, through citing various other works which use non-softmax functions. E.g. the literature on random Fourier features, or random features more generally. \n- Section 4.1 Is clipping standard in practice? If so, could you cite a paper or two?\n- I am confused by the terminology of \"pretraining\". Isn't this just standard training, before ICL happens? \n\nMinor typos:\n- \"and largest distancerespectively.\" between equations (1) and (2).\n- Definition 2.1. \"A attention layer\" should be \"An attention layer\".\n- The incorrect usage of the word \"Besides\" throughout is a bit jarring.\n- The grammar in theorem 3.1 is not quite right.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OzdVwHdidc", "reviewer": "Reviewer_nBsS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. Although linear regression tasks have been widely studied in understanding the dynamics of in-context learning, the perspective from the EM algorithm is novel.\n\n2. The authors provide some concrete constructions to factorize the M-step into several gradient ascent steps, which share the same spirit of the gradient descent dynamics in previous works.\n\n3. Presentation is good.\n\nWeaknesses:\n\n1. In-context learning requires the model to be able to generalize to a broader range outside the possible combinations of observed linear regression tasks. In that sense, the ability to accomplish MoR tasks is still different from the in-context learning ability. Understanding ICL via EM is questionable, especially when compared to the previous gradient descent perspective that can deal with any new linear regression task.\n\n2. The emergence of in-context learning ability in the mixture of linear regression task requires quite a number of different linear weights (e.g., see *Pretraining task diversity and the emergence of non-bayesian in-context learning for regression*). Under such conditions, the minimum gap between $\\beta_i$'s is quite small, leading to a very small SNR (defined in this paper). However, only after training on such a large corpus would the model be able to perform ICL. In that sense, the SNR arguments in this paper cannot fully convince me.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors examine the capability of Transformer structures in performing the mixture of (linear) regression (MoR) task. The authors provide a construction-based Transformer weight that performs the Expectation-Maximization (EM) algorithm by layer-wise gradient ascent for the M-step. This construction gives an approximation gap for the linear MoR tasks. On the other hand, the authors derive a generalization result for the pre-training of Transformers on MoR tasks. Following the simplified model (a one-layer linear-attention-only Transformer), the authors analyze the converging point of the gradient flow. Some numerical results are provided to support the efficiency of Transformers in MoR tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I will be grateful if the authors could explain more about the following two points:\n\n1. How do you separate the possibility of in-weight learning (IWL) from ICL with your EM explanations? I'm quite curious because it seems that the only possible $\\beta$'s in the test phase are those in the training phase.\n\n2. How do you view the SNR when the number of possible $\\beta_i$'s tends to infinity? See Weaknesses part 2.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)