Large language models (LLMs) based on the transformer architecture (Vaswani et al., 2017) have demonstrated remarkable in-context learning (ICL) abilities (Brown et al., 2020). When given a prompt consisting of examples of a learning task, these models can learn to solve this task for new test examples without any parameter updating. This behavior has been empirically demonstrated in state-of-the-art models on real-world tasks (OpenAI, 2023;Touvron et al., 2023). This impressive capacity of transformer-based models has inspired many recent works aiming to understand the ICL abilities of transformers. A more comprehensive literature review can be found in Appendix A. Garg et al. (2022) was the first to study the ICL abilities of transformers for various function classes. They empirically showed that transformers can learn linear regression models in context. Later on, a line of research was developed to theoretically explain how transformers perform in-context linear regression. For example, Akyürek et al. (2022); Von Oswald et al. (2023); Bai et al. (2024); Fu et al. (2023); Giannou et al. (2024) showed by construction that, some specially-designed transformers can perform linear regression in context. Moreover, some recent works such as Zhang et al. (2023a); Huang et al. (2023); Chen et al. (2024) studied the training dynamics of a singlelayer transformer for in-context linear regression. They proved the convergence of certain single-layer transformers during training and showed that the trained transformers are able to perform linear regression in context.
Building on the earlier works that largely focus on linear regression problems, several recent papers have started to investigate the ICL capabilities of transformers for classification problems. For instance, Bai et al. (2024) showed that, by construction, multi-layer transformers can be approximately viewed as multiple steps of gradient descents for logistic regression. Giannou et al. (2024) further showcased that the constructed transformers can approximately perform Newton's method for logistic regression. Lin & Lee (2024) studied the dual operating modes for in-context classification of Gaussian mixtures. However, their analyses were based on the idealized Bayes-optimal next-token predictor for ICL tasks and did not consider the training dynamics of transformer models. Some recent works have started to study the training dynamics of transformers for certain classification problems. For example, Li et al. (2024b) proved that a single-layer transformer can be trained to learn one-nearest neighbor in context. Li et al. (2024a) studied the training dynamics of a single-layer transformer for some binary classification tasks with finite, pairwise orthogonal patterns. Frei & Vardi (2024) analyzed a linear transformer model for in-context classification of Gaussian mixtures, assuming additional conditions such as a sufficiently large signal-to-noise ratio. However, all these works (Li et al., 2024a;b;Frei & Vardi, 2024) only considered the binary classification problems. The training dynamics of transformers for more general in-context classification problems beyond the specific settings and assumptions in Li et al. (2024a;b); Frei & Vardi (2024) remain largely underexplored.
In this work, we study the training dynamics of a singlelayer transformer for both binary and multi-class classification of Gaussian mixtures during ICL, a fundamental problem in machine learning. Our main contributions can be summarized as follows:
• We prove that with appropriately distributed training data (Assumptions 3.2, 4.2), a single-layer transformer trained via gradient descent will converge to its global minimizer at a linear rate (Theorems 3.3,4.3) for both in-context binary or multi-class classification problems. To the best of our knowledge, we are the first to prove the training convergence of transformers for in-context multi-class classification. Moreover, our analysis reveals that the trained single-layer transformer can be viewed as approximately implementing linear discriminant analysis (LDA). • Due to the non-linearity of our loss function, we cannot directly find the closed-form expression of the global minimizer. Instead, we prove an important property that the global minimizer consists of a constant plus an error term that is induced by the finite training prompt length (N ). We further show that the max norm of this error term is bounded, and converges to zero at a rate of O(1/N ). • With properly distributed testing prompts (Assumptions 3.5, 4.4), we establish an upper bound of the inference error (defined in Equation ( 3)) of the trained transformer and quantify the impact of the training and testing prompt lengths on this error. We further prove that when the lengths of training prompts (N ) and testing prompts (M ) approach infinity, this error converges to zero at a rate of O(1/N + 1/ √ M ) , and the prediction of the trained transformer has an identical distribution to that of the ground-truth label (Theorems 3.6, 4.5).
Notations. We denote [n] = {1, 2, . . . , n}. For a matrix A ∈ R m×n , we denote its Frobenius norm as ∥A∥ F , and its max norm as ∥A∥ max = max i∈[m],j∈ [n] |A ij |. We use A a,b (or A ab ) to represent the element of matrix A at the a-th row and b-th column, and use A a:c,b to represent a vector of dimension c -a + 1 whose i-th element is A (a+i-1),b . We denote the l 2 norm of a vector as ∥ • ∥ 2 . We denote the all-zero vector of size n as 0 n and the all-zero matrix of size m × n as 0 m×n . We use σ(x) := 1/(1 + exp(-x)) to denote the sigmoid function. We define softmax(•) : R k → (0, 1) k , and its i-th element as softmax(•) i , where softmax(x) i = exp(x i )/( k j=1 exp(x j )).
Given an input embedding matrix E ∈ R de×dn , a single head self-attention module F SA with width d e will output
where W V , W K , W Q ∈ R de×de are the value, key, and query weight matrices, respectively, ρ > 0 is a normalization factor, and f attn is an activation function for attention.
There are different choices of f attn ; for example Vaswani et al. (2017) adopts softmax.
In this work, similar to Zhang et al. (2023a); Wu et al. (2023), we set f attn (x) = x and define W KQ = (W K ) ¦ W Q ∈ R de×de . We use F to denote this simplified model. Then, the output of F with an input embedding matrix E ∈ R de×dn can be expressed as
In the following theoretical study and the subsequent experiments (Section 5.2), we show that this simplified transformer model has sufficient capability to approach the optimal classifier for the in-context classification of Gaussian mixtures.
We adopt a framework for in-context learning similar to that used in Bai et al. (2024). Under this framework, the model receives a prompt P = (D, x query ) comprising a set
∼ P and a query x query ∼ P x , where P is the joint distribution of (x, y) and P x is the marginal distribution of x. Here, x i ∈ R d is an in-context example, and y i is the corresponding label for x i . For instance, in regression tasks, y i ∈ R is a scalar. In this paper, we focus on classification tasks. Thus, the range of y i can be any set containing c different elements, such as {1, . . . , c}, for classification problems involving c classes. The objective is to generate an output y query that approximates the target y query ∼ P y|xquery .
Since y query is a discrete random variable, we use the total variation distance to measure the difference between y query and y query :
where R(y query ) is the range of y query .
When ∆(y query , y query ) = 0, y query has the same distribution as y query , which means the output of the model perfectly approximates y query .
Unlike standard supervised learning, each prompt P τ can be sampled from a different distribution P τ in ICL. We say that a model has the ICL capability if it can approximate y τ,query for a broad range of P τ 's with fixed parameters.
In this section, we study the learning dynamics of a singlelayer transformer for in-context binary classification. It is a special case of the general multi-class classification. As a result, the analysis is more concise. The general in-context multi-class classification problem is studied in Section 4.
We first introduce the prompt and the transformer structure we will use for in-context binary classification. The prompt for in-context binary classification is denoted as P = (x 1 , y 1 , . . . , x N , y N , x query ), where x i ∈ R d and y i ∈ {-1, 1}. We can convert this prompt P into its corresponding embedding matrix E(P ) in the following form:
Similar to Huang et al. (2023); Wu et al. (2023); Ahn et al. (2024), we set some of the parameters in our model to 0 or 1 to simplify the optimization problem, and consider the parameters of our model (W V , W KQ ) in the following sparse form:
where W ∈ R d×d . We set the normalization factor ρ equal to the length of the prompt N . Let F (E(P ); W ) be the output matrix of the transformer. We then read out the bottom-right entry of the output matrix through a sigmoid function, and denote this output as y out . The output y out of the transformer with prompt P and parameters W can be expressed as
We denote the prediction of our model for x query as y query , which is a random variable depending on y out . Consider generating a random variable u uniformly on [0, 1]. If u f y out , we output y query = 1; if u > y out , we output y query = -1. Then, we have P ( y query = 1) = y out , P ( y query = -1) = 1 -y out .
We study the binary classification of two Gaussian mixtures and use the following definition.
Definition 3.1. We say a data pair (x, y) ∼ P b (µ 0 , µ 1 , Λ) if y follows a Bernoulli distribution with P (y = -1) = P (y = 1) = 1/2 and f (x|y = -1) = N(µ 0 , Λ), f (x|y = 1) = N(µ 1 , Λ), where µ 0 , µ 1 ∈ R d and Λ ∈ R d×d is a positive definite matrix.
We consider the case of B training tasks indexed by τ ∈ [B].
Each training task τ is associated with a prompt P τ = (x τ,1 , y τ,1 , . . . , x τ,N , y τ,N , x τ,query ) and a corresponding label y τ,query . We make the following assumption in this section.
Assumption 3.2. For each learning task τ ∈ [B], we assume
(1) {x τ,i , y τ,i } N i=1 and {x τ,query , y τ,query } i.i.d.
, and U τ is uniformly distributed over the closed set of real unitary matrices such that
We denote the distribution of (µ τ,0 , µ τ,1 ) as P b Ω (Λ). Note that U τ,Λ = Λ 1/2 U τ Λ -1/2 can be viewed as a linear transformation that preserves the inner product of vectors in Λ -1weighted norm, and we have µ
) ) be the output of our transformer for task τ . We define the empirical risk over B independent tasks as
Taking the limit of infinite training tasks B → ∞, the expected training loss can be defined as
where the expectation is taken over
Applying gradient descent over the expected training loss (7), we have the following theorem.
Theorem 3.3. Under Assumption 3.2, the following statements hold.
(1) Optimizing the training loss L(W ) in Equation (7) with training prompt length N via gradient descent W t+1 = W t -η∇L(W t ), we have that for any t g 1,
where W 0 is the initial parameter and W * is the global minimizer of L(W ), and κ = l/α. Here α, l are constants satisfying
), and a = µ ¦ Λ -1 q for simplicity. Then we have
where S = 4∇ 2 L(2Λ -1 ), L(2Λ -1 ) = lim N →∞ L(2Λ -1 ), σ ′ (•) and σ ′′ (•) are the firstand second-order derivatives of σ(•), respectively, and the expectation is taken over
The detailed proof of Theorem 3.3 can be found in Appendix D. In the following, we provide a brief proof sketch to highlight the key ideas.
Proof sketch for Theorem 3.3. As a first step, we prove in Lemma D.2 that the expected loss function L(W ) in Equation ( 7) is strictly convex with respect to (w.r.t.) W and is strongly convex in any compact set of R d×d . Moreover, we prove L(W ) has one unique global minimizer W * . Since the loss function L(W ) we consider is highly non-linear, we cannot directly find the closed-form expression of W * , as is often done in the prior literature. This poses a significant challenge to our analysis.
We address this technical challenge via the following method. First, in Lemma D.3, by analyzing the Taylor expansion of L(W ), we prove that as N → ∞, our loss function L(W ) converges to L(W ) pointwisely (defined in Equation ( 25)), and the global minimizer W * converges to 2Λ -1 . Thus, we denote W * = 2(Λ -1 + G), and prove ∥G∥ max is bounded and scales as
Next, in Lemma D.4, by further analyzing the Taylor expansion of the equation ∇L(W * ) = 0 at the point 2Λ -1 , we establish a tighter bound ∥G∥ max = O(N -1 ). In Lemma D.5, we prove that our loss function is l-smooth and provide an upper bound for l. Thus, in a compact set R W , our loss function is α-strongly convex and l-smooth. Finally, leveraging the standard results from the convex optimization, we prove Theorem 3.3.
According to Theorem 3.3, we have
). Thus, we have the following corollary.
Corollary 3.4. If we optimize the expected loss L(W ) in Equation (7) via gradient descent with training prompt length N , initial parameters W 0 , and learning rate η = 1/l, then, under Assumption 3.2, after T g 2κ log(N ∥W 0 -W * ∥ F ) steps, the updated model W satisfies
where ∥ G∥ max = O(1/N ), κ = l/α, and α, l are constants defined in (9).
Theorem 3.3 and Corollary 3.4 show that training a singlelayer transformer with properly distributed data (Assumption 3.2) for binary classification via gradient descent can linearly converge to its global minimum W * = 2(Λ -1 +G). Furthermore, when the prompt length N grows, this global minimum W * will converge to 2Λ -1 at a rate of O(1/N ).
Next, we analyze the performance of the trained transformer in Equation ( 11) for in-context binary classification tasks. We make the following assumption.
Assumption 3.5. For an in-context test prompt P test = (x 1 , y 1 , . . . , x M , y M , x query ), we assume
With this assumption, for y query ∼ P b y|xquery (µ 0 , µ 1 , Λ), according to the Bayes' theorem, we have P (y query = 1|x query ) = f (x query |y query = 1)P (y query = 1)
If we test the trained transformer with parameters W in Equation ( 11) and P test , by a simple calculation, we have
Intuitively, when the training prompt length N → ∞, we have G → 0, and when the test prompt length M → ∞, we have 2 M M i=1 y i x ¦ i → (µ 1 -µ 0 ) ¦ . Thus, when N, M → ∞, P ( y query = 1) = y out → σ((µ 1 -µ 0 ) ¦ Λ -1 x query ) = P (y query = 1|x query ), and the prediction of the trained transformer y query perfectly matches with the distribution of the ground truth label y query .
By analyzing the Taylor expansion of y out at point σ((µ 1µ 0 ) ¦ Λ -1 x query ), we formally present the aforementioned intuition in the following theorem, which establishes an upper bound of the total variation distance between y query and y query . Theorem 3.6. Consider a test prompt P test satisfying Assumption 3.5, and let y query ∼ P b y|xquery (µ 0 , µ 1 , Λ). Let y query be the prediction of the trained transformer with parameters W in Equation (11). Then, for the inference error defined in Equation (3), we have
where µ = µ 1 -µ 0 , u = 2(µ 1 + µ 0 ), q = x query , and the expectation is taken over
The proof of Theorem 3.6 can be found in Appendix E. Since ∥ G∥ max = O(1/N ), Theorem 3.6 suggests that if we ignore the constants regarding µ 0 , µ 1 , Λ, x query , the expected total variation distance between y query and y query is at most O(1/N + 1/ √ M ). On the other hand, for data pair (x, y) ∼ P b (µ 0 , µ 1 , Λ), the distribution of y, P (y = 1|x) = σ((µ 1 -µ 0 ) ¦ Λ -1 x), can be characterized by a logistic regression model σ(w ¦ x + b) with parameters w = Λ -1 (µ 1 -µ 0 ) and b = 0. Therefore, when N, M → ∞, the prediction of the trained transformer is equivalent to the optimal logistic regressor for binary classification problems with distribution P b (µ 0 , µ 1 , Λ).
Note that different from Assumption 3.2 which states that µ τ,0 , µ τ,1 , x τ,query are sampled according to some specific distributions during training, Assumption 3.5 does not impose strong distributional constraints on µ 0 , µ 1 and x query , which shows the strong generalization ability of the trained transformer. Moreover, even if M → ∞, the distribution variation between y query and y query does not disappear unless N → ∞. Thus, the ICL ability of trained transformers for binary classification is limited by the finite length of training prompts. Similar behaviors have also been observed in Zhang et al. (2023a) for in-context linear regression. Remark 3.7. Theorem 3.6 requires Assumption 3.5 to hold. For example, we need the covariance matrix Λ in training and testing to be the same. A similar consistency requirement of Λ in training and testing had also been observed for in-context linear regression in Zhang et al. (2023a). Here, we discuss the consequences when Assumption 3.5 does not hold. For example, suppose the labels of our data in test prompts are not balanced where P (y = 1) = p 1 , P (y = -1) = p 0 . Besides, µ 0 , µ 1 do not have the same Λ -1 weighted norm, and the covariance matrix of test data satisfies Γ ̸ = Λ. Then, as N, M → ∞, we have
On the other hand, the distribution of the ground truth label is
x query . Then, we can see that unless ẑ = z or |σ(ẑ) -σ(z)| is sufficiently small, the transformer cannot correctly perform the in-context binary classification. Remark 3.8. Another important insight of our analysis is that the pre-trained single-layer transformer can be viewed as approximately implementing linear discriminant analysis (LDA). For example, suppose we are given {x i , y i } M i=1 and x query , and we need to predict the label y query for x query . LDA assumes that {x i , y i } M i=1 and {x query , y query } are i.i.d. samples, with P (y i = 1) = P (y i = -1), and the conditional probability density functions f (x i |y i = 1) and f (x i |y i = -1) are Gaussian with means µ 1 , µ -1 and same covariance Σ. Under these assumptions, it can be derived that the optimal decision criterion for x query is to predict
1 Σ -1 µ 1 ) > 0 and y query = -1, otherwise. LDA can estimate μ1 as the average of x i with y i = 1, estimate μ-1 as the average of x i with y i = -1, and estimate the covariance Σ from the within-class variances. For the singlelayer transformer, it can compute the in-context estimate μ1 -μ-1 = 1 M M i=1 y i x i , however, it is hard for the singlelayer transformer to estimate Σ in context. Thus, in our paper, we make the following assumptions (Assumptions 3.2 and 3.5). We assume the pre-train data and test data have the same covariance matrix Λ so that the transformer can learn an approximation of Λ during pre-training. Moreover, we assume the two class means µ 0 , µ 1 have the same Λ-weighted norm so that µ
Under these assumptions, the quadratic term cancels out, and the estimated decision criterion simplifies to
x query , which is very close to Equation (12) in our paper. When we use W to approximate 2Λ -1 , the estimated decision criterion becomes exactly Equation ( 12). Therefore, when W = 2Λ -1 and the in-context examples are balanced across classes, the transformer's decision criterion is the same as that of the LDA with exact knowledge of Λ. Our experiments also corroborate this theoretical findings. For example, in Figure 3, since the pre-trained transformer has already learned a relatively good approximation of Λ -1 , while LDA must estimate Λ -1 in context, the trained transformer significantly outperforms LDA when the number of in-context examples is small. As the context length increases, LDA's performance approaches that of the trained transformer.
We now extend the study of the learning dynamics of a single-layer transformer to in-context multi-class classification, generalizing the results of the previous section. We will present the detailed formulation and then focus on the main differences to binary classification.
We first introduce the prompt and the transformer structure that will be used for in-context multi-class classification. The prompt for in-context multi-class classification involving c g 2 classes can be expressed as
. . , e c }, and e i is the i-th standard unit vector of R c . Its embedding matrix can be formulated as
Similar to the binary case, we set some of the parameters in our model as 0 and 1 to simplify the optimization problem and consider the parameters of our model (W V , W KQ ) in the following sparse form:
where W ∈ R d×d . We set the normalization factor ρ equal to the length of the prompt N . We read out the bottom-right c-dimensional column vector from the output matrix with a softmax function as the output, denoted as y out . With parameters W and a prompt P = (x 1 , y 1 , . . . , x N , y N , x query ), the output can be expressed as
We denote the prediction of the model for x query as y query , which is a random variable depending on y out . Randomly sample a random variable u that is uniformly distributed on
, where ( y out ) j is the j-th element of y out , we let y query = e i . Thus, P ( y query = e i ) = ( y out ) i .
We focus on the multi-class classification of Gaussian mixtures and use the following definition.
Definition 4.1. We say a data pair (x, y) ∼ P m (µ, Λ) if
We consider the case of B training tasks indexed by τ ∈
Each training task τ is associated with a prompt P τ = (x τ,1 , y τ,1 , . . . , x τ,N , y τ,N , x τ,query ) and a corresponding label y τ,query . We make the following assumption in this section.
Assumption 4.2. For each learning task τ ∈ [B], we assume:
(1) {x τ,i , y τ,i } N i=1 and {x τ,query , y τ,query } i.i.d.
, and U τ,k are uniformly distributed over the closed set of real unitary matrices such that
We denote the distribution of µ τ as P m Ω (Λ). Note that U τ,k,Λ = Λ 1/2 U τ,k Λ -1/2 can be viewed as linear transformation that preserves the inner product of vectors in the Λ -1 weighted norm, and we have
Let y τ,out = softmax [F (E(P τ ); W )] (d+1):(d+c),(N +1) be the output of the transformer for task τ . We define the empirical risk over B independent tasks as
Taking the limit of infinite training tasks B → ∞, the expected training loss can be defined as
where the expectation is taken over
Applying gradient descent over the expected training loss in Equation ( 16), we have the following theorem.
Theorem 4.3. (Informal) Under Assumption 4.2, the following statements hold.
(1) Optimizing training loss L(W ) in Equation ( 16) with training prompt length N via gradient descent W t+1 = W t -η∇L(W t ), for any t g 1, we have 17) where W 0 is the initial parameter and W * is the global minimizer of L(W ), κ = l/α. Here, α, l are constants such that
where
The formal statement and complete proof of Theorem 4.3 can be found in Appendix F. Technically, the proof of Theorem 4.3 builds on that of Theorem 3.3, but the more complicated cross terms in the Taylor expansions of the softmax functions, which are due to the nature of multi-class classification, bring new challenges to the analysis. To address these issues, we derive new bounds on the expected errors of the cross terms in Lemma F.1, F.2, which may be of independent interest to other similar problems.
Theorem 4.3 shows that training a single-layer transformer with properly distributed data (Assumption 4.2) for incontext multi-class classification via gradient descent can linearly converge to its global minimum
When the prompt length N grows, this global minimum W * will converge to cΛ -1 at a rate of O(c/N ). Compared to the binary case, the new results establish the scaling behavior w.r.t. the number of classes c.
Assumption 4.4. For an in-context test prompt
With this assumption, for y query ∼ P m y|xquery (µ, Λ), according to the Bayes' theorem, we have If we test the trained transformer with parameters W in Equation ( 19) and prompt P test , by a simple calculation, we have
Note that, when the training prompt length N → ∞, we have G → 0, and when the test prompt length M → ∞,
, the prediction of the trained transformer y query matches the ground truth label y query .
By analyzing the Taylor expansion of y out at point softmax(µ ¦ Λ -1 x query ), we crystallize the aforementioned intuition in the following theorem, which establishes an upper bound of the total variation distance between y query and y query .
Theorem 4.5. (Informal) Let P test satisfy Assumption 4.4 and y query ∼ P m y|xquery (µ, Λ). Denote y query as the prediction of the trained transformer with parameter W in Equation (19). Then, for the inference error defined in Equation (3), we have
where the expectation is taken over {x i , y i } M i=1 i.i.d.
The formal statement and proof of Theorem 4.5 can be found in Appendix G. We can see that the convergence rate of the inference error in multi-class classification w.r.t. N and M is similar to that in the binary classification, except for the constant coefficient c. This suggests that classification tasks with more classes may have higher errors than those with fewer classes. On the other hand, for data pair (x, y) ∼ P m (µ, Λ), the distribution of y,
the prediction of the trained transformer is equivalent to the optimal softmax regressor for multi-class classification problems with distribution P m (µ, Λ). Note that different from Assumption 4.2 which states that µ τ , x τ,query are sampled according to some specific distributions during training, Assumption 4.4 does not impose strong distributional constraints on µ or x query , which shows the strong generalization ability of the trained transformer. We also discuss the consequences when Assumption 4.4 does not hold in Remark G.2, which highlights the necessity of Assumption 4.4. Moreover, even if M → ∞, the distribution variation between y query and y query does not disappear unless N → ∞. Thus, the ICL ability of the trained transformers for multi-class classification is limited by the finite length of training prompts. Similar behavior has also been observed in Zhang et al. (2023a) for in-context linear regression and in Section 3.2 for in-context binary classification.
In this section, we report the experiment results on multilayer, nonlinear transformers to investigate their similarities and differences to the single-layer, linear transformer we theoretically analyzed in the pervious sections. Detailed experimental settings and additional results can be found in Appendix H.
We train single-layer and multi-layer transformers for incontext classification of Gaussian mixtures with different numbers of Gaussian mixtures c, different lengths of training prompts N , and test them with different test prompt lengths M . The results are reported in Figure 1. We can see that for both single-layer and multi-layer transformers, the inference errors decrease as N and M increase, and they increase as c increases, which not only verifies our theoretical claims but also shows that, the simplified model we have studied indeed exhibits behavioral similarities to the more complex multi-layer, nonlinear transformers, and some of our observations for this simplified model also hold for more complex transformers.
Note that in Assumption 3.2, 4.2, 3.5, 4.4, we assume that the covariance Λ during pre-training and during inference are the same, and the means of all Gaussian components From Figure 2 (a), we can see that both models perform better when their µ τ,i have the same Λ -1 weighted norm ('same norm'), however, in the 'different norms' setting, the performance of '1-layer' deteriorates more significantly, while transformers with a more complex structure ('3-layer') show better robustness under this distribution shift. Similar situations also happen in Figure 2 (b), where '3-layer' also shows better tolerance to the covariance shifts than '1-layer'.
Experimental results in Figure 2 show the necessities of Assumptions 3.2, 4.2, 3.5, 4.4 for the single-layer transformers considered in this paper, and also demonstrate the better robustness of multi-layer, nonlinear transformers. Developing a better understanding of the robustness of more complex transformers is an intriguing direction for future research.
Additionally, we conduct experiments comparing the ICL performances of the transformers with other machine learning algorithms for the classification of three Gaussian mixtures. From Figure 3, we can see that all three transformer models significantly outperform the classical methods (softmax regression, linear discriminant analysis), demonstrating the strong ICL capacities of transformers.
We studied the learning dynamics of transformers for incontext classification of Gaussian mixtures, and showed that with properly distributed data, a single-layer transformer trained via gradient descent converges to its global minimum. Moreover, we established the upper bounds of the inference errors of the trained transformers and discussed how the training and test prompt lengths influence the performance of the model. Experimental results also corroborated the theoretical claims. There are some directions worth further exploring. One potential avenue is to investigate whether the assumptions regarding the training and test prompts can be relaxed. Additionally, we have only ex-amined single-layer transformers with linear attention and sparse parameters. The learning dynamics of multi-layer transformers with nonlinear attention (e.g., softmax) for incontext classification problems remain an interesting area for future investigation.
It has been observed that transformer-based models have impressive ICL abilities in natural language processing (Brown et al., 2020;Nye et al., 2021;Wei et al., 2022;Dasgupta et al., 2022;Zhang et al., 2022). Garg et al. (2022) first initiated the study of the ICL abilities of transformers in a mathematical framework and they empirically showed that transformers can in-context learn linear regression, two-layer ReLU networks, and decision trees. Subsequently, numerous works have been developed to explain the ICL capacities of transformers in solving in-context mathematical problems. These works mainly use two approaches: constructing specific transformers capable of performing certain in-context learning tasks, and studying the training dynamics of transformers for such tasks.
Constructions of transformers. Akyürek et al. (2022); Von Oswald et al. (2023) showed by construction that multi-layer transformers can be viewed as multiple steps of gradient descent for linear regression. Akyürek et al. (2022) also showed that constructed transformers can implement closed-form ridge regression. Guo et al. (2023) showed that constructed transformers can perform in-context learning with representations. Bai et al. (2024) proved that constructed transformers can perform various statistical machine learning algorithms through in-context gradient descent and showed that constructed transformers can perform in-context model selection. Lin et al. (2023) demonstrated that constructed transformers can approximate several in-context reinforcement learning algorithms. Fu et al. (2023); Giannou et al. (2024) further proved that constructed transformers can perform higher-order optimization algorithms like Newton's method. Pathak et al. (2023) showed that transformers can learn mixtures of linear regressions. Giannou et al. (2023) proved that looped transformers that can emulate various in-context learning algorithms. Cheng et al. (2023) showed that transformers can perform functional gradient descent for learning non-linear functions in context. Zhang et al. (2024) showed that a linear attention layer followed by a linear layer can learn and encode a mean signal vector for in-context linear regression.
Training dynamics of transformers. Mahankali et al. (2023); Ahn et al. (2024) proved that the global minimizer of the in-context learning loss of linear transformer can be equivalently viewed as one-step preconditioned gradient descent for linear regression. Zhang et al. (2023a) proved the convergence of gradient flow on a single-layer linear transformer and discussed how training and test prompt length will influence the prediction error of transformers for linear regression. Huang et al. (2023) proved the convergence of gradient descent on a single-layer transformer with softmax attention with certain orthogonality assumptions on the data features. Li et al. (2023c) showed that trained transformers can learn topic structure. Wu et al. (2023) analyzed the task complexity bound for pretraining single-layer linear transformers on in-context linear regression tasks. Tarzanagh et al. (2023) built the connections between single-layer transformers and support vector machines (SVMs). Nichani et al. (2024) showed that transformers trained via gradient descent can learn causal structure. Chen et al. (2024) proved the convergence of gradient flow on a multi-head softmax attention model for in-context multi-task linear regression. Kim & Suzuki (2024); Yang et al. (2024) proved that trained transformers can learn nonlinear features in context. Bu et al. (2024) studied the training convergence of transformer for tasks encoded with multiple cross-concept semantics. Li et al. (2024b) proved that a one-layer transformer can be trained to learn one-nearest neighbor for binary classification in context. Li et al. (2024a) studied the training dynamics of a single layer transformer for in-context classification problems. However, they only studied the binary classification tasks with finite patterns. They generated their data as x = µ j + κv k , where {µ j } M1 j=1 are in-domain-relevant patterns and {ν k } M2 k=1 are in-domain-irrelevant patterns, M 1 g M 2 and these patterns are all pairwise orthogonal. Thus, the possible distribution of their data is finite and highly limited. In contrast, our work explores the ICL capabilities of transformers for both binary and multi-class classification of Gaussian mixtures. Specifically, our data is drawn according to P b (µ 0 , µ 1 , Λ) or P m (µ, Λ), and the range and possible distributions of our data are infinite. Recently, (Frei & Vardi, 2024) also studied the the ICL of a linear transformer model for classifying Gaussian mixtures and showed that trained transformers can exhibit benign overfitting in-context. However, their analysis relies on additional assumptions, such as a sufficiently large signal-to-noise ratio. In contrast, our work does not rely on such assumptions. Additionally, compared to (Frei & Vardi, 2024), we consider a more general multi-class setting. We believe these distinctions highlight the independent contributions of our work. Some works also studied the ICL from other perspectives. To name a few, Xie et al. (2021) explained the ICL as implicit Bayesian inference; Wang et al. (2023) explained the LLMs as latent variable models; Zhang et al. (2023b) explained the ICL abilities of transformers as implicitly implementing a Bayesian model averaging algorithm; and Li et al. (2023b) studied the generalization and stability of the ICL abilities of transformers. Hahn & Goyal (2023) showed that ICL can arise through recombination of compositional structure found in linguistic data. Lin & Lee (2024) studied the dual operating modes for in-context classification of Gaussian mixtures. However, the analyses of both Hahn & Goyal (2023) and Lin & Lee (2024) are based on the idealized optimal next-token predictor for ICL tasks and did not discuss the training dynamics of transformer models.
Apart from studies focusing on the training convergence properties of transformers in in-context learning (ICL), many other works have investigated the convergence behavior of transformers when learning different tasks. For example, (Jelassi et al., 2022) and (Li et al., 2023a) theoretically analyzed how Vision Transformers learn and converge. (Wang et al., 2024) showed that transformers can provably learn sparse token selection, whereas fully-connected networks cannot. (Zhang et al., 2025) demonstrated that transformers can learn optimal variable selection in group-sparse classification. (Gao et al., 2024) studied the global convergence of large-scale transformers under gradient flow dynamics.
We denote X ∼ Bin(n, p) if a random variable X follows the binomial distribution with parameters n ∈ N and p ∈ [0, 1], which means
where x * = argmin x f (x).
Lemma C.2 ( (Bubeck, 2015)). Suppose f : R d → R is α-strongly convex and β-smooth for some 0 < α f β. Then, the gradient descent iterating w t+1 = w t -η∇f (w t ) with learning rate η = 1/β and initialization w 0 ∈ R d satisfies that for any t g 1,
where κ = β/α is the condition number of f , and w * = argmin w∈R d f (w) is the minimizer of f .
In this section, we present the proof of Theorem 3.3.
First, we prove in Lemma D.2 that the expected loss function L(W ) in ( 7) is strictly convex w.r.t. W and is strongly convex in a compact set of R d×d . Moreover, we prove L(W ) has one unique global minimizer W * . Then, in Lemma D.3, by analyzing the Taylor expansion of L(W ), we prove that as N → ∞, our loss function L(W ) point wisely converges to L(W ) (defined in (25)), and the global minimizer W * converge to 2Λ -1 . We denote W * = 2(Λ -1 + G), and prove ∥G∥ max = O(N -1/2 ). Next, in Lemma D.4, by further analyzing the Taylor expansion of the equation ∇L(W * ) = 0 at the point 2Λ -1 , we establish a tighter bound ∥G∥ max = O(N -1 ). In Lemma D.5, we prove that our loss function is l-smooth and provide an upper bound for l. Thus, in a compact set R W , our loss function is α-strongly convex and l-smooth. Finally, leveraging the standard results from convex optimization, we prove Theorem 3.3 in subsection D.4.
In this section, we use the following notations.
Recall the expected loss function ( 7) is
where
is the output of the transformer, and the label of the data follows the distribution
In this section, we introduce the following notations to analyze (7). We denote µ = µ τ,1 -µ τ,0 , µ 1 = µ τ,1 , µ 0 = µ τ,0 and q = x τ,query . Then with probability P (y τ,query = 1) = 1/2 we have q = µ 1 +v, and with probability P (y τ,query = 0) = 1/2 we have q = µ 0 + v, where v ∼ N(0, Λ). We define p = 2 N N i=1 y τ,i x τ,i . Since with probability P (y τ,i = 1) = 1/2 we have x τ,i = µ 1 + v i , and with probability P (y τ,i = 0) = 1/2 we have x τ,i = µ 0 + v i , where v i ∼ N(0, Λ), we known
Then, the expected loss function ( 7) can be expressed as
The gradient of the loss function ( 7) can be expressed as
Moreover, we define a function L(W ) as
In Lemma D.3, we show that as N → ∞, L(W ) will point wisely converge to L(W ).
Lemma D.1. Suppose N 1 ∼ Bin(N, 1/2). Defining h = N 1 /N -1/2, we have
We can compute the moment-generating function of h as follows:
Thus, we know the coefficients of t, t 2 , t 3 are 0, 1/(8N ), 0 respectively, and the coefficients of t n , n g 4 are O(1/N 2 ). We have
Moreover, according to the Jensen's inequality, we have
Lemma D.2. For the loss function L(W ) (7), we have
Additionally, L(W ) has one unique global minimizer on R d×d .
For L(W ) defined in (25), we also have ∇ 2 L(W ) { 0. For any compact set R W of R d×d , when W ∈ R W , we have ∇ 2 L(W ) { γI d for some γ > 0. Additionally, L(W ) has one unique global minimizer on R d×d .
Proof. We vectorize W as Vec(W ) ∈ R d 2 , where Vec(W ) i = W t1(i),t2(i) , t 1 (x) = +(x -1)/d, + 1, t 2 (x) = ((x -1) mod d) + 1. Then, we have
The Hessian matrix of the loss function ( 7) is
Considering z ∈ R d 2 such that z ̸ = 0, we have
where f pq (p, q) are the probability density function (PDF) function of p, q. Since for any p, q, σ(p ¦ W q/2)(1σ(p ¦ W q/2)) > 0, we have z ¦ ∇ 2 L(W )z g 0. Thus, ∇ 2 L(W ) °0 and L(W ) is convex.
Moreover, for any z ̸ = 0, we denote
Then, we define region Ω(a, b)
min
we have S(Ω, W ) > 0. Since with probability P (y τ,query = 1) = 1/2, q = µ 1 + v, with probability P (y τ,query = 0) = 1/2, q = µ 0 + v, where v ∼ N(0, Λ) and p = µ + hu + g, where g ∼ N(0, 4Λ/N ), v ∼ N(0, Λ), µ 0 ∼ N(0, I d ), the covariance matrices of p, q are positive definite and we have f pq (p, q) > 0 for all p, q ∈ R d . Moreover, Ω(a, b) are non-zero measures on R d×d . Thus, we have C(Ω) > 0. Then, for any z ̸ = 0, we have
Thus, we have ∇ 2 L(W ) { 0. L(W ) is strictly convex.
Moreover, for any compact set R W of R d×d , for any W ∈ R W , we have
Then, for any W ∈ R W , for any z ̸ = 0, we have
Thus, when W ∈ R W , where R W is a compact set, we have ∇ 2 L(W ) { C(Ω)S(Ω)I d /4 and the loss function L(W ) is γ-strongly convex, where γ = C(Ω)S(Ω)/4.
Because our loss function is strictly convex in R d×d , it has at most one global minimizer in R d×d . Next, we prove all level sets of our loss function are compact, i.e. V α = {W ∈ R d×d | L(W ) f α} is compact for all α. We prove it by contradiction. Suppose V α is not compact for some α. Since our loss function is continuous and convex, V α is an unbounded convex set. Since the dimension of
Thus, we have
We define
we have S(Ω 0 ) > 0. Since Ω 0 are non-zero measures for p, q, we have C(Ω 0 ) > 0. Then, we have
This contradicts the assumption L(W α + tW k ) f α. Thus, all level sets of the loss function L(W ) are compact, which means there exists a global minimizer for L(W ). Together with the fact that L(W ) is strictly convex, L(W ) has one unique global minimizer on R d×d .
Similarly, we can prove the same conclusions for L(W ).
Lemma D.3. Denoting the global minimizer of the loss function (7) as W * , we have W * = 2(Λ -1 + G), where
Proof. Let a = µ ¦ Λ -1 q, s = µ ¦ W q/2, r = (hu + g) ¦ W q/2. Performing the Taylor expansion on (7), we have
where ξ(s, r) are real numbers between s and s + r. According to Lemma D.1, we have E [r] = E (hu + g) ¦ W q/2 = 0. Thus, we have
Moreover, we have j,k,l∈[d] g i g k W kl q l W ij q j = (gg ¦ W q) ¦ W q and E[gg ¦ ] = 4Λ/N . C l is a constant independent of N and W . Thus, we have
This shows that L(W ) point wisely converges to L(W ).
According to Lemma D.2, L(W ) has one unique global minimizer. Consider the equation:
We can easily find that ∇ L(2Λ -1 ) = 0 and W = 2Λ -1 is the global minimizer of L(W ).
Here ρ W , C W are some positive finite constants. Then, we have
where
W is a constant independent of N and W . This shows that, for W ∈ R W , our loss function L(W ) uniformly converge to L(W ).
Denote W * as the global minimizer of the loss function L(W ) with prompt length N . Then, we show that, when N is sufficiently large, W * ∈ R W . We first denote ∂R W = {W | ∥W -2Λ -1 ∥ F = ρ W } and ∆ = min W ∈∂R W L(W ) -L(2Λ -1 ) > 0. Then, for N g 4C ′ l /∆, and for any W ∈ R W , we have
Then, we have
According to Lemma D.2, for W ∈ R W , we have ∇ 2 L(W ) { γI d , where γ is a positive constant independent of N . Thus, L(W ) is γ-strongly convex in R W . According to Lemma C.1, we have
Thus, when N → ∞, we have
Lemma D.4. The global minimizer of the loss function ( 7) is W * = 2(Λ -1 + G), where
Proof. According to Lemma D.2, the loss function L(W ) has a unique global minimizer W * . We have
The Taylor expansion of σ(a + b) at point a with an Lagrange form of remainder is
where ξ(a, b) are real numbers between a and a + b. Thus, our equation ( 27) become
Note that E[σ ′ (a)bpq ¦ ] = E µ,u,q σ ′ (a)E g,h bpq ¦ . For E g,h [bpq ¦ ], according to Lemma D.1 and g ∼ N(0, 4Λ/N ), we have
Then, we have
Similarly for E[σ ′′ (a)b 2 pq ¦ /2], we have
.
For each term in (i), it contains two G. Thus, their max norms are at most smaller than O(∥G∥ 2 max ). For each term in (ii), it contains one G and h 2 or contains one G and two g. According to E[h 2 ] = 1/(4N ) in Lemma D.1, the max norm of terms with one G and h 2 are smaller than O(∥G∥ max /N ). Defining ḡ = N 1/2 Λ -1/2 g/2, we have ḡ ∼ N(0, I d ) and g = 2N -1/2 Λ 1/2 ḡ. Thus, converting two g to ḡ, we have a coefficient of N -1 . Therefore, the max norms of terms with one G and two g are also smaller than O(∥G∥ max /N ). Therefore, for terms (i), (ii), we have
For term (iii), according to Lemma D.1 and g ∼ N(0, 4Λ/N ), we have
For E[σ ′′′ (ξ(a, b))b 3 pq ¦ /3!], we have
.
For terms in ( * ) containing two or three G, these terms' expected absolute values are at most smaller than O(∥G∥ 2 max ). For terms in ( * * ) containing one G, these terms must contain n 1 number of h and n 2 number of elements of g, where n 1 + n 2 = 2, 3, 4, n 1 , n 2 ∈ N. According to Lemma D.1, we know that for n 1 = 1, 2, 3, 4, E|h n1 | f O(N -n1/2 ). Defining ḡ = N 1/2 Λ -1/2 g/2, we have ḡ ∼ N(0, I d ) and g = 2N -1/2 Λ 1/2 ḡ. Converting g to ḡ, we have a coefficient of N -n2/2 . Thus, for terms in ( * * ), these terms' expected absolute values are at most smaller than O(∥G∥ max N -(n1+n2)/2 ) f O(∥G∥ max N -1 ). For terms in ( * * * ) without G, these terms must contain n 1 number of h and n 2 number of elements of g, we have n 1 + n 2 = 3, 4, n 1 , n 2 ∈ N. Similarly, these term's expected absolute values are at most smaller than O(N -(n1+n2)/2 ) f O(N -3/2 ). Therefore, we have
Moreover, we have
where
, where S ij = s t1(i),t2(i),t1(j),t2(j) = Eσ ′ (a)µ t1(i) q t2(i) µ t1(j) q t2(j) . Then ( 36) can be expressed as
Note that S = 4∇ 2 L(2Λ -1 ). According to Lemma D.2, S is positive definite. Thus, combining ( 28), ( 29), ( 30), ( 31), ( 32), ( 34), ( 35), (37), we have
Lemma D.5. The loss function (7) is l-smooth, where l f
Proof. The Hessian matrix of the loss function is
where (a) is due to the Cauchy-Schwarz inequality. Thus, ∇ 2 L(W ) ¯lI d and L(W ) is l-smooth, where l is a constant smaller than 1 D.4. Proof of Theorem 3.3
Proof. According to Lemma D.4, the global minimizer of L(W ) is W * = 2(Λ -1 + G), where
Here α > 0 is a positive constant number. Thus, L(W ) is α-strongly convex in R W . Moreover, according to Lemma D.5, L(W ) is l-smooth. Then according to Lemma C.2, applying gradient descent with η = 1/l, for any t g 1, we have
where κ = l/α.
In this section, we use the following notations. We denote
E.2. Proof of Theorem 3.6
Proof. The output of the trained transformer is
The probability of y query = 1 given x query is
Defining a = µ ¦ Λ -1 q, b = (µ + hu + g) ¦ Gq + (hu + g) ¦ Λ -1 q, we have
where ξ are real numbers between a and a + b. Thus, we have
We first consider the term σ ′ (a)E [|b|]. Defining ḡ = Λ -1/2 M 1/2 g/2, we have
where (a) is due to
For E[b 2 ], we have
Notice that terms in E [(µ + hu + g) ¦ Gq] 2 contain two G. Thus, they are at most smaller than
Terms in E [(hu + g) ¦ Λ -1 q] 2 /2 contain two h, or two g, or one h and one g. According to Lemma D.1, we have
Converting one g to ḡ, we have a coefficient of M -1/2 . Thus, terms in E [(hu + g) ¦ Λ -1 q] 2 /2 contain two h, or two g, or one h and one g are O(1/M ). Terms in E (µ+hu+g) ¦ Gq(hu+g) ¦ Λ -1 q contain at least one G and one h or one G and one g. Thus, they are at most smaller than
In this section, we present the proof of Theorem 4.3.
First, we prove in Lemma F.3 that the expected loss function L(W ) (16) is strictly convex w.r.t. W and is strongly convex in a compact set of R d×d . Moreover, we prove L(W ) has one unique global minimizer W * . Then, in Lemma F.4, by analyzing the Taylor expansion of L(W ), we prove that as N → ∞, our loss function L(W ) point wisely converges to L(W ) (defined in (44)), and the global minimizer W * converge to 2Λ -1 . Thus, we denote W * = 2(Λ -1 + G), and prove ∥G∥ max = O(N -1/4 ). Next, in Lemma F.5, by further analyzing the Taylor expansion of the equation ∇L(W * ) = 0 at the point 2Λ -1 , we establish a tighter bound ∥G∥ max = O(cN -1 ). In Lemma F.6, we prove that our loss function is l-smooth and provide an upper bound for l. Thus, in a compact set R W , our loss function is α-strongly convex and l-smooth. Finally, leveraging the standard results from the convex optimization, we prove Theorem 4.3 in subsection F.3.
In this section, we use the following notations.
Proof. Since (N 1 , N 2 , . . . , N c ) ∼ Multin(N, 1/c), the moment-generating function of (N 1 , N 2 , . . . , N c ) is
We can compute the moment-generating function of h = (h 1 , h 2 , . . . , h c ) as follows:
Observing the coefficients of h, we have
Iteratively applying the Hölder's inequality, we have
Moreover, we have
For any
Thus, with the results from Lemma F.1, for any n 1k , n 2ki satisfying
Moreover, according to the Jensen's inequality, we have
Thus, with the results from Lemma F.1, for any
Lemma F.3. For the loss function L(W ) (16), we have ∇ 2 L(W ) { 0. For any compact set R W , when W ∈ R W , we have ∇ 2 L(W ) { γI d for some γ > 0. Additionally, L(W ) has one unique global minimizer on R d×d .
For L(W ) defined in ( 44), we also have ∇ 2 L(W ) { 0. For any compact set R W , when W ∈ R W , we have ∇ 2 L(W ) { γI d for some γ > 0. Additionally, L(W ) has one unique global minimizer on R d×d .
Proof. We vectorize W as Vec(W ) ∈ R d 2 , where Vec(W
where a k = -log( l=1,...,c,l̸ =k exp ((p l -p k )W q/c)). For ∇a k , we have
and (∇softmax(P ¦ W q/c) k ) j = l=1,...,c,l̸ =k softmax(P ¦ W q/c) k softmax(P ¦ W q/c) l (p k -p l ) t1(j) q t2(j) /c. We can express the Hessian matrix of the loss function with the following form:
Since for any P, q, k, l, softmax(P ¦ W q/c) k softmax(P ¦ W q/c) l > 0, we have z ¦ ∇ 2 L(W )z g 0. Thus, ∇ 2 L(W ) °0 and L(W ) is convex.
Defining p = p 1 -p 2 , we have
z a pt1(a) q t2(a)
2 f P q (P, q)dP dq where f P q (P, q) are the PDF function of P, q. For any z ̸ = 0, we denote
we have S(Ω, W ) > 0. Since we have f P q (P, q) > 0 for all P, q and Ω(a, b) are non-zero measures for P, q. Thus, we have C(Ω) > 0. Then, for any z ̸ = 0, we have
Moreover, for any compact set R W of R d×d , for any W ∈ R W , we have
Then, for any W ∈ R W , for any z ̸ = 0, we have
. Performing the Taylor expansion on (16), we have
where the last inequality is due to Lemma F.1, F.2. C l is a constant independent of N and W . This shows that L(W ) point wisely converge to L(W ).
According to Lemma D.2, L(W ) has one unique global minimizer. Considering the equation:
We can easily find that ∇ L(cΛ -1 ) = 0 and W = cΛ -1 is the global minimizer of L(W ).
Here ρ W , C W are some positive finite constants. Then, we have
Denote W * as the global minimizer of L(W ) with prompt length N . Then, we show that, when N is sufficiently large,
l /∆) 2 , and for any W ∈ R W , we have
On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Then, we have
According to Lemma D.2, for W ∈ R W , we have ∇ 2 L(W ) { γI d , where γ is a positive constant independent of N . Thus, L(W ) is γ-strongly convex in R W . According to Lemma C.1, we have
Thus, when N → ∞, we have
Lemma F.5. The global minimizer of the loss function ( 16) is W * = c(Λ -1 + G). We have
where a = µ ¦ Λ -1 q, a k = µ ¦ k Λ -1 q, S = c 2 ∇ 2 L(cΛ -1 ). Ignoring constants other than c, N , we have ∥G∥ max f O(c/N ).
Proof. According to Lemma F.3, the loss function L(W ) has a unique global minimizer W * . We have ∂ζ k (a) ∂ 2 ζ k (a) ∂a l ∂a n φ1∈{µ l ,ch l µ l ,g l },φ2∈{µn,chnµn,gn} φ ¦ 1 Gqφ ¦ 2 Gqp k q ¦ (i) + φ1∈{µ l ,ch l µ l ,g l },φ2∈{chnµn,gn}
.
For terms (i) having two G, their max norms are at most smaller than O(∥G∥ 2 max ). For terms (ii) having one G, define ḡl = Λ -1/2 g l , these terms must contain n 1j number of h j and n 2ji number of (ḡ j ) i , we have j∈[c],i∈ [d] n 1j + n 2ji = n t , n t = 1, 2, 3. According to Lemma F.2, we know that for n t = 1, 2, 3,
Thus, the max norm of expectations of terms in (ii) are at most smaller than O(∥G∥ max N -1 ). Therefore, for terms (i), (ii), we have
For terms (iii) without G, we have
where the last inequity is due to Lemma F.1, F.2.
For the third term
φ ¦ 1 Gqφ ¦ 2 Gqφ ¦ 3 Λ -1 q(p k4 ) l q m ( * )
φ ¦ 1 Gqφ ¦ 2 Λ -1 qφ ¦ 3 Λ -1 q(p k4 ) l q m ( * * )
.
For terms in ( * ) having two or three G, these terms' expected absolute values are at most smaller than O(∥G∥ 2 max ). For terms in ( * * ) having one G, these terms must contain n 1j number of h j and n 2ji number of (ḡ j ) i , we have j∈[c],i∈ [d] n 1j +n 2ji = n t , n t = 2, 3, 4. According to Lemma F.2, for n t = 2, 3, 4, we have
Thus, these term's expected absolute values are at most smaller than O(∥G∥ max N -1 ). For terms in ( * * * ) without G, these terms must contain n 1j number of h j and n 2ji number of (ḡ j ) i , we have j∈[c],i∈ [d] n 1j + n 2ji = n t , n t = 3, 4.
According to Lemma F.2, for n t = 3, 4, we have
|h n 1k k (ḡ j )
where s ijnm = E c k=2 k-1 l=1 ζ k (a)ζ l (a)(µ k -µ l ) i (µ k -µ l ) n q m q j . We vectorize G as Vec(G) i = G t1(i),t2(i) . Define S ∈ R d 2 ×d 2 , where S ij = s t1(i),t2(i),t1(j),t2(j) = E c k=2 k-1 l=1 ζ k (a)ζ l (a)(µ k -µ l ) t1(i) q t2(i) (µ k -µ l ) t1(j) q t2(j) , (54) can be expressed as
Note that S = c 2 ∇ 2 L(cΛ -1 ). According to Lemma F.3, S is positive definite. Thus, combining (47), ( 48), ( 49), ( 50), ( 51), ( 52), ( 53), (55), we have
Ignoring constants other than c, N , we have ∥G∥ max f O(c/N ).
Lemma F.6. The loss function (7) is l-smooth, where l f 1
i∈[d 2 ] E[(p t1(i) q t2(i) ) 2 ].