On the Training Convergence of Transformers for In-Context Classification

## Abstract

Abstract While transformers have demonstrated impressive capacities for in-context learning (ICL) in practice, theoretical understanding of the underlying mechanism enabling transformers to perform ICL is still in its infant stage. This work aims to theoretically study the training dynamics of transformers for in-context classification tasks. We demonstrate that, for in-context classification of Gaussian mixtures under certain assumptions, a single-layer transformer trained via gradient descent converges to a globally optimal model at a linear rate. We further quantify the impact of the training and testing prompt lengths on the ICL inference error of the trained transformer. We show that when the lengths of training and testing prompts are sufficiently large, the prediction of the trained transformer approaches the Bayes-optimal classifier.
Experimental results corroborate the theoretical findings.

## 1 Introduction

Large language models (LLMs) based on the transformer architecture (Vaswani et al., ) have demonstrated remarkable in-context learning (ICL) abilities (Brown et al., ). 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, ; Touvron et al., ).

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. () 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. (); Von Oswald et al. (); Bai et al. (); Fu et al. (); Giannou et al. () showed *by construction* that, some specially-designed transformers can perform linear regression in context.
Moreover, some recent works like Zhang et al. (); Huang et al. (); Chen et al. () studied the training dynamics of a single-layer transformer for in-context linear regression. They proved the convergence of their single-layer transformers and showed their trained transformer 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 non-linear problems such as classification. For instance, Bai et al. () showed that, by construction, multi-layer transformers can be approximately viewed as multiple steps of gradient descents for logistic regression. Giannou et al. () further showcased that the constructed transformers can approximately perform Newton’s method for logistic regression.
Recently, Li et al. () studied the training dynamics of transformers for in-context binary classification. However, their analysis requires the data to be pairwise orthogonal and the possible distribution of their data is highly limited. The learning dynamics of transformers for more general in-context classification problems is not well understood.
Moreover, to the best of our knowledge, existing literature (Bai et al., ; Giannou et al., ; Li et al., ) studying the in-context classification of transformers focus only on binary classification. How transformers perform in-context multi-class classification remains unexplored.

In this work, we study the learning dynamics of a singly-layer transformer for both in-context binary and multi-class classification of Gaussian mixtures, a fundamental problem in machine learning. Our main contributions can be summarized as follows:

- •
To the best of our knowledge, we are the first to study the learning dynamics of transformers for in-context classification of Gaussian mixtures, and we are the first to prove the training convergence of transformers for in-context multi-class classification.
We prove that with appropriately distributed training data (Assumptions 3.1, 4.1), a single-layer transformer trained via gradient descent will converge to its global minimizer at a linear rate (Theorems 3.1, 4.1) for both in-context binary or multi-class classification problems.
- •
Due to the high 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 test prompts (Assumptions 3.2, 4.2), 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/\sqrt{M})$ (Theorems 3.2, 4.2), and the prediction of the trained transformer is Bayes-optimal, i.e., the optimal classifier given the data distribution.

## 2 Preliminaries

Notations.
We denote $[n]=\{1,2,\dots,n\}$. For a matrix $A\in\mathbb{R}^{m\times n}$, we denote its Frobenius norm as $\|A\|_{F}$, and its max norm as $\|A\|_{\max}=\max_{i\in[m],j\in[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 $\|\cdot\|_{2}$. We denote the all-zero vector of size $n$ as $0_{n}$ and the all-zero matrix of size $m\times n$ as $0_{m\times n}$. We use $\sigma(x):=1/(1+\exp(-x))$ to denote the sigmoid function. We define $\mathrm{softmax}(\cdot):\mathbb{R}^{k}\to(0,1)^{k}$, and its $i$-th element as $\mathrm{softmax}(\cdot)_{i}$, where $\mathrm{softmax}(x)_{i}=\exp(x_{i})/(\sum_{j=1}^{k}\exp(x_{j}))$.

### 2.1 Single-layer transformer

Given an input embedding matrix $E\in\mathbb{R}^{d_{e}\times d_{n}}$, a single head self-attention module $F_{SA}$ (Vaswani et al., ) with width $d_{e}$ will output

$$ $\displaystyle F_{SA}(E;W^{V},W^{K},W^{Q})=E+W^{V}E\cdot f_{\text{attn}}\left( \frac{(W^{K}E)^{\top}W^{Q}E}{\rho}\right),$ (1) $$

where $W^{V},W^{K},W^{Q}\in\mathbb{R}^{d_{e}\times d_{e}}$ are the value, key, and query weight matrices, respectively, $\rho>0$ is a normalization factor, and $f_{\text{attn}}$ is an activation function for attention. There are different choices of $f_{\text{attn}}$; for example Vaswani et al. () adopts $\mathrm{softmax}$.

In this work, similar to Zhang et al. (); Wu et al. (), we set $f_{\text{attn}}(x)=x$ and define $W^{KQ}=(W^{K})^{\top}W^{Q}\in\mathbb{R}^{d_{e}\times d_{e}}$. We use $F$ to denote this simplified model. Then, the output of $F$ with an input embedding matrix $E\in\mathbb{R}^{d_{e}\times d_{n}}$ can be expressed as

$$ $F(E;W^{V},W^{KQ})=E+W^{V}E\cdot\frac{E^{\top}W^{KQ}E}{\rho}.$ (2) $$

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 Bayes-optimal classifier for in-context classification of Gaussian mixtures.

### 2.2 In-context learning framework

We adopt a framework for in-context learning similar to that used in Bai et al. (). Under this framework, the model receives a prompt $P=(\mathcal{D},x_{\mathsf{query}})$ comprising a set of demonstrations $\mathcal{D}=\{(x_{i},y_{i})\}_{i\in[N]}\stackrel{{\scriptstyle\mathrm{i.i.d.}}
}{{\sim}}\mathcal{P}$ and a query $x_{\mathsf{query}}\sim\mathcal{P}_{x}$, where $\mathcal{P}$ is the joint distribution of $(x,y)$ and $\mathcal{P}_{x}$ is the marginal distribution of $x$. Here, $x_{i}\in\mathbb{R}^{d}$ is an in-context example, and $y_{i}$ is the corresponding label for $x_{i}$. For instance, in regression tasks, $y_{i}\in\mathbb{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,\dots,c\}$, for classification problems involving $c$ classes. The objective is to generate an output $\widehat{y}_{\mathsf{query}}$ that approximates the target $y_{\mathsf{query}}\sim\mathcal{P}_{y|x_{\mathsf{query}}}$.

Since $y_{\mathsf{query}}$ is a discrete random variable, we use the total variation distance to measure the difference between $\widehat{y}_{\mathsf{query}}$ and $y_{\mathsf{query}}$:

$$ $\displaystyle\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}})=\sup_{z \in R(y_{\mathsf{query}})}|\mathbb{P}\left(y_{\mathsf{query}}=z\right)-\mathbb {P}\left(\widehat{y}_{\mathsf{query}}=z\right)|,$ (3) $$

where $R(y_{\mathsf{query}})$ is the range of $y_{\mathsf{query}}$. When $\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}})=0$, $\widehat{y}_{\mathsf{query}}$ has the same distribution as $y_{\mathsf{query}}$, which means the output of the model perfectly approximates $y_{\mathsf{query}}$.

Unlike standard supervised learning, in ICL, each prompt $P_{\tau}$ can be sampled from a different distribution $\mathcal{P}_{\tau}$. We say that a model has the *ICL capability* if it can approximate $y_{\tau,\mathsf{query}}$ for a broad range of $\mathcal{P}_{\tau}$’s with fixed parameters.

## 3 In-context binary classification

In this section, we study the learning dynamics of a single-layer 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},\dots,x_{N},y_{N},x_{\mathsf{query}})}$, where $x_{i}\in\mathbb{R}^{d}$ and $y_{i}\in\{-1,1\}$.
We can convert this prompt $P$ into its corresponding embedding matrix $E(P)$ in the following form:

$$ $E=E(P)=\begin{pmatrix}x_{1}&x_{2}&\cdots&x_{N}&x_{\mathsf{query}}\\ y_{1}&y_{2}&\cdots&y_{N}&0\end{pmatrix}\in\mathbb{R}^{(d+1)\times(N+1)}.$ (4) $$

Similar to Huang et al. (); Wu et al. (); Ahn et al. (), we set some of the parameters in our model to 0 0 or $1$ to simplify the optimization problem, and consider the parameters of our model $(W^{V},W^{KQ})$ in the following sparse form:

$$ $\displaystyle W^{V}=\begin{pmatrix}0_{d\times d}&0_{d}\\ 0_{d}^{\top}&1\end{pmatrix},\qquad W^{KQ}$ $\displaystyle=\begin{pmatrix}W&0_{d}\\ 0_{d}^{\top}&0\end{pmatrix},$ (5) $$

where $W\in\mathbb{R}^{d\times d}$. We set the normalization factor $\rho$ 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 $\widehat{y}_{\mathsf{out}}$. The output $\widehat{y}_{\mathsf{out}}$ of the transformer with prompt $P$ and parameters $W$ can be expressed as

$$ $\displaystyle\widehat{y}_{\mathsf{out}}$ $\displaystyle=\sigma\left([F(E(P);W)]_{(d+1),(N+1)}\right)$ $\displaystyle=\sigma\Bigg{(}\begin{pmatrix}0_{d}^{\top}&1\end{pmatrix}\begin{ pmatrix}\frac{1}{N}\sum_{i=1}^{N}x_{i}x_{i}^{\top}+\frac{1}{N}x_{\mathsf{query }}x_{\mathsf{query}}^{\top}&\frac{1}{N}\sum_{i=1}^{N}x_{i}y_{i}\\ \frac{1}{N}\sum_{i=1}^{N}x_{i}^{\top}y_{i}&\frac{1}{N}\sum_{i=1}^{N}y_{i}^{2} \end{pmatrix}\begin{pmatrix}W&0_{d}\\ 0_{d}^{\top}&0\end{pmatrix}\begin{pmatrix}x_{\mathsf{query}}\\ 0\end{pmatrix}\Bigg{)}$ $\displaystyle=\sigma\left(\left(\frac{1}{N}\sum_{i=1}^{N}y_{i}x_{i}^{\top} \right)Wx_{\mathsf{query}}\right).$ $$

We denote the prediction of our model for $x_{\mathsf{query}}$ as $\widehat{y}_{\mathsf{query}}$, which is a random variable depending on $\widehat{y}_{\mathsf{out}}$.
Consider generating a random variable $u$ uniformly on $[0,1]$. If $u\leq\widehat{y}_{\mathsf{out}}$, we output $\widehat{y}_{\mathsf{query}}=1$; if $u>\widehat{y}_{\mathsf{out}}$, we output $\widehat{y}_{\mathsf{query}}=-1$. Then, we have $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=1\right)=\widehat{y}_{\mathsf{out}}$, $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=-1\right)=1-\widehat{y}_{\mathsf{
out}}$.

### 3.1 Training procedure

We study the binary classification of two Gaussian mixtures and use the following definition.

###### Definition 3.1

We say a data pair $(x,y)\sim\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$ if $y$ follows a Bernoulli distribution with $\mathbb{P}\left(y=-1\right)=\mathbb{P}\left(y=1\right)=1/2$ and $f(x|y=-1)=\mathsf{N}(\mu_{0},\Lambda)$, $f(x|y=1)=\mathsf{N}(\mu_{1},\Lambda)$, where $\mu_{0},\mu_{1}\in\mathbb{R}^{d}$ and $\Lambda\in\mathbb{R}^{d\times d}$ is a positive definite matrix.

We consider the case of $B$ training tasks indexed by $\tau\in[B]$. Each training task $\tau$ is associated with a prompt $P_{\tau}=(x_{\tau,1},y_{\tau,1},\dots,x_{\tau,N},y_{\tau,N},x_{\tau,\mathsf{
query}})$ and a corresponding label $y_{\tau,\mathsf{query}}$.
We make the following assumption in this section.

###### Assumption 3.1

For each learning task $\tau\in[B]$, we assume

- (1)
$\{x_{\tau,i},y_{\tau,i}\}_{i=1}^{N},\{x_{\tau,\mathsf{query}},y_{\tau,\mathsf{
query}}\}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}^{b}(\mu_{
\tau,0},\mu_{\tau,1},\Lambda).$
- (2)
$\mu_{\tau,0}$ is randomly sampled from $\mathsf{N}(0,I_{d})$, and $\mu_{\tau,1}=U_{\tau,\Lambda}\mu_{\tau,0}$ where $U_{\tau,\Lambda}=\Lambda^{1/2}U_{\tau}\Lambda^{-1/2}$, and $U_{\tau}$ is uniformly distributed over the closed set of real unitary matrices
such that $U_{\tau}U_{\tau}^{\top}=I_{d}$.

We denote the distribution of $(\mu_{\tau,0},\mu_{\tau,1})$ as $\mathcal{P}^{b}_{\Omega}(\Lambda)$.
Note that $U_{\tau,\Lambda}=\Lambda^{1/2}U_{\tau}\Lambda^{-1/2}$ can be viewed as a linear transformation that preserves the inner product of vectors in $\Lambda^{-1}$-weighted norm, and we have $\mu_{\tau,0}^{\top}\Lambda^{-1}\mu_{\tau,0}-\mu_{\tau,1}^{\top}\Lambda^{-1}\mu
_{\tau,1}=0$.

Let $\widehat{y}_{\tau,\mathsf{out}}=\sigma([F(E(P_{\tau});W)]_{(d+1),(N+1)})$ be the output of our transformer for task $\tau$. We define the empirical risk over $B$ independent tasks as

$$ $\widehat{L}(W)=\frac{1}{2B}\sum_{\tau=1}^{B}-(1+y_{\tau,\mathsf{query}})\log( \widehat{y}_{\tau,\mathsf{out}})-(1-y_{\tau,\mathsf{query}})\log(1-\widehat{y} _{\tau,\mathsf{out}}).$ (6) $$

Taking the limit of infinite training tasks $B\to\infty$, the expected training loss can be defined as

$$ $L(W)=\lim_{B\to\infty}\widehat{L}(W)=-\frac{1}{2}\mathbb{E}\left[(1+y_{\tau, \mathsf{query}})\log(\widehat{y}_{\tau,\mathsf{out}})+(1-y_{\tau,\mathsf{query }})\log(1-\widehat{y}_{\tau,\mathsf{out}})\right],$ (7) $$

where the expectation is taken over $(\mu_{\tau,0},\mu_{\tau,1})\sim\mathcal{P}^{b}_{\Omega}(\Lambda)$, $\{x_{\tau,i},y_{\tau,i}\}_{i=1}^{N},\{x_{\tau,\mathsf{query}},y_{\tau,\mathsf{
query}}\}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}^{b}(\mu_{
\tau,0},\mu_{\tau,1},\Lambda)$.

Applying gradient descent over the expected training loss (7), we have the following theorem.

###### Theorem 3.1

Under Assumption 3.1, the following statements hold.

- (1)
Optimizing the training loss $L(W)$ in (7) with training prompt length $N$ via gradient descent $W^{t+1}=W^{t}-\eta\nabla L(W^{t})$, we have that for any $t\geq 1$
$\displaystyle\|W^{t}-W^{*}\|_{F}^{2}\leq\exp(-t/\kappa)\|W^{0}-W^{*}\|_{F}^{2},$
(8)
where $W^{0}$ is the initial parameter and $W^{*}$ is the global minimizer of $L(W)$, and $\kappa=l/\alpha$. Here $\alpha,l$ are constants satisfying
$\displaystyle 0<\alpha\leq\lambda_{\min}(\nabla^{2}L(W))\leq\lambda_{\max}(
\nabla^{2}L(W))\leq l,\text{ for all }\,W\in R_{W},$
(9)
where $R_{W}=\{W\in\mathbb{R}^{d\times d}\,|\,\|W-W^{*}\|_{F}\leq\|W^{0}-W^{*}\|_{F}\}$.
- (2)
Denote $W^{*}=2(\Lambda^{-1}+G)$, $q=x_{\tau,\mathsf{query}}$, $\mu=\mu_{\tau,1}-\mu_{\tau,0}$, $u=2(\mu_{\tau,1}+\mu_{\tau,0})$, $a=\mu^{\top}\Lambda^{-1}q$ for simplicity. Then we have
$\displaystyle\|G\|_{\max}\leq$
$\displaystyle\frac{1}{N}\Big{\|}S^{-1}\mathbb{E}\big{[}\sigma^{\prime}(a)\left
(4qq^{\top}+\frac{1}{4}uu^{\top}\Lambda^{-1}qq^{\top}\right)$
$\displaystyle+\sigma^{\prime\prime}(a)\left(\frac{1}{8}(u^{\top}\Lambda^{-1}q)
^{2}\mu q^{\top}+2q^{\top}\Lambda^{-1}q\mu q^{\top})\big{]}\right)\Big{\|}_{
\max}+o\left(\frac{1}{N}\right),$
(10)
where $S=4\nabla^{2}\widetilde{L}(2\Lambda^{-1})$, $\widetilde{L}(2\Lambda^{-1})=\lim_{N\to\infty}L(2\Lambda^{-1})$, $\sigma^{\prime}(\cdot)$ and $\sigma^{\prime\prime}(\cdot)$ are the first- and second-order derivatives of $\sigma(\cdot)$, respectively, and the expectation is taken over $(\mu_{\tau,0}$, $\mu_{\tau,1})\sim\mathcal{P}^{b}_{\Omega}(\Lambda)$, $x_{\tau,\mathsf{query}}\sim\mathcal{P}^{b}_{x}(\mu_{\tau,0},\mu_{\tau,1},\Lambda)$.

The detailed proof of Theorem 3.1 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.1.
As a first step, we prove in Lemma D.2 that the expected loss function $L(W)$ in (7) is strictly convex with respect to (w.r.t.) $W$ and is strongly convex in any compact set of $\mathbb{R}^{d\times 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. 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\to\infty$, our loss function $L(W)$ converges to $\widetilde{L}(W)$ pointwisely (defined in (25)), and the global minimizer $W^{*}$ converges to $2\Lambda^{-1}$. Thus, we denote $W^{*}=2(\Lambda^{-1}+G)$, and prove $\|G\|_{\max}$ is bounded and scales as $\|G\|_{\max}=O(N^{-1/2})$. Next, in Lemma D.4, by further analyzing the Taylor expansion of the equation $\nabla L(W^{*})=0$ at the point $2\Lambda^{-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 $\alpha$-strongly convex and $l$-smooth. Finally, leveraging the standard results from the convex optimization, we prove Theorem 3.1.

According to Theorem 3.1, we have $W^{t}=W^{*}+H^{t}$ where $\|H^{t}\|_{\max}\leq\exp(-t/(2\kappa))\|W^{0}-W^{*}\|_{F}$. If we set $T\geq 2\kappa\log(N\cdot\|W^{0}-W^{*}\|_{F})$, we have $\|H^{T}\|_{\max}\leq 1/N$. Denoting $\widehat{W}=W^{T}$, we have $\widehat{W}=2(\Lambda^{-1}+G+H^{T}/2)=2(\Lambda^{-1}+\widehat{G})$, where $\widehat{G}=G+H^{T}/2,\|\widehat{G}\|_{\max}\leq\|G\|_{\max}+\|H^{T}\|_{\max}=
O(1/N)$. Thus, we have the following corollary.

###### Corollary 3.1

If we optimize the expected loss $L(W)$ in (7) via gradient descent with training prompt length $N$, initial parameters $W^{0}$, and learning rate $\eta=1/l$, then, under Assumption 3.1, after $T\geq 2\kappa\log(N\|W^{0}-W^{*}\|_{F})$ steps, the updated model $\widehat{W}$ satisfies

$$ $\displaystyle\widehat{W}=2(\Lambda^{-1}+\widehat{G}),$ (11) $$

where $\|\widehat{G}\|_{\max}=O(1/N)$, $\kappa=l/\alpha$,
and $\alpha,l$ are constants defined in (9).

Theorem 3.1 and Corollary 3.1 show that training a single-layer transformer with properly distributed data (Assumption 3.1) for binary classification via gradient descent can *linearly* converge to its global minimum $W^{*}=2(\Lambda^{-1}+G)$. Furthermore, when the prompt length $N$ grows, this global minimum $W^{*}$ will converge to $2\Lambda^{-1}$ at a rate of $O(1/N)$.

###### Definition 3.1

###### Assumption 3.1

###### Theorem 3.1

###### Corollary 3.1

### 3.2 In-context inference

Next, we analyze the performance of the trained transformer (11) for in-context binary classification tasks. We make the following assumption.

###### Assumption 3.2

For an in-context test prompt $P_{\mathsf{test}}=(x_{1},y_{1},\dots,x_{M},y_{M},x_{\mathsf{query}})$, we assume

- (1)
$\{x_{i},y_{i}\}_{i=1}^{M}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}
\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$, $x_{\mathsf{query}}\in\mathbb{R}^{d}$.
- (2)
$\mu_{0}^{\top}\Lambda^{-1}\mu_{0}=\mu_{1}^{\top}\Lambda^{-1}\mu_{1}.$

With this assumption, for $y_{\mathsf{query}}\sim\mathcal{P}^{b}_{y|x_{\mathsf{query}}}(\mu_{0},\mu_{1},\Lambda)$, according to the Bayes’ theorem, we have

$$ $\displaystyle\mathbb{P}\left(y_{\mathsf{query}}=1|x_{\mathsf{query}}\right)=$ $\displaystyle\frac{f(x_{\mathsf{query}}|y_{\mathsf{query}}=1)\mathbb{P}\left(y _{\mathsf{query}}=1\right)}{\sum_{z\in\{\pm 1\}}f(x_{\mathsf{query}}|y_{ \mathsf{query}}=z)\mathbb{P}\left(y_{\mathsf{query}}=z\right)}=\sigma((\mu_{1} -\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{query}}).$ $$

If we test the trained transformer with parameters $\widehat{W}$ in (11) and $P_{\mathsf{test}}$, by a simple calculation, we have

$$ $\displaystyle\widehat{y}_{\mathsf{out}}=\sigma\left(\left(\frac{2}{M}\sum_{i=1 }^{M}y_{i}x_{i}^{\top}\right)(\Lambda^{-1}+\widehat{G})x_{\mathsf{query}} \right).$ (12) $$

Intuitively, when the training prompt length $N\to\infty$, we have $\widehat{G}\to 0$, and when the test prompt length $M\to\infty$, we have $\frac{2}{M}\sum_{i=1}^{M}y_{i}x_{i}^{\top}\to(\mu_{1}-\mu_{0})^{\top}$. Thus, when $N,M\to\infty$, $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=1\right)=\widehat{y}_{\mathsf{out
}}\to\sigma((\mu_{1}-\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{query}})=\mathbb{P}
\left(y_{\mathsf{query}}=1|x_{\mathsf{query}}\right)$, and the prediction of the trained transformer $\widehat{y}_{\mathsf{query}}$ perfectly matches with the distribution of the ground truth label $y_{\mathsf{query}}$.

By analyzing the Taylor expansion of $\widehat{y}_{\mathsf{out}}$ at point $\sigma((\mu_{1}-\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{query}})$, we formally present the aforementioned intuition in the following theorem, which establishes an upper bound of the total variation distance between $y_{\mathsf{query}}$ and $\widehat{y}_{\mathsf{query}}$.

###### Theorem 3.2

Consider a test prompt $P_{\mathsf{test}}$ satisfying Assumption
3.2, and let $y_{\mathsf{query}}\sim\mathcal{P}^{b}_{y|x_{\mathsf{query}}}(\mu_{0},\mu_{1},\Lambda)$. Let $\widehat{y}_{\mathsf{query}}$ be the prediction of the trained transformer with parameters $\widehat{W}$ in (11). Then, for the inference error defined in (3), we have

$$ $\displaystyle\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}})]$ $\displaystyle\leq\sigma^{\prime}(\mu^{\top}\Lambda^{-1}q)\left[\|\widehat{G}\| _{\max}\sum_{i,j\in[d]}|\mu_{i}q_{j}|+\frac{1}{\sqrt{M}}\left(\frac{1}{2}|u^{ \top}\Lambda^{-1}q|+\frac{2\sqrt{2}}{\sqrt{\pi}}\sum_{i,j\in[d]}|\Lambda^{-1/2 }_{ij}q_{j}|\right)\right]+o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right),$ $$

where $\mu=\mu_{1}-\mu_{0}$,
$u=2(\mu_{1}+\mu_{0})$,
$q=x_{\mathsf{query}}$, and the expectation is taken over $\{x_{i},y_{i}\}_{i=1}^{M}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}
\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$.

