In-Context Deep Learning via Transformer Models
Authors: Weimin Wu
Maojiang Su
Jerry Yao-Chieh Hu
Zhao Song
Han Liu
Venue: ICML 2025 poster (2025)
Keywords: foundation model
transformer
in-context learning
Primary area: deep_learning->foundation_models
TL;DR: We investigate the transformer's capability for in-context learning (ICL) to simulate the training process of deep models.

Abstract: We investigate the transformer's capability for in-context learning (ICL) to simulate the training process of deep models. 
Our key contribution is providing a positive example of using a transformer to train a deep neural network by gradient descent in an implicit fashion via ICL. 
Specifically, we provide an explicit construction of a $(2N+4)L$-layer transformer capable of simulating $L$ gradient descent steps of an $N$-layer ReLU network through ICL.
We also give the theoretical guarantees for the approximation within any given error and the convergence of the ICL gradient descent.
Additionally, we extend our analysis to the more practical setting using Softmax-based transformers. 
We validate our findings on synthetic datasets for 3-layer, 4-layer, and 6-layer neural networks.
The results show that ICL performance matches that of direct training.

In-Context Deep Learning via Transformer Models
Weimin Wu 12
Maojiang Su 12 Jerry Yao-Chieh Hu 12 Zhao Song? Han Liu 124
Abstract
other hand, the high cost of pretraining these models often makes them prohibitive outside certain industrial labs (Jiang et al., 2024; Bi et al., 2024; Achiam et al., 2023). In this work, we aim to advance the "one-for-ne tone-formmnn"" modeling phi- losophy of foundation model paradigm (Bommasani et al., 2021) by considering the following research problem:
We investigate the transformer's capability for in- context learning (ICL) to simulate the training process of deep models. Our key contribution is providing a positive example of using a trans- former to train a deep neural network by gradient descent in an implicit fashion via ICL. Specif- ieally, we provide explicit construction of a (2N + L-layer transformer capable of simulat- ing L gradient descent steps of an N N-layer ReLU network through ICI. We also give the theoretical guarantees for the approximation within any given error and the convergence of the ICL gradient de- scent. Additionally, we extend our analysis to the more practical setting using Softmax-basel trans- formers. We validate our findings on synthetic datasets for 3-layer, 4-layer, and 6-yayer neural networks. The results show that ICL performance matches that of direct training.
Question 1. Is it possible train one deep model with the ICL of another foundation model?
The implication of Question 1 is profound: if true, one foundation model could lead to many others without per- taining. In this work, we provide an affirmative example for Question 1. Specifically, we show that transformers are ca- pable of simulating the training of a deep ReLU-based feed- forward neural network with provable guarantees through ICL. Our analysis assumes that we have well-pretrainee the transformer using the data generated by the deep net- work. We require the deep network to maintain consistent hyperparameters (e.g., model width and depth) during the pretraining and testing. However, during the testing, we vary the parameter distribution and input data distribution of the deep network to generate data for the transformer.
Introduction
We study transformers' ability to simulate the training pro- cess of deep models. This analysis is not only practical but also timely. On one hand, transformers and deep models (Brown et al., 2020; Radford et al., 2019) are so powerful, popular and form new machine learning paradigm foun- dation models. These large-scale machine learning models, trained on vast data, provide a general general-purpose purpose foundation for various tasks with minimal supervision (Team et al., 2023; Touvron et al., 2023; Zhang et al., 2022). On the
In ICL, the models learn solve new tasks during inference by using task-specific examples provided part of the input prompt, rather than through parameter updates (Shi et al., 2023b; Bubeck et al., 2023; Achiam et al., 2023; Bai et al., 2023; Min et al., 2022: Garg et al., 2022). Unlike standard supervised learning, ICL enables models to adapt to new tasks during inference using only the provided examples. In this work, the new task of our interest is algorithmic approximation via ICL (Bai et al., 2023; Zhang et al., 2024; Wang et al., 2024). Specifically, ICL capability to replace/simulate the standard supervised training algorithms for N N-layer networks. To be concrete, we formalize the learning problem of how transformers learn (i) given function and (ii) a machine learning algorithm (e.g.. gradient descent) via ICL, following (Bai et al., 2023).
aim to use transformer's ss
"Equal contribution 'Center for Foundation Models and Generative AI, Northwestern University, USA Department of Computer Science, Northwestemn University, USA University of California Berkeley. USA 'Department of Statistics and Data Science, Northwestern University, USA. Correspondence to: Weimin Wu nortteunnortwweeeee www.hhewttteeeee.eeee..e edu>>, Mao- jiang Su Jerry Yao-Cheih Hu . www.hwester...n.omm Zhao Song <magic.. 1in..l.demxx coms, Han Liu Chanliuttorttiwesteenn.cm
wtw.aagge200.....oeeeetetrrr. edu>>,
(i) ICL for Function f. Let f Rd + E be the func- tion of our interest. Suppose we have a dataset D,, = (x),,) se))) where (a) S Rdd and {utheen C R are the input and output of f. respectively. Let Øn-1 be the test input. The goal of ICL is to use a transformer, denoted by T. to predict Vn+1 based on the test input and the in- context dataset autoregresively: Øn-1 ~ T(D T(D Ta+1). The
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025, Copyright 2025 by the author(s).
Full version and future updates are on arxiv.

---

In-Context Deep Learning via Transformer Models
goal is for the prediction Un+1 be close to yn+1 = [(x).
4L-layer Softmax transformer to approximate L steps of gradient descent, and guarantee the approximation and the convergence (Theorem 2).
(ii) ICL for Gradient Descent of a Parametrized Model -w,,)) Bai et al. (2023) generalize (i) to include algo- rithmic approximations of Gradient Descent (GD) training algorithms and explore how transformers simulate gradi- ent descent during inference without parameter updates. They term the simulated GD algorithm "In-Context Gra- dient Descent (ICGD)." In essence, ICGD enables trans- formers approximate gradient descent on a loss function L. (w)) for a parameterized model (w: based on a dataset DR. Traditional gradient descent updates w iteratively as W = W/ ,, (w,). In contrast, ICGD uses a trans- former T to simulate these updates within a forward pass. Given example data D,, and test input F,+1, the transformer performs gradient steps in an implicit fashion by inferring parameter updates through its internal representations, using input context without explicit weight changes. Please see Section 2 for explicit formulation.
Experimental Validation. We validate our theory with ReLU- and for the N-layer networks (Theorem 1, Theorem 2, and Theorem 4). We assess the ICL capabilities of transform- ers by training 3-, 4-, and 6-layer networks in Section 5. The numerical results show that the performance of ICL matches that of training -V-layer networks. However, a minor limitation is that the trained transformers do not always achieve the theoretical construction.
specifically, ICGD
Organization. We present our main results in Theo- rem 1. Section 2 covers the preliminaries. Section presents the problem setup and the ICL approximation of GD steps for an N-layer feed-forward network with both HeLU--ransformet and Trannnxx-TTanffomer. Section 5 presents the experimental results, with additional details in Appendix F. The appendix includes related work (Ap- pendix A. detailed proofs for Section 3 (Appendix C), and an application to train diffusion models via ICL (Ap- pendix G).
In this work, we investigate the case where f(w,)) :) is a deep feed-forward neural network. We defer the detailed problem setting to Section 2. comparison to standard ICGD (Bai al., 2023), ICGD for deep feed-forward networks is not trivial. This is due to two technical challenges:
(C1) Analytical feasibility of gradient computation for these thick networks.
Notations. We use lower case letters denote vectors and upper case letters to denote matrices. The index set {1, I} is denoted by [4], where I E N For any matrices A E TRmx", let ép norm of A be induced by vector f--norm, defined as || == supfl|Azi] : E R" with ||>||p = 1}. We use A[i, to denote the element i-th row and j-th column of matrix A. For any matrices A Æ=xx and B E mmxn let O denotes the Hadamard product: (A O B) = A[i,5] B[ū,j]. For any matrices A E Rmxx and B E Rexg, let @ denote the Kronecker product:
(C2) Explicit construction capable of approximating ICGD for such layers and their gradients.
To this end, we present the first explicit expression for gradient computation of N-layer feed-forward network (Lemma 1). Importantly, its term-by-term tractability pro- vides key insights for the detailed construction of specific transformer to train this network via ICGD (Theorem 1).
1]B
A|L,n|B
Contributions. Our contributions are threefold:
B@
:
:
Approximation by ReLU-Transformen For simplic- ity, we begin with the ReLU-based transformer. For a broad class of smooth empirical risks, construct (2N + 4)L-layyer transformer to approximate L steps of in-context gradient descent on the N-layer feed-forward networks with the same input and output dimensions (The- orem 1). We then extend this to accommodate varying dimensions (Theorem 4). We also provide the theoretical guarantees for the approximation within any given error (Corollary 1.1) and the convergence of the ICL gradient descent (Lemma 14).
Am,11B
A[m, n)
2 Preliminaries: ICL and ICGD
We present the ideas we built upon: In-Context Gradient Descent (ICGD).
(i) ICL for Function f. Let f. : Rd + R be the func- tion of our interest. Suppose we have a dataset D, == ((r,, y1) vi))se where (ax {# S Rd and (yilve(m C R are the input and output of f. respectively. Let xn-1 be the test input. The goal of ICL is to use transformer, denoted by T, to predict yx++1 based on the test input and the in- context dataset autoregresively: Ta+1 ~ T(D,, 11). For convenience in our analysis, we adopt the ICL notation from (Bai et al., 2023). Specifically, we shorthand (D,, F1+1) into an input sequence (i.c., prompt) of length n + 1 and represent as a compact matrix H. E Dx(+11 =
Approximation by Transssore-Tassrmeerr We extend our aialysis to the max-ttaaxroooo to better reflect realistic applications. The key technique is to ensure a qualified approximation error at each point to achieve uni- versal approximation capabilities of the fettmax-baseee Transformer (Lemma 16). We give a construction of
2

---

In-Context Deep Learning via Transformer Models
11 M hn
1
in the form:
[::;; y4: w0 ; 0; 1; tzl. where, with w.!0) = 0,
D2
In
In+l
Th)) Projw (wll-)) VL (w + e e
H =
y/
1/2
yn
0
E (DDx(n-1),
91
42
On
qu+1
(2.3)
OD (d+33
is updated recursively, and ||c(|=1)|| ||2 I represents the approximation error at step (- 1.
9: =
E Dd++1.)
(2.1)
tz
Problem 1 aims to find a transformers T to perform L steps gradient descent on loss (w) in an implicit fashion (i.e., explicit parameter update). More precisely, Bai et al. (2023) configure 'I with T. identical blocks, each approx- imating one gradient descent step using T' layers. In this work, we investigate the case where f(w,) .) is an "N-layer neural network."
We use qi to fill in the remain D D- (d+1) entries in addition to 2y E I and y; E IR. The last entry 4 := 1(i < n + 1) of q is the position indicator to distinguish the n in-context examples and the test data. The problem of "ICL for f is to show the existence of transformer' T that, when given H, outputs T(H) E Dx(++11) of the same shape, and the "(d + 1, + 1) entry of T(H)" provides the prediction I-+1- The goal is for the prediction 7n to be close to Un+11 = f(x) measured by some proper loss.
Transformer.
We defer the standard definition of trans- former to Appendix B.I.
(ii) ICL for Gradient Descent of a Parametrized Model .w., We aim to use ICL to simulate the standard super- vised training procedure for N --layer neural networks. To achieve this, we introduce the concept of In-Context Gradi- ent Descent (ICGD) for a parameterized model. Consider machine learning model [w,,,) RD- x Rd - Rd. parametrized by w E DR, Given a dataset Dn, == ((r,, ve))e()) jid P. a typical learning task is to find pa- rameters w" such that (("",.) becomes closest to the true data distribution P. Then, for any test input #n+1, we pre- dict: Da+1 = fw",, ;n+1). To find w", Bai et al. (2023) configure transformer to implement gradient descent on (wc,-) through ICI., simulating optimization algorithms during inference without explicit parameter updates. We formalize this In-Context Gradient Descent (ICGD) prob- lem: using a pretrained model to simulate gradient descent on (wr w.r.t. the provided context En+1)
3 In-Context Gradient Descent on N-Layer Neural Networks
We now show that transformers is capable of implement- ing gradient descent on NN-layer neural networks through ICL. In Section 3.1, we define the N-layer ReLU neural network and state its ICGD problem. In Section 3.2, we derive explicit gradient descent expression for N-layyer NN. In Section 3.3, we construct ReLU-Tranfformrr executing gradient descent N-layer NN via ICL. In Section 4, we show the existence of Taannnnno---aaaoorerer capable of per- forming in-context gradient descent on N --layer NN.
3.1 Problem Setup: ICGD for N-Layer Neural Networks
To begin, we introduce the construction of our N-Layer Neural Network which aims to implement gradient de- scent on its empirical loss function.
Definition 1 (N-Layer Neural Network), An N N-Layer Neu- ral Network comprises N hidden layers and out- put layer, all constructed similarly. Let R + R be the activation function. For the hidden layers: for any E [n + 1],j E [N - 1]. and k E [K]. the output for the first j layers w.r.t. input x; E Rd, denoted by pred, (xx; E B is defined as recursive form:
Problem (In-Context Gradient Descent (ICGD) on Model f(w,-) (Bai et al., 2023)). Let > and L > 1. Consider a machine learning model f(w,z) : RR Rd - Rd parameterized by w E Du, Given a dataset D,, = iid ((r,, ((r,):)) 3e) ~ with (1), E Rd x Rd, define the empirical risk function:
pred, (z4; [)[ == r(vT,=;),
(w) = 2n Σ (..W).))..).
(2.2)
prod, (z;; == r("" pred, (Z:j - 1)); where U1, E Rd and Uys E R for j E (2...., N 1} are the k-th parameter vectors in the first layer and the j-th layer, respectively. For the output layer (N-th layer), the output for the first N layers (i.e the entire neural network) w.r.t. input x; E Rdd, denoted by pred, (I.; w, N) E IR,, is defined for any ho E [d] as follows:
where e: Rdd x 2d IR is loss function. Let W C p be a closed domain, and Projw denote the projection onto W. The problem of "ICGD on model f(w,,)" is to find transformer T with T. blocks, each approximating one step of gradient descent using T' layers. For any input (0) E (x+++) in the form of (2.1), the transformer ((H()) approximates L steps of gradient descent. Specifically, for I E [L] and i E [n- + 11, the output at layer T7 is: h
r(oT, pred, N - 1)), (3.1) where UNA E R are the k-th parameter vectors in the
pred, w, NV)
3

---

