SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In fact the choice  $b_{1} = 1$ ,  $b_{0} = 0$ , and  $A = -\frac{2n}{2(d + 2) + (n - 1)} \cdot I_{d}$  achieves this, and more crucially, satisfies the property that  $f_{\mathrm{lower}} = f$  for the corresponding parameters. Therefore, this shows that such choice is a global minimizer.
B Proofs for the multi-layer case
The proof is based on probabilistic methods (Alon and Spencer, 2016). According to Lemma 5, the objective function can be written as (for more details check the derivations in (25))
$$
\begin{array}{l} f \left(A _ {1}, A _ {2}\right) = \mathbb {E} \operatorname {T r} \left(\mathbb {E} \left[ \prod_ {i = 1} ^ {2} \left(I - X _ {0} ^ {\top} A _ {i} X _ {0} M\right) X _ {0} ^ {\top} w _ {\star} w _ {\star} ^ {\top} X _ {0} \prod_ {i = 1} ^ {2} \left(I - M X _ {0} ^ {T} A _ {i} X _ {0}\right) \right]\right) \\ = \mathbb {E} \operatorname {T r} \left(\mathbb {E} \left[ \prod_ {i = 2} ^ {1} \left(I - X _ {0} ^ {\top} A _ {i} X _ {0} M\right) X _ {0} ^ {\top} X _ {0} \prod_ {j = 1} ^ {2} \left(I - M X _ {0} ^ {T} A _ {j} X _ {0}\right) \right]\right), \\ \end{array}
$$
where we use the isotropy of  $w_{\star}$  and the linearity of trace to get the last equation. Suppose that  $A_0^*$  and  $A_1^*$  denote the global minimizer of  $f$  over symmetric matrices. Since  $A_1^*$  is a symmetric matrix, it admits the spectral decomposition  $A_1 = U D_1 U^\top$  where  $D_1$  is a diagonal matrix and  $U$  is an orthogonal matrix. Remarkably, the distribution of  $X_0$  is invariant to a linear transformation by an orthogonal matrix, i.e.,  $X_0$  has the same distribution as  $X_0 U^\top$ . This invariance yields
$$
f (U D _ {1} U ^ {\top}, A _ {2} ^ {*}) = f (D _ {1}, U ^ {\top} A _ {2} ^ {*} U).
$$
Thus, we can assume  $A_1^*$  is diagonal without loss of generality. To prove  $A_2^*$  is also diagonal, we leverage a probabilistic proof technique. Consider the random diagonal matrix  $S$  whose diagonal elements are either 1 or -1 with probability  $\frac{1}{2}$ . Since the input distribution is invariant to orthogonal transformations, we have
$$
f \left(D _ {1}, A _ {2} ^ {*}\right) = f \left(S D _ {1} S, S A _ {2} ^ {*} S\right) = f \left(D _ {1}, S A _ {2} ^ {*} S\right).
$$
Note that we use  $SD_{1}S = D_{1}$  in the last equation, which holds due to  $D_{1}$  and  $S$  are diagonal matrices and  $S$  has diagonal elements in  $\{+1, -1\}$ . Since  $f$  is convex in  $A_{2}$ , a straightforward application of Jensen's inequality yields
$$
f \left(D _ {1}, A _ {2} ^ {*}\right) = \mathbb {E} \left[ f \left(D _ {1}, S A _ {2} ^ {*} S\right) \right] \geq f \left(D _ {1}, \mathbb {E} \left[ S A _ {2} ^ {*} S \right]\right) = f \left(D _ {1}, \operatorname {d i a g} \left(A _ {2} ^ {*}\right)\right).
$$
Thus, there are diagonal  $D_{1}$  and  $\mathrm{diag}(A_2^*)$  for which  $f(D_{1},\mathrm{diag}(A_{2}^{*}))\leq f(A_{1}^{*},A_{2}^{*})$  holds for an optimal  $A_{1}^{*}$  and  $A_{2}^{*}$ . This concludes the proof.
Let us drop the factor of  $1 / n$  which was present in the original update (56). This is because the constant  $1 / n$  can be absorbed into  $A_{i}$ 's. Doing so does not change the theorem statement, but reduces notational clutter.
Let us consider the reformulation of the in-context loss  $f$  presented in Lemma 5. Specifically, let  $\overline{Z}_0$  be defined as
$$
\overline {{Z}} _ {0} = \left[ \begin{array}{c c c c c} x ^ {(1)} & x ^ {(2)} & \dots & x ^ {(n)} & x ^ {(n + 1)} \\ y ^ {(1)} & y ^ {(2)} & \dots & y ^ {(n)} & y ^ {(n + 1)} \end{array} \right] \in \mathbb {R} ^ {(d + 1) \times (n + 1)},
$$
where  $y^{(n + 1)} = \langle w_{\star},x^{(n + 1)}\rangle$  . Let  $\overline{Z}_i$  denote the output of the  $(i - 1)^{th}$  layer of the linear transformer (as defined in (56), initialized at  $\overline{Z}_0$  ). For the rest of this proof, we will drop the bar, and simply denote  $\overline{Z}_i$  by  $Z_{i}$  . Let  $X_{i}\in \mathbb{R}^{d\times (n + 1)}$  denote the first  $d$  rows of  $Z_{i}$  and let  $Y_{i}\in \mathbb{R}^{1\times (n + 1)}$  denote the  $(d + 1)^{th}$  row of  $Z_{k}$  . Under the sparsity pattern enforced in (8), we verify that, for any  $i\in \{0,\dots ,k\}$
$$
X _ {i} = X _ {0},
$$
$$
Y _ {i + 1} = Y _ {i} + Y _ {i} M X _ {i} ^ {\top} A _ {i} X _ {i} = Y _ {0} \prod_ {\ell = 0} ^ {i} \left(I + M X _ {0} ^ {\top} A _ {\ell} X _ {0}\right). \tag {25}
$$
This use of  $Z_{i}$  differs the original definition in (1). But we will not refer to the original definition anywhere in this proof.
where  $M = \begin{bmatrix} I_{n\times n} & 0\\ 0 & 0 \end{bmatrix}$ . We adopt the shorthand  $A = \{A_i\}_{i = 0}^k$
We adopt the shorthand  $A = \{A_i\}_{i=0}^k$ . Let  $S \subset \mathbb{R}^{(k+1) \times d \times d}$ , and  $A \in S$  if and only if for all  $i \in \{0, \ldots, k\}$ , there exists scalars  $a_i \in \mathbb{R}$  such that  $A_i = a_i \Sigma^{-1}$  and  $B_i = b_i I$ . We use  $f(A)$  to refer to the in-context loss of Theorem 3, that is,
$$
f (A) := f \left(\left\{Q _ {i} = \left[ \begin{array}{c c} A _ {i} & 0 \\ 0 & 0 \end{array} \right], P _ {i} = \left[ \begin{array}{c c} 0 _ {d \times d} & 0 \\ 0 & 1 \end{array} \right] \right\} _ {i = 0} ^ {k}\right).
$$
Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:
$$
f (A) = \mathbb {E} _ {\left(X _ {0}, w _ {\star}\right)} \left[ \operatorname {T r} \left(\left(I - M\right) Y _ {k + 1} ^ {\top} Y _ {k + 1} (I - M)\right) \right]. \tag {26}
$$
The theorem statement is equivalent to the following:
$$
\inf  _ {A \in \mathcal {S}} \sum_ {i = 0} ^ {k} \| \nabla_ {A i} f (A) \| _ {F} ^ {2} = 0, \tag {27}
$$
where  $\nabla_{A_i}f$  denotes derivative wrt the Frobenius norm  $\| A_i\| _F$ . Towards this end, we establish the following intermediate result: if  $A\in S$ , then for any  $R\in \mathbb{R}^{(k + 1)\times d\times d}$ , there exists  $\tilde{R}\in S$ , such that, at  $t = 0$ ,
$$
\frac {d}{d t} f (A + t \tilde {R}) \leq \frac {d}{d t} f (A + t R). \tag {28}
$$
In fact, we show that  $\tilde{R}_i\coloneqq r_iI$ , for  $r_i = \frac{1}{d}\operatorname {Tr}\left(\Sigma^{1 / 2}R_i\Sigma^{1 / 2}\right)$ . This implies (27) via the following simple argument: Consider the "S-constrained gradient flow": let  $A(t):\mathbb{R}^{+}\to \mathbb{R}^{(k + 1)\times d\times d}$  be defined as
$$
\frac {d}{d t} A _ {i} (t) = - r _ {i} (t) \Sigma^ {- 1}, \quad r _ {i} (t) := \mathrm {T r} (\Sigma^ {1 / 2} \nabla_ {A _ {i}} f (A (t)) \Sigma^ {1 / 2})
$$
for  $i = 0,\dots ,k$  .By (28),we verify that
$$
\frac {d}{d t} f (A (t)) \leq - \sum_ {i = 0} ^ {k} \| \nabla_ {A _ {i}} f (A (t)) \| _ {F} ^ {2}. \tag {29}
$$
We verify from its definition that  $f(A) \geq 0$ ; if the infimum in (27) fails to be zero, then inequality (29) will ensure unbounded descent as  $t \to \infty$ , contradicting the fact that  $f(A)$  is lower-bounded. This concludes the proof.
Proof outline. The remainder of the proof will be devoted to showing (28), which we outline as follows:
- In Step 1, we reduce the condition in (29) to a more easily verified layer-wise condition. Specifically, we only need to verify (29) when  $R_{i}$  are all zero except for  $R_{j}$  for some fixed  $j$  (see (30)) At the end of Step 1, we set up some additional notation, and introduce an important matrix  $G$ , which is roughly "a product of attention layer matrices". In (31), we study the evolution of  $f(A(t))$  when  $A(t)$  moves in the direction of  $R$ , as  $X_{0}$  is (roughly speaking) randomly transformed.
- In Step 2, we use the results of Step 2 to study  $G$  (see (32)) and  $\frac{d}{dt} G(A(t))$  (see (33)) under random transformation of  $X_0$ . The idea in (33) is that "randomly transforming  $X_0$ " has the same effect as "randomly transforming  $S$ " (recall  $S$  is the perturbation to  $B$ ).
- In Step 3, we apply the result from Step 2 to the expression of  $\frac{d}{dt} f(A(t))$  in (31). We verify that  $\tilde{R}$  in (28) is exactly the expected matrix after "randomly transforming  $S$ ". This concludes our proof.
1. Reduction to layer-wise condition. To prove (28), it suffices to show the following simpler condition: Let  $j \in \{0, \dots, k\}$ . Let  $R_j \in \mathbb{R}^{d \times d}$  be arbitrary matrices. For  $C \in \mathbb{R}^{d \times d}$ , let  $A(tC, j)$  denote
the collection of matrices, where  $[A(tC,j)]_j = A_j + tC$ , and for  $i \neq j$ ,  $A(tC,j)_i = A_i$ . We show that for all  $j \in \{0,\ldots,k\}$ ,  $R_j \in \mathbb{R}^{d \times d}$ , there exists  $\tilde{R}_j = r_j\Sigma^{-1}$ , such that, at  $t = 0$
$$
\frac {d}{d t} f (A (t \tilde {R} _ {j}, j)) \leq \frac {d}{d t} f (A (t R _ {j}, j)) \tag {30}
$$
We can verify that (28) is equivalent to (30) by noticing that for any  $R$ , at  $t = 0$ ,  $\frac{d}{dt} f(A + tR) = \sum_{j=0}^{k} \frac{d}{dt} f(A(tR_j, j))$ . We will now work towards proving (30) for some index  $j$  that is arbitrarily chosen but fixed throughout.
Let us define, for any  $C \in \mathbb{R}^{d \times d}$ ,  $G(X, A_j + C) \coloneqq X \prod_{i=0}^k (I - MX^\top [A(C,j)]_i X)$ . By (25) and (26),
$$
\begin{array}{l} f (A (t R _ {j}, j)) \\ = \mathbb {E} \left[ \operatorname {T r} \left(\left(I - M\right) Y _ {k + 1} ^ {\top} Y _ {k + 1} (I - M)\right) \right] \\ = \mathbb {E} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, A _ {j} + t R _ {j}) ^ {\top} w _ {\star} ^ {\top} w _ {\star} G (X _ {0}, A _ {j} + t R _ {j}) (I - M)\right) \right] \\ = \mathbb {E} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, A _ {j} + t R _ {j}) ^ {\top} \Sigma^ {- 1} G (X _ {0}, A _ {j} + t R _ {j}) (I - M)\right) \right] \\ \end{array}
$$
The second equality follows from plugging in (25). For the rest of this proof, let  $U$  denote a uniformly randomly sampled orthogonal matrix. Let  $U_{\Sigma} \coloneqq \Sigma^{1/2}U\Sigma^{-1/2}$ . Using the fact that  $X_0 \stackrel{d}{=} U_{\Sigma}X_0$ , we can verify
$$
\begin{array}{l} \left. \frac {d}{d t} f (A (t R _ {j}, j)) \right| _ {t = 0} \\ = \frac {d}{d t} \mathbb {E} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, A _ {j} + t R _ {j}) ^ {\top} \Sigma^ {- 1} G (X _ {0}, A _ {j} + t R _ {j}) (I - M)\right) \right] \Bigg | _ {t = 0} \\ = \frac {d}{d t} \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left((I - M) G (U _ {\Sigma} X _ {0}, A _ {j} + t R _ {j}) ^ {\top} \Sigma^ {- 1} G (U _ {\Sigma} X _ {0}, A _ {j} + t R _ {j}) (I - M)\right) \right] \Bigg | _ {t = 0} \\ = 2 \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left((I - M) G \left(U _ {\Sigma} X _ {0}, A _ {j}\right) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G \left(U _ {\Sigma} X _ {0}, A _ {j} + t R _ {j}\right) \Bigg | _ {t = 0} (I - M)\right) \right]. \tag {31} \\ \end{array}
$$
2.  $G$  and  $\frac{d}{dt} G$  under random transformation of  $X_0$ . We will now verify that  $G(U_{\Sigma}X_0, A_j) = U_{\Sigma}G(X_0, A_j)$ :
$$
\begin{array}{l} G (U _ {\Sigma} X _ {0}, A _ {j}) \\ = U _ {\Sigma} X _ {0} \prod_ {i = 0} ^ {k} \left(I + M X _ {0} ^ {T} U _ {\Sigma} ^ {\top} A _ {i} U _ {\Sigma} X _ {0}\right) \\ = U _ {\Sigma} G \left(X _ {0}, A _ {j}\right), \tag {32} \\ \end{array}
$$
where we use the fact that  $U_{\Sigma}^{\top}A_{i}U_{\Sigma} = U_{\Sigma}^{\top}(a_{i}\Sigma^{-1})U_{\Sigma} = A_{i}$ . Next, we verify that
$$
\begin{array}{l} \left. \frac {d}{d t} G \left(U _ {\Sigma} X _ {0}, A + t R _ {j}\right) \right| _ {t = 0} \\ = U _ {\Sigma} X _ {0} \left(\prod_ {i = 0} ^ {j - 1} \left(I + M X _ {0} ^ {T} A _ {i} X _ {0}\right)\right) M X _ {0} ^ {T} U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma} X _ {0} \prod_ {i = j + 1} ^ {k} \left(I + M X _ {0} ^ {T} A _ {i} X _ {0}\right) \\ = U _ {\Sigma} \frac {d}{d t} G \left(X _ {0}, A _ {j} + t U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma}\right) \tag {33} \\ \end{array}
$$
where the first equality again uses the fact that  $U_{\Sigma}^{\top}A_{i}U_{\Sigma} = A_{i}$ .
3. Putting everything together. Let us continue from (31). Plugging (32) and (33) into (31),
$$
\left. \frac {d}{d t} f (A (t R _ {j}, j)) \right| _ {t = 0}
$$
$$
\begin{array}{l} = 2 \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left(\left(I - M\right) G \left(U _ {\Sigma} X _ {0}, A _ {j}\right) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G \left(U _ {\Sigma} X _ {0}, A _ {j} + t R _ {j}\right) \Bigg | _ {t = 0} (I - M)\right) \right] \\ \stackrel {(i)} {=} 2 \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, A _ {j}) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G (X _ {0}, A _ {j} + t U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma}) \Big | _ {t = 0} (I - M)\right) \right] \\ = 2 \mathbb {E} _ {X _ {0}} \left[ \right. \operatorname {T r} \left( \right.(I - M) G (X _ {0}, A _ {j}) ^ {\top} \Sigma^ {- 1} \mathbb {E} _ {U} \left[ \right. \frac {d}{d t} G (X _ {0}, A _ {j} + t U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma}) \left. \right| _ {t = 0} \left. \right] (I - M)\left. \right)\left. \right] \\ \stackrel {(i i)} {=} 2 \mathbb {E} _ {X _ {0}} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, A _ {j}) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G (X _ {0}, A _ {j} + t \mathbb {E} _ {U} [ U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma} ]) \Big | _ {t = 0} (I - M)\right) \right] \\ = 2 \mathbb {E} _ {X _ {0}} \left[ \operatorname {T r} \left(\left(I - M\right) G \left(X _ {0}, A _ {j}\right) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G \left(X _ {0}, A _ {j} + t \cdot r _ {j} \Sigma^ {- 1}\right) \Bigg | _ {t = 0} (I - M)\right) \right] \\ = \left. \frac {d}{d t} f (A (t \cdot r _ {j} \Sigma^ {- 1}, j)) \right| _ {t = 0}, \\ \end{array}
$$
where  $r_j \coloneqq \frac{1}{d} \operatorname{Tr}\left(\Sigma^{1/2} R_j \Sigma^{1/2}\right)$ . In the above, (i) uses 1. (32) and (33), as well as the fact that  $U_{\Sigma}^{\top} \Sigma^{-1} U_{\Sigma} = \Sigma^{-1}$ . (ii) uses the fact that  $\left.\frac{d}{dt} G(X_0, A_j + tC)\right|_{t=0}$  is affine in  $C$ . To see this, one can verify from the definition of  $G$ , e.g. using similar algebra as (33), that  $\left.\frac{d}{dt} G(X_0, A_j + C)\right|$  is affine in  $C$ . Thus  $\mathbb{E}_U\left[G(X_0, A_j + tU_{\Sigma}^{\top} R_j U_{\Sigma})\right] = G(X_0, A_j + t\mathbb{E}_U\left[U_{\Sigma}^{\top} R_j U_{\Sigma}\right]$ .
The proof of Theorem 4 is similar to that of Theorem 3, and with a similar setup. However to keep the proof self-contained, we will restate the setup. Once again, we drop the factor of  $\frac{1}{n}$  which was present in the original update (56). This is because the constant  $1/n$  can be absorbed into  $A_i$ 's. Doing so does not change the theorem statement, but reduces notational clutter.
where  $y^{(n + 1)} = \langle w_{\star},x^{(n + 1)}\rangle$  . Let  $\overline{Z}_i$  denote the output of the  $(i - 1)^{th}$  layer of the linear transformer (as defined in (56), initialized at  $\overline{Z}_0$  ). For the rest of this proof, we will drop the bar, and simply denote  $\overline{Z}_i$  by  $Z_{i}$  . Let  $X_{i}\in \mathbb{R}^{d\times n + 1}$  denote the first  $d$  rows of  $Z_{i}$  and let  $Y_{i}\in \mathbb{R}^{1\times n + 1}$  denote the  $(d + 1)^{th}$  row of  $Z_{k}$  . Under the sparsity pattern enforced in (11), we verify that, for any  $i\in \{0,\dots ,k\}$
$$
X _ {i + 1} = X _ {i} + B _ {i} X _ {i} M X _ {i} ^ {\top} A _ {i} X _ {i}
$$
$$
Y _ {i + 1} = Y _ {i} + Y _ {i} M X _ {i} ^ {\top} A _ {i} X _ {i} = Y _ {0} \prod_ {\ell = 0} ^ {i} \left(I + M X _ {\ell} ^ {T} A _ {\ell} X _ {\ell}\right). \tag {34}
$$
We adopt the shorthand  $A = \{A_i\}_{i=0}^k$  and  $B = \{B_i\}_{i=0}^k$ . Let  $S \subset \mathbb{R}^{2 \times (k+1) \times d \times d}$ , and  $(A, B) \in S$  if and only if for all  $i \in \{0, \ldots, 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$ . Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:
$$
f (A, B) := \mathbb {E} _ {\left(X _ {0}, w _ {\star}\right)} \left[ \operatorname {T r} \left(\left(I - M\right) Y _ {k + 1} ^ {\top} Y _ {k + 1} (I - M)\right) \right]. \tag {35}
$$
(note that the only randomness in  $Z_0$  comes from  $X_0$  as  $Y_0$  is a deterministic function of  $X_0$ ). The theorem statement is equivalent to the following:
$$
\inf  _ {(A, B) \in \mathcal {S}} \sum_ {i = 0} ^ {k} \| \nabla_ {A _ {i}} f (A, B) \| _ {F} ^ {2} + \| \nabla_ {B _ {i}} f (A, B) \| _ {F} ^ {2} = 0 \tag {36}
$$
where  $\nabla_{A_i}f$  denotes derivative wrt the Frobenius norm  $\| A_i\| _F$
Our goal is to show that, if  $(A,B)\in S$ , then for any  $(R,S)\in \mathbb{R}^{2\times (k + 1)\times d\times d}$ , there exists  $(\tilde{R},\tilde{S})\in S$ , such that, at  $t = 0$
$$
\frac {d}{d t} f (A + t \tilde {R}, B + t \tilde {S}) \leq \frac {d}{d t} f (A + t R, B + t S). \tag {37}
$$
In fact, we show that  $\tilde{R}_i\coloneqq r_iI$ , for  $r_i = \frac{1}{d}\mathrm{Tr}\left(\Sigma^{1 / 2}R_i\Sigma^{1 / 2}\right)$  and  $\tilde{S}_i = s_iI$ , for  $s_i = \frac{1}{d}\mathrm{Tr}\left(\Sigma^{-1 / 2}S_i\Sigma^{1 / 2}\right)$ . This implies (36) via the following simple argument: Consider the "S-constrained gradient flow": let  $A(t):\mathbb{R}^{+}\to \mathbb{R}^{(k + 1)\times d\times d}$  and  $B(t):\mathbb{R}^{+}\to \mathbb{R}^{(k + 1)\times d\times d}$  be defined as
$$
\frac {d}{d t} A _ {i} (t) = - r _ {i} (t) \Sigma^ {- 1}, \quad r _ {i} (t) := \operatorname {T r} (\Sigma^ {1 / 2} \nabla_ {A _ {i}} f (A (t), B (t)) \Sigma^ {1 / 2})
$$
$$
\frac {d}{d t} B _ {i} (t) = - s _ {i} (t) \Sigma^ {- 1}, \quad s _ {i} (t) := \mathrm {T r} (\Sigma^ {- 1 / 2} \nabla_ {B _ {i}} f (A (t), B (t)) \Sigma^ {1 / 2}),
$$
for  $i = 0,\dots ,k$  .By (37),we verify that
$$
\frac {d}{d t} f (A (t), B (t)) \leq - \left(\sum_ {i = 0} ^ {k} \| \nabla_ {A _ {i}} f (A (t), B (t)) \| _ {F} ^ {2} + \| \nabla_ {B _ {i}} f (A (t), B (t)) \| _ {F} ^ {2}\right). \tag {38}
$$
We verify from its definition that  $f(A, B) \geq 0$ ; if (36) does not hold then (38) will ensure unbounded descent as  $t \to \infty$ , contradicting the fact that  $f(A, B)$  is lower-bounded. This concludes the proof.
Proof outline. The remainder of the proof will be devoted to showing (37), which we outline as follows:
- In Step 1, we reduce the condition in (37) to a more easily verified layer-wise condition. Specifically, we only need to verify (37) in one of the two cases: (I) when  $R_{i}, S_{i}$  are all zero except for  $R_{j}$  for some fixed  $j$  (see (40)), or (II) when  $R_{i}, S_{i}$  are all zero except for  $S_{j}$  for some fixed  $j$  (see (39)). We focus on the proof of (II), as the proof of (I) is almost identical. At the end of Step 1, we set up some additional notation, and introduce an important matrix  $G$ , which is roughly "a product of attention layer matrices". In (41), we study the evolution of  $f(A, B(t))$  when  $B(t)$  moves in the direction of  $S$ , as  $X_{0}$  is (roughly speaking) randomly transformed. This motivates the subsequent analysis in Steps 2 and 3 below.
- In Step 2, we study how outputs of each layer (34) changes when  $X_0$  is randomly transformed. There are two main results here: First we provide the expression for  $X_i$  in (42). Second, we provide the expression for  $\frac{d}{dt} X_i(B(t))$  in (43).
- In Step 3, we use the results of Step 2 to study  $G$  (see (47)) and  $\frac{d}{dt} G(B(t))$  (see (48)) under random transformation of  $X_0$ .
The idea in (48) is that "randomly transforming  $X_0$ " has the same effect as "randomly transforming  $S$ " (recall  $S$  is the perturbation to  $B$ ).
- In Step 4, we use the results from Steps 2 and 3 to the expression of  $\frac{d}{dt} f(A, B(t))$  in (41). We verify that  $\tilde{S}$  in (37) is exactly the expected matrix after "randomly transforming  $S$ ". This concludes our proof of (II).
- In Step 5, we sketch the proof of (I), which is almost identical to Steps 2-4.
1. Reduction to layer-wise condition. To prove (37), it suffices to show the following simpler condition: Let  $j \in \{0, \dots, k\}$ . Let  $R_j, S_j \in \mathbb{R}^{d \times d}$  be arbitrary matrices. For  $C \in \mathbb{R}^{d \times d}$ , let  $A(tC, j)$  denote the collection of matrices, where  $A(tC, j)_j = A_j + tC$ , and for  $i \neq j$ ,  $A(tC, j)_i = A_i$ . Define  $B(tC, j)$  analogously. We show that for all  $j \in \{0, \dots, k\}$  and all  $R_j, S_j \in \mathbb{R}^{d \times d}$ , there exists  $\tilde{R}_j = r_j\Sigma^{-1}$  and  $\tilde{S}_j = s_j\Sigma^{-1}$ , such that, at  $t = 0$ ,
$$
\frac {d}{d t} f (A (t \tilde {R} _ {j}, j), B) \leq \frac {d}{d t} f (A (t R _ {j}, j), B) \tag {39}
$$
$$
\text {a n d} \quad \frac {d}{d t} f (A, B (t \tilde {S} _ {j}, j)) \leq \frac {d}{d t} f (A, B (t S _ {j}, j)). \tag {40}
$$
We can verify that (37) is equivalent to (39)+(40) by noticing that for any  $(R,S)\in \mathbb{R}^{2\times (k + 1)\times d\times d}$ , at  $t = 0$ ,  $\frac{d}{dt} f(A + tR,B + tS) = \sum_{j = 0}^{k}\left(\frac{d}{dt} f(A(tR_j,j),B) + \frac{d}{dt} f(A,B(tS_j,j))\right)$ .
We will first focus on proving (40) (the proof of (39) is similar, and we present it in Step 5 at the end), for some index  $j$  that is arbitrarily chosen but fixed throughout. Notice that  $X_{i}$  and  $Y_{i}$  in (34) are in fact functions of  $A, B$  and  $X_{0}$ . For most of our subsequent discussion,  $A_{i}$  (for all  $i$ ) and  $B_{i}$  (for all  $i \neq j$ ) can be treated as constant matrices. We will however make the dependence on  $X_{0}$  and  $B_{j}$  explicit (as we consider the curve  $B_{j} + tS$ ), i.e. we use  $X_{i}(X,C)$  (resp  $Y_{i}(X,C)$ ) to denote the value of  $X_{i}$  (resp  $Y_{i}$ ) from (34), with  $X_{0} = X$ , and  $B_{j} = C$ .
By (35) and (34),
$$
\begin{array}{l} f (A, B (t S _ {j}, j)) \\ = \mathbb {E} \left[ \operatorname {T r} \left(\left(I - M\right) Y _ {k + 1} \left(X _ {0}, B _ {j} + t S\right) ^ {\top} Y _ {k + 1} \left(X _ {0}, B _ {j} + t S _ {j}\right) (I - M)\right) \right] \\ = \mathbb {E} \left[ \operatorname {T r} \left((I - M) G \left(X _ {0}, B _ {j} + t S _ {j}\right) ^ {\top} w _ {\star} ^ {\top} w _ {\star} G \left(X _ {0}, B _ {j} + t S _ {j}\right) (I - M)\right) \right] \\ = \mathbb {E} \left[ \operatorname {T r} \left((I - M) G \left(X _ {0}, B _ {j} + t S _ {j}\right) ^ {\top} \Sigma^ {- 1} G \left(X _ {0}, B _ {j} + t S _ {j}\right) (I - M)\right) \right] \\ \end{array}
$$
where  $G(X, C) \coloneqq X \prod_{i=0}^{k} \left( I - M X_i(X, C)^T A_i X_i(X, C) \right)$ . The second equality follows from plugging in (34).
For the rest of this proof, let  $U$  denote a uniformly randomly sampled orthogonal matrix. Let  $U_{\Sigma} \coloneqq \Sigma^{1/2}U\Sigma^{-1/2}$ . Using the fact that  $X_0 \stackrel{d}{=} U_{\Sigma}X_0$ , we can verify
$$
\begin{array}{l} \left. \frac {d}{d t} f (A, B (t S _ {j}, j)) \right| _ {t = 0} \\ = \frac {d}{d t} \mathbb {E} _ {X _ {0}} \left[ \operatorname {T r} \left((I - M) G (X _ {0}, B _ {j} + t S _ {j}) ^ {\top} \Sigma^ {- 1} G (X _ {0}, B _ {j} + t S _ {j}) (I - M)\right) \right] \Bigg | _ {t = 0} \\ = \frac {d}{d t} \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left(\left(I - M\right) G \left(U _ {\Sigma} X _ {0}, B _ {j} + t S _ {j}\right) ^ {\top} \Sigma^ {- 1} G \left(U _ {\Sigma} X _ {0}, B _ {j} + t S _ {j}\right) (I - M)\right) \right] \Bigg | _ {t = 0} \\ = 2 \mathbb {E} _ {X _ {0}, U} \left[ \operatorname {T r} \left((I - M) G \left(U _ {\Sigma} X _ {0}, B _ {j}\right) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G \left(U _ {\Sigma} X _ {0}, B _ {j} + t S _ {j}\right) \Bigg | _ {t = 0} (I - M)\right) \right]. \tag {41} \\ \end{array}
$$
2.  $X_{i}$  and  $\frac{d}{dt} X_{i}$  under random transformation of  $X_{0}$ . In this step, we prove that when  $X_{0}$  is transformed by  $U_{\Sigma}$ ,  $X_{i}$  for  $i \geq 1$  are likewise transformed in a simple manner. The first goal of this step is to show
$$
X _ {i} \left(U _ {\Sigma} X _ {0}, B _ {j}\right) = U _ {\Sigma} X _ {i} \left(X _ {0}, B _ {j}\right). \tag {42}
$$
We will prove this by induction. When  $i = 0$ , this clearly holds by definition. Suppose that (42) holds for some  $i$ . Then