The proof of Theorem 3.2 can be found in Appendix E.
Since $\|\widehat{G}\|_{\max}=O(1/N)$, Theorem 3.2 suggests that if we ignore the constants regarding $\mu_{0},\mu_{1},\Lambda,x_{\mathsf{query}}$, the expected total variation distance between $y_{\mathsf{query}}$ and $\widehat{y}_{\mathsf{query}}$ is at most $O(1/N+1/\sqrt{M})$.
On the other hand, for data pair $(x,y)\sim\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$, the Bayes-optimal classifier is $\mathbb{P}(y=1|x)=f(x|y)\mathbb{P}(y=1)/f(x)=\sigma((\mu_{1}-\mu_{0})^{\top}
\Lambda^{-1}x)$, which corresponds to the logistic regression model $\sigma(w^{\top}x+b)$ with parameters $w=\Lambda^{-1}(\mu_{1}-\mu_{0})$ and $b=0$. Therefore, when $N,M\to\infty$, the prediction of the trained transformer is Bayes-optimal, and is equivalent to the optimal logistic regressor for binary classification problems with distribution $\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$.
Note that different from Assumption 3.1 which states that $\mu_{\tau,0},\mu_{\tau,1},x_{\tau,\mathsf{query}}$ are sampled according to some specific distributions during training, Assumption 3.2 does not impose strong distributional constraints on $\mu_{0},\mu_{1}$ and $x_{\mathsf{query}}$, which shows the strong generalization ability of the trained transformer. We also discuss the consequences when Assumption 3.2 does not hold in Remark E.1, which highlights the necessity of Assumption 3.2.
Moreover, even if $M\to\infty$, the distribution variation between $y_{\mathsf{query}}$ and $\widehat{y}_{\mathsf{query}}$ does not disappear unless $N\to\infty$. 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. () for in-context linear regression.

###### Assumption 3.2

###### Theorem 3.2

## 4 In-context multi-class classification

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\geq 2$ classes can be expressed as ${P=(x_{1},y_{1},\dots,x_{N},y_{N},x_{\mathsf{query}})}$, where $x_{i}\in\mathbb{R}^{d}$, $y_{i}\in\{\textbf{e}_{1},\textbf{e}_{2},\dots,\textbf{e}_{c}\}$, and $\textbf{e}_{i}$ is the $i$-th standard unit vector of $\mathbb{R}^{c}$.
Its embedding matrix can be formulated as

$$ $E=E(P)=\begin{pmatrix}x_{1}&x_{2}&\cdots&x_{N}&x_{\mathsf{query}}\\ y_{1}&y_{2}&\cdots&y_{N}&0_{c}\end{pmatrix}\in\mathbb{R}^{(d+c)\times(N+1)}.$ (13) $$

Similar to the binary case, we set some of the parameters in our model as 0 0 and $1$ to simplify the optimization problem and consider the parameters of our model $(W^{V},W^{KQ})$ in the following sparse form:

$$ $\displaystyle W^{V}=\begin{pmatrix}0_{d\times d}&0_{d\times c}\\ 0_{c\times d}&I_{c}\end{pmatrix},\qquad W^{KQ}$ $\displaystyle=\begin{pmatrix}W&0_{d\times c}\\ 0_{c\times d}&0_{c\times c}\end{pmatrix},$ (14) $$

where $W\in\mathbb{R}^{d\times d}$. We set the normalization factor $\rho$ 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 $\widehat{y}_{\mathsf{out}}$. With parameters $W$ and a prompt $P=(x_{1},y_{1},\dots,x_{N},y_{N},x_{\mathsf{query}})$, the output can be expressed as

$$ $\widehat{y}_{\mathsf{out}}=\mathrm{softmax}\left([F(E(P);W)]_{(d+1):(d+c),(N+1 )}\right)=\mathrm{softmax}\left(\left(\frac{1}{N}\sum_{i=1}^{N}y_{i}x_{i}^{ \top}\right)Wx_{\mathsf{query}}\right).$ $$

We denote the prediction of the model for $x_{\mathsf{query}}$ as $\widehat{y}_{\mathsf{query}}$, which is a random variable depending on $\widehat{y}_{\mathsf{out}}$.
Randomly sample a random variable $u$ that is uniformly distributed on $[0,1]$. If $u\in\big{[}\sum_{j=1}^{i-1}(\widehat{y}_{\mathsf{out}})_{j},\sum_{j=1}^{i}(
\widehat{y}_{\mathsf{out}})_{j}\big{)}$, where $(\widehat{y}_{\mathsf{out}})_{j}$ is the $j$-th element of $\widehat{y}_{\mathsf{out}}$, we let $\widehat{y}_{\mathsf{query}}=\textbf{e}_{i}$. Thus, $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=\textbf{e}_{i}\right)=(\widehat{y
}_{\mathsf{out}})_{i}$.

### 4.1 Training procedure

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)\sim\mathcal{P}^{m}(\mu,\Lambda)$ if $\mathbb{P}\left(y=\textbf{e}_{i}\right)=1/c$ and $f(x|y=\textbf{e}_{i})=\mathsf{N}(\mu_{i},\Lambda)$ for $i\in[c]$, where $\mu=(\mu_{1},\dots,\mu_{c})\in\mathbb{R}^{d\times c}$ and $\Lambda\in\mathbb{R}^{d\times d}$ is a positive definite matrix.

We consider the case of $B$ training tasks indexed by $\tau\in[B]$. Each training task $\tau$ is associated with a prompt $P_{\tau}=(x_{\tau,1},y_{\tau,1},\dots,x_{\tau,N},y_{\tau,N},x_{\tau,\mathsf{
query}})$ and a corresponding label $y_{\tau,\mathsf{query}}$.
We make the following assumption in this section.

###### Assumption 4.1

For each learning task $\tau\in[B]$, we assume

- (1)
$\{x_{\tau,i},y_{\tau,i}\}_{i=1}^{N},\{x_{\tau,\mathsf{query}},y_{\tau,\mathsf{
query}}\}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}^{m}(\mu_{
\tau}=(\mu_{\tau,1},\dots,\mu_{\tau,c}),\Lambda)$.
- (2)
$\mu_{\tau,1}$ is sampled from $\mathsf{N}(0,I_{d})$, and $\mu_{\tau,k}=U_{\tau,k,\Lambda}\mu_{\tau,1},k=2,3,\dots,c$, where $U_{\tau,k,\Lambda}=\Lambda^{1/2}U_{\tau,k}\Lambda^{-1/2}$, and $U_{\tau,k}$ are uniformly distributed over the closed set of real unitary matrices
such that $U_{\tau,k}U_{\tau,k}^{\top}=I_{d}$.

We denote the distribution of $\mu_{\tau}$ as $\mathcal{P}^{m}_{\Omega}(\Lambda)$. Note that $U_{\tau,k,\Lambda}=\Lambda^{1/2}U_{\tau,k}\Lambda^{-1/2}$ can be viewed as linear transformations that preserve the inner product of vectors in the $\Lambda^{-1}$ weighted norm, and we have $\mu_{\tau,i}^{\top}\Lambda^{-1}\mu_{\tau,i}=\mu_{\tau,j}^{\top}\Lambda^{-1}\mu
_{\tau,j}$, for $i,j\in[c]$.
Let $\widehat{y}_{\tau,\mathsf{out}}=\mathrm{softmax}\left([F(E(P_{\tau});W)]_{(d+1
):(d+c),(N+1)}\right)$ be the output of the transformer for task $\tau$. We define the empirical risk over $B$ independent tasks as

$$ $\widehat{L}(W)=\frac{1}{B}\sum_{\tau=1}^{B}\sum_{k=1}^{c}-(y_{\tau,\mathsf{ query}})_{k}\log((\widehat{y}_{\tau,\mathsf{out}})_{k}).$ (15) $$

Taking the limit of infinite training tasks $B\to\infty$, the expected training loss can be defined as

$$ $L(W)=\lim_{B\to\infty}\widehat{L}(W)=-\mathbb{E}\left[\sum_{k=1}^{c}(y_{\tau, \mathsf{query}})_{k}\log((\widehat{y}_{\tau,\mathsf{out}})_{k})\right],$ (16) $$

where the expectation is taken over $\mu_{\tau}\sim\mathcal{P}^{m}_{\Omega}(\Lambda)$, $\{x_{\tau,i},y_{\tau,i}\}_{i=1}^{N},\{x_{\tau,\mathsf{query}},y_{\tau,\mathsf{
query}}\}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}^{m}(\mu_{
\tau},\Lambda)$.

Applying gradient descent over the expected training loss (16), we have the following theorem.

###### Theorem 4.1

(Informal)

Under Assumption 4.1, the following statements hold.

- (1)
Optimizing training loss $L(W)$ in (16) with training prompt length $N$ via gradient descent $W^{t+1}=W^{t}-\eta\nabla L(W^{t})$, for any $t\geq 1$, we have
$\displaystyle\|W^{t}-W^{*}\|_{F}^{2}\leq\exp(-t/\kappa)\|W^{0}-W^{*}\|_{F}^{2},$
(17)
where $W^{0}$ is the initial parameter and $W^{*}$ is the global minimizer of $L(W)$, $\kappa=l/\alpha$. Here, $\alpha,l$ are constants such that
$\displaystyle 0<\alpha\leq\lambda_{\min}(\nabla^{2}L(W))\leq\lambda_{\max}(
\nabla^{2}L(W))\leq l,\text{ for all }\,W\in R_{W},$
(18)
where $R_{W}=\{W\in\mathbb{R}^{d\times d}\,|\,\|W-W^{*}\|_{F}\leq\|W^{0}-W^{*}\|_{F}\}$.
- (2)
Denoting $W^{*}=c(\Lambda^{-1}+G)$, we have $\|G\|_{\max}=O(c/N)$.
- (3)
After $T\geq 2\kappa\log(N\cdot\|W^{0}-W^{*}\|_{F})$ steps, denoting the updated model $\widehat{W}$ satisfies
$\displaystyle\widehat{W}=c(\Lambda^{-1}+\widehat{G}),$
(19)
where $\|\widehat{G}\|_{\max}=O(c/N)$.

The formal statement and proof of Theorem 4.1 can be found in Appendix F.
Technically, the proof of Theorem 4.1 builds on that of Theorem 3.1, 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 derived 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.1 shows that training a single-layer transformer with properly distributed data (Assumption 4.1) for in-context multi-class classification via gradient descent can linearly converge to its global minimum $W^{*}=c(\Lambda^{-1}+G)$. When the prompt length $N$ grows, this global minimum $W^{*}$ will converge to $c\Lambda^{-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$.

###### Definition 4.1

###### Assumption 4.1

###### Theorem 4.1

### 4.2 In-context inference

###### Assumption 4.2

For an in-context test prompt $P_{\mathsf{test}}=(x_{1},y_{1},\dots,x_{M},y_{M},x_{\mathsf{query}})$, we assume

- (1)
$\{x_{i},y_{i}\}_{i=1}^{M}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}
\mathcal{P}^{m}(\mu,\Lambda)$, $\mu=(\mu_{1},\dots,\mu_{c})\in\mathbb{R}^{d\times c}$, $x_{\mathsf{query}}\in\mathbb{R}^{d}$.
- (2)
$\mu_{i}^{\top}\Lambda^{-1}\mu_{i}=\mu_{j}^{\top}\Lambda^{-1}\mu_{j}$, for $i,j\in[c]$.

With this assumption, for $y_{\mathsf{query}}\sim\mathcal{P}^{m}_{y|x_{\mathsf{query}}}(\mu,\Lambda)$, according to the Bayes’ theorem, we have

$$ $\displaystyle\mathbb{P}\left(y_{\mathsf{query}}=\textbf{e}_{k}|x_{\mathsf{ query}}\right)=\frac{f(x_{\mathsf{query}}|y_{\mathsf{query}}=\textbf{e}_{k}) \mathbb{P}\left(y_{\mathsf{query}}=\textbf{e}_{k}\right)}{\sum_{j=1}^{c}f(x_{ \mathsf{query}}|y_{\mathsf{query}}=\textbf{e}_{j})\mathbb{P}\left(y_{\mathsf{ query}}=\textbf{e}_{j}\right)}=\mathrm{softmax}(\mu^{\top}\Lambda^{-1}x_{ \mathsf{query}})_{k}.$ $$

If we test the trained transformer with parameters $\widehat{W}$ in (19) and prompt $P_{\mathsf{test}}$, by a simple calculation, we have

$$ $\displaystyle\widehat{y}_{\mathsf{out}}=\mathrm{softmax}\left(\left(\frac{c}{M }\sum_{i=1}^{M}y_{i}x_{i}^{\top}\right)(\Lambda^{-1}+\widehat{G})x_{\mathsf{ query}}\right).$ (20) $$

Note that, when the training prompt length $N\to\infty$, we have $\widehat{G}\to 0$, and when the test prompt length $M\to\infty$, we have $\frac{c}{M}\sum_{i=1}^{M}y_{i}x_{i}^{\top}\to\mu^{\top}$. Thus, when $N,M\to\infty$, $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=\textbf{e}_{k}\right)=(\widehat{y
}_{\mathsf{out}})_{k}\to\mathrm{softmax}(\mu^{\top}\Lambda^{-1}x_{\mathsf{
query}})_{k}=\mathbb{P}\left(y_{\mathsf{query}}=\textbf{e}_{k}|x_{\mathsf{
query}}\right)$, i.e., the prediction of the trained transformer $\widehat{y}_{\mathsf{query}}$ matches the ground truth label $y_{\mathsf{query}}$.

By analyzing the Taylor expansion of $\widehat{y}_{\mathsf{out}}$ at point $\mathrm{softmax}(\mu^{\top}\Lambda^{-1}x_{\mathsf{query}})$, we crystallize the aforementioned intuition in the following theorem, which establishes an upper bound of the total variation distance between $y_{\mathsf{query}}$ and $\widehat{y}_{\mathsf{query}}$.

###### Theorem 4.2

(Informal)
Let $P_{\mathsf{test}}$ satisfy Assumption
4.2 and $y_{\mathsf{query}}\sim\mathcal{P}^{m}_{y|x_{\mathsf{query}}}(\mu,\Lambda)$. Denote $\widehat{y}_{\mathsf{query}}$ as the prediction of the trained transformer with parameter $\widehat{W}$ in (19). Then, for the inference error defined in (3), we have

$$ $\displaystyle\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}} )]=O(c^{2}N^{-1}+c^{3/2}M^{-1/2}),$ $$

where the expectation is taken over $\{x_{i},y_{i}\}_{i=1}^{M}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}
\mathcal{P}^{m}(\mu,\Lambda)$.

The formal statement and proof of Theorem 4.2 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)\sim\mathcal{P}^{m}(\mu,\Lambda)$, the Bayes-optimal classifier is $\mathbb{P}(y=\textbf{e}_{k}|x)=f(x|y)\mathbb{P}(y=\textbf{e}_{k})/f(x)=\mathrm
{softmax}(\mu^{\top}\Lambda^{-1}x)_{k}$, which corresponds to a softmax regression model $\mathrm{softmax}(Wx+b)$ with parameters $W=\mu^{\top}\Lambda^{-1}$ and $b=0$. When $N,M\to\infty$, the prediction of the trained transformer is Bayes-optimal, and is equivalent to the optimal softmax regressor for multi-class classification problems with distribution $\mathcal{P}^{m}(\mu,\Lambda)$. Note that different from Assumption 4.1 which states that $\mu_{\tau},x_{\tau,\mathsf{query}}$ are sampled according to some specific distributions during training, Assumption 4.2 impose strong distributional constraints on $\mu$ or $x_{\mathsf{query}}$, which shows the strong generalization ability of the trained transformer. We also discuss the consequences when Assumption 4.2 does not hold in Remark G.1, which highlights the necessity of Assumption 4.2.
Moreover, even if $M\to\infty$, the distribution variation between $y_{\mathsf{query}}$ and $\widehat{y}_{\mathsf{query}}$ does not disappear unless $N\to\infty$. Thus, the ICL ability of the trained transformers for multi-class classification is limited by the finite length of training prompts. Similar behaviors have also been observed in Zhang et al. () for in-context linear regression and in Section 3.2 for in-context binary classification.

###### Assumption 4.2

###### Theorem 4.2

## 5 Experiments

In this section, we carry out experiments to verify the theoretical claims and compare the ICL performances of transformers with other machine learning algorithms. Detailed experimental settings can be found in Appendix H.

### 5.1 ICL performances of the single-layer transformers

Figure: (a) $c=10$
Refer to caption: extracted/5929270/fig/N.png

We first train single-layer transformers for in-context classification of Gaussian mixtures with different numbers of classes $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 from Figure 1 (a,b) that the inference errors decrease as $N$ and $M$ increase, and they increase as $c$ increases. In Figure 1 (c,d), we first fix the training prompt length (test prompt length) to a large number $2000$, and then vary the test prompt length (training prompt length) from 20 to 2000. The results show that, as $M$ and $N$ become sufficiently large, the inference error, which is an approximation of $\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}})]$ (see Appendix H for detailed definitions), decreases to near-zero. This indicates that the prediction of the trained transformer approaches the Bayes-optimal classifier. All these experimental results corroborate our theoretical claims.

### 5.2 Comparison of transformers with other machine learning algorithms

Figure: Figure 2: Inference errors for the classification of three Gaussian mixtures. Colored lines correspond to different models as follows: a single-layer transformer defined in Section 4 (1-layer, sparse), a single-layer transformer with full parameters (59) (1-layer, full), a 3-layer transformer with a GPT2 architecture (3-layer, GPT2), softmax regression (softmax), SVM with linear kernel (SVM, linear), SVM with Gaussian kernel (SVM, gaussian), 1-nearest neighbor (1-NN), and 3-nearest neighbor (3-NN). All three transformers are trained with prompt length $N=100$.
Refer to caption: extracted/5929270/fig/Compare.png

Additionally, we conduct experiments comparing the ICL performances of the transformers with other machine learning algorithms for the classification of three Gaussian mixtures. Again, the detailed experimental setting can be found in Appendix H. Form Figure 2, we can see that, when the prompt length is less than 100, all three transformer models have comparable or better performances than the traditional methods (softmax regression, SVM, $K$-nearest neighbor), which demonstrates the strong ICL capacities of transformers. When the prompt length is larger than 100, a 3-layer transformer with the GPT-2 architecture shows a decline in performance, probably because our transformer models were only trained with a small prompt length of $N=100$. Similar declined performance when the training prompt length is smaller than the test prompt length has also been observed for in-context linear regression tasks; see e.g. Zhang et al. ().
However, even when the prompt length is larger than 100, the single-layer transformer defined in Section 4 (1-layer, sparse) and the single-layer transformer with full parameters (59) (1-layer, full) still significantly outperform the traditional methods. This also indicates that, though the model (1-layer, sparse) we studied in this paper is relatively simple, this model is already sufficiently complex for in-context classification of Gaussian mixtures, both in theory and practice.

## 6 Conclusion

We studied the learning dynamics of transformers for in-context 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 regrading the training and test prompts can be relaxed. Additionally, we have only examined single-layer transformers with linear attention and sparse parameters. The learning dynamics of multi-layer transformers with nonlinear attention (e.g., softmax) for in-context classification problems remain an interesting area for future investigation.

## Appendix

The Appendix is organized as follows. In Section A, we provide a literature review of the related works that studied the ICL abilities of transformers. In Section B, we introduce the additional notations for the proofs in the Appendix. In Section C, we introduce some useful Lemmas we adopt from previous literature. In Sections D, E, F, G, we present the proofs of Theorem 3.1, 3.2, 4.1, 4.2 respectively. In Section H, we provide additional details of our experiments.

## Appendix A Related work

It has been observed that transformer-based models have impressive ICL abilities in natural language processing [Brown et al., , Nye et al., , Wei et al., , Dasgupta et al., , Zhang et al., ].
Garg et al. [] 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. [], Von Oswald et al. [] showed by construction that multi-layer transformers can be viewed as multiple steps of gradient descent for linear regression. Akyürek et al. [] also showed that
constructed transformers can implement closed-form ridge regression. Guo et al. [] showed that constructed transformers can perform in-context learning with representations. Bai et al. [] 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. [] demonstrated that constructed transformers can approximate several in-context reinforcement learning algorithms. Fu et al. [], Giannou et al. [] further proved that constructed transformers can perform higher-order optimization algorithms like Newton’s method. Pathak et al. [] showed that transformers can learn mixtures of linear regressions. Giannou et al. [] proved that looped transformers that can emulate various in-context learning algorithms. Cheng et al. [] showed that transformers can perform functional gradient descent for learning non-linear functions in context. Zhang et al. [] 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. [], Ahn et al. [] 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. [] 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. [] 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. [] showed that trained transformers can learn topic structure. Wu et al. [] analyzed the task complexity bound for pretraining single-layer linear transformers on in-context linear regression tasks. Tarzanagh et al. [] built the connections between single-layer transformers and support vector machines (SVMs). Nichani et al. [] showed that transformers trained via gradient descent can learn causal structure.
Chen et al. [] proved the convergence of gradient flow on a multi-head softmax attention model for in-context multi-task linear regression. Kim and Suzuki [], Yang et al. [] proved that trained transformers can learn nonlinear features in context.

Recently, Li et al. [] 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=\mu_{j}+\kappa v_{k}$, where $\{\mu_{j}\}_{j=1}^{M_{1}}$ are in-domain-relevant patterns and $\{\nu_{k}\}_{k=1}^{M_{2}}$ are in-domain-irrelevant patterns, $M_{1}\geq 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 $\mathcal{P}^{b}(\mu_{0},\mu_{1},\Lambda)$ or $\mathcal{P}^{m}(\mu,\Lambda)$, and the range and possible distributions of our data are *infinite*. Furthermore, the transformer architectures analyzed in their work also differ from those in our study, thereby highlighting the distinct contributions and independent interests of our work.

Some works also studied the ICL from other perspectives. To name a few, Xie et al. [] explained the ICL as implicit Bayesian inference; Wang et al. [] explained the LLMs as latent variable models; Zhang et al. [] explained the ICL abilities of transformers as implicitly implementing a Bayesian model averaging algorithm; and Li et al. [] studied the generalization and stability of the ICL abilities of transformers.

## Appendix B Additional notations

We denote $X\sim\mathsf{Bin}(n,p)$ if a random variable $X$ follows the binomial distribution with parameters $n\in\mathbb{N}$ and $p\in[0,1]$, which means $\mathbb{P}\left(X=k\right)=\frac{n!}{k!(n-k)!}p^{k}(1-p)^{n-k}$.
We denote $X\sim\mathsf{Multin}(n,p)$ if random variables $X=(X_{1},X_{2},\dots,X_{k})$ follow the Multinomial distribution with parameters $n\in\mathbb{N}$ and $p_{1}=p_{2}=\dots=p_{k}=1/k$, which means $\mathbb{P}\left(X=(x_{1},x_{2},\dots,x_{k})\right)=\frac{n!}{\prod_{i=1}^{k}x_
{k}!}k^{-n}$.
We denote $\zeta_{i}(x)=\mathrm{softmax}(x)_{i}=\exp(x_{i})/(\sum_{j=1}^{k}\exp(x_{j}))$ for simplicity. We define $\delta_{ii}=1$, $\delta_{ij}=0,i\neq j$. For $x\in\mathbb{N}$, we define $t_{1}(x)=\lfloor(x-1)/d\rfloor+1,t_{2}(x)=((x-1)\text{ mod }d)+1$.

## Appendix C Useful lemmas

###### Lemma C.1 ( [Karimi et al., 2016 ] )

If $f:\mathbb{R}^{d}\to\mathbb{R}$ is $\mu$-strongly convex, then

$$ $\displaystyle f(x)-\min_{x}f(x)$ $\displaystyle\geq\frac{\mu}{2}\|x^{*}-x\|_{2}^{2}$ $$

where $x^{*}=\mathop{\rm argmin}_{x}f(x)$.

###### Lemma C.2 ( [Bubeck, 2015 ] )

Suppose $f:\mathbb{R}^{d}\to\mathbb{R}$ is $\alpha$-strongly convex and $\beta$-smooth for some $0<\alpha\leq\beta$. Then, the gradient descent iterating $w^{t+1}=w^{t}-\eta\nabla f(w^{t})$ with learning rate $\eta=1/\beta$ and initialization $w^{0}\in\mathbb{R}^{d}$ satisfies that for any $t\geq 1$,

$$ $\displaystyle\|w^{t}-w^{*}\|_{2}^{2}\leq\exp(-t/\kappa)\|w^{0}-w^{*}\|_{2}^{2}$ $$

where $\kappa=\beta/\alpha$ is the condition number of $f$, and $w^{*}=\mathop{\rm argmin}_{w\in\mathbb{R}^{d}}f(w)$ is the minimizer of $f$.

###### Lemma C.1 ( [Karimi et al., 2016 ] )

###### Lemma C.2 ( [Bubeck, 2015 ] )

## Appendix D Training procedure for in-context binary classification

In this section, we present the proof of Theorem 3.1.

