SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Layer 2: Use a single head to compute  $\| \widehat{\mathbf{t}}\| _2^2\colon \mathbf{Q}_1^{(2)}\mathbf{h}_i^{(1)} = [\widehat{\mathbf{t}};\mathbf{0}_{D - d}],\mathbf{K}_1^{(2)}\mathbf{h}_j^{(1)} = [\widehat{\mathbf{t}};\mathbf{0}_{D - d}]$ , and  $\mathbf{V}_1^{(2)}\mathbf{h}_j^{(1)} = [1;\mathbf{0}_{D - 1}]$ . By manipulating the output dimension in  $\mathbf{V}^{(2)}$ , write the result  $\| \widehat{\mathbf{t}}\| _2^2$  into blank memory space with dimension 1 at every token  $i\in [N + 1]$ . After layer 2, we have  $\mathbf{h}_i^{(3)} = [\mathbf{x}_i;y_it_i;*;\| \widehat{\mathbf{t}}\| _2^2;*;1]$ .
Layer 3: Use 2 heads to implement two ReLU functions with bias:  $\| \widehat{\mathbf{t}} \|_2^2 \mapsto \frac{1}{B - A} (\sigma(\| \widehat{\mathbf{t}} \|_2^2 - A) - \sigma(\| \widehat{\mathbf{t}} \|_2^2 - B))$ . The two query (or key) matrices contain values  $A$  and  $B$ . In our problem we take
$$
A = (\lambda_ {\min} B _ {w} ^ {\star} / 4) ^ {2}, \quad B = (3 \lambda_ {\min} B _ {w} ^ {\star} / 4) ^ {2},
$$
so that the above ReLU function implements  $\Psi^{\mathrm{lin}}(\mathcal{D})$  exactly. Write the result into a blank memory space with dimension 1. We finish the proof by noting that  $\| \pmb {\theta}\| \leq \mathcal{O}(1 + \lambda_{\min}^2 (B_w^\star)^2)$ .
Statistical guarantee for  $\Psi^{\mathrm{lin}}$  We consider the following well-posedness assumption for the linear correlation test  $\Psi^{\mathrm{lin}}$ . Note that, similar as Assumption A, the assumption does not require the data to be generated from any true linear model, but rather only requires some properties about the best linear fit  $\mathbf{w}_{\mathbb{P}}^{\star}$ , as well as sub-Gaussianity conditions.
Assumption D (Well-posedness for linear correlation test). We say a distribution  $\mathsf{P}$  on  $\mathbb{R}^d\times \mathbb{R}$  is well-posed for linear independence tests, if  $(\mathbf{x},y)\sim \mathsf{P}$  satisfies
(1)  $\| \mathbf{x}\| _2\leq B_x$  and  $|y|\leq B_y$  almost surely;
(2) The covariance  $\pmb{\Sigma}_{\mathbb{P}} \coloneqq \mathbb{E}_{\mathbb{P}}[\mathbf{x}\mathbf{x}^{\top}]$  satisfies  $\lambda_{\min}\mathbf{I}_d \preceq \pmb{\Sigma}_{\mathbb{P}} \preceq \lambda_{\max}\mathbf{I}_d$ , with  $0 < \lambda_{\min} \leq \lambda_{\max}$ , and  $\kappa \coloneqq \lambda_{\max} / \lambda_{\min}$ .
(3) The whitened vector  $\pmb{\Sigma}_{\mathbb{P}}^{-1 / 2}\mathbf{x}$  is  $K^2$  -sub-Gaussian for some  $K\geq 1$
(4) The best linear predictor  $\mathbf{w}_{\mathsf{P}}^{\star} \coloneqq \mathbb{E}_{\mathsf{P}}[\mathbf{x}\mathbf{x}^{\top}]^{-1}\mathbb{E}_{\mathsf{P}}[\mathbf{x}y]$  satisfies  $\| \mathbf{w}_{\mathsf{P}}^{\star}\|_{2} \leq \overline{B_{w}^{\star}}$ .
(5) The label  $y$  is  $\sigma^2$ -sub-Gaussian.
(6) The residual  $z\coloneqq y - \langle \mathbf{x},\mathbf{w}_{\mathsf{P}}^{\star}\rangle$  is  $\sigma^2$  -sub-Gaussian with probability one (over  $\mathbf{x}$ ).
The following results state that  $\Psi^{\mathrm{lin}}$  achieves high power as long as the sample size is high enough, and the signal  $\| \mathbf{w}_{\mathbb{P}}^{\star}\|_{2}$  is either sufficiently high or sufficiently low.
Proposition 1.5 (Power of linear correlation test). Suppose distribution  $\mathsf{P}$  satisfies Assumption  $D$  with parameters  $\lambda_{\mathrm{min}},\lambda_{\mathrm{max}},\overline{B_w^{\star}}$ . Then, for the linear correlation test  $\Psi^{\mathrm{lin}}$  with parameters  $(\lambda_{\mathrm{min}},B_w^{\star})$  with  $B_w^\star \leq \overline{B_w^\star}$  and any  $N\geq \widetilde{\mathcal{O}}\Big(\max \left\{K^4,\frac{\lambda_{\mathrm{max}}K^2\sigma^2}{(B_w^\star)^2\lambda_{\mathrm{min}}^2}\right\} \cdot d\Big)$ , we have
1. If  $\| \mathbf{w}_{\mathsf{P}}^{\star}\|_{2}\geq B_{w}^{\star}$ , then with probability at least  $1 - \delta$  over  $\mathcal{D}$ , we have  $\Psi^{\mathrm{lin}}(\mathcal{D}) = 1$ .
2. If  $\| \mathbf{w}_{\mathsf{P}}^{\star}\|_{2}\leq \frac{\lambda_{\min}}{10\lambda_{\max}} B_{w}^{\star}$ , then with probability at least  $1 - \delta$  over  $\mathcal{D}$ , we have  $\Psi^{\mathrm{lin}}(\mathcal{D}) = 0$ .
Proof. For any  $\mathbb{P}$  satisfying Assumption D, note that  $\mathbb{E}[\mathbf{x}z] = \mathbb{E}[\mathbf{x}(y - \langle \mathbf{w}_{\mathbb{P}}^{\star},\mathbf{x}\rangle)] = \mathbf{0}$  by construction. Therefore, by standard sub-Gaussian and sub-exponential concentration combined with union bound, the following events hold simultaneously with probability at least  $1 - \delta$ :
$$
\begin{array}{l} 0. 9 \boldsymbol {\Sigma} _ {\mathrm {P}} \preceq \widehat {\boldsymbol {\Sigma}} = \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} \mathbf {x} _ {i} ^ {\top} \preceq 1. 1 \boldsymbol {\Sigma} _ {\mathrm {P}} \quad \text {a s} N \geq \widetilde {\mathcal {O}} (d K ^ {4}) \text {b y} (2 7), \\ \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} z _ {i} \right\| _ {2} \leq \lambda_ {\max } ^ {1 / 2} \cdot \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \boldsymbol {\Sigma} _ {\mathsf {P}} ^ {- 1 / 2} \mathbf {x} _ {i} z _ {i} \right\| _ {2} \leq \widetilde {\mathcal {O}} \left(\lambda_ {\max } ^ {1 / 2} \left(\frac {K \sigma \sqrt {d}}{\sqrt {N}} + \frac {K \sigma d}{N}\right)\right) \\ \leq \widetilde {\mathcal {O}} \left(\lambda_ {\max} ^ {1 / 2} K \sigma \sqrt {\frac {d}{N}}\right) \leq \frac {\lambda_ {\min} B _ {w} ^ {\star}}{8}, \mathrm {a s} N \geq \widetilde {\mathcal {O}} \left(\frac {\lambda_ {\max} d K ^ {2} \sigma^ {2}}{(B _ {w} ^ {\star}) ^ {2} \lambda_ {\min} ^ {2}}\right). \\ \end{array}
$$
On the above event, we have
$$
\left\| \widehat {\mathbf {t}} \right\| _ {2} = \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} \left(\left\langle \mathbf {x} _ {i}, \mathbf {w} _ {\mathsf {P}} ^ {\star} \right\rangle + \mathbf {z} _ {i}\right) \right\| _ {2} = \left\| \widehat {\boldsymbol {\Sigma}} \mathbf {w} _ {\mathsf {P}} ^ {\star} + \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} z _ {i} \right\| _ {2}.
$$
Therefore, in case 1, we have
$$
\left\| \widehat {\mathbf {t}} \right\| _ {2} \geq \left\| \widehat {\boldsymbol {\Sigma}} \mathbf {w} _ {\mathsf {P}} ^ {\star} \right\| _ {2} - \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} z _ {i} \right\| _ {2} \geq 0. 9 \lambda_ {\min} \| \mathbf {w} _ {\mathsf {P}} ^ {\star} \| _ {2} - \frac {\lambda_ {\min} B _ {w} ^ {\star}}{8} \geq \frac {3 \lambda_ {\min} B _ {w} ^ {\star}}{4}.
$$
In case 2, we have
$$
\left\| \widehat {\mathbf {t}} \right\| _ {2} \leq \left\| \widehat {\boldsymbol {\Sigma}} \mathbf {w} _ {\mathsf {P}} ^ {\star} \right\| _ {2} + \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} z _ {i} \right\| _ {2} \leq \lambda_ {\max} \cdot \frac {\lambda_ {\min} B _ {w} ^ {\star}}{1 0 \lambda_ {\max}} + \frac {\lambda_ {\min} B _ {w} ^ {\star}}{8} \leq \frac {\lambda_ {\min} B _ {w} ^ {\star}}{4}.
$$
The proof is finished by recalling the definition of  $\Psi^{\mathrm{lin}}$  in (45), so that  $\Psi^{\mathrm{lin}}(\mathcal{D}) = 1$  if  $\| \widehat{\mathbf{t}} \|_2 \geq 3\lambda_{\min} B_w^\star / 4$ , and  $\Psi^{\mathrm{lin}}(\mathcal{D}) = 0$  if  $\| \widehat{\mathbf{t}} \|_2 \leq \lambda_{\min} B_w^\star / 4$ .
Application: Confident linear regression By directly composing the linear correlation test in Lemma I.1 with the transformer construction in Corollary 5 (using an argument similar as the proof of Proposition I.4), and using the power of the linear correlation test Proposition I.5, we immediately obtain the following result, which outputs a prediction from (approximately) least squares if  $\widehat{\psi} \coloneqq \Psi^{\mathrm{lin}}(\mathcal{D}) = 1$ , and abstains from predicting if  $\widehat{\psi} = 0$ . This can be viewed as a form of "confident linear regression", where the model predicts only if it thinks the linear signal is strong enough.
Proposition 1.6 (Confident linear regression). For any  $B_w > 0$ ,  $0 < B_w^* \leq \overline{B_w^*}$ ,  $0 \leq \lambda_{\min} \leq \lambda_{\max}$ ,  $\varepsilon \leq B_x B_w / 10$ ,  $0 < \alpha \leq \beta$  with  $\kappa := \beta / \alpha$ , there exists a L-layer attention-only transformer with
$$
L \leq \mathcal {O} \left(\kappa \log \frac {B _ {x} B _ {w}}{\varepsilon}\right), \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq \mathcal {O} (1), \quad \| \boldsymbol {\theta} \| \leq \mathcal {O} \left(R + \frac {1}{\beta} + \lambda_ {\min } ^ {2} (B _ {w} ^ {\star}) ^ {2}\right)
$$
(with  $R \coloneqq \max \{B_x B_w, B_y, 1\}$ ) such that the following holds. Let  $N \geq \widetilde{\mathcal{O}}\left(\max \left\{K^4, \frac{\lambda_{\max} K^2 \sigma^2}{(B_w^\star)^2 \lambda_{\min}^2}\right\} \cdot d\right)$ . Suppose the input format is (3) with dimension  $D \geq 2d + 4$ . Let ICL instance  $(\mathcal{D}, \mathbf{x}_{N+1})$  be drawn from any distribution  $\mathsf{P}$  satisfying Assumption  $D$ . Then the transformer outputs a 2-dimensional prediction (within the test token  $\widetilde{\mathbf{h}}_{N+1}$ )
$$
\left(\widehat {y} _ {N + 1}, \widehat {\psi}\right) \in \mathbb {R} \times \{0, 1 \}
$$
such that the following holds:
1. If  $\| \mathbf{w}_{\mathsf{P}}^{\star}\|_{2}\geq B_{w}^{\star}$ , then with probability at least  $1 - \delta$  over  $\mathcal{D}$ , we have  $|\widehat{y}_{N + 1} - \langle \widehat{\mathbf{w}}_{\mathrm{LS}},\mathbf{x}_{N + 1}\rangle |\leq \varepsilon$ , and  $\widehat{\psi} = 1$  if  $\mathcal{D}$  is in addition well-conditioned for least squares (in the sense of (5) with  $\lambda = 0$ ).
2. If  $\| \mathbf{w}_{\mathsf{P}}^{\star}\|_{2}\leq \frac{\lambda_{\min}}{10\lambda_{\max}} B_{w}^{\star}$ , then with probability at least  $1 - \delta$  over  $\mathcal{D}$ , we have  $\widehat{y}_{N + 1} = 0$  and  $\widehat{\psi} = 0$ .
J Proof of Theorem 12: Noisy linear model with mixed noise levels
For each fixed  $k \in [K]$ , we consider the following data generating model  $\mathbb{P}_k$ , where we first sample  $\mathsf{P} = \mathsf{P}_{\mathbf{w}_{\star}, \sigma_k} \sim \pi$  from  $\mathbf{w}_{\star} \sim \mathsf{N}(\mathbf{0}, \mathbf{I}_d / d)$ , and then sample data  $\{(\mathbf{x}_i, y_i)\}_{i \in [N + 1]}$ $\stackrel{\mathrm{iid}}{\sim} \mathsf{P}_{\mathbf{w}_{\star}, \sigma_k}$  as
$$
\mathrm {P} _ {\mathbf {w} _ {\star}, \sigma_ {k}}: \mathbf {x} _ {i} \sim \mathrm {N} (\mathbf {0}, \mathbf {I} _ {d}), \quad y _ {i} = \left\langle \mathbf {x} _ {i}, \mathbf {w} _ {\star} \right\rangle + \varepsilon_ {i}, \quad \varepsilon_ {i} \sim \mathrm {N} (0, \sigma_ {k} ^ {2}).
$$
Also, recall that the Bayes optimal estimator on  $\mathbb{P}_k$  is given by  $\hat{y}_{N + 1}^{\mathrm{Bayes}} = \left\langle \mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}),\mathbf{x}_{N + 1}\right\rangle$  with ridge  $\lambda_{k} = \sigma_{k}^{2}d / N$ , and the Bayes risk on  $\mathbb{P}_k$  is given by
$$
\mathsf {B a y e s R i s k} _ {k} := \inf  _ {\mathcal {A}} \mathbb {E} _ {k} \left[ \frac {1}{2} (\mathcal {A} (\mathcal {D}) (\mathbf {x} _ {N + 1}) - y _ {N + 1}) ^ {2} \right] = \mathbb {E} _ {k} \left[ \frac {1}{2} \left(\widehat {y} _ {N + 1} ^ {\mathsf {B a y e s}} - y _ {N + 1}\right) ^ {2} \right].
$$
Recall that in Section 4.1.1, we consider a mixture law  $\mathbb{P}_{\pi}$  that generates data from  $\mathbb{P}_k$  with  $k\sim \Lambda$ . It is clear that we have (pushing  $\inf_{\mathcal{A}}$  into  $\mathbb{E}_{k\sim \Lambda}$  does not increase the value) we have
$$
\text {B a y e s R i s k} _ {\pi} \geq \mathbb {E} _ {k \sim \Lambda} [ \text {B a y e s R i s k} _ {k} ],
$$
i.e., the Bayes risk can only be greater if we consider a mixture of models. In other words, if a transformer can achieve near-Bayes ICL on each meta-task  $\mathbb{P}_k$ , then it can perform near-Bayes ICL on any meta-task  $\pi$  which is a mixture of  $\mathbb{P}_k$  with  $k \sim \Lambda$ . Therefore, to prove Theorem 12, it suffices to show the following (strengthened) result.
Theorem J.1 (Formal version of Theorem 12). Suppose that  $N \geq 0.1d$  and we write  $\sigma_{\max} = \max_k \{\sigma_k, 1\}$ ,  $\sigma_{\min} = \min_k \{\sigma_k, 1\}$ . Suppose in input format (3) we have  $D \geq \Theta(Kd)$ . Then there exists a transformer  $\theta$  with
$$
L \leq \mathcal {O} \left(\sigma_ {\min } ^ {- 2} \log (N / \sigma_ {\min })\right), \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq \mathcal {O} \left(K\right), \quad \max  _ {\ell \in [ L ]} D ^ {(\ell)} \leq \mathcal {O} (K ^ {2}),
$$
$$
\left\| \boldsymbol {\theta} \right\| \leq \mathcal {O} \left(\sigma_ {\max } K d \log (N)\right),
$$
such that for any  $k\in [K]$  , it holds that
$$
\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} + \widetilde {\mathcal {O}} \left(\frac {\sigma_ {\max } ^ {2}}{\sigma_ {\min } ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right)
$$
if we choose  $N_{\mathrm{val}} \coloneqq |\mathcal{D}_{\mathrm{val}}| \asymp N^{2/3}[\log K]^{1/3}$ .
The core of the proof of Theorem J.1 is to show that any estimator  $\widehat{\mathbf{w}}$  that achieves small validation loss  $\widehat{L}_{\mathrm{val}}$  must achieve small population loss.
Throughout the rest of this section, recall that we define  $N_{\mathrm{train}} = |\mathcal{D}_{\mathrm{train}}|, N_{\mathrm{val}} = |\mathcal{D}_{\mathrm{val}}|, \mathcal{I}_{\mathrm{train}} = \{i : (\mathbf{x}_i, y_i) \in \mathcal{D}_{\mathrm{train}}\}$ ,  $\mathcal{I}_{\mathrm{val}} = \{i : (\mathbf{x}_i, y_i) \in \mathcal{D}_{\mathrm{val}}\}$ , and  $\mathbf{X}_{\mathrm{train}} = [\mathbf{x}_i]_{i \in \mathcal{I}_{\mathrm{train}}}$ .
J.1 Proof of Theorem J.1
Fix parameters  $\delta, \varepsilon, \gamma > 0$  and a large universal constant  $C_0$ . Let us set
$$
\alpha = \max  \left\{0, 1 / 2 - \sqrt {d / N _ {\mathrm {t r a i n}}} \right\} ^ {2}, \quad \beta = 2 5,
$$
$$
B _ {w} ^ {\star} = 1 + C _ {0} \sqrt {\frac {\log (N)}{d}}, \qquad B _ {w} = C _ {0} (B _ {w} ^ {\star} + \sigma_ {\max} / \sigma_ {\min}),
$$
$$
B _ {x} = C _ {0} \sqrt {d \log (N)}, \qquad B _ {y} = C _ {0} (B _ {w} ^ {\star} + \sigma_ {\max}) \sqrt {\log (N)},
$$
Then, we define good events similarly to the proof of Corollary 6 (Appendix F.4):
$$
\mathcal {E} _ {\pi} = \left\{\| \mathbf {w} _ {\star} \| _ {2} \leq B _ {w} ^ {\star}, \| \varepsilon \| _ {2} \leq 2 \sigma_ {\max } \sqrt {N} \right\},
$$
$$
\mathcal {E} _ {w} = \left\{\alpha \leq \lambda_ {\min } \left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} / N _ {\text {t r a i n}}\right) \leq \lambda_ {\max } \left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} / N _ {\text {t r a i n}}\right) \leq \beta \right\},
$$
$$
\mathcal {E} _ {b, \text {t r a i n}} = \left\{\forall \left(\mathbf {x} _ {i}, y _ {i}\right) \in \mathcal {D} _ {\text {t r a i n}}, \| \mathbf {x} _ {i} \| _ {2} \leq B _ {x}, | y _ {i} | \leq B _ {y} \right\},
$$
$$
\mathcal {E} _ {b, \mathrm {v a l}} = \left\{\forall \left(\mathbf {x} _ {i}, y _ {i}\right) \in \mathcal {D} _ {\mathrm {v a l}}, \| \mathbf {x} _ {i} \| _ {2} \leq B _ {x}, | y _ {i} | \leq B _ {y} \right\},
$$
$$
\mathcal {E} _ {b, N + 1} = \left\{\| \mathbf {x} _ {N + 1} \| _ {2} \leq B _ {x}, | y _ {N + 1} | \leq B _ {y} \right\}.
$$
For the good event  $\mathcal{E} \coloneqq \mathcal{E}_{\pi} \cap \mathcal{E}_w \cap \mathcal{E}_{b,\text{train}} \cap \mathcal{E}_{b,\text{test}} \cap \mathcal{E}_{b,N+1}$ , we can show that  $\mathbb{P}(\mathcal{E}^c) \leq \mathcal{O}(N^{-10})$ . Further, by the proof of Lemma F.1 (see e.g. (34)), we know that  $\max_{k \in [K]} \left\| \mathbf{w}_{\text{ridge}}^{\lambda_k}(\mathcal{D}_{\text{train}}) \right\|_2 \leq B_w / 2$  holds under the good event  $\mathcal{E}$ .
For the ridge  $\lambda_{k} = \frac{d\sigma_{k}^{2}}{N_{\mathrm{train}}}$  and parameters  $(\alpha, \beta, \gamma, \varepsilon)$ , we consider the transformer  $\theta$  constructed in Theorem I.2, with a clipped prediction  $\widehat{y}_{N+1} = \widetilde{\mathrm{read}}_{\mathbf{y}}(\mathrm{TF}_{\theta}(\mathbf{H}))$ .
In the following, we upper bound the quantity  $\mathbb{E}_k(\widehat{y}_{N + 1} - y_{N + 1})^2$  for any fixed  $k$ . Similar to the proof of Corollary 6 (Appendix F.4), we decompose
$$
\mathbb {E} _ {k} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} = \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] + \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} ^ {c} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right],
$$
Part I. Recall that by our construction, when  $\mathcal{E}$  holds, we have  $\widehat{y}_{N + 1} = \mathrm{clip}_{B_y}(\langle \widehat{\mathbf{w}},\mathbf{x}_{N + 1}\rangle)$  and the statements of Theorem I.2 hold for  $\widehat{\mathbf{w}}$ . Thus, we have
$$
\begin{array}{l} \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] = \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} \left(\operatorname {c l i p} _ {B _ {y}} \left(\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} \rangle\right) - y _ {N + 1}\right) ^ {2} \right] \\ \leq \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} \left(\left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} \right\rangle - y _ {N + 1}\right) ^ {2} \right]. \\ \end{array}
$$
Let us consider the following risk functional
$$
L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\mathbf {w}) = \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P} _ {\mathbf {w} _ {\star}, \sigma_ {k}}} \left[ \frac {1}{2} \big (\langle \mathbf {w}, \mathbf {x} \rangle - y \big) ^ {2} \right] = \frac {1}{2} \Big (\| \mathbf {w} - \mathbf {w} _ {\star} \| _ {2} ^ {2} + \sigma_ {k} ^ {2} \Big).
$$
Then, under the good event  $\mathcal{E}_0\coloneqq \mathcal{E}_{\pi}\cap \mathcal{E}_w\cap \mathcal{E}_{b,\mathrm{train}}\cap \mathcal{E}_{b,\mathrm{test}}$  of  $(\mathbf{w}_{\star},\mathcal{D})$
$$
\begin{array}{l} \mathbb {E} _ {k} \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} _ {k} \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} _ {k} \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} _ {(\mathbf {x}, y) \sim \mathcal {P} _ {\mathbf {w} _ {\star}, \sigma_ {k}}} \left[ \left(\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} (\mathcal {D}) \rangle - y _ {N + 1}\right) ^ {2} \right] \\ = L _ {\text {v a l}, \mathbf {w} _ {*}} (\widehat {\mathbf {w}} (\mathcal {D})). \\ \end{array}
$$
By our construction, under the good event  $\mathcal{E}_0$ , we have
$$
L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} (\mathcal {D})) \leq L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} _ {k} (\mathcal {D} _ {\mathrm {t r a i n}})) + \max  _ {l \in [ K ]} \left| \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {l} (\mathcal {D} _ {\mathrm {t r a i n}})) - L _ {\mathrm {v a l}, \mathbf {w} _ {\star}} (\widehat {\mathbf {w}} _ {l} (\mathcal {D} _ {\mathrm {t r a i n}})) \right| + \gamma ,
$$
where  $\left\| \widehat{\mathbf{w}}_l(\mathcal{D}_{\mathrm{train}})\right) - \mathbf{w}_{\mathrm{ridge}}^{\lambda_l}(\mathcal{D}_{\mathrm{train}})\bigg\| _2\leq \varepsilon$  for each  $l\in [K]$ . Clearly,
$$
\begin{array}{l} 2 \mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {0} \right\} L _ {\text {v a l}, \mathbf {w} _ {\star}} \left(\widehat {\mathbf {w}} _ {k} \left(\mathcal {D} _ {\text {t r a i n}}\right)\right) \right] = \mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {0} \right\} \left(\left\| \widehat {\mathbf {w}} _ {k} \left(\mathcal {D} _ {\text {t r a i n}}\right) - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} + \sigma_ {k} ^ {2}\right) \right] \\ \leq \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} _ {0} \} \left(\left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda_ {k}} (\mathcal {D} _ {\text {t r a i n}}) - \mathbf {w} _ {\star} \right\| _ {2} ^ {2} + 2 \varepsilon \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda_ {k}} (\mathcal {D} _ {\text {t r a i n}}) - \mathbf {w} _ {\star} \right\| _ {2} + \varepsilon^ {2}\right) \right] + \sigma_ {k} ^ {2} \\ \leq \mathbb {E} _ {k} \left[ \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} + 2 \varepsilon \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} + \varepsilon^ {2} \right] + \sigma_ {k} ^ {2} \\ \leq 2 \operatorname {R i s k} _ {k, \text {t r a i n}} + 2 \varepsilon \sqrt {2 \operatorname {R i s k} _ {k , \text {t r a i n}}} + \varepsilon^ {2}, \\ \end{array}
$$
where we denote  $2\mathrm{Risk}_{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$ , and we also note that  $\mathrm{Risk}_{k,\mathrm{train}} \leq 1 + \sigma_k^2$  by definition. By Lemma J.1, we have
$$
\operatorname {R i s k} _ {k, \text {t r a i n}} \leq \text {B a y e s R i s k} _ {k} + \mathcal {O} \left(\left(\sigma_ {k} ^ {2} + 1\right) \frac {N _ {\text {v a l}}}{N}\right).
$$
We next deal with the term  $\varepsilon_{\mathrm{val}} := \max_{l \in [K]} \left| \widehat{L}_{\mathrm{val}}(\widehat{\mathbf{w}}_l(\mathcal{D}_{\mathrm{train}})) - L_{\mathrm{val}, \mathbf{w}_\star}(\widehat{\mathbf{w}}_l(\mathcal{D}_{\mathrm{train}})) \right|$ . Note that for the good event  $\mathcal{E}_{\mathrm{train}} := \mathcal{E}_\pi \cap \mathcal{E}_w \cap \mathcal{E}_{b,\mathrm{train}}$  of  $(\mathbf{w}_\star, \mathcal{D}_{\mathrm{train}})$ , we have
$$
\mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {0} \right\} \varepsilon_ {\text {v a l}} \right] \leq \mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {\text {t r a i n}} \right\} \varepsilon_ {\text {v a l}} \right] \leq \mathbb {E} _ {\mathbf {w} _ {\star}, \mathcal {D} _ {\text {t r a i n}} \sim \mathbb {P} _ {k}} \left[ 1 \left\{\mathcal {E} _ {\text {t r a i n}} \right\} \cdot \mathbb {E} _ {\mathcal {D} _ {\text {v a l}}} \left[ \varepsilon_ {\text {v a l}} \mid \mathbf {w} _ {\star}, \mathcal {D} _ {\text {t r a i n}} \right] \right].
$$
Thus, Lemma J.2 yields
$$
\mathbb {E} _ {k} \left[ 1 \left\{\mathcal {E} _ {0} \right\} \varepsilon_ {\mathrm {v a l}} \right] \leq \mathcal {O} \left(B _ {w} ^ {2}\right) \cdot \left[ \sqrt {\frac {\log K}{N _ {\mathrm {v a l}}}} + \frac {\log K}{N _ {\mathrm {v a l}}} \right].
$$
$$
\mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \right] \leq 2 \text {B a y e s R i s k} _ {k} + \mathcal {O} \left(\varepsilon \sigma_ {\max } + \varepsilon^ {2} + \frac {\sigma_ {\max } ^ {2} N _ {\mathrm {v a l}}}{N} + B _ {w} ^ {2} \sqrt {\frac {\log K}{N _ {\mathrm {v a l}}}} + \frac {B _ {w} ^ {2} \log K}{N _ {\mathrm {v a l}}}\right).
$$
Therefore, we can choose  $(\varepsilon ,N_{\mathrm{val}})$  so that  $N_{\mathrm{val}}\leq N / 2$  as
$$
N _ {\mathrm {v a l}} = \max  \left\{\left(\frac {B _ {w} ^ {2}}{\sigma_ {\mathrm {m a x}} ^ {2}} N\right) ^ {2 / 3} \log^ {1 / 3} (K), \log K \right\}, \qquad \varepsilon = \frac {\sigma_ {\mathrm {m a x}}}{N}.
$$
It is worth noting that such choice of  $N_{\mathrm{val}}$  is feasible as long as  $N \gtrsim \frac{B_w^4}{\sigma_{\max}^4} \log K$ . Under such choice, we obtain
$$
\frac {1}{2} \mathbb {E} _ {k} \left[ 1 \{\mathcal {E} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {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).
$$
Part II. Similar to the proof of Corollary 6, we have
$$
\mathbb {E} \left[ 1 \{\mathcal {E} ^ {c} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \right] \leq \mathcal {O} \left(\frac {B _ {y} ^ {2}}{N ^ {5}}\right) \leq \mathcal {O} \left(\frac {\sigma_ {\max } ^ {2}}{N ^ {4}}\right).
$$
Conclusion. Combining the both cases, we obtain