SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} _ {k} \left[ \frac {1}{2} \left(y _ {N + 1} - \widehat {y} _ {N + 1}\right) ^ {2} \right] \leq \text {B a y e s R i s k} _ {k} + \mathcal {O} \left(\sigma_ {\max } ^ {4 / 3} B _ {w} ^ {2 / 3} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right) \\ \leq \mathsf {B a y e s R i s k} _ {k} + \mathcal {O} \left(\frac {\sigma_ {\max } ^ {2}}{\sigma_ {\min } ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3} + \sigma_ {\max } ^ {4 / 3} \frac {\log^ {2 / 3} (N) \log^ {1 / 3} (K)}{d ^ {2 / 3} N ^ {1 / 3}}\right) \\ \leq \mathsf {B a y e s R i s k} _ {k} + \widetilde {\mathcal {O}} \left(\frac {\sigma_ {\max } ^ {2}}{\sigma_ {\min } ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right), \\ \end{array}
$$
where we plug in our choice of  $B_{y}$ . The bounds on  $M^{(\ell)}, D^{(\ell)}$  and  $\| \pmb{\theta} \|$  follow immediately from Theorem I.2. This completes the proof.
J.2 Derivation of the exact Bayes predictor
Let  $(\mathcal{D},\mathbf{x}_{N + 1},y_{N + 1})$  be  $(N + 1)$  observations from the data generating model  $\pi$  considered in Section 4.1.1. On observing  $(\mathcal{D},\mathbf{x}_{N + 1})$ , the Bayes predictor of  $y_{N + 1}$  is given by its posterior mean:
$$
\mathbb {E} _ {\pi} [ y _ {N + 1} | \mathcal {D}, \mathbf {x} _ {N + 1} ] = \mathbb {E} _ {\pi} [ \langle \mathbf {x} _ {N + 1}, \mathbf {w} _ {\star} \rangle + \varepsilon_ {N + 1} | \mathcal {D}, \mathbf {x} _ {N + 1} ] = \langle \mathbf {x} _ {N + 1}, \mathbb {E} _ {\pi} [ \mathbf {w} _ {\star} | \mathcal {D} ] \rangle .
$$
It thus remains to derive  $\mathbb{E}_{\pi}[\mathbf{w}_{\star}|\mathcal{D}]$ . Recall that our data generating model is given by  $k\sim \Lambda$ , By Bayes' rule, we have
$$
\mathbb {E} _ {\pi} \left[ \mathbf {w} _ {\star} | \mathcal {D} \right] = \sum_ {k ^ {\prime} \in [ K ]} \mathbb {P} _ {\pi} \left(k = k ^ {\prime} | \mathcal {D}\right) \cdot \mathbb {E} _ {\pi} \left[ \mathbf {w} _ {\star} | \mathcal {D}, k = k ^ {\prime} \right]. \tag {46}
$$
On  $k = k'$ , the data is generated from the noisy linear model  $\mathbf{w}_{\star} \sim \mathsf{N}(\mathbf{0}, \mathbf{I}_d / d)$ , and  $\mathbf{y} = \mathbf{X}\mathbf{w}_{\star} + \varepsilon$  where  $\varepsilon_i \stackrel{\mathrm{iid}}{\sim} \mathsf{N}(0, \sigma_{k'}^2)$ . It is a standard result that  $\mathbb{E}_{\pi}[\mathbf{w}_{\star}|\mathcal{D}, k = k']$  is given by the ridge estimator
$$
\begin{array}{l} \mathbb {E} _ {\pi} \left[ \mathbf {w} _ {\star} \mid \mathcal {D}, k = k ^ {\prime} \right] = \underbrace {\left(\mathbf {X} ^ {\top} \mathbf {X} + d \sigma_ {k ^ {\prime}} ^ {2}\right) ^ {- 1}} _ {\widehat {\Sigma} _ {k ^ {\prime}} ^ {- 1}} \mathbf {X} ^ {\top} \mathbf {y} =: \widehat {\mathbf {w}} _ {k ^ {\prime}} \\ = \left(\frac {\mathbf {X} ^ {\top} \mathbf {X}}{N} + \frac {d \sigma_ {k ^ {\prime}} ^ {2}}{N}\right) ^ {- 1} \frac {\mathbf {X} ^ {\top} \mathbf {y}}{N}. \\ \end{array}
$$
(Note that the sample covariance within  $\widehat{\Sigma}_{k'}$  is not normalized by  $N$ , which is not to be confused with remaining parts within the paper.) Therefore, the posterior mean (46) is exactly a weighted combination of  $K$  ridge regression estimators, each with regularization  $d\sigma_k^2 / N$ .
It remains to derive the mixing weights  $\mathbb{P}_{\pi}(k = k'|\mathcal{D})$  for all  $k' \in [K]$ . By Bayes' rule, we have
$$
\begin{array}{l} \mathbb {P} _ {\pi} (k = k ^ {\prime} | \mathcal {D}) \propto_ {k ^ {\prime}} \mathbb {P} _ {\pi} (k = k ^ {\prime}) \cdot \int_ {\mathbf {w} _ {\star}} p (\mathbf {w} _ {\star}) \cdot \mathfrak {p} _ {k ^ {\prime}, \mathbf {w} _ {\star}} (\mathcal {D} | \mathbf {w} _ {\star}) d \mathbf {w} _ {\star} \\ \propto \Lambda_ {k ^ {\prime}} \cdot \int_ {\mathbf {w}} \frac {1}{(2 \pi d) ^ {d / 2} (2 \pi \sigma_ {k ^ {\prime}} ^ {2}) ^ {N / 2}} \exp \left(- \frac {d \| \mathbf {w} \| _ {2} ^ {2}}{2} - \frac {\| \mathbf {X} \mathbf {w} - \mathbf {y} \| _ {2} ^ {2}}{2 \sigma_ {k ^ {\prime}} ^ {2}}\right) d \mathbf {w} \\ \propto \Lambda_ {k ^ {\prime}} \cdot \int_ {\mathbf {w}} \frac {1}{(2 \pi \sigma_ {k ^ {\prime}} ^ {2}) ^ {N / 2}} \exp \left(- \frac {1}{2} \mathbf {w} ^ {\top} \left(\frac {\mathbf {X} ^ {\top} \mathbf {X}}{\sigma_ {k ^ {\prime}} ^ {2}} + d \mathbf {I} _ {d}\right) \mathbf {w} + \left\langle \mathbf {w}, \frac {\mathbf {X} ^ {\top} \mathbf {y}}{\sigma_ {k ^ {\prime}} ^ {2}} \right\rangle - \frac {\| \mathbf {y} \| _ {2} ^ {2}}{2 \sigma_ {k ^ {\prime}} ^ {2}}\right) d \mathbf {w} \\ \propto \Lambda_ {k ^ {\prime}} \cdot \int_ {\mathbf {w}} \frac {1}{(2 \pi \sigma_ {k ^ {\prime}} ^ {2}) ^ {N / 2}} \exp \left(- \frac {1}{2 \sigma_ {k ^ {\prime}} ^ {2}} (\mathbf {w} - \widehat {\mathbf {w}} _ {k ^ {\prime}}) ^ {\top} \widehat {\boldsymbol {\Sigma}} _ {k ^ {\prime}} (\mathbf {w} - \widehat {\mathbf {w}} _ {k ^ {\prime}}) - \frac {1}{2 \sigma_ {k ^ {\prime}} ^ {2}} \left(\| \mathbf {y} \| _ {2} ^ {2} - \mathbf {y} ^ {\top} \mathbf {X} \widehat {\boldsymbol {\Sigma}} _ {k ^ {\prime}} ^ {- 1} \mathbf {X} ^ {\top} \mathbf {y}\right)\right) d \mathbf {w} \\ \propto \Lambda_ {k ^ {\prime}} \cdot \frac {\operatorname* {d e t} \left(\widehat {\boldsymbol {\Sigma}} _ {k ^ {\prime}} / \sigma_ {k ^ {\prime}} ^ {2}\right) ^ {- 1 / 2}}{\sigma_ {k ^ {\prime}} ^ {N}} \exp \left(- \frac {1}{2 \sigma_ {k ^ {\prime}} ^ {2}} \left(\| \mathbf {y} \| _ {2} ^ {2} - \mathbf {y} ^ {\top} \mathbf {X} \widehat {\boldsymbol {\Sigma}} _ {k ^ {\prime}} ^ {- 1} \mathbf {X} ^ {\top} \mathbf {y}\right)\right) \\ \end{array}
$$
$$
\propto \Lambda_ {k ^ {\prime}} \cdot \frac {1}{\sigma_ {k ^ {\prime}} ^ {N - d} \mathrm {d e t} (\mathbf {X} ^ {\top} \mathbf {X} + d \sigma_ {k ^ {\prime}} ^ {2} \mathbf {I} _ {d}) ^ {1 / 2}} \exp \left(- \frac {1}{2 \sigma_ {k ^ {\prime}} ^ {2}} \Big (\| \mathbf {y} \| _ {2} ^ {2} - \langle \mathbf {y}, \mathbf {X} \widehat {\mathbf {w}} _ {k ^ {\prime}} \rangle \Big)\right).
$$
Note that such mixing weights involve the determinant of the matrix  $\widehat{\mathbf{\Sigma}}_{k'} = \mathbf{X}^\top \mathbf{X} + d\sigma_{k'}^2\mathbf{I}_d$ , which depends on the data  $\mathbf{X}$  in a non-trivial fashion; Any transformer has to approximate these weights if their mechanism is to directly approximate the exact Bayesian predictor (46).
J.3 Useful lemmas
Lemma J.1. For  $2\mathsf{Ris}k_{k,\mathrm{train}} = \mathbb{E}_k\left\| \mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}_{\mathrm{train}}) - \mathbf{w}_\star \right\| _2^2 +\sigma_k^2$  there exists universal constant  $C$  such that
$$
\operatorname {R i s k} _ {k, \text {t r a i n}} \leq \operatorname {B a y e s R i s k} _ {k} + C \left(\sigma_ {k} ^ {2} + 1\right) \frac {N _ {\text {v a l}}}{N}.
$$
Proof. Recall that under  $\mathbb{P}_k$ , we have
$$
\mathbf {w} _ {\star} \sim \mathsf {N} (0, \mathbf {I} _ {d} / d), \quad y _ {i} = \left\langle \mathbf {x} _ {i}, \mathbf {w} _ {\star} \right\rangle + \varepsilon_ {i}, \quad \varepsilon_ {i} \sim \mathsf {N} (0, \sigma^ {2}).
$$
We denote  $\mathbf{y}_t = [y_i]_{i\in \mathcal{I}_{train}}$ , then by definition  $\mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}_{train}) = (\mathbf{X}_{train}^\top \mathbf{X}_{train} + d\sigma_k^2)^{-1}\mathbf{X}_{train}\mathbf{y}_t$  (with  $\lambda_{k} = d\sigma_{k}^{2} / N_{train}$ ). Thus, a simple calculation yields
$$
2 \operatorname {R i s k} _ {k, \text {t r a i n}} = \mathbb {E} _ {k} \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda_ {k}} \left(\mathcal {D} _ {\text {t r a i n}}\right) - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} + \sigma_ {k} ^ {2} = \sigma_ {k} ^ {2} \mathbb {E} \operatorname {t r} \left(\left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} + d \sigma_ {k} ^ {2}\right) ^ {- 1}\right) + \sigma_ {k} ^ {2},
$$
and analogously, 2BayesRisk  $k_{k} = \sigma_{k}^{2}\mathbb{E}\mathrm{tr}\big((\mathbf{X}^{\top}\mathbf{X} + d\sigma_{k}^{2}\mathbf{I}_{d})^{-1}\big) + \sigma_{k}^{2}$ . Therefore,
$$
\begin{array}{l} 2 \operatorname {R i s k} _ {k, \text {t r a i n}} - 2 \operatorname {B a y e s R i s k} _ {k} = \sigma_ {k} ^ {2} \mathbb {E} \operatorname {t r} \left(\left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} + d \sigma_ {k} ^ {2} \mathbf {I} _ {d}\right) ^ {- 1}\right) - \sigma_ {k} ^ {2} \mathbb {E} \operatorname {t r} \left(\left(\mathbf {X} ^ {\top} \mathbf {X} + d \sigma_ {k} ^ {2} \mathbf {I} _ {d}\right) ^ {- 1}\right) \\ \leq \sigma_ {k} ^ {2} N _ {\text {v a l}} \mathbb {E} _ {k} \left[ \lambda_ {\min } (\boldsymbol {\Sigma}) ^ {- 1} \right], \\ \end{array}
$$
where in the above inequality we denote  $\pmb{\Sigma} \coloneqq \mathbf{X}_{\mathrm{train}}^{\top} \mathbf{X}_{\mathrm{train}} + d \sigma_k^2 \mathbf{I}_d$  and use the following fact:
$$
\begin{array}{l} \operatorname {t r} \left(\boldsymbol {\Sigma} ^ {- 1}\right) - \operatorname {t r} \left(\left(\boldsymbol {\Sigma} + \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v}\right) ^ {- 1}\right) = \operatorname {t r} \left(\boldsymbol {\Sigma} ^ {- 1 / 2} \left(\mathbf {I} _ {d} - \left(\mathbf {I} _ {d} + \boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1 / 2}\right) ^ {- 1}\right) \boldsymbol {\Sigma} ^ {- 1 / 2}\right) \\ = \operatorname {t r} \left(\boldsymbol {\Sigma} ^ {- 1 / 2} \left(\mathbf {I} _ {d} + \boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1 / 2}\right) ^ {- 1} \boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1}\right) \\ = \left\langle \left(\mathbf {I} _ {d} + \boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1 / 2}\right) ^ {- 1} \boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1 / 2}, \boldsymbol {\Sigma} ^ {- 1} \right\rangle \\ \leq \operatorname {r a n k} \left(\boldsymbol {\Sigma} ^ {- 1 / 2} \mathbf {X} _ {v} ^ {\top} \mathbf {X} _ {v} \boldsymbol {\Sigma} ^ {- 1 / 2}\right) \lambda_ {\max } (\boldsymbol {\Sigma} ^ {- 1}) \leq N _ {\text {v a l}} \lambda_ {\min } (\boldsymbol {\Sigma}) ^ {- 1}. \\ \end{array}
$$
Case 1. We first suppose that  $N_{\mathrm{train}} \leq 16d$ . Then by definition  $\Sigma \succeq d\sigma_k^2\mathbf{I}_d$ , and hence
$$
\sigma_ {k} ^ {2} N _ {\mathrm {v a l}} \mathbb {E} _ {k} [ \lambda_ {\mathrm {m i n}} (\pmb {\Sigma}) ^ {- 1} ] \leq \frac {\sigma_ {k} ^ {2} N _ {\mathrm {v a l}}}{d \sigma_ {k} ^ {2}} \leq \frac {1 6 N _ {\mathrm {v a l}}}{N _ {\mathrm {t r a i n}}} \leq \frac {3 2 N _ {\mathrm {v a l}}}{N}.
$$
Case 2. When  $N_{\mathrm{train}} \geq 9d$ , then we consider the event  $\mathcal{E}_t \coloneqq \{\lambda_{\min}(\mathbf{X}_{\mathrm{train}}^\top \mathbf{X}_{\mathrm{train}} / N_{\mathrm{train}}) \geq \frac{1}{16}\}$ . By Lemma B.2 we have  $\mathbb{P}(\mathcal{E}_t^c) \leq \exp(-N_{\mathrm{train}}/8)$ . Therefore,
$$
\begin{array}{l} \sigma_ {k} ^ {2} N _ {\mathrm {v a l}} \mathbb {E} _ {k} \left[ \lambda_ {\min } (\boldsymbol {\Sigma}) ^ {- 1} \right] = \sigma_ {k} ^ {2} N _ {\mathrm {v a l}} \mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {t} \right\} \lambda_ {\min } (\boldsymbol {\Sigma}) ^ {- 1} \right] + \sigma_ {k} ^ {2} N _ {\mathrm {v a l}} \mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {t} ^ {c} \right\} \lambda_ {\min } (\boldsymbol {\Sigma}) ^ {- 1} \right] \\ \leq \frac {1 6 \sigma_ {k} ^ {2} N _ {\text {v a l}}}{N _ {\text {t r a i n}}} \cdot \mathbb {P} \left(\mathcal {E} _ {t}\right) + \frac {N _ {\text {v a l}}}{d} \cdot \mathbb {P} \left(\mathcal {E} _ {t} ^ {c}\right) \\ \leq \frac {3 2 \sigma_ {k} ^ {2} N _ {\mathrm {v a l}}}{N} + \frac {N _ {\mathrm {v a l}}}{d} \cdot \exp (- N / 1 6) = \mathcal {O} \left(\frac {\left(\sigma_ {k} ^ {2} + 1\right) N _ {\mathrm {v a l}}}{N}\right). \\ \end{array}
$$
Combining these two cases finishes the proof.
Lemma J.2. Condition on the event  $\mathcal{E}_{\mathrm{train}}$ , we have
$$
\mathbb {E} _ {\mathcal {D} _ {\mathrm {v a l}} \sim \mathbb {P} _ {k} | \mathbf {w} _ {\star}, \mathcal {D} _ {\mathrm {t r a i n}}} \left[ \max _ {l \in [ K ]} \left| \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {l}) - L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} _ {l}) \right| \right] \leq C B _ {w} ^ {2} \left[ \frac {\log (2 K)}{N _ {\mathrm {v a l}}} + \sqrt {\frac {\log (2 K)}{N _ {\mathrm {v a l}}}} \right],
$$
where we denote  $\widehat{\mathbf{w}}_l = \widehat{\mathbf{w}}_l(\mathcal{D}_{\mathrm{train}})$
Proof. We only need to work with a fixed pair of  $(\mathbf{w}_{\star}, \mathcal{D}_{\mathrm{train}})$  such that  $\mathcal{E}_{\mathrm{train}}$  holds. Hence, in the following we only consider the randomness of  $\mathcal{D}_{\mathrm{val}}$  conditional on such a  $(\mathbf{w}_{\star}, \mathcal{D}_{\mathrm{train}})$ .
Recall that for any  $\mathbf{w}$
$$
\widehat {L} _ {\mathrm {v a l}} (\mathbf {w}) = \frac {1}{2 | \mathcal {D} _ {\mathrm {v a l}} |} \sum_ {(\mathbf {x} _ {i}, y _ {i}) \in \mathcal {D} _ {\mathrm {v a l}}} \left(\left\langle \mathbf {x} _ {i}, \mathbf {w} \right\rangle - y _ {i}\right) ^ {2},
$$
and we have  $\mathbb{E}_{\mathcal{D}_v}[\widehat{L}_{\mathrm{val}}(\mathbf{w})] = L_{\mathrm{val},\mathbf{w}_*}(\mathbf{w})$ . For each  $i\in \mathcal{I}_{\mathrm{val}}$
$$
y _ {i} - \langle \mathbf {x} _ {i}, \widehat {\mathbf {w}} _ {l} \rangle = \varepsilon_ {i} - \langle \mathbf {x} _ {i}, \mathbf {w} _ {\star} - \widehat {\mathbf {w}} _ {l} \rangle \sim \operatorname {S G} \left(\sigma_ {k} ^ {2} + \| \mathbf {w} _ {\star} - \widehat {\mathbf {w}} _ {l} \| ^ {2}\right).
$$
Note that under  $\mathcal{E}_{\mathrm{train}}$ , we have  $\widehat{\mathbf{w}}_l \in \mathsf{B}_2(B_w)$  for all  $l \in [K]$ , and hence  $\sigma_k^2 + \| \mathbf{w}_\star - \widehat{\mathbf{w}}_l \|^2 \leq 5B_w^2$ . We then have  $(y_i - \langle \mathbf{x}_i, \widehat{\mathbf{w}}_l \rangle)^2$ 's are (conditional) i.i.d random variables in  $\mathrm{SE}(CB_w^4)$ . Then, by Bernstein's inequality, we have
$$
\mathbb {P} _ {\mathcal {D} _ {\mathrm {v a l}}} \left(\left| \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {l}) - L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} _ {l}) \right| \geq t\right) \leq 2 \exp \left(- c N _ {\mathrm {v a l}} \min  \left\{\frac {t ^ {2}}{B _ {w} ^ {2}}, \frac {t}{B _ {w}} \right\}\right),
$$
where  $c$  is a universal constant. Applying the union bound, we obtain
$$
\mathbb {P} _ {\mathcal {D} _ {\mathrm {v a l}}} \left(\max _ {l \in [ K ]} \left| \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {l}) - L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} _ {l}) \right| \geq t\right) \leq K \exp \left(- c N _ {\mathrm {v a l}} \min \left\{\frac {t ^ {2}}{B _ {w} ^ {2}}, \frac {t}{B _ {w}} \right\}\right).
$$
Taking integration completes the proof.
J.4 Generalized linear models with adaptive link function selection
Suppose that  $(g_{k}:\mathbb{R}\to \mathbb{R})_{k\in [K]}$  is a set of link functions such that  $g_{k}$  is non-decreasing and  $C^2$  -smooth for each  $k\in [K]$ . We consider the input format we introduce in Section 4.1 with  $|\mathcal{D}_{\mathrm{train}}| = \lceil N / 2\rceil, |\mathcal{D}_{\mathrm{val}}| = \lfloor N / 2\rfloor$ .
Theorem J.2 (GLMs with adaptive link function selection). For any fixed set of parameters defined in Assumption  $B$ , as long as  $N \geq \mathcal{O}(d)$ , there exists a transformer  $\pmb{\theta}$  with  $L \leq \mathcal{O}(\log(N))$  layers, input dimension  $D = \Theta(dK)$  and  $\max_{\ell \in [L]} M^{(\ell)} \leq \widetilde{\mathcal{O}}(d^3 N)$ , such that the following holds.
For any  $k^{\star} \in [K]$  and any distribution  $\mathsf{P}$  that is a generalized linear model of the link function  $g_{k^{\star}}$  and some parameter  $\beta$ , if Assumption B holds for each pair  $(\mathsf{P}, g_k)$ , then
$$
\mathbb {E} _ {(\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} \big [ (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \big ] \leq \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P}} \big [ (g _ {k ^ {\star}} (\langle \mathbf {x}, \boldsymbol {\beta} \rangle) - y) ^ {2} \big ] + \mathcal {O} \left(\frac {d}{N} + \sqrt {\frac {\log (K)}{N}}\right),
$$
or equivalently,  $\mathbb{E}_{(\mathcal{D},\mathbf{x}_{N + 1})\sim \mathsf{P}}[(\widehat{y}_{N + 1} - \mathbb{E}[y_{N + 1}|\mathbf{x}_{N + 1}]')^2 ]\leq \mathcal{O}\left(d / N + \sqrt{\log(K) / N}\right)$ .
Proof. For each  $k \in [K]$ , we consider optimizing the following training loss:
$$
\mathbf{w}_{\mathrm{GLM}}^{(k)}:= \operatorname *{arg  min}_{\mathbf{w}}\widehat{L}^{(k)}_{\mathrm{train}}(\mathbf{w}):= \frac{1}{N_{\mathrm{train}}}\sum_{(\mathbf{x}_{i},y_{i})\in \mathcal{D}_{\mathrm{train}}}\ell_{k}(\langle \mathbf{x}_{i},\mathbf{w}\rangle ,y_{i}),
$$
where  $\ell_k(t,y)\coloneqq -yt + \int_0^t g_k(s)ds$  is the convex (integral) loss associated with  $g_{k}$  (as in Section 3.1).
Also, for each predictor  $f: \mathbb{R}^d \to \mathbb{R}$ , we consider the squared validation loss  $\widehat{L}_{\mathrm{val}}$ :
$$
\widehat {L} _ {\mathrm {v a l}} (f) := \frac {1}{2 N _ {\mathrm {v a l}}} \sum_ {(\mathbf {x} _ {i}, y _ {i}) \in \mathcal {D} _ {\mathrm {v a l}}} \left(f (\mathbf {x} _ {i}) - y _ {i}\right) ^ {2}.
$$
Fix a large universal constant  $C_0$ . Let us set
$$
\alpha = \mu_ {g} \mu_ {x} / 8, \quad \beta = 8 L _ {g} K _ {x},
$$
$$
B _ {x} = C _ {0} K _ {x} \sqrt {d \log (N)}, \qquad B _ {y} = C _ {0} K _ {y} \sqrt {\log (N)},
$$
$$
\mathcal {E} _ {w} = \left\{\forall k \in [ K ], \forall \mathbf {w} \in \mathsf {B} _ {2} (B _ {w}), \alpha \leq \lambda_ {\min } (\nabla^ {2} \widehat {L} _ {\mathtt {t r a i n}} ^ {(k)} (\mathbf {w})) \leq \lambda_ {\max } (\nabla^ {2} \widehat {L} _ {\mathtt {t r a i n}} ^ {(k)} (\mathbf {w})) \leq \beta , \right\},
$$
$$
\mathcal {E} _ {r} = \left\{\forall k \in [ K ], \| \mathbf {w} _ {\mathrm {G L M}} ^ {(k)} \| _ {2} \leq B _ {w} / 2 \right\},
$$
$$
\mathcal {E} _ {b, N + 1} = \left\{\left\| \mathbf {x} _ {N + 1} \right\| _ {2} \leq B _ {x}, | y _ {N + 1} | \leq B _ {y} \right\}.
$$
Similar to the proof of Theorem G.2 (Appendix G.2), we know the good event  $\mathcal{E} \coloneqq \mathcal{E}_w \cap \mathcal{E}_r \cap \mathcal{E}_{b,\mathrm{train}} \cap \mathcal{E}_{b,\mathrm{test}} \cap \mathcal{E}_{b,N+1}$  holds with high probability:  $\mathsf{P}(\mathcal{E}^c) \leq \mathcal{O}(N^{-10})$ .
Similar to the proof of Theorem I.2, we can show that there exists a transformer  $\theta$  with prediction  $\widehat{y}_{N + 1} = \widetilde{\mathrm{read}}_{\mathrm{y}}(\mathrm{TF}_{\theta}(\mathbf{H}))$  (clipped by  $B_{y}$ ), such that (for any  $\mathsf{P}$ ) the following holds under  $\mathcal{E}$ :
(a) For each  $k \in [K]$ ,  $f_{k} = \mathcal{A}_{k}(\mathcal{D}_{\mathrm{train}})$  is a predictor such that  $\left|f_{k}(\mathbf{x}_{i}) - g_{k}(\langle \mathbf{x}_{i}, \mathbf{w}_{\mathrm{GLM}}^{(k)} \rangle)\right| \leq \varepsilon$  for all  $i \in [N + 1]$  (where  $\varepsilon$  is chosen as in Appendix G.2).
(b)  $\widehat{y}_{N + 1} = \mathrm{clip}_{B_y}(\widehat{f} (\mathbf{x}_{N + 1}))$ , where  $\widehat{f} = \mathcal{A}_{\mathsf{TF}}(\mathcal{D})$  is an aggregated predictor given by  $\widehat{f} = \sum_{k}\lambda_{k}f_{k}$ , such that  $(\lambda_{k})$  is a distribution supported on  $k\in [K]$  such that  $\widehat{L}_{\mathrm{val}}(f_k)\leq \min_{k'\in [K]}\widehat{L}_{\mathrm{val}}(f_{k'}) + \gamma$ .
Similar to the proof of Theorem G.2, for  $\mathcal{E}_0\coloneqq \mathcal{E}_w\cap \mathcal{E}_r\cap \mathcal{E}_{b,\mathrm{train}}\cap \mathcal{E}_{b,\mathrm{test}}$  , we have
$$
\mathbb {E} _ {(\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \leq \mathbb {E} _ {\mathcal {D} \sim \mathsf {P}} \left[ 1 \{\mathcal {E} _ {0} \} L _ {\mathrm {v a l}} (\widehat {f}) \right] + \mathcal {O} \left(\frac {B _ {y} ^ {2}}{N ^ {5}}\right),
$$
where we denote  $L_{\mathrm{val}}(f)\coloneqq \mathbb{E}_{(\mathbf{x},y)\sim \mathsf{P}}\Big[1\{\| \mathbf{x}\| _2\leq B_x\} (f(\mathbf{x}) - y)^2\Big]$  for each predictor  $f$ . By the definition of  $\widehat{f}$ , we then have (under  $\mathcal{E}_0$ )
$$
L _ {\mathrm {v a l}} (\widehat {f}) \leq L _ {\mathrm {v a l}} (f _ {k ^ {\star}}) + \max  _ {l} \left| \widehat {L} _ {\mathrm {v a l}} (f _ {l}) - L _ {\mathrm {v a l}} (f _ {l}) \right| + \gamma .
$$
For the first term, repeating the argument in the proof of Theorem G.2 directly yields that for  $\mathcal{E}_{\mathrm{train}} := \mathcal{E}_w \cap \mathcal{E}_r \cap \mathcal{E}_{b,\mathrm{train}}$ ,
$$
\mathbb {E} _ {\mathcal {D} _ {\text {t r a i n}} \sim \mathrm {P}} \left[ 1 \left\{\mathcal {E} _ {\text {t r a i n}} \right\} L _ {\text {v a l}} \left(f _ {k ^ {*}}\right) \right] \leq \mathbb {E} _ {(\mathbf {x}, y) \sim \mathrm {P}} \left(g _ {k ^ {*}} \left(\left\langle \mathbf {x}, \boldsymbol {\beta} \right\rangle\right) - y\right) ^ {2} + \mathcal {O} \left(d / N _ {\text {t r a i n}}\right).
$$
For the second term, similar to Lemma J.2, we can show that conditional on  $\mathcal{D}_{\mathrm{train}}$  such that  $\mathcal{E}_{\mathrm{train}}$  holds, it holds
$$
\mathbb {E} _ {\mathcal {D} _ {\mathrm {v a l}} \sim \mathsf {P} | \mathcal {D} _ {\mathrm {t r a i n}}} \left[ 1 \{\mathcal {E} _ {0} \} \max _ {l} \left| \widehat {L} _ {\mathrm {v a l}} (f _ {l}) - L _ {\mathrm {v a l}} (f _ {l}) \right| \right] \leq \mathcal {O} \left(K _ {y} ^ {2}\right) \cdot \left(\sqrt {\frac {\log K}{N _ {\mathrm {v a l}}}} + \frac {\log K}{N _ {\mathrm {v a l}}}\right).
$$
Combining these inequalities and suitably choosing  $\gamma$  complete the proof.
K Analysis of pretraining
Thus far, we have established the existence of transformers for performing various ICL tasks with good in-context statistical performance. We now analyze the sample complexity of pretraining these transformers from a finite number of training ICL instances.
K.1 Generalization guarantee for pretraining
Setup At pretraining time, each training ICL instance has form  $\mathbf{Z} := (\mathbf{H}, y_{N+1})$ , where  $\mathbf{H} := \mathbf{H}(\mathcal{D}, \mathbf{x}_{N+1}) \in \mathbb{R}^{D \times (N+1)}$  denote the input sequence formatted as in (3). We consider the square loss between the in-context prediction and the ground truth label:
$$
\ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z}) := \frac {1}{2} \left(y _ {N + 1} - \underbrace {\operatorname {c l i p} _ {B _ {y}} \big (\operatorname {r e a d} _ {\mathbf {y}} (\operatorname {T F} _ {\boldsymbol {\theta}} ^ {R} (\mathbf {H})) \big)} _ {\widetilde {\operatorname {r e a d} _ {\mathbf {y}}}}\right) ^ {2}.
$$
Above,  $\mathrm{clip}_{B_y}(t) \coloneqq \max \left\{\min \left\{t, B_y\right\}, -B_y\right\}$  is the standard clipping operator onto  $[-B_y, B_y]$ , and  $\mathrm{TF}_{\theta}^{R}$  the transformer architecture as in Definition 3 with clipping operators after each layer: let  $\mathbf{H}^{(0)} = \mathrm{clip}_{\mathbb{R}}(\mathbf{H})$
$$
\mathbf {H} ^ {(\ell)} = \operatorname {c l i p} _ {\mathsf {R}} \left(\operatorname {M L P} _ {\boldsymbol {\theta} _ {\operatorname {m i p}} ^ {(\ell)}} \left(\operatorname {A t t n} _ {\boldsymbol {\theta} _ {\operatorname {a t t n}} ^ {(\ell)}} \left(\mathbf {H} ^ {(\ell - 1)}\right)\right)\right) \text {f o r a l l} \ell \in [ L ], \quad \operatorname {c l i p} _ {\mathsf {R}} (\mathbf {H}) := \left[ \operatorname {P r o j} _ {\| \mathbf {h} \| _ {2} \leq \mathsf {R}} (\mathbf {h} _ {i}) \right] _ {i}.
$$
The clipping operator is used to control the Lipschitz constant of  $\mathrm{TF}_{\theta}$  with respect to  $\theta$ , and we typically choose a sufficiently large clipping radius  $R$  so that it does not modify the behavior of the transformer on any input sequence of our concern.