SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathsf {G} = \frac {1}{n} \sum_ {i} \left[ \begin{array}{c c} x ^ {(i)} x ^ {(i) ^ {\top}} & y ^ {(i)} x ^ {(i)} \\ y ^ {(i)} x ^ {(i) ^ {\top}} & y ^ {(i) ^ {2}} \end{array} \right].
$$
To verify that the gradient is equal to zero, let us first compute  $\mathbb{E}[w_{\star}[j]\mathsf{G}]$ . For each  $i = 1,\dots ,n$  note that  $\mathbb{E}[w_{\star}[j]x^{(i)}x^{(i)^{\top}}] = O$  because  $\mathbb{E}[w_{\star}] = 0$ . Moreover,  $\mathbb{E}[w_{\star}[j](y^{(i)})^2 ] = 0$  because  $w_{\star}$  is symmetric, i.e.,  $w_{\star}\stackrel {d}{=} - w_{\star}$ , and  $y^{(i)} = \langle w_{\star},x^{(i)}\rangle$ . Lastly, for  $k = 1,2,\ldots ,d$ , we have
$$
\mathbb {E} \left[ w _ {\star} [ j ] y ^ {(i)} x ^ {(i)} [ k ] \right] = \mathbb {E} \left[ w _ {\star} [ j ] \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ k ] \right] = \mathbb {E} \left[ w _ {\star} [ j ] ^ {2} x ^ {(i)} [ j ] x ^ {(i)} [ k ] \right] = \mathbb {1} _ {[ j = k ]} \tag {16}
$$
because  $\mathbb{E}[w_{\star}[i]w_{\star}[j]] = 0$  for  $i\neq j$ . Combining the above calculations, it follows that
$$
\boxed {\mathbb {E} [ w _ {\star} [ j ] \mathsf {G} ] = E _ {d + 1, j} + E _ {j, d + 1}.} \tag {17}
$$
In order to compute  $\mathbb{E}[\langle \mathsf{G},X\rangle \mathsf{G}]$ , let us compute  $\mathbb{E}[\langle \mathsf{G},E_{i,i^{\prime}}\rangle \mathsf{G}]$  for  $i,i^{\prime} = 1,\ldots ,d + 1$ . Without loss of generality,  $i\geq i'$ . First of all We now compute compute  $\mathbb{E}[\langle \mathsf{G},E_{d + 1,j}\rangle \mathsf{G}]$ . Note first that
$$
\langle \mathrm {G}, E _ {d + 1, j} \rangle = \sum_ {i} \left\langle w _ {\star}, x ^ {(i)} \right\rangle x ^ {(i)} [ j ].
$$
Hence, it holds that
$$
\mathbb {E} \left[ \langle \mathsf {G}, E _ {d + 1, j} \rangle \left(\sum_ {i} x ^ {(i)} x ^ {(i) ^ {\top}}\right) \right] = \mathbb {E} \left[ \left(\sum_ {i} \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ j ]\right) \left(\sum_ {i} x ^ {(i)} x ^ {(i) ^ {\top}}\right) \right] = O.
$$
because  $\mathbb{E}[w_{\star}] = 0$  . Next, we have
$$
\mathbb {E} \left[ \langle \mathrm {G}, E _ {d + 1, j} \rangle \left(\sum_ {i} y ^ {(i) ^ {2}}\right) \right] = \mathbb {E} \left[ \left(\sum_ {i} \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ j ]\right) \left(\sum_ {i} y ^ {(i) ^ {2}}\right) \right] = 0
$$
because  $w_{\star} \stackrel{d}{=} -w_{\star}$ . Lastly, we compute
$$
\mathbb {E} \left[ \langle \mathsf {G}, E _ {d + 1, j} \rangle \left(\sum_ {i} y ^ {(i)} x ^ {(i) ^ {\top}}\right) \right].
$$
To that end, note that for  $j \neq j'$ ,
$$
\mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle   x ^ {(i)} [ j ]   \langle w _ {\star}, x ^ {(i ^ {\prime})} \rangle   x ^ {(i ^ {\prime})} [ j ^ {\prime} ] \right] = \left\{ \begin{array}{l l} \mathbb {E} [ \langle x ^ {(i)}, x ^ {(i ^ {\prime})} \rangle   x ^ {(i)} [ j ]   x ^ {(i ^ {\prime})} [ j ^ {\prime} ] ] = 0 & \text {i f} i \neq i ^ {\prime}, \\ \mathbb {E} [ \| x ^ {(i)} \| ^ {2}   x ^ {(i)} [ j ]   x ^ {(i)} [ j ^ {\prime} ] ] = 0 & \text {i f} i = i ^ {\prime}, \end{array} \right.
$$
$$
\mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ j ] \langle w _ {\star}, x ^ {\left(i ^ {\prime}\right)} \rangle x ^ {\left(i ^ {\prime}\right)} [ j ] \right] = \left\{ \begin{array}{l l} \mathbb {E} \left[ \left(x ^ {(i)} [ j ]\right) ^ {2} \left(x ^ {\left(i ^ {\prime}\right)} [ j ]\right) ^ {2} \right] = 1 & \text {i f} i \neq i ^ {\prime}, \\ \mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle^ {2} \left(x ^ {(i)} [ j ]\right) ^ {2} \right] = d + 2 & \text {i f} i = i ^ {\prime}, \end{array} \right. \tag {18}
$$
where the last case follows from the fact that the fourth moment of Gaussian is 3 and
$$
\mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle^ {2} (x ^ {(i)} [ j ]) ^ {2} \right] = \mathbb {E} \left[ \| x ^ {(i)} \| ^ {2} (x ^ {(i)} [ j ]) ^ {2} \right] = 3 + d - 1 = d + 2.
$$
Combining the above calculations together, we arrive at
$$
\begin{array}{l} \mathbb {E} \left[ \langle \mathsf {G}, E _ {d + 1, j} \rangle \mathsf {G} \right] = \frac {1}{n ^ {2}} \cdot (n (n - 1) + (d + 2) n) \left(E _ {d + 1, j} + E _ {j, d + 1}\right) \\ = \left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right) \left(E _ {d + 1, j} + E _ {j, d + 1}\right). \tag {19} \\ \end{array}
$$
Therefore, combining (17) and (19), the results follow.
3. Combining global minima of each component. From Lemma 2, it follows that
is the unique global minimum of  $f_{j}$ . Hence,  $b$  and  $A = [a_{1} a_{1} \cdots a_{d}]$  achieve the global minimum of  $f(b, A) = \sum_{j=1}^{d} f_{j}(ba_{j}^{\top})$  if they satisfy
$$
b a _ {j} ^ {\top} = - \frac {1}{\left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right)} E _ {d + 1, j} \quad \text {f o r a l l} i = 1, 2, \dots , d.
$$
This can be achieve by the following choice:
$$
b ^ {\top} = \mathbf {e} _ {d + 1}, \quad a _ {j} = - \frac {1}{\left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right)} \mathbf {e} _ {j} \quad \mathrm {f o r} i = 1, 2, \ldots , d,
$$
where  $\mathbf{e}_j$  is the  $j$ -th coordinate vector. This choice precisely corresponds to
$$
b = \mathbf {e} _ {d + 1}, \quad A = - \frac {1}{\left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right)} \left[ \begin{array}{c} I _ {d} \\ 0 \end{array} \right].
$$
We next move on to the non-isotropic case.
A.3 Proof for the non-isotropic case
1. Diagonal covariance case. We first consider the case where  $x^{(i)}$  is sampled from  $\mathcal{N}(0, \Lambda)$  where  $\overline{\Lambda} = \mathrm{diag}(\lambda_1, \dots, \lambda_d)$  and  $w_{\star}$  is sampled from  $\mathcal{N}(0, I_d)$ . We prove the following generalization of Lemma 2.
Lemma 3. Suppose that  $x^{(i)}$  is sampled from  $\mathcal{N}(0, \Lambda)$  where  $\Lambda = \mathrm{diag}(\lambda_1, \dots, \lambda_d)$  and  $w_{\star}$  is sampled from  $\mathcal{N}(0, I_d)$ . Consider the following objective
$$
X _ {j} = - \frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot (\sum_ {k} \lambda_ {k})} E _ {d + 1, j},
$$
Proof of Lemma 3. Similarly to the proof of Lemma 2, it suffices to check that
$$
2 \mathbb {E} \left[ \langle \mathsf {G}, X _ {0} \rangle \mathsf {G} \right] + 2 \mathbb {E} \left[ w _ {\star} [ j ] \mathsf {G} \right] = 0,
$$
A similar calculation as the proof of Lemma 2 yields
$$
\mathbb {E} \left[ w _ {\star} [ j ] \mathsf {G} \right] = \lambda_ {j} \left(E _ {d + 1, j} + E _ {j, d + 1}\right). \tag {20}
$$
Here the factor of  $\lambda_{j}$  comes from the following generalization of (16):
$$
\mathbb {E} [ w _ {\star} [ j ] y ^ {(i)} x ^ {(i)} [ k ] ] = \mathbb {E} [ w _ {\star} [ j ] \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ k ] ] = \mathbb {E} \left[ w _ {\star} [ j ] ^ {2} x ^ {(i)} [ j ] x ^ {(i)} [ k ] \right] = \lambda_ {j} \mathbb {1} _ {[ j = k ]}.
$$
Next, we compute  $\mathbb{E}\left[\langle \mathsf{G}, E_{d+1,j} \rangle \mathsf{G}\right]$ . Again, we follow a similar calculation to the proof of Lemma 2 except that this time we use the following generalization of (18):
$$
\mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle x ^ {(i)} [ j ] \langle w _ {\star}, x ^ {(i ^ {\prime})} \rangle x ^ {(i ^ {\prime})} [ j ] \right] = \left\{ \begin{array}{l l} \mathbb {E} [ x ^ {(i)} [ j ] ^ {2} x ^ {(i ^ {\prime})} [ j ] ^ {2} ] = \lambda_ {j} ^ {2} & \mathrm {i f} i \neq i ^ {\prime}, \\ \mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle^ {2} x ^ {(i)} [ j ] ^ {2} \right] = \lambda_ {j} \sum_ {k} \lambda_ {k} + 2 \lambda_ {j} ^ {2} & \mathrm {i f} i = i ^ {\prime}, \end{array} \right.
$$
where the last line follows since
$$
\mathbb {E} \left[ \langle w _ {\star}, x ^ {(i)} \rangle^ {2} x ^ {(i)} [ j ] ^ {2} \right] = \mathbb {E} \left[ \| x ^ {(i)} \| ^ {2} x ^ {(i)} [ j ] ^ {2} \right] = \mathbb {E} \left[ x ^ {(i)} [ j ] ^ {2} \sum_ {k} x ^ {(i)} [ k ] ^ {2} \right] = \lambda_ {j} \sum_ {k} \lambda_ {k} + 2 \lambda_ {j} ^ {2}.
$$
$$
\begin{array}{l} \mathbb {E} \left[ \langle \mathsf {G}, E _ {d + 1, j} \rangle \mathsf {G} \right] = \frac {1}{n ^ {2}} \cdot \left(n (n - 1) \lambda_ {j} ^ {2} + n \lambda_ {j} \sum_ {k} \lambda_ {k} + 2 n \lambda_ {j} ^ {2}\right) \left(E _ {d + 1, j} + E _ {j, d + 1}\right) \\ = \left(\frac {n + 1}{n} \lambda_ {j} ^ {2} + \frac {1}{n} \left(\lambda_ {j} \sum_ {k} \lambda_ {k}\right)\right) \left(E _ {d + 1, j} + E _ {j, d + 1}\right). \tag {21} \\ \end{array}
$$
Therefore, combining (20) and (21), the results follow.
Now we finish the proof. From Lemma 2, it follows that
$$
X _ {j} = - \frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot (\sum_ {k} \lambda_ {k})} E _ {d + 1, j}
$$
is the unique global minimum of  $f_{j}$ . Hence,  $b$  and  $A = [a_{1} a_{1} \cdots a_{d}]$  achieve the global minimum of  $f(b, A) = \sum_{j=1}^{d} f_{j}(b, A_{j})$  if they satisfy
$$
b a _ {j} ^ {\top} = X _ {j} = - \frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot \left(\sum_ {k} \lambda_ {k}\right)} E _ {d + 1, j} \quad \text {f o r a l l} i = 1, 2, \dots , d.
$$
$$
b ^ {\top} = \mathbf {e} _ {d + 1}, \quad a _ {j} = - \frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot (\sum_ {k} \lambda_ {k})} \mathbf {e} _ {j} \quad \mathrm {f o r} i = 1, 2, \ldots , d,
$$
$$
b = \mathbf {e} _ {d + 1}, \quad A = - \left[ \begin{array}{c} \operatorname {d i a g} \left(\left\{\frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot \left(\sum_ {k} \lambda_ {k}\right)} \right\} _ {j}\right) \\ 0 \end{array} \right].
$$
2. Non-diagonal covariance case (the setting of Theorem 1). We finally prove the general result of Theorem 1, namely  $x^{(i)}$  is sampled from a Gaussian with covariance  $\Sigma = U\Lambda U^{\top}$  where  $\Lambda = \mathrm{diag}(\lambda_1,\dots ,\lambda_d)$  and  $w_{\star}$  is sampled from  $\mathcal{N}(0,I_d)$ . The proof works by reducing this case to the previous case. For each  $i$ , define  $\widetilde{x}^{(i)}\coloneqq U^T x^{(i)}$ . Then  $\mathbb{E}[\widetilde{x}^{(i)}(\widetilde{x}^{(i)})^\top ] = \mathbb{E}[U^\top (U\Lambda U^\top)U] = \Lambda$ . Now let us write the loss function (15) with this new coordinate system: since  $x^{(i)} = U\widetilde{x}^{(i)}$ , we have
$$
f (b, A) = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ (b ^ {\top} \mathsf {G} A + w _ {\star} ^ {\top}) U \widetilde {x} ^ {(n + 1)} \right] ^ {2} = \sum_ {j = 1} ^ {d} \lambda_ {j} \mathbb {E} _ {Z _ {0}, w _ {\star}} \left[ \big ((b ^ {\top} \mathsf {G} A + w _ {\star} ^ {\top}) U \big) [ j ] \right] ^ {2}.
$$
Hence, let us consider the vector  $(b^{\top}GA + w_{\star}^{\top})U$ . By definition of  $G$ , we have
$$
\begin{array}{l} (b ^ {\top} \mathsf {G} A + w _ {\star} ^ {\top}) U = \frac {1}{n} \sum_ {i} b ^ {\top} \left[ \begin{array}{c} x ^ {(i)} \\ \langle x ^ {(i)}, w _ {\star} \rangle \end{array} \right] ^ {\otimes 2} A U + w _ {\star} ^ {\top} U \\ = \frac {1}{n} \sum_ {i} b ^ {\top} \left[ \begin{array}{c} U \tilde {x} _ {i} \\ \langle U x ^ {(i)}, w _ {\star} \rangle \end{array} \right] ^ {\otimes 2} A U + w _ {\star} ^ {\top} U \\ = \frac {1}{n} \sum_ {i} b ^ {\top} \left[ \begin{array}{c c} U & 0 \\ 0 & 1 \end{array} \right] \left[ \begin{array}{c} \tilde {x} _ {i} \\ \langle U x ^ {(i)}, w _ {\star} \rangle \end{array} \right] ^ {\otimes 2} \left[ \begin{array}{c c} U ^ {\top} & 0 \\ 0 & 1 \end{array} \right] A U + w _ {\star} ^ {\top} U \\ = \frac {1}{n} \sum_ {i} \widetilde {b} ^ {\top} \left[ \left\langle x ^ {(i)}, \widetilde {w} _ {\star} \right\rangle \right] ^ {\otimes 2} \widetilde {A} + \widetilde {w} _ {\star} ^ {\top} \\ \end{array}
$$
where we define  $\widetilde{b}^{\top} := b^{\top}\left[ \begin{array}{cc}U & 0\\ 0 & 1 \end{array} \right],\widetilde{A} := \left[ \begin{array}{cc}U^{\top} & 0\\ 0 & 1 \end{array} \right]AU,$  and  $\widetilde{w}_{\star}:= U^{\top}w_{\star}$ . By the rotational symmetry,  $\widetilde{w}_{\star}$  is also distributed as  $\mathcal{N}(0,I_d)$ . Hence, this reduces to the previous case, and a global minimum is given as
$$
\widetilde {b} = \mathbf {e} _ {d + 1}, \quad \widetilde {A} = - \left[ \begin{array}{c} \operatorname {d i a g} \left(\left\{\frac {1}{\frac {n + 1}{n} \lambda_ {j} + \frac {1}{n} \cdot \left(\sum_ {k} \lambda_ {k}\right)} \right\} _ {j}\right) \\ 0 \end{array} \right].
$$
From the definition of  $\tilde{b}$ ,  $\tilde{A}$ , it thus follows that a global minimum is given by
$$
\begin{array}{r} b ^ {\top} = \mathbf {e} _ {d + 1}, \quad A = - \left[ \begin{array}{c} U \mathrm {d i a g} \left(\left\{\frac {1}{\frac {n + 1}{n} \lambda_ {i} + \frac {1}{n} \cdot \left(\sum_ {k} \lambda_ {k}\right)} \right\} _ {i}\right) U ^ {\top} \\ 0 \end{array} \right], \end{array}
$$
A.4 Proof for non-linear attentions (Theorem 5)
As mentioned in Theorem 5, we focus on the setting where the last row of  $Q$  is zero, i.e., let
$$
Q = \left[ \begin{array}{c c} A & a \\ 0 ^ {\top} & 0 \end{array} \right] \quad \text {f o r} A \in \mathbb {R} ^ {d \times d} \text {a n d} a \in \mathbb {R} ^ {d}.
$$
We first rewrite the loss function and simplify it following Subsection A.1. Moreover, for simple notation we will often write  $z, x$  instead of  $z^{(n + 1)}, x^{(n + 1)}$ .
1. Rewriting loss function.
Following Subsection A.1, let us write down the in-context loss (5). for the single-layer nonlinear attention denoted by  $f(P,Q)$ :
$$
f \left(P, Q\right) = \mathbb {E} _ {Z _ {0}, w _ {\star}} \left(\left[ Z _ {0} + \frac {1}{n} \mathrm {A t t n} _ {P, Q} ^ {\sigma} (Z _ {0}) \right] _ {d + 1, n + 1} + w _ {\star} ^ {\top} x\right) ^ {2}
$$
Recalling the definition of the ReLU attention  $\mathrm{Attn}_{P,Q}^{\sigma}(Z) \coloneqq PZM\sigma(Z^{\top}QZ)$ , the data matrix  $Z \coloneqq [z^{(1)} \cdots z^{(n+1)}]$ , and the mask matrix  $M$ , the term  $\mathrm{Attn}_{P,Q}^{\sigma}(Z_0)$  can be written as:
$$
\operatorname{Attn}_{P,Q}^{\sigma}(Z_{0}) = \underbrace{PZ_{0}}_{\mathbb{R}^{(d + 1)\times (n + 1)}}\cdot \left[ \begin{array}{cc}I_{n\times n} & 0\\ 0 & 0 \end{array} \right]\cdot \underbrace{\sigma\left(Z_{0}^{\top}QZ_{0}\right)}_{\mathbb{R}^{(n + 1)\times (n + 1)}}.
$$
Hence, it follows that the  $(d + 1,n + 1)$ -th entry of  $\mathrm{Attn}_{P,Q}^{\sigma}(Z_0)$  is equal to the product of the  $(d + 1)$ -th row of  $PZ_0$ , the mask matrix  $M$ , and the  $(n + 1)$ -th column of  $\sigma(Z_0^\top QZ_0)$ . Hence, let us write them down explicitly:
- Letting  $b^{\top}$  be the last row of the matrix  $P$ , it holds that the  $(d + 1)$ -th row of  $PZ_{0}$  is equal to  $[\langle b, z^{(i)} \rangle]_{i = 1, \dots, n + 1}$ .
- The  $(n + 1)$ -th column of  $\sigma\left(Z_0^\top QZ_0\right)$  is equal to  $\left[\sigma\left((z^{(i)})^\top Qz^{(n + 1)}\right)\right]_{i = 1,\ldots ,n + 1}$ . Letting  $A$  the first  $d$  columns of  $Q$ , this vector is equal to  $\left[\sigma\left((x^{(i)})^\top Ax^{(n + 1)}\right)\right]_{i = 1,\ldots ,n + 1}$  because the last row of  $Q$  is zero and the last row of  $z^{(n + 1)}$  is zero (since  $(z^{(n + 1)})^\top = [(x^{(n + 1)})^\top 0]$ ).
Thus, the product of  $[\langle b,z^{(i)}\rangle ]_{i = 1,\ldots ,n + 1}$ , the mask matrix  $M$ , and  $\left[\sigma \left((x^{(i)})^\top Ax^{(n + 1)}\right)\right]_{i = 1,\ldots ,n + 1}$  results in the following expression of the attention (writing  $z,x$  instead of  $z^{(n + 1)},x^{(n + 1)}$ ):
$$
\left[ \mathrm {A t t n} _ {P, Q} ^ {\sigma} (Z _ {0}) \right] _ {d + 1, n + 1} = \sum_ {i = 1} ^ {n} \left[ \langle b, z ^ {(i)} \rangle \cdot \sigma ((x ^ {(i)}) ^ {\top} A x) \right].
$$
Since  $[Z_0]_{d + 1,n + 1} = 0$ , we therefore have
$$
\left[ Z _ {0} + \frac {1}{n} \mathrm {A t t n} _ {P, Q} ^ {\sigma} (Z _ {0}) \right] _ {d + 1, n + 1} = \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \langle b, z ^ {(i)} \rangle \cdot \sigma ((x ^ {(i)}) ^ {\top} A x) \right].
$$
Therefore, it follows that the in-context loss  $f(P, Q)$  only depends on  $b$  and  $A$ . Henceforth, let us write  $f(b, A)$  instead of  $f(P, Q)$  following Subsection A.1. In particular, writing  $b^{\top} = [b_0^{\top}, b_1]$  for  $b_0 \in \mathbb{R}^d$  and  $b_1 \in \mathbb{R}$ , the loss function can be expressed as
$$
f (b, A) := \mathbb {E} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \left(\left\langle b _ {0}, x ^ {(i)} \right\rangle + b _ {1} y ^ {(i)}\right) \cdot \sigma \left(\left(x ^ {(i)}\right) ^ {\top} A x\right) \right] + \left\langle w _ {\star}, x \right\rangle\right) ^ {2}. \tag {22}
$$
2. Simplifying the loss function with symmetry.
Now, we use the fact that both  $x^{(i)}$ 's and  $w_{\star}$  are sampled from the isotropic Gaussian, i.e.,  $\mathcal{N}(0,I_d)$  in order to further simplify the loss function in (22). In particular, we use the following facts:
(a) For orthonormal matrices  $U, V \in \mathbb{R}^{d \times d}$ , it holds that  $Ux^{(i)}$ ,  $Vx$  and  $Uw_{\star}$  have the same distributions as  $\mathcal{N}(0, I_d)$ .
(b) Moreover, for a diagonal matrix  $\Xi = \mathrm{diag}(\xi_i)\in \mathbb{R}^{d\times d}$  with the diagonal entries being random signs  $\xi_{i}\sim \{\pm 1\}$ , it holds that  $\Xi x^{(i)},\Xi x$  and  $\Xi w_{\star}$  have the same distributions as  $\mathcal{N}(0,I_d)$ .
Now let us fix a matrix  $A \in \mathbb{R}^{d \times d}$  and  $b^{\top} = [b_0^{\top}, b_1]$  for  $b_0 \in \mathbb{R}^d$  and  $b_1 \in \mathbb{R}$ . Letting  $A = U\Sigma V^{\top}$  be the SVD of the matrix  $A$ , it follows that
$$
\begin{array}{l} f (b, A) = \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (b _ {0} ^ {\top} x ^ {(i)} + b _ {1} w _ {\star} ^ {\top} x ^ {(i)}) \cdot \sigma ((x ^ {(i)}) ^ {\top} A x) \right] + w _ {\star} ^ {\top} x \right] ^ {2} \\ \stackrel {(a)} {=} \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (b _ {0} ^ {\top} U x ^ {(i)} + b _ {1} w _ {\star} ^ {\top} U ^ {\top} U x ^ {(i)}) \cdot \sigma ((x ^ {(i)}) ^ {\top} U ^ {\top} A V x) \right] + w _ {\star} ^ {\top} U ^ {\top} V x \right] ^ {2} \\ \stackrel {(b)} {=} \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (b _ {0} ^ {\top} U \Xi x ^ {(i)} + b _ {1} w _ {\star} ^ {\top} x ^ {(i)}) \cdot \sigma ((x ^ {(i)}) ^ {\top} \Sigma x) \right] + w _ {\star} ^ {\top} \Xi U ^ {\top} V \Xi x \right] ^ {2} \\ \geq \mathbb {E} \left[ \mathbb {E} _ {\Xi} \left\{\frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (b _ {0} ^ {\top} U \Xi x ^ {(i)} + b _ {1} w _ {\star} ^ {\top} x ^ {(i)}) \cdot \sigma ((x ^ {(i)}) ^ {\top} \Sigma x) \right] + w _ {\star} ^ {\top} \Xi U ^ {\top} V \Xi x \right\} \right] ^ {2} \\ = \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ b _ {1} w _ {\star} ^ {\top} x ^ {(i)} \cdot \sigma ((x ^ {(i)}) ^ {\top} \Sigma x) \right] + w _ {\star} ^ {\top} \operatorname {d i a g} (U ^ {\top} V) x \right] ^ {2} \\ =: f _ {\text {l o w e r}} \left(b _ {1}, \Sigma , D := \operatorname {d i a g} \left(U ^ {\top} V\right)\right). \\ \end{array}
$$
where in the third line we use the fact that  $\Xi^{\top}\Sigma \Xi = \Xi \Sigma \Xi = \Sigma$ ; and the fourth line follows from the Jensen's inequality. Hence for the remainder of the proof, we will characterize the global minimizer of the lower bound, i.e.,  $f_{\mathrm{lower}}$  and then we will connect it back to the original objective.
3. Computation of the lower bound  $f_{\text{lower}}$ .
Let us now explicitly compute  $f_{\mathrm{lower}}$ . Let us rewrite the definition of  $f_{\mathrm{lower}}$ . In fact since,  $\sigma = \mathrm{ReLU}$  is homogenous, one can further simplify the lower bound by pushing the constant  $b_{1}$  inside and write  $b_{1}\Sigma$  as  $\Sigma$ . Hence, for two diagonal matrices  $\Sigma, D \in \mathbb{R}^{d\times d}$ ,  $f_{\mathrm{lower}}$  is defined as:
$$
f _ {\mathrm {l o w e r}} (\Sigma , D) := \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \langle w _ {\star}, x ^ {(i)} \rangle \cdot \sigma (x ^ {\top} \Sigma x ^ {(i)}) \right] + w _ {\star} ^ {\top} D x \right] ^ {2}.
$$
In particular,  $D$  is constrained to be the diagonal part of an orthogonal matrix (since  $D = \mathrm{diag}(U^\top V)$  in the above derivation). Now we focus on characterizing the global minimizers of  $f_{\mathrm{lower}}$ .
The main part of the argument is inspired by the elegant observation of Erdogdu et al. (2016), which says that the solution of least squares and generalized linear models are collinear for Gaussian inputs. We leverage the same proof technique (à la Stein's Lemma) to prove that the presence of ReLU only changes the scaling of global optimum.
First, since  $w_{\star}$  is isotropic Gaussian, we can take the expectation over  $w_{\star}$  to obtain
$$
f _ {\text {l o w e r}} (\Sigma , D) = \mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \sigma \left(x ^ {\top} \Sigma x ^ {(i)}\right) x ^ {(i)} \right] + D x \right] ^ {2},
$$
which after a careful expansion becomes
$$
\mathbb {E} \left[ \frac {1}{n ^ {2}} \sum_ {i, j} \sigma \left(x ^ {\top} \Sigma x ^ {(i)}\right) \sigma \left(x ^ {\top} \Sigma x ^ {(j)}\right) \langle x ^ {(i)}, x ^ {(j)} \rangle + \frac {2}{n} \sum_ {i} \sigma \left(x ^ {\top} \Sigma x ^ {(i)}\right) \langle x ^ {(i)}, D x \rangle \right] + \text {c o n s t .} \tag {23}
$$
In order to compute (23), we will rely on the aforementioned argument of Erdogdu et al. (2016). In particular, from integration by parts, or Stein's lemma (Erdogdu et al., 2016) (since  $x \sim \mathcal{N}(0, I_d)$ ), we have
$$
\mathbb {E} _ {x} [ \sigma (x ^ {\top} v) x ] = \mathbb {E} _ {x} [ \sigma^ {\prime} (x ^ {\top} v) ] v \quad \text {f o r a f i x e d} v \in \mathbb {R} ^ {d}.
$$
We use this to compute all the terms in (23) as follows:
- We first apply Stein's lemma to the first term of (23) for  $i \neq j$ . This results in
$$
\begin{array}{l} \mathbb {E} _ {x ^ {(i)}, x ^ {(j)}, x} \sigma (x ^ {\top} \Sigma x ^ {(i)}) \sigma (x ^ {\top} \Sigma x ^ {(j)}) \langle x ^ {(i)}, x ^ {(j)} \rangle \\ = \mathbb {E} _ {x ^ {(j)}, x} \left[ \mathbb {E} _ {x ^ {(i)}} \left[ \sigma^ {\prime} \left(x ^ {\top} \Sigma x ^ {(i)}\right) \right] \sigma \left(x ^ {\top} \Sigma x ^ {(j)}\right) x ^ {\top} \Sigma x ^ {(j)} \right] \\ \end{array}
$$
Using the fact that  $x^{(i)}$  is a symmetric random variable, one can compute the expectation above as follows: one the one hand, we know  $\mathbb{E}_{x^{(i)}}[\sigma'(x^{\top} \Sigma x^{(i)})] = \mathbb{E}_{x^{(i)}}[\sigma'(-x^{\top} \Sigma x^{(i)})]$ . On the other hand, we also know that for any scalar  $\alpha$ ,  $\sigma'(-\alpha) + \sigma'(\alpha) = 1$ . Therefore, we conclude that  $\mathbb{E}_{x^{(i)}}[\sigma'(x^{\top} \Sigma x^{(i)})] = 1/2$ . Thus, applying this technique twice, we obtain the following
$$
\mathbb {E} _ {x ^ {(i)}, x ^ {(j)}, x} \sigma (x ^ {\top} \Sigma x ^ {(i)}) \sigma (x ^ {\top} \Sigma x ^ {(j)}) \langle x ^ {(i)}, x ^ {(j)} \rangle = \frac {1}{4} \mathbb {E} _ {x} [ x ^ {\top} \Sigma^ {2} x ] = \frac {1}{4} \operatorname {T r} (\Sigma^ {2}).
$$
- Similarly, we can use Stein's lemma to the second term of (23) to conclude
$$
\mathbb {E} \sigma (x ^ {\top} \Sigma x ^ {(i)}) \langle x ^ {(i)}, D x \rangle = \frac {1}{2} \mathbb {E} _ {x} x ^ {\top} \Sigma D x = \frac {1}{2} \operatorname {T r} (\Sigma D).
$$
- Lastly, the computation of the first term of (23) for  $i = j$  is straightforward. Using the fact that  $\forall \alpha \in \mathbb{R}, \sigma^2(\alpha) + \sigma^2(-\alpha) = a^2$ , we get
$$
\begin{array}{l} \mathbb {E} \left[ \sigma^ {2} \left(x ^ {\top} \Sigma x ^ {(i)}\right) \| x ^ {(i)} \| ^ {2} \right] = \frac {1}{2} \mathbb {E} [ \left(x ^ {\top} \Sigma x ^ {(i)}\right) ^ {2} \| x ^ {(i)} \| ^ {2} ] \\ = \frac {1}{2} \mathbb {E} \left[ \left(x ^ {(i)}\right) ^ {\top} \Sigma^ {2} x ^ {(i)} \| x ^ {(i)} \| ^ {2} \right] = \frac {d + 2}{2} \operatorname {T r} \left(\Sigma^ {2}\right). \\ \end{array}
$$
Putting things all together (and ignoring the constant part in (23)), we have
$$
f _ {\text {l o w e r}} (\Sigma , D) = \frac {2 (d + 2) + (n - 1)}{4 n} \operatorname {T r} \left(\Sigma^ {2}\right) + \operatorname {T r} (\Sigma D). \tag {24}
$$
4. Connecting back to the original loss function.
One can in fact write (24) solely in terms of  $b_{1}$  and  $A$  as follows:
$$
\frac {2 (d + 2) + (n - 1)}{4 n} \operatorname {T r} (\Sigma^ {2}) + \operatorname {T r} (\Sigma D) = \frac {2 (d + 2) + (n - 1)}{4 n} \| b _ {1} A \| _ {F} ^ {2} + \operatorname {T r} (b _ {1} A).
$$
Since the latter is a convex function in the matrix  $b_{1}A$ , it follows that the minimizer corresponds to
$$
b _ {1} A = - \frac {2 n}{2 (d + 2) + (n - 1)} \cdot I _ {d}
$$