SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {1}{2} \left\| \widehat {\mathbf {w}} _ {\mathrm {L S}} - \mathbf {w} _ {\mathsf {P}} ^ {\star} \right\| _ {\boldsymbol {\Sigma} _ {\mathsf {P}}} ^ {2} 1 \left\{E _ {\text {c o v}} \right\} = \frac {1}{2} \left(\left(\mathbf {X} ^ {\top} \mathbf {X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {y} - \mathbf {w} _ {\mathsf {P}} ^ {\star}\right) ^ {\top} \boldsymbol {\Sigma} _ {\mathsf {P}} \left(\left(\mathbf {X} ^ {\top} \mathbf {X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {y} - \mathbf {w} _ {\mathsf {P}} ^ {\star}\right) 1 \left\{E _ {\text {c o v}} \right\} \\ = \frac {1}{2} \mathbf {z} ^ {\top} \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {X}) ^ {- 1} \boldsymbol {\Sigma} _ {\mathrm {P}} (\mathbf {X} ^ {\top} \mathbf {X}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {z} \cdot 1 \{E _ {\mathrm {c o v}} \} \\ = \frac {1}{2 N ^ {2}} \mathbf {z} ^ {\top} \mathbf {X} \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {- 1 / 2} \left(\boldsymbol {\Sigma} _ {\mathsf {P}} ^ {1 / 2} \hat {\boldsymbol {\Sigma}} ^ {- 1} \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {1 / 2}\right) ^ {2} \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {- 1 / 2} \mathbf {X} ^ {\top} \mathbf {z} \cdot 1 \{E _ {\mathrm {c o v}} \} \\ \leq \frac {2}{N ^ {2}} \left\| \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {- 1 / 2} \mathbf {X} ^ {\top} \mathbf {z} \right\| _ {2} ^ {2} 1 \{E _ {\mathrm {c o v}} \} = \frac {2}{N ^ {2}} \left\| \sum_ {i = 1} ^ {N} \widetilde {\mathbf {x}} _ {i} z _ {i} \right\| _ {2} ^ {2} 1 \{E _ {\mathrm {c o v}} \} \leq \frac {2}{N ^ {2}} \left\| \sum_ {i = 1} ^ {N} \widetilde {\mathbf {x}} _ {i} z _ {i} \right\| _ {2} ^ {2}. \\ \end{array}
$$
Note that  $\mathbb{E}[\widetilde{\mathbf{x}}_i z_i] = \pmb{\Sigma}_{\mathsf{P}}^{-1/2} \mathbb{E}[\mathbf{x}_i (y_i - \langle \mathbf{w}_{\mathsf{P}}^\star, \mathbf{x}_i \rangle)] = 0$ . Therefore, taking expectation on the above (over  $\mathcal{D}$ ), we get
$$
\mathbb {E} _ {\mathcal {D}} \left[ \frac {1}{2} \| \widehat {\mathbf {w}} _ {\mathrm {L S}} - \mathbf {w} _ {\mathsf {P}} ^ {\star} \| _ {\Sigma_ {\mathsf {P}}} ^ {2} 1 \{E _ {\mathrm {c o v}} \} \right] \leq \frac {2}{N ^ {2}} \mathbb {E} \left[ \left\| \sum_ {i = 1} ^ {N} \widetilde {\mathbf {x}} _ {i} z _ {i} \right\| _ {2} ^ {2} \right] = \frac {2}{N} \mathbb {E} \left[ \| \widetilde {\mathbf {x}} _ {1} z _ {1} \| _ {2} ^ {2} \right] = \frac {2}{N} \mathbb {E} \left[ z _ {1} ^ {2} \mathbf {x} _ {1} ^ {\top} \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {- 1} \mathbf {x} _ {1} \right] \tag {29}
$$
$$
\leq \frac {2 \sigma^ {2}}{N} \mathbb {E} \left[ \mathbf {x} _ {1} ^ {\top} \boldsymbol {\Sigma} _ {\mathrm {P}} ^ {- 1} \mathbf {x} _ {1} \right] = \frac {2 d \sigma^ {2}}{N}. \tag {30}
$$
Above, (i) follows by conditioning on  $\mathbf{x}_1$  and using Assumption A(5). Combining with (28), we obtain
$$
\mathbb {E} \left[ \frac {1}{2} \left(\operatorname {c l i p} _ {R} \left(\left\langle \widehat {\mathbf {w}} _ {\mathrm {L S}}, \mathbf {x} _ {N + 1} \right\rangle\right) - y _ {N + 1}\right) ^ {2} \right] \leq L _ {\mathsf {P}} \left(\mathbf {w} _ {\mathsf {P}} ^ {\star}\right) + \mathcal {O} \left(R ^ {2} \delta + \frac {d \sigma^ {2}}{N}\right).
$$
This proves (24).
Finally, we show  $\mathbb{P}(E_{\mathrm{cov}}\cap E_w)\geq 1 - \delta /10$  . Using (29) and  $\pmb{\Sigma}_{\mathbb{P}}\succeq \lambda_{\min}\mathbf{I}_d$  by assumption, we get
$$
\mathbb {E} \left[ \| \widehat {\mathbf {w}} _ {\mathrm {L S}} - \mathbf {w} _ {\mathsf {P}} ^ {\star} \| _ {2} ^ {2} 1 \{E _ {\mathrm {c o v}} \} \right] \leq \frac {4 d \sigma^ {2}}{N \lambda_ {\min }}.
$$
Therefore, using an argument similar to Chebyshev's inequality,
$$
\begin{array}{l} \mathsf {P} (E _ {\mathrm {c o v}} \cap E _ {w} ^ {c}) = \mathbb {E} \left[ 1 \{E _ {\mathrm {c o v}} \} \times 1 \{\| \widehat {\mathbf {w}} _ {\mathrm {L S}} \| _ {2} > \sqrt {\frac {2 0}{\delta} \cdot \frac {4 d \sigma^ {2}}{N \lambda_ {\mathrm {m i n}}}} + B _ {w} ^ {\star} \} \right] \\ \leq \mathbb {E} \left[ 1 \left\{E _ {\mathrm {c o v}} \right\} \times 1 \left\{\left\| \widehat {\mathbf {w}} _ {\mathrm {L S}} - \mathbf {w} _ {\mathsf {P}} ^ {\star} \right\| _ {2} > \sqrt {\frac {2 0}{\delta} \cdot \frac {4 d \sigma^ {2}}{N \lambda_ {\mathrm {m i n}}} \right\} \right] \\ \leq \mathbb {E} \left[ 1 \{E _ {\mathrm {c o v}} \} \times \frac {\| \widehat {\mathbf {w}} _ {\mathrm {L S}} - \mathbf {w} _ {\mathsf {P}} ^ {\star} \| _ {2} ^ {2}}{\frac {2 0}{\delta} \cdot \frac {4 d \sigma^ {2}}{N \lambda_ {\mathrm {m i n}}} \right] \leq \delta / 2 0. \\ \end{array}
$$
$$
\mathrm {P} \left(E _ {\text {c o v}} \cap E _ {w}\right) = \mathrm {P} \left(E _ {\text {c o v}}\right) - \mathrm {P} \left(E _ {\text {c o v}} \cap E _ {w} ^ {c}\right) \geq 1 - \delta / 2 0 - \delta / 2 0 \geq 1 - \delta / 1 0.
$$
F.3 Proof of Corollary 5
The proof follows by first checking the well-conditionedness of the data  $\mathcal{D}$  (cf. (5)) with high probability, then invoking Theorem 4 (for approximation least squares) and Proposition F.1 (for the statistical power of least squares).
First, as  $\mathsf{P}$  satisfies Assumption A, by Proposition F.1, as long as  $N\geq \mathcal{O}(K^4 (d + \log (1 / \delta)))$  we have with probability at least  $1 - \delta /10$  that event  $E_{\mathrm{cov}}\cap E_w$  holds. On this event, we have
$$
\frac {1}{2} \lambda_ {\min } \mathbf {I} _ {d} \preceq \frac {1}{2} \boldsymbol {\Sigma} _ {\mathrm {P}} \preceq \widehat {\boldsymbol {\Sigma}} = \mathbf {X} ^ {\top} \mathbf {X} / N \preceq 2 \boldsymbol {\Sigma} _ {\mathrm {P}} \preceq 2 \lambda_ {\max } \mathbf {I} _ {d},
$$
$$
\| \widehat {\mathbf {w}} _ {\mathrm {L S}} \| _ {2} \leq B _ {w} / 2 := \mathcal {O} \left(B _ {w} ^ {\star} + \sqrt {\frac {d \sigma^ {2}}{\delta N \lambda_ {\min}}}\right),
$$
and thus the dataset  $\mathcal{D}$  is well-conditioned (in the sense of (5)) with parameters  $\alpha = \lambda_{\mathrm{min}} / 2$ ,  $\beta = 2\lambda_{\mathrm{max}}$ , and  $B_w$  defined as above. Note that the condition number of  $\widehat{\Sigma}$  is upper bounded by  $\beta / \alpha = 4\lambda_{\mathrm{max}} / \lambda_{\mathrm{min}} \leq 4\kappa$ , where  $\kappa$  is the upper bound on the condition number of  $\Sigma_{\mathbb{P}}$  as in Assumption A(c).
Define parameters
$$
\varepsilon = \sqrt {\frac {d \sigma^ {2}}{N}}, \quad \delta = \frac {d \sigma^ {2}}{B _ {y} ^ {2} N} \wedge 1. \tag {31}
$$
Note that  $B_w \leq \mathcal{O}(B_w^\star + \sqrt{B_y^2 / \lambda_{\min}})$  by the above choice of  $\delta$ .
We can thus apply Theorem 4 in the unregularized case ( $\lambda = 0$ ) to obtain that, there exists a transformer  $\pmb{\theta}$  with  $\max_{\ell \in [L]} M^{(\ell)} \leq 3$ ,  $\| \pmb{\theta} \| \leq 4R + 4 / \lambda_{\max}$  (with  $R = \max \{B_x B_w, B_y, 1\}$ ), and number of layers
$$
L \leq \mathcal {O} \bigg (\kappa \log \frac {B _ {x} B _ {w}}{\varepsilon} \bigg) \leq \mathcal {O} \bigg (\kappa \log \left(B _ {x} \sqrt {\frac {N}{d \sigma^ {2}}} \left(B _ {w} ^ {\star} + \frac {B _ {y} ^ {2}}{\sqrt {\lambda_ {\mathrm {m i n}}}}\right)\right) \bigg),
$$
such that on  $E_{\mathrm{cov}} \cap E_w$  (so that  $\mathcal{D}$  is well-conditioned), we have (choosing the clipping radius in  $\widetilde{\mathrm{read}}_y(\cdot) = \mathrm{clip}_{B_y}(\mathrm{read}_y(\cdot))$  to be  $B_y$ ):
$$
\left| \widetilde {\operatorname {r e a d}} _ {\mathcal {Y}} \left(\mathrm {T F} _ {\boldsymbol {\theta}} ^ {0} (\mathbf {H})\right) - \operatorname {c l i p} _ {B _ {y}} \left(\left\langle \widehat {\mathbf {w}} _ {\mathrm {L S}}, \mathbf {x} _ {N + 1} \right\rangle\right) \right| \leq \left| \operatorname {r e a d} _ {\mathcal {Y}} \left(\mathrm {T F} _ {\boldsymbol {\theta}} ^ {0} (\mathbf {H})\right) - \left\langle \widehat {\mathbf {w}} _ {\mathrm {L S}}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq \varepsilon = \sqrt {\frac {d \sigma^ {2}}{N}}. \tag {32}
$$
We now bound the excess risk of the above transformer. Combining Proposition F.1 and (32), we have
$$
\mathbb {E} \left[ \left(\widetilde {\operatorname {r e a d}} _ {\mathrm {y}} \left(\operatorname {T F} _ {\boldsymbol {\theta}} ^ {0} (\mathbf {H})\right) - y _ {N + 1}\right) ^ {2} \right]
$$
$$
\begin{array}{l} = \mathbb {E} \left[ \left(\widetilde {\mathrm {r e a d}} _ {\mathbf {y}} (\mathrm {T F} _ {\pmb {\theta}} ^ {0} (\mathbf {H})) - y _ {N + 1}\right) ^ {2} 1 \{E _ {\mathrm {c o v}} \cap E _ {w} \} \right] + \mathbb {E} \left[ \left(\widetilde {\mathrm {r e a d}} _ {\mathbf {y}} (\mathrm {T F} _ {\pmb {\theta}} ^ {0} (\mathbf {H})) - y _ {N + 1}\right) ^ {2} 1 \{(E _ {\mathrm {c o v}} \cap E _ {w}) ^ {c} \} \right] \\ \leq 2 \mathbb {E} \left[ \left(\widetilde {\mathrm {r e a d}} _ {\mathbf {y}} (\mathrm {T F} _ {\pmb {\theta}} ^ {0} (\mathbf {H})) - \mathrm {c l i p} _ {B _ {y}} (\langle \widehat {\mathbf {w}} _ {\mathrm {L S}}, \mathbf {x} _ {N + 1} \rangle)\right) ^ {2} 1 \{E _ {\mathrm {c o v}} \cap E _ {w} \} \right] \\ + 2 \mathbb {E} \left[ \left(\operatorname {c l i p} _ {B _ {y}} \left(\langle \widehat {\mathbf {w}} _ {\mathrm {L S}}, \mathbf {x} _ {N + 1} \rangle\right) - y _ {N + 1}\right) ^ {2} 1 \{E _ {\mathrm {c o v}} \cap E _ {w} \} \right] + 2 B _ {y} ^ {2} \cdot \delta / 1 0 \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i)} {\leq} 2 \varepsilon^ {2} + L _ {\mathsf {P}} \left(\mathbf {w} _ {\mathsf {P}} ^ {\star}\right) + \mathcal {O} \left(B _ {y} ^ {2} \delta + \frac {d \sigma^ {2}}{N}\right) + \mathcal {O} \left(B _ {y} ^ {2} \delta\right) \\ \leq L _ {\mathsf {P}} \left(\mathbf {w} _ {\mathsf {P}} ^ {\star}\right) + \mathcal {O} \left(B _ {y} ^ {2} \delta + \frac {d \sigma^ {2}}{N}\right) \leq \mathcal {O} \left(\frac {d \sigma^ {2}}{N}\right). \\ \end{array}
$$
Above, (i) uses the approximation guarantee (32) as well as Proposition F.1(a) (with clipping radius  $B_{y}$ ). This proves the desired excess risk guarantee.
Finally, under the canonical choice of parameters (23), the bounds for  $L,M,\| \pmb {\theta}\|$  simplify to
$$
L \leq \mathcal {O} \left(\kappa \log \frac {N \kappa}{\sigma}\right), \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq 3, \quad \| \boldsymbol {\theta} \| \leq \mathcal {O} (\sqrt {\kappa d}), \tag {33}
$$
and the requirement for  $N$  simplifies to  $N \geq \mathcal{O}(d + \log(1 / \delta)) = \tilde{\mathcal{O}}(d)$  (as  $K = \Theta(1)$ ). This proves the claim about the required  $\bar{N}$  and  $L$ .
F.4 Proof of Corollary 6
Fix parameters  $\delta, \underline{\varepsilon} > 0$  to be specified later and a large universal constant  $C_0$ . Let us set
$$
\alpha = \max \left\{0, 1 / 2 - \sqrt {d / N} \right\} ^ {2}, \beta = 2 5,
$$
$$
B _ {w} ^ {\star} := 1 + 2 \sqrt {\frac {\log (4 / \delta)}{d}}, \qquad B _ {w} = C _ {0} (B _ {w} ^ {\star} + \sigma),
$$
$$
B _ {x} = C _ {0} \sqrt {d \log (N / \delta)}, \qquad B _ {y} = C _ {0} \left(B _ {w} ^ {\star} + \sigma\right) \sqrt {\log (N / \delta)}.
$$
Consider the following good events (below  $\varepsilon = [\varepsilon_i]_{i\in [N]}\in \mathbb{R}^N$  is given by  $\varepsilon_{i} = y_{i} - \langle \mathbf{w}_{\star},\mathbf{x}_{i}\rangle$ )
$$
\mathcal {E} _ {\pi} = \left\{\| \mathbf {w} _ {\star} \| _ {2} \leq B _ {w} ^ {\star}, \| \boldsymbol {\varepsilon} \| _ {2} \leq 2 \sqrt {N} \sigma \right\},
$$
$$
\mathcal {E} _ {w} = \left\{\alpha \leq \lambda_ {\min } \left(\mathbf {X} ^ {\top} \mathbf {X} / N\right) \leq \lambda_ {\max } \left(\mathbf {X} ^ {\top} \mathbf {X} / N\right) \leq \beta \right\},
$$
$$
\mathcal {E} _ {b} = \{\forall i \in [ N ], \| \mathbf {x} _ {i} \| _ {2} \leq B _ {x}, | y _ {i} | \leq B _ {y} \},
$$
$$
\mathcal {E} _ {b, N + 1} = \left\{\| \mathbf {x} _ {N + 1} \| _ {2} \leq B _ {x}, | y _ {N + 1} | \leq B _ {y} \right\},
$$
and we define  $\mathcal{E} \coloneqq \mathcal{E}_{\pi} \cap \mathcal{E}_w \cap \mathcal{E}_b \cap \mathcal{E}_{b,N+1}$ . Under the event  $\mathcal{E}$ , the problem (ICRidge) is well-conditioned and  $\| \mathbf{w}_{\mathrm{ridge}}^{\lambda} \| \leq B_w / 2$  (by Lemma F.1).
Therefore, Theorem 4 implies that for  $\kappa = \frac{\alpha + \lambda}{\beta + \lambda}$ , there exists a  $L = \lceil 2\kappa \log (B_w / \underline{\varepsilon}) \rceil + 1$ -layer transformer  $\pmb{\theta}$  with prediction  $\widehat{y}_{N+1} := \widetilde{\mathrm{read}}_{\mathbf{y}}(\mathrm{TF}_{\pmb{\theta}}^{0}(\mathbf{H}))$  (clipped by  $B_y$ ), such that under the good event  $\mathcal{E}$ , we have  $\widehat{y}_{N+1} = \mathrm{clip}_{B_y}(\langle \mathbf{x}_{N+1}, \widehat{\mathbf{w}} \rangle)$  and  $\| \widehat{\mathbf{w}} - \mathbf{w}_{\mathrm{ridge}}^{\lambda} \| \leq \underline{\varepsilon}$ .
In the following, we show that  $\theta$  is indeed the desired transformer (when  $\underline{\varepsilon}$  and  $\delta$  is suitably chosen). Notice that we have
$$
\mathbb {E} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} = \mathbb {E} \big [ 1 \{\mathcal {E} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \big ] + \mathbb {E} \big [ 1 \{\mathcal {E} ^ {c} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \big ],
$$
and we analyze these two parts separately.
Prediction risk under good event  $\mathcal{E}$ . We first note that
$$
\begin{array}{l} \mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] = \mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\operatorname {c l i p} _ {B _ {y}} \left(\left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} \right\rangle\right) - y _ {N + 1}\right) ^ {2} \right] \\ \leq \mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} \right\rangle - y _ {N + 1}\right) ^ {2} \right], \\ \end{array}
$$
where the inequality is because  $y_{N + 1}\in [-B_y,B_y]$  under the good event  $\mathcal{E}$ . Notice that by our construction, under the good event  $\mathcal{E}$ ,  $\widehat{\mathbf{w}} = \widehat{\mathbf{w}} (\mathcal{D})$  depends only on the dataset  $\mathcal{D}^7$ . Therefore, we have  $\| \widehat{\mathbf{w}} (\mathcal{D}) - \mathbf{w}_{\mathrm{ridge}}^{\lambda}(\mathcal{D})\| \leq \underline{\varepsilon}$  as long as the event  $\mathcal{E}_0\coloneqq \mathcal{E}_\pi \cap \mathcal{E}_w\cap \mathcal{E}_b$  holds for  $(\mathbf{w}_{\star},\mathcal{D})$ . Thus, under  $\mathcal{E}_0$ ,
$$
\begin{array}{l} \mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} \right\rangle - y _ {N + 1}\right) ^ {2} \mid \mathbf {w} _ {\star}, \mathcal {D} \right] = \mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} (\mathcal {D}) \right\rangle - y _ {N + 1}\right) ^ {2} \mid \mathbf {w} _ {\star}, \mathcal {D} \right] \\ \leq \mathbb {E} \left[ \left(\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} (\mathcal {D}) \rangle - y _ {N + 1}\right) ^ {2} \mid \mathbf {w} _ {\star}, \mathcal {D} \right] \\ = \mathbb {E} \left[ \left(\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} (\mathcal {D}) \rangle - \langle \mathbf {x} _ {N + 1}, \mathbf {w} _ {\star} \rangle\right) ^ {2} \mid \mathbf {w} _ {\star}, \mathcal {D} \right] + \sigma^ {2} \\ = \| \widehat {\mathbf {w}} (\mathcal {D}) - \mathbf {w} _ {\star} \| _ {2} ^ {2} + \sigma^ {2}, \\ \end{array}
$$
$$
\begin{array}{l} \left\| \widehat {\mathbf {w}} (\mathcal {D}) - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} \leq \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} + 2 \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} - \mathbf {w} _ {\star} \right\| _ {2} \left\| \widehat {\mathbf {w}} (\mathcal {D}) - \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} \right\| _ {2} + \left\| \widehat {\mathbf {w}} (\mathcal {D}) - \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} \right\| _ {2} ^ {2} \\ \leq \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} + 2 \underline {{\varepsilon}} \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} - \mathbf {w} _ {\star} \right\| _ {2} + \underline {{\varepsilon}} ^ {2}. \\ \end{array}
$$
Recall that  $2\mathrm{BayesRisk}_{\pi} = \mathbb{E}_{\mathbf{w}_{\star},\mathcal{D}}\| \mathbf{w}_{\mathrm{ridge}}^{\lambda} - \mathbf{w}_{\star}\|_{2}^{2} + \sigma^{2}$ . Note that  $2\mathrm{BayesRisk}_{\pi} \leq 1 + \sigma^{2}$  by definition. Therefore, we can conclude that
$$
\mathbb {E} \left[ 1 \{\mathcal {E} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq 2 \text {B a y e s R i s k} _ {\pi} + 2 \underline {{\varepsilon}} + \underline {{\varepsilon}} ^ {2}.
$$
Prediction risk under bad event  $\mathcal{E}^c$ . Notice that
$$
\mathbb {E} \left[ 1 \{\mathcal {E} ^ {c} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \right] \leq \sqrt {\mathbb {P} (\mathcal {E} ^ {c}) \mathbb {E} [ (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {4} ]}.
$$
We can upper bound  $\mathbb{P}(\mathcal{E}^c) = \mathbb{P}(\mathcal{E}_\pi^c \cup \mathcal{E}_w^c \cup \mathcal{E}_b^c \cup \mathcal{E}_{b,N+1}^c)$  by Lemma B.1, Lemma B.2 and the sub-Gaussian tail bound:
$$
\mathbb {P} \left(\mathcal {E} _ {\pi} ^ {c}\right) \leq \frac {\delta}{2} + \exp (- N / 8), \quad \mathbb {P} \left(\mathcal {E} _ {w} ^ {c}\right) \leq 2 \exp (- N / 8), \quad \mathbb {P} \left(\mathcal {E} _ {b} ^ {c} \cup \mathcal {E} _ {b, N + 1} ^ {c}\right) \leq \frac {\delta}{4}.
$$
Thus, as long as  $N \geq 8\log(12/\delta)$ , we have  $\mathbb{P}(\mathcal{E}^c) \leq \delta$ . Further, a simple calculation yields
$$
\mathbb {E} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {4} \leq 8 \mathbb {E} \widehat {y} _ {N + 1} ^ {4} + 8 \mathbb {E} y _ {N + 1} ^ {4} \leq 8 B _ {y} ^ {2} + 8 \mathbb {E} y _ {N + 1} ^ {4}.
$$
Notice that  $y_{N + 1}|\mathbf{w}_{\star}\sim \mathsf{N}(0,\| \mathbf{w}_{\star}\| _2^2 +\sigma^2)$ , hence  $\mathbb{E}y_{N + 1}^{4} = 3\mathbb{E}(\| \mathbf{w}_{\star}\| _2^2 +\sigma^2)^2\leq 3(3 + 2\sigma^2 +\sigma^4)\leq B_y^4$ . Thus, we can conclude that
$$
\mathbb {E} \left[ 1 \{\mathcal {E} ^ {c} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq 4 \sqrt {\delta} B _ {y}.
$$
Choosing  $\underline{\varepsilon}$  and  $\delta$ . Combining the inequalities above, we have
$$
\mathbb {E} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \leq 2 \text {B a y e s R i s k} _ {\pi} + \left[ 2 \underline {{\varepsilon}} \sqrt {2 \text {B a y e s R i s k} _ {\pi}} + \underline {{\varepsilon}} ^ {2} + 4 \sqrt {\delta} B _ {y} \right].
$$
To ensure  $\frac{1}{2}\mathbb{E}(\widehat{y}_{N + 1} - y_{N + 1})^2\leq \mathrm{BayesRisk}_\pi +\varepsilon$  , we only need to take  $(\underline{\varepsilon},\delta)$  so that the following constraints are satisfied:
$$
\underline {{\varepsilon}} = \frac {1}{2} \min  \left\{\varepsilon , \sqrt {\varepsilon} \right\}, \quad 4 \sqrt {\delta} B _ {y} \leq \frac {\varepsilon}{2}, \quad N \geq 8 \log (1 2 / \delta).
$$
Therefore, it suffices to take  $\delta = \frac{c_0}{\log^2(N)}\left(\frac{\varepsilon^2}{1 + \sigma^2}\right)^2$  for some small constant  $c_{0}$ , then as long as
$$
N \geq C \log \left(\frac {\sigma^ {2} + 1}{\varepsilon}\right) + C.
$$
our choice of  $\underline{\varepsilon}$  and  $\delta$  is feasible. Note that  $\kappa \leq \mathcal{O}\left(1 + \sigma^{-2}\right)$ , and hence under such choice of  $(\underline{\varepsilon},\delta)$ , we have  $L = O(\log (1 / \varepsilon))$  and  $\| \pmb {\theta}\| = \widetilde{O}\Big(\sqrt{d}\Big)$ . This is the desired result.
<sup>7</sup>We need this, as on  $\mathcal{E}^c$ , the transformer output at this location could in principle depend additionally on  $\mathbf{x}_{N + 1}$ , as (15) may not hold due to the potential unboundedness of its input. A similar fact will also appear in later proofs (for generalized linear models and Lasso).
Lemma F.1. Under the event  $\mathcal{E}_{\pi} \cap \mathcal{E}_w$ , we have  $\left\| \mathbf{w}_{\mathrm{ridge}}^{\lambda}\right\|_2 \leq \mathcal{O}(B_w^\star + \sigma)$ .
Proof of Lemma F.1. By the definition of  $\mathbf{w}_{\mathrm{ridge}}^{\lambda}$  and recall that  $\lambda = d\sigma^{2} / N$ , we have  $\mathbf{w}_{\mathrm{ridge}}^{\lambda} = (\mathbf{X}^{\top}\mathbf{X} + d\sigma^{2}\mathbf{I}_{d})^{-1}\mathbf{X}^{\top}\mathbf{y}$ .
Therefore, we only need to prove the following fact: for any  $\gamma > 0$  and  $\widehat{\beta} = (\mathbf{X}^{\top}\mathbf{X} + d\gamma \mathbf{I}_d)^{-1}\mathbf{X}^{\top}\mathbf{y}$ , we have
$$
\left\| \widehat {\boldsymbol {\beta}} \right\| _ {2} \leq B _ {w} ^ {\star} + 1 0 \sigma \left(1 + \gamma^ {- 1 / 2}\right). \tag {34}
$$
We now prove (34). Note that we have
$$
\| \widehat {\boldsymbol {\beta}} \| _ {2} = \| (\mathbf {X} ^ {\top} \mathbf {X} + d \gamma \mathbf {I} _ {d}) ^ {- 1} \mathbf {X} ^ {\top} (\mathbf {X} \mathbf {w} _ {\star} + \varepsilon) \| _ {2} \leq \| \mathbf {B} _ {1} \| _ {\mathrm {o p}} \| \mathbf {w} _ {\star} \| _ {2} + \| \mathbf {B} _ {2} \| _ {\mathrm {o p}} \| \varepsilon \| _ {2}
$$
where  $\mathbf{B}_1 = \mathbf{X}^\top \mathbf{X}(\mathbf{X}^\top \mathbf{X} + d\gamma \mathbf{I}_d)^{-1}$ ,  $\mathbf{B}_2 = (\mathbf{X}^\top \mathbf{X} + d\gamma \mathbf{I}_d)^{-1}\mathbf{X}^\top$ . Note that  $\| \mathbf{B}_1\|_{\mathrm{op}} \leq 1$  clearly holds, and under  $\mathcal{E}_{\pi}$  we also have  $\| \varepsilon \|_2 \leq 2\sqrt{N}\sigma$ . Therefore, it remains to bound the term  $\| \mathbf{B}_2\|_{\mathrm{op}}$ .
Consider the SVD decomposition of  $\mathbf{X} = U\Sigma V$ ,  $\Sigma = \mathrm{diag}(\lambda_1,\dots ,\lambda_d)$ , and  $U\in \mathbb{R}^{N\times d}$ ,  $V\in \mathbb{R}^{d\times d}$  are orthonormal matrices. Then  $\mathbf{B}_2 = V^\top (\Sigma^2 +d\gamma \mathbf{I}_d)^{-1}\Sigma U^\top$ , and hence
$$
\| \mathbf {B} _ {2} \| _ {\mathrm {o p}} = \left\| (\boldsymbol {\Sigma} ^ {2} + d \gamma \mathbf {I} _ {d}) ^ {- 1} \boldsymbol {\Sigma} \right\| _ {\mathrm {o p}} = \max _ {i} \frac {\lambda_ {i}}{\lambda_ {i} ^ {2} + d \gamma}.
$$
When  $N \leq 36d$ , we directly have  $\| \mathbf{B}_2 \|_{\mathrm{op}} \leq \frac{1}{2} (d\gamma)^{-1/2} \leq 3(N\gamma)^{-1/2}$ . Otherwise, we have  $N \geq 36d$ , and then for each  $i \in [d]$ ,  $\lambda_i \geq \sqrt{\lambda_{\min}(\mathbf{X}^\top \mathbf{X})} \geq \sqrt{\alpha N} \geq \sqrt{N}/3$ . Hence, in this case we also have  $\| \mathbf{B}_2 \|_{\mathrm{op}} \leq \max_i \lambda_i^{-1} \leq 3N^{-1/2}$ . Combining the both cases completes the proof of (34).
G In-context learning of generalized linear models
As a natural generalization of linear regression, we now show that transformers can recover learn generalized linear models (GLMs) [53] (which includes logistic regression for linear classification as an important special case), by implementing the corresponding convex risk minimization algorithm in context, and achieve near-optimal excess risk under standard statistical assumptions.
Let  $g: \mathbb{R} \to \mathbb{R}$  be a link function that is non-decreasing and  $C^2$ -smooth. We consider the following convex empirical risk minimization (ERM) problem
$$
\mathbf {w} _ {\mathrm {G L M}} := \underset {\mathbf {w} \in \mathbb {R} ^ {d}} {\arg \min } \widehat {L} _ {N} (\mathbf {w}) := \frac {1}{N} \sum_ {i = 1} ^ {N} \ell \left(\left\langle \mathbf {x} _ {i}, \mathbf {w} \right\rangle , y _ {i}\right), \tag {ICGLM}
$$
where  $\ell(t, y) \coloneqq -yt + \int_0^t g(s) ds$  is the convex (integral) loss associated with  $g$ . A canonical example of (ICGLM) is logistic regression, in which  $g(t) = \sigma_{\log}(t) \coloneqq (1 + e^{-t})^{-1}$  is the sigmoid function, and the resulting  $\ell(t, y) = \ell_{\log}(t, y) = -yt + \log(1 + e^t)$  is the logistic loss.
The following result (proof in Appendix G.1) shows that, as long as the empirical risk  $\widehat{L}_N$  satisfies strong convexity and bounded solution conditions (similar as in Theorem 4), transformers can approximately implement the ERM predictor  $g(\langle \mathbf{x}_{N + 1},\mathbf{w}_{\mathrm{GLM}}\rangle)$ , with  $\mathbf{w}_{\mathrm{GLM}}$  given by (ICGLM).
Theorem G.1 (Implementing convex risk minimization for GLMs). For any  $0 < \alpha < \beta$  with  $\kappa \coloneqq \frac{\beta}{\alpha}$ ,  $B_w > 0, B_x > 0$ ,  $\kappa_w \coloneqq L_g B_x^2 / \alpha + 1$  and  $\varepsilon < B_w / 2$ , there exists an attention-only transformer  $\mathrm{TF}_{\theta}^0$  with