### D.1 Proof sketch

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 $\mathbb{R}^{d\times 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\to\infty$, our loss function $L(W)$ point wisely converges to $\widetilde{L}(W)$ (defined in (25)), and the global minimizer $W^{*}$ converge to $2\Lambda^{-1}$. We denote $W^{*}=2(\Lambda^{-1}+G)$, and prove $\|G\|_{\max}=O(N^{-1/2})$. Next, in Lemma D.4, by further analyzing the Taylor expansion of the equation $\nabla L(W^{*})=0$ at the point $2\Lambda^{-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 $\alpha$-strongly convex and $l$-smooth. Finally, leveraging the standard results from convex optimization, we prove Theorem 3.1 in subsection D.4.

In this section, we use the following notations.

### D.2 Notations

Recall the expected loss function (7) is

$$ $\displaystyle L(W)=-\frac{1}{2}\mathbb{E}\left[(1+y_{\tau,\mathsf{query}})\log (\widehat{y}_{\tau,\mathsf{out}})+(1-y_{\tau,\mathsf{query}})\log(1-\widehat{y }_{\tau,\mathsf{out}})\right],$ (21) $$

where

$$ $\displaystyle\widehat{y}_{\tau,\mathsf{out}}=\sigma\left(\left(\frac{2}{N}\sum _{i=1}^{N}y_{\tau,i}x_{\tau,i}^{\top}\right)\frac{W}{2}x_{\tau,\mathsf{query}}\right)$ $$

is the output of the transformer, and the label of the data follows the distribution

$$ $\displaystyle\mathbb{P}\left(y_{\tau,\mathsf{query}}=1|x_{\tau,\mathsf{query}} \right)=\sigma((\mu_{\tau,1}-\mu_{\tau,0})^{\top}\Lambda^{-1}x_{\tau,\mathsf{ query}})).$ $$

In this section, we introduce the following notations to analyze (7). We denote
$\mu=\mu_{\tau}$, $\mu_{1}=\mu_{\tau,1}$, $\mu_{0}=\mu_{\tau,0}$ and $q=x_{\tau,\mathsf{query}}$. Then with probability $\mathbb{P}\left(y_{\tau,\mathsf{query}}=1\right)=1/2$ we have $q=\mu_{1}+v$, and with probability $\mathbb{P}\left(y_{\tau,\mathsf{query}}=0\right)=1/2$ we have $q=\mu_{0}+v$, where $v\sim\mathsf{N}(0,\Lambda)$. We define $p=\frac{2}{N}\sum_{i=1}^{N}y_{\tau,i}x_{\tau,i}$. Since with probability $\mathbb{P}\left(y_{\tau,i}=1\right)=1/2$ we have $x_{\tau,i}=\mu_{1}+v_{i}$, and with probability $\mathbb{P}\left(y_{\tau,i}=0\right)=1/2$ we have $x_{\tau,i}=\mu_{0}+v_{i}$, where $v_{i}\sim\mathsf{N}(0,\Lambda)$, we known $p=2N_{1}\mu_{1}/N-2N_{0}\mu_{0}/N+g$, where $g=\frac{2}{N}\sum_{i=1}^{N}v_{i}$, $g\sim\mathsf{N}(0,4\Lambda/N)$, $N_{1}\sim\mathsf{Bin}(N,1/2)$.
Defining $h=N_{1}/N-1/2$, $u=2(\mu_{1}+\mu_{0})$, we have $N_{0}/N=1/2-h$ and

$$ $\displaystyle p=\mu+hu+g.$ (22) $$

Then, the expected loss function (7) can be expressed as

$$ $\displaystyle L(W)=\mathbb{E}[-\sigma(\mu^{\top}\Lambda^{-1}q)\log(\sigma(p^{ \top}Wq/2))-(1-\sigma(\mu^{\top}\Lambda^{-1}q))\log(1-\sigma(p^{\top}Wq/2))].$ (23) $$

The gradient of the loss function (7) can be expressed as

$$ $\displaystyle\nabla L(W)=\frac{1}{2}\mathbb{E}[(\sigma(p^{\top}Wq/2)-\sigma( \mu^{\top}\Lambda^{-1}q))pq^{\top}].$ (24) $$

Moreover, we define a function $\widetilde{L}(W)$ as

$$ $\displaystyle\widetilde{L}(W)=\mathbb{E}[-\sigma(\mu^{\top}\Lambda^{-1}q)\log( \sigma(\mu^{\top}Wq/2))-(1-\sigma(\mu^{\top}\Lambda^{-1}q))\log(1-\sigma(\mu^{ \top}Wq/2))].$ (25) $$

In Lemma D.3, we show that as $N\to\infty$, $L(W)$ will point wisely converge to $\widetilde{L}(W)$.

### D.3 Lemmas

###### Lemma D.1

Suppose $N_{1}\sim\mathsf{Bin}(N,1/2)$. Defining $h=N_{1}/N-1/2$, we have

$$ $\displaystyle\mathbb{E}[h]=0$ $\displaystyle\mathbb{E}[h^{2}]=\frac{1}{4N}$ $\displaystyle\mathbb{E}[h^{3}]=0$ $\displaystyle\mathbb{E}[h^{n}]=O(N^{-2}),\text{ for }n\geq 4$ $\displaystyle\mathbb{E}[|h|]\leq\frac{1}{2N^{1/2}}$ $\displaystyle\mathbb{E}[|h^{3}|]=O(N^{-3/2}).$ $$

Proof
Since $N_{1}\sim\mathsf{Bin}(N,1/2)$, the moment-generating function of $N_{1}$ is

$$ $\displaystyle M_{N_{1}}(t)=\left(\frac{1}{2}+\frac{1}{2}\exp(t)\right)^{N}.$ $$

We can compute the moment-generating function of $h$ as follows:

$$ $\displaystyle M_{h}(t)=$ $\displaystyle\exp{\left(-\frac{t}{2}\right)}M_{N_{1}}\left(\frac{t}{N}\right)= \left(\frac{\exp{\frac{-t}{2N}}+\exp{\frac{t}{2N}}}{2}\right)^{N}=\left(\cosh \left(\frac{t}{2N}\right)\right)^{N}$ $\displaystyle=$ $\displaystyle\left(1+\frac{t^{2}}{8N^{2}}+\sum_{i=2}^{\infty}\frac{t^{2i}}{(2i )!(2N)^{2i}}\right)^{N}.$ $$

Thus, we know the coefficients of $t$, $t^{2}$, $t^{3}$ are $0,1/(8N),0$ respectively, and the coefficients of $t^{n},n\geq 4$ are $O(1/N^{2})$. We have

$$ $\displaystyle\mathbb{E}[h]=0$ $\displaystyle\mathbb{E}[h^{2}]=\frac{1}{4N}$ $\displaystyle\mathbb{E}[h^{3}]=0$ $\displaystyle\mathbb{E}[h^{n}]=O(1/N^{2}),\text{ for }n\geq 4.$ $$

Moreover, according to the Jensen’s inequality, we have

$$ $\displaystyle\mathbb{E}[|h|]\leq\left(\mathbb{E}[h^{2}]\right)^{1/2}=\frac{1}{ 2N^{1/2}}$ $\displaystyle\mathbb{E}[|h^{3}|]\leq\left(\mathbb{E}[h^{4}]\right)^{3/4}=O(N^{ -3/2}).$ $$

$\blacksquare$

###### Lemma D.2

For the loss function $L(W)$ (7), we have $\nabla^{2}L(W)\succ 0$. For any compact set $R_{W}$ of $\mathbb{R}^{d\times d}$, when $W\in R_{W}$, we have $\nabla^{2}L(W)\succ\gamma I_{d}$ for some $\gamma>0$.
Additionally, $L(W)$ has one unique global minimizer on $\mathbb{R}^{d\times d}$.

For $\widetilde{L}(W)$ defined in (25), we also have $\nabla^{2}\widetilde{L}(W)\succ 0$. For any compact set $R_{W}$ of $\mathbb{R}^{d\times d}$, when $W\in R_{W}$, we have $\nabla^{2}\widehat{L}(W)\succ\gamma I_{d}$ for some $\gamma>0$.
Additionally, $\widetilde{L}(W)$ has one unique global minimizer on $\mathbb{R}^{d\times d}$.

Proof
We vectorize $W$ as $\text{Vec}(W)\in\mathbb{R}^{d^{2}}$, where $\text{Vec}(W)_{i}=W_{t_{1}(i),t_{2}(i)}$, $t_{1}(x)=\lfloor(x-1)/d\rfloor+1,t_{2}(x)=((x-1)\text{ mod }d)+1$. Then, we have

$$ $\displaystyle(\nabla L(W))_{i}=\mathbb{E}_{p,q}\left[\frac{1}{2}(\sigma(p^{ \top}Wq/2)-\sigma(\mu^{\top}\Lambda^{-1}q))p_{t_{1}(i)}q_{t_{2}(i)}\right].$ (26) $$

The Hessian matrix of the loss function (7) is

$$ $\displaystyle(\nabla^{2}L(W))_{ij}=\mathbb{E}_{p,q}\left[\frac{1}{4}\sigma(p^{ \top}Wq/2)(1-\sigma(p^{\top}Wq/2))p_{t_{1}(i)}q_{t_{2}(i)}p_{t_{1}(j)}q_{t_{2} (j)}\right].$ $$

Considering $z\in\mathbb{R}^{d^{2}}$ such that $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z=$ $\displaystyle\mathbb{E}_{q,p}\left[\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^ {\top}Wq/2))\sum_{ab}z_{a}z_{b}p_{t_{1}(a)}q_{t_{2}(a)}p_{t_{1}(b)}q_{t_{2}(b) }\right]$ $\displaystyle=$ $\displaystyle\int\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top}Wq/2))\left( \sum_{a\in[d^{2}]}z_{a}p_{t_{1}(a)}q_{t_{2}(a)}\right)^{2}f_{pq}(p,q)dpdq,$ $$

where $f_{pq}(p,q)$ are the probability density function (PDF) function of $p,q$. Since for any $p,q$, $\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top}Wq/2))>0$, we have $z^{\top}\nabla^{2}L(W)z\geq 0$. Thus, $\nabla^{2}L(W)\succeq 0$ and $L(W)$ is convex.

Moreover, for any $z\neq 0$, we denote $z_{ij}=z_{((i-1)d+j)},i,j\in[d]$. Suppose $a,b\in\arg\max_{i,j}|z_{ij}|$, we consider a set of constants
$\{c_{1pi},c_{2pi}\},\{c_{1qi},c_{2qi}\},i,j\in[d]$, where
$c_{1pa}=d,c_{2pa}=d+1$, $c_{1qb}=d,c_{2qb}=d+1$, and $c_{1pi}=1/16,c_{2pi}=1/8,i\neq a$, $c_{1qj}=1/16,c_{2qj}=1/8,j\neq b$. Then, for any $c_{pi}\in[c_{1pi},c_{2pi}],c_{qj}\in[c_{1qj},c_{2qj}]$. We have

$$ $\displaystyle\left|\sum_{i,j\in[d]}z_{ij}c_{pi}c_{qj}\right|\geq\left[d^{2}-2( d+1)(d-1)/8-(d-1)^{2}/64\right]\max_{ij}|z_{ij}|\geq d^{2}\max_{ij}|z_{ij}|/2.$ $$

Then, we define region $\Omega(a,b)\triangleq\{p=\sum_{i}c_{pi}\textbf{e}_{i}$, $q=\sum_{j}c_{qj}\textbf{e}_{j}$,$c_{pi}\in[c_{1pi},c_{2pi}],c_{qj}\in[c_{1qj},c_{2qj}]\}$. We have

$$ $\displaystyle\min_{\Omega(a,b)}\left(\sum_{c\in[d^{2}]}z_{c}p_{t_{1}(c)}q_{t_{ 2}(c)}\right)^{2}\geq d^{4}\max_{ij}|z_{ij}|^{2}/4\geq\|z\|_{2}^{2}/4.$ $$

Defining

$$ $\displaystyle C(\Omega)=\min_{a\in[d],b\in[d]}\int_{\Omega(a,b)}f_{pq}(p,q)dpdq,$ $\displaystyle S(\Omega,W)=\min_{a\in[d],b\in[d]}\min_{\Omega(a,b)}\left\{\frac {1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top}Wq/2))\right\},$ $$

we have $S(\Omega,W)>0$. Since with probability $\mathbb{P}\left(y_{\tau,\mathsf{query}}=1\right)=1/2$, $q=\mu_{1}+v$, with probability $\mathbb{P}\left(y_{\tau,\mathsf{query}}=0\right)=1/2$, $q=\mu_{0}+v$, where $v\sim\mathsf{N}(0,\Lambda)$ and
$p=\mu+hu+g$, where
$g\sim\mathsf{N}(0,4\Lambda/N),v\sim\mathsf{N}(0,\Lambda),\mu_{0}\sim\mathsf{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\in\mathbb{R}^{d}$. Moreover, $\Omega(a,b)$ are non-zero measures on $\mathbb{R}^{d\times d}$. Thus, we have $C(\Omega)>0$. Then, for any $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z\geq$ $\displaystyle\int_{\Omega(a,b)}\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{ \top}Wq/2))\left(\sum_{l}z_{l}p_{t_{1}(l)}q_{t_{2}(l)}\right)^{2}f_{pq}(p,q)dpdq$ $\displaystyle\geq$ $\displaystyle C(\Omega)S(\Omega,W)\|z\|_{2}^{2}/4$ $\displaystyle>$ $\displaystyle 0.$ $$

Thus, we have $\nabla^{2}L(W)\succ 0$. $L(W)$ is strictly convex.

Moreover, for any compact set $R_{W}$ of $\mathbb{R}^{d\times d}$, for any $W\in R_{W}$, we have

$$ $\displaystyle S(\Omega)=\min_{W\in R_{W}}\min_{a\in[d],b\in[d]}\min_{\Omega(a, b)}\left\{\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top}Wq/2))\right\}>0.$ $$

Then, for any $W\in R_{W}$, for any $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z\geq$ $\displaystyle\int_{\Omega(a,b)}\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{ \top}Wq/2))\left(\sum_{l}z_{l}p_{t_{1}(l)}q_{t_{2}(l)}\right)^{2}f_{pq}(p,q)dpdq$ $\displaystyle\geq$ $\displaystyle\frac{1}{4}C(\Omega)S(\Omega)\|z\|_{2}^{2}.$ $$

Thus, when $W\in R_{W}$, where $R_{W}$ is a compact set, we have $\nabla^{2}L(W)\succ C(\Omega)S(\Omega)I_{d}/4$ and the loss function $L(W)$ is $\gamma-$strongly convex, where $\gamma=C(\Omega)S(\Omega)/4$.

Because our loss function is strictly convex in $\mathbb{R}^{d\times d}$, it has at most one global minimizer in $\mathbb{R}^{d\times d}$. Next, we prove all level sets of our loss function are compact, i.e. $V_{\alpha}=\{W\in\mathbb{R}^{d\times d}\,|\,L(W)\leq\alpha\}$ is compact for all $\alpha$. We prove it by contradiction. Suppose $V_{\alpha}$ is not compact for some $\alpha$. Since our loss function is continuous and convex, $V_{\alpha}$ is an unbounded convex set.
Since the dimension of $V_{\alpha}$ is $d^{2}$, consider a point $W^{\alpha}\in V_{\alpha}$, there must exists a $W^{k}\neq 0_{d\times d}$ such that $\{W^{\alpha}+tW^{k}\,|\,t=[0,\infty)\}\in V_{\alpha}$. For this $W^{k}\neq 0_{d\times d}$, there must exist a set of constants $0<c_{3pi}<c_{4pi},0<c_{3qj}<c_{4qj}$ such that for any $c_{pi}\in[c_{3pi},c_{4pi}],c_{qj}\in[c_{3qj},c_{4qj}]$, we have

$$ $\displaystyle|\sum_{ij}c_{pi}c_{qj}W_{ij}^{k}|\neq 0.$ $$

Thus, we have

$$ $\displaystyle\lim_{t\to\infty}|\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{ k})|=\infty.$ $$

We define $\Omega_{0}=\{p=\sum_{i}c_{pi}\textbf{e}_{i}$, $q=\sum_{j}c_{qj}\textbf{e}_{j}$, $c_{pi}\in[c_{3pi},c_{4pi}],c_{qj}\in[c_{3qj},c_{4qj}],\|\mu\|_{2}^{2}\leq\sum_
{i}c_{4pi}^{2}+c_{4qj}^{2}\}$. Then, defining

$$ $\displaystyle C(\Omega_{0})=\int_{\Omega_{0}}f_{pq}(p,q)dpdq,$ $\displaystyle S(\Omega_{0})=\min_{\Omega_{0}}\left\{\min\{\sigma(\mu^{\top} \Lambda^{-1}q),(1-\sigma(\mu^{\top}\Lambda^{-1}q))\}\right\},$ $$

we have $S(\Omega_{0})>0$. Since $\Omega_{0}$ are non-zero measures for $p,q$, we have $C(\Omega_{0})>0$.
Then, we have

$$ $\displaystyle\lim_{t\to\infty}L(W^{\alpha}+tW^{k})$ $\displaystyle=$ $\displaystyle\lim_{t\to\infty}\mathbb{E}[-\sigma(\mu^{\top}\Lambda^{-1}q)\log( \sigma(p^{\top}(W^{\alpha}+tW^{k})q/2))-(1-\sigma(\mu^{\top}\Lambda^{-1}q)) \log(1-\sigma(p^{\top}(W^{\alpha}+tW^{k})q/2))]$ $\displaystyle\geq$ $\displaystyle\lim_{t\to\infty}\int_{\Omega_{0}}[-\sigma(\mu^{\top}\Lambda^{-1} q)\log(\sigma(\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/2))]f_{pq}(p, q)dpdq$ $\displaystyle+\lim_{t\to\infty}\int_{\Omega_{0}}[-(1-\sigma(\mu^{\top}\Lambda^ {-1}q))\log(1-\sigma(\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/2))]f_ {pq}(p,q)dpdq$ $\displaystyle\geq$ $\displaystyle C(\Omega_{0})S(\Omega_{0})\cdot\min_{\Omega_{0}}\left\{\lim_{t \to\infty}[-\log(\sigma(\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/2)) ]\right\}$ $\displaystyle+$ $\displaystyle C(\Omega_{0})S(\Omega_{0})\cdot\min_{\Omega_{0}}\left\{\lim_{t \to\infty}[-\log(1-\sigma(\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/2 ))]\right\}$ $\displaystyle=$ $\displaystyle\infty.$ $$

This contradicts the assumption $L(W^{\alpha}+tW^{k})\leq\alpha$. 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 $\mathbb{R}^{d\times d}$.

Similarly, we can prove the same conclusions for $\widetilde{L}(W)$.
$\blacksquare$

###### Lemma D.3

Denoting the global minimizer of the loss function (7) as $W^{*}$, we have $W^{*}=2(\Lambda^{-1}+G)$, where $\|G\|_{\max}=O(N^{-1/2})$.

Proof
Let $a=\mu^{\top}\Lambda^{-1}q$, $s=\mu^{\top}Wq/2$, $r=(hu+g)^{\top}Wq/2$. Performing the Taylor expansion on (7), we have

$$ $\displaystyle L(W)=$ $\displaystyle\mathbb{E}\left[-\sigma(a)\log(\sigma(s+r))-(1-\sigma(a))\log(1- \sigma(s+r))\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[-\sigma(a)\log(\sigma(s))-(1-\sigma(a))\log(1- \sigma(s))\right]-\mathbb{E}\left[\left(\sigma(a)(1-\sigma(s))-(1-\sigma(a)) \sigma(s))\right)r\right]$ $\displaystyle+\mathbb{E}\left[\sigma(\xi(s,r))(1-\sigma(\xi(s,r)))r^{2}/2\right]$ $\displaystyle=$ $\displaystyle\widetilde{L}(W)-\mathbb{E}\left[\left(\sigma(a)(1-\sigma(s))-(1- \sigma(a))\sigma(s))\right)r\right]+\mathbb{E}\left[\sigma(\xi(s,r))(1-\sigma( \xi(s,r)))r^{2}/2\right],$ $$

where $\xi(s,r)$ are real numbers between $s$ and $s+r$. According to Lemma D.1, we have $\mathbb{E}\left[r\right]=\mathbb{E}\left[(hu+g)^{\top}Wq/2\right]=0$. Thus, we have

$$ $\displaystyle\mathbb{E}\left[\left(\sigma(a)(1-\sigma(s))-(1-\sigma(a))\sigma( s))\right)r\right]=\mathbb{E}_{\mu,u,q}\left[\left(\sigma(a)(1-\sigma(s))-(1- \sigma(a))\sigma(s))\right)\mathbb{E}_{g,h}\left[r\right]\right]=0.$ $$

Moreover, we have

