In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention

# In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention
Source: [https://arxiv.org/html/2607.15819](https://arxiv.org/html/2607.15819)
\(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\\ell\_\{1\}regularization when the target output is the least squares estimate\.

Keywordin\-context learning, linear self\-attention, a simple linear regression, layer normalization

## 1Introduction

In\-context learning is a remarkable property of transformers, which form the basis of large language models such as GPT\-3\[[3](https://arxiv.org/html/2607.15819#bib.bib3)\], 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](https://arxiv.org/html/2607.15819#bib.bib5),[11](https://arxiv.org/html/2607.15819#bib.bib11),[1](https://arxiv.org/html/2607.15819#bib.bib1),[4](https://arxiv.org/html/2607.15819#bib.bib4),[2](https://arxiv.org/html/2607.15819#bib.bib2)\]\.

\[[5](https://arxiv.org/html/2607.15819#bib.bib5)\]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](https://arxiv.org/html/2607.15819#bib.bib11)\]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](https://arxiv.org/html/2607.15819#bib.bib1)\]proved that a transformer can implement a gradient descent algorithm and a closed form solution for ridge regression\.\[[4](https://arxiv.org/html/2607.15819#bib.bib4)\]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](https://arxiv.org/html/2607.15819#bib.bib5),[11](https://arxiv.org/html/2607.15819#bib.bib11),[1](https://arxiv.org/html/2607.15819#bib.bib1),[10](https://arxiv.org/html/2607.15819#bib.bib10),[4](https://arxiv.org/html/2607.15819#bib.bib4)\]do not consider the training phase,\[[13](https://arxiv.org/html/2607.15819#bib.bib13)\]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](https://arxiv.org/html/2607.15819#bib.bib2)\]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](https://arxiv.org/html/2607.15819#bib.bib1)\],\[[2](https://arxiv.org/html/2607.15819#bib.bib2)\]precisely evaluated the prediction performance in terms of network size, and shown a near\-optimal predictive power\.\[[2](https://arxiv.org/html/2607.15819#bib.bib2)\]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](https://arxiv.org/html/2607.15819#bib.bib1)\]\. We then present a numerical example, in which this implementation is mainly used in the transformer trained with theℓ1\\ell\_\{1\}regularization when the target is the least squares estimate\.

Previous works\[[5](https://arxiv.org/html/2607.15819#bib.bib5),[11](https://arxiv.org/html/2607.15819#bib.bib11),[1](https://arxiv.org/html/2607.15819#bib.bib1),[4](https://arxiv.org/html/2607.15819#bib.bib4),[2](https://arxiv.org/html/2607.15819#bib.bib2)\]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](https://arxiv.org/html/2607.15819#bib.bib1)\]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](https://arxiv.org/html/2607.15819#bib.bib1)\]\. After showing this construction, we present numerical experiments to show that our implementation is actually used in the transformer trained withℓ1\\ell\_\{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\.

## 2Problem setting

### 2\.1Notations

In this paper,𝐎I,J\{\\mathbf\{O\}\}\_\{I,J\}is theI×JI\\times Jzero matrix\. Note that we use this notation also for vectors, where eitherIIorJJequals11\. For anI×JI\\times Jmatrix𝐀\{\\mathbf\{A\}\},𝐀​\[i,j\]\{\\mathbf\{A\}\}\[i,j\]is the\(i,j\)\(i,j\)\-entry of𝐀\{\\mathbf\{A\}\},𝐀​\[i,:\]\{\\mathbf\{A\}\}\[i,:\]is theii\-th row vector of𝐀\{\\mathbf\{A\}\}and𝐀​\[:,j\]\{\\mathbf\{A\}\}\[:,j\]is thejj\-th column vector of𝐀\{\\mathbf\{A\}\}fori=1,…,Ii=1,\\ldots,Iandj=1,…,Jj=1,\\ldots,J\. When𝐀\{\\mathbf\{A\}\}is anI×1I\\times 1vector,𝐀​\[i\]\{\\mathbf\{A\}\}\[i\]denotes itsii\-th entry\.

### 2\.2Setting of in\-context learning of transformer

We next explain the transformer\-based in\-context learning of a simple linear regression task\. LetMMbe 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 eachmmin\{1,…,M\}\\\{1,\\ldots,M\\\}, a set ofNNrandom 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,\\ldots,N\\\}, are generated\. We denote themm\-th prompt \(input to the transformer\) by𝐏m\{\\mathbf\{P\}\}\_\{m\}, which is an\(N\+1\)×3\(N\+1\)\\times 3matrix and whosenn\-th row is

𝐏m​\[n,:\]=\[1xm,nym,n\],\\displaystyle\{\\mathbf\{P\}\}\_\{m\}\[n,:\]=\\begin\{bmatrix\}1&x\_\{m,n\}&y\_\{m,n\}\\end\{bmatrix\},\(1\)where we defineym,N\+1:=0y\_\{m,N\+1\}:=0andxm,N\+1:=umx\_\{m,N\+1\}:=u\_\{m\}that is a prediction point; e\.g\., see\[[1](https://arxiv.org/html/2607.15819#bib.bib1)\]\. Thus, the input sequence length isN\+1N\+1and the number of in\-context samples isNN\. We assume that, for eachxm,nx\_\{m,n\},ym,ny\_\{m,n\}is generated by

ym,n=θm,0\+θm,1​xm,n\+εm,n,\\displaystyle y\_\{m,n\}=\\theta\_\{m,0\}\+\\theta\_\{m,1\}x\_\{m,n\}\+\\varepsilon\_\{m,n\},\(2\)whereεm,1,…,εm,N\\varepsilon\_\{m,1\},\\ldots,\\varepsilon\_\{m,N\},m=1,…,Mm=1,\\ldots,Mare i\.i\.d\. additive noises from a probability distribution with mean0and varianceσ2<∞\\sigma^\{2\}<\\infty\. Our goal is to obtain the least squares prediction atx=umx=u\_\{m\}for the prompt𝐏m\{\\mathbf\{P\}\}\_\{m\}\. Therefore, for eachmm, we need to calculate the least squares solution using

𝐃m=\{\(xm,n,ym,n\):n=1,…,N\},\\displaystyle\{\\mathbf\{D\}\}\_\{m\}=\\\{\(x\_\{m,n\},y\_\{m,n\}\):n=1,\\ldots,N\\\},\(3\)which is the set of in\-context samples\. The important point is that the regression lines can be different for eachmm\. Thus, we may assume that\(θm,0,θm,1\)\(\\theta\_\{m,0\},\\theta\_\{m,1\}\)are sampled from a probability distribution for eachmm\. However, in the construction of the transformer, we do not make any specific assumptions on the underlying probability distribution of\(θm,0,θm,1\)\(\\theta\_\{m,0\},\\theta\_\{m,1\}\)\. Although we also do not make any specific assumptions on the underlying probability distribution ofxm,nx\_\{m,n\}andεm,n\\varepsilon\_\{m,n\}, we will make an assumption on the in\-context samples later\.

### 2\.3Training data for the transformer

For themm\-th training data𝐃m\{\\mathbf\{D\}\}\_\{m\}defined in \([3](https://arxiv.org/html/2607.15819#S2.E3)\), we define

x¯m\\displaystyle\{\\overline\{x\}\}\_\{m\}:=1N​∑n=1Nxm,n\\displaystyle:=\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}x\_\{m,n\}\(4\)y¯m\\displaystyle\{\\overline\{y\}\}\_\{m\}:=1N​∑n=1Nym,n\\displaystyle:=\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}y\_\{m,n\}\(5\)Vm\\displaystyle V\_\{m\}:=1N​∑n=1N\(xm,n−x¯m\)2\\displaystyle:=\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}\(x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\)^\{2\}\(6\)Cm\\displaystyle C\_\{m\}:=1N​∑n=1N\(xm,n−x¯m\)​\(ym,n−y¯m\)\.\\displaystyle:=\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}\(x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\)\(y\_\{m,n\}\-\{\\overline\{y\}\}\_\{m\}\)\.\(7\)
For a simple regression problem, it is easy to see that the prediction atx=umx=u\_\{m\}using the least squares solution based on𝐃m\{\\mathbf\{D\}\}\_\{m\}is given by

y^m​\(um\):=y¯m\+CmVm​\(um−x¯m\)\\displaystyle\{\\widehat\{y\}\}\_\{m\}\(u\_\{m\}\):=\{\\overline\{y\}\}\_\{m\}\+\\frac\{C\_\{m\}\}\{V\_\{m\}\}\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\(8\)\(see e\.g\.,\[[14](https://arxiv.org/html/2607.15819#bib.bib14)\]\)\.

Below, we construct the transformer that receives𝐏m\{\\mathbf\{P\}\}\_\{m\}as an input and outputsy^m​\(um\)\{\\widehat\{y\}\}\_\{m\}\(u\_\{m\}\)form=1,…,Mm=1,\\ldots,M\. Thus, the transformer calculates the least squares estimate at the prediction pointumu\_\{m\}using in\-context samples𝐃m\{\\mathbf\{D\}\}\_\{m\}at eachmm\.

### 2\.4Linear self\-attention

We define a linear self\-attention \(LSA\), which receives an\(N\+1\)×D\(N\+1\)\\times Dmatrix𝐐\{\\mathbf\{Q\}\}as an input and outputs\(N\+1\)×D\(N\+1\)\\times Dmatrix defined by

LSA𝚯​\(𝐐\)\\displaystyle\{\\rm LSA\}\_\{\{\\boldsymbol\{\\Theta\}\}\}\(\{\\mathbf\{Q\}\}\):=\(𝐐𝐖3\)​\(𝐐𝐖1⊤\)⊤​\(𝐐𝐖2\)\\displaystyle:=\(\{\\mathbf\{Q\}\}\{\\mathbf\{W\}\}\_\{3\}\)\(\{\\mathbf\{Q\}\}\{\\mathbf\{W\}\}\_\{1\}^\{\\top\}\)^\{\\top\}\(\{\\mathbf\{Q\}\}\{\\mathbf\{W\}\}\_\{2\}\)=\(𝐐𝐖3\)​\(𝐖1​𝐐⊤​𝐐𝐖2\),\\displaystyle=\(\{\\mathbf\{Q\}\}\{\\mathbf\{W\}\}\_\{3\}\)\(\{\\mathbf\{W\}\}\_\{1\}\{\\mathbf\{Q\}\}^\{\\top\}\{\\mathbf\{Q\}\}\{\\mathbf\{W\}\}\_\{2\}\),\(9\)where𝚯=\{𝐖1,𝐖2,𝐖3\}\{\\boldsymbol\{\\Theta\}\}=\\\{\{\\mathbf\{W\}\}\_\{1\},\{\\mathbf\{W\}\}\_\{2\},\{\\mathbf\{W\}\}\_\{3\}\\\}is an ordered set of parameters and𝐖1\{\\mathbf\{W\}\}\_\{1\},𝐖2\{\\mathbf\{W\}\}\_\{2\}and𝐖3\{\\mathbf\{W\}\}\_\{3\}areD×DD\\times Dmatrices\. The operation using this LSA is precisely discussed in\[[6](https://arxiv.org/html/2607.15819#bib.bib6)\]\.

### 2\.5Transformer structure

The transformer structure considered in this paper is illustrated in Fig\.[1](https://arxiv.org/html/2607.15819#S2.F1), 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](https://arxiv.org/html/2607.15819#S2.F1)\(a\)\)\. The transformer block consists of multi\-head LSA \(MHLSA\) block with layer normalization followed by a skip connection \(Fig\.[1](https://arxiv.org/html/2607.15819#S2.F1)\(b\)\)\. It receives themm\-th prompt𝐏m\{\\mathbf\{P\}\}\_\{m\}defined in \([1](https://arxiv.org/html/2607.15819#S2.E1)\) and, for simplicity, outputs a scalar value whose target is the least squares estimate at a prediction pointumu\_\{m\}, which isy^m​\(um\)\{\\widehat\{y\}\}\_\{m\}\(u\_\{m\}\)given by \([8](https://arxiv.org/html/2607.15819#S2.E8)\)\.

![Refer to caption](https://arxiv.org/html/2607.15819v1/x1.png)

\(a\) Main stream

![Refer to caption](https://arxiv.org/html/2607.15819v1/x2.png)

\(b\) Transformer block

Figure 1:Structure of transformerLet𝐖in\{\\mathbf\{W\}\}\_\{\\rm in\}and𝒃in\{\\boldsymbol\{b\}\_\{\\rm in\}\}be aD×3D\\times 3input weight matrix andD×1D\\times 1input bias vector respectively\. The embedding of the prompt𝐏m\{\\mathbf\{P\}\}\_\{m\}by an affine transformation with\(𝐖in,𝒃in\)\(\{\\mathbf\{W\}\}\_\{\\rm in\},\{\\boldsymbol\{b\}\_\{\\rm in\}\}\)is denoted by𝐐m,1\{\\mathbf\{Q\}\}\_\{m,1\}whose size is\(N\+1\)×D\(N\+1\)\\times D\. More precisely, we define

𝐐m,1​\[n,:\]⊤:=𝐖in​𝐏m​\[n,:\]⊤\+𝒃in\\displaystyle\{\\mathbf\{Q\}\}\_\{m,1\}\[n,:\]^\{\\top\}:=\{\\mathbf\{W\}\}\_\{\\rm in\}\{\\mathbf\{P\}\}\_\{m\}\[n,:\]^\{\\top\}\+\{\\boldsymbol\{b\}\_\{\\rm in\}\}\(10\)forn=1,…,N\+1n=1,\\ldots,N\+1\. The number of layers is denoted byLL\. For themm\-th training data, the input to thell\-th layer is denoted by𝐐m,l\{\\mathbf\{Q\}\}\_\{m,l\},l=1,…,Ll=1,\\ldots,L\. Thus, the above𝐐m,1\{\\mathbf\{Q\}\}\_\{m,1\}is the input to the first layer\.

The parameter of the LSA block of thekk\-th head in thell\-th layer is denoted by𝚯l,k=\{𝐖l,k,1,𝐖l,k,2,𝐖l,k,3\}\{\\boldsymbol\{\\Theta\}\}\_\{l,k\}=\\\{\{\\mathbf\{W\}\}\_\{l,k,1\},\{\\mathbf\{W\}\}\_\{l,k,2\},\{\\mathbf\{W\}\}\_\{l,k,3\}\\\}\. The output of the MHLSA in thell\-th layer formm\-th training data is

𝐇m,l,k:=LSA𝚯l,k​\(𝐐m,l\)\\displaystyle\{\\mathbf\{H\}\}\_\{m,l,k\}:=\{\\rm LSA\}\_\{\{\\boldsymbol\{\\Theta\}\}\_\{l,k\}\}\(\{\\mathbf\{Q\}\}\_\{m,l\}\)\(11\)𝐇m,l:=∑k=1K𝐇m,l,k,\\displaystyle\{\\mathbf\{H\}\}\_\{m,l\}:=\\sum\_\{k=1\}^\{K\}\{\\mathbf\{H\}\}\_\{m,l,k\},\(12\)whereKKis the number of heads\. The size of𝐇m,l\{\\mathbf\{H\}\}\_\{m,l\}is\(N\+1\)×D\(N\+1\)\\times 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\\displaystyle\\mu\_\{m,l,n\}:=1D​∑d=1D𝐇m,l​\[n,d\]\\displaystyle:=\\frac\{1\}\{D\}\\sum\_\{d=1\}^\{D\}\{\\mathbf\{H\}\}\_\{m,l\}\[n,d\]\(13\)σm,l,n2\\displaystyle\\sigma^\{2\}\_\{m,l,n\}:=1D​∑d=1D\(𝐇m,l​\[n,d\]−μm,l,n\)2\.\\displaystyle:=\\frac\{1\}\{D\}\\sum\_\{d=1\}^\{D\}\(\{\\mathbf\{H\}\}\_\{m,l\}\[n,d\]\-\\mu\_\{m,l,n\}\)^\{2\}\.\(14\)Then, the output of the layer normalization with these parameters is denoted by𝐇¯m,l\\overline\{\\mathbf\{H\}\}\_\{m,l\}whose\(n,d\)\(n,d\)\-entry is

𝐇¯m,l​\[n,d\]:=\(𝐇m,l​\[n,d\]−μm,l,n\)/σm,l,n\\displaystyle\\overline\{\\mathbf\{H\}\}\_\{m,l\}\[n,d\]:=\(\{\\mathbf\{H\}\}\_\{m,l\}\[n,d\]\-\\mu\_\{m,l,n\}\)/\\sigma\_\{m,l,n\}\(15\)ford=1,…,Dd=1,\\ldots,Dat eachnn\. Finally, the output of thell\-th layer is given by

𝐐m,l\+1=𝐐m,l\+𝐇¯m,l,\\displaystyle\{\\mathbf\{Q\}\}\_\{m,l\+1\}=\{\\mathbf\{Q\}\}\_\{m,l\}\+\\overline\{\\mathbf\{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](https://arxiv.org/html/2607.15819#bib.bib1)\]\)\.

𝐐m,L\{\\mathbf\{Q\}\}\_\{m,L\}is the output of the last layer\. It is then flattened and obtain𝒒m\{\\boldsymbol\{q\}\}\_\{m\}which is aD​\(N\+1\)×1D\(N\+1\)\\times 1vector\. Let𝐖out\{\\mathbf\{W\}\}\_\{\\rm out\}andboutb\_\{\\rm out\}be aD​\(N\+1\)×1D\(N\+1\)\\times 1output weight matrix and output bias respectively\. We then obtain the transformer output by

Qout​\(𝐏m\):=𝐖out⊤​𝒒m\+bout\\displaystyle Q\_\{\\rm out\}\(\{\\mathbf\{P\}\}\_\{m\}\):=\{\\mathbf\{W\}\}\_\{\\rm out\}^\{\\top\}\{\\boldsymbol\{q\}\}\_\{m\}\+b\_\{\\rm out\}\(17\)which is a scalar\.

We refer to this transformer asTF​\(D,L,K\)\{\\rm TF\}\(D,L,K\), whereDD,LLandKKare the model dimension, the number of layers and heads respectively\.

## 3A 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\}\-\{\\overline\{x\}\}\_\{m\}\|/Nand\|y¯m\|/N\|\{\\overline\{y\}\}\_\{m\}\|/Nbe negligible in the calculation below\. If we naturally assume that\|um−x¯m\|\|u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\|and\|y¯m\|\|\{\\overline\{y\}\}\_\{m\}\|do not increase asNNincreases then this is satisfied whenNNis sufficiently large\. With this case in mind, we assume thatNNis sufficiently large so that\|um\|/N\|u\_\{m\}\|/N,\|x¯m\|/N\|\{\\overline\{x\}\}\_\{m\}\|/Nand\|y¯m\|/N\|\{\\overline\{y\}\}\_\{m\}\|/Nare negligible for anymm\. We will address this assumption in a later section

We here fix the structure of the transformer, in which the model dimensionD=4D=4, the number of layersL=2L=2and the number of headsK=2K=2\. Therefore, the transformer isTF​\(4,2,2\)\{\\rm TF\}\(4,2,2\)\. Now, themm\-th prompt toTF​\(4,2,2\)\{\\rm TF\}\(4,2,2\)is𝐏m\{\\mathbf\{P\}\}\_\{m\}defined in \([1](https://arxiv.org/html/2607.15819#S2.E1)\)\.

### 3\.1Input embedding

LetRRbe a real positive number\. We assume thatRRcan be chosen sufficiently large below\. As in\[[1](https://arxiv.org/html/2607.15819#bib.bib1)\],RRis important for approximate implementation of the closed form solution\.

In the input embedding, we set

𝐖in=\[10001R​N0001R​N2​NN\+1−1R​N−1R​N\]\\displaystyle\{\\mathbf\{W\}\}\_\{\\rm in\}=\\begin\{bmatrix\}1&0&0\\\\ 0&\\frac\{1\}\{RN\}&0\\\\ 0&0&\\frac\{1\}\{RN\}\\\\ \\frac\{\\sqrt\{2\}N\}\{N\+1\}&\-\\frac\{1\}\{RN\}&\-\\frac\{1\}\{RN\}\\end\{bmatrix\}\(18\)and𝒃in=𝐎4,1\{\\boldsymbol\{b\}\_\{\\rm in\}\}=\{\\mathbf\{O\}\}\_\{4,1\}\. By \([10](https://arxiv.org/html/2607.15819#S2.E10)\), thenn\-th row of the input embedding is

𝐐m,1​\[n,:\]⊤=\[11R​N​xm,n1R​N​ym,n2​NN\+1−1R​N​xm,n−1R​N​ym,n\],\\displaystyle\{\\mathbf\{Q\}\}\_\{m,1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\\\\ \\frac\{1\}\{RN\}x\_\{m,n\}\\\\ \\frac\{1\}\{RN\}y\_\{m,n\}\\\\ \\frac\{\\sqrt\{2\}N\}\{N\+1\}\-\\frac\{1\}\{RN\}x\_\{m,n\}\-\\frac\{1\}\{RN\}y\_\{m,n\}\\end\{bmatrix\},\(19\)which is the input to the transformer\.

### 3\.2Output of the first layer

We here construct the first layer\.

We consider the first head\. We have

ξ1,1,1\\displaystyle\\xi\_\{1,1,1\}:=𝐐m,1​\[:,1\]⊤​𝐐m,1​\[:,1\]=N\+1\.\\displaystyle:=\{\\mathbf\{Q\}\}\_\{m,1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,1\}\[:,1\]=N\+1\.\(20\)and

ξ1,1,2\\displaystyle\\xi\_\{1,1,2\}:=𝐐m,1​\[:,1\]⊤​𝐐m,1​\[:,2\]\\displaystyle:=\{\\mathbf\{Q\}\}\_\{m,1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,1\}\[:,2\]=1R​N​\(N​x¯m\+um\)≃1R​x¯m,\\displaystyle=\\frac\{1\}\{RN\}\(N\{\\overline\{x\}\}\_\{m\}\+u\_\{m\}\)\\simeq\\frac\{1\}\{R\}\{\\overline\{x\}\}\_\{m\},\(21\)where the last approximation holds by the assumption\. Sinceym,N\+1=0y\_\{m,N\+1\}=0, we have

ξ1,1,3\\displaystyle\\xi\_\{1,1,3\}:=𝐐m,1​\[:,1\]⊤​𝐐m,1​\[:,3\]\\displaystyle:=\{\\mathbf\{Q\}\}\_\{m,1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,1\}\[:,3\]=1R​N​N​y¯m=1R​y¯m\.\\displaystyle=\\frac\{1\}\{RN\}N\{\\overline\{y\}\}\_\{m\}=\\frac\{1\}\{R\}\{\\overline\{y\}\}\_\{m\}\.\(22\)By the same way as above, we also have

ξ1,1,4\\displaystyle\\xi\_\{1,1,4\}:=𝐐m,1​\[:,1\]⊤​𝐐m,1​\[:,4\]\\displaystyle:=\{\\mathbf\{Q\}\}\_\{m,1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,1\}\[:,4\]≃2N−1Rx¯m\+−1Ry¯m\.\\displaystyle\\simeq\\sqrt\{2\}N\-\\frac\{1\}\{R\}\{\\overline\{x\}\}\_\{m\}\+\-\\frac\{1\}\{R\}\{\\overline\{y\}\}\_\{m\}\.\(23\)Therefore, we have

𝐐m,1⊤​𝐐m,1=\[ξ1,1,1ξ1,1,2ξ1,1,3ξ1,1,4−−−−−−−−−−−−\],\\displaystyle\{\\mathbf\{Q\}\}\_\{m,1\}^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,1\}=\\begin\{bmatrix\}\\xi\_\{1,1,1\}&\\xi\_\{1,1,2\}&\\xi\_\{1,1,3\}&\\xi\_\{1,1,4\}\\\\ \-&\-&\-&\-\\\\ \-&\-&\-&\-\\\\ \-&\-&\-&\-\\\\ \\end\{bmatrix\},\(24\)where the irrelevant elements are indicated by “−\-”\. We set𝚯1,1=\{𝐖1,1,1,𝐖1,1,2,𝐖1,1,3\}\{\\boldsymbol\{\\Theta\}\}\_\{1,1\}=\\\{\{\\mathbf\{W\}\}\_\{1,1,1\},\{\\mathbf\{W\}\}\_\{1,1,2\},\{\\mathbf\{W\}\}\_\{1,1,3\}\\\}in which

𝐖1,1,1\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,1\}=\[1000000000000000\]\\displaystyle=\\begin\{bmatrix\}1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(25\)𝐖1,1,2\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,2\}=\[NN\+10000−120000−120000−12\]\\displaystyle=\\begin\{bmatrix\}\\frac\{N\}\{N\+1\}&0&0&0\\\\ 0&\-\\frac\{1\}\{\\sqrt\{2\}\}&0&0\\\\ 0&0&\-\\frac\{1\}\{\\sqrt\{2\}\}&0\\\\ 0&0&0&\\frac\{\-1\}\{\\sqrt\{2\}\}\\\\ \\end\{bmatrix\}\(26\)𝐖1,1,3\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,3\}=\[1000000000000000\]\.\\displaystyle=\\begin\{bmatrix\}1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\.\(27\)Then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\) and \([11](https://arxiv.org/html/2607.15819#S2.E11)\), thenn\-th row of the output of the first head is

𝐇m,1,1​\[n,:\]⊤=\[N−12​R​x¯m−12​R​y¯m−N\+12​R​x¯m\+12​R​y¯m\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{m,1,1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}N\\\\ \-\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{x\}\}\_\{m\}\\\\ \-\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{y\}\}\_\{m\}\\\\ \-N\+\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{x\}\}\_\{m\}\+\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{y\}\}\_\{m\}\\end\{bmatrix\}\.\(28\)
On the other hand, for the second head, we set𝚯1,2=\{𝐖1,2,1,𝐖1,2,2,𝐖1,2,3\}\{\\boldsymbol\{\\Theta\}\}\_\{1,2\}=\\\{\{\\mathbf\{W\}\}\_\{1,2,1\},\{\\mathbf\{W\}\}\_\{1,2,2\},\{\\mathbf\{W\}\}\_\{1,2,3\}\\\}in which

𝐖1,2,1=𝐖1,2,2=𝐖1,2,3=𝐎d,d\.\\displaystyle\{\\mathbf\{W\}\}\_\{1,2,1\}=\{\\mathbf\{W\}\}\_\{1,2,2\}=\{\\mathbf\{W\}\}\_\{1,2,3\}=\{\\mathbf\{O\}\}\_\{d,d\}\.\(29\)Then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\) and \([11](https://arxiv.org/html/2607.15819#S2.E11)\), thenn\-th row of the output of the second head is

𝐇m,1,2​\[n,:\]=𝐎1,4\.\\displaystyle\{\\mathbf\{H\}\}\_\{m,1,2\}\[n,:\]=\{\\mathbf\{O\}\}\_\{1,4\}\.\(30\)
Therefore, by \([12](https://arxiv.org/html/2607.15819#S2.E12)\), thenn\-th row of the output of the MHLSA in the first layer is

𝐇m,1​\[n,:\]⊤\\displaystyle\{\\mathbf\{H\}\}\_\{m,1\}\[n,:\]^\{\\top\}=𝐇m,1,1​\[n,:\]⊤\\displaystyle=\{\\mathbf\{H\}\}\_\{m,1,1\}\[n,:\]^\{\\top\}=\[N−12​R​x¯m−12​R​y¯m−N\+12​R​x¯m\+12​R​y¯m\]\\displaystyle=\\begin\{bmatrix\}N\\\\ \-\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{x\}\}\_\{m\}\\\\ \-\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{y\}\}\_\{m\}\\\\ \-N\+\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{x\}\}\_\{m\}\+\\frac\{1\}\{\\sqrt\{2\}R\}\{\\overline\{y\}\}\_\{m\}\\end\{bmatrix\}\(31\)
Now, we consider the layer normalization\. By \([13](https://arxiv.org/html/2607.15819#S2.E13)\), \([14](https://arxiv.org/html/2607.15819#S2.E14)\) and \([3\.2](https://arxiv.org/html/2607.15819#S3.Ex5)\), for𝐇m,1​\[n,:\]\{\\mathbf\{H\}\}\_\{m,1\}\[n,:\], we have

μm,1,n\\displaystyle\\mu\_\{m,1,n\}=0\\displaystyle=0\(32\)σm,1,n2\\displaystyle\\sigma^\{2\}\_\{m,1,n\}=14​\(N2\+O​\(1R2\)\+N2\+O​\(NR\)\)≃N22,\\displaystyle=\\frac\{1\}\{4\}\\left\(N^\{2\}\+O\\left\(\\frac\{1\}\{R^\{2\}\}\\right\)\+N^\{2\}\+O\\left\(\\frac\{N\}\{R\}\\right\)\\right\)\\simeq\\frac\{N^\{2\}\}\{2\},\(33\)since we can set a sufficiently large value forRR; e\.g\.,R=NpR=N^\{p\}, whereppis a sufficiently large positive integer\.

By \([15](https://arxiv.org/html/2607.15819#S2.E15)\), \([3\.2](https://arxiv.org/html/2607.15819#S3.Ex5)\), \([32](https://arxiv.org/html/2607.15819#S3.E32)\) and \([33](https://arxiv.org/html/2607.15819#S3.E33)\), thenn\-th row of the output of the layer normalization is

𝐇¯m,1​\[n,:\]⊤=\[2−1R​N​x¯m−1R​N​y¯m−2\+1R​N​x¯m,n\+1R​N​y¯m,n\]\.\\displaystyle\\overline\{\\mathbf\{H\}\}\_\{m,1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}\\sqrt\{2\}\\\\ \-\\frac\{1\}\{RN\}\{\\overline\{x\}\}\_\{m\}\\\\ \-\\frac\{1\}\{RN\}\{\\overline\{y\}\}\_\{m\}\\\\ \-\\sqrt\{2\}\+\\frac\{1\}\{RN\}\{\\overline\{x\}\}\_\{m,n\}\+\\frac\{1\}\{RN\}\{\\overline\{y\}\}\_\{m,n\}\\end\{bmatrix\}\.\(34\)
Then, by \([16](https://arxiv.org/html/2607.15819#S2.E16)\), \([19](https://arxiv.org/html/2607.15819#S3.E19)\) and \([34](https://arxiv.org/html/2607.15819#S3.E34)\), thenn\-th row of the output of the first layer is

𝐐m,2​\[n,:\]⊤=\[1\+21R​N​Δ​xm,n1R​N​Δ​ym,n−2N\+1\],\\displaystyle\{\\mathbf\{Q\}\}\_\{m,2\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\+\\sqrt\{2\}\\\\ \\frac\{1\}\{RN\}\\Delta x\_\{m,n\}\\\\ \\frac\{1\}\{RN\}\\Delta y\_\{m,n\}\\\\ \-\\frac\{\\sqrt\{2\}\}\{N\+1\}\\end\{bmatrix\},\(35\)where

Δ​xm,n\\displaystyle\\Delta x\_\{m,n\}=xm,n−x¯m\\displaystyle=x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\(36\)Δ​ym,n\\displaystyle\\Delta y\_\{m,n\}=ym,n−y¯m\.\\displaystyle=y\_\{m,n\}\-\{\\overline\{y\}\}\_\{m\}\.\(37\)and theO​\(1/\(R​N\)\)O\(1/\(RN\)\)terms in𝐐m,2​\[n,4\]\{\\mathbf\{Q\}\}\_\{m,2\}\[n,4\]is omitted sinceRRis sufficiently large and, indeed, they do not affect the final outcome in our construction\.

### 3\.3Output of the second layer

We next construct the second layer\.

By \([35](https://arxiv.org/html/2607.15819#S3.E35)\), we have

ξ2,2,2\\displaystyle\\xi\_\{2,2,2\}:=Qm,2​\[:,2\]⊤​Qm,2​\[:,2\]\\displaystyle:=Q\_\{m,2\}\[:,2\]^\{\\top\}Q\_\{m,2\}\[:,2\]=1R2​N2​\{N​Vm\+\(um−x¯m\)2\}\\displaystyle=\\frac\{1\}\{R^\{2\}N^\{2\}\}\\\{NV\_\{m\}\+\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)^\{2\}\\\}≃1R2​N​Vm\\displaystyle\\simeq\\frac\{1\}\{R^\{2\}N\}V\_\{m\}\(38\)ξ2,2,3\\displaystyle\\xi\_\{2,2,3\}:=Qm,2​\[:,2\]⊤​Qm,2​\[:,3\]\\displaystyle:=Q\_\{m,2\}\[:,2\]^\{\\top\}Q\_\{m,2\}\[:,3\]=1R2​N2​\{N​Cm\+\(um−x¯m\)​\(0−y¯m\)\}\\displaystyle=\\frac\{1\}\{R^\{2\}N^\{2\}\}\\\{NC\_\{m\}\+\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\(0\-\{\\overline\{y\}\}\_\{m\}\)\\\}≃1R2​N​Cm\\displaystyle\\simeq\\frac\{1\}\{R^\{2\}N\}C\_\{m\}\(39\)under the assumption and, thus, we have

𝐐m,2⊤​𝐐m,2=\[−−−−−ξ2,2,2ξ2,2,3−−−−−−−−−\]\.\\displaystyle\{\\mathbf\{Q\}\}\_\{m,2\}^\{\\top\}\{\\mathbf\{Q\}\}\_\{m,2\}=\\begin\{bmatrix\}\-&\-&\-&\-\\\\ \-&\\xi\_\{2,2,2\}&\\xi\_\{2,2,3\}&\-\\\\ \-&\-&\-&\-&\\\\ \-&\-&\-&\-&\\\\ \\end\{bmatrix\}\.\(40\)
We consider the first head\. We set𝚯2,1=\{𝐖2,1,1,𝐖2,1,2,𝐖2,1,3\}\{\\boldsymbol\{\\Theta\}\}\_\{2,1\}=\\\{\{\\mathbf\{W\}\}\_\{2,1,1\},\{\\mathbf\{W\}\}\_\{2,1,2\},\{\\mathbf\{W\}\}\_\{2,1,3\}\\\}, in which

𝐖2,1,1\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,1\}=\[0100000000000000\]\\displaystyle=\\begin\{bmatrix\}0&1&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(41\)𝐖2,1,2\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,2\}=\[000000000R2​N/2−R2​N/200000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 0&0&0&0\\\\ 0&R^\{2\}N/\\sqrt\{2\}&\-R^\{2\}N/\\sqrt\{2\}&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(42\)𝐖2,1,3\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,3\}=\[0000100000000000\]\.\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\.\(43\)Then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), \([11](https://arxiv.org/html/2607.15819#S2.E11)\) and \([40](https://arxiv.org/html/2607.15819#S3.E40)\), thenn\-th row of the output of the first head is

𝐇m,2,1​\[n,:\]⊤=\[012​R​N​Cm​Δ​xm,n−12​R​N​Cm​Δ​xm,n0\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{m,2,1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}0\\\\ \\frac\{1\}\{\\sqrt\{2\}RN\}C\_\{m\}\\Delta x\_\{m,n\}\\\\ \-\\frac\{1\}\{\\sqrt\{2\}RN\}C\_\{m\}\\Delta x\_\{m,n\}\\\\ 0\\\\ \\end\{bmatrix\}\.\(44\)
On the other hand, we consider the second head\. We set𝚯2,2=\{𝐖2,2,1,𝐖2,2,2,𝐖2,2,3\}\{\\boldsymbol\{\\Theta\}\}\_\{2,2\}=\\\{\{\\mathbf\{W\}\}\_\{2,2,1\},\{\\mathbf\{W\}\}\_\{2,2,2\},\{\\mathbf\{W\}\}\_\{2,2,3\}\\\}in which

𝐖2,2,1\\displaystyle\{\\mathbf\{W\}\}\_\{2,2,1\}=\[0100000000000000\]\\displaystyle=\\begin\{bmatrix\}0&1&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(45\)𝐖2,2,2\\displaystyle\{\\mathbf\{W\}\}\_\{2,2,2\}=\[0000R2​N00−R2​N00000000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ R^\{2\}N&0&0&\-R^\{2\}N\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(46\)𝐖2,2,3\\displaystyle\{\\mathbf\{W\}\}\_\{2,2,3\}=\[1/\(1\+2\)00000000000\]\.\\displaystyle=\\begin\{bmatrix\}1/\(1\+\\sqrt\{2\}\)&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\.\(47\)Then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), \([11](https://arxiv.org/html/2607.15819#S2.E11)\) and \([40](https://arxiv.org/html/2607.15819#S3.E40)\), thenn\-th row of the output of the second head is

𝐇m,2,2​\[n,:\]⊤=\[Vm00−Vm\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{m,2,2\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}V\_\{m\}\\\\ 0\\\\ 0\\\\ \-V\_\{m\}\\\\ \\end\{bmatrix\}\.\(48\)
Therefore, by \([12](https://arxiv.org/html/2607.15819#S2.E12)\), thenn\-th row of the output of the MHLSA in the second layer is

𝐇m,2​\[n,:\]⊤=\[Vm12​R​N​Cm​Δ​xm,n−12​R​N​Cm​Δ​xm,n−Vm\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{m,2\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}V\_\{m\}\\\\ \\frac\{1\}\{\\sqrt\{2\}RN\}C\_\{m\}\\Delta x\_\{m,n\}\\\\ \-\\frac\{1\}\{\\sqrt\{2\}RN\}C\_\{m\}\\Delta x\_\{m,n\}\\\\ \-V\_\{m\}\\\\ \\end\{bmatrix\}\.\(49\)
Now, we consider the layer normalization\. By \([13](https://arxiv.org/html/2607.15819#S2.E13)\), \([14](https://arxiv.org/html/2607.15819#S2.E14)\) and \([49](https://arxiv.org/html/2607.15819#S3.E49)\), for𝐇m,2​\[n,:\]\{\\mathbf\{H\}\}\_\{m,2\}\[n,:\], we have

μm,2,n\\displaystyle\\mu\_\{m,2,n\}=0\\displaystyle=0\(50\)σm,2,n2\\displaystyle\\sigma^\{2\}\_\{m,2,n\}=14​\(2​Vm2\+O​\(1R2​N2\)\)≃12​Vm2\\displaystyle=\\frac\{1\}\{4\}\\left\(2V\_\{m\}^\{2\}\+O\\left\(\\frac\{1\}\{R^\{2\}N^\{2\}\}\\right\)\\right\)\\simeq\\frac\{1\}\{2\}V\_\{m\}^\{2\}\(51\)sinceRRis sufficiently large\. By \([15](https://arxiv.org/html/2607.15819#S2.E15)\), thenn\-th row of the output of the layer normalization is

𝐇¯m,2​\[n,:\]⊤=\[21R​N​CmVm​Δ​xm,n−1R​N​CmVm​Δ​xm,n−2\]\.\\displaystyle\\overline\{\\mathbf\{H\}\}\_\{m,2\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}\\sqrt\{2\}\\\\ \\frac\{1\}\{RN\}\\frac\{C\_\{m\}\}\{V\_\{m\}\}\\Delta x\_\{m,n\}\\\\ \-\\frac\{1\}\{RN\}\\frac\{C\_\{m\}\}\{V\_\{m\}\}\\Delta x\_\{m,n\}\\\\ \-\\sqrt\{2\}\\\\ \\end\{bmatrix\}\.\(52\)The division performed by the layer normalization is actually used for this purpose in\[[1](https://arxiv.org/html/2607.15819#bib.bib1)\]\.

Then, by \([16](https://arxiv.org/html/2607.15819#S2.E16)\), \([35](https://arxiv.org/html/2607.15819#S3.E35)\) and \([52](https://arxiv.org/html/2607.15819#S3.E52)\), thenn\-th row of the output of the second layer is

𝐐m,3​\[n,:\]⊤=\[1\+2​21R​N​Δ​xm,n\+1R​N​CmVm​Δ​xm,n1R​N​Δ​ym,n−1R​N​CmVm​Δ​xm,n−2​\(N\+2\)/\(N\+1\)\]\.\\displaystyle\{\\mathbf\{Q\}\}\_\{m,3\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\+2\\sqrt\{2\}\\\\ \\frac\{1\}\{RN\}\\Delta x\_\{m,n\}\+\\frac\{1\}\{RN\}\\frac\{C\_\{m\}\}\{V\_\{m\}\}\\Delta x\_\{m,n\}\\\\ \\frac\{1\}\{RN\}\\Delta y\_\{m,n\}\-\\frac\{1\}\{RN\}\\frac\{C\_\{m\}\}\{V\_\{m\}\}\\Delta x\_\{m,n\}\\\\ \-\\sqrt\{2\}\(N\+2\)/\(N\+1\)\\\\ \\end\{bmatrix\}\.\(53\)

### 3\.4Model output

Sincexm,N\+1=umx\_\{m,N\+1\}=u\_\{m\}andym,N\+1=0y\_\{m,N\+1\}=0, we have

𝐐m,3​\[N\+1,3\]=−1R​N​\(y¯m\+CmVm​\(um−x¯m\)\)\\displaystyle\{\\mathbf\{Q\}\}\_\{m,3\}\[N\+1,3\]=\-\\frac\{1\}\{RN\}\\left\(\{\\overline\{y\}\}\_\{m\}\+\\frac\{C\_\{m\}\}\{V\_\{m\}\}\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\\right\)\(54\)by \([53](https://arxiv.org/html/2607.15819#S3.E53)\)\. Now,𝐐m,3\{\\mathbf\{Q\}\}\_\{m,3\}is flattened into𝒒m\{\\boldsymbol\{q\}\}\_\{m\}\. Then, we have𝒒m​\[D​\(N\+1\)−1\]=𝐐m,3​\[N\+1,3\]\{\\boldsymbol\{q\}\}\_\{m\}\[D\(N\+1\)\-1\]=\{\\mathbf\{Q\}\}\_\{m,3\}\[N\+1,3\]\. Hence, if we set

𝐖out​\[i\]=\{−R​Ni=D​\(N\+1\)−10otherwise\\displaystyle\{\\mathbf\{W\}\}\_\{\\rm out\}\[i\]=\\begin\{cases\}\-RN&i=D\(N\+1\)\-1\\\\ 0&\{\\rm otherwise\}\\end\{cases\}\(55\)then the transformer output for themm\-th prompt is, approximately,

Qout​\(𝐏m\)=y¯m\+CmVm​\(um−x¯m\)\\displaystyle Q\_\{\\rm out\}\(\{\\mathbf\{P\}\}\_\{m\}\)=\{\\overline\{y\}\}\_\{m\}\+\\frac\{C\_\{m\}\}\{V\_\{m\}\}\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\(56\)by \([17](https://arxiv.org/html/2607.15819#S2.E17)\) and \([54](https://arxiv.org/html/2607.15819#S3.E54)\)\. This is consistent with the least squares estimate atumu\_\{m\}in \([8](https://arxiv.org/html/2607.15819#S2.E8)\)\.

### 3\.5Discussion

In our construction, we need to choose a large value forNNthat is the number of in\-context samples\. This is because, for example as in \([3\.2](https://arxiv.org/html/2607.15819#S3.Ex2)\), 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](https://arxiv.org/html/2607.15819#S2.Ex1)\)\. 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](https://arxiv.org/html/2607.15819#bib.bib1)\]\. 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](https://arxiv.org/html/2607.15819#bib.bib6)\]\. Therefore, a largeNNassumption 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 forRRis 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](https://arxiv.org/html/2607.15819#A1), 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\.

## 4Numerical experiment

We here show several numerical experiments to demonstrate that our construction in the previous section is actually used in the trained transformer\.

### 4\.1Setting 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 usingtorch\.nn\.Linear\. For the training data𝐃m\{\\mathbf\{D\}\}\_\{m\}defined by \([3](https://arxiv.org/html/2607.15819#S2.E3)\), we setσ2=0\.04\\sigma^\{2\}=0\.04for the noise variance and, at eachmm, the inputs including prediction points are randomly drawn from a standard normal distributionN​\(0,1\)N\(0,1\)\. And, we generateθm,0\\theta\_\{m,0\}andθm,1\\theta\_\{m,1\}according toN​\(0,1\)N\(0,1\)independently at eachmm\. The length of the input sequence isN\+1=51N\+1=51\(the number of in\-context samples isN=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 isTF​\(4,2,2\)\{\\rm TF\}\(4,2,2\)\. For the training, the number of training data isM=5000M=5000, the number of validation data is10001000, the batch size is500500, the learning rate is0\.0050\.005, the maximum number of training epochs is2000020000\. We apply the layer normalization without an element\-wise affine setting intorch\.nn\.LayerNorm\. As mentioned in Sec\.[3\.5](https://arxiv.org/html/2607.15819#S3.SS5), there is a lack of uniqueness in the LSA representation\. Therefore, we apply theℓ1\\ell\_\{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 every100100epochs\.

### 4\.2Verifying the use of layer normalization

We execute1010different training runs under the above setting, in which both of data and initialization differ for each run\. We generate10001000new 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 for10001000new data\.

Simultaneously, for eachmm\-th prompt in the new data, we can calculate𝐇m,2\{\\mathbf\{H\}\}\_\{m,2\}defined in \([12](https://arxiv.org/html/2607.15819#S2.E12)\)\. The variance along the embedding used in the layer normalization is the variance of𝐇m,2​\[n,:\]\{\\mathbf\{H\}\}\_\{m,2\}\[n,:\]for thenn\-th sequence, which isσm,2,n2\\sigma\_\{m,2,n\}^\{2\}defined in \([14](https://arxiv.org/html/2607.15819#S2.E14)\)\. Also, we can calculateVmV\_\{m\}defined in \([6](https://arxiv.org/html/2607.15819#S2.E6)\) as the variance of in\-context input samples in the new data\. In our construction, the division byVmV\_\{m\}in \([8](https://arxiv.org/html/2607.15819#S2.E8)\) is executed using the division by the standard deviation in the layer normalization in \([52](https://arxiv.org/html/2607.15819#S3.E52)\)\. To see this, we calculate the correlation coefficient between1/σm,2,511/\\sigma\_\{m,2,51\},m=1,…,1000m=1,\\ldots,1000and1/Vm1/V\_\{m\},m=1,…,1000m=1,\\ldots,1000in each run\.

The results for1010runs are summarized in Table[1](https://arxiv.org/html/2607.15819#S4.T1)\. We can see that the correlation coefficient is larger than0\.90\.9in the first and ninth runs\. This implies that the division byVmV\_\{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
### 4\.3Input and output weights

In Fig\.[2](https://arxiv.org/html/2607.15819#S4.F2)\(a\) and \(b\), for the first run, we show the absolute values of the input and output weights which are denoted by𝐖in\{\\mathbf\{W\}\}\_\{\\rm in\}and𝐖out\{\\mathbf\{W\}\}\_\{\\rm out\}in this paper\. The weight values are shown on a logarithmic scale\.𝐖in\{\\mathbf\{W\}\}\_\{\\rm in\}is a3×43\\times 4matrix\. 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𝐖out\{\\mathbf\{W\}\}\_\{\\rm out\}is a204204\-dimensional vector, we here convert it into a4×514\\times 51matrix whose row and column correspond to the indices of embeddings and sequences respectively\. We refer to the entry in rowiiand columnjjof 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 thejj\-th sequence of theii\-th embedding\. We refer to the cell connected to the\(i,j\)\(i,j\)weight as the\(i,j\)\(i,j\)cell\.

![Refer to caption](https://arxiv.org/html/2607.15819v1/x3.png)

\(a\) Input weights

![Refer to caption](https://arxiv.org/html/2607.15819v1/x4.png)

\(b\) Output weights

Figure 2:Absolute weight valuesIn Fig\.[2](https://arxiv.org/html/2607.15819#S4.F2)\(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 thatyyis 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 thatxxis 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](https://arxiv.org/html/2607.15819#A1)\.

In Fig\.[2](https://arxiv.org/html/2607.15819#S4.F2)\(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\.4Regression 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](https://arxiv.org/html/2607.15819#S2.E8)\) by using a regression analysis\. We generate10001000new data for the prompt to the trained transformer, which are i\.i\.d\. samples from the same distribution as the training data\. Then, for eachmm\-th new prompt, we can obtain𝐐m,3\{\\mathbf\{Q\}\}\_\{m,3\}in \([53](https://arxiv.org/html/2607.15819#S3.E53)\) and focus onzm:=𝐐m,3​\[51,3\]z\_\{m\}:=\{\\mathbf\{Q\}\}\_\{m,3\}\[51,3\], which is the output of the\(3,51\)\(3,51\)cell in the second layer\. We then fitzmz\_\{m\}by

fm:=α0\+α1​\(um−x¯m\)\+α2​y¯m\+α3​CmVm​\(um−x¯m\),\\displaystyle f\_\{m\}:=\\alpha\_\{0\}\+\\alpha\_\{1\}\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\+\\alpha\_\{2\}\{\\overline\{y\}\}\_\{m\}\+\\alpha\_\{3\}\\frac\{C\_\{m\}\}\{V\_\{m\}\}\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\),\(57\)where\(α0,α1,α2,α3\)\(\\alpha\_\{0\},\\alpha\_\{1\},\\alpha\_\{2\},\\alpha\_\{3\}\)are regression coefficients\.

Table 2:Results of regression analysisThe coefficient estimates are summarized in Table[2](https://arxiv.org/html/2607.15819#S4.T2), in which we show the standardized regression coefficients in the brackets\. In Table[2](https://arxiv.org/html/2607.15819#S4.T2), the estimates ofα0\\alpha\_\{0\}andα1\\alpha\_\{1\}are small and those ofα2\\alpha\_\{2\}andα3\\alpha\_\{3\}are large\. Here, the correlation betweenzmz\_\{m\}andfmf\_\{m\}was0\.994790\.99479\. Thus,zmz\_\{m\}is found to be the form of the last two terms infmf\_\{m\}\. It is the same form of the least squares solution in \([8](https://arxiv.org/html/2607.15819#S2.E8)\)\. And, the correlation betweenzmz\_\{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\.5Prediction

We finally show the fitting curve of the trained transformer in the first run to visually confirm that our conclusion is valid\. We generate100100new data for the prompts to the trained transformer, which are\{\(x1,y1\),…,\(xn,y50\),\(uk,0\)\}\\\{\(x\_\{1\},y\_\{1\}\),\\ldots,\(x\_\{n\},y\_\{50\}\),\(u\_\{k\},0\)\\\}fork=1,…,100k=1,\\ldots,100, where\(xi,yi\)\(x\_\{i\},y\_\{i\}\)is generated from the same distribution as the training data anduku\_\{k\}is the equidistant points in\[−3,3\]\[\-3,3\]\. In other words, we obtain the transformer output at100100different prediction points under the same in\-context samples\. Therefore, we can plot a fitting curve for the in\-context samples\.

In Fig\.[3](https://arxiv.org/html/2607.15819#S4.F3)\(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\}\),\\ldots,\(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](https://arxiv.org/html/2607.15819#S4.F3)\(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](https://arxiv.org/html/2607.15819#S3.E54)\) 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\.

![Refer to caption](https://arxiv.org/html/2607.15819v1/x5.png)

\(a\) Transformer output

![Refer to caption](https://arxiv.org/html/2607.15819v1/x6.png)

\(b\) Output through\(3,51\)\(3,51\)cell

Figure 3:Fitting curve

## 5Conclusion 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](https://arxiv.org/html/2607.15819#bib.bib1)\]\. 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\\ell\_\{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](https://arxiv.org/html/2607.15819#S4.F3), 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](https://arxiv.org/html/2607.15819#S4.F2)\(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](https://arxiv.org/html/2607.15819#S4.T1), the trained transformer in the fourth run may not implement the division byVmV\_\{m\}using the layer normalization in the second layer\. Nevertheless, the test error is low\. Furthermore, in training without theℓ1\\ell\_\{1\}regularization, similar results are obtained\. Thus, there may be different types of solver for obtaining the least squares estimate forTF​\(4,2,2\)\{\\rm 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 inTF​\(4,2,2\)\{\\rm 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 a33\-layer transformer with linear attention; e\.g\., see[B](https://arxiv.org/html/2607.15819#A2)\. 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 largeNNassumption is almost accurate; i\.e\., for example, in \([3\.2](https://arxiv.org/html/2607.15819#S3.Ex2)\), the mean ofNNsamples is almost consistent with that ofN\+1N\+1samples\. 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](https://arxiv.org/html/2607.15819#S3.SS5), a largeNNapproximation 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, pp\.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, pp\.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, pp\.1\-55, 2023\.
- \[14\]N\. Ravishanker, D\. K\. Dey, “A first course in linear model theory”, Chapman & Hall/CRC, 2002\.

## Appendix AAn example of a lack of uniqueness

We show here a simple example, in which LSA produces the same output for the different inputs\.

Let𝐐\{\\mathbf\{Q\}\}be anN×3N\\times 3matrix whosenn\-th row is\[1,an,bn\]\[1,a\_\{n\},b\_\{n\}\]and definea:=∑n=1Nana:=\\sum\_\{n=1\}^\{N\}a\_\{n\}andb:=∑n=1Nbnb:=\\sum\_\{n=1\}^\{N\}b\_\{n\}\. In \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), we set this𝐐\{\\mathbf\{Q\}\}as the input matrix and we set𝚯=\{𝐖1,𝐖2,𝐖3\}\{\\boldsymbol\{\\Theta\}\}=\\\{\{\\mathbf\{W\}\}\_\{1\},\{\\mathbf\{W\}\}\_\{2\},\{\\mathbf\{W\}\}\_\{3\}\\\}, where

𝐖1=𝐖3=\[100000000\],𝐖2=\[000010001\]\.\\displaystyle\{\\mathbf\{W\}\}\_\{1\}=\{\\mathbf\{W\}\}\_\{3\}=\\begin\{bmatrix\}1&0&0\\\\ 0&0&0\\\\ 0&0&0\\\\ \\end\{bmatrix\},~\{\\mathbf\{W\}\}\_\{2\}=\\begin\{bmatrix\}0&0&0\\\\ 0&1&0\\\\ 0&0&1\\\\ \\end\{bmatrix\}\.\(58\)Then, thenn\-th row ofLSA𝚯​\(𝐐\)\{\\rm LSA\}\_\{\{\\boldsymbol\{\\Theta\}\}\}\(\{\\mathbf\{Q\}\}\)is\[0,a,b\]\[0,a,b\]for anynn\. On the other hand, let𝐐\{\\mathbf\{Q\}\}be anN×3N\\times 3matrix whosenn\-th row is\[1,an,an\+bn\]\[1,a\_\{n\},a\_\{n\}\+b\_\{n\}\]\. In \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), by setting

𝐖1=𝐖3=\[100000000\],𝐖2=\[00001−1001\],\\displaystyle\{\\mathbf\{W\}\}\_\{1\}=\{\\mathbf\{W\}\}\_\{3\}=\\begin\{bmatrix\}1&0&0\\\\ 0&0&0\\\\ 0&0&0\\\\ \\end\{bmatrix\},~\{\\mathbf\{W\}\}\_\{2\}=\\begin\{bmatrix\}0&0&0\\\\ 0&1&\-1\\\\ 0&0&1\\\\ \\end\{bmatrix\},\(59\)thenn\-th row ofLSA𝚯​\(𝐐\)\{\\rm LSA\}\_\{\{\\boldsymbol\{\\Theta\}\}\}\(\{\\mathbf\{Q\}\}\)is\[0,a,b\]\[0,a,b\]for anynn\. This is the same as in the first case\.

## Appendix BImplementation of gradient descent method

### B\.1Gradient descent method for a simple regression

The least squares estimates atumu\_\{m\}is given by \([8](https://arxiv.org/html/2607.15819#S2.E8)\)\. Here, we consider to obtain this solution by using a gradient descent method\. We define

fθm​\(um\):=y¯m\+θm​\(u−x¯m\),\\displaystyle f\_\{\\theta\_\{m\}\}\(u\_\{m\}\):=\{\\overline\{y\}\}\_\{m\}\+\\theta\_\{m\}\(u\-\{\\overline\{x\}\}\_\{m\}\),\(60\)whereθm\\theta\_\{m\}is a parameter to be adjusted\. Thus,θ^m=Cm/Vm\{\\widehat\{\\theta\}\}\_\{m\}=C\_\{m\}/V\_\{m\}is a closed form solution toθm\\theta\_\{m\}\. We also define

Sm​\(θ\):=12​∑n=1N\(ym,n−fθ​\(xm,n\)\)2\.\\displaystyle S\_\{m\}\(\\theta\):=\\frac\{1\}\{2\}\\sum\_\{n=1\}^\{N\}\(y\_\{m,n\}\-f\_\{\\theta\}\(x\_\{m,n\}\)\)^\{2\}\.\(61\)We have

∂Sm​\(θ\)∂θ\\displaystyle\\frac\{\\partial S\_\{m\}\(\\theta\)\}\{\\partial\\theta\}=−∑n=1N\(ym,n−fθ​\(xm,n\)\)​\(xm,n−x¯m\)\\displaystyle=\-\\sum\_\{n=1\}^\{N\}\(y\_\{m,n\}\-f\_\{\\theta\}\(x\_\{m,n\}\)\)\(x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\)=−∑n=1N\(\(ym,n−y¯m\)−θ​\(xm,n−x¯m\)\)​\(xm,n−x¯m\)\\displaystyle=\-\\sum\_\{n=1\}^\{N\}\(\(y\_\{m,n\}\-\{\\overline\{y\}\}\_\{m\}\)\-\\theta\(x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\)\)\(x\_\{m,n\}\-\{\\overline\{x\}\}\_\{m\}\)=−N​\(Cm−θ​Vm\)\.\\displaystyle=\-N\(C\_\{m\}\-\\theta V\_\{m\}\)\.\(62\)By settingθm​\(0\)\\theta\_\{m\}\(0\)as an initial value, at thett\-th step, the update equation of the gradient descent forθm​\(t\)\\theta\_\{m\}\(t\)is given by

θm​\(t\)\\displaystyle\\theta\_\{m\}\(t\)=θm​\(t−1\)−η​∂Sm​\(θ\)∂θ\|θ=θm​\(t−1\)\\displaystyle=\\theta\_\{m\}\(t\-1\)\-\\eta\\left\.\\frac\{\\partial S\_\{m\}\(\\theta\)\}\{\\partial\\theta\}\\right\|\_\{\\theta=\\theta\_\{m\}\(t\-1\)\}=θm​\(t−1\)\+N​η​\(Cm−θm​\(t−1\)​Vm\),\\displaystyle=\\theta\_\{m\}\(t\-1\)\+N\\eta\(C\_\{m\}\-\\theta\_\{m\}\(t\-1\)V\_\{m\}\),\(63\)whereη\>0\\eta\>0is a learning rate\.

We here give a solution to the second step\. We set

θm​\(0\)=am\\displaystyle\\theta\_\{m\}\(0\)=a\_\{m\}\(64\)and have

θm​\(1\)\\displaystyle\\theta\_\{m\}\(1\)=θm​\(0\)\+N​η​\(Cm−θm​\(0\)​Vm\)\\displaystyle=\\theta\_\{m\}\(0\)\+N\\eta\(C\_\{m\}\-\\theta\_\{m\}\(0\)V\_\{m\}\)=am\+N​η​\(Cm−am​Vm\)\\displaystyle=a\_\{m\}\+N\\eta\(C\_\{m\}\-a\_\{m\}V\_\{m\}\)\(65\)By a simple calculation, we then have

θm​\(2\)\\displaystyle\\theta\_\{m\}\(2\)=θm​\(1\)\+N​η​\(Cm−θm​\(1\)​Vm\)\\displaystyle=\\theta\_\{m\}\(1\)\+N\\eta\(C\_\{m\}\-\\theta\_\{m\}\(1\)V\_\{m\}\)=am\+2​N​η​Cm−2​N​η​am​Vm−N2​η2​Cm​Vm\+N2​η2​am​Vm2\.\\displaystyle=a\_\{m\}\+2N\\eta C\_\{m\}\-2N\\eta a\_\{m\}V\_\{m\}\-N^\{2\}\\eta^\{2\}C\_\{m\}V\_\{m\}\+N^\{2\}\\eta^\{2\}a\_\{m\}V\_\{m\}^\{2\}\.\(66\)and we obtain the estimate atumu\_\{m\}by

fθm​\(2\)​\(um\)=y¯m\+θm​\(2\)​\(um−x¯m\)\\displaystyle f\_\{\\theta\_\{m\}\(2\)\}\(u\_\{m\}\)=\{\\overline\{y\}\}\_\{m\}\+\\theta\_\{m\}\(2\)\(u\_\{m\}\-\{\\overline\{x\}\}\_\{m\}\)\(67\)with thisθm​\(2\)\\theta\_\{m\}\(2\)\.

In this section, we show an implementation of this22\-step gradient descent byTF​\(4,3,2\)\{\\rm TF\}\(4,3,2\)under the assumption thatNNis large or theO​\(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 subscriptmmthat denotes an index of training data for the transformer\.

### B\.2Input embedding

We define

𝐖in=\[1000α100α2β20α3β3\]\\displaystyle\{\\mathbf\{W\}\}\_\{\\rm in\}=\\begin\{bmatrix\}1&0&0\\\\ 0&\\alpha\_\{1\}&0\\\\ 0&\\alpha\_\{2\}&\\beta\_\{2\}\\\\ 0&\\alpha\_\{3\}&\\beta\_\{3\}\\\\ \\end\{bmatrix\}\(68\)and𝒃in=𝐎D,1\{\\boldsymbol\{b\}\_\{\\rm in\}\}=\{\\mathbf\{O\}\}\_\{D,1\}\. By \([10](https://arxiv.org/html/2607.15819#S2.E10)\), we then have the embedding whoseii\-th row is

𝐐1​\[n,:\]⊤=\[1α1​xnα2​xn\+β2​ynα3​xn\+β3​yn\]\.\\displaystyle\{\\mathbf\{Q\}\}\_\{1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\\\\ \\alpha\_\{1\}x\_\{n\}\\\\ \\alpha\_\{2\}x\_\{n\}\+\\beta\_\{2\}y\_\{n\}\\\\ \\alpha\_\{3\}x\_\{n\}\+\\beta\_\{3\}y\_\{n\}\\\\ \\end\{bmatrix\}\.\(69\)𝐐1\{\\mathbf\{Q\}\}\_\{1\}is the\(N\+1\)×D\(N\+1\)\\times Dmatrix and is the input toTF​\(4,3,2\)\{\\rm TF\}\(4,3,2\)\.

#### B\.2\.1The first layer

We show a construction of the first head\. SincexN\+1=ux\_\{N\+1\}=uandyN\+1=0y\_\{N\+1\}=0, we have

𝐐1​\[:,1\]⊤​𝐐1​\[:,3\]\\displaystyle\{\\mathbf\{Q\}\}\_\{1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{1\}\[:,3\]=∑n=1N\+11⋅\(α2​xn\+β2​yn\)\\displaystyle=\\sum\_\{n=1\}^\{N\+1\}1\\cdot\\left\(\\alpha\_\{2\}x\_\{n\}\+\\beta\_\{2\}y\_\{n\}\\right\)=α2​N​x¯\+β2​N​y¯\+α2​u\\displaystyle=\\alpha\_\{2\}N\{\\overline\{x\}\}\+\\beta\_\{2\}N\{\\overline\{y\}\}\+\\alpha\_\{2\}u≃α2​N​x¯\+β2​N​y¯\\displaystyle\\simeq\\alpha\_\{2\}N\{\\overline\{x\}\}\+\\beta\_\{2\}N\{\\overline\{y\}\}\(70\)𝐐1​\[:,1\]⊤​𝐐1​\[:,4\]\\displaystyle\{\\mathbf\{Q\}\}\_\{1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{1\}\[:,4\]=∑n=1N\+11⋅\(α3​xn\+β3​yn\)\\displaystyle=\\sum\_\{n=1\}^\{N\+1\}1\\cdot\\left\(\\alpha\_\{3\}x\_\{n\}\+\\beta\_\{3\}y\_\{n\}\\right\)=α3​N​x¯\+β3​N​y¯\+α3​u\\displaystyle=\\alpha\_\{3\}N\{\\overline\{x\}\}\+\\beta\_\{3\}N\{\\overline\{y\}\}\+\\alpha\_\{3\}u≃α3​N​x¯\+β3​N​y¯\\displaystyle\\simeq\\alpha\_\{3\}N\{\\overline\{x\}\}\+\\beta\_\{3\}N\{\\overline\{y\}\}\(71\)by omittingO​\(1\)O\(1\)terms under the assumption thatNNis large\. If we set

𝐖1,1,1\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,1\}=\[−1N000000000000000\]\\displaystyle=\\begin\{bmatrix\}\-\\frac\{1\}\{N\}&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(72\)𝐖1,1,2\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,2\}=\[0000000000100001\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&1&0\\\\ 0&0&0&1\\\\ \\end\{bmatrix\}\(73\)𝐖1,1,3\\displaystyle\{\\mathbf\{W\}\}\_\{1,1,3\}=\[1000000000000000\]\\displaystyle=\\begin\{bmatrix\}1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(74\)then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), we have

𝐇1,1​\[n,:\]⊤=\[00−α2​x¯−β2​y¯−α3​x¯−β3​y¯\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{1,1\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}0\\\\ 0\\\\ \-\\alpha\_\{2\}\{\\overline\{x\}\}\-\\beta\_\{2\}\{\\overline\{y\}\}\\\\ \-\\alpha\_\{3\}\{\\overline\{x\}\}\-\\beta\_\{3\}\{\\overline\{y\}\}\\\\ \\end\{bmatrix\}\.\(75\)
We next show a construction of the second head\. We have

𝐐1​\[:,1\]⊤​𝐐1​\[:,2\]≃α1​N​x¯\\displaystyle\{\\mathbf\{Q\}\}\_\{1\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{1\}\[:,2\]\\simeq\\alpha\_\{1\}N\{\\overline\{x\}\}\(76\)whenNNis large\. Thus, if we set

𝐖1,2,1\\displaystyle\{\\mathbf\{W\}\}\_\{1,2,1\}=\[−1N000000000000000\]\\displaystyle=\\begin\{bmatrix\}\-\\frac\{1\}\{N\}&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(77\)𝐖1,2,2\\displaystyle\{\\mathbf\{W\}\}\_\{1,2,2\}=\[0000010000000000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 0&1&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(78\)𝐖1,2,3\\displaystyle\{\\mathbf\{W\}\}\_\{1,2,3\}=\[1000000000000000\]\\displaystyle=\\begin\{bmatrix\}1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(79\)then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), we have

𝐇1,2​\[n,:\]=\[0−α1​x¯00\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{1,2\}\[n,:\]=\\begin\{bmatrix\}0&\-\\alpha\_\{1\}\{\\overline\{x\}\}&0&0\\\\ \\end\{bmatrix\}\.\(80\)
As a result, the output of the first layer is

𝐐2=\(𝐇1,1\+𝐇1,2\)\+𝐐1,\\displaystyle\{\\mathbf\{Q\}\}\_\{2\}=\(\{\\mathbf\{H\}\}\_\{1,1\}\+\{\\mathbf\{H\}\}\_\{1,2\}\)\+\{\\mathbf\{Q\}\}\_\{1\},\(81\)whosenn\-th row is given by

𝐐2​\[n,:\]⊤=\[1α1​\(xn−x¯\)α2​\(xn−x¯\)\+β2​\(yn−y¯\)α3​\(xn−x¯\)\+β3​\(yn−y¯\)\]\\displaystyle\{\\mathbf\{Q\}\}\_\{2\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\\\\ \\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\\\\ \\alpha\_\{2\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{2\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \\alpha\_\{3\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{3\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \\end\{bmatrix\}\(82\)since we omit the layer normalization\.

### B\.3The second layer

We show a construction of the firs head\. We have

𝐐2​\[:,1\]⊤​𝐐2​\[:,2\]\\displaystyle\{\\mathbf\{Q\}\}\_\{2\}\[:,1\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{2\}\[:,2\]=∑n=1N\+1α1​\(xn−x¯\)​\{α2​\(xn−x¯\)\+β2​\(yn−y¯\)\}\\displaystyle=\\sum\_\{n=1\}^\{N\+1\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\\left\\\{\\alpha\_\{2\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{2\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\right\\\}≃N​ξ2,1\\displaystyle\\simeq N\\xi\_\{2,1\}\(83\)for a largeNN, where

ξ2,1:=α1​α2​V\+α1​β2​C\.\\displaystyle\\xi\_\{2,1\}:=\\alpha\_\{1\}\\alpha\_\{2\}V\+\\alpha\_\{1\}\\beta\_\{2\}C\.\(84\)If we set

𝐖2,1,1\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,1\}=\[01N00000000000000\]\\displaystyle=\\begin\{bmatrix\}0&\\frac\{1\}\{N\}&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(85\)𝐖2,1,2\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,2\}=\[0000000000010000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&1\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(86\)𝐖2,1,3\\displaystyle\{\\mathbf\{W\}\}\_\{2,1,3\}=\[0000100000000000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(87\)then, by \([2\.4](https://arxiv.org/html/2607.15819#S2.Ex1)\), we have

𝐇2,1​\[n,:\]=\[000ξ2,1​α1​\(xn−x¯\)\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{2,1\}\[n,:\]=\\begin\{bmatrix\}0&0&0&\\xi\_\{2,1\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\\\\ \\end\{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,\\displaystyle\{\\mathbf\{W\}\}\_\{2,2,1\}=\{\\mathbf\{W\}\}\_\{2,2,2\}=\{\\mathbf\{W\}\}\_\{2,2,3\}=\{\\mathbf\{O\}\}\_\{4,4\},\(89\)then we formally have

𝐇2,2=𝐎N\+1,4\.\\displaystyle\{\\mathbf\{H\}\}\_\{2,2\}=\{\\mathbf\{O\}\}\_\{N\+1,4\}\.\(90\)
Therefore, the the output of the second layer is

𝐐3=\(𝐇2,1\+𝐇2,2\)\+𝐐2,\\displaystyle\{\\mathbf\{Q\}\}\_\{3\}=\(\{\\mathbf\{H\}\}\_\{2,1\}\+\{\\mathbf\{H\}\}\_\{2,2\}\)\+\{\\mathbf\{Q\}\}\_\{2\},\(91\)whosenn\-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¯\)\]\.\\displaystyle\{\\mathbf\{Q\}\}\_\{3\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\\\\ \\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\\\\ \\alpha\_\{2\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{2\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \(\\xi\_\{2,1\}\\alpha\_\{1\}\+\\alpha\_\{3\}\)\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{3\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \\end\{bmatrix\}\.\(92\)

### B\.4The third layer

We show a construction of the first head\. We have

𝐐3​\[n,2\]⊤​𝐐3​\[n,4\]\\displaystyle\{\\mathbf\{Q\}\}\_\{3\}\[n,2\]^\{\\top\}\{\\mathbf\{Q\}\}\_\{3\}\[n,4\]=∑n=1N\+1\{\(ξ2,1​α1\+α3\)​\(xn−x¯\)\+β3​\(yn−y¯\)\}​α1​\(xn−x¯\)\\displaystyle=\\sum\_\{n=1\}^\{N\+1\}\\left\\\{\(\\xi\_\{2,1\}\\alpha\_\{1\}\+\\alpha\_\{3\}\)\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{3\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\right\\\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)≃N​ξ3,1\\displaystyle\\simeq N\\xi\_\{3,1\}\(93\)for a largeNN, where

ξ3,1\\displaystyle\\xi\_\{3,1\}=α1​\(ξ2,1​α1\+α3\)​V\+α1​β3​C\.\\displaystyle=\\alpha\_\{1\}\(\\xi\_\{2,1\}\\alpha\_\{1\}\+\\alpha\_\{3\}\)V\+\\alpha\_\{1\}\\beta\_\{3\}C\.\(94\)
If we set

𝐖3,1,1\\displaystyle\{\\mathbf\{W\}\}\_\{3,1,1\}=\[01/N00000000000000\]\\displaystyle=\\begin\{bmatrix\}0&1/N&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(95\)𝐖3,1,2\\displaystyle\{\\mathbf\{W\}\}\_\{3,1,2\}=\[0000000000000010\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ 0&0&1&0\\\\ \\end\{bmatrix\}\(96\)𝐖3,1,3\\displaystyle\{\\mathbf\{W\}\}\_\{3,1,3\}=\[0000100000000000\]\\displaystyle=\\begin\{bmatrix\}0&0&0&0\\\\ 1&0&0&0\\\\ 0&0&0&0\\\\ 0&0&0&0\\\\ \\end\{bmatrix\}\(97\)then we have

𝐇3,1​\[n,:\]=\[00ξ3,1​α1​\(xn−x¯\)0\]\.\\displaystyle\{\\mathbf\{H\}\}\_\{3,1\}\[n,:\]=\\begin\{bmatrix\}0&0&\\xi\_\{3,1\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)&0\\\\ \\end\{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,\\displaystyle\{\\mathbf\{W\}\}\_\{3,2,1\}=\{\\mathbf\{W\}\}\_\{3,2,2\}=\{\\mathbf\{W\}\}\_\{3,2,3\}=\{\\mathbf\{O\}\}\_\{4,4\},\(99\)then we formally have

𝐇3,2=𝐎N\+1,4\.\\displaystyle\{\\mathbf\{H\}\}\_\{3,2\}=\{\\mathbf\{O\}\}\_\{N\+1,4\}\.\(100\)Thus, the the output of the third layer is

𝐐4=\(𝐇3,1\+𝐇3,2\)\+𝐐3,\\displaystyle\{\\mathbf\{Q\}\}\_\{4\}=\(\{\\mathbf\{H\}\}\_\{3,1\}\+\{\\mathbf\{H\}\}\_\{3,2\}\)\+\{\\mathbf\{Q\}\}\_\{3\},\(101\)whosenn\-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¯\)\],\\displaystyle\{\\mathbf\{Q\}\}\_\{4\}\[n,:\]^\{\\top\}=\\begin\{bmatrix\}1\\\\ \\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\\\\ \\xi\_\{3,1\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\alpha\_\{2\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{2\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \\xi\_\{2,1\}\\alpha\_\{1\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\alpha\_\{3\}\(x\_\{n\}\-\{\\overline\{x\}\}\)\+\\beta\_\{3\}\(y\_\{n\}\-\{\\overline\{y\}\}\)\\\\ \\end\{bmatrix\},\(102\)where we can write

ξ3,1\\displaystyle\\xi\_\{3,1\}=α1​\(ξ2,1​α1\+α3\)​V\+α1​β3​C\\displaystyle=\\alpha\_\{1\}\(\\xi\_\{2,1\}\\alpha\_\{1\}\+\\alpha\_\{3\}\)V\+\\alpha\_\{1\}\\beta\_\{3\}C=\(α12​ξ2,1\+α1​α3\)​V\+α1​β3​C\\displaystyle=\(\\alpha\_\{1\}^\{2\}\\xi\_\{2,1\}\+\\alpha\_\{1\}\\alpha\_\{3\}\)V\+\\alpha\_\{1\}\\beta\_\{3\}C=\(α13​α2​V\+α13​β2​C\+α1​α3\)​V\+α1​β3​C\\displaystyle=\(\\alpha\_\{1\}^\{3\}\\alpha\_\{2\}V\+\\alpha\_\{1\}^\{3\}\\beta\_\{2\}C\+\\alpha\_\{1\}\\alpha\_\{3\}\)V\+\\alpha\_\{1\}\\beta\_\{3\}C=α1​α3​V\+α1​β3​C\+α13​α2​V2\+α13​β2​C​V\\displaystyle=\\alpha\_\{1\}\\alpha\_\{3\}V\+\\alpha\_\{1\}\\beta\_\{3\}C\+\\alpha\_\{1\}^\{3\}\\alpha\_\{2\}V^\{2\}\+\\alpha\_\{1\}^\{3\}\\beta\_\{2\}CV\(103\)by \([84](https://arxiv.org/html/2607.15819#A2.E84)\) and \([94](https://arxiv.org/html/2607.15819#A2.E94)\)\.

𝐐4\{\\mathbf\{Q\}\}\_\{4\}is flattended into𝒒\{\\boldsymbol\{q\}\}which is a4​\(N\+1\)×14\(N\+1\)\\times 1vector\. Then, by setting

𝐖out​\[n\]=\{1n=4​N\+30otherwise,\\displaystyle\{\\mathbf\{W\}\}\_\{\\rm out\}\[n\]=\\begin\{cases\}1&n=4N\+3\\\\ 0&\{\\rm otherwise\},\\end\{cases\}\(104\)the prediction of the transformer is obtained by

y^​\(u\)\\displaystyle\{\\widehat\{y\}\}\(u\)=𝐐4​\[N\+1,3\]\\displaystyle=\{\\mathbf\{Q\}\}\_\{4\}\[N\+1,3\]=−β2​y¯\+ξ4,1​\(u−x¯\),\\displaystyle=\-\\beta\_\{2\}\{\\overline\{y\}\}\+\\xi\_\{4,1\}\(u\-\{\\overline\{x\}\}\),\(105\)where

ξ4,1\\displaystyle\\xi\_\{4,1\}=α2\+α12​α3​V\+α12​β3​C\+α14​α2​V2\+α14​β2​C​V\.\\displaystyle=\\alpha\_\{2\}\+\\alpha\_\{1\}^\{2\}\\alpha\_\{3\}V\+\\alpha\_\{1\}^\{2\}\\beta\_\{3\}C\+\\alpha\_\{1\}^\{4\}\\alpha\_\{2\}V^\{2\}\+\\alpha\_\{1\}^\{4\}\\beta\_\{2\}CV\.\(106\)

### B\.5Correspondence to gradient descent

By the correspondence between \([B\.4](https://arxiv.org/html/2607.15819#A2.Ex28)\) and \([67](https://arxiv.org/html/2607.15819#A2.E67)\), we have

\{−β2=1−α2=aα12​α3=−2​n​η​aα12​β3=2​n​ηα14​α2=n2​η2​aα14​β2=−n2​η2\\displaystyle\\begin\{cases\}\-\\beta\_\{2\}&=1\\\\ \-\\alpha\_\{2\}&=a\\\\ \\alpha\_\{1\}^\{2\}\\alpha\_\{3\}&=\-2n\\eta a\\\\ \\alpha\_\{1\}^\{2\}\\beta\_\{3\}&=2n\\eta\\\\ \\alpha\_\{1\}^\{4\}\\alpha\_\{2\}&=n^\{2\}\\eta^\{2\}a\\\\ \\alpha\_\{1\}^\{4\}\\beta\_\{2\}&=\-n^\{2\}\\eta^\{2\}\\\\ \\end\{cases\}\(107\)whereaais an initial value of the22\-step gradient decent in \([64](https://arxiv.org/html/2607.15819#A2.E64)\)\. By solving these equations, we have

\(α1,α2,β2,α3,β3\)=\(n​η,a,−1,−2​a,2\)\.\\displaystyle\(\\alpha\_\{1\},\\alpha\_\{2\},\\beta\_\{2\},\\alpha\_\{3\},\\beta\_\{3\}\)=\(\\sqrt\{n\\eta\},a,\-1,\-2a,2\)\.\(108\)Therefore, the output of the transformer is approximately consistent with the output of the linear function after the22\-step gradient descent\.