SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Above, the parameter  $\pmb{\theta} = (\pmb{\theta}_{\mathrm{mattn}}^{(1:L)},\pmb{\theta}_{\mathrm{mlp}}^{(1:L)})$  is the parameter consisting of the attention layers  $\pmb{\theta}_{\mathrm{mattn}}^{(\ell)} = \{(\mathbf{V}_m^{(\ell)},\mathbf{Q}_m^{(\ell)},\mathbf{K}_m^{(\ell)})\}_{m\in [M^{(\ell)}]}\subset \mathbb{R}^{D\times D}$  and the MLP layers  $\pmb{\theta}_{\mathrm{mlp}}^{(\ell)} = (\mathbf{W}_1^{(\ell)},\mathbf{W}_2^{(\ell)})\in \mathbb{R}^{D^{(\ell)}\times D}\times \mathbb{R}^{D\times D^{(\ell)}}$ . We will frequently consider "attention-only" decoder-based transformers with  $\mathbf{W}_1^{(\ell)},\mathbf{W}_2^{(\ell)} = \mathbf{0}$ , which we denote as  $\mathrm{DTF}_{\pmb{\theta}}^0 (\cdot)$  for shorthand, with  $\pmb{\theta} = \pmb{\theta}^{(1:L)}\coloneqq \pmb{\theta}_{\mathrm{mattn}}^{(1:L)}$ .
We also use (2) to define the norm of  $\mathrm{DTF}_{\theta}$ .
C.2 In-context learning with decoder-based transformers
We consider using decoder-based TFs to perform ICL. We encode  $(\mathcal{D},\mathbf{x}_{N + 1})$ , which follows the generating rule as described in Section 2.2, into an input sequence  $\mathbf{H}\in \mathbb{R}^{D\times (2N + 1)}$ . In our theory, we use the following format, where the first two rows contain  $(\mathcal{D},\mathbf{x}_{N + 1})$  which alternating between  $[\mathbf{x}_i;0]\in \mathbb{R}^{d + 1}$  and  $[\mathbf{0}_{d\times 1};y_i]\in \mathbb{R}^{d + 1}$  (the same setup as adopted in [31, 2]); The third row contains fixed vectors  $\{\mathbf{p}_i\}_{i\in [N + 1]}$  with ones, zeros, the example index, and indicator for being the covariate token (similar to a positional encoding vector):
$$
\mathbf {H} = \left[ \begin{array}{c c c c c c} \mathbf {x} _ {1} & \mathbf {0} & \dots & \mathbf {x} _ {N} & \mathbf {0} & \mathbf {x} _ {N + 1} \\ 0 & y _ {1} & \dots & 0 & y _ {N} & 0 \\ \mathbf {p} _ {1} & \mathbf {p} _ {2} & \dots & \mathbf {p} _ {2 N - 1} & \mathbf {p} _ {2 N} & \mathbf {p} _ {2 N + 1} \end{array} \right], \quad \mathbf {p} _ {i} := \left[ \begin{array}{c} \mathbf {0} _ {D - (d + 4)} \\ \lceil i / 2 \rceil \\ 1 \\ \mathrm {m o d} (i + 1, 2) \end{array} \right] \in \mathbb {R} ^ {D - (d + 1)}. \tag {11}
$$
(11) is different from out input format (3) for encoder-based TFs. The main difference is that  $(\mathbf{x}_i, y_i)$  are in different tokens in (11), whereas  $(\mathbf{x}_i, y_i)$  are in the same token in (3). The reason for the former (i.e., different tokens in decoder) is that we want to avoid every  $[x_i; 0]$  token seeing the information of  $y_i$ , since we will evaluate the loss at every token. The reason for the latter (i.e., the same token in encoder) is for presentation convenience: since we only evaluate the loss at the last token, it is not necessary to alternate between  $[x_i; 0]$  and  $[0; y_i]$  to avoid information leakage.
We then feed  $\mathbf{H}$  into a decoder TF to obtain the output  $\widetilde{\mathbf{H}} = \mathrm{DTF}_{\boldsymbol{\theta}}(\mathbf{H}) \in \mathbb{R}^{D \times (2N + 1)}$  with the same shape, and read out the prediction  $\widehat{y}_{N + 1}$  from the  $(d + 1, 2N + 1)$ -th entry of  $\widetilde{\mathbf{H}} = [\widetilde{\mathbf{h}}_i]_{i \in [2N + 1]}$  (the entry corresponding to the last missing test label):  $\widehat{y}_{N + 1} = \mathrm{read}_y(\widetilde{\mathbf{H}}) \coloneqq (\widetilde{\mathbf{h}}_{2N + 1})_{d + 1}$ . The goal is to predict  $\widehat{y}_{N + 1}$  that is close to  $y_{N + 1} \sim \mathsf{P}_{y|\mathbf{x}_{N + 1}}$  measured by proper losses.
The benefit of using the decoder architecture is that, during the pre-training phase, one can construct the training loss function by using all the predictions  $\{\widehat{y}_j\}_{j\in [N + 1]}$ , where  $\widehat{y}_j$  gives the  $(d + 1,2j - 1)$ -th entry of  $\widetilde{\mathbf{H}} = [\widetilde{\mathbf{h}}_i]_{i\in [2N + 1]}$  for each  $j\in [N + 1]$  (the entry corresponding to the missing test label of the  $2j - 1$ th token):  $\widehat{y}_j = \mathrm{read}_{y,j}(\widetilde{\mathbf{H}})\coloneqq (\widetilde{\mathbf{h}}_{2j - 1})_{d + 1}$ . Given a loss function  $\ell :\mathbb{R}\times \mathbb{R}\to \mathbb{R}$  associated to a single response, the training loss associated to the whole input sequence can be defined by  $\ell (\mathbf{H}) = \sum_{j = 1}^{N + 1}\ell (y_j,\widehat{y}_j)$ . This potentially enables less training sequences in the pre-training stage, and some generalization bound analysis justifying this benefit was provided in [47].
We discuss how our theoretical results upon encoder TFs can be converted to those of the decoder TFs. Taking the implementation of (ICGD) (a key mechanism that enables most basic ICL algorithms such as ridge regression; cf. Appendix D.1) as an example, this conversion is enabled by the following facts: (a) the input format (11) of decoders can be converted to the input format (3) of encoders by a 2-layer decoder TF; (b) the encoder TF that implements (ICGD) with input format (3), by a slight parameter modification, can be converted to a decoder TF that implements the (ICGD) algorithm with a converted input format.
Input format conversion Despite the difference between the input format (11) and (3), we show that there exists a 2-layer decoder TF that can convert the input format (11) to format (3). The proof can be found in Appendix C.4.
Proposition C.1 (Input format conversion). There exists a 2-layer decoder TF DTF with 3 heads per layer, hidden dimension 2 and  $\| \pmb{\theta} \| \leq 12$  such that upon taking input  $\mathbf{H}$  of format (11), it outputs  $\widetilde{\mathbf{H}} = \mathrm{DTF}(\mathbf{H})$  with
$$
\widetilde {\mathbf {H}} = \left[ \begin{array}{c c c c c c} \mathbf {x} _ {1} & \mathbf {x} _ {1} & \dots & \mathbf {x} _ {N} & \mathbf {x} _ {N} & \mathbf {x} _ {N + 1} \\ 0 & y _ {1} & \dots & 0 & y _ {N} & 0 \\ \mathbf {p} _ {1} & \mathbf {p} _ {2} & \dots & \mathbf {p} _ {2 N - 1} & \mathbf {p} _ {2 N} & \mathbf {p} _ {2 N + 1} \end{array} \right]. \tag {12}
$$
In particular, format (12) contains format (3) as a submatrix, by restricting to the  $\{1, 2, \ldots, D - 1, D - 2, D\}$  rows and  $\{2, 4, \ldots, 2N - 2, 2N, 2N + 1\}$  columns.
Generalization TF constructions to decoder architecture The construction in Theorem D.1 can be generalized to using the input format (12) along with a decoder TF, by using the scratch pad within the last token to record the gradient descent iterates. Further, if we slightly change the normalization in MSK from  $1 / i$  to  $1 / ((i - 1)\lor 1)$ , then the same construction performs (ICGD) (with training examples  $\{1,\ldots ,j\}$ ) at every token  $i = 2j + 1$  (corresponding to predicting at  $\mathbf{x}_{j + 1}$ ). Building on this extension, all our constructions in Section 3 and Section 4.2 can be generalized to decoder TFs.
C.4 Proof of Proposition C.1
For the simplicity of presentation, we write  $c_{i} = \lceil i / 2 \rceil$ ,  $t_{i} = \mathrm{mod}(i + 1,2)$ ,  $\mathbf{u}_i = \mathbf{h}_i[1:d] \in \mathbb{R}^{d + 1}$  be the vector of first  $d$  entries of  $\mathbf{h}_i$ , and let  $v_{i} = \mathbf{h}_{i}[d + 1]$  be the  $(d + 1)$ -th entry of  $\mathbf{h}_i$ . With such notations, the input sequence  $\mathbf{H} = [\mathbf{h}_i]_i$  can be compactly written as
$$
\mathbf {h} _ {i} = \left[ \mathbf {u} _ {i}; v _ {i}; \mathbf {0} _ {D - d - 4}; c _ {i}; 1; t _ {i} \right].
$$
In the following, we construct the desired  $\pmb{\theta} = (\pmb{\theta}^{(1)},\pmb{\theta}^{(2)})$  as follows.
Step 1: construction of  $\pmb{\theta}^{(1)} = (\pmb{\theta}_{\mathrm{matn}}^{(1)},\pmb{\theta}_{\mathrm{mlp}}^{(1)})$ , so that  $\mathrm{MLP}_{\pmb{\theta}_{\mathrm{mlp}}^{(1)}}\circ \mathrm{MAttn}_{\pmb{\theta}_{\mathrm{matn}}^{(1)}}$  maps
$$
\mathbf {h} _ {i} \xrightarrow {\mathrm {M A t t n} _ {\boldsymbol {\theta} _ {\mathrm {m a t t n}} ^ {(1)}}} \quad \mathbf {h} _ {i} ^ {\prime} = [ \mathbf {u} _ {i}; v _ {i}; \mathbf {0} _ {D - d - 6}; t _ {i} (c _ {i} ^ {2} + 0. 5); t _ {i} c _ {i}; c _ {i}; 1; t _ {i} ]
$$
In other words, when  $2 \nmid i$ ,  $\mathbf{u}_i = \mathbf{x}_{(i-1)/2}$ ; when  $2 \mid i$ ,  $\mathbf{u}_i = \mathbf{0}_d$ .
$$
\begin{array}{c c} \stackrel {{\mathrm {M L P}}} {{\longrightarrow}} & \mathbf {h} _ {i} ^ {(1)} = [ \mathbf {u} _ {i}; v _ {i}; \mathbf {0} _ {D - d - 6}; t _ {i} c _ {i} ^ {2}; t _ {i} c _ {i}; c _ {i}; 1; t _ {i} ]. \end{array}
$$
For  $m\in \{0,1\}$ , we define matrices  $\mathbf{Q}_m^{(1)},\mathbf{K}_m^{(1)},\mathbf{V}_m^{(1)}\in \mathbb{R}^{D\times D}$  such that
$$
\mathbf {Q} _ {0} ^ {(1)} \mathbf {h} _ {i} = \mathbf {Q} _ {1} ^ {(1)} \mathbf {h} _ {i} = \left[ \begin{array}{c} t _ {i} \\ \mathbf {0} \end{array} \right], \qquad \mathbf {K} _ {0} ^ {(1)} \mathbf {h} _ {j} = \mathbf {K} _ {1} ^ {(1)} \mathbf {h} _ {j} = \left[ \begin{array}{c} c _ {j} \\ \mathbf {0} \end{array} \right], \qquad \mathbf {V} _ {0} ^ {(1)} \mathbf {h} _ {j} = \left[ \begin{array}{c} \mathbf {0} _ {D - 4} \\ 3 c _ {j} \\ \mathbf {0} _ {3} \end{array} \right], \qquad \mathbf {V} _ {1} ^ {(1)} \mathbf {h} _ {j} = \left[ \begin{array}{c} \mathbf {0} _ {D - 3} \\ 2 \\ \mathbf {0} _ {2} \end{array} \right],
$$
for all  $i,j$ . By the structure of  $\mathbf{h}_i$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max _ {m} \left\| \mathbf {Q} _ {m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \max _ {m} \left\| \mathbf {K} _ {m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \sum_ {m} \left\| \mathbf {V} _ {m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq 5.
$$
Now, for every  $i$
$$
\frac {1}{i} \sum_ {j = 1} ^ {i} \sum_ {m \in \{0, 1 \}} \sigma \Big (\Big \langle \mathbf {Q} _ {m} ^ {(1)} \mathbf {h} _ {i}, \mathbf {K} _ {m} ^ {(1)} \mathbf {h} _ {j} \Big \rangle \Big) \mathbf {V} _ {m} ^ {(1)} \mathbf {h} _ {j} = \frac {1}{i} \sum_ {j = 1} ^ {i} t _ {i} \cdot [ \mathbf {0} _ {D - 4}; 3 c _ {j} ^ {2}; 2 c _ {j}; 0; 0 ].
$$
Notice that  $t_i \neq 0$  only when  $2 \mid i$ , we then compute for  $i = 2k$  that
$$
\sum_ {j = 1} ^ {i} 3 c _ {j} ^ {2} = 3 \cdot \frac {k (k - 1) (2 k - 1)}{3} + 3 k ^ {2} = 2 k ^ {3} + k, \quad \sum_ {j = 1} ^ {i} 2 c _ {j} = 2 \cdot k (k - 1) + 2 k = 2 k ^ {2}.
$$
Therefore, the  $\pmb{\theta}_{\mathrm{mattn}}^{(1)} = \{(\mathbf{Q}_m^{(1)},\mathbf{K}_m^{(1)},\mathbf{V}_m^{(1)}\in \mathbb{R}^{D\times D})\}_{m\in \{0,1\}}$  we construct above is indeed the desired attention layer. The existence of the desired  $\pmb{\theta}_{\mathrm{mlp}}^{(1)}$  is clear, and  $\pmb{\theta}_{\mathrm{mlp}}^{(1)} = (\mathbf{W}_1^{(1)},\mathbf{W}_2^{(1)})$  can further be chosen so that  $\| \mathbf{W}_1^{(1)}\|_{\mathrm{op}}\leq 1,\| \mathbf{W}_2^{(1)}\|_{\mathrm{op}}\leq 1$
Step 2: construction of  $\pmb{\theta}^{(2)}$ . For every  $m \in \{-1,0,1\}$ , we define matrices  $\mathbf{Q}_m^{(2)}, \mathbf{K}_m^{(2)}, \mathbf{V}_m^{(2)} \in \mathbb{R}^{D \times D}$  such that
$$
\mathbf {Q} _ {0} ^ {(2)} \mathbf {h} _ {i} ^ {(1)} = \mathbf {Q} _ {1} ^ {(2)} \mathbf {h} _ {i} ^ {(1)} = \mathbf {Q} _ {- 1} ^ {(2)} \mathbf {h} _ {i} ^ {(1)} = \left[ \begin{array}{c} t _ {i} c _ {i} ^ {2} \\ t _ {i} c _ {i} \\ \mathbf {0} \end{array} \right],
$$
$$
\mathbf {K} _ {0} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} 1 \\ - c _ {j} \\ \mathbf {0} \end{array} \right], \qquad \mathbf {K} _ {1} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} 1 \\ - (c _ {j} + 1) \\ \mathbf {0} \end{array} \right], \qquad \mathbf {K} _ {1} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} 1 \\ - (c _ {j} - 1) \\ \mathbf {0} \end{array} \right],
$$
$$
\mathbf {V} _ {0} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} - 4 \mathbf {u} _ {j} \\ \mathbf {0} _ {D - d} \end{array} \right], \qquad \mathbf {V} _ {1} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \mathbf {V} _ {- 1} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} 2 \mathbf {u} _ {j} \\ \mathbf {0} _ {D - d} \end{array} \right],
$$
for all  $i,j$ . By the structure of  $\mathbf{h}_i^{(1)}$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max _ {m} \left\| \mathbf {Q} _ {m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 1, \quad \max _ {m} \left\| \mathbf {K} _ {m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 2, \quad \sum_ {m} \left\| \mathbf {V} _ {m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 8.
$$
Now, for every  $i,j$ , we have
$$
\begin{array}{l} \sum_ {m \in \{- 1, 0, 1 \}} \sigma \Big (\left\langle \mathbf {Q} _ {m} ^ {(2)} \mathbf {h} _ {i} ^ {(1)}, \mathbf {K} _ {m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} \right\rangle \Big) \mathbf {V} _ {m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} \\ = \left\{- 2 \sigma \left(t _ {i} c _ {i} ^ {2} - t _ {i} c _ {i} c _ {j}\right) + \sigma \left(t _ {i} c _ {i} ^ {2} - t _ {i} c _ {i} (c _ {j} + 1)\right) + \sigma \left(t _ {i} c _ {i} ^ {2} - t _ {i} c _ {i} (c _ {j} - 1)\right) \right\} \cdot 2 [ \mathbf {u} _ {j}; \mathbf {0} _ {D - d} ] \\ = \left\{- 2 \sigma \left(c _ {i} - c _ {j}\right) + \sigma \left(\left(c _ {i} - c _ {j}\right) - 1\right) + \sigma \left(\left(c _ {i} - c _ {j}\right) + 1\right) \right\} \cdot 2 c _ {i} t _ {i} \left[ \mathbf {u} _ {j}; \mathbf {0} _ {D - d} \right] \\ = \mathbb {I} \left(c _ {i} = c _ {j}\right) \cdot 2 c _ {i} t _ {i} \left[ \mathbf {u} _ {j}; \mathbf {0} _ {D - d} \right], \\ \end{array}
$$
where the last equality follows from the fact that
$$
- 2 \sigma (x) + \sigma (x - 1) + \sigma (x + 1) = \left\{ \begin{array}{l l} 0, & x \geq 1 \text {o r} x \leq - 1, \\ x + 1, & x \in [ - 1, 0 ], \\ 1 - x, & x \in [ 0, 1 ]. \end{array} \right.
$$
$$
\begin{array}{l} \frac {1}{i} \sum_ {j = 1} ^ {i} \sum_ {m \in \{- 1, 0, 1 \}} \sigma \left(\left\langle \mathbf {Q} _ {m} ^ {(2)} \mathbf {h} _ {i} ^ {(1)}, \mathbf {K} _ {m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} \right\rangle\right) \mathbf {V} _ {m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \frac {1}{i} \sum_ {j = 1} ^ {i} 2 \mathbb {I} (c _ {i} = c _ {j}) c _ {i} t _ {i} [ \mathbf {u} _ {j}; \mathbf {0} _ {D - d} ] \\ = \left\{ \begin{array}{l l} [ \mathbf {x} _ {k}; \mathbf {0} _ {D - d} ], & i = 2 k \\ \mathbf {0} _ {D}, & \text {o t h e r w i s e} \end{array} \right.. \\ \end{array}
$$
Therefore, the  $\pmb{\theta}_{\mathrm{mattn}}^{(2)} = \{(\mathbf{Q}_m^{(2)},\mathbf{K}_m^{(2)},\mathbf{V}_m^{(2)}\in \mathbb{R}^{D\times D})\}_{m\in \{-1,0,1\}}$  we construct above maps
$$
\mathbf {h} _ {i} ^ {(1)} \quad \rightarrow \quad \mathbf {h} _ {i} ^ {\prime \prime} = \left[ \mathbf {x} _ {\lceil i / 2 \rceil}; v _ {i}; \mathbf {0} _ {D - d - 6}; t _ {i} c _ {i} ^ {2}; t _ {i} c _ {i}; c _ {i}; 1; t _ {i} \right].
$$
Finally, we only need to take a MLP layer  $\theta_{\mathrm{mlp}}^{(2)} = (\mathbf{W}_1^{(2)},\mathbf{W}_2^{(2)})$  with hidden dimension 2 that maps
$$
\mathbf {h} _ {i} ^ {\prime \prime} \quad \rightarrow \quad \mathbf {h} _ {i} ^ {(2)} = \left[ \mathbf {x} _ {\lceil i / 2 \rceil}; v _ {i}; \mathbf {0} _ {D - d - 6}; 0; 0; c _ {i}; 1; t _ {i} \right],
$$
which clearly exists and can be chosen so that  $\| \mathbf{W}_1^{(2)}\|_{\mathrm{op}}\leq 1,\| \mathbf{W}_2^{(2)}\|_{\mathrm{op}}\leq 1$
Combining the two steps above, we complete the proof of Proposition C.1.
D Mechanism: In-context gradient descent
Technically, the constructions in Section 3.1-3.2 rely on a new efficient construction for transformers to implement in-context gradient descent and its variants, which we present as follows. We begin by presenting the result for implementing (vanilla) gradient descent on convex empirical risks.
Compact notation of input We will often use shorthand  $y_{i}^{\prime} \in \mathbb{R}$  defined as  $y_{i}^{\prime} = y_{i}$  for  $i \in [N]$  and  $y_{N + 1}^{\prime} = 0$  to simplify our notation, with which the input sequence  $\mathbf{H} \in \mathbb{R}^{D \times (N + 1)}$  can be compactly written as  $\mathbf{h}_i = [\mathbf{x}_i; y_i'; \mathbf{p}_i] = [\mathbf{x}_i; y_i'; \mathbf{0}_{D - d - 3}; 1; t_i]$  for  $i \in [N + 1]$ , where  $t_i := 1 \{i < N + 1\}$  is the indicator for the training examples.
D.1 Gradient descent on convex empirical risk
Let  $\ell(\cdot, \cdot): \mathbb{R}^2 \to \mathbb{R}$  be a loss function. Let  $\widehat{L}_N(\mathbf{w}) := \frac{1}{N} \sum_{i=1}^{N} \ell(\mathbf{w}^\top \mathbf{x}_i, y_i)$  denote the empirical risk with loss function  $\ell$  on dataset  $\{(\mathbf{x}_i, y_i)\}_{i \in [N]}$ , and
$$
\mathbf {w} _ {\mathrm {G D}} ^ {t + 1} := \mathbf {w} _ {\mathrm {G D}} ^ {t} - \eta \nabla \widehat {L} _ {N} \left(\mathbf {w} _ {\mathrm {G D}} ^ {t}\right) \tag {ICGD}
$$
denote the gradient descent trajectory on  $\widehat{L}_N$  with initialization  $\mathbf{w}_{\mathrm{GD}}^0\in \mathbb{R}^d$  and learning rate  $\eta >0$ . We require the partial derivative of the loss  $\partial_s\ell :(s,t)\mapsto \partial_s\ell (s,t)$  (as a bivariate function) to be approximable by a sum of relus, defined as follows.
Definition D.1 (Approximability by sum of relus). A function  $g: \mathbb{R}^k \to \mathbb{R}$  is ( $\varepsilon_{\mathrm{approx}}, R, M, C$ -approximable by sum of relus, if there exists a "  $(M, C)$ -sum of relus" function
$$
f _ {M, C} (\mathbf {z}) = \sum_ {m = 1} ^ {M} c _ {m} \sigma \big (\mathbf {a} _ {m} ^ {\top} [ \mathbf {z}; 1 ] \big) \mathrm {w i t h} \sum_ {m = 1} ^ {M} | c _ {m} | \leq C, \max _ {m \in [ M ]} \| \mathbf {a} _ {m} \| _ {1} \leq 1, \mathbf {a} _ {m} \in \mathbb {R} ^ {k + 1}, c _ {m} \in \mathbb {R},
$$
such that  $\sup_{\mathbf{z}\in [-R,R]^k}\left|g(\mathbf{z}) - f_{M,C}(\mathbf{z})\right|\leq \varepsilon_{\mathrm{approx}}$
Definition D.1 is known to contain broad class of functions. For example, any mildly smooth  $k$ -variate function is approximable by a sum of relus for any  $(\varepsilon_{\mathrm{approx}}, R)$ , with mild bounds on  $(M, C)$  (Proposition B.1, building on results of Bach [4]). Also, any function that is a  $(M, C)$ -sum of relus itself (which includes all piecewise linear functions) is by definition  $(0, \infty, M, C)$ -approximable by sum of relus.
We show that  $L$  steps of (ICGD) can be approximately implemented by an  $(L + 1)$ -layer transformer.
Theorem D.1 (Convex ICGD). Fix any  $B_w > 0$ ,  $L > 1$ ,  $\eta > 0$ , and  $\varepsilon \leq B_w / (2L)$ . Suppose that
1. The loss  $\ell (\cdot ,\cdot)$  is convex in the first argument;
2.  $\partial_s\ell$  is  $(\varepsilon ,R,M,C)$  -approximable by sum of relus with  $R = \max \{B_xB_w,B_y,1\}$
Figure 4: Illustration of our main mechanism for implementing basic ICL algorithms: One attention layer implements a single (ICGD) iterate (Proposition E.1 & Theorem D.1). Top: the attention mechanism as in Definition 1. Bottom: A single (ICGD) iterate. Middle: Linear algebraic illustration of the attention layer for implementing a GD update.
Then, there exists an attention-only transformer  $\mathrm{TF}_{\theta}^{0}$  with  $(L + 1)$  layers,  $\max_{\ell \in [L]} M^{(\ell)} \leq M$  heads within the first  $L$  layers, and  $M^{(L + 1)} = 2$  such that for any input data  $(\mathcal{D}, \mathbf{x}_{N + 1})$  such that
$$
\sup_{\| \mathbf{w}\|_{2}\leq B_{w}}\lambda_{\max}(\nabla^{2}\widehat{L}_{N}(\mathbf{w}))\leq 2 / \eta ,\quad \exists \mathbf{w}^{\star}\in \operatorname *{arg  min}_{\mathbf{w}\in \mathbb{R}^{d}}\widehat{L}_{N}(\mathbf{w}) \text{such that}\| \mathbf{w}^{\star}\|_{2}\leq B_{w} / 2,
$$
$\mathrm{TF}_{\pmb{\theta}}^{0}(\mathbf{H}^{(0)})$  approximately implements (ICGD) with initialization  $\mathbf{w}_{\mathrm{GD}}^{0} = \mathbf{0}$ :
1. (Parameter space) For every  $\ell \in [L]$ , the  $\ell$ -th layer's output  $\mathbf{H}^{(\ell)} = \mathrm{TF}_{\pmb{\theta}^{(1:\ell)}}^0 (\mathbf{H}^{(0)})$  approximates  $\ell$  steps of (ICGD): We have  $\mathbf{h}_i^{(\ell)} = [\mathbf{x}_i;y_i';\widehat{\mathbf{w}}^\ell ;\mathbf{0}_{D - 2d - 3};1;t_i]$  for every  $i\in [N + 1]$ , where
$$
\left\| \widehat {\mathbf {w}} ^ {\ell} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell} \right\| _ {2} \leq \varepsilon \cdot (\ell \eta B _ {x}).
$$
Note that the bound scales as  $\mathcal{O}(\ell)$ , a linear error accumulation.
2. (Prediction space) The final output  $\mathbf{H}^{(L + 1)} = \mathrm{TF}_{\pmb{\theta}}^{0}(\mathbf{H}^{(0)})$  approximates the prediction of  $L$  steps of (ICGD): We have  $\mathbf{h}_{N + 1}^{(L + 1)} = [\mathbf{x}_{N + 1};\widehat{y}_{N + 1};\widehat{\mathbf{w}}^L;\mathbf{0}_{D - 2d - 3};1;t_i]$ , where  $\widehat{y}_{N + 1} = \left\langle \widehat{\mathbf{w}}^L,\mathbf{x}_{N + 1}\right\rangle$  so that
$$
\left| \hat {y} _ {N + 1} - \left\langle \mathbf {w} _ {\mathrm {G D}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq \varepsilon \cdot \left(L \eta B _ {x} ^ {2}\right).
$$
Further, the transformer admits norm bound  $\| \pmb {\theta}\| \leq 2 + R + 2\eta C$
The proof can be found in Appendix E.2. Theorem D.1 substantially generalizes that of von Oswald et al. [86] (which only does GD on square losses with a linear self-attention), and is simpler than the ones in Akyurek et al. [2] and Giannou et al. [32]. See Figure 4 for a pictorial illustration of the basic component of the construction, which implements a single step of gradient descent using a single attention layer (Proposition E.1).
Technically, we utilize the stability of convex gradient descent as in the following lemma (proof in Appendix E.3) to obtain the linear error accumulation in Theorem D.1; the error accumulation will become exponential in  $L$  in the non-convex case in general; see Lemma D.3(b).
Lemma D.1 (Composition of error for approximating convex GD). Suppose  $f: \mathbb{R}^d \to \mathbb{R}$  is a convex function. Let  $\mathbf{w}^\star \in \arg \min_{\mathbf{w} \in \mathbb{R}^d} f(\mathbf{w})$ ,  $R \geq 2\|\mathbf{w}^\star\|_2$ , and assume that  $\nabla f$  is  $L_f$ -smooth on  $\mathsf{B}_2^d(R)$ . Let sequences  $\{\widehat{\mathbf{w}}^\ell\}_{\ell \geq 0} \subset \mathbb{R}^d$  and  $\{\mathbf{w}_{\mathrm{GD}}^\ell\}_{\ell \geq 0} \subset \mathbb{R}^d$  be given by  $\widehat{\mathbf{w}}^0 = \mathbf{w}_{\mathrm{GD}}^0 = \mathbf{0}$ ,
$$
\left\{ \begin{array}{l} \widehat {\mathbf {w}} ^ {\ell + 1} = \widehat {\mathbf {w}} ^ {\ell} - \eta \nabla f (\widehat {\mathbf {w}} ^ {\ell}) + \varepsilon^ {\ell}, \qquad \| \varepsilon^ {\ell} \| _ {2} \leq \varepsilon , \\ \mathbf {w} _ {\mathrm {G D}} ^ {\ell + 1} = \mathbf {w} _ {\mathrm {G D}} ^ {\ell} - \eta \nabla f (\mathbf {w} _ {\mathrm {G D}} ^ {\ell}), \end{array} \right.
$$
for all  $\ell \geq 0$ . Then as long as  $\eta \leq 2 / L_f$ , for any  $0 \leq L \leq R / (2\varepsilon)$ , it holds that  $\| \widehat{\mathbf{w}}^L - \mathbf{w}_{\mathrm{GD}}^L \|_2 \leq L\varepsilon$  and  $\| \widehat{\mathbf{w}}^L \|_2 \leq \frac{R}{2} + L\varepsilon \leq R$ .
D.2 Proximal gradient descent for regularized convex losses
Proximal gradient descent (PGD) is a variant of gradient descent that is suitable for minimizing regularized risks [66], in particular those with a non-smooth regularizer such as the  $\ell_1$  norm. In this section, we show that transformers can approximate PGD with similar quantitative guarantees as for GD in Appendix D.1.
Let  $\ell(\cdot, \cdot): \mathbb{R}^2 \to \mathbb{R}$  be a loss function. Let  $\widehat{L}_N(\mathbf{w}) := \frac{1}{N} \sum_{i=1}^{N} \ell(\mathbf{w}^\top \mathbf{x}_i, y_i) + \mathcal{R}(\mathbf{w})$  denote the regularized empirical risk with loss function  $\ell$  on dataset  $\{(x_i, y_i)\}_{i \in [N]}$  and regularizer  $\mathcal{R}$ . To minimize  $\widehat{L}_N$ , we consider the proximal gradient descent trajectory on  $\widehat{L}_N$  with initialization  $\mathbf{w}_{\mathrm{GD}}^0 = \mathbf{0} \in \mathbb{R}^d$  and learning rate  $\eta > 0$ :
$$
\mathbf {w} _ {\mathrm {P G D}} ^ {t + 1} := \mathbf {p r o x} _ {\eta \mathcal {R}} \left(\mathbf {w} _ {\mathrm {P G D}} ^ {t} - \eta \nabla \widehat {L} _ {N} ^ {0} \left(\mathbf {w} _ {\mathrm {P G D}} ^ {t}\right)\right), \tag {ICPGD}
$$
where we denote  $\widehat{L}_N^0 (\mathbf{w})\coloneqq \frac{1}{N}\sum_{i = 1}^{N}\ell (\mathbf{w}^\top \mathbf{x}_i,y_i)$
To approximate (ICPGD) by transformers, in addition to the requirement on the loss  $\ell$  as in Theorem D.1, we additionally require the proximal operator  $\mathbf{prox}_{\eta \mathcal{R}}(\cdot)$  to be approximable by an MLP layer (as a vector-valued analog of Definition D.1) defined as follows.