$$ $\displaystyle\mathbb{E}\left[\sigma(\xi(s,r))(1-\sigma(\xi(s,r))r^{2}/2\right]$ $\displaystyle\leq\mathbb{E}\left[r^{2}\right]$ $\displaystyle=\mathbb{E}[h^{2}u^{\top}Wqu^{\top}Wq+g^{\top}Wqg^{\top}Wq]$ $\displaystyle\overset{(a)}{=}\mathbb{E}[u^{\top}Wqu^{\top}Wq/(4N)+4(\Lambda Wq )^{\top}Wq/N]$ $\displaystyle\leq C_{l}\|W\|_{\max}^{2}/N,$ $$

where $(a)$ is due to Lemma D.1, $g^{\top}Wqg^{\top}Wq=\sum_{i,j,k,l\in[d]}g_{i}W_{ij}q_{j}g_{k}W_{kl}q_{l}=\sum
_{i,j,k,l\in[d]}g_{i}g_{k}W_{kl}q_{l}W_{ij}q_{j}=(gg^{\top}Wq)^{\top}Wq$ and $\mathbb{E}[gg^{\top}]=4\Lambda/N$. $C_{l}$ is a constant independent of $N$ and $W$. Thus, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq C_{l}\|W\|_{\max}^{2}/N.$ $$

This shows that $L(W)$ point wisely converges to $\widetilde{L}(W)$.

According to Lemma D.2, $\widetilde{L}(W)$ has one unique global minimizer. Consider the equation:

$$ $\displaystyle\nabla\widetilde{L}(W)=\mathbb{E}[\sigma(\mu^{\top}Wq/2)-\sigma( \mu^{\top}\Lambda^{-1}q)]=0.$ $$

We can easily find that $\nabla\widetilde{L}(2\Lambda^{-1})=0$ and $W=2\Lambda^{-1}$ is the global minimizer of $\widetilde{L}(W)$.

Considering a compact set $R_{W}=\{W\,|\,\|W-2\Lambda^{-1}\|_{F}\leq\rho_{W}\}$, we have $\|W\|_{\max}\leq C_{W}$ for $W\in R_{W}$. Here $\rho_{W},C_{W}$ are some positive finite constants. Then, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq C^{\prime}_{l}/N,\,W\in R_ {W},$ $$

where $C^{\prime}_{l}=C_{l}C_{W}^{2}$ is a constant independent of $N$ and $W$. This shows that, for $W\in R_{W}$, our loss function $L(W)$ uniformly converge to $\widetilde{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^{*}\in R_{W}$. We first denote $\partial R_{W}=\{W\,|\,\|W-2\Lambda^{-1}\|_{F}=\rho_{W}\}$ and $\Delta=\min_{W\in\partial R_{W}}\widetilde{L}(W)-\widetilde{L}(2\Lambda^{-1})>0$. Then, for $N\geq 4C^{\prime}_{l}/\Delta$, and for any $W\in R_{W}$, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq\Delta/4,$ $$

This means

$$ $\displaystyle\min_{W\in\partial R_{W}}L(W)-\min_{W\in R_{W}}L(W)$ $\displaystyle\geq$ $\displaystyle\min_{W\in\partial R_{W}}L(W)-L(2\Lambda^{-1})$ $\displaystyle\geq$ $\displaystyle\min_{W\in\partial R_{W}}\widetilde{L}(W)-\widetilde{L}(2\Lambda^ {-1})-\Delta/2$ $\displaystyle\geq$ $\displaystyle\Delta/2$ $\displaystyle>$ $\displaystyle 0.$ $$

Since $L(W)$ is strictly convex, we have $W^{*}=\mathop{\rm argmin}_{W}L(W)\in R_{W}$.

Then, we have

$$ $\displaystyle|\widetilde{L}(W^{*})-L(W^{*})|\leq C^{\prime}_{l}/N$ $\displaystyle|\widetilde{L}(2\Lambda^{-1})-L(2\Lambda^{-1})|\leq C^{\prime}_{l }/N$ $$

and

$$ $\displaystyle\widetilde{L}(W^{*})\leq L(W^{*})+C^{\prime}_{l}/N\leq L(2\Lambda ^{-1})+C^{\prime}_{l}/N\leq\widetilde{L}(2\Lambda^{-1})+2C^{\prime}_{l}/N.$ $$

According to Lemma D.2, for $W\in R_{W}$, we have $\nabla^{2}\widetilde{L}(W)\succ\gamma I_{d}$, where $\gamma$ is a positive constant independent of $N$. Thus, $\widetilde{L}(W)$ is $\gamma$-strongly convex in $R_{W}$. According to Lemma C.1, we have

$$ $\displaystyle\|W^{*}-2\Lambda^{-1}\|_{F}^{2}\leq\frac{2}{\gamma}(\widetilde{L} (W^{*})-\widetilde{L}(2\Lambda^{-1}))\leq\frac{4C^{\prime}_{l}}{\gamma N}.$ $$

Thus, when $N\to\infty$, we have $W^{*}\to 2\Lambda^{-1}$. Denoting $W^{*}=2(\Lambda^{-1}+G)$, we have $\|G\|_{\max}=O(1/\sqrt{N})$.
$\blacksquare$

###### Lemma D.4

The global minimizer of the loss function (7) is $W^{*}=2(\Lambda^{-1}+G)$, where

$$ $\displaystyle\|G\|_{\max}\leq$ $\displaystyle\frac{1}{N}\|S^{-1}(\mathbb{E}[\sigma^{\prime}(a)(4qq^{\top}+uu^{ \top}\Lambda^{-1}qq^{\top}/4)$ $\displaystyle+\sigma^{\prime\prime}(a)((u^{\top}\Lambda^{-1}q)^{2}\mu q^{\top} /8+2q^{\top}\Lambda^{-1}q\mu q^{\top})])\|_{\max}+o(1/N),$ $$

$a=\mu^{\top}\Lambda^{-1}q$, $S=4\nabla^{2}\widetilde{L}(2\Lambda^{-1})$.

Proof
According to Lemma D.2, the loss function $L(W)$ has a unique global minimizer $W^{*}$. We have

$$ $\displaystyle\nabla L(W^{*})=\mathbb{E}\left[(\sigma(p^{\top}W^{*}q/2)-\sigma( \mu^{\top}\Lambda^{-1}q))pq^{\top}\right]=0.$ (27) $$

Let $W^{*}=2(\Lambda^{-1}+G)$, $a=\mu^{\top}\Lambda^{-1}q$, $b=(\mu+hu+g)^{\top}Gq+(hu+g)^{\top}\Lambda^{-1}q$. We have

$$ $\displaystyle p^{\top}W^{*}q/2$ $\displaystyle=$ $\displaystyle(\mu+hu+g)^{\top}(\Lambda^{-1}+G)q$ $\displaystyle=$ $\displaystyle(\mu+hu+g)^{\top}Gq+(hu+g)^{\top}\Lambda^{-1}q+\mu^{\top}\Lambda^ {-1}q$ $\displaystyle=$ $\displaystyle a+b.$ $$

The Taylor expansion of $\sigma(a+b)$ at point $a$ with an Lagrange form of remainder is

$$ $\displaystyle\sigma(a+b)pq^{\top}=\sigma(a)pq^{\top}+\sigma^{\prime}(a)bpq^{ \top}+\frac{\sigma^{\prime\prime}(a)}{2}b^{2}pq^{\top}+\frac{\sigma^{\prime \prime\prime}(\xi(a,b))}{3!}b^{3}pq^{\top},$ $$

where $\xi(a,b)$ are real numbers between $a$ and $a+b$.
Thus, our equation (27) become

$$ $\displaystyle\mathbb{E}_{\mu,u,g,h,q}\left[\sigma^{\prime}(a)bpq^{\top}+\frac{ \sigma^{\prime\prime}(a)}{2}b^{2}pq^{\top}+\frac{\sigma^{\prime\prime\prime}( \xi(a,b))}{3!}b^{3}pq^{\top}\right]=0.$ (28) $$

Note that $\mathbb{E}[\sigma^{\prime}(a)bpq^{\top}]=\mathbb{E}_{\mu,u,q}\left[\sigma^{
\prime}(a)\mathbb{E}_{g,h}\left[bpq^{\top}\right]\right]$. For $\mathbb{E}_{g,h}[bpq^{\top}]$, according to Lemma D.1 and $g\sim\mathsf{N}(0,4\Lambda/N)$, we have

$$ $\displaystyle\mathbb{E}_{g,h}[bpq^{\top}]$ $\displaystyle=\mathbb{E}[\mu^{\top}Gq\mu q^{\top}+g^{\top}\Lambda^{-1}qgq^{ \top}+g^{\top}Gqgq^{\top}+h^{2}u^{\top}Gquq^{\top}+h^{2}u^{\top}\Lambda^{-1} quq^{\top}]$ $\displaystyle=\mu\mu^{\top}Gqq^{\top}+4qq^{\top}/N+4\Lambda Gqq^{\top}/N+uu^{ \top}Gqq^{\top}/(4N)+uu^{\top}\Lambda^{-1}qq^{\top}/(4N).$ (29) $$

Then, we have

$$ $\displaystyle\|\mathbb{E}_{\mu,u,q}[\sigma^{\prime}(a)(4\Lambda Gqq^{\top}/N+ uu^{\top}Gqq^{\top}/(4N))]\|_{\max}\leq c_{1}\|G\|_{\max}/N,$ $$

where $c_{1}=\max_{ij}\left|\mathbb{E}\left[\sum_{kl}4\sigma^{\prime}(a)\left(\Lambda
_{ik}q_{l}q_{j}\right)+\sum_{kl}\sigma^{\prime}(a)\left(u_{i}u_{k}q_{l}q_{j}/4
\right)\right]\right|$ is a constant independent of $N$. According to Lemma D.3, $\|G\|_{\max}=O(1/\sqrt{N})=o(1)$, we have

$$ $\displaystyle\|\mathbb{E}_{\mu,u,q}[\sigma^{\prime}(a)(4\Lambda Gqq^{\top}/N+ uu^{\top}Gqq^{\top}/(4N))]\|_{\max}=o(1/N),$ (30) $$

Similarly for $\mathbb{E}[\sigma^{\prime\prime}(a)b^{2}pq^{\top}/2]$, we have

$$ $\displaystyle\mathbb{E}_{g,h}[b^{2}pq^{\top}]$ $\displaystyle=\underbrace{\mathbb{E}[\mu^{\top}Gq\mu^{\top}Gq\mu q^{\top}+h^{2 }u^{\top}Gqu^{\top}Gq\mu q^{\top}+g^{\top}Gqg^{\top}Gq\mu q^{\top}+2h^{2}u^{ \top}Gq\mu^{\top}Gquq^{\top}+2g^{\top}Gq\mu^{\top}Gqgq^{\top}]}_{(i)}$ $\displaystyle+\underbrace{\mathbb{E}[2h^{2}u^{\top}Gqu^{\top}\Lambda^{-1}q\mu q ^{\top}+2g^{\top}Gqg^{\top}\Lambda^{-1}q\mu q^{\top}+2h^{2}\mu^{\top}Gqu^{\top }\Lambda^{-1}quq^{\top}+2\mu^{\top}Gqg^{\top}\Lambda^{-1}qgq^{\top}]}_{(ii)}$ $\displaystyle+\underbrace{\mathbb{E}[h^{2}u^{\top}\Lambda^{-1}qu^{\top}\Lambda ^{-1}q\mu q^{\top}+g^{\top}\Lambda^{-1}qg^{\top}\Lambda^{-1}q\mu q^{\top}]}_{( iii)}.$ $$

Each term in $(i)$ contains two $G$’s. Thus, their max norms are at most $O(\|G\|_{\max}^{2})$. For each term in $(ii)$, it contains one $G$ and $h^{2}$ or contains one $G$ and two $g$.
According to $\mathbb{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 $\bar{g}=N^{1/2}\Lambda^{-1/2}g/2$, we have $\bar{g}\sim\mathsf{N}(0,I_{d})$ and $g=2N^{-1/2}\Lambda^{1/2}\bar{g}$. Thus, converting two $g$ to $\bar{g}$, 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

$$ $\displaystyle\|\mathbb{E}[\sigma^{\prime\prime}(a)(i)/2]\|_{\max}\leq O(\|G\|_ {\max}^{2})=o(\|G\|_{\max}),$ (31) $\displaystyle\|\mathbb{E}[\sigma^{\prime\prime}(a)(ii)/2]\|_{\max}\leq O(\|G\| _{\max}/N)=o(1/N).$ (32) $$

For term $(iii)$, according to Lemma D.1 and $g\sim\mathsf{N}(0,4\Lambda/N)$, we have

$$ $\displaystyle\|\mathbb{E}[\sigma^{\prime\prime}(a)(iii)/2]\|_{\max}$ $\displaystyle=$ $\displaystyle\|\mathbb{E}\left[\sigma^{\prime\prime}(a)(h^{2}u^{\top}\Lambda^{ -1}qu^{\top}\Lambda^{-1}q\mu q^{\top}+g^{\top}\Lambda^{-1}qg^{\top}\Lambda^{-1 }q\mu q^{\top})/2\right]\|_{\max}$ (33) $\displaystyle=$ $\displaystyle\frac{1}{N}\|\mathbb{E}\left[\sigma^{\prime\prime}(a)((u^{\top} \Lambda^{-1}q)^{2}\mu q^{\top}/8+2q^{\top}\Lambda^{-1}q\mu q^{\top})\right]\|_ {\max}.$ (34) $$

For $\mathbb{E}[\sigma^{\prime\prime\prime}(\xi(a,b))b^{3}pq^{\top}/3!]$, we have

$$ $\displaystyle\|\mathbb{E}[\sigma^{\prime\prime\prime}(\xi(a,b))b^{3}pq^{\top}/ 3!]\|_{\max}\leq\max_{z\in\mathbb{R}}|\sigma^{\prime\prime\prime}(z)|/3!\cdot \max_{ij}\mathbb{E}\left[\left|b^{3}p_{i}q_{j}\right|\right]$ $\displaystyle\leq$ $\displaystyle O(1)\cdot\max_{ij}\mathbb{E}\bigg{[}\underbrace{\sum_{\phi_{1}, \phi_{2},\phi_{3}\in\{\mu,hu,g\}}\left|\phi_{1}^{\top}Gq\phi_{2}^{\top}Gq\phi_ {3}^{\top}Gqp_{i}q_{j}\right|}_{(*)}$ $\displaystyle+\underbrace{\sum_{\phi_{1},\phi_{2}\in\{\mu,hu,g\},\phi_{3}\in\{ hu,g\}}\left|\phi_{1}^{\top}Gq\phi_{2}^{\top}Gq\phi_{3}^{\top}\Lambda^{-1}qp_{ i}q_{j}\right|}_{(*)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{\mu,hu,g\},\phi_{2},\phi_{3}\in\{ hu,g\}}\left|\phi_{1}^{\top}Gq\phi_{2}^{\top}\Lambda^{-1}q\phi_{3}^{\top} \Lambda^{-1}qp_{i}q_{j}\right|}_{(**)}$ $\displaystyle+\underbrace{\sum_{\phi_{1},\phi_{2},\phi_{3}\in\{hu,g\}}\left| \phi_{1}^{\top}\Lambda^{-1}q\phi_{2}^{\top}\Lambda^{-1}q\phi_{3}^{\top}\Lambda ^{-1}qp_{i}q_{j}\right|}_{(***)}\bigg{]}.$ $$

For terms in $(*)$ containing two or three $G$, these terms’ expected absolute values are at most smaller than $O(\|G\|_{\max}^{2})$. 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}\in\mathbb{N}$. According to Lemma D.1, we know that for $n_{1}=1,2,3,4$, $\mathbb{E}|h^{n_{1}}|\leq O(N^{-n_{1}/2})$. Defining $\bar{g}=N^{1/2}\Lambda^{-1/2}g/2$, we have $\bar{g}\sim\mathsf{N}(0,I_{d})$ and $g=2N^{-1/2}\Lambda^{1/2}\bar{g}$. Converting $g$ to $\bar{g}$, we have a coefficient of $N^{-n_{2}/2}$. Thus, for terms in $(**)$, these terms’ expected absolute values are at most smaller than $O(\|G\|_{\max}N^{-(n_{1}+n_{2})/2})\leq 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}\in\mathbb{N}$. Similarly, these term’s expected absolute values are at most smaller than $O(N^{-(n_{1}+n_{2})/2})\leq O(N^{-3/2})$. Therefore, we have

$$ $\displaystyle\|\mathbb{E}[\sigma^{\prime\prime\prime}(\xi(a,b))b^{3}pq^{\top}/ 3!]\|_{\max}$ $\displaystyle\leq\max_{ij}\mathbb{E}\left[\left|b^{3}p_{i}q_{j}\right|\right] \cdot\max_{z}|\sigma^{\prime\prime\prime}(z)|/3!$ $\displaystyle=O(\|G\|_{\max}^{2})+O(\|G\|_{\max}/N)+O(1/N^{-3/2})$ $\displaystyle=o(\|G\|_{\max})+o(1/N).$ (35) $$

Moreover, we have

$$ $\displaystyle\bigg{\{}\mathbb{E}_{\mu,u,q}[\sigma^{\prime}(a)\mu\mu^{\top}Gqq^ {\top}]\bigg{\}}_{ij}=\sum_{kl}s_{ijkl}G_{kl},$ (36) $$

where $s_{ijkl}=\mathbb{E}\sigma^{\prime}(a)\mu_{i}\mu_{k}q_{l}q_{j}$. We vectorize $G$ as $\text{Vec}(G)_{i}=G_{t_{1}(i),t_{2}(i)}$. Define $S\in\mathbb{R}^{d^{2}\times d^{2}}$, where $S_{ij}=s_{t_{1}(i),t_{2}(i),t_{1}(j),t_{2}(j)}=\mathbb{E}\sigma^{\prime}(a)\mu
_{t_{1}(i)}q_{t_{2}(i)}\mu_{t_{1}(j)}q_{t_{2}(j)}$. Then
(36) can be expressed as

$$ $\displaystyle\bigg{\{}\mathbb{E}_{\mu,v}[\sigma^{\prime}(a)\mu\mu^{\top}Gqq^{ \top}]\bigg{\}}=SG.$ (37) $$

Note that $S=4\nabla^{2}\widetilde{L}(2\Lambda^{-1})$. According to Lemma D.2, $S$ is positive definite. Thus,
combining (28), (29), (30), (31), (32), (34), (35), (37), we have

$$ $\displaystyle\|G\|_{\max}\leq\frac{1}{N}\|S^{-1}\left(\mathbb{E}[\sigma^{ \prime}(a)(4qq^{\top}+uu^{\top}\Lambda^{-1}qq^{\top}/4)+\sigma^{\prime\prime}( a)((u^{\top}\Lambda^{-1}q)^{2}\mu q^{\top}/8+2q^{\top}\Lambda^{-1}q\mu q^{\top })]\right)\|_{\max}+o(1/N).$ $$

$\blacksquare$

###### Lemma D.5

The loss function (7) is $l$-smooth, where $l\leq\frac{1}{4}\sum_{i\in[d^{2}]}\mathbb{E}[(p_{t_{1}(i)}q_{t_{2}(i)})^{2}]$.

Proof
The Hessian matrix of the loss function is

$$ $\displaystyle(\nabla^{2}L(W))_{ij}=\frac{1}{4}\mathbb{E}[\sigma(p^{\top}Wq/2)( 1-\sigma(p^{\top}Wq/2))p_{t_{1}(i)}q_{t_{2}(i)}p_{t_{1}(j)}q_{t_{2}(j)}].$ $$

Considering $z\in\mathbb{R}^{d^{2}}$ such that $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z=$ $\displaystyle\mathbb{E}\left[\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top} Wq/2))\sum_{ab}z_{a}z_{b}p_{t_{1}(a)}q_{t_{2}(a)}p_{t_{1}(b)}q_{t_{2}(b)}\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[\frac{1}{4}\sigma(p^{\top}Wq/2)(1-\sigma(p^{\top} Wq/2))\left(\sum_{a\in[d^{2}]}z_{a}p_{t_{1}(a)}q_{t_{2}(a)}\right)^{2}\right]$ $\displaystyle\leq$ $\displaystyle\mathbb{E}\left[\frac{1}{4}\left(\sum_{a\in[d^{2}]}z_{a}p_{t_{1}( a)}q_{t_{2}(a)}\right)^{2}\right]$ $\displaystyle\overset{(a)}{\leq}$ $\displaystyle\frac{1}{4}\|z\|_{2}^{2}\sum_{i\in[d^{2}]}\mathbb{E}[(p_{t_{1}(i) }q_{t_{2}(i)})^{2}]$ $$

where $(a)$ is due to the Cauchy–Schwarz inequality.
Thus, $\nabla^{2}L(W)\preceq lI_{d}$ and $L(W)$ is $l$-smooth, where $l$ is a constant smaller than $\frac{1}{4}\sum_{i\in[d^{2}]}\mathbb{E}[(p_{t_{1}(i)}q_{t_{2}(i)})^{2}]$.
$\blacksquare$

###### Lemma D.1

###### Lemma D.2

###### Lemma D.3

###### Lemma D.4

###### Lemma D.5

### D.4 Proof of Theorem 3.1

Proof
According to Lemma D.4, the global minimizer of $L(W)$ is $W^{*}=2(\Lambda^{-1}+G)$, where

$$ $\displaystyle\|G\|_{\max}\leq$ $\displaystyle\frac{1}{N}\|S^{-1}(\mathbb{E}[\sigma^{\prime}(a)(4qq^{\top}+ \frac{1}{4}uu^{\top}\Lambda^{-1}qq^{\top})$ $\displaystyle+\sigma^{\prime\prime}(a)(\frac{1}{8}(u^{\top}\Lambda^{-1}q)^{2} \mu q^{\top}+2q^{\top}\Lambda^{-1}q\mu q^{\top})])\|_{\max}+o(1/N).$ (38) $$

Define $R_{W}=\{W\in\mathbb{R}^{d\times d}\,|\,\|W-W^{*}\|_{F}\leq\|W^{0}-W^{*}\|_{F}\}$. $R_{W}$ is a compact set. Then, according to Lemma D.2, for $W\in R_{W}$, we have $\nabla^{2}L(W)\succeq\alpha I_{d}$. Here $\alpha>0$ is a positive constant number. Thus, $L(W)$ is $\alpha$-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 $\eta=1/l$, for any $t\geq 1$, we have

$$ $\displaystyle\|W^{t}-W^{*}\|_{F}^{2}\leq\exp(-t/\kappa)\cdot\|W^{0}-W^{*}\|_{F }^{2},$ $$

where $\kappa=l/\alpha$.
$\blacksquare$

## Appendix E In-context inference of binary classification

### E.1 Notations

In this section, we use the following notations. We denote
$\mu=\mu_{1}-\mu_{0}$,
$u=2(\mu_{1}+\mu_{0})$,
$q=x_{\mathsf{query}}$. Define $p=\frac{2}{M}\sum_{i=1}^{M}y_{i}x_{i}$. Since with probability $\mathbb{P}\left(y_{i}=1\right)=1/2$, $x_{i}=\mu_{1}+v_{i}$, with probability $\mathbb{P}\left(y_{i}=0\right)=1/2$, $x_{i}=\mu_{0}+v_{i}$, where $v_{i}\sim\mathsf{N}(0,\Lambda)$, we have $p=2M_{1}\mu_{1}/M-2M_{0}\mu_{0}/M+g$,
where $g=\frac{2}{M}\sum_{i=1}^{M}v_{i}$, $g\sim\mathsf{N}(0,4\Lambda/M)$, $M_{1}\sim\mathsf{Bin}(M,1/2)$.
Defining $h=M_{1}/N-1/2$, $u=2(\mu_{1}+\mu_{0})$, we have $M_{0}/N=1/2-h$ and

$$ $\displaystyle p=\mu+hu+g.$ (39) $$

### E.2 Proof of Theorem 3.2

Proof
The output of the trained transformer is

$$ $\displaystyle\widehat{y}_{\mathsf{out}}=\sigma\left(\left(\frac{2}{M}\sum_{i=1 }^{M}y_{i}x_{i}^{\top}\right)(\Lambda^{-1}+\widehat{G})x_{\mathsf{query}} \right)=\sigma(p^{\top}(\Lambda^{-1}+\widehat{G})q).$ (40) $$

The probability of $y_{\mathsf{query}}=1$ given $x_{\mathsf{query}}$ is

$$ $\mathbb{P}\left(y_{\mathsf{query}}=1|x_{\mathsf{query}}\right)=\sigma((\mu_{1} -\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{query}})=\sigma(\mu^{\top}\Lambda^{-1}q).$ $$

Defining $a=\mu^{\top}\Lambda^{-1}q$, $b=(\mu+hu+g)^{\top}\widehat{G}q+(hu+g)^{\top}\Lambda^{-1}q$, we have

$$ $\displaystyle p^{\top}(\Lambda^{-1}+\widehat{G})q$ $\displaystyle=$ $\displaystyle(\mu+hu+g)^{\top}(\Lambda^{-1}+\widehat{G})q$ $\displaystyle=$ $\displaystyle(\mu+hu+g)^{\top}\widehat{G}q+(hu+g)^{\top}\Lambda^{-1}q+\mu^{ \top}\Lambda^{-1}q=a+b,$ $$

and

$$ $\displaystyle\mathbb{E}\left[\sigma(p^{\top}(\Lambda^{-1}+\widehat{G})q)\right ]=\mathbb{E}\left[\sigma(a+b)\right]=\mathbb{E}[\sigma(a)+\sigma^{\prime}(a)b+ \sigma^{\prime\prime}(\xi(a,b))b^{2}/2],$ $$

where $\xi$ are real numbers between $a$ and $a+b$.
Thus, we have

$$ $\displaystyle\mathbb{E}[\left|\sigma(a+b)-\sigma(a)\right|]$ $\displaystyle\leq$ $\displaystyle\mathbb{E}[\left|\sigma^{\prime}(a)b+\sigma^{\prime\prime}(\xi(a, b))b^{2}/2\right|]$ $\displaystyle\leq$ $\displaystyle\sigma^{\prime}(a)\mathbb{E}[\left|b\right|]+\mathbb{E}[b^{2}]$ $$

We first consider the term $\sigma^{\prime}(a)\mathbb{E}[\left|b\right|]$.
Defining $\bar{g}=\Lambda^{-1/2}M^{1/2}g/2$, we have

$$ $\displaystyle\sigma^{\prime}(a)\mathbb{E}[\left|b\right|]$ $\displaystyle\leq$ $\displaystyle\sigma^{\prime}(a)\left[|\mu^{\top}\widehat{G}q|+\mathbb{E}[|hu^{ \top}\widehat{G}q|]+\mathbb{E}[|g^{\top}\widehat{G}q|]+\mathbb{E}[|hu^{\top} \Lambda^{-1}q|]+\mathbb{E}[|g^{\top}\Lambda^{-1}q|]\right]$ $\displaystyle\overset{(a)}{\leq}$ $\displaystyle\sigma^{\prime}(a)\left[|\mu^{\top}\widehat{G}q|+\frac{1}{2M^{1/2 }}|u^{\top}\widehat{G}q|+\frac{2}{M^{1/2}}\mathbb{E}[|\bar{g}^{\top}\Lambda^{1 /2}\widehat{G}q|]+\frac{1}{2M^{1/2}}|u^{\top}\Lambda^{-1}q|+\frac{2}{M^{1/2}} \mathbb{E}[|\bar{g}^{\top}\Lambda^{-1/2}q|]\right]$ $\displaystyle\overset{(b)}{\leq}$ $\displaystyle\sigma^{\prime}(a)\left[\|\widehat{G}\|_{\max}\sum_{i,j\in[d]}| \mu_{i}q_{j}|+\frac{1}{M^{1/2}}\left(\frac{1}{2}|u^{\top}\Lambda^{-1}q|+\frac{ 2\sqrt{2}}{\sqrt{\pi}}\sum_{i,j\in[d]}|\Lambda^{-1/2}_{ij}q_{j}|\right)\right] +o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right),$ $$

where $(a)$ is due to $\mathbb{E}[|h|]\leq 1/(2M^{1/2})$ in Lemma D.1. $(b)$ is because that $\bar{g}_{i}\sim\mathsf{N}(0,1)$ and $\mathbb{E}[|\bar{g}_{i}|]=\sqrt{2}/\sqrt{\pi}$, for $i\in[d]$.

For $\mathbb{E}[b^{2}]$, we have

$$ $\displaystyle\mathbb{E}[b^{2}]\leq\mathbb{E}\Big{[}[(\mu+hu+g)^{\top}\widehat{ G}q]^{2}\Big{]}+\mathbb{E}\Big{[}[(hu+g)^{\top}\Lambda^{-1}q]^{2}\Big{]}+2 \mathbb{E}\Big{[}(\mu+hu+g)^{\top}\widehat{G}q(hu+g)^{\top}\Lambda^{-1}q\Big{]}.$ $$

Notice that terms in $\mathbb{E}\Big{[}[(\mu+hu+g)^{\top}\widehat{G}q]^{2}\Big{]}$ contain two $\widehat{G}$. Thus, they are at most smaller than $O(\|\widehat{G}\|^{2}_{\max})=O(1/N^{2})$. Terms in $\mathbb{E}\Big{[}[(hu+g)^{\top}\Lambda^{-1}q]^{2}\Big{]}/2$ contain two $h$, or two $g$, or one $h$ and one $g$. According to Lemma D.1, we have $\mathbb{E}[|h|]=O(1/\sqrt{M})$, $\mathbb{E}[h^{2}]=1/(4M)$. Moreover, $g=2M^{-1/2}\Lambda^{1/2}\bar{g}$. Converting one $g$ to $\bar{g}$, we have a coefficient of $M^{-1/2}$. Thus, terms in $E\Big{[}[(hu+g)^{\top}\Lambda^{-1}q]^{2}\Big{]}/2$ contain two $h$, or two $g$, or one $h$ and one $g$ are $O(1/M)$. Terms in $E\Big{[}(\mu+hu+g)^{\top}\widehat{G}q(hu+g)^{\top}\Lambda^{-1}q\Big{]}$ contain at least one $\widehat{G}$ and one $h$ or one $\widehat{G}$ and one $g$. Thus, they are at most smaller than $O(\|\widehat{G}\|_{\max}/\sqrt{M})=O(1/(N\sqrt{M}))$.
Therefore, we have $E[b^{2}|]/2=O(1/N^{2}+1/M+1/(N\sqrt{M}))=o(1/N+1/\sqrt{M})$.

Finally, we have