In-Context Deep Learning via Transformer Models
N-th layer and w E 22d (NN denotes the vector containing all parameters in the neural network,
first j layers. By (3.2), we have:
0,
j = 0
:= [eT, . "Ix .
NN. (3.2)
dk,
1 ==1
Dj =
(j - 1)K2 + dK,
2 25 j S N -
Remark 1 (Prediction Function for j-th layer on i-th Data: Ps (i)) For simplicity, we abbreviate the output from the first j-th layer of the N-layer neural networks NN with input
(NN-2)K2 + +dd,,
j = N.
The [v]
parameter
vector
w
3; as p)()),
T "IK "IR! N- -N 1x UN
T
No
3i E R2d pc(i)= =- prod, (24: IRK (r; w, .) Rd.
for j =
follows (3.2). Define di (N), dp.(N)
(p.(N)
for [N 1] for j N.
RN,,, For any j E [N], let A.(i) denote the derivative of (PP. 34) with respect to the parameters in the j-th layer: A,(5) <x(
(3.3)
DyJ. where d. tj selects elements from the a-th to b-th position in Φ;.
Additionally, we define
-..)) (p:( E (NN-1)
Pt :=:=
For activation function r(t), let ('()) be its derivative. De- fine r R as:
We formalize the problem of using transformer simulate gradient descent algorithms for training the N-layer NN de- fined in Definition 1, by optimizing loss (2.2). Specifically, we consider the ICGD (Problem 1) with the parameterized model f(w, -) = (red,(=:,, N).
r(()) == r' r'(u
Define ry == R((j))
1)] and R2( as:
(1 p.()),, N - 2 ddiag(r' (rr"("((((),,,,11) ((()),. (" (())) j= N
Problem 2 (ICGD on N-Layer Neural Networks). Let the
N-layer neural networks, activation function r, and predic- tion function (j) for all layers follow Definition I and Re- mark 1. Assume we under the identical setting as Problem 1,
where Ri( E R kx for j E to..... N. 2 and
R((j) E Rdxd for j = N 1
considering model f(w, -) := pred,(;; w, N) and specify- ing W is a closed domain such that for any j E [N -1 and 1: E [
For any j E [N let V3 denote the parameters in the j-th
layer as:
W c [v,u] E Px 11/52112 1l2 < B The problem "ICGD on N --layer neural networks" to find a TL layers transformer T, capable of implementing L steps gradient descent as in Problem 1.
(3.4)
U1p Ul
E RRxx
j j=1
V3 Ujio...,
Ujk E RRKK j 2.....
N -1
E ERuxK j j- N.
UN UNa
Remark 2 (Why Bounded Domain W?). For using sum of ReLU to approximate functions like r, which is illustrated in the consequent section, we need to avoid gradient exploding. Therefore, we require W to be a bounded domain, and utilize Projw project into bounded domain W.
Definition 2 splits the gradient of (w) into N parts. This makes V Vw.Cn()) more interpretable and tractable, since all parts follows a recursion formula according to chain rule. With above notations, we calculate the gradient descent step (2.3) of N -llayer neural network as follows:
3.2 Explicit Gradient Descent of N-Layer Neural Networks
Lemma 1 (Decomposition of One Gradient Descent Step). Fix any BG, O. Suppose loss function (w) on data points {(x;, vi) follows (2.2). Suppose closed domain W and projection function Projw follows (3.4). Let A,)),,(( R2),,, V, be as defined in Definition 2. Then the explicit form of gradient VL,, (w) becomes
Intuitively, Problem 2 asks whether there exists a trans- former capable of simulating the gradient descent algorithm on the loss function of an --layer neural network. We answer Problem 2 by providing an explicit construction for such a transformer T in Theorem 1. To facilitate our proof, we first introduce the necessary notations for explicit expression of the gradient Cn (w).
A((11)
VL, (w)
Σ 2n i=1
:
(3.5)
Definition 2 (Abbreviations). Fix i E [n + 1). and consider an N-layer neural network with activation function and prediction function (i) as defined in Definition 1.
A.(NN)
where A;(j) denote the derivative of ((p((N), 9:) with re- spect the parameters in the j-th layer.
. Let D, E R denote the total number of parameters in the
A₂()) =
4

---

In-Context Deep Learning via Transformer Models
(N - 11VN ... ..UR 1) |KXK 2 pili - p:lp-1113 - -( ( )T, j N
with Σh= leal I C, max Il | lalll S 1, E R and Ch E IR, such that
(P)(N)
sup IØ(2) ['m,(()) =E(-R.
Capprox-
(R [Loxd @ Pi(N Pi(N-1) I)
(llm)(M)1)) j N.
Overview of Our Proof Strategy. Lemma and Defini- tion 4 motivate the following strategy: term-by-term ap- proximation for our gradient descent step (3.5), Please see Figure 1 for a high-level visualization.
Proof Sketch. Using the chain rule and product rule, we decompose the gradient as follows: VaLn (w) =
2n
1
2i=1
P (M)pp.()) Thus, we only need
Step 1. Given (x),w), we use N attention layers to approximate the output of the first j layers with input Tis Pc(j) = pred, (x4;,) E TR (Definition 1) for any j E [N]. Then we use attention layer to approximate chain- rule intermediate terms r - [)[]]] == r'vv] pe(j -)) (Definition 2) for any i E [m]. j E [N] and l E [K]: Lemma 2 and Lemma 3.
to compute (N) By Definition 1 and the chain rule, we
prove that 2p:(N2 satisfies the recursive formulation (C.4).
8w
Combining these, we derive the explicit form of gradient
V V. (w). and the gradient step follows directly. Please see Appendix C.1 for detailed proof.
M
is hard to calculate the elements in (;) in a straightfor- ward mannar, we calculate each parts of it successively. We define the intermediate terms st (j) and as follows
Step 2. Given (r(, Pi, w), we use an MLP layer ap- proximate (NV), y1) (Definition 3), for i E [r=]. and use N element-wii wiiee multiplication layers to approximate s;(j) (Definition 3), for any j E [N]: Lemma 4 and Lemma 5. Moreover, Lemma shows the closeness re- sult for approximating s; (j), which leads to the final error accumulation in Theorem 1.
Definition 3 (Definition of intermediate terms). Let A,()), r,()) R)), Vj be as defined in Definition 2. For any y E Rd, we define vector function u(t, y) == ( @8(t.w) )T= 2t
Rd x Rd > Ted. Moreover, for any j E [M]. E [n + 1).
we define 8, (j) as
Step 3. Given (pi, r'ys Iis((j), w)., we use an attention
R (jj - 1)V RN(N- - 2 RN(N- - 1)
layer to approximate w - V (w)). Then we use an MLP layer to approximate Projw (wr) And implement- ing L steps gradient descent by (2.V + +4) neural network NN, constructed based on Step 1 and 2. Fi- nally, we arrive our main result: Theorem 1. Furthermore, Lemma 14 shows closeness results to the true gradient descent path.
st())=
u(p.(N)), vi) E ERK, . 1i N
RN(N - 10NN-11-w(pp (()):((pp))))) E Rd,
j= N.
Let @ denotes Hadamard product. For any j E [N-1),
[N + 1]. Definition 3 leads to
sx()) rl(jj- - 1) O (V,T1 +,(i+1)), (3.6)
Moreover, by Definition 3, holds
Step 1. We start with approximation for pt(j).
T T @ p. 11] s((i,, j*M N, Taxa & PN - 1) (N), j= N.
Lemma 2 (Approximate ps(i)). Let upper bounds
A.0)
(3.7)
By, B, > such that for any & E [K], j E [N] and i E [n].
| I B,, and ||æ:||2 S Bz. For any j E [N], i E [m].
define
3.3 Transformers Approximate Gradient Descent of N-Layer Neural Networks In-Context
B
max Ir((), B" == Ba, and B, max BJ.
| B. B
For using neural networks to approximate (2.2), which con- tains smooth functions changeable, we need to approximate these smooth functions by simple combination of activa- tion functions. Our key approximation theory is the sum of ReLUs for any smooth function (Bai et al., 2023).
Let function r(t) be (c,, R1. M, C₁) approximable for R1 max B, 1}, Mj S Ō(C9c=2), where C1 de- pends only on Rú and the C*-smoothness of Then, for any 6, > O, there exist N attention layers Attno, Attnox such that for any input h4 E RDD takes from (2.1), they map h4 = [1r; M45 w;; (1);;; P.li 1): 0; 1;4,]
Definition 4 (Approximability by Sum of ReLUs, Definition 12 of (Bai et al., 2023)). Let E R. We say function g: R + IR is (Capprox; R, H. -pprroxxiieeee by sum of ReLUs if there exist "(H, Cu- of ReLUs" function fH,o(z) defined as
Attne, hJ [.:; MN w. (.); ... :.) 1; 1; where (j( is approximation for Pi(j) (Definition 1). In the expressions of h, and h. the dimension of 0 differs. Specifically, the in h2 is larger than in he The dimensional difference between these 0 vectors equals the dimension of Di(j). Suppose function is Smoohh in bounded domain
H fn,c(z) = -
(a 1)).
hml
5

---

In-Context Deep Learning via Transformer Models
Inputs
TF
TF3
TF3
EWML
TF
Lemma 2
Lemma 3
Lemma
Lemma
Theorem
{( (s))eel)).u
u(pi(NV), me)
st(i)
(rojw - v/ Ln (w))
Figure 1: One Step In-Context Gradient Descent (ICGD) with (2.V + -ayer Transformer. This illustration presents the backpropagation process within an ICGD in a transformer model with 2NV + 4 layers. simulates single gradient descent step for an N-layer neural network, trained with loss L., and datasets {(rs, M)]]ejnj The term pı(j) denotes the output after the j-th layer for input æ;. The terms r(),), (NV, ys), and sx(j) are intermediate gradient terms of gradient VCn (w) from the chain rule. The expression Proiw( v7 Cn (w) shows one gradient descent step. Here, is the learning rate, and W denotes the bounded domain for the -llayer NN parameters wo.
W. then for any E [n + 11, j E [N], P.( such that pi(j) = p.( + ((,,
for any 6, O, there exist an attention layer Att such that for any input h; E I2 takes from (3.13), it maps
(3.8)
h4 ( ysi we; Pr; 0; 1:4]
E C K L,) I 1
||e(i,))l22 - CE N-1 L.
Attmo h [a,; yes Pi; F/: 0; 1;
N.
(3.9)
where P,(i) is approximation for r((i) (Definition 2) and r [+)0;; ..: P.(N -1] E R +d Similar to Lemma 2, in the expressions h; and hú the dimension of 0 differs. In addition, let E, be defined in (3.12), for any
Additionally, for any j E [N], the norm of parameters Be, defined as (B.1) such that Be, I 1 + KCG-
Proof. Please see Appendix C.2 for a detailed proof.
i [n 1). [N], k: E [K], F((i) such that
Pl - 1)[K - - 1)[k] + jii,j,k), le(t,j, Sc. + L, E.
Notice that the form of error accumulation in Lemma 2 is complicated. For the ease of later presentations, define the upper bound of coefficient in (3.8)
(3.14)
where e, denotes the error generated in approximating by sum of ReLUs r follows (C.5), Additionally, the norm of parameters B defined as (B.1) such that B8 S + K(N 1)C2.
[cl2,
E, = max JE[N]
(3.10)
Er
1-1
N L
max JE[N]
Proof Sketch. By Lemma 2, we obtain p:(i), the approx- imation for p:(j) (3.3). Using pr(j), we construct an At- tention layer to approximate r(()). We then establish up- per bounds for the errors [F( - [k]]] by applying inequality and Lemma 2. Finally we present the norms (B.1) of the Transformers constructed. Please see Appendix C.3 for a detailed proof.
1=0
1=0
(3.11)
such that (3.8) becomes
pi(i) = p+()) + e(i.j), [|e(t,))ll S E,er.
(3.12)
Moreover, we abbreviate P. [p.(1); ... :: E
such that the output of Attno,
Attno
is
Let Attng (j E [N]) be as defined in Lemma 2, then Lemma implies that for the input takes from Problem 2, the output of Attng, o Attne is
h = [ Mi we; P;; 0; 1;4].
(3.13)
Then, the next lemma approximates rf( base on (.()) obtained in Lemma 2.
h2 [mo;; Pe: F/: 0; 1;44)
(3.15)
Lemma (Approximate r(())) Let upper bounds Be, B2 > 0 such that for any [K],j e [N] and e [n]. ||05. lla < B., and l2:ll22 S B,, For any j [N], E [n].
Step 2. Now, we construct an approximation for (NV,, y1) = (N) (N)
define
Lemma 4 (Approximate u(p; 91)). Let upper bounds By, Bz, > 0 such that for any l: E [K].] j [N] and E [n]. |/ B,, and || 2|2 < B,, For any k E [d], sup- pose function u(t. y) be (e,, R3, M C approxiabblee for Rs max B,, By. 1), M3 where C.. depends only 12 and the 8 -moothhnsss of
B = max. (1()),, BO, = Bz, and Br == max B
Suppose function r() is (Err, R2, M2, C2) approxxaale for R2 max max{B,., Br, 1). M2 S Ö(C3e-²) where C2 depends only on R₂ and the -moothhnss of r'. Then,
O((C))&g
6

---

In-Context Deep Learning via Transformer Models
(t,, y) [K]. Then, there exists MLP layer MLP 0x42 such that for any input sequences h; E I takes from (3.15), it maps
L-layer neural networks
EWML == EWML8, ..... EWML8
Remark 3 (Necessary for Element-Wise Multiplication
h4
| 1::::: Ps: ; 0; 1:
Layer). As we shall show in subsequent sections, ELML is capable of implementing multiplication in h,- Specifically, it allows us to multiply some elements in h2 Lemma 5. By Definition 7. it is impossible for transformer layers to achieve our goal without any other assumptions.
MIPeNx2, h; i ye w: Pe: F 9:: 0; 1:t,),
where 9: E Rd is approximation for (NV) v.). For any & [d], assume u(p, (NV), 34) is L-- Lipschitz continu- ous. Then the approximation 9i such that,
9: [k] = u(p.(N), [.) + e(i, k).
(3.16)
Similar to (B.1), we define the norm for L-layer transformer FWMT & as:
with lc(2,k) k)] S 6, + c, Additionally, the parameters 42 such that
By == max (110 Il, I| lb. IV (3.19)
12 < max R3 + 1, Cz).
Then, given the approximations for pı(i) and r(()), we use N element- wemeett-wise multiplication layer (Definition 6) to approximate $;()), the chain-rule intermediate terms defined as Definition 3.
Proof Sketch. By Definition 1, we provide term-by-term approximations for p:( as forward propagation. Specifi- cally, we construct Attention layers to implement forward propagation algorithm. Then we establish upper bounds for the errors |p:(j) p+(i)llz 2 inductively. Finally, we present the norms (B.1) 1) of the Transformers constructed. Please see Appendix C.4 for detailed proof.
Lemma 5 (Approximate 8; (i)). Recall that sc()) = r/ 1) O (V s, + 1)) follows Definition 3. Let the initial input take from (3.17). Then, there exist N element- wise multiplication layers: EWML. WWW.Lgyis....... 13:-... that for input sequences, j [N),
D Let Attno, (j E [N + +1) be as defined in Lemma and Lemma 3, then for any input sequences h, TR. takes from (2.1), the output of Attne, o o 4t &ML is (3.17)
EWML 12 such
h4 = [m u:: w; Pe; F g.: (N): : 8,( + 0): they
0: 1;t,] EW 2 (h)
map
h = [ yo: w; p. Fl: g:: 0; 1;
[T,; yes w; Pr: F/s 96, (N); ........... 5. 1;4].
where s; s;()) is defined as recursive form: for
Before introducing our next approximation lemma, we define an element-wise multiplication layer, since atten- tion mechanisms and MLPs are unable to compute self- products (e.g.. output zy from input [x: y). To enable self- multiplication, we introduce a function 7. This function, for any square matrix, preserves the diagonal elements and sets all others zero.
the approximation any i E [n + 11, j E [N].
5.()) = [r/ 1) @ (VJ+1 3. + 1)), P2(N - 1) o gir
1 [N - 1]
j j= N.
(3.20)
Additionally, for any j E [N], BoN defined in (B.1) satisfics Bax S 1.
Definition 5 (Operator Function 7). For any square matrix A E R define
(A) = diag(([[ A|m, ni) E g=x=
Proof. Please see Appendix C.5 for detailed proof.
By Definition 5, we introduce the following element- wise multiplication layer, capable of performing self- multiplication operations such the Hadamard product.
Let Attno, (j E IN 1]), ML,M be as defined in Lemma 2, Lemma 3 and Lemma 4 respectively. Define s; := [NN(N; :.; E R then for any input sequences h; E RR' takes from Problem 2, the output of neural network
Definition 6 (Element-wise Multiplication Layer). Iet be defined as Definition 5. An element-wise multiplication layer with m heads is denoted as A-) with parameters On any input sequence H E
Attng, o... Attno oMLP 0N 0 EWML v+3
0 {Qm; Km, )m m D
o.. Ww..L..ML.... 02N+2:
(3.21)
is
(WMLO(()) = H + Σ (V 7((2ml)) ( H)).
h; [mn yn w; P. 7/s 81; O; 4:4.
(3.22)
For the sake of simplicity, we consider ReLU Attention layer and MLP layer are both a special kind of transformer. In this way, by Definition 9, (3.21) becomes
(3.18)
where Qms Km, Vn E DxD and 7(-) is operator func- tion follows Definition 5. In vector form, for for each token by E RR in H, it outputs [E))L = h4 + IM 7(1Qm Km hi)) - Vnnha In addition, we define
TE o EWML
Next
calculate the error accumulation [B.()){K] 8, (R]] based on Lemma 3 and Lemma 4.
7

---

In-Context Deep Learning via Transformer Models
Lemma 6 (Error for 5,(i), Suppose the upper bounds
and f. And this neural network such that for any input sequences (0),, take from (2.1), (H)) implements L steps in-context gradient descent on risk Equ (2.2): For
Bv, B2 > 0 such that for any h E [ E [N] and E [m],
|| ll < B., and læelle I Bg. Let r( E R such
(1,++,,P(i))) follows Definition 2. Let that r+ == r'(y
every / E [L]. | the (2N + 4)l-th layer outputs h: ((2N +4)() [x:: ys: w0; 0; 1; ty] for every i E [n + 11, and approxima- tion gradients (7()) such that
sc(i) = R ( 11 ....... R2 22VV ( 1)u fol-
lows Definition 3. Let (.). gr, 3.) be the approximations for r)))), w(p::.(pp(.N.)).. us). si( follows Lemma 3, Lemma 4 and Lemma 5 respectively. Let B, be the upper bound of r[(])[]]] and [A] as defined in Lemma 3. Let B₁ be the upper bound of gi and (NV: (N). ve) as defined in Lemma 4.
w(t) Projwww
V/ V (w"-1) + )
where T(() 0, and ||e || S 7/e is an error term.
Proof Sketch. Let the first 2N + 2 layers of NN₂ are Trans- formers and EWMLs constructed in Lemma 2, Lemma 3, Lemma 4, and Lemma 5. Explicitly, we design the last two layers to implement the gradient descent step (Lemma D). We then establish the upper bounds for er- ror || (w)) - (w)(()) [2, where (w), derived from the outputs of NN approximates V.nC. (w). Next, for any 6 > O, we select appropriate parameters 6t, 6, and || V V. (w((-1) v. (all- 1) S
Then for any E [n + 1 E [NT] k E [K],
[.([)[]] - Bs:
[+) - [A]] S e + Ey ep . E where B, is the upper bound of s; (kk]] and E Ey', E are the coefficients of er, <,, et in the upper bounds of [s:([][M}] s, (j)[ respectively.
Proof. Please see Appendix C.6 for a detailed proof.
e,' to ensure that
holds for any L [L].
Lemma 6 offers the explicit form of the error [+: [b]]] which is crucial for calculating the error V (w) - V (w)ll in Theorem I.
Please see Appendix C.7 for a detailed proof.
We summarize and visualize the backpropagation process within an ICGD in a transformer model with 2.N + 4 layers in Figure 1. As direct result, the neural networks NN, constructed earlier is able to approximate the true gradient descent trajectory {w(D)120 defined by w&D = O and wit1 = w 11 nn Cn (wan)) ) for any | Z 0. Consequently, Theorem motivates us to investigate the error accumulation under setting
Step 3. Combining the above, we prove the existence of a neural network, that implements in-context GD steps on our N-layer neural network. And finally we arrive our main result: neural network T for Problem 2.
Theorem 1 (In- Unn-Context Gradient Descent on N-layer NNs). Fix any Bg, 7/, < 0, T. 2 1. For any input sequences takes from (2.1), their exist upper bounds B,g By such that for any l22 I Bz. Assume functions r(t),
( = Projw (w)l-1)) 1/ L. ( +el1)))
i E [m]. |lua112 S By, ||æa||2
(1) and u(1,, [k] are L,, Ly, Lx-Lipschitz continuous.
where ( = 0, and ||e |e(-_))||2 S Tee represents error terms. Moreover, Corollary 1.1 shows NN, constructed in The- orem implements L steps ICGD with exponential error accumulation to the true GD paths.
Suppose W is a closed domain such that for any j E [N-
and k E [K].
W c [w Px PN 11 S B.}.
and Projw project w into bounded domain W. Assume Projw = ML.P, for some MLP layer with hidden dimen- sion D., parameters |0] S C.. If functions r(t), (1) and u(t, v) [k] are -mmsmoothnesss then for any e >> O, there ists a transformer model NN, with (2N T. hidden layers consists of L neural network blocks TFN+ EWML TF?.
Corollary 1.1 (Error for implementing ICGD on N-layer neural network). Fix L 2 1, under the same setting Theo- rem 1, (2NN + L-yLer neural networks NN approximates the true gradient descent trajectory (wap the error accumulation || w ll S L (1 + nby)/e where L, denotes the Lipschitz constant of C, (w) within W.
R Rxx with
NN. TFN+2 22 EWML, N TF? such that the heads number M', parameter dimensions D and the parameter norms BB suffice
Proof. Please see Appendix C.8 for a detailed proof.
4 In-Context Deep Learning with Softmax Transformers
max M < Ō()-)), [)((2N+4)L)
D' < O(NK2) + Dw?
In this section, we extend our analysis from ReLU- transformers to more practical dermmax-transforman for ICGD of N-layer neural network (Appendix E). Specifi- cally, we establish the existence of Softmax- transsoorrr capable of performing ICGD for N N-layer neural networks in Theorem and give more details in Appendix E.
max [2
max BB < 0()) + Cw +1, 1e((2N+1)L]
where Ō(-) hides the constants that depend d, K, N. the radius parameters Bz, By, B, and the smoothness of
8

---

In-Context Deep Learning via Transformer Models
1.0
1.0
0.8
0.8
R-Squared
R-Squared
0.6
0.6
N(-2.1)
N(-2.1)
0.4
0.9N(-2, 0.1N(2, I)
0.4
0.9N( -2. 1) 0.
0.2
1. + 0.3N(2,0)
0.2
0.7N(-2.1)) 0.33222
1.5 +..NN(2, I)
0.5N 2. 2) 1.
0.0
6-Layer NN
0.0
6-Layer
-0.2
0
50 In-context Examples
25
75
50 In-context Examples
25
75
(a)
(b) Softmax-Transfommn Transformer
Figure 2: Performance of ICL in Rernnsrafefrrm and Softmax-Transformen ICL learns 6-layer NN and achieves R-squared values comparable to those from training with prompt samples.
5 Numerical Studies
Theorem 2 (In-Context Gradient Descent of Softmaxx-rannformeet Transformer). Fix any Bw, 7), 0, 2 1. For any input sequences takes from (2.1), their exist upper bounds Br,, By such that for any E [m]. |lyx/|mmx < By, |æz/|lmax S Bg. Suppose W is a closed domain such that || I Bm, and Projw project w into bounded domain W. Assume Projw = MLP, for some MLP layer. Define loss function with T-LLipschitz gradient. Let C., (w) = 1 -. f(w, (I,,I;18
In this section, we conduct experiments to verify the capabil- ity of ICL to learn feed-forward neural networks, and give details in Appendix F. We conduct the experiments based on 3-, 4- and 6-layer NN using both ReLU- and Soft Transformer. The main objective is to validate the perfor- mance of ICL matches that of training N-layer networks, i.e., the results in Theorem 1, Theorem 2, and Theorem 4. However, minor limitation that the trained transformers do not always achieve the theoretical construction.
Iww Ft, yi)
I,, 34) denote the empirical loss function, then there exists a Softmax-transte transforerr NNg. such that for any input sequences H1)), take from (2.1), NNo (H)))) implements L steps in-context gradient descent on C. (w)): For every E [L]. the 4/-th layer outputs n [ y4; m(:); 0; I; for every
Specifically, we sample the input of feed-forward net- work I E Rd from the Gaussian mixture distribution: WiN (-2, La) + (2, Ia), where wI, E IR and d=20. We consider the network f R + R as a 3-, 4-, or 6-layer NN. We generate the true output by y = f(x)) For the pertaining data, we use 50 in-context examples, and sam- ple them from N--2, Ia). For the testing data, we use 75 in-context examples, and sample them from four distribu- tions: (i) W1 = 1, 0, (ii) W1 = 0.9, = 0.1, (iii) W1 0.7, w2 = 0.3, (iv) w1 0.5, w2 = 0.5. We show the results of 6-layer NN in Figure 2.
[n + 1], and approximation gradients ( with (7)) = 0 such that
w Projo (wr(11) V,, (t) ((-1)),
where || 71 an error term.
Proof Sketch. By our assumption Projw = MLP we only need to find a transformer to implement gradient de-
scent = = w - 7/ Vn (w). For any input takes form (F.2), let function = Dx= + Dx maps w to w-n/V, (w) and preserve other elements. By Lemma there exists transformer block fsottmax capable of upproximating f with any desired small error. Therefore, fsoftmax MLP suffices our requirements. Please see Appendix E.2 for a proof. D
6 Conclusion
We provide an explicit characterization of the ICL capabili- ties of both ReLU- and toftmax-trrnssfrmeer in approximat- ing the gradient descent training process of a N N-layer feed- forward neural network. Our results include approximation (Theorem and Theorem and convergence (Corollary 1.1) guarantees. We also provide experimental validation.
Remark 4 (Note the Use of the Universal Approxima- tion Lemma 16). We assume the transformer parameters are independent of inputs when invoking the universal approxi- mation lemma (Lemma 16). This is because the lemma guar- antees the existence a transformer that can approximate any L-Lipschitz permutation equivariant function over a bounded input domain. Although we can invoke the lemma to show that transformer can represent global minimizers of the target N --layer neural network, such an approach only proves existence without providing explicit construction. Our goal is to show how the transformer performs gradient descent step by step as explicitly as possible.
Extensions. We further extend our analysis from N -laayer networks with the same input and output dimensions to scenarios with arbitrary dimensions (Appendix D).
Applications. We apply our results to learn the score func- tion of the diffusion model through ICL in Appendix G.
Related Work and Limitations. Please the related works, a detailed comparison with (Wang et al., 2024). broader impact, and limitations in Appendix A.

---

In-Context Deep Learning via Transformer Models
Acknowledgments
Advances in neural information processing systems, 33:
1877-1901, 2020.
The authors would like to thank Zhijia Li, Mimi Gallagher, Sara Sanchez, Dino Feng and Andrew Chen for helpful discussions; Hude Liu, Hong- Yu Chen, Jennifer Zhang, and Teng-Yun Yun Hsiao for collaborations on related topics; and Jiayi Wang for facilitating experimental deployments. JH also thanks the Red Maple Family for their support. The authors also thank the anonymous reviewers and program chairs for their constructive comments.
Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, E.,, Lee, P., Lee, Y. T., Li, Y., Lundberg, S., et al. Sparks of artificial general intel- ligence: Early experiments with gpt-4. arXiv preprint (tti:2303.12712,, 2023.
Chan, S. et al. Tutorial on diffusion models for imaging and vision. Foundations and Trends@ in Computer Graphics and Vision, 16(4):322 471, 2024.
JH is partially supported by the Walter P. Murphy Fellow- ship. HL is partially supported by NIH R Vie and Dolby. This research was supported in part through the computational resources and staff contributions provided for the Quest high performance computing facility at Northwestern University which is jointly supported by the Office of the Provost, the Office for Research, and North- western University Information Technology. The content is solely the responsibility of the authors and does not neces- sarily represent the official views of the funding agencies.
Chen, M., Du, J., Pasumuru, R., Mihaylov, T., Iyer, S., Stoy- anov, V., and Kozareva, 7. Improving in-context few-shot learning via self-supervised training. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 3558-3573, 2022.
Chen, M., Huang, K., Zhao, T., and Wang. M. Score approx- imation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Con- ference on Machine Learning, pp. 4672-4712. PMLR, 2023.
Impact Statement
By the theoretical nature of this paper, we do not expect immediate negative social impact.
Chen, M., Mei, S., Fan, J., and Wang, M. An overview of diffusion models: Applications, guided genera- tion, statistical rates and optimization. arXiv preprint MXXi:2404.0077/, 2024.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774. 2023.
Dai, D., Sun, Y., Dong. L., Hao, Y., Ma, S., Sui, Z., and Wei, F. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2022.
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization large language models. Advances Neural Information Pro- cessing Systems, 35:38546-38556, 2022.
Dai, Z., Yang, Z., Yang, Y., Carbonell, J. G., Le, Q. and Salakhutdinov, R. Transformer-xt xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pp.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Trans- formers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural infor- mation processing systems, 36:57 114551225-57111 2023.
2019.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? case study of sim- ple function classes. Advances in Neural Information Processing Systems, 35:30583-30580 2022.
Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C.,, Ding, H., Dong, K., Du, Q., Fu, Z., et al. Deepseek I1m: Scaling open-source language models with longtermism. CoRR, 2024.
Gu, Y., Dong, L., Wei, F., and Huang, M. Pre-training learn in context. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 4849-4870, 2020.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosse- lut, A., Brunskill, E., et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.0725.. 2021.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute- optimal large language models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, pp. 30016-30030, 2022.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., et al. Language models are few-shot learners.
10

---

In-Context Deep Learning via Transformer Models
Hu, J. Y.-C., Wang, W.PP. Gilani, A., Li, C., Song, Z., and Liu, H. Fundamental limits of prompt tuning trans- formers: Universality, capacity and efficiency. In The Thirteenth International Conference on Learning Repre- sentations, 2025a.
Shi, Z., Wei, J., Xu, Z., and Liang, Y. Why larger language models do in-context learning differently? In Forty-first International Conference on Machine Learning, 2023b.
Shin, S., Lee, S. W., Ahn, H., Kim, S., Kim, H. S., Kim, B., Cho, K., Lee, G., Park, W., Ha, J. W., et al. On the effect of pretraining corpora in-context learning by large-scale language model. In 2022 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, NAACI. 2022, pp. 5168-5186, 2022.
Hu, J. Y.-C., Wu, W., Lee, Y.-C., Huang, Y.-C,, Chen, M., and Liu, H. On statistical rates of conditional diffusion transformers: Approximation. estimation and minimax optimality. In The Thirteenth International Conference on Learning Reprexentations, 2025b.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. I., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint ar arXiv:2401.0400 0440888, 2024.
Song, Y., Garg, S., Shi, J., and Ermon. S. Sliced score matching: A scalable approach to density and score es- timation. In Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020a.
Song, Y., Sohl-Dickstein, J.. Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020b.
Kajitsuka, T. and Sato, I. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? In The Twelfth International Conference on Learning Representations (ICLR), 2024.
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Li, S.,, Song, Z., Xia, Y., Yu, T., and Zhou, of in-context learning and weight shifting for softmax regression. Advances in Neurul Information Processing Systems, 37:62584-62611 2024.
T. The closeness
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint 1xXXi::232.13977 2023.
Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zeutlemoee, L. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 11048- 11064, 2022.
Vincent, P. A connection between score matching and de- noising autoencoders. Neural computation, 23(7): 1674. 2011.
Panigrahi, A., Malladi, S., Xia, M., and Arora, S. Trainable transformer in transformer. In International Conference on Machine Learning, pp. 39448-39492. PMLR, 2024.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp.
Panwar, M., Ahuja, K., and Goyal, N. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2023.
35151-35174, PMLR, 2023.
Wang. Z., Jiang, B., and Li, In-context learning on function classes unveiled for transformers. In Forty-first International Conference on Machine Learning, 2024.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 188)99 2019.
Wies, N., Levine, Y., and Shashua, A. The learnability of in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Safran, I. and Shamir, O. Depth-width tradeoff's in approxi- mating natural functions with neural networks. In Inter- national conference on machine learning, pp. 2979-2987. PMLR. 2017.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02088 2021.
Shi, W., Min, S., Lomeli, M., Zhou, C., Li, M., Lin, X. V., Smith, N. A., Zettlemoyer, L., Yih, W.-t,, and Lewis, M. In-context pretraining: Language modeling beyond document boundaries. In The Twelfth International Con- ference on Learning Representations, 2023a.
Yang, L., Zhang, Z., Song, Y., Hong, S., Xu, R., Zhao, Y., Zhang, W., Cui, B., and Yang, M.-H. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023.

---

In-Context Deep Learning via Transformer Models
Yoo, K. M., Kim, J., Kim, H. J., Cho, H., Jo, H., Lee, S.- W., Lee, S.-G., and Kim, T. Ground-truth labels matter: A deeper look into input-label demonstrations. In 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, pp. 2422-2437, 2022.
Zhang, R., Frei, S., and Bartlett, P. I.. Trained transform- learn linear models in-context. Journal of Machine Learning Research, 25(49):1-555 2024.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint 2050022205.01668, 2022.
12

---

In-Context Deep Learning via Transformer Models
Supplementary Material
A Related Work, Broader Impact, Further Discussion and Limitations
14
A.I
Related Work
14
A.2
Further Discussion
15
A.3
Limitations
15
Supplementary Theoretical Backgrounds
17
B.1 Transformers
17
B.2 ReLU Provably Approximates Smooth k- Variiabl Functions
17
Proofs Main Text
18
C.I Proof of Lemma
18
C.2 Proof of Lemma
19
C.3 Proof of Lemma
22
C.4 Proof of Lemma 4
23
C.5 Proof of Lemma'
25
C.6 Proof of Lemma 6
26
C.7 Proof of Theorem
28
C88 Proof of Corollary
Extension: Different Input and Output Dimensions
33
Extension: Softmax Transformer
35
E.1 Axillary Lemma: Universal Approximation of Softmax Transformer
................................
35
E.2 In-Context Gradient Descent with Softmax Transformer
................................
35
E.3 Proof of Theorem 2
36
E.4 Proof of Lemma 16
36
F
Experimental Details
44
F.1
Experiments for Objectives and 2
44
F.1.1
Performance of ReLU Transformer.
45
F.1.2
Performance of Softmax Transformer.
45
F2
Experiments for Objective 3
45
F3
Experiments for Objective 4
46
Application: ICL for Diffusion Score Approximation
48
G.1 Score Matching Generative Diffusion Models
48
G.2 ICL for Score Approximation..
48
13

---

In-Context Deep Learning via Transformer Models
A Related Work, Broader Impact, Further Discussion and Limitations
In this section, we show the related works, broader impact and limitations.
A.1 Related Work
In-Context Learning. Large language models (LLMs) demonstrate the in-context learning (ICL) ability (Brown et al., 2020), an ability to flexibly adjust their prediction based on additional data given context. In recent years, a number of studies investigate enhancing ICL capabilittes (Chen al., 2022; Gu et al., 2023; Shi al., 2023a), exploring influencing factors (Shin et al., 2022; Yoo et al., 2022), and interpreting ICL theoretically (Xie al., 2021; Wies et al., 2024; Panwar et al., 2023; Li et al., 2024; Bai et al., 2023; Dai et al., 2022). The works most relevant to ours are as follows. Von Oswald et al. (2023) show that linear attention-only Transformers with manually set parameters closely resemble models trained via gradient descent. Bai et al. (2023) provide a more efficient construction for in-context gradient descent and establish quantitative error bounds for simulating multi-step gradient descent. However, these results focus on simple ICL algorithms or specific tasks like least squares, ridge regression, and gradient descent on two-layer neural networks. These algorithms are inadequate for practical applications. For example: (i) Approximating the diffusion score function requires neural networks with multiple layers (Chen et al., 2023). (ii) Approximating the indicator function requires least 3-layer networks (Safran & Shamir, 2017). Therefore, the explicit construction of transformers implement in-context gradient descent (ICGD) on deep models is necessary to better align with real-world in-context settings. Our work achieves this by analyzing the gradient descent on N-layer neural networks through the use of ICL. We provide a more efficient construction for in-context gradient descent. Furthermore, we extend our analysis to mottmax-transfommer in Appendix E to better align with real-world uses.
In-Context Gradient Descent on Deep Models (Wang al., 2024; Panigrahi et al., 2024). A work similar to ours is (Wang et al., 2024). It constructs a family of transformers with flexible activation functions implement multiple steps of ICGD deep neural networks. This work emphasizes the generality of activation functions and demonstrates the theoretical feasibility of such constructions. Our work adopts a different approach by enhancing the efficiency of transformers and better aligning with practical applications. Here we highlight key differences:
More Structured and Efficient Transformer Architecture. While Wang et al. (2024) use O(N* transformer to approximate L gradient descent steps N-layer neural networks, our approach achieves more efficient simulation for ICGD. We approximate specific terms in the gradient expression to reduce computational costs, requiring only a (2N- + 4--yyer transformer for L gradient descent steps. Our method focuses on selecting and approximating the most impactful intermediate terms in the explicit gradient descent expression (Lemmas 3 to 5), optimizing layer complexity to (NL)
Less Restrictive Input and Output Dimensions for N-layer Neural Networks. Wang et al. (2024) simplify the output of --llaeer networks to a scalar. Our work expands this by considering cases where output dimensions exceed one, as detailed in Appendix D. This includes scenarios where input and output dimensions differ.
More Practical Transformer Model. Wang et al. (2024) discuss activation functions in the attention layer that meet a general decay condition (Wang et al., 2024, Definition 2.3) without considering the Softmax activation function. We extend our analysis to include Softmax-tramm transformers. Our analysis reflects more realistic applications, as detailed in Appendix E.
More Advanced and Complicated Applications. Wang et al. (2024) discuss the applications to functions, including indicators, linear, and smooth functions. We explore more advanced and complicated scenarios, i.e., the score function in diffusion models discussed in Appendix G. The score function (Chen et al., 2023) falls outside the smooth function class. This enhancement broadens the applicability of our results.
Another work similar to ours is (Panigrahi et al., 2024). It proposes a new efficient construction, Transformer in Transformer (TINT), allow a transformer simulate and finetune more complex models (e.g., one transformer). The main distinction between ours and (Panigrahi et al., 2024) lies in the different aims: Our approach focuses using standard transformer for the simulator (with a minor modification: the "element-wise multiplication layer"), and we provide a theoretical understanding of how a standard transformer can learn the ICGD of an N-layer network using ICL. In contrast, the work (Panigrahi et al., 2024) aims to build even stronger transformers by introducing several structural modifications that enable running gradient descent on auxiliary transformers. While it demonstrates in-context gradient descent for a more advanced model, i.e., one transformer, our work offers the following potential advantages:
14

---

In-Context Deep Learning via Transformer Models
Explicit Transformer Construction. We provide an explicit construction of the transformer, whereas the work (Panigrahi et al., 2024) does not detail the explicit construction of model parameters within their transformer.
Exact Gradient Descent. We compute the exact and explicit gradient descent for an N -llayer network (Lemma 1). Building on this, we employ the transformer's ICL perform gradient descent on all parameters. However, the work (Panigrahi al., 2024) stops the gradient computation through attention scores in the self-attention layer and only updates the value parameter in the self-attention module. Additionally, it uses Taylor expansion to approximate the gradient.
Rigorous Error and Convergence Guarantees. We provide rigorous gradient descent approximation errors (for multiple steps) and convergence guarantees for the ICGD on an N-layer network (Corollary 1.1 and Lemma 14). However, the work (Panigrahi et al., 2024) only presents the gradient approximation error for each specific part of the parameters in a single step.
Attention Layer Better Aligned with Practice, Our analysis is based on ReLU-attention (Theorem 1) or Softmax- attention (Theorem 5), whereas the work (Panigrahi et al., 2024) utilizes linear attention. Our choice of attention layer better aligns with practical applications.
A.2 Further Discussion
We provide an interpretation and example of how to explicitly instantiate the constants for ReLU approximations in Lemmas 2 to 4. The key reason that the function approximated by the sum of Rel.Us is simple our context, such as the Sigmoid activation function. For such simple functions, it straightforward derive an explicit construction.
Here, we take the Sigmoid activation function example and propose explicit construction method. Let r(2) denote the Sigmoid function.
Segment the Input Domain. For example, divide the domain [-10, 10) smaller intervals such as [-10, 99, [-9, -8]. .... 10).
Approximate Each Segment Locally Using a Linear Function via Linear Interpolation. For instance, in the domain [9, 10), approximate r(z) using a linear function 01 + C1, where a1 and Cj are calculated as follows: a1 = (r(10) - (())) - -9), and c₁ = r(9) - a1 * 9.
Approximate Linear Function 01 + (2 E [9, 10)) Using Sum of ReLU Terms. This step involves two substeps, which are straightforward to implement: (i) Approximate the indicator function for z E [9, 10) using a sum of ReLU terms. (ii) Approximate the constant e1 using the sum of ReLU. This is because bias terms are not included in the sum of ReLU terms in Definition 4. The bias term e₁ must be approximated using an additional sum of ReLU terms.
Combine All the Sum of ReLU Approximators Across All Segments. Finally, integrate the approximations for all segments to construct the complete approximation.
Estimation of the Parameters in Definition& 4. Capprox 0.625, R 10, H 80, and C 25. Furthermore, to achieve higher precision in the approximation, it is sufficient to use finer segmentations.
A.3 Limitations
Our work has the following six limitations:
Although we provide a theoretical guarantee for the ICL of the Softmax-Transforme Transformer to approximate gradient descent in N-layer NN, characterizing the weight matrices construction in Transfor-rr remains challenging. This motivates us to rethink transformer universality and explore more accurate proof techniques for ICL in Softmax- Transformer, which we leave for future work.
The hidden dimension and MLP dimension of the transformer in Theorem are both Ō(N + Dw. which is very large. The reason for the large dimensions is that we use ICL to perform ICGD on the N-layer network, we need to allow the transformer to realize the N-layer network parameters. This means that is reasonable for the input dimension to be so large. However, is possible to reduce the hidden dimension and MI.P dimension of the transformer through smarter construction. We leave this for future work.
The generalization capabilities are limited compared with traditional transformers. In our setting, the pretraining task refers to using in-context examples generated by an N-layer network for a given N. Specifically, during pretraining, the distribution of the -lyeyer network parameters predetermined (e.g., (O, T)). The input data distribution of N -layerr
15

---

In-Context Deep Learning via Transformer Models
network for generating the in-context examples is also predetermined (e.g., ((-2,7) I)). The generalization capabilities include the following aspects: (i) Varying the input data distribution for the N-layer network to generate the in-context examples. For example, change the input data distribution from (- 1) to 0.9.V I) + 0.1 f) during the testing in Appendix F.I. (ii) Varying the distribution of the N-layer network parameters. For example, we change the distribution from (O, 7) to N(.55, /) in Appendix F.2. The above points lead to differences between the distributions of in-context examples during pretraining and testing. However, we must generate the in-context examples by the -layeer network with the same hyperparameters, including the network width and depth. We leave the theoretical analysis of broader generalization capabilities for future work.
In theory, the FLOPs (Hoffmann et al., 2022) required to perform one forward pass of the transformer are greater than those required for the direct training of an N-layer network. (i) For the forward pass of the transformer, the FLOPs for in-context learning (ICL) are O K /e)), where the approximation error the sum of ReLUU. (ii) For direct training of the N --layer network, the FLOPs without ICL. are OnLLNK3) Therefore, the FLOPs required for ICL exceed those needed for direct training of the N-layer network. However, experimental results in Appendix F demonstrate that the transformer with ICL can achieve the performance of trained 6-layer network using fewer FLOPs in practice (3.3 billion vs. 7.6 billion FLOPs). This finding encourages further exploration of more efficient architectures. We also leave this topic for future research.
The empirically trained transformer differs from the transformer constructed in our theoretical analysis. Our experiments confirm the existence of transformer capable of simulating gradient descent (GD) steps for N-layer neural networks through in-context learning (ICL). Despite this discrepancy, the limitation does not affect the primary contribution: establishing the theoretical existence of this transformer by explicit construction.
There are two minor differences between the transformer used in the theoretical analysis and a standard transformer: (i) The transformer used in the theoretical analysis incorporates element-wise multiplication layer, specialized variant of self-attention that retains only the diagonal score and allows efficient implementation. (ii) It does not alternate self-attention and MLP layers. We emphasize that this also qualifies standard transformer because we view either an attention or an MLP layer as equivalent to an attention plus MLP layer due to the residual connections.
16

