SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Before stating the theorem, we introduce the following simplified notation: let  $A \coloneqq \{A_i\}_{i=0}^{L-1} \in \mathbb{R}^{L \times d \times d}$ . We use  $f(A)$  to denote the in-context loss of  $f\left(\{P_i, Q_i\}_{i=0}^{L-1}\right)$  as defined in (5), when  $Q_i$  depends on  $A_i$ , and  $P_i$  is a constant matrix, as described in (8).
Theorem 3. Assume that  $x^{(i)} \stackrel{iid}{\sim} \mathcal{N}(0, \Sigma)$  and  $w_{\star} \sim \mathcal{N}(0, \Sigma^{-1})$ , for  $i = 1, \dots, n$ , and for some  $\Sigma \succ 0$ . Consider the optimization of in-context loss for a  $k$ -layer transformer with the parameter configuration in Eq. (8) given by:
$$
\min_{\substack{\{A_{i}\}_{i = 0}^{L - 1}}}f\left(A\right).
$$
Let  $\mathcal{S} \subset \mathbb{R}^{L \times d \times d}$  be defined as follows:  $A \in \mathcal{S}$  if and only if for all  $i = 0, \dots, L - 1$ , there exists scalars  $a_i \in \mathbb{R}$  such that  $A_i = a_i \Sigma^{-1}$ . Then
$$
\inf  _ {(A, B) \in \mathcal {S}} \sum_ {i = 0} ^ {L - 1} \| \nabla_ {A _ {i}} f (A, B) \| _ {F} ^ {2} = 0,
$$
where  $\nabla_{A_i}f$  denotes derivative wrt the Frobenius norm  $\| A_i\| _F$ .
As discussed in the motivation above, under the setting of  $A_{i} = a_{i}\Sigma^{-1}$ , the linear transformer implements an algorithm that is reminiscent of Newton's method (as well as a number of other adaptive algorithms such as the full-matrix variant of Adagrad); these can converge significantly faster than vanilla gradient descent when the problem is ill-conditioned. The proposed parameters  $A_{i}$  in Theorem 3 are also similar to  $A_{i}$ 's in Theorem 1 when  $n$  is large. However, in contrast to Theorem 1, there is no trade-off with statistical robustness; this is because  $w_{\star}$  has covariance matrix  $\Sigma^{-1}$  in the Theorem 3, while Theorem 1 has isotropic  $w_{\star}$ .
Unlike our prior results, Theorem 3 only guarantees that the set  $S$  of transformer prameters satisfying  $\left\{A_i \propto \Sigma^{-1}\right\}_{i=0}^{L-1}$  essentially contains critical points of the in-context loss. However, in the next section, we show experimentally that this choice of  $A_i$ 's does indeed seem to be recovered by training.
We defer the proof of Theorem 3 to Subsection B.2. Due to the complexity of the transformer function, even verifying critical points can be challenging. We show that the in-context loss can be equivalently written as (roughly) a matrix polynomial involving the weights at each layer. By
<sup>2</sup>A subtle issue is that the infimum may not be attained, so it is possible that  $S$  contains points with arbitrarily small gradient, but does not contain a point with exactly 0 gradient.
exploiting invariances in the underlying distribution of prompts, we construct a flow, contained entirely in  $S$ , whose objective value decreases as fast as gradient flow. Since  $f$  is lower bounded, we conclude that there must be points in  $S$  whose gradient is arbitrarily small.
4.3 Experimental validations for Theorem 3
We present here an empirical verification of our results in Theorem 3. We consider the ICL loss for linear regression. The dimension is  $d = 5$ , and the number of training samples in the prompt is  $n = 20$ . Both  $x^{(i)} \sim \mathcal{N}(0, \Sigma)$  and  $w_{\star} \sim \mathcal{N}(0, \Sigma^{-1})$ , where  $\Sigma = U^T D U$ , where  $U$  is a uniformly random orthogonal matrix, and  $D$  is a fixed diagonal matrix with entries  $(1, 1, 0.25, 0.0625, 1)$ .
We optimize  $f$  for a three-layer linear transformer using ADAM, where the matrices  $A_0, A_1$ , and  $A_2$  are initialized by i.i.d. Gaussian matrices. Each gradient step is computed from a minibatch of size 20000, and we resample the minibatch every 100 steps. We clip the gradient of each matrix to 0.01. All plots are averaged over 5 runs with different  $U$  (i.e.  $\Sigma$ ) sampled each time.
Figure 1d plots the average loss. We observe that the training converges to an almost 0 value, suggesting the convergence to global minimum. The parameters at convergence match the stationary point introduced in Theorem 3, and indeed appear to be globally optimal.
To quantify the similarity between  $A_0, A_1, A_2$  and  $\Sigma^{-1}$  (up to scaling), we use the normalized Frobenius norm distance:  $\mathrm{Dist}(M, I) \coloneqq \min_{\alpha} \frac{\| M - \alpha \cdot I \|}{\| M \|_F}$ , (equivalent to choosing  $\alpha \coloneqq \frac{1}{d} \sum_{i=1}^{d} M[i, i]$ ). This is essentially the projection distance of  $\frac{M}{\| M \|_F}$  onto the space of scaled identity matrices.
We plot  $\operatorname{Dist}(A_i, I)$ , averaged over 5 runs, against iteration in Figures 1a, 1b, 1c. In each plot, the blue line represents  $\operatorname{Dist}(\Sigma^{1/2} A_i \Sigma^{1/2}, I)$ , and we verify that the optimal parameters are converging to the critical point introduced in Theorem 3, which implements preconditioned gradient descent. The red line represents  $\operatorname{Dist}(A_i, I)$ ; it remains constant indicating that the trained transformer is not implementing plain gradient descent. Figures 2a–2c visualize each  $\Sigma^{1/2} A_i \Sigma^{1/2}$  matrix at the end of training to further validate that the learned parameter is as described in Theorem 3.
(a)  $\operatorname{Dist}(\Sigma^{1/2}A_0\Sigma^{1/2}, I)$
Figure 1: Plots for verifying convergence of general linear transformer, defined in Theorem 3. Figure (d) shows convergence of loss to 0. Figures (a),(b),(c) illustrate convergence of  $A_{i}$ 's to identity. More specifically, the blue line represents  $\mathrm{Dist}(\Sigma^{1/2}A_i\Sigma^{1/2},I)$ , which measures the convergence to the critical point introduced in Theorem 3 (corresponding to  $\Sigma^{-1}$ -preconditioned gradient descent). The red line represents  $\mathrm{Dist}(A_i,I)$ ; it remains constant indicating that the trained transformer is not implementing plain gradient descent.
(b)  $\mathrm{Dist}(\Sigma^{1 / 2}A_1\Sigma^{1 / 2},I)$
(c)  $\operatorname{Dist}(\Sigma^{1/2}A_2\Sigma^{1/2}, I)$
(d)  $\log (\mathrm{Loss})$
5 Multi-layer transformers beyond standard optimization methods
In this section, we study the more general setting of
$$
P _ {i} = \left[ \begin{array}{l l} B _ {i} & 0 \\ 0 & 1 \end{array} \right], \quad Q _ {i} = \left[ \begin{array}{l l} A _ {i} & 0 \\ 0 & 0 \end{array} \right] \quad \text {w h e r e} A _ {i}, B _ {i} \in \mathbb {R} ^ {d \times d}. \tag {11}
$$
Note that  $A_i, B_i$  are not constrained to be symmetric. Similar to Section 4, we introduce the following simplified notation: let  $A := \{A_i\}_{i=0}^{L-1} \in \mathbb{R}^{L \times d \times d}$  and  $B := \{B_i\}_{i=0}^{L-1} \in \mathbb{R}^{L \times d \times d}$ . We use  $f(A, B)$  to denote the in-context loss of  $f(\{P_i, Q_i\}_{i=0}^{L-1})$  as defined in (5), when  $P_i$  and  $Q_i$  depend on  $B_i$  and  $A_i$  as described in (11).
(a) Visualization of  $\Sigma^{1 / 2}A_0\Sigma^{1 / 2}$
Figure 2: Visualization of learned weights for the setting of Theorem 3. We visualize each  $\Sigma^{1/2}A_i\Sigma^{1/2}$  matrix at the end of training. Note that the optimized weights match the stationary point discussed in Theorem 3.
(b) Visualization of  $\Sigma^{1 / 2}A_1\Sigma^{1 / 2}$
(c) Visualization of  $\Sigma^{1 / 2}A_2\Sigma^{1 / 2}$
With this relaxed parameter configuration, it turns out transformers can learn algorithms beyond the conventional preconditioned gradient descent. The next theorem asserts the possibility of learning a novel preconditioned gradient method. Let  $L$  be a fixed but arbitrary number of layers.
Theorem 4. Let  $\Sigma$  denote any PSD matrix. Assume that  $x^{(i)} \stackrel{iid}{\sim} \mathcal{N}(0, \Sigma)$  and  $w_{\star} \sim \mathcal{N}(0, \Sigma^{-1})$ , for  $i = 1, \ldots, n$ , and for some  $\Sigma \succ 0$ . Consider the optimization of in-context loss for a  $L$ -layer linear transformer with the parameter configuration in Eq. (11) given by:
$$
\min_{\substack{\{A_{i},B_{i}\}_{i = 0}^{L - 1}}}f\left(A,B\right).
$$
Let  $\mathcal{S} \subset \mathbb{R}^{2 \times L \times d \times d}$  be defined as follows:  $(A, B) \in \mathcal{S}$  if and only if for all  $i \in \{0, \dots, k\}$ , there exists scalars  $a_i, b_i \in \mathbb{R}$  such that  $A_i = a_i \Sigma^{-1}$  and  $B_i = b_i I$ . Then
$$
\inf  _ {(A, B) \in \mathcal {S}} \sum_ {i = 0} ^ {L - 1} \| \nabla_ {A _ {i}} f (A, B) \| _ {F} ^ {2} + \| \nabla_ {B _ {i}} f (A, B) \| _ {F} ^ {2} = 0, \tag {12}
$$
In words, parameter matrices in  $S$  implement the following algorithm:  $\left\{A_{i} = a_{i}\Sigma^{-1}\right\}_{i = 0}^{L - 1}$  plays the role of a distribution-dependent preconditioner for the gradient steps. At the same time,  $B_{i} = b_{i}I$  transforms the covariates themselves to make the Gram matrix have better condition number with each iteration. When the  $\Sigma = I$ , the algorithm implemented by  $A_{i}\propto I,b_{i}\propto I$  is exactly the GD++ algorithm proposed in (von Oswald et al., 2023) (up to stepsize).
The result in (12) says that the set  $S$  essentially contains critical points of the in-context loss  $f(A, B)$ . In the next section, we provide empirical evidence that the trained transformer parameters do in fact converge to a point in  $S$ .
5.1 Experimental validations for Theorem 4
The experimental setup is similar to Subsection 4.3: we consider ICL for linear regression with  $n = 10, d = 5$ , with  $x^{(i)} \sim \mathcal{N}(0, \Sigma)$  and  $w_{\star} \sim \mathcal{N}(0, \Sigma^{-1})$ , where  $\Sigma = U^T D U$ , where  $U$  is a uniformly random orthogonal matrix, and  $D$  is a fixed diagonal matrix with entries (1, 1, 0.25, 0.0625, 1). We train a three-layer linear transformer, under the constraints in (11) which is less restrictive than (8) in Subsection 4.3. We train the matrices  $A_0, A_1, A_2, B_0, B_1$  using ADAM with the same setup as in Section Subsection 4.3. We repeat this experiment 5 times with different random seeds, each time we sample a different  $U$  (i.e.  $\Sigma$ ).
<sup>3</sup>Once again, similar to the case of Theorem 3, the infimum may not be attained, so it is possible that  $S$  contains points with arbitrarily small gradient, but does not contain a point with exactly 0 gradient.
Note that the objective function does not depend on  $B_{2}$
In Figure 3c, we plot the in-context loss through the iterations of ADAM; the loss appears to be converging to 0, suggesting that parameters are converging to the global minimum.
(a)  $\operatorname{Dist}(B_0, I)$
(b)  $\operatorname{Dist}(B_1, I)$
(c)  $\log (\mathrm{Loss})$
(d) Distances for  $A_0$
Figure 3: Plots for verifying convergence of general linear transformer, defined in Theorem 4. Figure (c) shows convergence of loss to 0. Figures (a),(b) illustrate convergence of  $B_0$ ,  $B_1$  to identity. Figures (d),(e),(f) illustrate convergence of  $A_i$ 's to  $\Sigma^{-1}$ .
(e) Distances for  $A_{1}$
(f) Distances for  $A_{2}$
We next verify that the parameters at convergence are consistent with Theorem 4. We will once again use  $\mathrm{Dist}(M, I)$  to measure the distance from  $M$  to the identity matrix, up to scaling (see Subsection 4.3 for definition of  $\mathit{Dist}$ ). Figures 3a and 3b show that  $B_0$  and  $B_1$  are close to identity, as  $\mathrm{Dist}(B_i, I)$  appears to be decreasing to 0. Figures 3d, 3e and 3f plot  $\mathrm{Dist}(A_i, I)$  (red line) and  $\mathrm{Dist}(\Sigma^{1/2} A_i \Sigma^{1/2}, I)$  (blue line); the results here suggest that  $A_i$  is converging to  $\Sigma^{-1}$ , up to scaling. In Figures 3a and 3b, we observe that  $B_0$  and  $B_1$  also converge to the identity matrix (without left and right multiplication by  $\Sigma^{1/2}$ ), consistent with Theorem 4.
We visualize each of  $B_0, B_1$  in Figure 4 and  $A_0, A_1, A_2$  in Figure 5a-5c at the end of training. We highlight two noteworthy observations:
1. Let  $X_{k} \in \mathbb{R}^{d \times n}$  denote the first  $d$  rows of  $Z_{k}$ , which are the output at layer  $k - 1$  defined in (4). Then the update to  $X_{k}$  is  $X_{k + 1} = X_{k} + B_{k}X_{k}MX_{k}^{T}A_{k}X_{k} \approx X_{k + 1} = X_{k}\left(I - |a_{k}b_{k}|MX_{k}^{T}X_{k}\right)$ , where  $M$  is a mask defined in (2). As noted by von Oswald et al. (2023), this may be motivated by curvature correction.
2. As seen in Figures 5a-5c in the Appendix,  $\| A_0 \| \leq \| A_1 \| \leq \| A_2 \|$  that implies the transformer implements gradient descent with a small stepsize at the beginning and a large stepsize at the end. This makes intuitive sense as  $X_2$  is better-conditioned compared to  $X_1$ , due to the choice of  $B_0, B_1$ . This can be contrasted with the plots in Figures (2a)-(2c), where similar trends are not as pronounced because  $B_i$ 's are constrained to be 0.
We take a first step toward proving that transformers can learn algorithms when trained over a set of random problem instances. Specifically, we investigate the possibility of learning gradient based methods when training on the in-context loss for linear regression. For a single layer transformer, we prove that the global minimum corresponds to a single iteration of preconditioned gradient descent. For multiple layers, we show that certain parameters that correspond to the critical points of the in-context loss can be interpreted as a broad family of adaptive gradient-based algorithms.
Figure 4: Visualization of optimized weight matrices  $B_{0}$  (left) and  $B_{1}$  (right). One can see that the weight pattern matches the stationary point analyzed in Theorem 4. Matrices  $A_{0}, A_{1}$  and  $A_{2}$  are similar to Figure 2, and are visualized in Figure 5 in Appendix D.
We discuss below two interesting future directions.
Beyond linear attention. The standard transformer architecture comes with nonlinear activations in attention. Hence, the natural question here is to ask the effect of nonlinear activations for our main results. Empirically, von Oswald et al. (2023) have observed that for linear regression task, softmax activations generally degrade the prediction performance, and in particular, softmax transformers typically need more attention heads to match their performance with that of linear transformers.
As a first step analysis, we consider the nonlinear attention defined as
$$
\operatorname {A t t n} _ {P, Q} ^ {\sigma} (Z) := P Z M \sigma (Z ^ {\top} Q Z) \quad \text {w h e r e} \sigma : \mathbb {R} \to \mathbb {R} \text {i s a p p l i e d e n t r y - w i s e}.
$$
The following result is an analog of Theorem 1 for single-layer nonlinear attention. It characterizes a global minimizer for this setting with ReLU activation. Here, our choice of ReLU activation was motivated by Wortsman et al. (2023) who observed that ReLU attention matches the performance of softmax attention for vision transformers.
Theorem 5. Consider the single layer nonlinear attention setting with  $\sigma = \mathrm{ReLU}$ . Assume that vector  $x^{(i)}$  is sampled from  $\mathcal{N}(0, I_d)$ . Moreover, assume that  $w_{\star}$  is sampled from  $\mathcal{N}(0, I_d)$ . Consider the parameter configuration  $P_0, Q_0$  where we additionally assume that the last row of  $Q_0$  is zero. Then, the following parameters form a global minimizer of the corresponding in-context loss:
$$
P _ {0} = \left[ \begin{array}{c c} 0 _ {d \times d} & 0 \\ 0 & 1 \end{array} \right], \quad Q _ {0} = - \frac {1}{\frac {1}{2} \frac {n - 1}{n} + (d + 2) \frac {1}{n}} \cdot \left[ \begin{array}{c c} I _ {d} & 0 \\ 0 & 0 \end{array} \right].
$$
The proof of Theorem 5 involves an instructive argument and leverages tools from (Erdogdu et al., 2016); we defer it to Subsection A.4. Thus, for isotropic Gaussian data, the structure of global minimum under ReLU attention is similar to the global minimum with linear attention, established in Theorem 1 (specifically the minimizer for the isotropic date given in (7)).
Refined landscape analysis for multilayer transformer. Theorem 4 proves that a stationary point of the in-context loss corresponds to implementing a preconditioned gradient method. However, we do not prove that all critical points of the non-convex objective lead to similar optimization methods. In fact, in Lemma 4 in Appendix B, we prove that the in-context loss can have multiple critical points. It will be interesting to analyze the set of all critical points and try to understand their algorithmic interpretations, as well as quantify their (sub)optimality.
We thank Ekin Akyurek, Johannes von Oswald, Alex Gu and Joshua Robinson for helpful discussions. Kwangjun Ahn was supported by the ONR grant (N00014-20-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. Kwangjun Ahn also acknowledges support from the Kwanjeong Educational Foundation. Xiang Cheng acknowledges support from NSF CCF-2112665 (TILOS AI Research Institute). Hadi Daneshmand acknowledges support from NSF TRIPODS program (award DMS-2022448). Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute).
A Proofs for the single layer case 13
A.1 Rewriting the loss function 13
A.2 Warm-up:proof for the isotropic data 14
A.3 Proof for the non-isotropic case 16
A.4 Proof for non-linear attentions (Theorem 5) 18
B Proofs for the multi-layer case 21
B.1 Proof of Theorem 2 21
B.2 Proof of Theorem 3 21
B.3 Proof of Theorem 4 24
B.4 Equivalence under permutation 33
C Auxiliary Lemmas 33
C.1 Proof of Lemma 1 (Equivalence to Preconditioned Gradient Descent) 33
C.2 Reformulating the in-context loss 35
D Additional experimental results 36
A Proofs for the single layer case
In this section, we prove our characterization of global minima for the single layer case (Theorem 1). We begin by simplifying the loss into a more concrete form. Throughout the proof, we will write  $P, Q$  instead of  $P_0, Q_0$  for brevity.
A.1 Rewriting the loss function
Recall the in-context loss (5) for the single layer case  $f(P, Q)$  is defined as:
$$
f \left(P, Q\right) = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \left(Z _ {0} + \frac {1}{n} \operatorname {A t t n} _ {P, Q} \left(Z _ {0}\right)\right) _ {(d + 1), (n + 1)} + w _ {\star} ^ {\top} x ^ {(n + 1)} \right] ^ {2}
$$
From the definition of attention given in (3), one can further spell out the expression  $Z_0 + \frac{1}{n} \operatorname{Attn}_{P,Q}(Z_0)$  using the notation  $Z_0 = [z^{(1)} z^{(2)} \dots z^{(n+1)}]$  as follows:
$$
\begin{array}{l} [ z ^ {(1)} \dots z ^ {(n + 1)} ] + \frac {1}{n} P [ z ^ {(1)} \dots z ^ {(n + 1)} ] M \left([ z ^ {(1)} \dots z ^ {(n + 1)} ] ^ {\top} Q [ z ^ {(1)} \dots z ^ {(n + 1)} ]\right) \\ = \left[ z ^ {(1)} \dots z ^ {(n + 1)} \right] + \frac {1}{n} P \left(\sum_ {i = 1} ^ {n} z ^ {(i)} z ^ {(i) ^ {\top}}\right) Q \left[ z ^ {(1)} \dots z ^ {(n + 1)} \right]. \\ \end{array}
$$
Thus, the last column of the above matrix can be expressed as
$$
\left[ \begin{array}{c} x ^ {(n + 1)} \\ 0 \end{array} \right] + \frac {1}{n} P \left(\sum_ {i = 1} ^ {n} z ^ {(i)} z ^ {(i) ^ {\top}}\right) Q \left[ \begin{array}{c} x ^ {(n + 1)} \\ 0 \end{array} \right],
$$
where note that the summation is for  $i = 1,2,\ldots ,n$  due to the mask matrix  $M$ . Therefore, letting  $b^{\top}$  be the last row of  $P$ , and  $A\in \mathbb{R}^{d + 1,d}$  be the first  $d$  columns of  $Q$  (as we did in (14)), then  $\left[Z_0 + \frac{1}{n}\mathrm{Attn}_{P,Q}(Z_0)\right]_{(d + 1),(n + 1)}$  can be written as
$$
\frac {1}{n} b ^ {\top} \left(\sum_ {i = 1} ^ {n} z ^ {(i)} z ^ {(i) ^ {\top}}\right) A \left[ \begin{array}{c} x ^ {(n + 1)} \\ 0 \end{array} \right], \tag {13}
$$
in other words,  $f(P, Q)$  only depends on the parameter  $b$  and  $A$ . Henceforth, we will write  $f(P, Q)$  as  $f(b, A)$ . Let us summarize our conclusion so far since it's crucial for the analysis to follow.
Conclusion so far: A careful inspection reveals that the in-context loss only depends on the last row of  $P$  and the first  $d$  columns of  $Q$ . Thus, consider the following parametrization
$$
P = \left[ \begin{array}{l} 0 \\ b ^ {\top} \end{array} \right] \quad \text {a n d} \quad Q = \left[ \begin{array}{l l} A & 0 \end{array} \right], \text {w h e r e} b \in \mathbb {R} ^ {d + 1} \text {a n d} A \in \mathbb {R} ^ {(d + 1) \times d}. \tag {14}
$$
Now with this parametrization, the in-context loss can be written as  $f(b,A) \coloneqq f([0 b]^{\top},[A 0])$ .
Now, let us spell out  $f(b,A)$  based on (13) as follows:
$$
\begin{array}{l} f (b, A) = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ b ^ {\top} \underbrace {\frac {1}{n} \sum_ {i} z ^ {(i)} z ^ {(i)} ^ {\top}} _ {=: \mathbf {G}} A x ^ {(n + 1)} + w _ {\star} ^ {\top} x ^ {(n + 1)} \right] ^ {2} \\ =: \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ b ^ {\top} \mathsf {G} A x ^ {(n + 1)} + w _ {\star} ^ {\top} x ^ {(n + 1)} \right] ^ {2} = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ (b ^ {\top} \mathsf {G} A + w _ {\star} ^ {\top}) x ^ {(n + 1)} \right] ^ {2}, \tag {15} \\ \end{array}
$$
where we used the notation  $G \coloneqq \frac{1}{n} \sum_{i} z^{(i)} z^{(i)}^{\top}$  to simplify. We now analyze the global minima of this loss function. To illustrate the proof idea clearly, we begin with the proof for the simpler case of isotropic data.
A.2 Warm-up: proof for the isotropic data
As a warm-up, we first prove the result for the special case where  $x^{(i)}$  is sampled from  $\mathcal{N}(0, I_d)$ .
1. Decomposing the loss function into components. Writing  $A = [a_{1} a_{1} \cdots a_{d}]$ , and use the fact that  $\mathbb{E}[x^{(n + 1)}[j] x^{(n + 1)}[j']] = 0$  for  $j \neq j'$  and  $\mathbb{E}[x^{(n + 1)}[j]^{2}] = 1$ , we get
$$
f \left(b, A\right) = \sum_ {j = 1} ^ {d} \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ b ^ {\top} \mathsf {G} a _ {j} + w _ {\star} [ j ] \right] ^ {2} \mathbb {E} [ x ^ {(n + 1)} [ j ] ^ {2} ] = \sum_ {j = 1} ^ {d} \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ b ^ {\top} \mathsf {G} a _ {j} + w _ {\star} [ j ] \right] ^ {2}.
$$
The key idea is to characterize the global minima of each component in the summation separately. Another key idea is to reparametrize the cost function given the following identity:
$$
\mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ b ^ {\top} \mathsf {G} a _ {j} + w _ {\star} [ j ] \right] ^ {2} = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \mathrm {T r} (\mathsf {G} a _ {j} b ^ {\top}) + w _ {\star} [ j ] \right] ^ {2} = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \left\langle \mathsf {G}, b a _ {j} ^ {\top} \right\rangle + w _ {\star} [ j ] \right] ^ {2},
$$
where we use the notation  $\langle X,Y\rangle \coloneqq \mathrm{Tr}(XY^{\top})$  for two matrices  $X$  and  $Y$  here and below. Given the above identity, we define each component in the summation as follows.
$$
\boxed {f _ {j} (X) := \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \langle \mathrm {G}, X \rangle + w _ {\star} [ j ] \right] ^ {2} \quad \text {f o r} X \in \mathbb {R} ^ {(d + 1) \times (d + 1)}.}
$$
2. Characterizing global minima of each component. To characterize the global minima of each objective, we prove the following result.
Lemma 2 (Global minima of each component). Suppose that  $x^{(i)}$  is sampled from  $\mathcal{N}(0, I_d)$  and  $w_{\star}$  is sampled from  $\mathcal{N}(0, I_d)$ . Consider the following objective (here,  $\langle X, Y \rangle \coloneqq \mathrm{Tr}(XY^{\top})$  for two matrices  $X$  and  $Y$ )
$$
f _ {j} (X) = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \langle \mathsf {G}, X \rangle + w _ {\star} [ j ] \right] ^ {2}.
$$
Then a global minimum is given as
$$
X _ {j} = - \frac {1}{\left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right)} E _ {d + 1, j},
$$
where  $E_{i_1,i_2}$  is the matrix whose  $(i_1,i_2)$ -th entry is 1, and the other entries are zero.
Proof of Lemma 2. Note first that  $f_{j}$  is convex in  $X$ . Hence, in order to show that matrix  $X_{j}$  is the global optimum of  $f_{j}$ , it suffices to show that the gradient vanishes at that point, in other words,  $\nabla f_{j}(X_{j}) = 0$ . To verify this, let us compute the gradient of  $f_{j}$ : for a matrix  $X$ ,
$$
\nabla f _ {j} (X) = 2 \mathbb {E} \left[ \langle \mathrm {G}, X \rangle \mathrm {G} \right] + 2 \mathbb {E} \left[ w _ {\star} [ j ] \mathrm {G} \right],
$$
where we recall that  $G$  is defined as