$$ $\displaystyle\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}} )]=\mathbb{E}[|\widehat{y}_{\mathsf{out}}-\mathbb{P}\left(y_{\mathsf{query}}=1 |x_{\mathsf{query}}\right)|]=\mathbb{E}[|\sigma(a+b)-\sigma(a)|]\leq\sigma^{ \prime}(a)\mathbb{E}[\left|b\right|]+\mathbb{E}[b^{2}]$ $\displaystyle\leq$ $\displaystyle\sigma^{\prime}(a)\left[\|\widehat{G}\|_{\max}\sum_{i,j\in[d]}| \mu_{i}q_{j}|+\frac{1}{M^{1/2}}\left(\frac{1}{2}|u^{\top}\Lambda^{-1}q|+\frac{ 2\sqrt{2}}{\sqrt{\pi}}\sum_{i,j\in[d]}|\Lambda^{-1/2}_{ij}q_{j}|\right)\right] +o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right).$ $$

$\blacksquare$

###### Remark E.1

We note that Theorem 3.2 requires Assumption 3.2 to hold. For example, we need the covariance $\Lambda$ in training and testing to be the same. A similar consistency requirement of the covariance $\Lambda$ in training and testing had also been observed for in-context linear regression in Zhang et al. [].

Here, we discuss the consequences when Assumption 3.2 does not hold. For example, suppose the labels of our data in test prompts are not balanced where $\mathbb{P}\left(y=1\right)=p_{1},\mathbb{P}\left(y=-1\right)=p_{0}$. Besides, $\mu_{0},\mu_{1}$ do not have the same $\Lambda^{-1}$ weighted norm, and the covariance matrix of test data is $\Gamma\neq\Lambda$. Then, as $N,M\to\infty$, we have

$$ $\frac{2}{M}\sum_{i=1}^{M}y_{i}x_{i}^{\top}\to 2(p_{1}\mu_{1}-p_{0}\mu_{0})^{ \top},$ $$

and

$$ $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=1\right)\to\sigma(2(p_{1}\mu_{1}- p_{0}\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{query}}).$ $$

On the other hand, the distribution of the ground truth label is

$$ $\mathbb{P}\left(y_{\mathsf{query}}=1\right)=\sigma((\mu_{1}-\mu_{0})^{\top} \Gamma^{-1}x_{\mathsf{query}}+(\mu_{1}^{\top}\Lambda^{-1}\mu_{1}-\mu_{0}^{\top }\Lambda^{-1}\mu_{0})/2+\log(p_{1}/p_{0})).$ $$

Define $z\triangleq(\mu_{1}-\mu_{0})^{\top}\Gamma^{-1}x_{\mathsf{query}}+(\mu_{1}^{
\top}\Lambda^{-1}\mu_{1}-\mu_{0}^{\top}\Lambda^{-1}\mu_{0})/2+\log(p_{1}/p_{0})$ and $\hat{z}\triangleq 2(p_{1}\mu_{1}-p_{0}\mu_{0})^{\top}\Lambda^{-1}x_{\mathsf{
query}}$. Then, we can notice that unless $\hat{z}=z$ or $\left|\sigma(\hat{z})-\sigma(z)\right|$ is sufficiently small, the transformer cannot correctly perform the in-context binary classification.

###### Remark E.1

## Appendix F Training procedure for in-context multi-class classification

In this section, we present the proof of Theorem 4.1.

### F.1 Proof sketch

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 $\mathbb{R}^{d\times 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\to\infty$, our loss function $L(W)$ point wisely converges to $\widetilde{L}(W)$ (defined in (44)), and the global minimizer $W^{*}$ converge to $2\Lambda^{-1}$. Thus, we denote $W^{*}=2(\Lambda^{-1}+G)$, and prove $\|G\|_{\max}=O(N^{-1/4})$. Next, in Lemma F.5, by further analyzing the Taylor expansion of the equation $\nabla L(W^{*})=0$ at the point $2\Lambda^{-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 $\alpha$-strongly convex and $l$-smooth. Finally, leveraging the standard results from the convex optimization, we prove Theorem 4.1 in subsection F.3.

In this section, we use the following notations.

### F.2 Notations

Recall the expected loss function (16) is

$$ $\displaystyle L(W)=-\mathbb{E}\left[\sum_{k=1}^{c}(y_{\tau,\mathsf{query}})_{k }\log((\widehat{y}_{\tau,\mathsf{out}})_{k})\right],$ (41) $$

where

$$ $\displaystyle(\widehat{y}_{\tau,\mathsf{out}})_{k}=\mathrm{softmax}\left(\frac {1}{c}\left(\frac{c}{N}\sum_{i=1}^{N}y_{\tau,i}x_{\tau,i}^{\top}\right)Wx_{ \tau,\mathsf{query}}\right)_{k}$ $$

is the output of the transformer, and the label of the data follows the distribution

$$ $\displaystyle\mathbb{P}\left(y_{\tau,\mathsf{query}}=\textbf{e}_{k}|x_{\tau, \mathsf{query}}\right)=\mathrm{softmax}(\mu_{\tau}^{\top}\Lambda^{-1}x_{\tau, \mathsf{query}}))_{k}.$ $$

In this section, we introduce the following notations to analyze (16). We denote $\mu_{k}=\mu_{\tau,k}$, $\mu=(\mu_{1},\mu_{2},\dots,\mu_{k})\in\mathbb{R}^{d\times c}$ and $q=x_{\tau,\mathsf{query}}$. Then with probability $\mathbb{P}\left(y_{\tau,\mathsf{query}}=\textbf{e}_{k}\right)=1/c$, $q=\mu_{k}+v$, where $v\sim\mathsf{N}(0,\Lambda)$.
We define $p_{k}=\frac{c}{N}\sum_{i=1}^{N}(y_{\tau,i})_{k}x_{\tau,i}\in\mathbb{R}^{d}$ and $P=(p_{1},p_{2},\dots,p_{c})\in\mathbb{R}^{d\times c}$. We have $P^{\top}=\frac{c}{N}\sum_{i=1}^{N}y_{i}x_{\tau,i}^{\top}\in\mathbb{R}^{c\times
d}$. Since with probability $\mathbb{P}\left(y_{\tau,i}=\textbf{e}_{k}\right)=1/c$ we have $x_{\tau,i}=\mu_{k}+v_{i}$, where $v_{i}\sim\mathsf{N}(0,\Lambda)$, we known
$p_{k}=\frac{c}{N}\sum_{i=1}^{N}(y_{\tau,i})_{k}x_{\tau,i}=cN_{k}\mu_{k}/N+g_{k}$, where $g_{k}=\frac{c}{N}\sum_{i\in\{i|y_{\tau,i}=\textbf{e}_{k}\}}v_{i}$, $g_{k}\sim\mathsf{N}(0,c^{2}N_{k}\Lambda/N^{2})$ and $(N_{1},N_{2},\dots,N_{c})\sim\mathsf{Multin}(n,1/c)$. Defining $h_{k}=N_{k}/N-1/c$, we have $N_{k}/N=1/c+h_{k}$ and $p_{k}=\mu_{k}+ch_{k}\mu_{k}+g_{k}$. Defining $\bar{g}_{k}=\Lambda^{-1/2}g_{k}$, we have $\bar{g}_{k}\sim\mathsf{N}(0,c^{2}N_{k}I_{d}/N^{2})$. Defining $\mu_{h}=(h_{1}\mu_{1},h_{2}\mu_{2},\dots,h_{k}\mu_{k})\in\mathbb{R}^{d\times c}$ and $g=(g_{1},g_{2},\dots,g_{k})\in\mathbb{R}^{d\times c}$, we have $P=\mu+c\mu_{h}+g$.

Then, the expected loss function (16) can be expressed as

$$ $\displaystyle L(W)=\mathbb{E}\left[\sum_{k=1}^{c}-\mathrm{softmax}(\mu^{\top} \Lambda^{-1}q)_{k}\log(\mathrm{softmax}(P^{\top}Wq/c)_{k})\right].$ (42) $$

The gradient of the loss function (16) can be expressed as

$$ $\displaystyle\nabla L(W)=\mathbb{E}\left[\sum_{k=1}^{c}\left[(\mathrm{softmax} (P^{\top}Wq/c)_{k}-\mathrm{softmax}(\mu^{\top}\Lambda^{-1}q)_{k})p_{k}q^{\top} /c\right]\right].$ (43) $$

Moreover, we define a function $\widetilde{L}(W)$ as

$$ $\displaystyle\widetilde{L}(W)=\mathbb{E}[\sum_{k=1}^{c}-\mathrm{softmax}(\mu^{ \top}\Lambda^{-1}q)_{k}\log(\mathrm{softmax}(\mu^{\top}Wq/c)_{k})].$ (44) $$

In Lemma F.4, we show that as $N\to\infty$, $L(W)$ will point wisely converge to $\widetilde{L}(W)$.

###### Lemma F.1

Suppose $(N_{1},N_{2},\dots,N_{c})\sim\mathsf{Multin}(N,1/c)$. Defining $h_{k}=N_{k}/N-1/c$, we have

$$ $\displaystyle\mathbb{E}[h_{k}]=0$ $\displaystyle\mathbb{E}[h_{k}^{2}]=\frac{1}{N}\left(\frac{1}{c}-\frac{1}{c^{2} }\right)$ $\displaystyle\mathbb{E}[h_{i}h_{j}]=-\frac{1}{Nc^{2}},i\neq j$ $\displaystyle\mathbb{E}[\prod_{k=1}^{c}h_{k}^{n_{k}}]=O\left(N^{-2}\right), \sum_{k}n_{k}\geq 3$ $\displaystyle\mathbb{E}[|h_{j}|]\leq N^{-1/2}c^{-1/2}(1-1/c)^{1/2}$ $\displaystyle\mathbb{E}[|h_{i}h_{j}|]=O(N^{-1})$ $\displaystyle\mathbb{E}[|h_{i}h_{j}h_{k}|]=O\left(N^{-3/2}\right)$ $\displaystyle\mathbb{E}[|h_{i}h_{j}h_{k}h_{l}|]=O\left(N^{-2}\right),$ $$

where $i,j,k,l\in[c]$.

Proof
Since $(N_{1},N_{2},\dots,N_{c})\sim\mathsf{Multin}(N,1/c)$, the moment-generating function of $(N_{1},N_{2},\dots,N_{c})$ is

$$ $\displaystyle M_{N}(t)=\left(\frac{1}{c}\sum_{i=1}^{c}\exp(t_{i})\right)^{N}$ $$

We can compute the moment-generating function of $h=(h_{1},h_{2},\dots,h_{c})$ as follows:

$$ $\displaystyle M_{h}(t)=$ $\displaystyle\exp\left(-\sum_{i=1}^{c}t_{i}/c\right)M_{N}(t/N)=\left(\frac{1}{ c}\sum_{i=1}^{c}\exp\left(\frac{1}{N}\left(t_{i}-\sum_{j=1}^{c}t_{j}/c\right) \right)\right)^{N}$ $\displaystyle=$ $\displaystyle\left[1+\frac{1}{Nc}\left(\sum_{i=1}^{c}t_{i}-c\sum_{j=1}^{c}t_{j }/c\right)+\frac{1}{2N^{2}c}\left(\sum_{i=1}^{c}\left(t_{i}-\sum_{j=1}^{c}t_{j }/c\right)^{2}\right)\right.$ $\displaystyle\left.+\sum_{k=3}^{\infty}\frac{1}{k!N^{k}c}\left(\sum_{i=1}^{c} \left(t_{i}-\sum_{j=1}^{c}t_{j}/c\right)^{k}\right)\right]^{N}$ $\displaystyle=$ $\displaystyle\left[1+\sum_{i=1}^{c}\frac{1}{2N}(1/c-1/c^{2})t_{i}^{2}-\sum_{i \neq j\in[c]}\frac{1}{2Nc^{2}}t_{i}t_{j}+\sum_{k=3}^{\infty}\frac{1}{k!N^{k}c} \left(\sum_{i=1}^{c}\left(t_{i}-\sum_{j=1}^{c}t_{j}/c\right)^{k}\right)\right]$ $$

Observing the coefficients of $h$, we have

$$ $\displaystyle\mathbb{E}[h_{k}]=0$ $\displaystyle\mathbb{E}[h_{k}^{2}]=\frac{1}{N}\left(\frac{1}{c}-\frac{1}{c^{2} }\right)$ $\displaystyle\mathbb{E}[h_{i}h_{j}]=-\frac{1}{Nc^{2}},i\neq j$ $\displaystyle\mathbb{E}[\prod_{k=1}^{c}h_{k}^{n_{k}}]=O\left(N^{-2}\right), \sum_{k}n_{k}\geq 3,$ $$

where $i,j,k\in[c]$.

Iteratively applying the Hölder’s inequality, we have

$$ $\displaystyle\mathbb{E}[|h_{j}|]\leq\left(\mathbb{E}[h_{j}^{2}]\right)^{1/2}=N ^{-1/2}c^{-1/2}(1-1/c)^{1/2}$ $\displaystyle\mathbb{E}[|h_{i}h_{j}|]\leq\left(\mathbb{E}[h_{i}^{2}h_{j}^{2}] \right)^{1/2}=O(N^{-1})$ $\displaystyle\mathbb{E}[|h_{i}|^{3}]\leq\mathbb{E}[|h_{i}|^{4}]^{3/4}=(N^{-3/2})$ $\displaystyle\mathbb{E}[|h_{i}h_{j}h_{k}|]\leq\mathbb{E}[|h_{i}|^{3}]^{1/3} \mathbb{E}[|h_{j}|^{3}]^{1/3}\mathbb{E}[|h_{k}|^{3}]^{1/3}=O\left(N^{-3/2}\right)$ $\displaystyle\mathbb{E}[|h_{i}h_{j}h_{k}h_{l}|]\leq\mathbb{E}[|h_{i}|^{4}]^{1/ 4}\mathbb{E}[|h_{j}|^{4}]^{1/4}\mathbb{E}[|h_{k}|^{4}]^{1/4}\mathbb{E}[|h_{l}| ^{4}]^{1/4}=O\left(N^{-2}\right)$ $$

where $i,j,k,l\in[c]$.
$\blacksquare$

###### Lemma F.2

Suppose $g_{k}\sim\mathsf{N}(0,c^{2}N_{k}\Lambda/N^{2})$ and $(N_{1},N_{2},\dots,N_{c})\sim\mathsf{Multin}(N,1/c)$, define $\bar{g}_{k}=\Lambda^{-1/2}g_{k}$ and $N_{k}/N=1/c+h_{k}$, we have

$$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}]=0$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}(\bar{g}_{l})_{j}]=\delta_{kl}\delta_ {ij}c/N$ $\displaystyle\mathbb{E}[(\bar{g}_{k_{1}})_{i_{i}}(\bar{g}_{k_{2}})_{i_{2}}( \bar{g}_{k_{3}})_{i_{3}}]=0$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}^{4}]=\mathbb{E}[3c^{2}/N^{2}(1+ch_{k })^{2}]=O(N^{-2})$ $\displaystyle\mathbb{E}[h_{m}(\bar{g}_{k})_{i}(\bar{g}_{l})_{j}]=\mathbb{E}[c^ {2}\delta_{kl}\delta_{ij}h_{m}h_{k}/N]=O(N^{-2})$ $\displaystyle\mathbb{E}[h_{m}h_{l}(\bar{g}_{k})_{i}]=0$ $$

where $i,j,i_{1},i_{2},i_{3}\in[d],k,l,m,k_{1},k_{2},k_{3}\in[c]$.

For any $n_{1k}$, $n_{2ki}$ satisfying $\sum_{k\in[c]}n_{1k}+\sum_{k\in[c],i\in[d]}n_{2ki}=1,2,3$, we have

$$ $\displaystyle\mathbb{E}[\prod_{k\in[c],i\in[d]}h_{k}^{n_{1k}}(\bar{g}_{k})_{i} ^{n_{2ki}}]=O(N^{-1})$ $$

Moreover, we have

$$ $\displaystyle\mathbb{E}[|(\bar{g}_{k})_{i}|]\leq\mathbb{E}[(\bar{g}_{k})_{i}^{ 2}]^{1/2}=N^{-1/2}c^{1/2}$ $\displaystyle\mathbb{E}[|(\bar{g}_{k})_{i}|^{3}]\leq\mathbb{E}[(\bar{g}_{k})_{ i}^{4}]^{3/4}=O(N^{-3/2})$ $$

where $i\in[d],k\in[c]$.

For any $n_{1k}$, $n_{2ki}$ satisfying $\sum_{k\in[c]}n_{1k}+\sum_{k\in[c],i\in[d]}n_{2ki}=n$, $n=1,2,3,4$, we have

$$ $\displaystyle\mathbb{E}[\prod_{k\in[c],i\in[d]}|h_{k}^{n_{1k}}(\bar{g}_{k})_{i }^{n_{2ki}}|]=O(N^{-n/2})$ $$

Proof
Since $g_{k}\sim\mathsf{N}(0,c^{2}N_{k}\Lambda/N^{2})$ and $\bar{g}_{k}\sim\mathsf{N}(0,c^{2}N_{k}I_{d}/N^{2})=\mathsf{N}(0,(c/N+c^{2}h_{k
}/N)I_{d})$, we have

$$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}]=0$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}(\bar{g}_{l})_{j}]=\delta_{kl}\delta_ {ij}c/N$ $\displaystyle\mathbb{E}[(\bar{g}_{k_{1}})_{i_{i}}(\bar{g}_{k_{2}})_{i_{2}}( \bar{g}_{k_{3}})_{i_{3}}]=0$ $\displaystyle\mathbb{E}[(\bar{g}_{k})_{i}^{4}]=\mathbb{E}[3c^{2}/N^{2}(1+ch_{k })^{2}]=O(N^{-2})$ $\displaystyle\mathbb{E}[h_{m}(\bar{g}_{k})_{i}(\bar{g}_{l})_{j}]=\mathbb{E}[c^ {2}\delta_{kl}\delta_{ij}h_{m}h_{k}/N]=O(N^{-2})$ $\displaystyle\mathbb{E}[h_{m}h_{l}(\bar{g}_{k})_{i}]=0$ $$

where $i,j,i_{1},i_{2},i_{3}\in[d],k,l,m,k_{1},k_{2},k_{3}\in[c]$. Thus, with the results from Lemma F.1, for any $n_{1k}$, $n_{2ki}$ satisfying $\sum_{k\in[c]}n_{1k}+\sum_{k\in[c],i\in[d]}n_{2ki}=1,2,3$, we have

$$ $\displaystyle\mathbb{E}[\prod_{k\in[c],i\in[d]}h_{k}^{n_{1k}}(\bar{g}_{k})_{i} ^{n_{2ki}}]=O(N^{-1})$ $$

Moreover, according to the Jensen’s inequality, we have

$$ $\displaystyle\mathbb{E}[|(\bar{g}_{k})_{i}|]\leq\mathbb{E}[(\bar{g}_{k})_{i}^{ 2}]^{1/2}=N^{-1/2}c^{1/2}$ $\displaystyle\mathbb{E}[|(\bar{g}_{k})_{i}|^{3}]\leq\mathbb{E}[(\bar{g}_{k})_{ i}^{4}]^{3/4}=O(N^{-3/2})$ $$

where $i\in[d],k\in[c]$.
Thus, with the results from Lemma F.1, for any $n_{1k}$, $n_{2ki}$ satisfying $\sum_{k\in[c]}n_{1k}+\sum_{k\in[c],i\in[d]}n_{2ki}=n$, $n=1,2,3,4$, we have

$$ $\displaystyle\mathbb{E}[\prod_{k\in[c],i\in[d]}|h_{k}^{n_{1k}}(\bar{g}_{k})_{i }^{n_{2ki}}|]\leq\prod_{k\in[c],i\in[d]}\mathbb{E}[|h_{k}^{n}|]^{n_{1k}/n} \mathbb{E}[|(\bar{g}_{k})_{i}^{n}|]^{n_{2ki}/n}=O(N^{-n/2}).$ $$

$\blacksquare$

###### Lemma F.3

For the loss function $L(W)$ (16), we have $\nabla^{2}L(W)\succ 0$. For any compact set $R_{W}$, when $W\in R_{W}$, we have $\nabla^{2}L(W)\succ\gamma I_{d}$ for some $\gamma>0$.
Additionally, $L(W)$ has one unique global minimizer on $\mathbb{R}^{d\times d}$.

For $\widetilde{L}(W)$ defined in (44), we also have $\nabla^{2}\widetilde{L}(W)\succ 0$. For any compact set $R_{W}$, when $W\in R_{W}$, we have $\nabla^{2}\widehat{L}(W)\succ\gamma I_{d}$ for some $\gamma>0$.
Additionally, $\widetilde{L}(W)$ has one unique global minimizer on $\mathbb{R}^{d\times d}$.

Proof
We vectorize $W$ as $\text{Vec}(W)\in\mathbb{R}^{d^{2}}$, where $\text{Vec}(W)_{i}=W_{t_{1}(i),t_{2}(i)}$, $t_{1}(x)=\lfloor(x-1)/d\rfloor+1,t_{2}(x)=((x-1)\text{ mod }d)+1$.
Then, we have

$$ $\displaystyle(\nabla L(W))_{i}=\mathbb{E}\left[\sum_{k=1}^{c}\left[(\mathrm{ softmax}(P^{\top}Wq/c)_{k}-\mathrm{softmax}(\mu^{\top}\Lambda^{-1}q)_{k})(p_{k })_{t_{1}(i)}q_{t_{2}(i)}/c\right]\right]$ (45) $$

Note that

$$ $\displaystyle\mathrm{softmax}(P^{\top}Wq/c)_{k}=\sigma(a_{k})$ $\displaystyle\nabla\mathrm{softmax}(P^{\top}Wq/c)_{k}=\sigma(a_{k})(1-\sigma(a _{k}))\nabla a_{k},$ $$

where $a_{k}=-\log(\sum_{l=1,\dots,c,l\neq k}\exp\left((p_{l}-p_{k})Wq/c\right))$.
For $\nabla a_{k}$, we have

$$ $\displaystyle\nabla a_{k}=$ $\displaystyle\frac{\sum_{l=1,\dots,c,l\neq k}\exp\left((p_{l}-p_{k})^{\top}Wq/ c\right)(p_{k}-p_{l})q^{\top}/c}{\sum_{l=1,\dots,c,l\neq k}\exp\left((p_{l}-p_ {k})^{\top}Wq/c\right)}$ $\displaystyle=$ $\displaystyle\frac{\sum_{l=1,\dots,c,l\neq k}\exp\left(p_{l}^{\top}Wq/c\right) (p_{k}-p_{l})q^{\top}/c}{\sum_{l=1,\dots,c,l\neq k}\exp\left(p_{l}^{\top}Wq/c \right)}.$ $$

Then we have

$$ $\displaystyle\nabla\mathrm{softmax}(P^{\top}Wq/c)_{k}=$ $\displaystyle\mathrm{softmax}(P^{\top}Wq/c)_{k}\frac{\sum_{l=1,\dots,c,l\neq k }\exp\left(p_{l}^{\top}Wq\right)(p_{k}-p_{l})q^{\top}/c}{\sum_{n=1,\dots,c} \exp\left(p_{n}^{\top}Wq/c\right)}$ $\displaystyle=$ $\displaystyle\sum_{l=1,\dots,c,l\neq k}\mathrm{softmax}(P^{\top}Wq/c)_{k} \mathrm{softmax}(P^{\top}Wq/c)_{l}(p_{k}-p_{l})q^{\top}/c$ $$

and

$$ $\displaystyle(\nabla\mathrm{softmax}(P^{\top}Wq/c)_{k})_{j}=\sum_{l=1,\dots,c, l\neq k}\mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}(p _{k}-p_{l})_{t_{1}(j)}q_{t_{2}(j)}/c.$ $$

We can express the Hessian matrix of the loss function with the following form:

$$ $\displaystyle(\nabla^{2}L(W))_{ij}=$ $\displaystyle\mathbb{E}\left[\sum_{k=1}^{c}\sum_{l=1,\dots,c,l\neq k}\mathrm{ softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}(p_{k})_{t_{1}(i)} q_{t_{2}(i)}(p_{k}-p_{l})_{t_{1}(j)}q_{t_{2}(j)}/c^{2}\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[\sum_{k=2}^{c}\sum_{l=1}^{k-1}\mathrm{softmax}(P^ {\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}(p_{k}-p_{l})_{t_{1}(i)}q_{t_ {2}(i)}(p_{k}-p_{l})_{t_{1}(j)}q_{t_{2}(j)}/c^{2}\right].$ $$

Considering $z\in\mathbb{R}^{d^{2}}$ such that $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z=$ $\displaystyle\mathbb{E}\left[\frac{1}{c^{2}}\sum_{k=2}^{c}\sum_{l=1}^{k-1} \mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}\sum_{ab}z _{a}z_{b}(p_{k}-p_{l})_{t_{1}(a)}q_{t_{2}(a)}(p_{k}-p_{l})_{t_{1}(b)}q_{t_{2}( b)}\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[\frac{1}{c^{2}}\sum_{k=2}^{c}\sum_{l=1}^{k-1} \mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}\left(\sum _{a\in[d^{2}]}z_{a}(p_{k}-p_{l})_{t_{1}(a)}q_{t_{2}(a)}\right)^{2}\right]$ $$

Since for any $P,q$, $k,l$, $\mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}>0$, we have $z^{\top}\nabla^{2}L(W)z\geq 0$. Thus, $\nabla^{2}L(W)\succeq 0$ and $L(W)$ is convex.

