Paper deep dive
In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention
Katsuyuki Hagiwara
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/20/2026, 3:47:44 AM
Summary
This paper investigates the in-context learning capabilities of transformers with linear self-attention for simple linear regression tasks. It demonstrates that such transformers can approximate the closed-form least squares solution using layer normalization, rather than relying on gradient descent algorithms. The authors construct a specific transformer architecture (TF(4,2,2)) and provide numerical evidence that when trained with L1 regularization, the model primarily learns to implement this analytical solution.
Entities (7)
Relation Signals (5)
Linear self-attention → partof → Transformer
confidence 96% · construct a transformer with linear self-attention
Transformer → implements → Least Squares Estimate
confidence 95% · construct a transformer with linear self-attention, which in-context learns the least squares estimate
In-Context Learning → exhibitedby → Transformer
confidence 94% · In-context learning is a remarkable property of transformers
Layer Normalization → enables → Closed Form Solution
confidence 93% · the closed form (analytical) solution is approximately obtained by using layer normalization
Transformer → trainedwith → L1 Regularization
confidence 88% · our implementation is mainly used in the transformer trained with l1 regularization
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solver for linear and non-linear regression problems, in which the most of them implement gradient descent algorithm. However, it is still unclear whether those implementations have actually been acquired through training. In this paper, we construct a transformer with linear self-attention, which in-context learns the least squares estimate in a simple regression task. The point here is that the closed form (analytical) solution is approximately obtained by using layer normalization rather than an approximate solution based on gradient descent algorithm. Then, we show an experimental example, in which our implementation is mainly used in the transformer trained with l1 regularization when the target output is the least squares estimate.
Tags
Links
- Source: https://arxiv.org/abs/2607.15819v1
- Canonical: https://arxiv.org/abs/2607.15819v1
Trouble viewing inline? Open PDF directly →
Full Text
59,732 characters extracted from source content.
Expand or collapse full text
In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention Katsuyuki Hagiwara (Faculty of Education, Mie University 1577 Kurima-Machiya-cho, Tsu, 514-8507, Japan hagi@edu.mie-u.ac.jp) Abstract In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solver for linear and non-linear regression problems, in which the most of them implement gradient descent algorithm. However, it is still unclear whether those implementations have actually been acquired through training. In this paper, we construct a transformer with linear self-attention, which in-context learns the least squares estimate in a simple regression task. The point here is that the closed form (analytical) solution is approximately obtained by using layer normalization rather than an approximate solution based on gradient descent algorithm. Then, we show an experimental example, in which our implementation is mainly used in the transformer trained with ℓ1 _1 regularization when the target output is the least squares estimate. Keyword in-context learning, linear self-attention, a simple linear regression, layer normalization 1 Introduction In-context learning is a remarkable property of transformers, which form the basis of large language models such as GPT-3 [3], and it has been the focus of recent research. Through in-context learning, given a prompt containing examples of a task and a new query input, the trained language model can generate the corresponding output for the new query in a one-shot manner. It is natural to think that transformers acquire, through training, the algorithms to solve tasks via in-context learning. In this regard, it has been demonstrated that transformers are capable of implementing various algorithms, particularly in regression tasks [5, 11, 1, 4, 2]. [5] empirically studied the in-context learning abilities of transformers for various function classes in machine learning, including the linear function class. In particular, for linear functions, a trained transformer performs similarly to the least squares solution. [11] provided an explicit construction of a linear self-attention layer that implements a single step of the gradient descent algorithm on the mean squared error loss. Additionally, they empirically showed that several self-attention layers can iteratively perform curvature correction, improving upon the plain gradient descent algorithm. [1] proved that a transformer can implement a gradient descent algorithm and a closed form solution for ridge regression. [4] also pointed out the correspondence between the linear version of attention and the gradient descent algorithm, claiming that transformers perform implicit fine-tuning. They also empirically investigated the similarity between in-context learning and explicit fine-tuning. While the works of [5, 11, 1, 10, 4] do not consider the training phase, [13] investigated the learning dynamics of a gradient flow in a simplified transformer architecture when the training prompts consist of random instances of linear regression datasets, concluding that transformers trained by a gradient flow in-context learn a class of linear functions. More recently, [2] showed that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, and Lasso. In contrast to [1], [2] precisely evaluated the prediction performance in terms of network size, and shown a near-optimal predictive power. [2] also demonstrated the algorithm selection ability of transformers, such as regularization selection according to validation error for ridge regression. In these works, however, it still remains unclear whether those algorithms are actually obtained through the training process of transformers. In this paper, for a simple regression task, we first construct a transformer with linear-self attention, which implements an approximate closed form (analytical) solution to the least squares estimate according to [1]. We then present a numerical example, in which this implementation is mainly used in the transformer trained with the ℓ1 _1 regularization when the target is the least squares estimate. Previous works [5, 11, 1, 4, 2] have investigated regression problems, finding that the solver implemented by a transformer is basically the gradient descent algorithm. This is natural since the attention mechanism in a transformer computes the product of the inputs, which is exactly what is required for the gradient descent method. In a linear regression problem, however, the least squares estimate is analytically obtained by calculating the matrix inverse, which requires division. Consequently, the gradient descent algorithm implicitly calculates this division by repeatedly performing multiplication and addition. Among these works, [1] showed that a transformer can implement a close form solution to a ridge regression problem, in which the division is implemented using layer normalization. Our construction of the least squares estimate for a simple regression problem is based on this insight. The transformer in our setting comprises a linear transformation of the input, a stack of transformer blocks, a flatten layer, and an output linear transformation, in which the transformer block consists of multi-head linear self-attention blocks with layer normalization followed by a skip connection. Thus, the layer normalization is applied to the sum of the output of the linear self-attention blocks and a skip connection sum is applied to the layer normalized output. This transformer receives the in-context samples and the query input (prediction point) as a prompt, and outputs the least squares estimate for the query input. Specifically, the number of layers, the number of heads, and the model dimension in our construction are 2, 2, and 4 respectively, which is very small. The closed form solution to the least squares estimate for a simple regression problem requires division by the variance of the input data in the in-context samples. We provide a specific construction to approximately calculate this closed form solution under a natural input form, in which the layer normalization is used for performing division according to [1]. After showing this construction, we present numerical experiments to show that our implementation is actually used in the transformer trained with ℓ1 _1 regularization when the target output is the least squares estimate. In other words, through training, the transformer mainly acquires the calculation of the closed form solution rather than the steps of the gradient descent algorithm. This paper is organized as follows. Section 2 formulates our problem setting. Section 3 presents the construction of a transformer that calculates the least squares estimate based on in-context samples. Section 4 provides a numerical example of training to demonstrate that this construction is actually valid. Finally, Section 5 concludes the paper and discusses future work. 2 Problem setting 2.1 Notations In this paper, I,JO_I,J is the I×JI× J zero matrix. Note that we use this notation also for vectors, where either I or J equals 11. For an I×JI× J matrix A, [i,j]A[i,j] is the (i,j)(i,j)-entry of A, [i,:]A[i,:] is the i-th row vector of A and [:,j]A[:,j] is the j-th column vector of A for i=1,…,Ii=1,…,I and j=1,…,Jj=1,…,J. When A is an I×1I× 1 vector, [i]A[i] denotes its i-th entry. 2.2 Setting of in-context learning of transformer We next explain the transformer-based in-context learning of a simple linear regression task. Let M be the number of training data for the transformer. Let (x,y)(x,y) be a pair of input-output variables in a simple regression problem. At each m in 1,…,M\1,…,M\, a set of N random samples of (x,y)(x,y), which is denoted by (xm,n,ym,n):n=1,…,N\(x_m,n,y_m,n):n=1,…,N\, are generated. We denote the m-th prompt (input to the transformer) by mP_m, which is an (N+1)×3(N+1)× 3 matrix and whose n-th row is m[n,:]=[1xm,nym,n], P_m[n,:]= bmatrix1&x_m,n&y_m,n bmatrix, (1) where we define ym,N+1:=0y_m,N+1:=0 and xm,N+1:=umx_m,N+1:=u_m that is a prediction point; e.g., see [1]. Thus, the input sequence length is N+1N+1 and the number of in-context samples is N. We assume that, for each xm,nx_m,n, ym,ny_m,n is generated by ym,n=θm,0+θm,1xm,n+εm,n, y_m,n= _m,0+ _m,1x_m,n+ _m,n, (2) where εm,1,…,εm,N _m,1,…, _m,N, m=1,…,Mm=1,…,M are i.i.d. additive noises from a probability distribution with mean 0 and variance σ2<∞σ^2<∞. Our goal is to obtain the least squares prediction at x=umx=u_m for the prompt mP_m. Therefore, for each m, we need to calculate the least squares solution using m=(xm,n,ym,n):n=1,…,N, D_m=\(x_m,n,y_m,n):n=1,…,N\, (3) which is the set of in-context samples. The important point is that the regression lines can be different for each m. Thus, we may assume that (θm,0,θm,1)( _m,0, _m,1) are sampled from a probability distribution for each m. However, in the construction of the transformer, we do not make any specific assumptions on the underlying probability distribution of (θm,0,θm,1)( _m,0, _m,1). Although we also do not make any specific assumptions on the underlying probability distribution of xm,nx_m,n and εm,n _m,n, we will make an assumption on the in-context samples later. 2.3 Training data for the transformer For the m-th training data mD_m defined in (3), we define x¯m x_m :=1N∑n=1Nxm,n := 1N _n=1^Nx_m,n (4) y¯m y_m :=1N∑n=1Nym,n := 1N _n=1^Ny_m,n (5) Vm V_m :=1N∑n=1N(xm,n−x¯m)2 := 1N _n=1^N(x_m,n- x_m)^2 (6) Cm C_m :=1N∑n=1N(xm,n−x¯m)(ym,n−y¯m). := 1N _n=1^N(x_m,n- x_m)(y_m,n- y_m). (7) For a simple regression problem, it is easy to see that the prediction at x=umx=u_m using the least squares solution based on mD_m is given by y^m(um):=y¯m+CmVm(um−x¯m) y_m(u_m):= y_m+ C_mV_m(u_m- x_m) (8) (see e.g., [14]). Below, we construct the transformer that receives mP_m as an input and outputs y^m(um) y_m(u_m) for m=1,…,Mm=1,…,M. Thus, the transformer calculates the least squares estimate at the prediction point umu_m using in-context samples mD_m at each m. 2.4 Linear self-attention We define a linear self-attention (LSA), which receives an (N+1)×D(N+1)× D matrix Q as an input and outputs (N+1)×D(N+1)× D matrix defined by LSA() LSA_ (Q) :=(3)(1⊤)⊤(2) :=(QW_3)(QW_1 ) (QW_2) =(3)(1⊤2), =(QW_3)(W_1Q QW_2), (9) where =1,2,3 =\W_1,W_2,W_3\ is an ordered set of parameters and 1W_1, 2W_2 and 3W_3 are D×D× D matrices. The operation using this LSA is precisely discussed in [6]. 2.5 Transformer structure The transformer structure considered in this paper is illustrated in Fig. 1, in which grey blocks are operation blocks. The transformer here consists of input linear transformation, stacking transformer blocks, flatten and output linear transformation (Fig. 1 (a)). The transformer block consists of multi-head LSA (MHLSA) block with layer normalization followed by a skip connection (Fig. 1 (b)). It receives the m-th prompt mP_m defined in (1) and, for simplicity, outputs a scalar value whose target is the least squares estimate at a prediction point umu_m, which is y^m(um) y_m(u_m) given by (8). (a) Main stream (b) Transformer block Figure 1: Structure of transformer Let inW_ in and in b_ in be a D×3D× 3 input weight matrix and D×1D× 1 input bias vector respectively. The embedding of the prompt mP_m by an affine transformation with (in,in)(W_ in, b_ in) is denoted by m,1Q_m,1 whose size is (N+1)×D(N+1)× D. More precisely, we define m,1[n,:]⊤:=inm[n,:]⊤+in Q_m,1[n,:] :=W_ inP_m[n,:] + b_ in (10) for n=1,…,N+1n=1,…,N+1. The number of layers is denoted by L. For the m-th training data, the input to the l-th layer is denoted by m,lQ_m,l, l=1,…,Ll=1,…,L. Thus, the above m,1Q_m,1 is the input to the first layer. The parameter of the LSA block of the k-th head in the l-th layer is denoted by l,k=l,k,1,l,k,2,l,k,3 _l,k=\W_l,k,1,W_l,k,2,W_l,k,3\. The output of the MHLSA in the l-th layer for m-th training data is m,l,k:=LSAl,k(m,l) H_m,l,k:= LSA_ _l,k(Q_m,l) (11) m,l:=∑k=1Km,l,k, H_m,l:= _k=1^KH_m,l,k, (12) where K is the number of heads. The size of m,lH_m,l is (N+1)×D(N+1)× D. The number of heads is common for all layers. We employ the layer normalization along the model (embedding) dimension. We define μm,l,n _m,l,n :=1D∑d=1Dm,l[n,d] := 1D _d=1^DH_m,l[n,d] (13) σm,l,n2 σ^2_m,l,n :=1D∑d=1D(m,l[n,d]−μm,l,n)2. := 1D _d=1^D(H_m,l[n,d]- _m,l,n)^2. (14) Then, the output of the layer normalization with these parameters is denoted by ¯m,l H_m,l whose (n,d)(n,d)-entry is ¯m,l[n,d]:=(m,l[n,d]−μm,l,n)/σm,l,n H_m,l[n,d]:=(H_m,l[n,d]- _m,l,n)/ _m,l,n (15) for d=1,…,Dd=1,…,D at each n. Finally, the output of the l-th layer is given by m,l+1=m,l+¯m,l, Q_m,l+1=Q_m,l+ H_m,l, (16) where the first term of the right-hand side comes from a skip connection. Note that, in our model, the layer normalization is applied before the skip connection(see e.g. [1]). m,LQ_m,L is the output of the last layer. It is then flattened and obtain m q_m which is a D(N+1)×1D(N+1)× 1 vector. Let outW_ out and boutb_ out be a D(N+1)×1D(N+1)× 1 output weight matrix and output bias respectively. We then obtain the transformer output by Qout(m):=out⊤m+bout Q_ out(P_m):=W_ out q_m+b_ out (17) which is a scalar. We refer to this transformer as TF(D,L,K) TF(D,L,K), where D, L and K are the model dimension, the number of layers and heads respectively. 3 A construction of transformer We here gives a construction of the transformer that calculates the least squares estimate at the prediction point. To achieve this, we require that |um−x¯m|/N|u_m- x_m|/N and |y¯m|/N| y_m|/N be negligible in the calculation below. If we naturally assume that |um−x¯m||u_m- x_m| and |y¯m|| y_m| do not increase as N increases then this is satisfied when N is sufficiently large. With this case in mind, we assume that N is sufficiently large so that |um|/N|u_m|/N, |x¯m|/N| x_m|/N and |y¯m|/N| y_m|/N are negligible for any m. We will address this assumption in a later section We here fix the structure of the transformer, in which the model dimension D=4D=4, the number of layers L=2L=2 and the number of heads K=2K=2. Therefore, the transformer is TF(4,2,2) TF(4,2,2). Now, the m-th prompt to TF(4,2,2) TF(4,2,2) is mP_m defined in (1). 3.1 Input embedding Let R be a real positive number. We assume that R can be chosen sufficiently large below. As in [1], R is important for approximate implementation of the closed form solution. In the input embedding, we set in=[10001RN0001RN2N+1−1RN−1RN] W_ in= bmatrix1&0&0\\ 0& 1RN&0\\ 0&0& 1RN\\ 2NN+1&- 1RN&- 1RN bmatrix (18) and in=4,1 b_ in=O_4,1. By (10), the n-th row of the input embedding is m,1[n,:]⊤=[11RNxm,n1RNym,n2N+1−1RNxm,n−1RNym,n], Q_m,1[n,:] = bmatrix1\\ 1RNx_m,n\\ 1RNy_m,n\\ 2NN+1- 1RNx_m,n- 1RNy_m,n bmatrix, (19) which is the input to the transformer. 3.2 Output of the first layer We here construct the first layer. We consider the first head. We have ξ1,1,1 _1,1,1 :=m,1[:,1]⊤m,1[:,1]=N+1. :=Q_m,1[:,1] Q_m,1[:,1]=N+1. (20) and ξ1,1,2 _1,1,2 :=m,1[:,1]⊤m,1[:,2] :=Q_m,1[:,1] Q_m,1[:,2] =1RN(Nx¯m+um)≃1Rx¯m, = 1RN(N x_m+u_m) 1R x_m, (21) where the last approximation holds by the assumption. Since ym,N+1=0y_m,N+1=0, we have ξ1,1,3 _1,1,3 :=m,1[:,1]⊤m,1[:,3] :=Q_m,1[:,1] Q_m,1[:,3] =1RNNy¯m=1Ry¯m. = 1RNN y_m= 1R y_m. (22) By the same way as above, we also have ξ1,1,4 _1,1,4 :=m,1[:,1]⊤m,1[:,4] :=Q_m,1[:,1] Q_m,1[:,4] ≃2N−1Rx¯m+−1Ry¯m. 2N- 1R x_m+- 1R y_m. (23) Therefore, we have m,1⊤m,1=[ξ1,1,1ξ1,1,2ξ1,1,3ξ1,1,4−], Q_m,1 Q_m,1= bmatrix _1,1,1& _1,1,2& _1,1,3& _1,1,4\\ -&-&-&-\\ -&-&-&-\\ -&-&-&-\\ bmatrix, (24) where the irrelevant elements are indicated by “−-”. We set 1,1=1,1,1,1,1,2,1,1,3 _1,1=\W_1,1,1,W_1,1,2,W_1,1,3\ in which 1,1,1 W_1,1,1 =[1000000000000000] = bmatrix1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (25) 1,1,2 W_1,1,2 =[N+10000−120000−120000−12] = bmatrix NN+1&0&0&0\\ 0&- 1 2&0&0\\ 0&0&- 1 2&0\\ 0&0&0& -1 2\\ bmatrix (26) 1,1,3 W_1,1,3 =[1000000000000000]. = bmatrix1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix. (27) Then, by (2.4) and (11), the n-th row of the output of the first head is m,1,1[n,:]⊤=[N−12Rx¯m−12Ry¯m−N+12Rx¯m+12Ry¯m]. H_m,1,1[n,:] = bmatrixN\\ - 1 2R x_m\\ - 1 2R y_m\\ -N+ 1 2R x_m+ 1 2R y_m bmatrix. (28) On the other hand, for the second head, we set 1,2=1,2,1,1,2,2,1,2,3 _1,2=\W_1,2,1,W_1,2,2,W_1,2,3\ in which 1,2,1=1,2,2=1,2,3=d,d. W_1,2,1=W_1,2,2=W_1,2,3=O_d,d. (29) Then, by (2.4) and (11), the n-th row of the output of the second head is m,1,2[n,:]=1,4. H_m,1,2[n,:]=O_1,4. (30) Therefore, by (12), the n-th row of the output of the MHLSA in the first layer is m,1[n,:]⊤ H_m,1[n,:] =m,1,1[n,:]⊤ =H_m,1,1[n,:] =[N−12Rx¯m−12Ry¯m−N+12Rx¯m+12Ry¯m] = bmatrixN\\ - 1 2R x_m\\ - 1 2R y_m\\ -N+ 1 2R x_m+ 1 2R y_m bmatrix (31) Now, we consider the layer normalization. By (13), (14) and (3.2), for m,1[n,:]H_m,1[n,:], we have μm,1,n _m,1,n =0 =0 (32) σm,1,n2 σ^2_m,1,n =14(N2+O(1R2)+N2+O(NR))≃N22, = 14 (N^2+O ( 1R^2 )+N^2+O ( NR ) ) N^22, (33) since we can set a sufficiently large value for R; e.g., R=NpR=N^p, where p is a sufficiently large positive integer. By (15), (3.2), (32) and (33), the n-th row of the output of the layer normalization is ¯m,1[n,:]⊤=[2−1RNx¯m−1RNy¯m−2+1RNx¯m,n+1RNy¯m,n]. H_m,1[n,:] = bmatrix 2\\ - 1RN x_m\\ - 1RN y_m\\ - 2+ 1RN x_m,n+ 1RN y_m,n bmatrix. (34) Then, by (16), (19) and (34), the n-th row of the output of the first layer is m,2[n,:]⊤=[1+21RNΔxm,n1RNΔym,n−2N+1], Q_m,2[n,:] = bmatrix1+ 2\\ 1RN x_m,n\\ 1RN y_m,n\\ - 2N+1 bmatrix, (35) where Δxm,n x_m,n =xm,n−x¯m =x_m,n- x_m (36) Δym,n y_m,n =ym,n−y¯m. =y_m,n- y_m. (37) and the O(1/(RN))O(1/(RN)) terms in m,2[n,4]Q_m,2[n,4] is omitted since R is sufficiently large and, indeed, they do not affect the final outcome in our construction. 3.3 Output of the second layer We next construct the second layer. By (35), we have ξ2,2,2 _2,2,2 :=Qm,2[:,2]⊤Qm,2[:,2] :=Q_m,2[:,2] Q_m,2[:,2] =1R2N2NVm+(um−x¯m)2 = 1R^2N^2\NV_m+(u_m- x_m)^2\ ≃1R2NVm 1R^2NV_m (38) ξ2,2,3 _2,2,3 :=Qm,2[:,2]⊤Qm,2[:,3] :=Q_m,2[:,2] Q_m,2[:,3] =1R2N2NCm+(um−x¯m)(0−y¯m) = 1R^2N^2\NC_m+(u_m- x_m)(0- y_m)\ ≃1R2NCm 1R^2NC_m (39) under the assumption and, thus, we have m,2⊤m,2=[−ξ2,2,2ξ2,2,3−]. Q_m,2 Q_m,2= bmatrix-&-&-&-\\ -& _2,2,2& _2,2,3&-\\ -&-&-&-&\\ -&-&-&-&\\ bmatrix. (40) We consider the first head. We set 2,1=2,1,1,2,1,2,2,1,3 _2,1=\W_2,1,1,W_2,1,2,W_2,1,3\, in which 2,1,1 W_2,1,1 =[0100000000000000] = bmatrix0&1&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (41) 2,1,2 W_2,1,2 =[000000000R2N/2−R2N/200000] = bmatrix0&0&0&0\\ 0&0&0&0\\ 0&R^2N/ 2&-R^2N/ 2&0\\ 0&0&0&0\\ bmatrix (42) 2,1,3 W_2,1,3 =[0000100000000000]. = bmatrix0&0&0&0\\ 1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix. (43) Then, by (2.4), (11) and (40), the n-th row of the output of the first head is m,2,1[n,:]⊤=[012RNCmΔxm,n−12RNCmΔxm,n0]. H_m,2,1[n,:] = bmatrix0\\ 1 2RNC_m x_m,n\\ - 1 2RNC_m x_m,n\\ 0\\ bmatrix. (44) On the other hand, we consider the second head. We set 2,2=2,2,1,2,2,2,2,2,3 _2,2=\W_2,2,1,W_2,2,2,W_2,2,3\ in which 2,2,1 W_2,2,1 =[0100000000000000] = bmatrix0&1&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (45) 2,2,2 W_2,2,2 =[0000R2N00−R2N00000000] = bmatrix0&0&0&0\\ R^2N&0&0&-R^2N\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (46) 2,2,3 W_2,2,3 =[1/(1+2)00000000000]. = bmatrix1/(1+ 2)&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix. (47) Then, by (2.4), (11) and (40), the n-th row of the output of the second head is m,2,2[n,:]⊤=[Vm00−Vm]. H_m,2,2[n,:] = bmatrixV_m\\ 0\\ 0\\ -V_m\\ bmatrix. (48) Therefore, by (12), the n-th row of the output of the MHLSA in the second layer is m,2[n,:]⊤=[Vm12RNCmΔxm,n−12RNCmΔxm,n−Vm]. H_m,2[n,:] = bmatrixV_m\\ 1 2RNC_m x_m,n\\ - 1 2RNC_m x_m,n\\ -V_m\\ bmatrix. (49) Now, we consider the layer normalization. By (13), (14) and (49), for m,2[n,:]H_m,2[n,:], we have μm,2,n _m,2,n =0 =0 (50) σm,2,n2 σ^2_m,2,n =14(2Vm2+O(1R2N2))≃12Vm2 = 14 (2V_m^2+O ( 1R^2N^2 ) ) 12V_m^2 (51) since R is sufficiently large. By (15), the n-th row of the output of the layer normalization is ¯m,2[n,:]⊤=[21RNCmVmΔxm,n−1RNCmVmΔxm,n−2]. H_m,2[n,:] = bmatrix 2\\ 1RN C_mV_m x_m,n\\ - 1RN C_mV_m x_m,n\\ - 2\\ bmatrix. (52) The division performed by the layer normalization is actually used for this purpose in [1]. Then, by (16), (35) and (52), the n-th row of the output of the second layer is m,3[n,:]⊤=[1+221RNΔxm,n+1RNCmVmΔxm,n1RNΔym,n−1RNCmVmΔxm,n−2(N+2)/(N+1)]. Q_m,3[n,:] = bmatrix1+2 2\\ 1RN x_m,n+ 1RN C_mV_m x_m,n\\ 1RN y_m,n- 1RN C_mV_m x_m,n\\ - 2(N+2)/(N+1)\\ bmatrix. (53) 3.4 Model output Since xm,N+1=umx_m,N+1=u_m and ym,N+1=0y_m,N+1=0, we have m,3[N+1,3]=−1RN(y¯m+CmVm(um−x¯m)) Q_m,3[N+1,3]=- 1RN ( y_m+ C_mV_m(u_m- x_m) ) (54) by (53). Now, m,3Q_m,3 is flattened into m q_m. Then, we have m[D(N+1)−1]=m,3[N+1,3] q_m[D(N+1)-1]=Q_m,3[N+1,3]. Hence, if we set out[i]=−RNi=D(N+1)−10otherwise W_ out[i]= cases-RN&i=D(N+1)-1\\ 0& otherwise cases (55) then the transformer output for the m-th prompt is, approximately, Qout(m)=y¯m+CmVm(um−x¯m) Q_ out(P_m)= y_m+ C_mV_m(u_m- x_m) (56) by (17) and (54). This is consistent with the least squares estimate at umu_m in (8). 3.5 Discussion In our construction, we need to choose a large value for N that is the number of in-context samples. This is because, for example as in (3.2), we need to eliminate the effect of the prediction point (um,0)(u_m,0) on the calculation of the mean of the in-context input samples. In other words, it is necessary to treat the in-context samples and prediction point separately here. This is not possible since the attention mechanism calculates the inner product of the two sequences, which implies that information on the sequence dependent property disappears; e.g., see (2.4). In large language models, this fact is generally important since we may need masking and/or positional encoding to control the sequence dependent property that cannot be controlled by the attention mechanism. Note that if we appropriately design the prompt (input to the transformer) then we may not need to worry about this approximation problem. For example, in the prompt, the in-context samples and prediction point are embedded separately in the different sequences. The input employed in our setting is just a natural one that is also employed in [1]. In other words, the prompt (input) design is possible to affect the algorithm that is obtained by training. This is also pointed out in [6]. Therefore, a large N assumption may not be essential for performing division using the layer normalization. The construction under the different prompt design is left as a future work. On the other hand, a large value for R is significantly required for performing division by layer normalization. Note that this enforces the absolute values of weights to be either close to zero or extremely large. Our construction may be relatively compact. Unfortunately, the representation of LSA is not unique in the sense that, as shown in A, it produces the same output for different inputs. Therefore, it might be possible to create the same output using a more verbose expression, in which the most weight values are not zeros. 4 Numerical experiment We here show several numerical experiments to demonstrate that our construction in the previous section is actually used in the trained transformer. 4.1 Setting of experiment The transformer defined in this paper is coded by PyTorch, in which we omit the biases in the input and output linear transformation using torch.n.Linear. For the training data mD_m defined by (3), we set σ2=0.04σ^2=0.04 for the noise variance and, at each m, the inputs including prediction points are randomly drawn from a standard normal distribution N(0,1)N(0,1). And, we generate θm,0 _m,0 and θm,1 _m,1 according to N(0,1)N(0,1) independently at each m. The length of the input sequence is N+1=51N+1=51 (the number of in-context samples is N=50N=50). The target output (teacher) for the transformer is the least squares estimate at the prediction point. This is a more direct instruction than the output data at the prediction point. The transformer is TF(4,2,2) TF(4,2,2). For the training, the number of training data is M=5000M=5000, the number of validation data is 10001000, the batch size is 500500, the learning rate is 0.0050.005, the maximum number of training epochs is 2000020000. We apply the layer normalization without an element-wise affine setting in torch.n.LayerNorm. As mentioned in Sec. 3.5, there is a lack of uniqueness in the LSA representation. Therefore, we apply the ℓ1 _1 regularization to the input and output weights, by which non-contributing weights are set to zero and important weights may be highlighted. This may make one expect a relatively compact expression as obtained in our construction. We employ the transformer that minimizes the validation error calculated every 100100 epochs. 4.2 Verifying the use of layer normalization We execute 1010 different training runs under the above setting, in which both of data and initialization differ for each run. We generate 10001000 new data of pairs of the prompt and the least squares estimate, in which the prompts are i.i.d. samples from the same distribution as the prompt in the training data. Then, for each prompt, we obtain the output of the trained transformer and can calculate the squared test error between the output and the least squares estimate at the prediction point. Thus, in each run, we can obtain the average test error for 10001000 new data. Simultaneously, for each m-th prompt in the new data, we can calculate m,2H_m,2 defined in (12). The variance along the embedding used in the layer normalization is the variance of m,2[n,:]H_m,2[n,:] for the n-th sequence, which is σm,2,n2 _m,2,n^2 defined in (14). Also, we can calculate VmV_m defined in (6) as the variance of in-context input samples in the new data. In our construction, the division by VmV_m in (8) is executed using the division by the standard deviation in the layer normalization in (52). To see this, we calculate the correlation coefficient between 1/σm,2,511/ _m,2,51, m=1,…,1000m=1,…,1000 and 1/Vm1/V_m , m=1,…,1000m=1,…,1000 in each run. The results for 1010 runs are summarized in Table 1. We can see that the correlation coefficient is larger than 0.90.9 in the first and ninth runs. This implies that the division by VmV_m is approximately calculated using the layer normalization for these two runs. Since, as seen in our construction, there exist approximation errors in doing this, these correlation coefficients may be satisfactory. Additionally, the trained transformer in the first run shows a better predictive performance in terms of the average test error. Then, we further analyze the trained transformer in the first run below. Note that the trained transformer in the fourth run shows the best predictive performance while the correlation is low. This indicates that an alternative way of the construction exists. We discuss this point in the conclusion. Table 1: Average squared test error and correlation between reciprocals of standard deviation in layer normalization and variance of in-context input samples No. Average test error Correlation 1 0.00072 0.90608 2 0.00139 -0.00510 3 0.00088 0.60064 4 0.00063 0.50725 5 0.00267 0.20170 6 0.00165 0.55747 7 0.00092 0.53876 8 0.00074 0.10411 9 0.00105 0.92632 10 0.00131 0.73416 4.3 Input and output weights In Fig. 2 (a) and (b), for the first run, we show the absolute values of the input and output weights which are denoted by inW_ in and outW_ out in this paper. The weight values are shown on a logarithmic scale. inW_ in is a 3×43× 4 matrix. The row and column are the indices of inputs and embeddings respectively. For the input weights, the indices of inputs (1,2,3)(1,2,3) correspond to the variables (1,x,y)(1,x,y) in the prompt to the transformer; i.e., a constant term, input and output in a simple regression problem. Although outW_ out is a 204204-dimensional vector, we here convert it into a 4×514× 51 matrix whose row and column correspond to the indices of embeddings and sequences respectively. We refer to the entry in row i and column j of these weight matrices as the (i,j)(i,j) weight value. Especially, the output weight is the weight for the output of the cell in the final layer; i.e., (i,j)(i,j) weight is for the cell that is located in the j-th sequence of the i-th embedding. We refer to the cell connected to the (i,j)(i,j) weight as the (i,j)(i,j) cell. (a) Input weights (b) Output weights Figure 2: Absolute weight values In Fig. 2 (a), for the first embedding that corresponds to the first column in the figure, the (3,1)(3,1) value is larger than (1,1)(1,1) and (2,1)(2,1) values, which implies that y is dominant in this embedding. For the second embedding, the (1,2)(1,2) value is relatively larger than the other values, which implies that the constant term is dominant in this embedding. For the third embedding, the (2,3)(2,3) and (3,3)(3,3) values are larger than the (1,3)(1,3) value, which implies that (x,y)(x,y) are dominant in this embedding. Lastly, for the fourth embedding, the (2,4)(2,4) value is larger than the other values, which implies that x is dominant in this embedding. Although the non-dominant weights are not exactly zero, this result may support that the inputs are extracted in the embedding separately for the further calculation. Note that this representation may vary depending on the run since LSA is possible to give the same output for different inputs; e.g., see A. In Fig. 2 (b), we can see that the weight values of the third embedding cells are relatively large and, especially, the (3,51)(3,51) value is extremely large compared to the other values. Therefore, the (3,51)(3,51) cell is the dominant factor for the transformer output. Note that the contribution of the third embedding cells are relatively large compared to the other embeddings. The effect of those cells will be mentioned in the conclusion. 4.4 Regression analysis We here check that, in the first run, the output of the (3,51)(3,51) cell in the second layer has the form of (8) by using a regression analysis. We generate 10001000 new data for the prompt to the trained transformer, which are i.i.d. samples from the same distribution as the training data. Then, for each m-th new prompt, we can obtain m,3Q_m,3 in (53) and focus on zm:=m,3[51,3]z_m:=Q_m,3[51,3], which is the output of the (3,51)(3,51) cell in the second layer. We then fit zmz_m by fm:=α0+α1(um−x¯m)+α2y¯m+α3CmVm(um−x¯m), f_m:= _0+ _1(u_m- x_m)+ _2 y_m+ _3 C_mV_m(u_m- x_m), (57) where (α0,α1,α2,α3)( _0, _1, _2, _3) are regression coefficients. Table 2: Results of regression analysis α0 _0 α1 _1 α2 _2 α3 _3 0.00250 -0.00159 -0.2436 -0.22498 (-0.00453) (-0.72452) (-0.69403) The coefficient estimates are summarized in Table 2, in which we show the standardized regression coefficients in the brackets. In Table 2, the estimates of α0 _0 and α1 _1 are small and those of α2 _2 and α3 _3 are large. Here, the correlation between zmz_m and fmf_m was 0.994790.99479. Thus, zmz_m is found to be the form of the last two terms in fmf_m. It is the same form of the least squares solution in (8). And, the correlation between zmz_m and the least squares solution was −0.99396-0.99396, which implies that the output of the (3,51)(3,51) cell is almost consistent with the least quares solution except the sign. The sign can be changed by adjusting the corresponding output weight as in our construction. 4.5 Prediction We finally show the fitting curve of the trained transformer in the first run to visually confirm that our conclusion is valid. We generate 100100 new data for the prompts to the trained transformer, which are (x1,y1),…,(xn,y50),(uk,0)\(x_1,y_1),…,(x_n,y_50),(u_k,0)\ for k=1,…,100k=1,…,100, where (xi,yi)(x_i,y_i) is generated from the same distribution as the training data and uku_k is the equidistant points in [−3,3][-3,3]. In other words, we obtain the transformer output at 100100 different prediction points under the same in-context samples. Therefore, we can plot a fitting curve for the in-context samples. In Fig. 3 (a), we show the in-context samples (open circle) and the transformer output (black solid line). In this figure, we also show the least squares estimate of a simple regression model (gray solid line), which is obtained for the in-context samples, (x1,y1),…,(xn,y50)(x_1,y_1),…,(x_n,y_50). We can see that the transformer output is well consistent with the least squares estimate of a simple regression model. In Fig. 3 (b), we show the transformer output when the output weight values are set to zero except the (3,51)(3,51) weight value. In other words, we check the output only through the (3,51)(3,51) cell. We can see that the output through this cell is almost consistent with the least squares estimate, which implies that the (3,51)(3,51) cell significantly contributes the entire output and, as in our construction, the output of the (3,51)(3,51) cell approximates (54) well. Note that, however, the output through the (3,51)(3,51) cell is not entirely consistent with the transformer output. This point will be mentioned in the conclusion. (a) Transformer output (b) Output through (3,51)(3,51) cell Figure 3: Fitting curve 5 Conclusion and future work In this paper, we constructed the transformer that outputs the least squares estimate of a simple regression problem. To achieve this, we approximately represent the closed form solution by using the layer normalization, which is an insight of [1]. For in-context learning of regression tasks, most studies have considered implementations of the gradient descent algorithms. This is because multiplication, which is a basic operation of the attention mechanism, is suitable for implementing the gradient descent algorithm. However, the experimental evidence that the gradient descent algorithm is implemented through training is still unclear. In contrast, in this paper, we experimentally showed that there exist the examples, in which our implementation based on the layer normalization is mainly used in the transformer trained with the ℓ1 _1 regularization when the target output is the least sqaures estimate. Althogh this paper considered a very simple case, it may be a first attempt to step into the training of transformer that implements in-context learning. Therefore, we further have several things to address. • Unfortunately, as see in Fig. 3, the output constructed based on the layer normalization is not entirely consistent with the transformer output. Thus, the trained transformer has the other mechanism to further fit the least squared estimate well. In Fig. 2 (b), we can see that the output weight values that are assigned to the third embedding cells are relatively larger than the weight values in the other embedding cells. Therefore, the outputs of the third embedding cells may serve to complement the approximation error in our construction. Also, as in Table 1, the trained transformer in the fourth run may not implement the division by VmV_m using the layer normalization in the second layer. Nevertheless, the test error is low. Furthermore, in training without the ℓ1 _1 regularization, similar results are obtained. Thus, there may be different types of solver for obtaining the least squares estimate for TF(4,2,2) TF(4,2,2). Hence, we need the investigation to reveal the other mechanism as part of our future work. • Note that the gradient descent based implementation may not be comparable in TF(4,2,2) TF(4,2,2) since the number of layers of this transformer may not be enough to implement the sufficient number of gradient descent steps. For example, the two gradient descent steps can be implemented by a 33-layer transformer with linear attention; e.g., see B. The detailed comparison to the gradient descent algorithm is left as a future work. • In our experiment, the prediction point is generated from the same distribution of the training inputs. Therefore, the approximation under a large N assumption is almost accurate; i.e., for example, in (3.2), the mean of N samples is almost consistent with that of N+1N+1 samples. Therefore, we need to investigate the case where the prediction point is out-of-distribution. Note that this relates to the generalization capability of transformers. On the other hand, as noted in Sec. 3.5, a large N approximation may be relaxed by the choice of the prompt format (design). The construction under the other prompt format and the numerical investigation are parts of our future work. • Although, in our experiment, we employed the least sqaures estimate as the target output in training, the output sample at the prediction point may be natural as the target output and it may be easily collectible. The experimental analysis of this case is also left as a future work. References [1] E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, D. Zhou, “What learning algorithm is in-context learning? investigations with linear models”, arXiv:2211.15661, 2022. [2] Y. Bai , F. Chen, H. Wang, C. Xiong, S. Mei, “Transformers as Statisticians: Provable in-context learning with in-context algorithm selection”, arXiv:2306.04637, 2023. [3] T. Brown, et al., “Language models are few-shot learners”. arXiv:2005.14165, 2020. [4] D. Dai, Y. Sun, L. Dong, Y. Hao, Z. Sui, F. Wei, “Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers”, arXiv:2212.10559, 2023. [5] S. Garg, D. Tsipras, P. Liang, G. Valiant, “What can transformers learn in-context? a case study of simple function classes”, arXiv:2208.01066, 2022. [6] K. Hagiwara, ‘Effect of bias matrix in linear self-attention for in-context learning”, IEICE Trans. D, vol.E109-D, No.8, 2026. [7] K. He, X. Zhang, S. Ren, J. Sun, “Deep residual learning for image recognition”, arXiv:1512.03385, 2015. [8] K. Hornik, M. Stinchcombe, H. White, “Multilayer feedforward networks are universal approximators”, Neural Networks, vol.2, no.5, p.359–366, 1989. [9] A. Krizhevsky, I. Sutskever, G. E. Hinton, “ImageNet classification with deep convolutional neural networks”, Communications of the ACM, vol.80, no.6, p.86-90, 2012. [10] A. Mahankali, T. B. Hashimoto, T. Ma, “One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention”, arXiv:2307.03576, 2023. [11] J. von Oswald, et al., “Transformers learn in-context by gradient descent”, arXiv:2212.07677, 2022. [12] A. Vaswani, et al., “Attention is All you Need”, arXiv:1706.03762v7, 2017. [13] R. Zhang, S. Frei, P. L. Bartlett, “Trained transformers learn linear models in-context”, Journal of Machine Learning Research, vol.25, no.49, p.1-55, 2023. [14] N. Ravishanker, D. K. Dey, “A first course in linear model theory”, Chapman & Hall/CRC, 2002. Appendix A An example of a lack of uniqueness We show here a simple example, in which LSA produces the same output for the different inputs. Let Q be an N×3N× 3 matrix whose n-th row is [1,an,bn][1,a_n,b_n] and define a:=∑n=1Nana:= _n=1^Na_n and b:=∑n=1Nbnb:= _n=1^Nb_n. In (2.4), we set this Q as the input matrix and we set =1,2,3 =\W_1,W_2,W_3\, where 1=3=[100000000],2=[000010001]. W_1=W_3= bmatrix1&0&0\\ 0&0&0\\ 0&0&0\\ bmatrix,~W_2= bmatrix0&0&0\\ 0&1&0\\ 0&0&1\\ bmatrix. (58) Then, the n-th row of LSA() LSA_ (Q) is [0,a,b][0,a,b] for any n. On the other hand, let Q be an N×3N× 3 matrix whose n-th row is [1,an,an+bn][1,a_n,a_n+b_n]. In (2.4), by setting 1=3=[100000000],2=[00001−1001], W_1=W_3= bmatrix1&0&0\\ 0&0&0\\ 0&0&0\\ bmatrix,~W_2= bmatrix0&0&0\\ 0&1&-1\\ 0&0&1\\ bmatrix, (59) the n-th row of LSA() LSA_ (Q) is [0,a,b][0,a,b] for any n. This is the same as in the first case. Appendix B Implementation of gradient descent method B.1 Gradient descent method for a simple regression The least squares estimates at umu_m is given by (8). Here, we consider to obtain this solution by using a gradient descent method. We define fθm(um):=y¯m+θm(u−x¯m), f_ _m(u_m):= y_m+ _m(u- x_m), (60) where θm _m is a parameter to be adjusted. Thus, θ^m=Cm/Vm θ_m=C_m/V_m is a closed form solution to θm _m. We also define Sm(θ):=12∑n=1N(ym,n−fθ(xm,n))2. S_m(θ):= 12 _n=1^N(y_m,n-f_θ(x_m,n))^2. (61) We have ∂Sm(θ)∂θ ∂ S_m(θ)∂θ =−∑n=1N(ym,n−fθ(xm,n))(xm,n−x¯m) =- _n=1^N(y_m,n-f_θ(x_m,n))(x_m,n- x_m) =−∑n=1N((ym,n−y¯m)−θ(xm,n−x¯m))(xm,n−x¯m) =- _n=1^N((y_m,n- y_m)-θ(x_m,n- x_m))(x_m,n- x_m) =−N(Cm−θVm). =-N(C_m-θ V_m). (62) By setting θm(0) _m(0) as an initial value, at the t-th step, the update equation of the gradient descent for θm(t) _m(t) is given by θm(t) _m(t) =θm(t−1)−η∂Sm(θ)∂θ|θ=θm(t−1) = _m(t-1)-η . ∂ S_m(θ)∂θ |_θ= _m(t-1) =θm(t−1)+Nη(Cm−θm(t−1)Vm), = _m(t-1)+Nη(C_m- _m(t-1)V_m), (63) where η>0η>0 is a learning rate. We here give a solution to the second step. We set θm(0)=am _m(0)=a_m (64) and have θm(1) _m(1) =θm(0)+Nη(Cm−θm(0)Vm) = _m(0)+Nη(C_m- _m(0)V_m) =am+Nη(Cm−amVm) =a_m+Nη(C_m-a_mV_m) (65) By a simple calculation, we then have θm(2) _m(2) =θm(1)+Nη(Cm−θm(1)Vm) = _m(1)+Nη(C_m- _m(1)V_m) =am+2NηCm−2NηamVm−N2η2CmVm+N2η2amVm2. =a_m+2Nη C_m-2Nη a_mV_m-N^2η^2C_mV_m+N^2η^2a_mV_m^2. (66) and we obtain the estimate at umu_m by fθm(2)(um)=y¯m+θm(2)(um−x¯m) f_ _m(2)(u_m)= y_m+ _m(2)(u_m- x_m) (67) with this θm(2) _m(2). In this section, we show an implementation of this 22-step gradient descent by TF(4,3,2) TF(4,3,2) under the assumption that N is large or the O(1/N)O(1/N) term is negligible. Here, we omit the layer normalization to simplify the construction and focus on the role of multiplication. We here drop the subscript m that denotes an index of training data for the transformer. B.2 Input embedding We define in=[1000α100α2β20α3β3] W_ in= bmatrix1&0&0\\ 0& _1&0\\ 0& _2& _2\\ 0& _3& _3\\ bmatrix (68) and in=D,1 b_ in=O_D,1. By (10), we then have the embedding whose i-th row is 1[n,:]⊤=[1α1xnα2xn+β2ynα3xn+β3yn]. Q_1[n,:] = bmatrix1\\ _1x_n\\ _2x_n+ _2y_n\\ _3x_n+ _3y_n\\ bmatrix. (69) 1Q_1 is the (N+1)×D(N+1)× D matrix and is the input to TF(4,3,2) TF(4,3,2). B.2.1 The first layer We show a construction of the first head. Since xN+1=ux_N+1=u and yN+1=0y_N+1=0, we have 1[:,1]⊤1[:,3] Q_1[:,1] Q_1[:,3] =∑n=1N+11⋅(α2xn+β2yn) = _n=1^N+11· ( _2x_n+ _2y_n ) =α2Nx¯+β2Ny¯+α2u = _2N x+ _2N y+ _2u ≃α2Nx¯+β2Ny¯ _2N x+ _2N y (70) 1[:,1]⊤1[:,4] Q_1[:,1] Q_1[:,4] =∑n=1N+11⋅(α3xn+β3yn) = _n=1^N+11· ( _3x_n+ _3y_n ) =α3Nx¯+β3Ny¯+α3u = _3N x+ _3N y+ _3u ≃α3Nx¯+β3Ny¯ _3N x+ _3N y (71) by omitting O(1)O(1) terms under the assumption that N is large. If we set 1,1,1 W_1,1,1 =[−1N000000000000000] = bmatrix- 1N&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (72) 1,1,2 W_1,1,2 =[0000000000100001] = bmatrix0&0&0&0\\ 0&0&0&0\\ 0&0&1&0\\ 0&0&0&1\\ bmatrix (73) 1,1,3 W_1,1,3 =[1000000000000000] = bmatrix1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (74) then, by (2.4), we have 1,1[n,:]⊤=[00−α2x¯−β2y¯−α3x¯−β3y¯]. H_1,1[n,:] = bmatrix0\\ 0\\ - _2 x- _2 y\\ - _3 x- _3 y\\ bmatrix. (75) We next show a construction of the second head. We have 1[:,1]⊤1[:,2]≃α1Nx¯ Q_1[:,1] Q_1[:,2] _1N x (76) when N is large. Thus, if we set 1,2,1 W_1,2,1 =[−1N000000000000000] = bmatrix- 1N&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (77) 1,2,2 W_1,2,2 =[0000010000000000] = bmatrix0&0&0&0\\ 0&1&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (78) 1,2,3 W_1,2,3 =[1000000000000000] = bmatrix1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (79) then, by (2.4), we have 1,2[n,:]=[0−α1x¯00]. H_1,2[n,:]= bmatrix0&- _1 x&0&0\\ bmatrix. (80) As a result, the output of the first layer is 2=(1,1+1,2)+1, Q_2=(H_1,1+H_1,2)+Q_1, (81) whose n-th row is given by 2[n,:]⊤=[1α1(xn−x¯)α2(xn−x¯)+β2(yn−y¯)α3(xn−x¯)+β3(yn−y¯)] Q_2[n,:] = bmatrix1\\ _1(x_n- x)\\ _2(x_n- x)+ _2(y_n- y)\\ _3(x_n- x)+ _3(y_n- y)\\ bmatrix (82) since we omit the layer normalization. B.3 The second layer We show a construction of the firs head. We have 2[:,1]⊤2[:,2] Q_2[:,1] Q_2[:,2] =∑n=1N+1α1(xn−x¯)α2(xn−x¯)+β2(yn−y¯) = _n=1^N+1 _1(x_n- x) \ _2(x_n- x)+ _2(y_n- y) \ ≃Nξ2,1 N _2,1 (83) for a large N, where ξ2,1:=α1α2V+α1β2C. _2,1:= _1 _2V+ _1 _2C. (84) If we set 2,1,1 W_2,1,1 =[01N00000000000000] = bmatrix0& 1N&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (85) 2,1,2 W_2,1,2 =[0000000000010000] = bmatrix0&0&0&0\\ 0&0&0&0\\ 0&0&0&1\\ 0&0&0&0\\ bmatrix (86) 2,1,3 W_2,1,3 =[0000100000000000] = bmatrix0&0&0&0\\ 1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (87) then, by (2.4), we have 2,1[n,:]=[000ξ2,1α1(xn−x¯)]. H_2,1[n,:]= bmatrix0&0&0& _2,1 _1(x_n- x)\\ bmatrix. (88) We next show a construction of the second head. If we set 2,2,1=2,2,2=2,2,3=4,4, W_2,2,1=W_2,2,2=W_2,2,3=O_4,4, (89) then we formally have 2,2=N+1,4. H_2,2=O_N+1,4. (90) Therefore, the the output of the second layer is 3=(2,1+2,2)+2, Q_3=(H_2,1+H_2,2)+Q_2, (91) whose n-th row is given by 3[n,:]⊤=[1α1(xn−x¯)α2(xn−x¯)+β2(yn−y¯)(ξ2,1α1+α3)(xn−x¯)+β3(yn−y¯)]. Q_3[n,:] = bmatrix1\\ _1(x_n- x)\\ _2(x_n- x)+ _2(y_n- y)\\ ( _2,1 _1+ _3)(x_n- x)+ _3(y_n- y)\\ bmatrix. (92) B.4 The third layer We show a construction of the first head. We have 3[n,2]⊤3[n,4] Q_3[n,2] Q_3[n,4] =∑n=1N+1(ξ2,1α1+α3)(xn−x¯)+β3(yn−y¯)α1(xn−x¯) = _n=1^N+1 \( _2,1 _1+ _3)(x_n- x)+ _3(y_n- y) \ _1(x_n- x) ≃Nξ3,1 N _3,1 (93) for a large N, where ξ3,1 _3,1 =α1(ξ2,1α1+α3)V+α1β3C. = _1( _2,1 _1+ _3)V+ _1 _3C. (94) If we set 3,1,1 W_3,1,1 =[01/N00000000000000] = bmatrix0&1/N&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (95) 3,1,2 W_3,1,2 =[0000000000000010] = bmatrix0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&1&0\\ bmatrix (96) 3,1,3 W_3,1,3 =[0000100000000000] = bmatrix0&0&0&0\\ 1&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ bmatrix (97) then we have 3,1[n,:]=[00ξ3,1α1(xn−x¯)0]. H_3,1[n,:]= bmatrix0&0& _3,1 _1(x_n- x)&0\\ bmatrix. (98) We next show a construction of the second head. If we set 3,2,1=3,2,2=3,2,3=4,4, W_3,2,1=W_3,2,2=W_3,2,3=O_4,4, (99) then we formally have 3,2=N+1,4. H_3,2=O_N+1,4. (100) Thus, the the output of the third layer is 4=(3,1+3,2)+3, Q_4=(H_3,1+H_3,2)+Q_3, (101) whose n-th row is given by 4[n,:]⊤=[1α1(xn−x¯)ξ3,1α1(xn−x¯)+α2(xn−x¯)+β2(yn−y¯)ξ2,1α1(xn−x¯)+α3(xn−x¯)+β3(yn−y¯)], Q_4[n,:] = bmatrix1\\ _1(x_n- x)\\ _3,1 _1(x_n- x)+ _2(x_n- x)+ _2(y_n- y)\\ _2,1 _1(x_n- x)+ _3(x_n- x)+ _3(y_n- y)\\ bmatrix, (102) where we can write ξ3,1 _3,1 =α1(ξ2,1α1+α3)V+α1β3C = _1( _2,1 _1+ _3)V+ _1 _3C =(α12ξ2,1+α1α3)V+α1β3C =( _1^2 _2,1+ _1 _3)V+ _1 _3C =(α13α2V+α13β2C+α1α3)V+α1β3C =( _1^3 _2V+ _1^3 _2C+ _1 _3)V+ _1 _3C =α1α3V+α1β3C+α13α2V2+α13β2CV = _1 _3V+ _1 _3C+ _1^3 _2V^2+ _1^3 _2CV (103) by (84) and (94). 4Q_4 is flattended into q which is a 4(N+1)×14(N+1)× 1 vector. Then, by setting out[n]=1n=4N+30otherwise, W_ out[n]= cases1&n=4N+3\\ 0& otherwise, cases (104) the prediction of the transformer is obtained by y^(u) y(u) =4[N+1,3] =Q_4[N+1,3] =−β2y¯+ξ4,1(u−x¯), =- _2 y+ _4,1(u- x), (105) where ξ4,1 _4,1 =α2+α12α3V+α12β3C+α14α2V2+α14β2CV. = _2+ _1^2 _3V+ _1^2 _3C+ _1^4 _2V^2+ _1^4 _2CV. (106) B.5 Correspondence to gradient descent By the correspondence between (B.4) and (67), we have −β2=1−α2=aα12α3=−2nηaα12β3=2nηα14α2=n2η2aα14β2=−n2η2 cases- _2&=1\\ - _2&=a\\ _1^2 _3&=-2nη a\\ _1^2 _3&=2nη\\ _1^4 _2&=n^2η^2a\\ _1^4 _2&=-n^2η^2\\ cases (107) where a is an initial value of the 22-step gradient decent in (64). By solving these equations, we have (α1,α2,β2,α3,β3)=(nη,a,−1,−2a,2). ( _1, _2, _2, _3, _3)=( nη,a,-1,-2a,2). (108) Therefore, the output of the transformer is approximately consistent with the output of the linear function after the 22-step gradient descent.