---

In-Context Deep Learning via Transformer Models
B
Supplementary Theoretical Backgrounds
Here present some ideas we built on.
B.1 Transformers
Lastly, we introduce key components for constructing a transformer for ICGD: ReL.U-Attention, MLP, and element-wise multiplication layers. We begin with the ReLU-Attention layer.
Definition (ReLLU-Attention Layer). For any input sequence H E RDx, an M-head ReL.U-attention layer with parameters (Qm: Km, J In]mm(a)) outputs
1 M
Attug (H) H
Σ m==
(Vm o(QQmH) (mH)),
where Qm, K,, Vn DxD and (-)) element-wise ReL.U activation function. In vector form, for each token h; E DR in H, it outputs (H)): - h4 + Em
Km hm) ) Vmhhs
Notably, Definition 7 uses normalized ReLU activation a/n, instead of the standard Softmax. We adopt this for technical convenience following (Bai et al., 2023). Next we define the MLP layer.
Definition 8 (MLP Layer). For any input sequence H E R Dxn, an d'-hidden dimensions MLP layer with parameters
[WH)) 0 = (W1, W2) outputs MLP,(H) == H + (2 H), where W1 Rd x W2 Dxd' and a(-) R - IR is element- wise ReLU activation function. In vector form, for each token h; E TRD H, it outputs ML (L)P(HH): == h, + W20 (W,(W,,)))
Then, we consider a transformer architecture with 2 transformer layers, each consisting of self-attention layer followed by an MLP layer.
Definition (Transformer). For any input sequence H e R Dxn, an I--laee transformer with parameters {f OMLP} outputs
TF) MLP (L) mip
o Attn
MLP
001) Attn (()) (H),
atin
where 0 (( OMLP) consists of Attention layers AuL {(Q K Vth he( mE(My and MLP layers 0MLP {(W], he) Above, for any 1 E [L], m
[M'], Q K Vi E DDD and (W), W3) E Rd' D x Dxd this section, we consider ReLU Attention layer and MI.P layer are both a special kind of 1-layer transformer, which for technical convenience.
For later proof use, we define the norm for L-layer transformer TF's as:
Bg max [E[L]
Σ 11 + [| + ||W2l1 22
(B.1)
mE[M]
The choice of operation norm and max/sum operation is for convenience in later proof only, as our result depends only on Bg.
B.2 ReLU Provably Approximates Smooth Variable Functions
Following lemma expresses that the smoothness enables the approximability of sum of ReLU.
Lemma 7 (Approximating Smooth k- V--aaiialle Functions, modified from Proposition A.1 of (Bai al., 2023)). For any C > 0, R 2 1. If function g R R such that for s:= [ 1)/2] 1, is C function on B%, (R), and for all i E (0...
BEB (R) sup (V'((2)]le & L,,
max L < Ci. OSSS
then function is (e, R, H, Cpaapproxiemale by sum of ReLUs (Definition 4) with H S C(k) log(1 + Ct/e)/e² and C' < C1 where ((k) is a constant that depends only on k.
17

---

In-Context Deep Learning via Transformer Models
Proofs of Main Text
C.1
Proof of Lemma
Lemma 8 (Lemma Restated: Decomposition of One Gradient Descent Step). Fix any B 0. Suppose loss function C,(w) data points {xs, Ji)he follows (2.2). Suppose closed domain W and projection function Proiw(w) follows
(3.4). Let A,()), r)) R,()). V3 be as defined in Definition 2. Then the explicit form of gradient VL,, (w) becomes
A((1)
(w))
2n 1=1
[A((N)
where Aj) denote the derivative of ((p:(N), 3/) with respect to the parameters in the j-th layer,
R4( 1). T pi - 1)T (R.(N- Laxa @ pr(N 1TT (8) (p:(N)
(R, - 1) VN
(M(p.(N). dp((N)
)T, j # N
A:())
= N.
Proof of Lemma
We start with calculating
V (w)). By chain rule and (2.2),
1
VwCn(W)
(N
(P:(N) (p.(N))
(By (2.2) and chain rule)
2n
dw
R4x1
Thus only need to calculate aw (NV) For vector x and function : IR, use r(z) denote the vector that i-th coordinate is r(xz). Let R,(i), Vj follows Definition 2, then it holds
BaXK
R
(p:(N)
Ør( VN
Pi(N- - 2))
BBy Definition 1)
dw
Ow
ReXDN
BaXDN
Ōr(VN Pi(N - 1)) OVN (N 1) OVN (N
(By chain rule)
dw
Raxa
BeXON
OV p: p.(N - 1) dw
di ( (NN, (N 11)), .))..... ( PNN
(By Definition 2)
= RN - OVN PNN-11
(C.1)
Ow
Notice that for any l E [d], UN2 is part of w, thus
Dx (x-1+(K-1KK K
Dy DN-1-
OUNz dw
I
0
1 R dx
DN.
(C.2)
Therefore, letting @ denotes Kronecker product, it holds
OVN PNNN 11
Dw
Op-(N dhe
+ +P.(N - 1)T
80x1
UN
:
(By chain rule and product rule)
(P)
+ PN(N 1)T
SUNA 8w
Ow
18

---

In-Context Deep Learning via Transformer Models
(N - 1) + [ODN-niIRXK Op.(N - }
Vv=
(C.3)
dw
where the last step follows from the definition of Vy (i.e., Definition 2) and (C.2). Substituting (C.3) into (C.1), we obtain
(N, Ra(N- (VN
Op.(N 1)
[OD laxa 0 p:NN ppN -1)J)-
Ow
Ow
Similarly, for any [N], we prove
0p.(i)
10 11.11, Op.(j -
R2(j
1)T; 10D Dpg_iilxx @ peli -)) DD-DD,))
(C.44
Ow
dw
By the recursion formula (C.4), for any [N - 1), we calculate A,) follows,
A(() (( (B((A((N))))
) [D; [Dj-1:
(By Definition 2)
(V:(N)
) T. (
[Di-1: Di1
By transpose property)
Ow
Op,(N)
(p.(NN
)T [1, Dy-1 Dy) (l(,, )T
Ow
Op((N)
(l(p:(N), )T Op.(N)
(R - 1) VN .......((j-- - 1) [ & p. p)()=
(By (C.4)
where [x a b] denotes sub-matrix of M, which includes all the columns but only the rows from the a-th row to the b-th row of A. Similarly, for N, it holds
A(NN) = (R.(N 1) [Taxa @ PN(N - 1) (l(p:(N), )T
Op:(N)
Thus we completes the proof.
C.2 Proof of Lemma
Lemma 9 (Lemma 2 Restated: Approximate p;(j)), Let upper bounds B,, B, > 0 such that for any k [K],j [N] and i [n], ||u lla & B,, and | I By. For any j E [N], E [n]. define
B. max SB
(r()), B = B and By max B
Let function r(t) be (Er, R1, M1, C Japproxximable for R₁ max B, 11, M1 O(C9E, where C₁ depends only on R₁ and the -m-smoothness of r. Then, for any 6 there exist N attention layers Attng, . ..... Attug such that for any input h; E RD takes from (2.1), ), they map
Attne hr = ys: w. P.) 0;
h [x U6 wi :A 1);0;144)
where p, (z) is approximation for p:(j) (Definition 1). In the expressions of h; and h; the dimension of 0 differs. Specifically, the O in h; is larger than in h, The dimensional difference between these 0 vectors equals the dimension of p.(j). Suppose function r is smooth in bounded domain W, then for any i E [n + 11, E [N], p.(i) such that
K2/2 VKE, K1/2 T. BB) Vdee,
1
(.()) p.() + eii.)), (11,)))2= <
j j= N
19

---

In-Context Deep Learning via Transformer Models
Additionally, for any j [N], the norm of parameters Bo, defined as (B.1) such that
Ba, S + KC₁.
Proof of Lemma 2. First need to give approximation for activation function r(t). By our assumption and Definition 4, r(1) is (er, R1, Mi, C1) >aappoxximable by sum of ReLUs, there exists:
Mi r(1) Σ
M2
m σ (a [4]]]) with Σ
cm
Ci, ||ata]1 S Vm [Mi]
(C.5)
such that SuPtE(-RR R.) (r(() r(]] e. Let pc(0) = (0( zyz Similar pr(i) follows Definition 1, we pick pz(j) such that for any [
A(O)[K] Ps 1)).
(C.6)
Fix any [M], suppose the input sequences h4 (x w; Pe (s(......::: : 1); 0; 1:4]. Then for every E [Mi], (or & E [d] if j N), define matrices 0₁ K m. k V
m DXD such that for all i e [n 1].
a [1] p:(j 1) a [2
Vik
Om =
KJ l4 =
Vingchs = Gn
(C.7)
where ej, denotes the position unit vector of element [k(] because this position only depends on j. k. Since input h4 = [105 yes wi p(()): ij - 1); 0; 1; 42). those matrices indeed exist. In fact, it is simple to check that
mm [ (
2 m.
om [2]
0
Ix (j. k)
0
0
0
KJ m.k
0
0
0
0
0
0
VJ
0
0
Cm
(j. k))
0
(C.8)
m,k
0
0
are suffice to (C.7). Ix, et (j,k) represents their positions are related to variables parentheses. Addition, by (B.1), notice that they have operator norm bounds
max 1. l11 1 k.m,k
max [K] mn 11,
ΣIVA
VI S KCi.
max
j.m,k
kmm
Consequently, for any j E [N), BB, 11+
By our construction follows (C.7), a simple calculation shows that
Σ
all'In hi, K2 ho)) Van
K
Sy ( [0J,p.li 1): 1))
By construction (C.7))
m=11
K Σ (F(0).p.(11
(By definition follows (C.5)
(::P:)):0.
(By definition of pr(i) follows (C.6)
20

---

In-Context Deep Learning via Transformer Models
Therefore, by definition of ReLU Attention layer follows Definition the output hi becomes
h,3 = [Attno,(Hh) (h:))
a+1
h4 +
Σ
E
a( hi K2 (a)) Vm hs
n s=1
mE [eI]
n+1
h4 +
(n 1):00 Σ
n +
[m:; UC: w: :..)
:Bs(j 0);0; 1;4,] + (0,
= [x MA3 wi (.(;......... : O):
Therefore, let the attention layer 02 = {1Omm k m. Vi k)
(n. (k.m)) we construct Attng,, such that
h [ ysi wi; (1): ... : 1;4] Attao h [
[T ysi w. D. (1): (.)
In addition, by setting R max[BBBB,r. 1} the lemma then follows directly by induction on j. For the base case 1, it holds
[ (L)] = =)) m.)l
(By Definition 1)
Ey
(By definition of F follows (C.5))
Suppose the claim holds for iterate j 1 and function r is L, smooth in bounded domain W. Then for iterate j.
[.((i)[k] [+(]
$ [p:0) - r(o - 1)) +rr(0]p0000 - 1)) p(())
(By triangle inequality)
Se, L 11 - -1)-p((i-1)
(By (C.5) and Mauchy-Schwarzz inequality)
1-2
S + KI, Σ K2/2 B,)) 1=0
(By inductive hypothesis)
6, KM L,
Thus, it holds
(IP(() PP-0)ll
$ VK(G Σ R(BLLB)
1=
This finishes the induction. Then for the output layer j = N, it holds
Il.(NN) - PP:(Nl22
- PN()))11
$ Valce cr
Σ R1/TT1BB))
1=0
Thus complete the proof.
21

---

In-Context Deep Learning via Transformer Models
C.3 Proof of Lemma
Lemma 10 (Lemma 3 Restated: Approximate r(())), Let upper bounds Bu: Bz 0 such that for any A E [K], [N] and i E [n], | I B,, and |2 I By. For any j E [N],e E [n)] define
B max
(()) B. : Bz and By = max B BBB.B.
Suppose function (') is (Gr,, R2, M2, C2) approximaae for R2 max max Br, 1), M2 I Ō(C3ct, where C2 depends only on R2 and the C*-smoothness of r'. Then, for any e, > 0, there exist an attention layer Atino such that for any input h4 E IR takes from (3.13), it maps
Attnen+1 h = [ ys: wi Ps: F 0; 1;
h = a ye w: Pe: 1:4]
where r(( is approximation for ri) (Definition 2) and := (0) :.. : -1) (N Similar to Lemma 2, the expressions h2 and hi. the dimension of 0 differs. addition, let E, be defined in (3.12), for any [n 1].
j e [N]. & E [K], r() such that
r0 - 1) =
(i.j,k), [c(,,, S 6, + Ey
where 6, denotes the error generated approximating r by sum ReLUs F follows (C.5), Additionally, the norm of parameters Bo defined as (B.1) such that BaN+: I K(N - 1)C2.
Proof of Lemma 3. By Definition 2, recall that for any j E [N], E In + 1] E []]
MO))- = ("(
(C.9)
Therefore we need to give approximation for r'. By our assumption and Definition 4, r'()) is (ep,, RS, M2, C2)- approximable by sum relus. In other words, there exists:
Ma
M2
((1) Σ 2n ( ( [t: 1)) with Σ
< C2, | S Vm E [M2],
(C.10)
m==1
m==1
such that SUPpEE[-Ra.ff R2] (1) r' r()( Gr. Similar to (C.9), we pick r(() such that
T00K) = ((
(C.11)
To ensure (C.11), we construct our attention layer as follows: for every j E [N], E [Ma], [K], we define matrices @ K-N+1 NN++1 E IR DDD such that
j.n,k J.m,k
2
[1]
(j( 1) 2 [2]
U/e
N+1 hi
KN
KN++1 Po
VN++ 2 e3
(C.12)
for all i [n + 1] and ei denotes the position unit vector of element F(())(k). Since input h4 = [Iss [Is wi p.: 1:t] similar (C.8), those matrices indeed exist. In addition, they have operator norm bounds
max 1QQ 1 S 1, j.m.k
max
J
M.NNtlll 1
Σ j.m.k
CI smmbll S - 1ON--1CCK 1)C2.
Consequently, by definition of parameter norm follows (B. 1), Boxe
S K(N-KNNN10 1)C2.
22

---

In-Context Deep Learning via Transformer Models
simple calculation shows that
Σ
(1Q jam,k hi, J.m,k hx) )
j.m,k
JEIN MM KEIK]
N
K
M
2 o - 1);
By construction follows (C.12))
N
K
ΣΣ
(F
[1] 1)
(By definition of F follows (C.5))
j=1 1=1
[0: F((0);
:F - 1):0]
(By definition of F( follows (C.11))
00: 0]
(By definition of F!)
Therefore, by definition of ReLUU Attention layer follows Definition 7, the output hy becomes
h4 [Attno (h:)]
n+1
h4 Σ
E
a1Q2mmm K. ha) V A
n2 +
jEIN m (Ma)kkkkk
n+1
h2
(n 00:4;
n +
Jan ys: wi Pa; 1: 00
Jan Ms; w; p;; ; 1;]:
Next, we calculate the error accumulation in this approximation layer. By our assumption, R2 [B, 13. Thus, for any [N], k [K], i [n + 1), it holds
P 1) - Ro.
As our assumption, we suppose function is smooth in bounded domain W. Combining above, the upper bound of accumulation [k)] (k(] becomes
is
[())( +(()][]] F()) -
-
(By triangle inequality)
I e,r + Ly 111 11 - - e, Ly E,
(By (C.10) and Cauchy-Schwarr Schwarz inequality) (By definition E, follows (3.12))
Thus we complete the proof.
C.4 Proof of Lemma
Lemma 11 (Lemma 4 Restated: Approximate 0, ( 9:)). Let upper bounds B,, B., > 0 such that for any k E
[K],e E [N] and E [n), || ll2 I B,,, and ||x.|22 < By. For any i E [d]. suppose function u(t. [b] be (4, Rg, M Chy-- approximable for R& max B, Br, By: 1), M3 O((C))² 2), where C. depends only on R25 and the -mmsmoothness of u(t, [k). Then, there exists an MLP layer MLP Bw such that for any input sequences h; E IR takes from (3.15). it maps
MLP hi [a4; ysi wi P.i F/: gs: 0; 1;t]],
h, = [ wi
wi::;pp:ppp F/F 0; 1;
where E Rd is approximation for [(pp(N) 21). For any & E [d], assume u(px(N), v:) is Lr- Lipschitz continuous.
23

---

In-Context Deep Learning via Transformer Models
Then the approximation such that,
[k] u[p:(N),, (k) + k(), with l<(2, S a + L,E,G.
Additionally, the parameters 02 such that Ba I max + Cz).
Proof of Lemma 4. By assumption and Definition 4, for any k [d], function u|k](tt,) v) is (a,, Rg, M2 Cz)- approximable by sum of relus, there exists:
M
M
9 (4, = Σ
m
[4: 11)) with Σ 3. 3.
((a
C
3.k
C3.
1anlll2 S 1, Vmn E [M3],
(C.13)
m=1
such that suP(1.))(l-Rs.R) (t, ut, v) S ej. Then construct our MLP layer.
Let M3 == Ex=1 My we pick matrices WiN+1 R M N+N+
R Dx Ma such that for any i E [n + 1], E [Ms].
ai [1] (NN) a 3. [2]
y +
3.1 [3]
Ra(1 - t,)
3.1 "Ng [1] (N: + 2. [2]
y: +
Mi
[3 Rs (1 t,)
W
E RR Ma
41 [11 -::(N) oj (2) Yi +
Ra - 4))
3.d [1] (N) +
3.d [2] Ui +
3.d [3] R3 (.) M
Mg
W
[j. ml
3.6 Cm
i D M m M
(C.14)
where DK denotes the position of element 9i [k] Since input h; I wi Pi; F/: O; 1; t.], similar to (C.8), those matrices indeed exist. Furthermore, by (B.1), they have operator norm bounds
W + < 1,
| ||| S C3
Consequently, Bax42 S max Rs + 1, C3).
By our construction (C. (C.14), simple calculation shows that
d M
W (W- = Σ Σ
af 3.1
a
[p:(N): ys: 1})
Rs(1 -t)) cm eDj
k= 1 m==1
o 91 (P(P) vi)
= 1 =
(F(F)(N). 0
For E [d], we let gı(k]] 1(1, 1} 9x (F,(N,,, y ye Je for i E [n + 1). Hence, MLP oN maps
MLP 2 hie [an yi wi, P.; F/: g3 0; 1; 4))
h;
bas ue w. p4; F/s 0; 1:4]
Next, we calculate the error generated in this approximation. By setting R3 max{B,B, By, 1}, for any e [n 1] it holds
Pc(N) < Ra, Vi S Ra
24

---

In-Context Deep Learning via Transformer Models
Moreover, our assumption, we suppose function 0 is Ly-smooth in bounded domain W. Therefore, by the definition of the function g, for each i [n)) the error becomes
[9 [i] (P) [e][
< 19: [k]] (L), [A)]]] + [V)P: ((p) yo][A]]
(By triangle inequality)
S (Nl)p:(N)) PP((Nll2
(By the definition of gs follows (C.13) and Ly-smooth assumption)
S ft + Ee
(By the definition of E, follows (3.12))
Combining above, we complete the proof.
C,5 Proof of Lemma
Lemma 12 (Lemma 5 Restated: Approximate &(()), Recall that s;(j) = rí0 1) o (VJ s,( + 1)) fol-
lows Definition 3. Let the initial input takes from (3.17). Then, there exist N element-wise multiplication layers: EWMLe M N.EEEWELL 22 such that for input sequences, [N].
= [
(N): (j + 1;44]
= [x: yes wi 84 (N): ; 1; 4:). where the approximation 8;( is defined as recursive form: for any [n + 1).j [N 1).
they map EWMLegg -- 1 (h.)
FiUi - 1 OV 5( - 1)), jE NN-1] Fi(N 1)
- 1] j N.
(G)
Additionally, for any j [N], B8 defined in (B.1) satisfies Bax+24> S 1.
Proof of Lemma 5. By Lemma and Lemma we obtain pi( and r(()), the approximation for pi(i) (3.3) and r/() respectively. Using pi()) and P((i), we construct N element-wise multiplication layers approximate s;()).
We give the construction of parameters directly. For every [N 11, E [K]. we define matrices Q2N+3-1 Kk k2N+3-j Vv
+2N+3-j E DDD such that for all i [n + 1].
Ug+1, [k]
Q2 Q24431
12 +3
3.0 3,0+1)
+ 1) 0
V V2N+3-1
= 1)[k] Cj
(C.15)
Us+ik [k]
where c3,, denotes the position unit vector of element (k))
Since input h; [ar;; ye; w; w;P;;; 7/; g3 (NN): 5. 1); 0; 1; 4s). similar (C.8), those matrices indeed exist. Thus, it is straightforward check that
Σ kE[R]
71Q22+ 2333-1 h;, K2 22N-3-1 2
K
=
Σ
(V [k, x] 3/j 1))F(( - 1),
(By definition of EWML layer follows Definition 6)
k=1
ri( 1) [4. s(++1)
(By definition of c3,x)
G0 - 1M | * s +1)
0
25

---

In-Context Deep Learning via Transformer Models
= M/ - 1 (vz s.(s + 1))
(By definition of hadamard product)
= 0; 0..
(By detinition of 3;(5) follows (3.20)
Therefore, by the definition of EWML layer follows Definition 6, the output hy becomes
T:s [Attnoan+s 43-5 (h:)
hu +
Σ
Nis K K2N+3-5
V2 -22-3311 1/s
m2
m.
ME(
h [0:
[ y4s wz Ps; 1i 1);...:.8.6) 11;0;;1;1 (0;8:(i);0)
[ar M4 uz; Pai Fli 94: (N :);; 1;4]-
Finally come back approximate the initial approximation s, (2 1) @ 9s- Notice that 9i and r - 1) already in the input hs [1: Mai wi Pa; F 96: 0; 1;], thus it is simple to construct EWML similar to (C.15), such that it maps.
1W (ay ye: w: 0; 9.
(ass yet wi Ps F 0; 1:t,]
Since don't 't using the sum of ReLU to approximate any variables, these step don't generate extra error. Besides, by (3.19), matrices have operator norm bounds
max 1 j.k
N+2+211 ||1 S
max J.k
N+2+j I
|1 S
max V N+2+2111 1.
Consequently, for any j E [N], Ba
I 1. Thus we complete the proof.
C.6 Proof of Lemma
Lemma 13 (Lemma Restated: Error for gi):(j)). Suppose the upper bounds B,, B, > 0 such that for any k [K],
[N] and i E [n]. ||U la S B,, and | ll2 Bg. Let r((i) E RK such that [K] r' ,, P;(j)) follows Definition 2.
Let sj) = R 1)V R. - 2 (N - 1)u follows Definition 3. Let F)( Gs. 5,(j) be the approximations
for (i), (N), 3a). ()) follows Lemma 3, Lemma and Lemma respectively. Let B, be the upper bound of r [
and [k]] as defined in Lemma 3. Let B, be the upper bound of gs[k] and u(p((N), w: [k] as defined in Lemma Then for any i [n + 1),5 [N], e (K].
{(){ B., [+){]] st()[aal S E, + Ey + E.q.
where
P max VK, Vdl
B, max JE(N] ((P B.BBB)NJB,,Il.
N N-j-1
E = max SE[N]
Σ
(B. B, L,,,),,,
1=0
N-j-1
E. max JE[N]
(PB.B.l
E = max JEIN]
(BB, N
26

---

In-Context Deep Learning via Transformer Models
Above, B, is the upper bound of (+(5)[k] and Eg Ey E2 are the coefficients of er, 4 6r in the upper bounds of [.([][[ kk]] respectively.
Proof of Lemma 6. By Lemma 5, we manage to approximate s;(s) by 5e()). By triangle inequality, we have
[(())[b] - 180--88 [+]]
$ 1)
- 1)]]]] | s,(n + 1)| + In - [A] [( 8, + - (1h+1, 8:(n ()]--
We bound these four terms separately, By Lemma 3, (n [)[k}] rn(n [)] is bounded by epr + Ly E,er We then use induction to establish upper bounds for [k] and [s:(]] - [
We first use induction to prove the first two statements. To begin with, we illustrate the recursion formula for 8,()) By (3.20), recall that for any j E [N].
s(0+ 1)), - 1 (VJ (i)
j [N 1] j N.
FI(N 1) @gh gh.
We consider applying induction to prove the first statement:
{())[a (P. B.)
As for the base case, j N:
(+V P((N [)[ gi(k] S B.
Therefore, if the statement holds for j + 1, by (3.20) and our assumption, it holds
3. (n)[kk] = (n - 1)[ (uft1, (n + 1))
(By recursion formula (3.20))
I (n(n 1)[k] || lla ||:: (n + 1)ll
(y Cauchy-schwarz inequality)
S (((n [)[ ||un+1, max (VK, vdy max (s: 1][]]]
< B.) max{ KK, Vd)- (max v K, vd) B.) -n-1111- Br
(By inductive hypothesis)
(P. B, B.))* N. B.
(By definition of P follows Lemma
Thus, by the principle of induction, the first statement is true for all integers j [M]. Moreover, by the definition of B, followsLemma 6, we know B, is the upper bound of (k))). Next we apply induction to prove the second statement:
N-n-1
ss ()M s((M S (cp
L,
Σ
B,P
1=0
+ (B. (Br E, + By Lr )er + Bier + B
For the base case, j = N:
1ss [K] [b]]
= (N - -1)(k] [k] - ,2 - [){{} (N),
(By definition (3.20) and (3.6)
< (N - -1)(k] - ri(N -|)k]] [9.[k]] + |7((NV - 1)[+|| [g.kk] u(p.(N).. u)]]]]
(By triangle inequality)
< (ep- + L, B,,,,, + (e + L, E,
(By (3.14) and (3. (3.16)
= (B E, + L, E, )e, Bie B
Therefore, if the statement holds for j + by (3.20) and our assumption, it holds
(n)))()]] [b]]
27

---

In-Context Deep Learning via Transformer Models
F((n 1)[R] (v,+1, f ()) 2((n 1)[4) (1,+1, (n 1)
(By the recursion formula (3.6) and (3.20))
S (n - -
+ (n -]][k] ((++ (e 1))
(By triangle inequality)
S (6. + L,, B,, +1 8; 1)ll
(By error accumulation of approximating r' follows (3.14)
( + L, B, PB,B, + B,P max (1 + 1)[A}] - 8: + 1)[A]
N-n-2
( L. Bo B. 1=0
E (B,
BB.. Br + BiEN + Breil}
(By inductive hypothesis)
N-n-1
(ep B,,,
Σ (B,BBP))
t=
+
+ Bp F. + BiC + B.
Thus, by the principle of induction, the second statement is true for all integers j E [N - 1]. By the definition of E, follows Lemma 6, it is simple to check that
[() s < EJe, + Eger E.e.
Thus we complete the proof.
C.7 Proof of Theorem
Theorem 3 (Theorem Restated: In-context gradient descent on N-layer NNs). Fix any B., 0, 2 1. For any input
sequences takes from (2. I), their exist upper bounds Bz, B, such that for any i E [n). ||v:l2 S By, |æv]| S Bz Assume functions r(t), (t) and u(t, [k] are L,, Lr,, Li-Lisschtz continuous. Suppose W is closed domain such that for any j E [N 1] and ls e [K].
W c { = [1je]] E RN [|, 1l2 S B.}.
and Projw project w into bounded domain W. Assume Projw = MLP3 for some MLP layer with hidden dimension Dw parameters ||0]] S C.- If functions r(t), (t) and u(t, [k] are C*-smoothnesss, then for any >> 0, there exists a transformer model NN, with (2N L hidden layers consists of L neural network blocks TE EWMLE TF3.
NN, == TFX+ o FWMI. o TF2 o .... TF o EWMT. o TF.
such that the heads number M embedding dimensions D', and the parameter norms By suffice
max (E((2N+4)L
M < (Õ
D' s O(NK3) DW,, tE((2N+4))
max
max IE[(2N+4))
By S O(1) C.r +1,
where Ö(-) hides the constants that depend on d, K, N. the radius parameters BR, By, B, and the smoothness of and f. And this neural network such that for any input sequences H( take from (2.1), NN, (H implements L steps in-context
gradient descent on risk Eqn (2.2): For every E [L]. the (2.N + 4)I-th layer outputs h ((2N+4))) [a:; ya w()); 0; 1: tt] for every i In + 1). and approximation gradients ( such that
Projww (1(1 (V + ((11))
( O,
where ||(|=1)||22 S nc is an error term.
Proof Sketch. Let the first 2N + layers of NNg are Transformers and EWMLs constructed in Lemma 2, Lemma 3, Lemma 4, and Lemma 5. Explicitly, we design the last two layers to implement the gradient descent step (Lemma 1).
28

---

In-Context Deep Learning via Transformer Models
We then establish the upper bounds for error | ( Vn (w:)l2, where V VuLn (w), derived from the outputs of NNo, approximates (w Next, for any > 0, select appropriate parameters e,, fr and &,t ensure that wn -V V Ww I holds for any E [L].
Proof of Theorrmm I. We consider the first N + transformer layers TFA are layers in Lemma 2 Leemma and Lemma 4. Then we let the middle element-wise multiplication layers EWML be layers in Lemma 5. We only need to check approximability conditions. By Lemma 7 and our assumptions, for any <r; Ey', e1, holds
Function r(t) is (c., R1, M1. C1) approximaalle for R1 = max B. Br: 13, M1 & O(C9G3), where C1 depends only on Ri and the Co-smoothness of r(t).
Function r' (t) is (E,r, R2, M2, C2) appprxxxxieeee for R2 = max 1}. M2 S 0(C24=), where C2 depends only on R2 and the C²-smoothness of r' (1).
Function a ((t y) is (EL, R3. M3, -ppaooxixeee for R3 max 13. M3 S Ö(C where C3 depends only on R3 and the -o-smoothness of u(t, [k].
which suffice approximability conditions in Lemma 2, Lemma and Lemma 4.
Now we construct the last two layers to implement v7 Ln (w) and Projw(w). First we construct a attention layer to ,2N V2N
approximate w v Ln (w). For every m E [ E [N], k E [K], we consider matrices Q mmi.k m.j.k E RDD such that
0
7(n + 1)
1.j.k
0
K2N+3
=
0
12N+3
2n
Pi(j
1)
0
h4 = [+([[[=]]]
11n + 1)
2
h4
K-2N433 2.
22N+3
h4 =
p -
1)
(C.16)
2,j.k
2n
0
Furthermore, we define approximation gradient V nn Cn (w) as follows,
V t Cw) :=
2(122N43 h4 K2N (2) V he
17 1) 1= JEE
K
0 Pe(j 10-1) 1) 0
2n Σ
ΣΣ
(+((())
(By our construction (C.16)
K ΣΣΣ
0 Pelj - 1)
(By f(z) =(x) (((x))- - o(-2))
2n
(=1 1=
0
n-1 N
0
Kxk @ pr(j - 1) 3(()
(By definition of Kronecker product)
2m
0
(1
Σ
(By 8m ( = 0 follows Lemma 5)
2m
(N(N)
where A,0) == I @ pe(j - Be(i) denotes the approximation for Az(j). Therefore, by the definition of ReLU
attention layer follows Definition 7, for any i E [n + +1,
h (Attnosnss (ha)
29

---

In-Context Deep Learning via Transformer Models
h; +
a ha K h4) V m.j.k
0
A((1)
[an ts: wo PG 19r399 0;1;4]
Σ
:
2n
A.(N)
0
[mn w v/
0;
(By definition of V En (w))
Since we do not use approximation technique like Definition 4, this step do not generate extra error. Besides,by (B. matrices have operator norm bounds
max J
22 S 1,
max J.m.k
K2N+ J.m.l
S 1,
Σ j.nk
Jam,k
S 2
Consequently, BoaN 13 11 NK Fix any then pick appropriate er, c, such that
|e2l11)12 = v/|l w En (w Vw S ne.
By Definition and Lemma 6, for any j E [N - 1], i E [n)]. it holds
|A.0)) - A1
[8:())(()p s. 1)ll2
(By Definition 2 and definition of A,)))
( st(ii|Il II 1)ll2 (s:0) [p 1) pp 1)ll
(By triangle inequality)
S[(EE + ey + E, B,E,er).
(By (3.12) and Lemma
where B, is the upper bound of [((]][]] and EN BY Ey are the coefficients of c, Er in the upper bounds of [x) [x]] follow Lemma 6, respectively. We can drive similar results as j = N. Actually, by P max{ VK, Vd) follows Lemma 6, above inequality also holds for j N. Therefore, the error in total such that for
any w,
V W.Cn(W (w) WwLn(w)l2
A((1)
A1(1)
1
1
:
l|2
(By definition of [w( and Ca (1))
2m
2n
A.(N)
A((N)
0
0
N
((,()) 1.))
max Isten
N
< 2
C + E + E.AVPB + B,E,Gr).
(By the error accumulation results derived before)
Let Ct, Cr,, Cr denotes coefficients in front of E,, Er, enr respectively. Then holds
G = N..
C, = NP3 F + NPB.F..
30

---

In-Context Deep Learning via Transformer Models
Cp NN
Thus, to ensure | (w) - V VuLn(w)ll S, 6, we only need to select El, 6ri e as
2E q
2e
2e c 3C,
G 30,
3G1
Therefore, we only need to pick the last MLP layer MLP2N+4 such that it maps
MLP [may: (Vw)ll/ww
[ 7V
By assumption the map Projw, this is easy.
Finally, we analyze how many embedding dimensions of Transformers are needed to implement the above ICGD. Recall that
Je. U1 Rd w E IR (N PP E ((N-)) 7/ (N k 9e E IR 8 E IR (R 1)
Therefore, max( } embedding dimensions of Transformer are required to implement ICGD on deep models. Combining the above, we complete the proof.
Remark 5 (Modest Assumptions). Our assumptions remain modest. For example, we require that the loss function I(-), the activation funetion r(-), and its derivative r' (.) are C'-smmothnesss, Many settings meet these conditions, including those using the sigmoid activation function r(-) and the squared loss function.
C.8 Proof of Corollary 1.1
Corollary 3.1 (Corollary 1.1 Restated: Error for implementing ICGD on N-layer neural network). Fix L 2 1, under the same setting as Theorem 1, (2NN 4)T.-ly neural networks NN, approximates the true gradient descent trajectory (wenhtzo E Dv with the error accumulation
|| wanll2 S (1
where L denotes the Lipschitz constant of (w) within W.
First we introduce helper lemma.
Lemma 14 (Error for Approximating GD, Lemma G.1 of (Bai al., 2023)). Let W Rd is a convex bounded domain and Projw projects all vectors into W. Suppose f W R and Vf is Ly-Liscchitz W. Fix any >> 0, let sequences (w]) h20 Rd and (weDhi2o Fd are given by के w&D 0, then for all / 0,
w Projw (w V, (w (w +
|
me,
WOD = Projw 111 - VV (w
To show the convergence, we define the gradient mapping at w with step size 17 as,
Projw VL (w))
G Gh.,
"
Then if S L,,, for all I 2 1, convergence holds
min [-(--1]
||G ( < 1
G W (T))
8(ʃ(0) in new [(w) 10c².
n
Moreover, for any / 120 0, the error accumulation is
||
w I 1 nL,)e.
31

---

In-Context Deep Learning via Transformer Models
Lemma 14 shows Theorem leads to exponential error accumulation in the general case. Moreover, Lemma 14 also provides convergence of approximating GD. Then we proof Corollary 1.1.
Proof. For any small by Theorem 1, the neural network NN, implements each gradient descent step with error bounded by Then we simply apply Lemma 14 to complete the proof.
32

---

In-Context Deep Learning via Transformer Models
D Extension: Different Input and Output Dimensions
In this section, explore the ICGD Nllayer neural networks under the setting where the dimensions of input z, and label Vi can be different. Specifically, we consider our prompt datasets {(z4, v:) hee()) where E IR and yi E Rd*, We start with our new N-layer neural network.
Definition 10 (N Layer Neural Network). An L-Layer Neural Network comprises N 1 hidden layers and output layer, all constructed similarly. Let R -> IR be the activation function. For the hidden layers: for any E [n + 11,5 E [N 1], and & E [K], the output for the first layers w.r.t. input æ; e Rd, denoted by pred, (x;; E IR is defined as recursive form:
pred, (24 1)(A] == r(v F.),
v.),
and (+ = pred, pred, (z:; -));
where U1, IRd and Vin E R for j E : NN 1 are the k-th parameter vectors in the first layer and the j-th layer, respectively. For the output layer (N-th layer), the output for the first N layers (i.e the entire neural network) w.r.t. input a; RR,,, denoted by pred w, N) Rdv, is defined for any & E [dyl follows:
pred w, [)] r(v,, pred, (I,;; --::--11 1)),
where UN,, E R the k-th parameter vectors in the N-th layer and E Rld_+d,))+NN2221 containing all parameters in the neural network,
denotes the vector
T
W = 1i
3u
UN UN
1 Nay
Notice that our new N-layer neural network only modify the output layer compared to Definition 1. Intuitively, this results in minimal change in output, which allows our framework in Section 3.3 to function across varying input/output dimensions. Theoretically, we derive the explicit form of gradient (w))
Lemma 15 (Decomposition of One Gradient Descent Step). Fix any B,, 0. Suppose the empirical loss function (w) on data points ((x,, v.) (Ejn) is defined as
(w)
1 2n
Σ
((((w,
where e R& x IR IR is loss function,
where f(w, .22),1) is the output of N -lyayer neural networks (Definition 10) with modified output layer. Suppose closed domain W and projection function Projw (w) follows (3.4). Let A,(),,(()), Re(j), V, be as defined in Definition 2 (with modified dimensions), then the explicit form of gradient VL, (w) becomes
A.(1)
(L) 2m ==1
[A.(N)
where A,() denote the derivative of é (p: ()) with respect to the parameters in the j-th layer,
A,0) = (R - 1). VN
Ra(j 1). [ 2 pi 1) 1)T
j N j == N.
(Ra(N-1) La,xd, PPNNN - 1) 1)T))T (l).. )T, Op:(N)
Proof. Simply follow the proof of Lemma 1. We show the different terms compared to Definition 2:
33

---

In-Context Deep Learning via Transformer Models
Let Dj E R denote the total number of parameters in the first j layers.
0,
j = 0
dg
j=1
Dj
8 dgK, (N 2) + (dz + dy)
(j -1)
2 2 -1
j j=N, N,
The intermediate term (N - 1),
(N 1) diag(r' (x), (() (.()) gdxdd
The parameters matrices of the first and the last layers:
T
E Rxdds j=1
U1
VJ
E ERA j= N.
UN1
UNug
Thus we complete the proof.
Lemma 15 shows that the explicit form of gradient V (w)) holds the same structure as Lemma 1. Therefore, it is simple to follow our framework in Section 3.3 to approximate V C, (w) term by term. Finally, we introduce the generalized version of main result Theorem
Theorem (In-Context Gradient Descent on N N-layer NNs), Fix any B., e > 0, 2 1. For any input sequences takes from (2.1), where {<,, ve)n) and z; _d= and y E Rdy, their exist upper bounds Ba, B, such that for any E [n]. llıal12 S By ||æ|||2 S B,- Assume functions r(t), r' (t) and utt, [)] are L,, L,,, Li-Lisschtzz continuous. Suppose W
closed domain such that for any j E [N -N--1 1] and [K].
W { w [1ya] PN ||05 ll2 B.)
and Projw project into bounded domain W. Assume Projw ML.P, for some MLP layer with hidden dimension D., parameters ||0|| < C.. If functions r(t) and u(t, [k) are C*-smoothnesss, then for any 0, there exists transformer model NNg with (2N 4) hidden layers consists of L neural network blocks TF EWML, TF?
NNg = TF o EWML o TF3 o... TF o EWML TF.
such that the heads number M parameter dimensions D¹, and the parameter norms By suffice
M < EŌ
D' O'K2)) + DEN
max By I xO(n) + Cw +1, (2(
1E[[2N+1))
(E(2N+1)L)
where Ö(-) hides the constants that depend on d, K, N, the radius parameters Bg. By: B, and the smoothness of and And this neural network such that for any input sequences ( take from (2.1), (No(H(()) implements steps in-context gradient descent on risk (w) follows Lemma 15: For every I E [L]. the (2N + 4)/-th layer outputs h3 ((2N+4))) [ 04: ((); 0; 1; t.] for every i E [n + 1]. and approximation gradients wf()) such that
Projw (w-(1) -
1/ (n + ((11)),
T ( 0,
where ||el t7e is an error term.
34

---

In-Context Deep Learning via Transformer Models
E Extension: Softmax Transformer
In this part, demonstrate the existence of pretrained Softmax transformers capable of implementing ICGD on an N-layer neural network. First, we introduce our main technique: the universal approximation property of softmax transformers in Appendix E.1. Then, we prove the existence of pretrained softmax transformers that implement ICGD on N N-layer neural networks in Appendix E.2.
E.1 Axillary Lemma: Universal Approximation of Softmax Transformer
Softmax- ttention Layer. We replace modified normalized ReLU activation o/n in ReLU attention layer (Definition 7) by standard softmax. Thus, for any input sequence H R Dxn single head attention layer outputs
Atin (H) H 1 (VH) Softmax (KH))
(E.1)
where W Q, K, V E DxD E RRd are the weight matrices. Then we introduce the softmax transformer block, which consists of two feed-forward neural network layers and a single-head elf-attention layer with the softmax function.
Definition 11 (Transformer Block Tsoftmax)- For any input sequences H E RR Dxn, let FF(H) = H + W22 ReI + b, + 1, be the Feed-Forward layer, where d' is hidden dimensions, W1 E R xD W2 E Dxd' b E RR, and b2 E IRd. We configure a transformer block with Softmax-attention layer as Tsoftmax == (FFo Attne RdxL - m J.
Universal Approximation of Softmax-TTanstorme Transformer. We show the universal approximation theorem for Transformer blocks (Definition 11). Specifically, Transformer blocks Tsoftmax universal approximators for continuous permutation equivariant functions on bounded domain.
Lemma 16 (Universal Approximation of Tsoftmax) Let f(-) R -> d be any L-Lipschitz permutation equivariant function supported on [0, J... We denote the discrete input domain of [0, Bz] ] by grid GD with granularity D E N defined as GD = {B./ D, 2Bz/D. D..... B. d C RRdx, For any > O, there exists transformer network fsottmax E TSoftmax. such that for any Z E [0, B,] B,]dxn, approximate f(Z) as:
|| (Z) (Z) f(Z)l22 S
Proof Sketch. First, we use piece-wise constant function to approximate f and derive an upper bound based on its L-Lipschitz property. Next, we demonstrate how the feed-forward neural network Jz(F) quantizes the continuous input domain into the discrete domain G through multiple-step function, using ReLU functions create piece-wise linear approximation. Then, we apply the self-attention layer ((SA) on JiFF) establishing a bounded output region for (SA) F(FF), Finally, employ second feed-forward network ((FF) to predict fsoitmax (Z) and assess the approximation error relative to the actual output / ((Z). See Appendix E.4 for detailed proof.
E.2 In-Context Gradient Descent with Softmax Transformer
In-Context Gradient Descent with Softmax Transformer. By applying universal approximation theory (Lemma 16), we now illustrate how to use Transformer block Toottmax (Definition 11) and MLP layers (Definition 8) to implement ICGD on general risk function C,.ww) (w)
Theorem (Theorem Restated: In- In-Context Gradient Descent on General Risk Function). Fix any B M, 0, 1. For any input sequences takes from (2.1), their exist upper bounds Bz, By such that for any i E [n]. || [[Villaxx By ||ax||max < B,. Suppose W is closed domain such that || < B,, and Projw project w into bounded domain W. Assume Projw = MLP, for some MLP layer. Define I(w, I4, 34) as a loss function with L-Lipschitz gradient. Let (w: = Em Eww, x; y.) denote the empirical loss function, then there exists a
NN2, such that for any input sequences (0), take from (2. 1). NNo ())) implements I steps in-context gradient descent on (w)): For every / [L]. the 4/-th layer outputs h; (4)) = [xs; yes ()) O; 1;4:] for every E [n + 1]. and approximation gradients (7) such that
( Proj (w Cn (111)) (0) 0.
where [e(l22)||2 5 t7e is an error term.
35

---

In-Context Deep Learning via Transformer Models
E.3
Proof of Theorem 2
Proof of Theorem 2. We only need to construct 4 layers transformer capable of implementing single step gradient descent. With out loss of generality, we assume w E R Pw, Recall that the input sequences H E R Dx(x+1) takes form
w
F1
22
Ta+1 0
=
111
1/2
E R
Dx(x))))
qi =
111
E C D- (d-1)
(E.2)
91
Q2
On
4n+1
Let function F:R.xM Pxn +
Dxn output
w
17 (w)
f(H) =
1
y1
22
1/
In
Yn
In+i
9: =
0
1 E R Dd++1)
q
Q2
On
gn+1
By Lemma 16, for any K 0, there exists a transformer network fsoftmax Tsuftmax: such that for any in- put H [-B, B dxL, we have [[SSootmaxx (If) s(ff)l2 I Therefore, by the equivalence of matrix norms, If (H) f() I holds without loss of generality. Above B max By. Bw 1} denotes the up- per bound for every elements in H. Thus, we obtain w from the identical position of w in fsottmax (H) Suppose we choose
VDw then it holds
10 (w (/),lll v Dw ||
(w VVV
fsottuax f(H)) luax
Dw
VDw
S..
Finally, by our assumption, there exists an MLP layer such that for any i [n 1J, maps
w - VL,, 0; 1:4]
MLP
[aav Projw (w - v/ [V,CC 0;
Therefore, four-layer transformer /Sottmax MLP is capable of implementing one-step gradient descent through ICL. As direct corollary, there exist 4T.-layer transformer consists of identical blocks fsoftmax MT.P to approximate steps gradient descent algorithm. Each block approximates one-step gradient descent algorithm general risk function (w:)
E.4 Proof of Lemma 16
In this section, we introduce a helper lemma Lemma to prove Lemma 16. At the beginning, we assume all input sequences are separated by a certain distance.
Definition 12 (Token-wise Separateness, Definition of (Kajitsuka & Sato, 2024)). Let N 2 and z(),.. Z(N) E Rdxn be input sequences. Then, z(1) z(),..., Z(N) are called token-wise (Fmin: Tmax, 8) -eeparate the following three conditions hold.
For any i [N] and & E [n]. 7z,0)
Mmin holds.
2
For any i [N] and l e [m]. (
| Tmax holds.
For any í.j E [N] and k, 1 E [n]] with 2 + 2!!i |. 2.10) || > 8 holds.
Note that we refer to (11), (N) token-wiss (Fmax: c) -epparated instead if the sequences satisfy the last two conditions.
Then we introduce the definition of contextual mapping. Intuitively, a contextual mapping can provide every input sequence with unique id, which enables us to construct approximation for labels.
36

---

In-Context Deep Learning via Transformer Models
Definition 13 (Contextual mapping, Definition of (Kajitsuka & Sato, 2024)). Let input sequences , 2(m) E Rddx,
Then, map q Rdx= + Rdx= is called an (r, ccontexttll mapping the following two conditions hold:
For any i [M] and k E [m]. | (Z
holds.
For any i.j [M] and k, I E [n], if Z * z.p then |9 (z(").. 100--(ZZ6) (Z 6 holds.
In partieular, ()() for E [N] called context of 2
Next, we show that softmax-based l-layer attention block with low-rank weight matrices is contextual mapping for almost all input sequences.
Lemma 17 (Softmax attention contextual mapping. Theorem 2 of (Kajitsuka & Sato, 2024)), Let 2(1,..., ".... 2(N) E RRdx be input sequences with no duplicate word token in each sequence, that is,
- 7!77
for any E [M] and k, E [n]. Also assume that Z ..... 2 are token-wise (rmin- Tmax: c) separated. Then, there exist weight matrices W(O) E Iadxs and V, K, Q E R such that the ranks of V, K and Q are all 1, and 1-layer single head attention with softmax, ie., ((SA) with h is (r, ccontexuall mapping for the input sequences 2(1), Z(N) Rddx with and 8 defined by
C
T mmax 4
2(lg n)²2² r
cxp (-(12) +1)
1)4 (2 - 3) dr max
2 max
(IV)+ (2 log n 3)
4ermin
Applying Lemma 17, extends Proposition of (Kajitsuka & Sato, 2024) to our Lemma 16'. We provide explicit upper bound of error [Softmaa ( Z) f(Z)|l2 and analysis with function f of broader supported domain.
Lemma 18 (Lemma 16 Restated: Universal Approximation of Tsottmax). Let f()) == Rdxx + Rdxn be any I--ipschitz jdxn permutation equivariant function supported on [0, 3___,x,, We denote the discrete input domain of [0, B,]ax= by grid GD with granularity D N defined Gn {B,/ / D, 2B./D,.... B, B dxx C Rdx=a For any > 0, there exists transformer network fsottmax Tsoftmax (Definition I)), such that for any Z E [0, B, dxxn approximate f(Z) as:
|| (Z) f(Z)ll2 S n.
Proof. We begin our 3-step proof.
Approximation of f by piece-wise constant function. Since f is continuous function compact set, f has maximum and minimum values on the domain. By scaling with J ((FF) and F((F)) f assumed be normalized: for any Z ddxn [0, B dxn
f(z)= = 0.
and for any Z E [0, dx
By S ((() s By
Let D E N be the granularity of a grid Gp:
B2 2B.
Cp = D D
B2 dxn C Rdxn
Thii extension builds on the results of (Hu et al., 2025a), which extend the rank-1 requirement any rank for attention weights. Additionally, Hu et al. (2025b) apply similar techniques to analyze the statistical rates diffusion transformers (DiTs).
37

---

In-Context Deep Learning via Transformer Models
where each coordinate only take discrete value Bz/D,2Ba ..... Bg- Now with continuous input Z, we approximate by using a piece-wise constant function f evaluating the nearest grid point I. of Z in the following way:
((2) - (ent
(E.3)
LEGn
Additionally I. [-111/ 0)dx, denote as Q(Z) L. Now bound the piece-wise constant approximation error ||ʃ fI] follows.
Define set PD ( LLL+ E Gp). is set of regions of size ( )dn whose vertexes are the points in GD.
For any subset U & Pp. the maximal difference of f and f in this region is:
max ZEU
[f(Z) F(Z) = max ZEU
(2) f(Q(2))]l2
[l f(Z')
EU
S L. max EU
1 7112
(By f is T--Lipschizz function)
(Dz)
L. v dn
(D' j2
(Z, Z' are same w--wide (d n)-dimension U.)
dn
(E.4)
D
Quantization of input using F(FF). In the second step, we use F(FF) to quantize the continuous input domain into Gn. This process is achieved by multiple-step function, and use ReLU functions to approximate this multiple-step functions. This ReLU function be easily implemented by one-layer feed-forwart network.
First for any small 8 > 0 and E R, we construct a 6-approximated step function using ReLU functions:
2 <<0 0
OR
OR
0 S < 8B,.
(E.5)
D
3D
B.
D
where mne-hidden-layer feed-forward neural network is able to implement this. By shifting (E.5) by Ba, for any E [D- 1]. we have:
LB D
FR
IB_
5D
OR
B
LB_
8D
2<
tB, < 2
8B, + tB.
(E.6)
D
AD
B
dB, + 1B <
D
when 8 is small the above function approximates to step function:
2 < 18g
quant (2)
D
B-
1B
<
D
D
By adding up (E.6) at every E [D- 1), we have approximated multiple-step function
4B
OR
B
tB
5D
3D
(E.7)
D
38

---

In-Context Deep Learning via Transformer Models
quant (t) D (z)
(when 5 is small.)
t=0
= quant p (z)
0
2<0
B
05 8x
D
(E.8)
Ba
B. S
D
Note that the error of approximation at 2 here estimated as:
OR
R2
6D
OR
B₃
18
8D
quantp(z) 1
B₁
(F.9)
D
D
and for matrix Z E Rdxn
D-1
OR
Q(Z)
6D
OR
B, E
Q(2) 8D
quant p (2)
D
1=0
|xxxx
y2
(Z E Rdxx)
D
Bz
D
Subtract the last step function from (E.7) we get the desired result:
D-1
on
BB
on D
B2
1B. oD
B
B
(o
0R
(E.10)
8
5
This equation approximate quantization of input domain [0, B&] into {Ba/Da D.... addition to the quantization of input domain (0, Bz]. we add a penalty term input out of [0, Bs] in the following way:
B,} and making IR \0,0 B₂] In
(2
B²)
(z
Bx) 8
1 i
BI --
[ -
BBo
B
Bzon
(E.11)
8
- Ba: =50 0
penalty(z)
0. 0< Bz
Bai Ba <.
Both (0) and (E.11) can be realized by the one-layer feed-forward neural network. Also, straightforward show that generate both of them input Z Rdxx.
Combining both components together, the first ferw-forward neural network layer J(FF) approximates the following function F (FF) (Z):
d (Z) quant dxn (Z) + penalty (Zz,))
F(FF) = (FF))
(F.12)
Note how we generalize penalty(-) to multi-dimennional occasions in the above equation. Whenever an input sequence Z has one entry Ze,x out of [0, B_]dxxx jdxn we penalize the whole input sequence by adding a -B, to all entries. This makes all
39

---

In-Context Deep Learning via Transformer Models
entries of this quantization lower bounded by -dn
(E.12) quantizes inputs in [0, B,] ]xx with granularity D, while every element of the output is non-positive for inputs outside
[0, B,] 1,,]xxm In particular, the norm of the output is upper-bounded when every entry in Z is out of [0, Bz), this adds -dn B,
penalties to all entries;
max ZER
F (Z) = Vd. (-dnBB)))))
(One column is d- d-dimension.)
S dn vdB.,
(F.13)
for any & E [n]]
Estimating the Influence of Self-Attention F(SA), Define Gn C GD as:
G = = E GD Vk, [ [n], L, * L.1.
(E.14)
is a set of all the input sequences that don't have have identical tokens after quantization.
Within this set, the elements are at least f separated by the quantization. Thus Lemma 17 allows construct self-attention (SA) to be a contextual mapping for such input sequences.
Since when D is sufficiently large, originally different tokens will still be different after quantization. In this context, we omit p for simplicity.
From the proof of Lemma 17 in (Kajitsuka & Sato, 2024), follow their way to construct self-attention and have following equation:
| (A(AA)(2) (Z):k Zink
max KEE(n]
IZ.
(E.15)
4vdD
for any k E [n] and Z _dx__
Combining this upper-bound with (E.13) we have
(S)) o ((FF) (Z)
- -F(FF) (Z)
(Z).
max
[(FF)
(Z.,))l22
4vdD
kEin
x dnv dB.
(By (E.13)
4VdD
dn
(E.16)
4D
We show that if we take large enough D, every element of the output for Z E RRdx [0, BN dxn is upper-b by
F(SA) ((FF) (Z) <
B.
(vt [d], &: KE [n])-
(E.17)
4D
To show 17) holds, we consider the opposite occasion that there exists a a ((SA) F(FF) (Z) 2 Bz/4D. Then we divide the case into two sub cases:
1. The whole ((FF) (Z) receives less than 2 penalties. In this occasion, since every entry consists of two counterparts in
(Z) E [0, Bz] and aggregated with a penalty part Σd= E penalty (Zr,x) S (E.12): the quantization part quant dxn (Z)
-2B,, for every entry have (FF) (Z) B,,.
40

---

In-Context Deep Learning via Transformer Models
This yields that:
|| (84) o F(PP) (2), -F(FF) (2) 1l 2 | F((F) (Z) F(FF) (2) ta ll2
B
(-Bx)
4D
dn
Bar
(for a large enough D)
4D
thus we derive contradiction towards (E.16) from the assumption, proving it to be incorrect.
2. The whole Fl (PFF) (Z) receives only penalty. In this case all entries in Z is penalized by BB, and satisfies:
F(PF) (Z), E [-B,,
(E.18)
By (E.15), this further denotes:
((SA) Fs o ((FF) (Z).,t -F(FF) (Z)
1
[I (Z) ll2
max
(By (E.15))
4vdD
K'E(n)
1
Vd B
(By (E.18)
4vdD
B₂
(E.19)
4D
Yet by our assumption, there exists such an entry F(SA) o F(FF) (Z) 2 B,/4D, which since F(FF) (Z) S0. yields:
| ((( ((P (Z)..o -((F
(Z) A0
2
((SA)
J(PM)
(Z) -(PP) (Z) (Z)
Br
4D
4D
The final conclusion contradict the former result, suggesting the prerequisite to be fallacious.
Joining the incorrectness of the two sub-cases of the opposite occasion, we confirm the upper bound when input Z is outside
[0, _,]dx= in (E.17).
For the input Z inside [0, B, jxxn B_]dx=, we now show is lower-bounded by
F(
(SA) o ((FF) (Z),,
3B
(Vt E [d], is E [m]).
(E.20)
4D
By our construction, every entry Z in [0, By) satisfies:
(FFF) (Z), E
1 Bz).
B.
(E.21)
By (5.15):
|f (PF) (Z)., ((FFF) (Z)
max | (FF) (Z) 112
(By (E.15)
4 AVdD
K'Enn)

---

In-Context Deep Learning via Transformer Models
Vd x B
(d-dimensional vector with each entry maximum value B..)
4vdn
BR
(E.22)
4D
This yields:
(=(8)) o F1 (2)r, -F(PF)
[ o F(FF)) (2) -(FF) (Z),,
1
Brr
(F.23)
4D
Finally, we have:
f! x(FF)
(2),k
F1 FFF (2)
| o F((F) (2) - -F(PP) (2aa112
B₁ - | (SA)) x(( (2),, - -J(PP) (Z),, ll2
>
(By (E.21).)
D
B₁ B₁
(By (E.23)
4D
3B2
4D
Hence we finally finish the proof for the upper bound of ((SA), o F(FF) (Z),, for Z outside [o, B,] in (E.17) and lower bound for Z inside [0, By] in (E.20).
Approximation Error. Now, we can conclude our work by constructing the final feed-forward network F (FF)) It receives the output of the self-attention layer and maps the ones in &n C (3Bz/4D, oo)dx to the corresponding value of the target function, and the rest in (-o0, / 4D)dx= to O.
In order to adapt to the L2 norm, we use continuous and Lipschitz function to map the input Z to its targeted corresponding output (((()))
According piece-wise linear approximation, function F2 F(PP) corresponding fL),, and for arbitrary input Z, its output suffices:
(FF) exists such that for any input L Gn. it maps it to
((FF) (Z) E
min
J(L),
max J(L)).
(E.24)
11- S
Bg
I 2 Mass S
Bg (SA)
Next we estimate the difference between J: (FF) o ((SA) x((FF) and F (FF)
. oF
(FF)
The difference is caused by the difference between ((()) and ((FF) By (E.9), this difference bounded by 1 in every dimension, for any input Z e dd
| (FF) (Z) (F [(FFF) <
VdnB. D
By (B.1
|| ( (2) - =1(5)) o (2)112
< || [SS)] F((FF) (2) - -F(PP) (Z) + (FFF)
(Z) - (2)12
+ +1 (z) ((A)) o F((FF) (7)122
(By triangle inequality)
42

---

In-Contey, Deep Learning via Transformer Models
ddnBz +2. 4D
vnBBg
(By All2 S
|| and (F.19))
In the section quantization of the input, used piece-wing wise linear functions (E.7) approximate piece-wise wisecesss--conttntt functions (E.8), this creates a deviation for the inputs the boundaries of the constant regions. Consider Z as one of these
inputs whose value deviated from F F(FF) F( ( (Q(Z)) Let f(I,) denote the value given to F² F(FF) F.S F((SA)
F F(FF)()) (Z). Because the deviation take the output to a grid most / D /nB,g/20 away from its original grid.
under the quantization of the output, ((L)) most deviate from its original output F F ((FF)) (Z) by the
distance of /dnB,/D D + /rB,/21 aggregated with times of the maximal distance within grid. They sum up to be:
|F! (51)) ((PF)_ -F(FF) ((A)) F( ((FF)
( 2VdnBB + VnBz
+2 VdnBr,
112 L.
2D
D
6v + VnBB
L
2D
Lastly, by condition we neglect the Gp GD part. This yields:
Fl
F(P o ((SA)
( =
Thus, adding up the errors yields:
If F2 ((F (((A)
(SA) F |
Iiff fff-f12+ + 11 1 (FF) F F(SA) F(FF)
(FF)
(By triangle inequality)
6VdnBz+ VnBr VdnB,
(By (E.4)
2D
L(8v dn + Br
2D
For any K > 0, we select large enough D, such that
LBs
(8v dn + vV)= S K.
2D
This completes the proof.
43

---

In-Context Deep Learning via Transformer Models
F
Experimental Details
In this section, conduct experiments to verify the capability of ICL to learn deep feed-forward neural networks. We conduct the experiments based on 3-layer NN, 4-layer NN and 6-layer NN using both Transformer based on the GPT-2 backbone.
and Softmax-
Experimental Objectives. Our objectives include the following three parts:
Objective 1. Validating the performance of ICL matches that of training N-layer networks, i.e., the results in Theorem 1, Theorem 4, and Theorem 5.
Objective Validating the ICL performance in scenarios where the testing distribution diverges from the pretraining one or where prompt lengths exceed those used in pretraining.
Objective 3. Validating the ICL performance in scenarios where the distribution of parameters in the N N-layer network diverges from that of the pretraining phase.
Objective 4. Validating that deeper transformer achieves better ICL performance, supporting the idea that scaling up the transformer enables it to perform more ICGD steps.
Computational Resource. We conduct all experiments using 1 NVIDIA A100 GPU with 80GB of memory. Our code based on the PyTorch implementation of the in-context learning for the transformer (Garg et al., 2022) https: //github.com/deee wwwo///sppiinnnooettteeoo -ott--eooteet---ninngg.
F.1 Experiments for Objectives 1 and 2
In this section, we conduct experiments to validate Objectives 1 and 2. We sample the input of feed-forward network E d from the Gaussian mixture distribution: ii N--2, La) + W2 (2, Ia), where W1, E I. We consider three kinds of network f R + R. (i) 3-layer NN, (ii) 4-layer NN, and (iii) 6-layer NN. We generate the true output by = fz) In our setting, we use
20.
Model Architecture. The sole difference between ReLU-Transfommer and in the attention layer. Both models comprise 12 transformer blocks, each with attention heads, and share the same hidden and MLP dimensions of 256.
is the activation function
Transformer Pretraining. We pretrain the ReLU-Transformer and TranssoomeTr based on the GPT-2 backbone. In our setting, we sample the pertaining data from ( Ia), i.e., w) = and 02 O. Following the pre-training method in (Garg et al., 2022), we use the batch size as 64. To construct each sample in a batch, we use the following steps (take the generation for the i-th sample as an example):
1. Initialize the parameters in fi with standard Gaussian distribution, i.e., (),
2. Generate n queries { (i.e., input of f.) from the Gaussian mixture model w (N(2,, Ia) + W2N Ia). Here we take n 51.
3. For each query Fi.j use Vi.j = fa (((..)) to calculate the true output.
This generates training sample for the transformer model with inputs
Ti, Vi,1, "* F1,50 V1,50: F1,51],
and training target
0, = [M1,1 9t.50 31,51]
We use the MSE loss hetween prediction and true value of o; The pretraining process iterates for 500k steps.
Testing Method. We generate samples similar to the pretraining process. The batch size 64, and the number batch 100, i.c., we have 6400 samples totally. For each sample, we extend the value n from 51 to 76 learn the performance of in-context learning when the prompt length is longer than used in pretraining. The input the model becomes
[i Vi.1, Fi,75 Ni,75, Z1,70]
44

---

In-Context Deep Learning via Transformer Models
We assess performance using the mean R-squared value for all 6400 samples.
Baseline. We use the 3-layer, 4-layer, and 6-layer feed-forward neural networks with 200 hidden dimensions baselines by training them with in-context examples. Specially, given testing sample (take the i-th sample as example), which R-1 and test query zi,k- We use {ma. Vi.j R-1 j=1 to train the network with MSE loss for 100 epochs. We select the highest R-squared value from each epoch the testing measure and calculate the average across all 6400 samples.
includes prompts {.a.g Vi.j |=
F.1.1 PERFORMANCE OF RELU TRANSFORMER.
We use four different Gaussian mixture distributions wI N (-2, + (2N Ia) for the testing data: (i) wi = 1, = (ii) W1 0.9, 4/ = 0.1, (iii) w1 0.7, w2 0.3, (iv) W1 0.5, W2 0.5. Here the distribution in the first setting matches the distribution in pretraining. We show the results in Figure 3.
1.0
1.0
1.0
Squared
0.5
0.5
R-Squared
0.8
0.0
NC-2.0
R-Squares
0.0
NI-2.R 0.1N2.1
NC-2.0
0.5
0.9N 0.N12.
09N1
0.4
0.3N22.
-1.0
0./NI-2.0 0.3N2
0,5M 0.5.2.0.
-1.0
0.70(-2.11 0.3N12.
0.5NI
2.0 0.5M2.
0.2
0.7M1-2.0- d.3M2.11
2.
0.5N2..0
-1.5
3-Layer NN
-1.5
4-Layer
-0.2
6-Layer NN
25 In-context Examples
50 In-context Examples
25 In-context Examples
75
(a) 3-Layer NN
(b) 4-Layer NN
(c) 6-Layer NN
Figure 3: Performance ICL in ReL.U-Transformr:: ICL learns 3-layer, 4-layer, and 6-layer NN and achieves R-squared values comparable those from training with prompt samples. The results also show the ICL performance declines as the testing distribution diverges from the pretraining one.
F.1.2 PERFORMANCE OF SOFTMAX TRANSFORMER.
We use four different Gaussian mixture distribution N Ia) + W2, Ia) for the testing data: (i) w1 = 1, 0. (ii) w1 = 0.9, W2 0.1, (iii) W1 = 0.7, w2 = 0.3, (iv) W1 0.5, w/ = 0.5. Here the distribution in the first setting matches the distribution in pretraining. We show the results in Figure 4.
1.0
1.0
1.0
0.8
0.8
R-Squared
0.6
Squared
0.5
R-Squared
0.6
AMOMA
0.0
0.2
0.4
0.9N 2. + 0 0
0.
M-2.0
0.5 + 0.5N22.0
0.3M2.1
-0.5
0.0
AI-2.11
0.9N1-2.0 0.1NI2.11
0.5N1-2.0 0.5N2.11
07M(-2.11
0.2
0.4
2.-20
0.9NL 2. D.IN(2.
0.5N(-2.1) 0.55M22.)
17N1-2,000.3W22
Layer
-1.0
Layer
0.0
6-Layer NN
0.4
25
50
75
0
25
50
75
25
50
75
In-context Examples
In-context Examples
In-context Examples
(a) 3-Layer NN
(b) 4-Layer NN
(c) 6-Layer NN
Figure 4: Performance of ICL in Softmax-Transformer ICL learns 3-layer, 4-layer, and 6-layer NN and achieves R-squared values comparable to those from training with prompt samples. The results also show the ICL performance declines as the testing distribution diverges from the pretraining one. Note that performance decreases when the prompt length exceeds the pretraining length (i.e., 50), a well-known issue (Dai et al., 2019; Anil et al., 2022). We believe this is due to the absolute positional encodings in GPT-2, as noted in (Zhang et al., 2024)
The results in Appendix F.1.1 and Appendix F.1.2 show that the performance of ICL in the transformer matches that of training N-layer networks, regardless of whether the prompt lengths are within or exceed those used in pretraining. Furthermore, the ICL performance declines the testing distribution diverges from the pretraining one.
F2 Experiments for Objective
In this section, we conduct experiments to validate Objective 3. For these experiments, we use testing data that identical to the training data, which follows a distribution of N (-,, We vary the distribution of parameters in the N-layer network. During the training process, we set the distribution as ((, 7). In the testing process, we examine different
45

---

In-Context Deep Learning via Transformer Models
distributions, including N(O, /). (-0055, I), and N(0.5, I). All other model hyperparameters and experimental details remain consistent with those described in Appendix F.1. We evaluate the ICL performance of both the ReL.U-Tranffommer and the performance in the transformer matches that of training N -llayer networks, regardless of whether the parameter distribution in the N-layer network diverges from that of the pretraining phase.
for 4-layer networks, shown in Figure 5 and Figure 6. The results demonstrate that the ICL
1.0
1.0
0.8
R.guureed
0.6
0.6
0.4
0.4
0.4
0.2
ICL of Transformer
0.2
ICL of Transformer
0.2
ICL of Tansformer
4-Layer NN (Param.
4-Layer iParam. MC-.5..0
4-Layer (Param.
0.0
0.0
0.0
25
50
75
25 In-context Examples
50
75
In-cantext Examples
In-context Examples
(a) Parameters ~ ()
(b) Parameters ~ (--0.,,
(c) Parameters ~ (0.5,
Figure 5: Performance of ICL Across Various N N-layer Network Parameter Distributions for the ReLU-Transformer ICL learns 4-layer NN and achieves R-squared values comparable those from training with prompt samples, even when the parameter distribution in the N-layer network during testing diverges from that in the pretraining phase (N(0, T)).
1.0
1.0
R-Squared
R-Squared
0.8
0.6
0.6
0.4
0.4
0.2
ICL Transformer
Transformer 4-Layer NN (Param. N N(-05, 01
0.2
of Transforer
4-Layer NN (Param. Noo.
4-Layer NN (Param. NIO.S,)
0.0
0.0
25
50
75
25
50
75
0
25 In-context Examples
50
75
In-context Examples
In-context Examples
(a) Parameters (0,
(b) Parameters N(-0.5,
(c) Parameters ~ N(
Figure Performance of ICL Across Various N-layer Network Parameter Distributions for the Softmax-Transtormn ICL learns 4-layer NN and achieves R-squared values comparable those from training with prompt samples, even when the parameter distribution in the N-layer network during testing diverges from that in the pretraining phase (N(O. 7)).
F.3 Experiments for Objective 4
In this section, we conduct experiments to validate Objective 4. For these experiments, we use testing data identical to the pertaining data from 2,-2,14 I2). We vary the number of layers in the transformer architecture, testing configurations with 4, 6, 8 and 10 layers. All other model hyperparameters and experimental details remain consistent with those described in Appendix F.1. We evaluate the ICL performance of both the ReLU-Tranfformer and the Softmax- Transformer with 15, 30, and 45 in-context examples, as shown in Figure 7. The results show that deeper transformer achieves better ICL performance, supporting the idea that scaling up the transformer enables it perform more ICGD steps.
46

---

In-Context Deep Learning via Transformer Models
0.91
0.88
0.87
0.90
R-Squared
R-Squared
0.86
0.89
0.85
0.88
0.84
15 In-context Examples
0.87
15 In-context Examples
0.83
30 In-context Examples
0.86
30 In-context Examples
0.82
45 In-context Examples
0.85
45 In-context Examples
8
10
4
8 Softmax-Transformer Layer Depth
6
10
ReLU-Transformer Layer Depth
(a) ReI.U--Trnsffrrmr
(b) Softmax-Transformer
Figure 7: Performance of ICL Across Varying Transformer Depths: We use the number of in-context examples as 15, 30, or 45 for both the Real--rrasformer and the Softmax-Tranne Transformer. The results show that deeper transformer achieves better ICL performance, supporting the idea that scaling up transformer enables to perform more ICGD steps.
47

---

In-Context Deep Learning via Transformer Models
G Application: ICL for Diffusion Score Approximation
In this part, we give an important application of our work, i.e., learn the score function of diffusion models by the in-context learning of transformer models. We give the preliminaries about score matching generative diffusion models in Appendix G.1. Then, we give the analysis for ICL to approximate the diffusion score function in Appendix G.2.
G.1 Score Matching Generative Difftusion Models
Diffusion Model. Let ze E Rd be initial data following target data distribution To ~ Po. In essence, a diffusion generative model consists of two stochastic process in R4:
forward process gradually add noise to the initial data (e.g., images): 20 - 21 + ZT.
A backward process gradually remove noise from pure noise: yr WT-1
yoo.
Importantly, the backward process is the reversed forward process, i.e., V/ R Jy-t for i E O...., This allows the backward process to reconstruct the initial data from noise, and hence generative. To achieve this time-reversal, diffusion model learns the reverse process by ensuring the backward conditional distributions mirror the forward ones. The most prevalent technique for aligning these conditional dynamics is through "score matching" - strategy training a model to match score function, i.e., the gradients of the log marginal density of the forward process (Song et al., 2020b;a; Vincent, 2011). To be precise, let Pt, Pe()) denote the distribution function and destiny function of z₁. The score function is given by V (oggp(() J In this work, we focus on leveraging the in-context learning (ICL) capability of transformers to emulate the score-matching training process.
Score Matching Loss. We introduce the basic setting of mcore-matching as follows?. To estimate the score function, we use the following loss to train a score network swr (-,t) () with parameters W:
min [' (1))
[[SS
log Pog dt, where (t)) is weight function,
(G.1)
and To is a small value for stabilizing training and preventing the score function from diverging. In practice, as (l unknown, we minimize the following equivalent loss (Vincent, 2011).
min W
(t)
Em,ino
[Ils (Je, t) - logp(1)/20)1)) dr,
(G.2)
To
where p[z)(zo) is distribution of o, conditioned on zo.
G.2 ICL for Score Approximation
We first give the problem setup about the ICI. for score approximation as the following:
Problem 3 (In-Context Learning (ICL) for Score Function V log logge()) pe(-)) Consider the score function log pr(-) for any
t 2 0. Given dataset D,, = {(1,, [e) where {z*]re(m) S Rd and yı V V Pt, (z,) C Rd (t, 2 0), and test input En+1, the goal of "ICL for Score Function" is to find transformer T to predict ya+1 based 2n++1 and the in-context dataset D,,, In essence, the desired transformer T serves the trained score network sw (-, (.)
To solve Problem 3, follow steps: (i) Approximate the diffusion score funetion V log pc(-) with multi-layer feed-forward network with ReLU activation functions under the given training dataset D,, (ii) Approximate the gradient descent used to train this network by the in-context learning of the Transformer until convergence, using the same training set D,, as the prompts of ICL.
For the first step, we follow the score approximation results based on a multi-layer feed-forward network with ReLU activation in (Chen et al., 2023), stated as next lemma.
Lemma 19 (Score Approximation by Feed-Forward Networks, Theorem of (Chen al., 2023)). Given approximation enror > 0, for any initial data distribution Po. there exist a multi-layer feed-forward network with ReLU activation,
R denotes distributional equivalence.
*Please also see Appendix A.1 and (Chen al., 2024; Chan et al., 2024; Yang et al., 2023) for overviews.
48

---

In-Context Deep Learning via Transformer Models
fw,,r, Rw Rd x R R. Then for any [To, T), we have [l, ., Pol (l)()))))))))))) S O(c).
With the approximation result, we reduce the Problem to Problem 2, where the loss function is (G.1). Following Theorem 1, we show that the in-context learning of transformer models can approximate the score function of diffusion model.
49


Reviews:
[{"review_id": "izm4zdN3jv", "reviewer": "Reviewer_wyVr", "summary": "This paper studies the expressive power of transformer models to simulate gradient descent on other architectures like N layer feedforward networks using in-context learning. The authors corroborate their study with experiments on synthetic datasets that show that the in-context learning performance of Transformers matches direct training of deep networks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1**: How would the bound vary if the transformer's size is allowed to increase beyond (2N+4)L? Can it be made polynomial, by making the transformer size grow exponential?\n\n**Response:**\n\nThank you for your question. The bound does not change if the transformer's size exceeds $(2N+4)L$, and the error cannot be reduced to a polynomial level. Here are some clarifications:\n\n* Each $(2N+4)$-layer transformer simulates one gradient descent (GD) step, and we stack such transformers $L$ times.\n\n* In non-convex optimization, each step’s trajectory depends on the previous steps. This dependency causes the exponential accumulation of the error.\n\n> **Q2**: Can some form of approximate self-attention or MLP layers be used in their place? How would the size of the constructed model vary then?\n\n**Response:**\n\nThank you for your insightful question. The answer is **No**. Here are some clarifications.\n\n* **Explicit construction of transformer**: We aim to provide an explicit transformer construction that simulates $L$ gradient descent steps. Although self-attention or MLP layers can approximate Hadamard products through their universal approximation capabilities, such approximations would result in non-explicit constructions.\n* **Specialized variant of self-attention**: EWML can be viewed as a specialized variant of self-attention that retains only the diagonal scores.\n\n---\n\nWe hope our responses address your concerns and look forward to further feedback.", "claims_and_evidence": "The authors show clean theoretical claims to show expressivity of transformer models. The authors extensively discuss the proof strategies for each theoretical claim. Furthermore, they conduct experiments in the appendix to support their theoretical claims.", "methods_and_evaluation_criteria": "The experiments in appendix have been conducted on toy datasets to corroborate their theoretical claims.", "theoretical_claims": "I glanced through the proofs of each theoretical claim. Even though I didn't go through specific details, the theoretical proofs and statements are reasonable. Furthermore, the authors clearly corroborate their theoretical observations with clean empirical experiments.", "experimental_designs_or_analyses": "The experiments in appendix have been conducted on toy datasets to corroborate their theoretical claims.", "relation_to_broader_scientific_literature": "The paper is relevant to the current interests of the scientific community. Exploring the strengths of in-context learning of transformer models has been a topic of interest for the last few years and this paper takes an important step towards this direction.", "essential_references_not_discussed": "The authors have extensively discussed the important references.", "other_strengths_and_weaknesses": "The main strength of this paper lies in its motivation to study the strengths of in-context learning abilities of real world transformer models. The paper first discusses the pros and cons of existing studies on similar topics and attempts to specialize the framework to understand expressive power for simulating multi-step gradient descent of N-layer feedforward networks.", "other_comments_or_suggestions": "I don't see any typos.", "questions_for_authors": "Please find my questions below.\n\n- In corollary 1.1, the error bound is exponential in the number of steps L. How would the bound vary if the transformer's size is allowed to increase beyond (2N+4)L? Can it be made polynomial, by making the transformer size grow exponential?\n \n- Theorem 1 requires Element-wise Multiplication Layers (EWMLs) to approximate Hadamard products. Can some form of approximate self-attention or MLP layers be used in their place? How would the size of the constructed model vary then?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z1Jf8e4F0Y", "reviewer": "Reviewer_Dau7", "summary": "The paper introduces an approach that harnesses the transformer's ability to emulate the in-context learning for training process of deep models. Its key contribution is the demonstration of a practical instance where a transformer is used to simulate the training process of a deep neural network. Furthermore, the paper appears to extend the work presented in \"Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection,\" which provides a solid theoretical analysis.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We greatly appreciate your attention to detail and recognition of our theoretical and experimental contributions! Your constructive comments and encouraging words are also highly appreciated!", "claims_and_evidence": "The paper poses the question, \"Is it possible to train one deep model using the in-context learning (ICL) capability of another foundation model?\" and answers it affirmatively. The paper offers theoretical guarantees for achieving approximation within any desired error margin and for the convergence of ICL-based gradient descent. The analysis centers on a practical setting involving Softmax-based transformers. The method is evaluated on synthetic datasets for 3-layer, 4-layer, and 6-layer neural networks, with results showing that ICL performance is on par with that of direct training.", "methods_and_evaluation_criteria": "The method was primarily motivated by theoretical considerations. The authors systematically formulate in-context learning (ICL), carefully defining the neural network's mathematical model. Their approach relies on a recursive application of the chain rule, approximating derivatives sequentially. The authors suggest optimizing a transformer to mimic in-context gradient descent across different configurations.", "theoretical_claims": "The paper's contribution mainly lies in its mathematical formulation and theoretical analysis. The ultimate goal is to demonstrate that bounds can be established for the approximation.", "experimental_designs_or_analyses": "The experiments primarily aim to validate the correctness of the theorems. Although they are not conducted on large neural networks, they are sufficient in my view.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "From my perspective the paper provides adequate citations.", "other_strengths_and_weaknesses": "Overall, the paper is clearly written and easy to follow. The proposed method has been validated through both theoretical analysis and experimental results.", "other_comments_or_suggestions": "Same as previous comments.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xbNNy87NY5", "reviewer": "Reviewer_aV3f", "summary": "This paper shows that one can construct weights of a ReLU-activation transformer that can simulate L steps of gradient descent on an N-layer ReLU network using in-context examples.\n\n## update after rebuttal\n\nWhile I have concerns with clarity and usefulness of the construction, I appreciate the author's honesty and thus will stick to my raised score of 2.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Thanks for your detailed review. We have revised our draft and addressed all concerns. The revised version (changes marked in BLUE) is available in this [anonymous Dropbox folder](https://www.dropbox.com/scl/fo/cv5unyj4vew76op0xjftj/ALWvpjkaQ-_83Je7h0APsgs?rlkey=z80f15tgc59a3zl2m7owaklbe&st=05cr7uzt&dl=0).**\n\n> **Q1**: Do trained transformers actually achieve the construction?\n\n**Response:**\n\nThank you for the question. **The trained transformers do not always achieve the construction.** We apologize for any confusion caused, and acknowledge the limitation you mentioned. However, this limitation does not affect the primary contributions of our work. Here are some clarifications.\n\nThe main contribution of this paper is to provide an **explicit construction** demonstrating the **existence** of a transformer capable of simulating gradient descent (GD) for $N$-layer neural networks via in-context learning (ICL). The experimental design empirically validates that the trained transformer can indeed simulate GD steps for $N$-layer neural networks, supporting our theoretical results on the existence of such a transformer. Although there is a discrepancy between the theoretically constructed transformer and the empirically used one, this difference does not weaken the central point of our work, which is establishing the theoretical existence of such a transformer.\n\nWe acknowledge the limitation you mentioned and have incorporated it into the revised version: `line 744-748`.\n\n> **C1**: I think the paper suffers a lot from reduced clarity. It seems that most of the content of the paper is in the appendix, including empirical comparisons as well as discussion, which to me makes for a less clear read. I would suggest the authors rewrite the main body of text to include more intuitions and takeaways (e.g. when some assumptions may hold in practice rather than just in theory).\n\n**Response:**\n\nThank you for your thoughtful suggestions. We agree that the main text can be better organized to improve readability and self-containment. In response, we have made the following modifications in the revised version:\n\n* **The results on softmax attention and the experimental results have been incorporated into the main body**—softmax results in Section 3.4 and experimental results in Section 4.\n\n* **The proof sketches have been moved to Appendix D**, just before the detailed proof of the lemmas.\n\n* **We include a remark concerning the practicality of our assumptions (Remark 4: `line 1577-1579`)**. Our assumptions remain modest. For example, we require that the loss function $l$, the activation function $r$, and its derivative be $C^4$-smooth. This condition is met by numerous network architectures, including those using the sigmoid activation function $r$ and the squared loss function $l$.\n\n---\n\nWe hope our responses address your concerns and look forward to further feedback.", "claims_and_evidence": "The theory appears sound but the empirics are lacking.\n\n> Experiments. We support our theory with numerical validations in Appendix G\n\nAs far as I can tell, the experiments in Appendix G have nothing to do with constructing a transformer? Instead, they follow the more classical setup (of e.g., Garg et al) in seeing if a transformer can infer the function f in context. No analogy is drawn to whether or not the solution found by the transformer is similar to the construction shown in the main text, beyond the fact that on certain evaluation data the performance matches doing GD on a ReLU network (but this could also just be performance being \"good\" rather than the \"same\").", "methods_and_evaluation_criteria": "See above.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "The paper appears to cite relevant literature. I'm not sure how relevant the finding is/if such a construction is actually of interest to the broader community.", "essential_references_not_discussed": "Not to my knowledge.", "other_strengths_and_weaknesses": "I think the paper suffers a lot from reduced clarity. It seems that most of the content of the paper is in the appendix, including empirical comparisons as well as discussion, which to me makes for a less clear read. I would suggest the authors re-write the main body of text to include more intuitions and takeaways (e.g. when some assumptions may hold in practice rather than just in theory).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Do trained transformers actually achieve the construction?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CHHk62qPAb", "reviewer": "Reviewer_Gaqd", "summary": "This paper studies the representation power of the transformer model in performing in-context learning (ICL). In particular, this paper focuses on the implementation of in-context gradient descent with respect a general deep neural network. The proposed construction is flexible that it can either use a ReLU attention or the more common softmax attention, and it is efficient that it uses $O(N)$ for each gradient update step of a $N$-layer network.\n\nIn my opinion, the contributions of this paper is significant because it works for an arbitrary depth neural network and can be implemented using a softmax attention. The later point is especially valuable because much of the literature on ICL cannot handle the softmax activation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Thanks for your detailed review. We have revised our draft (changes marked in BLUE) in this [anonymous Dropbox folder](https://www.dropbox.com/scl/fo/cv5unyj4vew76op0xjftj/ALWvpjkaQ-_83Je7h0APsgs?rlkey=z80f15tgc59a3zl2m7owaklbe&st=05cr7uzt&dl=0).**\n\n> **C1**: Reconstruction of the Manuscript.\n\n**Response**\n\nThank you for the suggestions. We agree with this. In response, we have made the following modifications in the revised version:\n\n* We have moved the definitions of transformers and attention into the introduction while keeping the full standard definition in Appendix D due to space constraints.\n\n* The results on softmax attention and the experimental results have been incorporated into the main body—softmax results in Section 3.4 and experimental results in Section 4.\n\n* To improve clarity, we simplify Problem 2 and clearly state that it describes an MLP with $N$ layers.\n\n* The proof sketches in Section 3.3 have been moved to Appendix D, just before the detailed proof of the lemmas.\n\n* The complex constants in Lemma 6 have been shifted to Appendix D to improve readability.\n\n* While we acknowledge that Definition 1 may appear verbose, it establishes the notation used throughout the paper, particularly in the proofs, ensuring precision and consistency. Therefore, we have retained it as is.\n\nAdditionally, we appreciate the suggestion to include a **visualization of the backpropagation process**. We have included it in Appendix A.2 of the revised version. We also redisplay it **[anonymous figure](https://imgur.com/a/jfsYTQI)**.\n\n> **C2**: Difference with Standard Transformer\n\n**Response**\n\nThank you for your suggestions. We have added this to the last point of limitations (`line 749-754`) and restate it here:\n\nThere are two minor differences between the transformer used in the theoretical analysis and a standard transformer: (i) The transformer used in the theoretical analysis incorporates an element-wise multiplication layer, a specialized variant of self-attention that retains only the diagonal score and allows efficient implementation. (ii) It does not alternate self-attention and MLP layers. We emphasize that this also qualifies as a standard transformer because we view either an attention or an MLP layer as equivalent to an attention plus MLP layer due to the residual connections.\n\n> **C3** Example of Explicit Construction of ReLU\n\n**Response**\n\nThank you for your insightful comments. Here, we provide clarifications on why the constants can be explicitly instantiated and illustrate this using a simple example.\n\nThe key reason is that the function approximated by the sum of ReLUs is relatively simple in our context, such as the Sigmoid activation function. For such simple functions, it is straightforward to derive an explicit construction. \n\nHere, we take the Sigmoid activation function as an example and propose one explicit construction method. Let $r(z)$ denote the Sigmoid function.\n\n1. **Segment the input domain.** For example, divide the domain $[-10, 10]$ smaller intervals such as $[-10, -9], [-9, -8], \\dots, [9, 10]$.\n\n2. **Approximate each segment locally using a linear function via linear interpolation.** For instance, in the domain $[9,10]$, approximate $r(z)$ using a linear function $a_1 z + c_1$, where $a_1$ and $c_1$ are calculated as follows: (i) $a_1 = (r(10)-r(9)) / (10-9)$. (ii) $c_1 = r(9) - a_1 * 9$.\n\n3. **Approximate linear function $a_1 z + c_1 (z \\in [9,10])$ using a sum of ReLU terms.** This step involves two substeps, which are straightforward to implement: (i) Approximate the indicator function for $z \\in [9,10]$ using a sum of ReLU terms. (ii) Approximate the constant $c_1$ using the sum of ReLU. This is because bias terms are not included in the sum of ReLU terms in Definition 4. The bias term $c_1$​ must be approximated using an additional sum of ReLU terms.\n\n4. **Combine all the sum of ReLU approximators across all segments.** Finally, integrate the approximations for all segments to construct the complete approximation.\n\n5. **Estimation of the parameters in Definition 4.** $\\epsilon_{approx} = 0.625$, $R=10$, $H=80$, and $C=25$.\n\nFurthermore, to achieve higher precision in the approximation, it is sufficient to use finer segmentations.\n\n> **Q1**: Difference of Dimensions\n\n**Response**\n\nThank you for your question. We apologize for any confusion caused. Here are some clarifications:\n\n* In Section 2.(i), the value of $D$ is derived from prior research on in-context learning (ICL) for functions, indicated as $D=\\Theta(d)$. This basis does not pertain to our focus on models, such as an $N$-layer neural network. We have removed these parts in the revised version to make things easier to read.\n\n* Regarding the zero values in the rows of q_i, these primarily consist of intermediate terms such as $\\bar{p}_i(j), \\bar{r}^{‘}_i(j), \\bar{s}_t(j)$, which are essential for calculating partial derivatives and contain gradient information. These rows are initially set to zero.", "claims_and_evidence": "Although this work is a substantial addition to the literature, it is plagued by very poor readability. I believe that some major restructuring would be necessary before publication.\n\n1. Several parts of the Appendix absolutely should be in the main body. The paper in its current form does not have a self-contained main body. The definition of transformer and attention have to be part of the intro. The results on softmax attention should also be discussed in the main text because this is in my opinion the most interesting result of the paper. And some of the experiments need to be in the main body because they are mentioned in the abstract. My understanding is that the readers are expected to understand the key arguments of the paper without checking the appendix, and this paper currently falls short.\n2. The paper is overly verbose and many of the details are disruptive to the flow without contributing much to the exposition. For example, \"Problem 1\" and \"Problem 2\" are basically the same thing. Definition 1 is just an over-complicated description of a MLP. And the \"proof sketch\" in Section 3.3 don't need this level of precision for a general audience -- they are simply too dense to parse. I think the details of Section 3.3 (especially those constant) can be striped down and moved to the appendix.\n3. I think an illustration would be very helpful. For example, one can make a diagram marking all of the partial derivatives in back-propagation of the neural network and then highlight which layer of the transformer model implements each of these partial derivatives.\n\nGiven the amount of \"fat\" in this paper, I think my proposed changes can fit into the page limit without too much efforts.\n\nAnd some more technical concerns:\n1. While my first impression is that the proposed construction is for a standard transformer architecture, this is actually not the case. The authors introduced a *piecewise multiplication layer,\" which is fine because it can be implemented efficiently in code. Also, the attention and MLP layers are not stacked in alternating orders. The authors should highlight these differences rather burying them under a wall of texts.\n2. In Lemmas 2-4, the constants for ReLU approximations are not explicitly instantiated. I think it would be very helpful if the author could gives the values of these parameters for a simple example such as ReLU network with quadratic loss.\n\nLastly, I really like that Appendix B.3 clearly discussed the limitations of this paper. I appreciate the authors' honesty.\n\nIn light of these observations, I *do not* recommend acceptance of this paper as-is. However, I do think the overall ideas in this paper can be valuable to the community and I hope the authors can take account of my feedback to improve this paper. In particular, I want to see concrete plans from the authors on how to address my concerns so the revised paper is more accessible, but I fully understand that I am asking for a lot of changes and it is okay if the author do not have enough time during the rebuttal period to finish the revision. I am happy to raise my score if my concerns are adequately addressed.", "methods_and_evaluation_criteria": "n/a", "theoretical_claims": "The theoretical derivations of the paper seem to be precise and the authors devoted significant efforts into tracking the various algebraic blobs. So I have confidence that the results do not have any fatal flaws.\n\nI have one question: in the second point Section B.2 seems to imply that the intermediate computations on the partial derivatives are passed down through the forward pass of the transformer layers, but Section 2.i implies that the tokens' embedding dimension is only of order $\\Theta(d)$. Can you clarify the differences and if possible, clearly state those values in the revision? Furthermore, it would be useful to discuss what values were written into the rows containing $q_i$'s, where are initialized to be 0.", "experimental_designs_or_analyses": "The experiments are clear and validates the main theoretical claims of this paper. I have no complaints about them.", "relation_to_broader_scientific_literature": "see above", "essential_references_not_discussed": "The citations in this paper is sufficiently comprehensive.", "other_strengths_and_weaknesses": "see above", "other_comments_or_suggestions": "see above", "questions_for_authors": "see above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The main contribution of the paper is to show that transformers can simulate gradient descent on MLPs within a single forward pass. The idea of using one deep neural network to simulate the training dynamics of another is an interesting and timely topic. I read the paper carefully to better understand the contribution, but I believe the writing can be significantly improved. For example, I found the proof of Theorem 6—which presents one of the core results—difficult to follow. In particular, it is unclear why the transformer parameters are **assumed to be independent** of the input when invoking the universal approximation theorem. One of my central questions is why the authors do not simply invoke the universal approximation theorem to show that the transformer can represent global minimizers of the target MLP. Despite these concerns, and in light of the overall positive reviews, I vote for acceptance. However, I strongly encourage the authors to improve the clarity of their presentation and elaborate on the issues mentioned above.
Decision: Accept (poster)