Defining $\tilde{p}=p_{1}-p_{2}$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z$ $\displaystyle\geq$ $\displaystyle\mathbb{E}\left[\frac{1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{1} \mathrm{softmax}(P^{\top}Wq/c)_{2}\left(\sum_{a\in[d^{2}]}z_{a}(p_{1}-p_{2})_{ t_{1}(a)}q_{t_{2}(a)}\right)^{2}\right]$ $\displaystyle=$ $\displaystyle\int\frac{1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{1}\mathrm{ softmax}(P^{\top}Wq/c)_{2}\left(\sum_{a\in[d^{2}]}z_{a}\tilde{p}_{t_{1}(a)}q_{ t_{2}(a)}\right)^{2}f_{Pq}(P,q)dPdq$ $$

where $f_{Pq}(P,q)$ are the PDF function of $P,q$. For any $z\neq 0$, we denote $z_{ij}=z_{((i-1)d+j)}$, suppose $a,b\in\arg\max_{i,j}|z_{ij}|$, we consider a set of constants
$\{c_{1pi},c_{2pi}\},\{c_{1qi},c_{2qi}\},i,j\in[d]$, where
$c_{1pa}=d,c_{2pa}=d+1$, $c_{1qb}=d,c_{2qb}=d+1$, and $c_{1pi}=1/16,c_{2pi}=1/8,i\neq a$, $c_{1qj}=1/16,c_{2qj}=1/8,j\neq b$. Then, for any $c_{pi}\in[c_{1pi},c_{2pi}],c_{qj}\in[c_{1qj},c_{2qj}]$, we have

$$ $\displaystyle\left|\sum_{i,j\in[d]}z_{ij}c_{pi}c_{qj}\right|\geq\left[d^{2}-2( d+1)(d-1)/8-(d-1)^{2}/64\right]\max_{ij}|z_{ij}|\geq d^{2}\max_{ij}|z_{ij}|/2.$ $$

Then, we define region $\Omega(a,b)=\{\tilde{p}=\sum_{i}c_{pi}\textbf{e}_{i}$, $q=\sum_{j}c_{qj}\textbf{e}_{j}$,$c_{pi}\in[c_{1pi},c_{2pi}],c_{qj}\in[c_{1qj},c_{2qj}],\|P\|_{F}^{2}\leq c^{2}(
\sum_{i}c_{2pi}^{2}+c_{2qj}^{2})\}$. We have

$$ $\displaystyle\min_{\Omega(a,b)}\left(\sum_{l\in[d^{2}]}z_{l}\tilde{p}_{t_{1}(l )}q_{t_{2}(l)}\right)^{2}\geq d^{4}\max_{ij}|z_{ij}|^{2}/4\geq\|z\|_{2}^{2}/4.$ $$

Defining

$$ $\displaystyle C(\Omega)=\min_{a\in[d],b\in[d]}\int_{\Omega(a,b)}f_{Pq}(P,q)dPdq,$ $\displaystyle S(\Omega,W)=\min_{a\in[d],b\in[d]}\min_{\Omega(a,b)}\left\{\frac {1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{1}\mathrm{softmax}(P^{\top}Wq/c)_{2} \right\},$ $$

we have $S(\Omega,W)>0$. Since we have $f_{Pq}(P,q)>0$ for all $P,q$ and $\Omega(a,b)$ are non-zero measures for $P,q$. Thus, we have $C(\Omega)>0$. Then, for any $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z$ $\displaystyle\geq$ $\displaystyle\int_{\Omega(a,b)}\frac{1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{ 1}\mathrm{softmax}(P^{\top}Wq/c)_{2}\left(\sum_{l\in[d^{2}]}z_{l}\tilde{p}_{t_ {1}(l)}q_{t_{2}(l)}\right)^{2}f_{Pq}(P,q)dPdq$ $\displaystyle\geq$ $\displaystyle C(\Omega)S(\Omega,W)\|z\|_{2}^{2}/4>0$ $$

Thus, we have $\nabla^{2}L(W)\succ 0$. $L(W)$ is strictly convex.

Moreover, for any compact set $R_{W}$ of $\mathbb{R}^{d\times d}$, for any $W\in R_{W}$, we have

$$ $\displaystyle S(\Omega)=\min_{W\in R_{W}}\min_{a\in[d],b\in[d]}\min_{\Omega(a, b)}\left\{\frac{1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{1}\mathrm{softmax}(P^ {\top}Wq/c)_{2}\right\}>0.$ $$

Then, for any $W\in R_{W}$, for any $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z$ $\displaystyle\geq$ $\displaystyle\int_{\Omega(a,b)}\frac{1}{c^{2}}\mathrm{softmax}(P^{\top}Wq/c)_{ 1}\mathrm{softmax}(P^{\top}Wq/c)_{2}\left(\sum_{l\in[d^{2}]}z_{l}\tilde{p}_{t_ {1}(l)}q_{t_{2}(l)}\right)^{2}f_{Pq}(P,q)dPdq$ $\displaystyle\geq$ $\displaystyle C(\Omega)S(\Omega)\|z\|_{2}^{2}/4.$ $$

Thus, when $W\in R_{W}$, $R_{W}$ is a compact set, we have $\nabla^{2}L(W)\succ C(\Omega)S(\Omega)I_{d}/4$, our loss function is $\gamma-$strongly convex, where $\gamma=C(\Omega)S(\Omega)/4$.

Because our loss function is strictly convex in $\mathbb{R}^{d\times d}$, it has at most one global minimizer in $\mathbb{R}^{d\times d}$. Next, we prove all level sets of our loss function are compact, i.e. $V_{\alpha}=\{W\in\mathbb{R}^{d\times d}\,|\,L(W)\leq\alpha\}$ is compact for all $\alpha$. We prove it by contradiction. Suppose $V_{\alpha}$ is not compact for some $\alpha$. Since our loss function is continuous and convex, $V_{\alpha}$ is an unbounded convex set. Since the dimension of $V_{\alpha}$ is $d^{2}$, consider a point $W^{\alpha}\in V_{\alpha}$, there must exists a $W^{k}\neq 0_{d\times d}$ that $\{W^{\alpha}+tW^{k}\,|\,t=[0,\infty)\}\in V_{\alpha}$.
For this $W^{k}\neq 0_{d\times d}$, there must exist a set of constants $0<c_{3pi}<c_{4pi},0<c_{3qj}<c_{4qj}$ such that for any $c_{pi}\in[c_{3pi},c_{4pi}],c_{qj}\in[c_{3qj},c_{4qj}]$, we have

$$ $\displaystyle|\sum_{ij}c_{pi}c_{qj}W_{ij}^{k}|\neq 0.$ $$

Thus, we have

$$ $\displaystyle\lim_{t\to\infty}|\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{ k})|=\infty.$ $$

We define $\Omega_{0}=\{\tilde{p}=\sum_{i}c_{pi}\textbf{e}_{i}$, $q=\sum_{j}c_{qj}\textbf{e}_{j}$, $c_{pi}\in[c_{3pi},c_{4pi}],c_{qj}\in[c_{3qj},c_{4qj}],\|P\|_{F}^{2}\leq c^{2}(
\sum_{i}c_{4pi}^{2}+c_{4qj}^{2}),\|\mu\|_{F}^{2}\leq c^{2}(\sum_{i}c_{4pi}^{2}
+c_{4qj}^{2})\}$. Then, defining

$$ $\displaystyle C(\Omega_{0})=\int_{\Omega_{0}}f_{Pq}(P,q)dPdq,$ $\displaystyle S(\Omega_{0})=\min_{\Omega_{0}}\left\{\min\{\mathrm{softmax}(\mu ^{\top}Wq/c)_{1},\mathrm{softmax}(\mu^{\top}Wq/c)_{2}\}\right\}$ $$

we have $S(\Omega_{0})>0$. Since $\Omega_{0}$ are non-zero measures for $P,q$, we have $C(\Omega_{0})>0$.
Then, we have

$$ $\displaystyle\lim_{t\to\infty}L(W^{\alpha}+tW^{k})$ $\displaystyle=$ $\displaystyle\lim_{t\to\infty}\mathbb{E}[\sum_{l=1}^{c}-\mathrm{softmax}(\mu^{ \top}\Lambda^{-1}q)_{l}\log(\mathrm{softmax}(P^{\top}(W^{\alpha}+tW^{k})q/c)_{ l})]$ $\displaystyle\geq$ $\displaystyle\lim_{t\to\infty}\int_{\Omega_{0}}[-\mathrm{softmax}(\mu^{\top} \Lambda^{-1}q)_{1}\log(\mathrm{softmax}(P^{\top}(W^{\alpha}+tW^{k})q/c)_{1})]f _{Pq}(P,q)dPdq$ $\displaystyle+\lim_{t\to\infty}\int_{\Omega_{0}}[-\mathrm{softmax}(\mu^{\top} \Lambda^{-1}q)_{2}\log(\mathrm{softmax}(P^{\top}(W^{\alpha}+tW^{k})q/c)_{2})]f _{Pq}(P,q)dPdq$ $\displaystyle\geq$ $\displaystyle\lim_{t\to\infty}\int_{\Omega_{0}}[-\mathrm{softmax}(\mu^{\top} \Lambda^{-1}q)_{1}\log(\sigma(\tilde{p}^{\top}(W^{\alpha}+tW^{k})q/c))]f_{Pq}( P,q)dPdq$ $\displaystyle+\lim_{t\to\infty}\int_{\Omega_{0}}[-\mathrm{softmax}(\mu^{\top} \Lambda^{-1}q)_{2}\log(\sigma(-\tilde{p}^{\top}(W^{\alpha}+tW^{k})q/c))]f_{Pq} (P,q)dPdq$ $\displaystyle\geq$ $\displaystyle C(\Omega_{0})S(\Omega_{0})\cdot\min_{\Omega_{0}}\left\{\lim_{t \to\infty}[-\log(\sigma(\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/c)) ]\right\}$ $\displaystyle+$ $\displaystyle C(\Omega_{0})S(\Omega_{0})\cdot\min_{\Omega_{0}}\left\{\lim_{t \to\infty}[-\log(\sigma(-\sum_{ij}c_{pi}c_{qj}(W_{ij}^{\alpha}+tW_{ij}^{k})/c) )]\right\}$ $\displaystyle=$ $\displaystyle\infty$ $$

This contradicts the assumption $L(W^{\alpha}+tW^{k})\leq\alpha$. Thus, all level sets of the loss function $L(W)$ are compact, which means there exists a a global minimizer for $L(W)$. Together with the fact that $L(W)$ is strictly convex, $L(W)$ has one unique a global minimizer on $\mathbb{R}^{d\times d}$.

Similarly, we can prove the same conclusions for $\widetilde{L}(W)$.
$\blacksquare$

###### Lemma F.4

Denoting the global minimizer of our loss function (16) as $W^{*}$, we have $W^{*}=c(\Lambda^{-1}+G)$, where $\|G\|_{\max}=O(N^{-1/4})$.

Proof
Let $a=\mu^{\top}\Lambda^{-1}q$, $s=\mu^{\top}Wq/c$, $r=(\mu_{h}+g)^{\top}Wq/c$,
$a_{k}=\mu_{k}^{\top}\Lambda^{-1}q$, $s_{k}=\mu_{k}^{\top}Wq/c$, $r_{k}=(ch_{k}\mu_{k}+g_{k})^{\top}Wq/c$.
Performing the Taylor expansion on (16), we have

$$ $\displaystyle L(W)=$ $\displaystyle\mathbb{E}\left[\sum_{k=1}^{c}-\zeta_{k}(a)\log(\zeta_{k}(s+r))\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[\sum_{k=1}^{c}-\zeta_{k}(a)\log(\zeta_{k}(s))- \sum_{k,l=1}^{c}\zeta_{k}(a)R_{kl}(s,r)r_{l}\right]$ $\displaystyle=$ $\displaystyle\widetilde{L}(W)-\mathbb{E}\left[\sum_{k,l=1}^{c}\zeta_{k}(a)R_{ kl}(s,r)r_{l}\right]$ $$

where $|R_{kl}(s,r)|\leq\sup_{y}|\frac{\partial\log(\zeta_{k}(y))}{\partial y_{l}}|
\sup_{y}|\frac{1}{\zeta_{k}(y)}\frac{\partial\zeta_{k}(y)}{\partial y_{l}}|=
\sup_{y}|\delta_{kl}-\zeta_{l}(y)|\leq 1$. Thus, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq$ $\displaystyle c\sum_{l=1}^{c}\mathbb{E}\left[|r_{l}|\right]$ $\displaystyle\leq$ $\displaystyle\sum_{l=1}^{c}c\mathbb{E}\left[|h_{l}\mu_{l}^{\top}Wq|\right]+ \mathbb{E}\left[|g_{l}^{\top}Wq|\right]$ $\displaystyle\leq$ $\displaystyle O(1)\|W\|_{\max}\mathbb{E}[|h_{l}|]+O(1)\|W\|_{\max}\mathbb{E}[| (\bar{g}_{l})_{i}|]$ $\displaystyle\leq$ $\displaystyle C_{l}\|W\|_{\max}N^{-1/2}$ $$

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 $\widetilde{L}(W)$.

According to Lemma D.2, $\widetilde{L}(W)$ has one unique global minimizer. Considering the equation:

$$ $\displaystyle\nabla\widetilde{L}(W)=\mathbb{E}\left[\sum_{k=1}^{c}-\mathrm{ softmax}(\mu^{\top}\Lambda^{-1}q)_{k}\log(\mathrm{softmax}(\mu^{\top}Wq/c)_{k} )\right]=0$ $$

We can easily find that $\nabla\widetilde{L}(c\Lambda^{-1})=0$ and $W=c\Lambda^{-1}$ is the global minimizer of $\widetilde{L}(W)$.

Considering a compact set $R_{W}=\{W\,|\,\|W-2\Lambda^{-1}\|_{F}\leq\rho_{W}\}$, we have $\|W\|_{\max}\leq C_{W}$ for any $W\in R_{W}$. Here $\rho_{W},C_{W}$ are some positive finite constants. Then, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq C^{\prime}_{l}N^{-1/2},\,W \in R_{W}$ $$

where $C^{\prime}_{l}=C_{l}C_{W}$ is a constant independent of $N$ and $W$. This shows that, for any $W\in R_{W}$, $L(W)$ uniformly converge to $\widetilde{L}(W)$.

Denote $W^{*}$ as the global minimizer of $L(W)$ with prompt length $N$. Then, we show that, when $N$ is sufficiently large, $W^{*}\in R_{W}$.
We first denote $\partial R_{W}=\{W\,|\,\|W-c\Lambda^{-1}\|_{F}=\rho_{W}\}$ , $\Delta=\min_{W\in\partial R_{W}}\widetilde{L}(W)-\widetilde{L}(c\Lambda^{-1})>0$.
Then, for $N\geq(4C^{\prime}_{l}/\Delta)^{2}$, and for any $W\in R_{W}$, we have

$$ $\displaystyle\left|\widetilde{L}(W)-L(W)\right|\leq\Delta/4$ $$

$$ $\displaystyle\min_{W\in\partial R_{W}}L(W)-\min_{W\in R_{W}}L(W)\geq\min_{W\in \partial R_{W}}L(W)-L(c\Lambda^{-1})\geq\Delta/2>0$ $$

Since $L(W)$ is strictly convex, we have $W^{*}=\mathop{\rm argmin}_{W}L(W)\in R_{W}$.

Then, we have

$$ $\displaystyle|\widetilde{L}(W^{*})-L(W^{*})|\leq C^{\prime}_{l}/N$ $\displaystyle|\widetilde{L}(c\Lambda^{-1})-L(c\Lambda^{-1})|\leq C^{\prime}_{l }/N$ $$

$$ $\displaystyle\widetilde{L}(W^{*})\leq L(W^{*})+C^{\prime}_{l}/N\leq L(c\Lambda ^{-1})+C^{\prime}_{l}/N\leq\widetilde{L}(c\Lambda^{-1})+2C^{\prime}_{l}N^{-1/2}$ $$

According to Lemma D.2, for $W\in R_{W}$, we have $\nabla^{2}\widetilde{L}(W)\succ\gamma I_{d}$, where $\gamma$ is a positive constant independent of $N$. Thus, $\widetilde{L}(W)$ is $\gamma$-strongly convex in $R_{W}$. According to Lemma C.1, we have

$$ $\displaystyle\|W^{*}-c\Lambda^{-1}\|_{F}^{2}\leq\frac{2}{\gamma}(\widetilde{L} (W^{*})-\widetilde{L}(c\Lambda^{-1}))\leq\frac{4C^{\prime}_{l}}{\gamma N^{1/2}}$ $$

Thus, when $N\to\infty$, we have $W^{*}\to c\Lambda^{-1}$. Denoting $W^{*}=c(\Lambda^{-1}+G)$, we have $\|G\|_{\max}=O(N^{-1/4})$.
$\blacksquare$

###### Lemma F.5

The global minimizer of the loss function (16) is $W^{*}=c(\Lambda^{-1}+G)$. We have

$$ $\displaystyle\|G\|_{\max}\leq$ $\displaystyle\frac{1}{N}\Bigg{\|}S^{-1}\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}\frac {\partial\zeta_{k}(a)}{\partial a_{l}}(c\delta_{kl}-1)\mu_{k}\mu_{l}^{\top} \Lambda^{-1}qq^{\top}+\sum_{k=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{k} }cqq^{\top}$ $\displaystyle+\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l} \partial a_{n}}(c\delta_{ln}-1)\mu_{l}^{\top}\Lambda^{-1}q\mu_{n}^{\top} \Lambda^{-1}q\mu_{k}q^{\top}/2+\sum_{k,l=1}^{c}\frac{\partial^{2}\zeta_{k}(a)} {\partial a_{l}^{2}}cq^{\top}\Lambda^{-1}q\mu_{k}q^{\top}/2\Bigg{]}\Bigg{\|}_{\max}$ $\displaystyle+o(1/N),$ $$

where $a=\mu^{\top}\Lambda^{-1}q$, $a_{k}=\mu_{k}^{\top}\Lambda^{-1}q$, $S=c^{2}\nabla^{2}\widetilde{L}(c\Lambda^{-1})$.
Ignoring constants other than $c,N$, we have $\|G\|_{\max}\leq O(c/N)$.

Proof
According to Lemma F.3, the loss function $L(W)$ has a unique global minimizer $W^{*}$. We have

$$ $\displaystyle\nabla L(W^{*})=\mathbb{E}\left[\sum_{k=1}^{c}\left[(\zeta_{k}(P^ {\top}W^{*}q/c)-\zeta_{k}(\mu^{\top}\Lambda^{-1}q))p_{k}q^{\top}/c\right] \right]=0.$ (46) $$

Let $W^{*}=c(\Lambda^{-1}+G)$, $a=\mu^{\top}\Lambda^{-1}q$, $a_{k}=\mu_{k}^{\top}\Lambda^{-1}q$, $b=(\mu+c\mu_{h}+g)^{\top}Gq+(c\mu_{h}+g)^{\top}\Lambda^{-1}q$, $b_{k}=(\mu_{k}+ch_{k}\mu_{k}+g_{k})^{\top}Gq+(ch_{k}\mu_{k}+g_{k})^{\top}
\Lambda^{-1}q$. The Taylor expansion of $\zeta_{k}(a+b)$ at point $a$ is

$$ $\displaystyle\zeta_{k}(a+b)=\zeta_{k}(a)+\sum_{l=1}^{c}\frac{\partial\zeta_{k} (a)}{\partial a_{l}}b_{l}+\sum_{l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{ \partial a_{l}\partial a_{n}}b_{l}b_{n}/2!+\sum_{l,n,m=1}^{c}R_{klnm}(a,b)b_{l }b_{n}b_{m}/3!,$ $$

where $|R_{klnm}(a,b)|\leq\sup_{x}|\frac{\partial^{3}\zeta_{k}(x)}{\partial x_{l}
\partial x_{n}\partial x_{m}}|$.
Thus, our equation (46) become

$$ $\displaystyle\mathbb{E}\left[\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}b_{l}p_{k}q^{\top}+\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta(a) }{\partial a_{l}\partial a_{n}}b_{l}b_{n}p_{k}q^{\top}/2!+\sum_{k,l,n,m=1}^{c} R_{klnm}(a,b)b_{l}b_{n}b_{m}p_{k}q^{\top}/3!\right]=0.$ (47) $$

For the first term $\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{l}}b_{l}p_{k}q^{\top}$, according to Lemma F.1, we have

$$ $\displaystyle\mathbb{E}\left[\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}b_{l}p_{k}q^{\top}\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\left[\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}\left[\mu_{l}^{\top}Gq\mu_{k}q^{\top}+c^{2}h_{l}h_{k}\mu_{l}^{ \top}Gq\mu_{k}q^{\top}+c^{2}h_{l}h_{k}\mu_{l}^{\top}\Lambda^{-1}q\mu_{k}q^{ \top}+g_{l}^{\top}\Lambda^{-1}qg_{k}q^{\top}+g_{l}^{\top}Gqg_{k}q^{\top}\right ]\right]$ $\displaystyle=$ $\displaystyle\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}\left(\mu_{k}\mu_{l}^{\top}Gqq^{\top}+(c\delta_{kl}-1)\mu_{k} \mu_{l}^{\top}Gqq^{\top}/N+(c\delta_{kl}-1)\mu_{k}\mu_{l}^{\top}\Lambda^{-1}qq ^{\top}/N\right)$ $\displaystyle+\sum_{k=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{k}}\left( cqq^{\top}/N+c\Lambda Gqq^{\top}/N\right)\Bigg{]}.$ (48) $$

According to Lemma F.4, $O(\|G\|_{\max})=O(N^{-1/4})=o(1)$, we have

$$ $\displaystyle\Bigg{\|}\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}\frac{\partial\zeta_{k }(a)}{\partial a_{l}}(c\delta_{kl}-1)\mu_{k}\mu_{l}^{\top}Gqq^{\top}/N+\sum_{k =1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{k}}c\Lambda Gqq^{\top}/N\Bigg{] }\Bigg{\|}_{\max}$ $\displaystyle\leq$ $\displaystyle O(\|G\|_{\max}/N)=o(1/N)$ (49) $$

For the second term $\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l}\partial a_{n}
}b_{l}b_{n}p_{k}q^{\top}/2!$, we have

$$ $\displaystyle\mathbb{E}\left[\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)} {\partial a_{l}\partial a_{n}}b_{l}b_{n}p_{k}q^{\top}/2!\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\Biggr{[}\sum_{k,l,n=1}^{c}\frac{\partial^{2 }\zeta_{k}(a)}{\partial a_{l}\partial a_{n}}\Biggr{(}\underbrace{\sum_{\phi_{1 }\in\{\mu_{l},ch_{l}\mu_{l},g_{l}\},\phi_{2}\in\{\mu_{n},ch_{n}\mu_{n},g_{n}\} }\phi_{1}^{\top}Gq\phi_{2}^{\top}Gqp_{k}q^{\top}}_{(i)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{\mu_{l},ch_{l}\mu_{l},g_{l}\}, \phi_{2}\in\{ch_{n}\mu_{n},g_{n}\}}2\phi_{1}^{\top}Gq\phi_{2}^{\top}\Lambda^{- 1}qp_{k}q^{\top}}_{(ii)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{ch_{l}\mu_{l},g_{l}\},\phi_{2}\in \{ch_{n}\mu_{n},g_{n}\}}\phi_{1}^{\top}\Lambda^{-1}q\phi_{2}^{\top}\Lambda^{-1 }qp_{k}q^{\top}}_{(iii)}\Biggr{)}\Biggr{]}.$ $$

For terms $(i)$ having two $G$, their max norms are at most smaller than $O(\|G\|_{\max}^{2})$. For terms $(ii)$ having one $G$, define $\bar{g}_{l}=\Lambda^{-1/2}g_{l}$, these terms must contain $n_{1j}$ number of $h_{j}$ and $n_{2ji}$ number of $(\bar{g}_{j})_{i}$, we have
$\sum_{j\in[c],i\in[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$,

$$ $\displaystyle\mathbb{E}[\prod_{j\in[c],i\in[d]}h_{j}^{n_{1j}}(\bar{g}_{j})_{i} ^{n_{2ji}}]=O(N^{-1})$ $$

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

$$ $\displaystyle\|\mathbb{E}[(i)]\|_{\max}\leq O(\|G\|_{\max}^{2})=o(\|G\|_{\max})$ (50) $\displaystyle\|\mathbb{E}[(ii)]\|_{\max}\leq O(\|G\|_{\max}/N)=o(1/N)$ (51) $$

For terms $(iii)$ without $G$, we have

$$ $\displaystyle\|\mathbb{E}[(iii)]\|_{\max}$ $\displaystyle=$ $\displaystyle\Biggr{\|}\mathbb{E}\Biggr{[}\sum_{k,l,n=1}^{c}\frac{\partial^{2} \zeta_{k}(a)}{\partial a_{l}\partial a_{n}}c^{2}h_{l}h_{n}\mu_{l}^{\top} \Lambda^{-1}q\mu_{n}^{\top}\Lambda^{-1}q\mu_{k}q^{\top}/2+\sum_{k,l=1}^{c} \frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l}^{2}}g_{l}^{\top}\Lambda^{-1}qg_ {l}^{\top}\Lambda^{-1}q\mu_{k}q^{\top}/2$ $\displaystyle+\sum_{k,l=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l} \partial a_{k}}ch_{l}\mu_{l}^{\top}\Lambda^{-1}qg_{k}^{\top}\Lambda^{-1}qg_{k} q^{\top}+\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l} \partial a_{n}}c^{3}h_{l}h_{n}h_{k}\mu_{l}^{\top}\Lambda^{-1}q\mu_{n}^{\top} \Lambda^{-1}q\mu_{k}q^{\top}/2\Biggr{]}\Biggr{\|}_{\max}$ $\displaystyle\leq$ $\displaystyle\frac{1}{2N}\Biggr{\|}\mathbb{E}\left[\sum_{k,l,n=1}^{c}\frac{ \partial^{2}\zeta_{k}(a)}{\partial a_{l}\partial a_{n}}(c\delta_{ln}-1)\mu_{l} ^{\top}\Lambda^{-1}q\mu_{n}^{\top}\Lambda^{-1}q\mu_{k}q^{\top}+\sum_{k,l=1}^{c }\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l}^{2}}cq^{\top}\Lambda^{-1}q\mu_ {k}q^{\top}\right]\Biggr{\|}_{\max}$ $\displaystyle+O(1/N^{2})$ (52) $$

where the last inequity is due to Lemma F.1, F.2.

For the third term $\sum_{k,l,n,m=1}^{c}R_{klnm}(a,b)b_{l}b_{n}b_{m}p_{k}q^{\top}/3!$, we have

$$ $\displaystyle\Biggr{\|}\mathbb{E}\left[\sum_{k,l,n,m=1}^{c}R_{klnm}(a,b)b_{l}b _{n}b_{m}p_{k}q^{\top}/3!\right]\Biggr{\|}_{\max}$ $\displaystyle\leq$ $\displaystyle O(1)\max_{l,m\in[d]}\mathbb{E}[\sum_{k_{1},k_{2},k_{3},k_{4}\in[ c]}\left|b_{k_{1}}b_{k_{2}}b_{k_{3}}(p_{k_{4}})_{l}q_{m}\right|]$ $\displaystyle\leq$ $\displaystyle O(1)\mathbb{E}\sum_{k_{1},k_{2},k_{3},k_{4}\in[c]}\Biggr{[} \underbrace{\sum_{\phi_{1}\in\{\mu_{k_{1}},ch_{k_{1}}\mu_{k_{1}},g_{k_{1}}\}, \phi_{2}\in\{\mu_{k_{2}},ch_{k_{2}}\mu_{k_{2}},g_{k_{2}}\},\phi_{3}\in\{\mu_{k _{3}},ch_{k_{3}}\mu_{k_{3}},g_{k_{3}}\}}\phi_{1}^{\top}Gq\phi_{2}^{\top}Gq\phi _{3}^{\top}Gq(p_{k_{4}})_{l}q_{m}}_{(*)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{\mu_{k_{1}},ch_{k_{1}}\mu_{k_{1}} ,g_{k_{1}}\},\phi_{2}\in\{\mu_{k_{2}},ch_{k_{2}}\mu_{k_{2}},g_{k_{2}}\},\phi_{ 3}\in\{ch_{k_{3}}\mu_{k_{3}},g_{k_{3}}\}}\phi_{1}^{\top}Gq\phi_{2}^{\top}Gq \phi_{3}^{\top}\Lambda^{-1}q(p_{k_{4}})_{l}q_{m}}_{(*)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{\mu_{k_{1}},ch_{k_{1}}\mu_{k_{1}} ,g_{k_{1}}\},\phi_{2}\in\{ch_{k_{2}}\mu_{k_{2}},g_{k_{2}}\},\phi_{3}\in\{ch_{k _{3}}\mu_{k_{3}},g_{k_{3}}\}}\phi_{1}^{\top}Gq\phi_{2}^{\top}\Lambda^{-1}q\phi _{3}^{\top}\Lambda^{-1}q(p_{k_{4}})_{l}q_{m}}_{(**)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{ch_{k_{1}}\mu_{k_{1}},g_{k_{1}}\} ,\phi_{2}\in\{ch_{k_{2}}\mu_{k_{2}},g_{k_{2}}\},\phi_{3}\in\{ch_{k_{3}}\mu_{k_ {3}},g_{k_{3}}\}}\phi_{1}^{\top}\Lambda^{-1}q\phi_{2}^{\top}\Lambda^{-1}q\phi_ {3}^{\top}\Lambda^{-1}q(p_{k_{4}})_{l}q_{m}}_{(***)}\Biggr{]}.$ $$

For terms in $(*)$ having two or three $G$, these terms’ expected absolute values are at most smaller than $O(\|G\|_{\max}^{2})$. For terms in $(**)$ having one $G$, these terms must contain $n_{1j}$ number of $h_{j}$ and $n_{2ji}$ number of $(\bar{g}_{j})_{i}$, we have
$\sum_{j\in[c],i\in[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

$$ $\displaystyle\mathbb{E}[\prod_{j\in[c],i\in[d]}|h_{k}^{n_{1k}}(\bar{g}_{j})_{i }^{n_{2ji}}|]=O(N^{-n_{t}/2})=O(N^{-1})$ $$

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 $(\bar{g}_{j})_{i}$, we have
$\sum_{j\in[c],i\in[d]}n_{1j}+n_{2ji}=n_{t},n_{t}=3,4$.
According to Lemma F.2, for $n_{t}=3,4$, we have

$$ $\displaystyle\mathbb{E}[\prod_{j\in[c],i\in[d]}|h_{k}^{n_{1k}}(\bar{g}_{j})_{i }^{n_{2ji}}|]=O(N^{-n_{t}/2})=O(N^{-3/2})$ $$

Thus, these term’s expected absolute values are at most smaller than $O(N^{-3/2})$.
Therefore, we have

$$ $\displaystyle\Biggr{\|}\mathbb{E}\left[\sum_{k,l,n,m=1}^{c}R_{klnm}(a,b)b_{l}b _{n}b_{m}p_{k}q^{\top}/3!\right]\Biggr{\|}_{\max}$ $\displaystyle\leq$ $\displaystyle O(1)\max_{l,m\in[d]}\mathbb{E}[\sum_{k_{1},k_{2},k_{3},k_{4}\in[ c]}\left|b_{k_{1}}b_{k_{2}}b_{k_{3}}(p_{k_{4}})_{l}q_{m}\right|]$ $\displaystyle\leq$ $\displaystyle O(\|G\|^{2}_{max})+O(\|G\|_{\max}N^{-1})+O(N^{-3/2})$ $\displaystyle\leq$ $\displaystyle o(\|G\|_{\max})+o(1/N).$ (53) $$

Moreover, we have

$$ $\displaystyle\Bigg{\{}\mathbb{E}\left[\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}( a)}{\partial a_{l}}\mu_{k}\mu_{l}^{\top}Gqq^{\top}\right]\Bigg{\}}_{ij}$ $\displaystyle=$ $\displaystyle\Bigg{\{}\mathbb{E}\left[\sum_{k=1}^{c}\zeta_{k}(a)(1-\zeta_{k}(a ))\mu_{k}\mu_{k}^{\top}Gqq^{\top}-\sum_{k,l=1,k\neq l}^{c}\zeta_{k}(a)\zeta_{l }(a)\mu_{k}\mu_{l}^{\top}Gqq^{\top}\right]\Bigg{\}}_{ij}$ $\displaystyle=$ $\displaystyle\Bigg{\{}\mathbb{E}\left[\sum_{k,l=1,k\neq l}^{c}\zeta_{k}(a) \zeta_{l}(a)\mu_{k}(\mu_{k}-\mu_{l})^{\top}Gqq^{\top}\right]\Bigg{\}}_{ij}$ $\displaystyle=$ $\displaystyle\Bigg{\{}\mathbb{E}\left[\sum_{k=2}^{c}\sum_{l=1}^{k-1}\zeta_{k}( a)\zeta_{l}(a)(\mu_{k}-\mu_{l})(\mu_{k}-\mu_{l})^{\top}Gqq^{\top}\right]\Bigg{ \}}_{ij}$ $\displaystyle=$ $\displaystyle\sum_{n,m=1}^{d}s_{ijnm}G_{nm},$ (54) $$

where $s_{ijnm}=\mathbb{E}\left[\sum_{k=2}^{c}\sum_{l=1}^{k-1}\zeta_{k}(a)\zeta_{l}(a
)(\mu_{k}-\mu_{l})_{i}(\mu_{k}-\mu_{l})_{n}q_{m}q_{j}\right]$.
We vectorize $G$ as $\text{Vec}(G)_{i}=G_{t_{1}(i),t_{2}(i)}$. Define $S\in\mathbb{R}^{d^{2}\times d^{2}}$, where $S_{ij}=s_{t_{1}(i),t_{2}(i),t_{1}(j),t_{2}(j)}=\mathbb{E}\left[\sum_{k=2}^{c}
\sum_{l=1}^{k-1}\zeta_{k}(a)\zeta_{l}(a)(\mu_{k}-\mu_{l})_{t_{1}(i)}q_{t_{2}(i
)}(\mu_{k}-\mu_{l})_{t_{1}(j)}q_{t_{2}(j)}\right]$,
(54) can be expressed as

$$ $\displaystyle\mathbb{E}\left[\sum_{k,l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}\mu_{k}\mu_{l}^{\top}Gqq^{\top}\right]=SG.$ (55) $$

Note that $S=c^{2}\nabla^{2}\widetilde{L}(c\Lambda^{-1})$. According to Lemma F.3, $S$ is positive definite. Thus, combining (47), (48), (49), (50), (51), (52), (53), (55), we have

$$ $\displaystyle\|G\|_{\max}\leq$ $\displaystyle\frac{1}{N}\Bigg{\|}S^{-1}\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}\frac {\partial\zeta_{k}(a)}{\partial a_{l}}(c\delta_{kl}-1)\mu_{k}\mu_{l}^{\top} \Lambda^{-1}qq^{\top}+\sum_{k=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{k} }cqq^{\top}$ $\displaystyle+\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l} \partial a_{n}}(c\delta_{ln}-1)\mu_{l}^{\top}\Lambda^{-1}q\mu_{n}^{\top} \Lambda^{-1}q\mu_{k}q^{\top}/2+\sum_{k,l=1}^{c}\frac{\partial^{2}\zeta_{k}(a)} {\partial a_{l}^{2}}cq^{\top}\Lambda^{-1}q\mu_{k}q^{\top}/2\Bigg{]}\Bigg{\|}_{\max}$ $\displaystyle+o(1/N).$ $$

Ignoring constants other than $c,N$, we have $\|G\|_{\max}\leq O(c/N)$.
$\blacksquare$

###### Lemma F.6

The loss function (7) is $l$-smooth, where $l\leq\frac{1}{c^{2}}\sum_{k=2}^{c}\sum_{l=1}^{k-1}\sum_{i\in[d^{2}]}\mathbb{E}
[((p_{k}-p_{l})_{t_{1}(i)}q_{t_{2}(i)})^{2}]$.

Proof
The Hessian matrix of the loss function is

$$ $\displaystyle(\nabla^{2}L(W))_{ij}=\mathbb{E}\left[\sum_{k=2}^{c}\sum_{l=1}^{k -1}\mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}(p_{k}- p_{l})_{t_{1}(i)}q_{t_{2}(i)}(p_{k}-p_{l})_{t_{1}(j)}q_{t_{2}(j)}/c^{2}\right].$ $$

Considering $z\in\mathbb{R}^{d^{2}}$ such that $z\neq 0$, we have

$$ $\displaystyle z^{\top}\nabla^{2}L(W)z=$ $\displaystyle\mathbb{E}\left[\frac{1}{c^{2}}\sum_{k=2}^{c}\sum_{l=1}^{k-1} \mathrm{softmax}(P^{\top}Wq/c)_{k}\mathrm{softmax}(P^{\top}Wq/c)_{l}\left(\sum _{a\in[d^{2}]}z_{a}(p_{k}-p_{l})_{t_{1}(a)}q_{t_{2}(a)}\right)^{2}\right]$ $\displaystyle\overset{(a)}{\leq}$ $\displaystyle\frac{1}{c^{2}}\|z\|_{2}^{2}\sum_{k=2}^{c}\sum_{l=1}^{k-1}\sum_{i \in[d^{2}]}\mathbb{E}[((p_{k}-p_{l})_{t_{1}(i)}q_{t_{2}(i)})^{2}]$ $$

where $(a)$ is due to the Cauchy–Schwarz inequality. Thus, $\nabla^{2}L(W)\preceq lI_{d}$ and $L(W)$ is $l$-smooth, where $l$ is a constant smaller than $\frac{1}{c^{2}}\sum_{k=2}^{c}\sum_{l=1}^{k-1}\sum_{i\in[d^{2}]}\mathbb{E}[((p_
{k}-p_{l})_{t_{1}(i)}q_{t_{2}(i)})^{2}]$.
$\blacksquare$

###### Theorem F.1 (Formal statement of Theorem 4.1 )

The following statements hold.

- (1)
Optimizing training loss $L(W)$ (16) with training prompt length $N$ via gradient descent $W^{t+1}=W^{t}-\eta\nabla L(W^{t})$, we have for any $t$
$\displaystyle\|W^{t}-W^{*}\|_{F}^{2}\leq\exp(-t/\kappa)\|W^{0}-W^{*}\|_{F}^{2},$
where $W^{0}$ is the initial parameter and $W^{*}$ is the global minimizer of $L(W)$, $\kappa=l/\alpha$. $\alpha,l$ are constants such that
$\displaystyle 0<\alpha\leq\lambda_{\min}(\nabla^{2}L(W))\leq\lambda_{\max}(
\nabla^{2}L(W))\leq l,\text{ for all }\,W\in R_{W},$
(56)
where $R_{W}=\{W\in\mathbb{R}^{d\times d}\,|\,\|W-W^{*}\|_{F}\leq\|W^{0}-W^{*}\|_{F}\}$.
- (2)
Denoting $W^{*}=c(\Lambda^{-1}+G)$, we have
$\displaystyle\|G\|_{\max}$
$\displaystyle\leq\frac{1}{N}\Bigg{\|}S^{-1}\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}
\frac{\partial\zeta_{k}(a)}{\partial a_{l}}(c\delta_{kl}-1)\mu_{k}\mu_{l}^{
\top}\Lambda^{-1}qq^{\top}+\sum_{k=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a
_{k}}cqq^{\top}$
$\displaystyle+\frac{1}{2}\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{
\partial a_{l}\partial a_{n}}(c\delta_{ln}-1)\mu_{l}^{\top}\Lambda^{-1}q\mu_{n
}^{\top}\Lambda^{-1}q\mu_{k}q^{\top}+\frac{1}{2}\sum_{k,l=1}^{c}\frac{\partial
^{2}\zeta_{k}(a)}{\partial a_{l}^{2}}cq^{\top}\Lambda^{-1}q\mu_{k}q^{\top}
\Bigg{]}\Bigg{\|}_{\max}$
$\displaystyle+o(1/N)$
$\displaystyle=O(c/N)$
where $S=c^{2}\nabla^{2}\widetilde{L}(2\Lambda^{-1})$, $\widetilde{L}(2\Lambda^{-1})=\lim_{N\to\infty}L(2\Lambda^{-1})$. The expectation is taken over $\mu_{\tau}\sim\mathcal{P}^{m}_{\Omega}(\Lambda)$, $x_{\tau,\mathsf{query}}\sim\mathcal{P}^{m}_{x}(\mu_{\tau},\Lambda)$.
- (3)
After $T\geq 2\kappa\log(N\cdot\|W^{0}-W^{*}\|_{F})$ gradient steps, denoting $\widehat{W}$ as the final model, we have
$\displaystyle\widehat{W}=c(\Lambda^{-1}+\widehat{G}),$
(57)
where $\|\widehat{G}\|_{\max}=O(c/N)$.

###### Lemma F.1

###### Lemma F.2

###### Lemma F.3

###### Lemma F.4

###### Lemma F.5

###### Lemma F.6

###### Theorem F.1 (Formal statement of Theorem 4.1 )

### F.3 Proof of Theorem 4.1

Proof
According to Lemma F.5, the global minimizer of $L(W)$ is $W^{*}=c(\Lambda^{-1}+G)$, where

$$ $\displaystyle\|G\|_{\max}\leq$ $\displaystyle\frac{1}{N}\Bigg{\|}S^{-1}\mathbb{E}\Bigg{[}\sum_{k,l=1}^{c}\frac {\partial\zeta_{k}(a)}{\partial a_{l}}(c\delta_{kl}-1)\mu_{k}\mu_{l}^{\top} \Lambda^{-1}qq^{\top}+\sum_{k=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{k} }cqq^{\top}$ $\displaystyle+\sum_{k,l,n=1}^{c}\frac{\partial^{2}\zeta_{k}(a)}{\partial a_{l} \partial a_{n}}(c\delta_{ln}-1)\mu_{l}^{\top}\Lambda^{-1}q\mu_{n}^{\top} \Lambda^{-1}q\mu_{k}q^{\top}/2+\sum_{k,l=1}^{c}\frac{\partial^{2}\zeta_{k}(a)} {\partial a_{l}^{2}}cq^{\top}\Lambda^{-1}q\mu_{k}q^{\top}/2\Bigg{]}\Bigg{\|}_{\max}$ $\displaystyle+o(1/N).$ $$

Ignoring constants other than $c,N$, we have $\|G\|_{\max}\leq O(c/N)$.

Define $R_{W}=\{W\in\mathbb{R}^{d\times d}\,|\,\|W-W^{*}\|_{F}\leq\|W^{0}-W^{*}\|_{F}\}$, and $R_{W}$ is a compact set. Then, according to Lemma F.3, for $W\in R_{W}$, we have $\nabla^{2}L(W)\succeq\alpha I_{d}$. Here $\alpha>0$ is a positive constant number.
Thus, $L(W)$ is $\alpha$-strongly convex in $R_{W}$. Moreover, according to Lemma F.6, $L(W)$ is $l$-smooth. Then according to Lemma C.2, applying gradient descent with $\eta=1/l$, for any $t\geq 1$, we have

$$ $\displaystyle\|W^{t}-W^{*}\|_{F}^{2}\leq\exp(-t/\kappa)\cdot\|W^{0}-W^{*}\|_{F }^{2},$ $$

where $\kappa=l/\alpha$.

After $T\geq 2\kappa\log(N\cdot\|W^{0}-W^{*}\|_{F})$ gradient steps, we have $\widehat{W}=W^{T}=c(\Lambda^{-1}+G+H^{T}/c)=2(\Lambda^{-1}+\widehat{G})$, where $\widehat{G}=G+H^{T}/c$, $\|H^{T}\|_{\max}\leq\exp(-T/\kappa)\cdot\|W^{0}-W^{*}\|_{F}^{2}\leq 1/N$.
Thus, $\|\widehat{G}\|_{\max}\leq\|G\|_{\max}+\|H^{T}\|_{\max}=O(c/N)$.
$\blacksquare$

## Appendix G In-context inference of multi-class classification

### G.1 Notations

In this section, we use the following notations.
We denote
$\mu=(\mu_{1},\mu_{2},\dots,\mu_{c})$,
$q=x_{\mathsf{query}}$.
Define $p_{k}=\frac{c}{M}\sum_{i=1}^{M}(y_{i})_{k}x_{i}$, and define $P=(p_{1},p_{2},\dots,p_{c})\in\mathbb{R}^{d\times c}$. We have $P^{\top}=\frac{c}{M}\sum_{i=1}^{M}y_{i}x_{\tau,i}^{\top}\in\mathbb{R}^{c\times
d}$.
Since with probability $\mathbb{P}\left(y_{\tau,i}=\textbf{e}_{k}\right)=1/c$, $x_{\tau,i}=\mu_{k}+v_{i}$, where $v_{i}\sim\mathsf{N}(0,\Lambda)$, we have
$p_{k}=\frac{c}{M}\sum_{i=1}^{M}(y_{\tau,i})_{k}x_{\tau,i}=cM_{k}\mu_{k}/M+g_{k}$, where $g_{k}=\frac{c}{M}\sum_{i\in\{i|y_{\tau,i}=\textbf{e}_{k}\}}v_{i}$, $g_{k}\sim\mathsf{N}(0,c^{2}M_{k}\Lambda/M^{2})$ and $(M_{1},M_{2},\dots,M_{c})\sim\mathsf{Multin}(M,1/c)$.
Defining $h_{k}=M_{k}/M-1/c$, we have $M_{k}/M=1/c+h_{k}$ and $p_{k}=\mu_{k}+ch_{k}\mu_{k}+g_{k}$.

###### Theorem G.1 (Formal statement of Theorem 4.2 )

Let $\widehat{y}_{\mathsf{query}}$ be the prediction of the trained transformer with parameters $\widehat{W}$ in (19) and $P_{\mathsf{test}}$ satisfying Assumption
4.2, and let $y_{\mathsf{query}}\sim\mathcal{P}^{m}_{y|x_{\mathsf{query}}}(\mu,\Lambda)$. Then, for the inference error defined in (3), we have

$$ $\displaystyle\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}})]$ $\displaystyle\leq$ $\displaystyle\max_{k\in[c]}\left\{\sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}\left[\|\widehat{G}\|_{\max}\sum_{i,j\in[d]}|(\mu_{l})_{i}q_{j} |+\frac{1}{M^{1/2}}\left(\sqrt{c(1-1/c)}|\mu_{l}^{\top}\Lambda^{-1}q|+\sqrt{c} \sum_{i,j\in[d]}|\Lambda^{-1/2}_{ij}q_{j}|\right)\right]\right\}$ $\displaystyle+o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right),$ $$

where $a=\mu^{\top}\Lambda^{-1}q$, $a_{k}=\mu_{k}^{\top}\Lambda^{-1}q$. The expectation is taken over $\{x_{i},y_{i}\}_{i=1}^{M}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}
\mathcal{P}^{m}(\mu,\Lambda)$.

###### Theorem G.1 (Formal statement of Theorem 4.2 )

### G.2 Proof of Theorem 4.2

Proof
The output of the trained transformer is

$$ $\displaystyle\widehat{y}_{\mathsf{out}}=\mathrm{softmax}\left(\left(\frac{c}{M }\sum_{i=1}^{M}y_{i}x_{i}^{\top}\right)(\Lambda^{-1}+\widehat{G})x_{\mathsf{ query}}\right)=\mathrm{softmax}(P^{\top}(\Lambda^{-1}+\widehat{G})q)$ (58) $$

The probability of $y_{\mathsf{query}}=\textbf{e}_{k}$ given $x_{\mathsf{query}}$ is

$$ $\mathbb{P}\left(y_{\mathsf{query}}=\textbf{e}_{k}|x_{\mathsf{query}}\right)= \mathrm{softmax}(\mu^{\top}\Lambda^{-1}x_{\mathsf{query}})_{k}=\mathrm{softmax }(\mu^{\top}\Lambda^{-1}q)_{k}$ $$

Defining $a=\mu^{\top}\Lambda^{-1}q$, $b=(\mu+\mu_{h}+g)^{\top}\widehat{G}q+(\mu_{h}+g)^{\top}\Lambda^{-1}q$, $a_{k}=\mu_{k}^{\top}\Lambda^{-1}q$, $b_{k}=(\mu_{k}+ch_{k}\mu_{k}+g_{k})^{\top}\widehat{G}q+(ch_{k}\mu_{k}+g_{k})^{
\top}\Lambda^{-1}q$,
we have

$$ $\displaystyle\mathbb{E}\left[\mathrm{softmax}(P^{\top}(\Lambda^{-1}+\widehat{G })q)_{k}\right]=\mathbb{E}\left[\zeta_{k}(a+b)\right]=\mathbb{E}[\zeta_{k}(a)+ \sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{l}}b_{l}+\sum_{l,n=1}^{c }R_{kln}(a,b)b_{l}b_{n}/2]$ $$

where $|R_{kln}(a,b)|\leq\sup_{x}|\frac{\partial^{2}\zeta_{k}(x)}{\partial x_{l}
\partial x_{n}}|$. Thus, we have

$$ $\displaystyle\mathbb{E}[\left|\zeta_{k}(a+b)-\zeta_{k}(a)\right|]\leq\mathbb{E }\left[\sum_{l=1}^{c}\left|\frac{\partial\zeta_{k}(a)}{\partial a_{l}}b_{l} \right|\right]+\mathbb{E}\left[\left|\sum_{l,n=1}^{c}R_{kln}(a,b)b_{l}b_{n}/2 \right|\right].$ $$

We first consider the term $\mathbb{E}\left[\sum_{l=1}^{c}\left|\frac{\partial\zeta_{k}(a)}{\partial a_{l}
}b_{l}\right|\right]$.
Defining $\bar{g}_{l}=\Lambda^{-1/2}g_{l}$, we have

$$ $\displaystyle\mathbb{E}\left[\sum_{l=1}^{c}\left|\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}b_{l}\right|\right]$ $\displaystyle\leq$ $\displaystyle\sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{l}}\left(| \mu_{l}^{\top}\widehat{G}q|+\mathbb{E}[|ch_{l}\mu_{l}^{\top}\widehat{G}q|]+ \mathbb{E}[|g_{l}^{\top}\widehat{G}q|]+\mathbb{E}[|ch_{l}\mu_{l}^{\top}\Lambda ^{-1}q|]+\mathbb{E}[|g_{l}^{\top}\Lambda^{-1}q|]\right)$ $\displaystyle\overset{(a)}{\leq}$ $\displaystyle\sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{l}}\left(| \mu_{l}^{\top}\widehat{G}q|+\frac{\sqrt{c(1-1/c)}}{M^{1/2}}|\mu_{l}^{\top} \widehat{G}q|+\mathbb{E}[|\bar{g}_{l}^{\top}\Lambda^{1/2}\widehat{G}q|]+\frac{ \sqrt{c(1-1/c)}}{M^{1/2}}|\mu_{l}^{\top}\Lambda^{-1}q|+\mathbb{E}[|\bar{g}_{l} ^{\top}\Lambda^{-1/2}q|]\right)$ $\displaystyle\overset{(b)}{\leq}$ $\displaystyle\sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{\partial a_{l}}\left[\| \widehat{G}\|_{\max}\sum_{i,j\in[d]}|(\mu_{l})_{i}q_{j}|+\frac{1}{M^{1/2}} \left(\sqrt{c(1-1/c)}|\mu_{l}^{\top}\Lambda^{-1}q|+\sqrt{c}\sum_{i,j\in[d]}| \Lambda^{-1/2}_{ij}q_{j}|\right)\right]$ $\displaystyle+o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right),$ $$

where $(a)$ is due to Lemma F.1 that $\mathbb{E}[|h|]\leq M^{-1/2}c^{-1/2}(1-1/c)^{1/2}$. $(b)$ is because that $\bar{g}_{l}\sim\mathsf{N}(0,c^{2}M_{l}I_{d}/M^{2})$, $\mathbb{E}[\left|(\bar{g}_{l})_{i}\right|]\leq\mathbb{E}[(\bar{g}_{l})_{i}^{2}
]^{1/2}=(c/M)^{1/2}$, for $l\in[c],i\in[d]$.

For $\mathbb{E}\left[\left|\sum_{l,n=1}^{c}R_{kln}(a,b)b_{l}b_{n}/2\right|\right]$, we have

$$ $\displaystyle\mathbb{E}\left[\left|\sum_{l,n=1}^{c}R_{kln}(a,b)b_{l}b_{n}/2 \right|\right]=O(1)\mathbb{E}\Biggr{[}\sum_{l,n=1}^{c}\Biggr{(}\underbrace{ \sum_{\phi_{1}\in\{\mu_{l},ch_{l}\mu_{l},g_{l}\},\phi_{2}\in\{\mu_{n},ch_{n} \mu_{n},g_{n}\}}\left|\phi_{1}^{\top}\widehat{G}q\phi_{2}^{\top}\widehat{G}q \right|}_{(i)}$ $\displaystyle+\underbrace{\sum_{\phi_{1}\in\{\mu_{l},ch_{l}\mu_{l},g_{l}\}, \phi_{2}\in\{ch_{n}\mu_{n},g_{n}\}}\left|2\phi_{1}^{\top}\widehat{G}q\phi_{2}^ {\top}\Lambda^{-1}q\right|}_{(ii)}+\underbrace{\sum_{\phi_{1}\in\{ch_{l}\mu_{l },g_{l}\},\phi_{2}\in\{ch_{n}\mu_{n},g_{n}\}}\left|\phi_{1}^{\top}\Lambda^{-1} q\phi_{2}^{\top}\Lambda^{-1}q\right|}_{(iii)}\Biggr{)}\Biggr{]}.$ $$

For terms $(i)$ having two $\widehat{G}$, they are at most smaller than $O(\|\widehat{G}\|_{\max}^{2})=O(1/N^{2})$. For terms $(ii)$ having one $G$, these terms must contain $n_{1j}$ number of $h_{j}$ and $n_{2ji}$ number of $(\bar{g}_{j})_{i}$, we have
$\sum_{j\in[c],i\in[d]}n_{1j}+n_{2ji}=n_{t},n_{t}=1,2$. According to Lemma F.2, we know that for $n_{t}=1,2$,

$$ $\displaystyle\mathbb{E}[\prod_{j\in[c],i\in[d]}\left|h_{j}^{n_{1j}}(\bar{g}_{j })_{i}^{n_{2ji}}\right|]=O(M^{-1/2}).$ $$

Thus, terms in (ii) are at most smaller than $O(\|G\|_{\max}M^{-1/2})=O(1/(N\sqrt{M}))$.
For terms $(iii)$ without $G$, these terms must contain $n_{1j}$ number of $h_{j}$ and $n_{2ji}$ number of $(\bar{g}_{j})_{i}$, we have
$\sum_{j\in[c],i\in[d]}n_{1j}+n_{2ji}=n_{t},n_{t}=2$.
According to Lemma F.2, for $n_{t}=2$, we have

$$ $\displaystyle\mathbb{E}[\prod_{j\in[c],i\in[d]}|h_{k}^{n_{1k}}(\bar{g}_{j})_{i }^{n_{2ji}}|]=O(M^{-n_{t}/2})=O(M^{-1}).$ $$

Thus, these term are $O(M^{-1})$. Therefore, we have $\mathbb{E}\left[\left|\sum_{l,n=1}^{c}R_{kln}(a,b)b_{l}b_{n}/2\right|\right]=O
(1/N^{2}+1/M+1/(N\sqrt{M}))=o(1/N+1/\sqrt{M})$.

Finally, we have

$$ $\displaystyle\mathbb{E}[\Delta(y_{\mathsf{query}},\widehat{y}_{\mathsf{query}} )]=\max_{k}\{\mathbb{E}[|\mathrm{softmax}(a+b)_{k}-\mathrm{softmax}(a)_{k}|]\}$ $\displaystyle\leq$ $\displaystyle\max_{k\in[c]}\left\{\sum_{l=1}^{c}\frac{\partial\zeta_{k}(a)}{ \partial a_{l}}\left[\|\widehat{G}\|_{\max}\sum_{i,j\in[d]}|(\mu_{l})_{i}q_{j} |+\frac{1}{M^{1/2}}\left(\sqrt{c(1-1/c)}|\mu_{l}^{\top}\Lambda^{-1}q|+\sqrt{c} \sum_{i,j\in[d]}|\Lambda^{-1/2}_{ij}q_{j}|\right)\right]\right\}$ $\displaystyle+o\left(\frac{1}{N}+\frac{1}{\sqrt{M}}\right).$ $$

$\blacksquare$

###### Remark G.1

We note that Theorem 4.2 requires Assumption 4.2 to hold.
For example, we need the covariance $\Lambda$ in training and testing to be the same. A similar consistency requirement of the covariance $\Lambda$ in training and testing had also been observed for in-context linear regression in Zhang et al. [] and for in-context binary classification in the previous section 3.2.

Here, we discuss the consequences when Assumption 4.2 does not hold. For example, suppose the labels of our data in test prompts are not balanced $\mathbb{P}\left(y=\textbf{e}_{k}\right)=p_{k}$, $\mu$ do not have the same $\Lambda^{-1}$ weighted norm $\mu_{k}^{\top}\Lambda^{-1}\mu_{k}\triangleq\Psi_{k}$, and the covariance matrix of test data is $\Gamma\neq\Lambda$, then as $N,M\to\infty$, we have

$$ $\frac{c}{M}\sum_{i=1}^{M}y_{i}x_{i}^{\top}\to c(p_{1}\mu_{1},p_{2}\mu_{2}, \dots,p_{c}\mu_{c})^{\top},$ $$

and

$$ $\mathbb{P}\left(\widehat{y}_{\mathsf{query}}=1\right)\to\mathrm{softmax}(c(p_{ 1}\mu_{1},p_{2}\mu_{2},\dots,p_{c}\mu_{c})^{\top}\Lambda^{-1}x_{\mathsf{query} }).$ $$

Denote $\Psi=(\Psi_{1},\dots,\Psi_{c})^{\top}$, $\Phi=(\log(p_{1}),\dots,\log(p_{c}))^{\top}$ and $z=\mu^{\top}\Gamma^{-1}x_{\mathsf{query}}-\Psi/2+\Phi$. Then distribution of the ground truth label is

$$ $\mathbb{P}\left(y_{\mathsf{query}}=\textbf{e}_{k}\right)=\mathrm{softmax}(z)_{ k}.$ $$

Define $\hat{z}=c(p_{1}\mu_{1},p_{2}\mu_{2},\dots,p_{c}\mu_{c})^{\top}\Lambda^{-1}x_{
\mathsf{query}}$.
Then, unless $\hat{z}=z$ or $\|\mathrm{softmax}(\hat{z})-\mathrm{softmax}(z)\|_{2}$ is sufficiently small, the transformer cannot correctly perform the in-context multi-class classification.

###### Remark G.1

## Appendix H Experiment details

For all tasks, we set $d=20$ and we randomly generate a covariance matrix $\Lambda=\text{diag}(\lambda_{1},\dots,\lambda_{d})$, where $\lambda_{i}=|\hat{\lambda}_{i}|$ and $\hat{\lambda}_{i}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{N}(3,1)$. For each training dataset with different training prompt lengths $N$, and different class numbers $c$, we randomly generate $B$ training samples. Training prompts $P_{\tau},\tau\in[B]$ and their corresponding labels $y_{\tau,\mathsf{query}}$ are generated according to Assumption 4.1. Moreover, we also generate testing datasets. For example, for each testing dataset, we first randomly generate 20 pairs of $(\mu_{j},x_{j,\mathsf{query}},y_{j,\text{prob}}),j\in[20]$, where $(\mu_{j})\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}^{m}_{
\Omega}(\Lambda)$, $x_{j,\mathsf{query}}\sim\mathcal{P}^{m}_{x}(\mu_{j},\Lambda)$. $y_{j,\text{prob}}=\mathrm{softmax}(\mu_{j}^{\top}\Lambda^{-1}x_{j,\mathsf{
query}})$ are the corresponding probability distributions of the ground truth label $y_{j,\mathsf{query}}$. For each $j$, we generate 100 testing prompts $P_{jk}=(x_{jk,1},y_{jk,1},\dots,x_{jk,M},y_{jk,M},x_{j,\mathsf{query}})$, where $(x_{jk,i},y_{jk,i})\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{P}
^{m}(\mu_{j}\Lambda),j\in[20],k\in[100],i\in[M]$.
We denote a model’s output for testing prompts $P_{jk}$ as $\widehat{y}_{jk}$.
We calculate its inference error with $\frac{1}{20\times 100}\sum_{j\in[20],k\in[100]}\max_{l\in[c]}\left|\widehat{(}
y_{jk})_{l}-(y_{j,\text{prob}})_{l}\right|$, which serves an approximation of the expected total variation distance we defined in (3).

For experiments in Figure 1, we train the single-layer transformers with the sparse-form parameters and structures defined in Section 4. We set the size of the training dataset to $B=10,000$ and set the batch size to 50. We train the transformers using SGD with learning rate $\{0.1,0.5,1\}$ for 10 epochs, and get the best model on each training dataset. Then, we test these trained models on different testing datasets. Each experiment is repeated 10 times with different random seeds. The mean results and standard deviation error bars of these 10 experiments are plotted in Figure 1.

For experiments in Figure 2, the structure of the transformer with full parameters (1-layer, full) is defined as

$$ $F(E;W^{V},W^{KQ})=E+W^{V}E\cdot\frac{E^{\top}W^{KQ}E}{\rho},$ (59) $$

where $W^{V},W^{KQ}\in\mathbb{R}^{(d+c)\times(d+c)}$ are the parameters for optimization. For the (3-layer, GPT2) model, we use the GPT2 architecture [Radford, ] with 64 embedding sizes, 3 layers and 2 heads as implemented by HuggingFace [Wolf et al., ], and we use the similar embedding method proposed by Garg et al. []. Note that, for the GPT2 model trained on prompt length $N=100$, it will produce predictions for the labels $\widehat{y}_{i}$, $i\in[101]$. Assume the ground truth labels for the prompt is $y_{i}$, $i\in[101]$,
usually, for decoder models like GPT2, we should define the loss function as $\text{Loss}=\sum_{i=1}^{101}\text{CrossEntropyLoss}(\widehat{y}_{i},y_{i})$, which can be viewed to some extent as the model is trained with prompts length ranging from 0 to 100. However, for in-context classification of Gaussian mixtures, longer prompts are significantly more helpful for training models than shorter prompts. For example, considering an extreme case where the classification task involves $c$ classes and the training prompt length is smaller than $c$, then it is nearly impossible for the model to learn to classify $c$ classes Gaussian mixtures in context with these short prompts. Our Theorems 3.1 and 4.1 also clearly demonstrate that the training error is inversely proportional to the prompt length with a scaling of $O(1/N)$. Thus, if we define the loss function for GPT2 as $\text{Loss}=\sum_{i=1}^{101}\text{CrossEntropyLoss}(\widehat{y}_{i},y_{i})$, the underlying short prompts will be detrimental to the model training.
Therefore, in our experiments, we define the loss function for our GPT2 model as $\text{Loss}=\sum_{i=41}^{101}\text{CrossEntropyLoss}(\widehat{y}_{i},y_{i})$, which can be viewed to some extent as the model that is trained with prompts length ranging from 40 to 100.
For all three transformer models, we set the size of the training dataset to $B=400,000$ and set the batch size to 50. We train the (1-layer, sparse) and (1-layer, full) using Adam with learning rate 0.001 for 5 epochs, and train the GPT2 model using Adam with learning rate 0.0001 for 5 epochs.
Each experiment is repeated 3 times with different random seeds. The mean results and standard deviation error bars of these 3 experiments are plotted in Figure 2.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 8
presentation: 3
questions: Does the same training dynamics results apply to tasks beyond linear regression/classification? Would it be different for other mathematical/statistical tasks, for example, time series, etc.
soundness: 4
strengths: 1. Rigorous theory: The paper provides a detailed theoretical analysis of the convergence properties of transformers for in-context classification tasks, demonstrating that under certain conditions, a single-layer transformer trained by gradient descent achieves global optimality at a linear rate.

2. Experimental validation: The theoretical claims are corroborated by experimental results. The paper's experiments on binary and multi-class classification tasks with Gaussian mixtures verify the theoretical predictions, showing that the transformers' prediction accuracy improves as the training and testing prompt lengths increase.
summary: The paper investigates the training convergence of transformers for in-context classification tasks. It demonstrates that a single-layer transformer trained by gradient descent converges to a globally optimal model at a linear rate for in-context classification of Gaussian mixtures. Experimental results confirm the theoretical findings, showing that the trained transformers perform well in binary and multi-class classification tasks.
weaknesses: The paper primarily focuses on single-layer transformers with simplified linear attention mechanisms. While it provides valuable insights into the convergence properties and error bounds for these models, the findings may not fully extend to more complex multi-layer transformers with softmax or relu attention mechanisms. And the training dynamics of multi-layer transformers could be different to single-layer -- just like multi-layer MLP can be different than linear models.

Nonetheless, I don't regard this as a strong weakness since the field is evolving and to my best knowledge, most studies are still on one-layer transformers.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.The condition (2) in Assumption 3.1 seems unusual to me. Could the authors provide more clarification on this assumption?

2.Some papers[1,2,3] have highlighted emergent behaviors in the training dynamics of in-context learning. However, this paper asserts that the transformer will converge to its global minimizer at a linear rate, which appears to contradict those findings. Can the authors discuss this further?

[1] In-context learning and induction heads

[2] Breaking through the learning plateaus of in-context learning in Transformer

[3] Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality.
soundness: 3
strengths: 1. The structure of this paper is very clear.
2. Analyzing the training dynamics of in-context learning is crucial.
3. The findings regarding infinite lengths of training prompts (N) and testing prompts (M) are interesting
summary: This paper study the learning dynamics of transformers for in-context classification of Gaussian mixtures and prove the training convergence of in-context multi-class classification. The author presents three key findings: 1) a proof that a single-layer transformer trained via gradient descent converges to a globally optimal solution at a linear rate; 2) an analysis of how the lengths of training and testing prompts influence the inference error in in-context learning; and 3) evidence that, with sufficiently long training and testing prompts, the predictions of the trained transformer approach those of the Bayes-optimal classifier. Some of these results are validated through experiments.
weaknesses: 1) This paper claims to be the first to explore the learning dynamics of transformers for in-context classification of Gaussian mixtures and to prove the convergence of training in multi-class classification. However, I find the significance of this assertion unclear, as the paper lacks sufficient detail. Specifically: 1) many prior works have analyzed in-context learning assuming $x $ comes from Gaussian distributions; what additional insights do the results on Gaussian mixtures provide? 2) Why is extending results from binary to multi-class classification considered essential and non-trivial?

2) Additionally, I have concerns regarding the techniques used: 
- The introduction of $ \tilde{L} $ appears to be a key element in proving Theorem 3.1, but its intuition is unclear, and I'm uncertain how it addresses the challenges posed by the non-linear loss function. 
- The paper heavily relies on Taylor expansion in its proofs, and I question whether this expansion can accurately approximate the original function. More detail is needed on this aspect.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. For data distribution, why is it essential to preserve the inner product of vectors in the $\Lambda^{-1}$-weighted norm? Is this primarily a technical consideration? It would be helpful if the authors could provide further clarification on the role of data distribution in the analysis.

2. For the inference stage, while $\mu_0$ and $\mu_1$ are not subject to additional constraints, $\Lambda$ remains fixed, imposing a strong assumption on the underlying structure of the data distribution. Do the authors have insights on how these results might extend to scenarios with a varying $\Lambda$ during inference?

3. The derived inference bound scales with $N$ and $M$ similarly to in-context regression [1]. Could the authors clarify the distinctive aspects of the multi-classification setting in this context? (This also points to weakness.)

4. For the multi-classification setting, what is the order of the number of classes $c$? On line 431, the authors mention that $c$ is treated as a constant coefficient—would a larger order of $c$ impact the analysis?








[1] Trained Transformers Learn Linear Models In-Context. Zhang et al., 2023
soundness: 3
strengths: 1. This work is the first to examine transformers for in-context multi-class classification from a training dynamics perspective.
2. The end-to-end connection between in-context inference and training offers a novel  insight.
3. Experimental results are provided to support theoretical statements.
summary: This work studied the training dynamics of a one-layer linear transformer trained via GD for in-context multi-classification tasks. They established  convergence guarantees for in-context trianing and also provided an in-context inference error bound, which scales with the length of both training and testing prompts.
weaknesses: My major concern lies in the analytical novelty of this paper compared to prior work on in-context regression [1]. While this study focuses on the multi-class classification problem, its model and analytical approach appear to share many similarities with [1]. It remains unclear how technically straightforward it is to generalize the results of [1] to multi-classification. Additionally, this paper restricts itself to the linear attention setting, simplifying the analysis and making it somewhat less impactful than [2], which addresses binary classification with strict data assumptions but in the more realistic softmax attention setting. Therefore, a thorough discussion clarifying the technical distinctions and contributions of this work relative to these previous studies would be helpful.




[1] Trained Transformers Learn Linear Models In-Context. Zhang et al., 2023

[2] Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. Li et al., 2024

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Binary classification is a special case of the multi-class classification. Why write up both?
- ICL in LLMs is a type of domain adaptation setting, where knowledge about a particular task is scarce in the pre-training corpus. The model needs to be "kindled" with ICL demos to get it to perform better on this task. In contrast, the presented theory first trains the transformers on samples of these tasks, and then requires this data to be "properly distributed" over the space of variations. It is not surprising that the presented results hold in this setting. How do you think the setup needs to change to reflect that realistic ICL in LLMs setting?
- With respect to the following: 
> probably because our transformer models were only trained with a small prompt length of N = 100.

  3 layer GPT2 seems like a small model, why not test with higher N?
soundness: 3
strengths: - The relation between prompt length and error is somewhat interesting because as far as I know, previous works studying this meta-learning capabilities of the transformer model (misnamed as ICL) do not talk about this learnability constraint for open ended problems. However, other works have looked at something similar [[link]](https://arxiv.org/pdf/2303.07971).
- The other claim about being the first to study multi-class classification maybe true. Its significance is unclear.
summary: Yet another ICL theory paper that does not study actual ICL.

This paper, like many previous works, studies the training dynamics of a simple (single layer, linear) transformer model trained using ICL objective. They work with Gaussian data to study the training convergence rates; and impact of prompt length on error. With these settings, they find the rather unamusing results of linear convergence rates, and Bayes optimal behavior with asymptotic prompt lengths. No connection is made to emergent ICL in real LLMs whatsover. The authors should look at this recent [ICML position paper](https://arxiv.org/pdf/2310.08540) to find the distinction between the two.

**Disclaimer**: I have not read the proofs in detail to verify that they are correct (hence this review is rated at confidence level 4). My review is based on the assumption that the proofs are correct.
weaknesses: - The setting is too unrealistic to say anything about real ICL. For example,
  - Training on hard-coded ICL prompts, when LLMs are trained on next-word prediction (ICL structure is generally not present in the pretraining corpus). This is a major setup difference which makes them incompatible.
  - Studying single layer transformers with no non-linear activation functions. This is a good intellectual curiosity but its relevance and usefulness in understanding ICL remains unclear (even classic deep learning theory struggles to present useful insights by studying 2-layer networks). In this paper itself, we see a deviation from expectations when a 3-layer GPT2 architecture with softmax is tested (section 5.2).
  - Gaussian data that presents fixed one-token length inputs and outputs. I don't have a problem with Gaussian data, but the framework should be flexible enough to even somewhat resemble real ICL (where the inputs and outputs both can be variable lengths).

Meta-review: This paper provides a theoretical analysis of in-context learning for classification tasks. The authors use a fairly standard setting of 1-layer linear attention. However the distinction of the work arises from its focus on binary and multiclass classification tasks and the associated novel analysis. The authors study these problems under a gaussian mixture dataset model. There was an insightful discussion between the reviewers and the authors on whether this work captures the in-context capabilities of modern LLMs. While the AC acknowledges that the 1-layer linear attention model of this work (and many other in-context learning theory works) is simplistic compared to modern LLMs, purely theoretical contributions are valuable and welcome to ICLR and can provide a stepping stone toward more sophisticated models.

While most reviewers found the technical contribution of the paper to be decent, the final recommendation is reject for the following reasons:

1) **Technical concerns:** There are some basic flaws that requires a second review. Firstly, the definition of Bayes-optimal classifier in Line 294 is incorrect. Bayes-optimal classifier is deterministic given the input so the correct classifier is not a probability but is obtained by applying the sign function. Related to this, Theorem 3.2 states a total variation distance guarantee between $y_{query}$ and $yhat_{query}$ which goes to zero as $N,M\rightarrow\infty$. However the authors use a less strict TV distance definition rather than the conventional one. Namely it does not mean that $y_{query}$ and $yhat_{query}$ are the same random variables or $yhat_{query}$ is the optimal decision, they just have identical distributions. By the definition of authors as stated in Line 270, $yhat_{query}$ is not the Bayes optimal decision because it introduces noise on the optimal classifier during sampling from $yhat_{out}$.

2) **Finite sample Bayes optimality of GMMs:** The authors discuss Bayes optimality only in the asymptotic sense as prompt length is infinite. In reality, even for finite prompt length, under suitable assumptions, 1-step gradient descent can be finite sample Bayes optimal estimator for binary GMMs. For instance, see Section 2.2 of [Mignacco et al. ICML 2020](https://proceedings.mlr.press/v119/mignacco20a/mignacco20a.pdf). This means one layer attention can do optimal classification under finite prompt length. This work does not discuss or capture this important aspect.

3) **Related work:** The second point above brings me to the related work section which needs substantial improvement.

- The authors are **missing any reference** to the literature on gaussian mixture models or classification with GMMs even if their results are heavily relying on gaussian mixture assumption. This even includes a work on in-context learning which similarly utilizes GMM data such as [Dual Operating Modes of In-Context Learning (ICML'24)](https://arxiv.org/pdf/2402.18819). I find it a bit unfortunate that most citations are to ICL/LLM papers within the last 2-3 years and not much to classical ML literature. Note that, once we make the assumption in Eq (14), we end up with the statistical properties of 1 step gradient estimator on gaussian mixture data (Line 151) which is essentially the plug-in estimator in Section 2.2 of [Mignacco et al. ICML 2020](https://proceedings.mlr.press/v119/mignacco20a/mignacco20a.pdf). In-context learning essentially constitutes a proxy for this fundamental model. I recommend that the authors should consider providing a thorough discussion on GMM literature (prior works on meta learning with GMMs, finite sample learning, Bayes optimal rates, multiclass etc). This would also provide better motivation for their assumptions.

- There is also no related work section in the main body. I would advise inserting 0.5 page of (shortened) related work in the final manuscript. I believe 10 pages provide enough space to do so.


4) **Technical clarity (minor concern):** Some of the notation should be introduced more clearly. For instance, if I am not mistaken, the $G$ matrix in Theorem 3.1 is essentially defined in terms of $W^*$, that is, its context is missing. Please go over the manuscript carefully to ensure technical clarity throughout.

Decision: Reject