SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In (i) above, we the following facts: 1.  $X_{i}(U_{\Sigma}X_{0},B_{j}) = U_{\Sigma}X_{i}(X_{0},B_{j})$  from (49), 2.  $A_{i} = a_{i}\Sigma^{-1}$  so that  $U_{\Sigma}^{\top}A_{i}U_{\Sigma} = A_{i}$ , 3.  $U_{\Sigma}U_{\Sigma}^{-1} = U_{\Sigma}^{-1}U_{\Sigma} = I$ . (ii) follows from (50).
We will now simplify  $\heartsuit$
$$
\begin{array}{l} = U _ {\Sigma} ^ {- 1} U _ {\Sigma} X _ {0} \left(\prod_ {i = 0} ^ {j - 1} \left(I + M X _ {i} \left(U _ {\Sigma} X _ {0}, A _ {j}\right) ^ {T} A _ {i} X _ {i} \left(U _ {\Sigma} X _ {0}, A _ {j}\right)\right)\right) \\ \cdot M X _ {j} \left(U _ {\Sigma} X _ {0}, A _ {j}\right) ^ {T} R _ {j} X _ {j} \left(U _ {\Sigma} X _ {0}, A _ {j}\right) \\ \cdot \left(\prod_ {i = j + 1} ^ {k} \left(I + M X _ {i} \left(U _ {\Sigma} X _ {0}, A _ {j}\right) ^ {T} A _ {i} X _ {i} \left(U _ {\Sigma} X _ {0}, A _ {j}\right)\right)\right) \\ \stackrel {(i)} {=} X _ {0} \left(\prod_ {i = 0} ^ {j - 1} \left(I + M X _ {i} (X _ {0}, A _ {j}) ^ {T} A _ {i} X _ {i} (X _ {0}, A _ {j})\right)\right) M X _ {j} (X _ {0}, A _ {j}) ^ {\top} U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma} X _ {j} (X _ {0}, A _ {j}) \\ \cdot \left(\prod_ {i = j + 1} ^ {k} \left(I + M X _ {i} \left(X _ {0}, A _ {j}\right) ^ {T} A _ {i} X _ {i} \left(X _ {0}, A _ {j}\right)\right)\right), \tag {55} \\ \end{array}
$$
where  $(i)$  uses the fact that  $X_{i}(U_{\Sigma}X_{0},B_{j}) = U_{\Sigma}X_{i}(X_{0},B_{j})$  from (49) and the fact that  $A_{i} = a_{i}\Sigma^{-1}$ .
By expanding  $\frac{d}{dt} G(X_0, A_j + tU_{\Sigma}^\top R_j U_{\Sigma})$ , we verify that
$$
\frac {d}{d t} G (X _ {0}, A _ {j} + t U _ {\Sigma} ^ {\top} R _ {j} U _ {\Sigma}) = (5 4) + (5 5) = \spadesuit + \heartsuit = U _ {\Sigma} ^ {- 1} \frac {d}{d t} G (U _ {\Sigma} X _ {0}, A _ {j} + t R _ {j}),
$$
this concludes the proof of (53).
The remainder of the proof is similar to what was done in (41) in Step 4:
$$
\begin{array}{l} \left. \frac {d}{d t} f (A (t R _ {j}, j), B \right| _ {t = 0} \\ = 2\mathbb{E}_{X_{0},U}\left[\operatorname {Tr}\left((I - M)G(U_{\Sigma}X_{0},A_{j})^{\top}\Sigma^{-1}\frac{d}{dt} G(U_{\Sigma}X_{0},A_{j} + tR_{j})\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] \\ \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((I - M) G (X _ {0}, A _ {j}) ^ {\top} \Sigma^ {- 1} \frac {d}{d t} G (X _ {0}, A _ {j} + t \cdot r _ {j} \Sigma^ {- 1}) \Big | _ {t = 0} (I - M)\right) \right] \\ = \left. \frac {d}{d t} f (A (t \cdot r _ {j} \Sigma^ {- 1}, j), B) \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. (49) and (53), 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 using a simple induction argument, that  $\left.\frac{d}{dt} X_i(X_0, A_j + tC)\right|$  is affine in  $C$  for all  $i$ . We can then verify from the definition of  $G$ , e.g., using similar algebra as the proof of (53), that  $\frac{d}{dt} G(X_0, A_j + C)$  is affine in  $X_i(X_0, A_j + tC)$  and  $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])$ .
This concludes the proof of (39), and hence of the whole theorem.
B.4 Equivalence under permutation
Lemma 4. Consider the same setup as Theorem 3. Let  $A = \{A_i\}_{i=0}^k$ , with  $A_i = a_i \Sigma^{-1}$ . Let
Let  $i, j \in \{0, \ldots, k\}$  be any two arbitrary indices, and let  $\tilde{A}_i = A_j$ ,  $\tilde{A}_j = A_i$ , and let  $\tilde{A}_{\ell} = A_{\ell}$  for all  $\ell \in \{0, \ldots, k\} \setminus \{i, j\}$ . Then  $f(A) = f(\tilde{A})$
Proof. Following the same setup leading up to (26) in the proof of Theorem 3, we verify that the in-context loss is
$$
f (A) = \mathbb {E} \left[ \operatorname {T r} \left(\left(I - M\right) G \left(X _ {0}, A\right) ^ {\top} \Sigma^ {- 1} G \left(X _ {0}, A\right) \left(I - M\right)\right) \right]
$$
where  $G(X_0, A) \coloneqq X_0 \prod_{\ell=0}^k \left(I + MX_0^T A_\ell X_0\right)$ .
Consider any fixed index  $\ell$ . We will show that
$$
\left(I + M X _ {0} ^ {T} A _ {\ell} X _ {0}\right) \left(I + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0}\right) = \left(I + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0}\right) \left(I + M X _ {0} ^ {T} A _ {\ell} X _ {0}\right).
$$
The lemma can then be proven by repeatedly applying the above, so that indices of  $A_{i}$  and  $A_{j}$  are swapped.
To prove the above equality,
$$
\begin{array}{l} \left(I + M X _ {0} ^ {T} A _ {\ell} X _ {0}\right) \left(I + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0}\right) \\ = I + M X _ {0} ^ {T} A _ {\ell} X _ {0} + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0} + M X _ {0} ^ {T} A _ {\ell} X _ {0} M X _ {0} ^ {T} A _ {\ell + 1} X _ {0} \\ = I + M X _ {0} ^ {T} A _ {\ell} X _ {0} + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0} + M X _ {0} ^ {T} a _ {\ell} \Sigma^ {- 1} X _ {0} M X _ {0} ^ {T} a _ {\ell + 1} \Sigma^ {- 1} X _ {0} \\ = I + M X _ {0} ^ {T} A _ {\ell} X _ {0} + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0} + M X _ {0} ^ {T} a _ {\ell + 1} \Sigma^ {- 1} X _ {0} M X _ {0} ^ {T} a _ {\ell} \Sigma^ {- 1} X _ {0} \\ = \left(I + M X _ {0} ^ {T} A _ {\ell + 1} X _ {0}\right) \left(I + M X _ {0} ^ {T} A _ {\ell} X _ {0}\right). \\ \end{array}
$$
This concludes the proof. Notice that we crucially used the fact that  $A_{\ell}$  and  $A_{\ell +1}$  are the same matrix up to scaling.
C.1 Proof of Lemma 1 (Equivalence to Preconditioned Gradient Descent)
Consider fixed samples  $x^{(1)}, \ldots, x^{(n)}$ , and fixed  $w_{\star}$ . Let  $P = \{P_i\}_{i=0}^k$ ,  $Q = \{Q_i\}_{i=0}^k$  denote fixed weights. Let  $Z_i$  evolve as described in (4). Let  $X_i$  denote the first  $d$  rows of  $Z_k$  (under (8),  $X_i = X_0$  for all  $I$ ) and let  $Y_i$  denote the  $(d+1)^{th}$  row of  $Z_i$ . Let  $g(x,y,k): \mathbb{R}^d \times \mathbb{R} \times \mathbb{Z} \to \mathbb{R}$  be a function defined as follows: let  $x^{n+1} = x$  and let  $y_0^{n+1} = y$ , then  $g(x,y,k) := y_k^{n+1}$ . Note that  $y_k^{n+1} = [Y_k]_{n+1}$ .
We verify that, under (8), the formula for updating  $y_{k}^{(n + 1)}$  is given by
$$
Y _ {k + 1} = Y _ {k} - \frac {1}{n} Y _ {k} M X _ {0} ^ {\top} A _ {k} X _ {0}.
$$
where  $M$  is a mask given by  $\begin{bmatrix} I & 0 \\ 0 & 0 \end{bmatrix}$ . We can verify the following facts
1.  $g(x,y,k) = g(x,0,k) + y$ . To see this, notice first that for all  $i\in \{1,\dots ,n\}$ ,
$$
y _ {k + 1} ^ {(i)} = y _ {k} ^ {(i)} - \frac {1}{n} \sum_ {j = 1} ^ {n} x ^ {(i) ^ {T}} A _ {k} x ^ {(j)} y _ {k} ^ {(j)}.
$$
In other words,  $y_{k}^{(i)}$  does not depend on  $y_{t}^{(n + 1)}$  for any  $t$ . Next, for  $y_{k}^{(n + 1)}$  itself,
$$
y _ {k + 1} ^ {(n + 1)} = y _ {k} ^ {(n + 1)} - \frac {1}{n} \sum_ {j = 1} ^ {n} x ^ {(n + 1) ^ {T}} A _ {k} x ^ {(j)} y _ {k} ^ {(j)},
$$
which depends on  $y_{k}^{n + 1}$  only additively. We can verify under a simple induction that  $g(x,y,k + 1) - y = g(x,y,k) - y$ .
2.  $g(x,0,k)$  is linear in  $x$ . To see this, notice first that for  $j \neq n + 1$ ,  $y_k^{(j)}$  is does not depend on  $x_t^{(n + 1)}$  for all  $t,j,k$ . Consequently, the update formula for  $y_{k + 1}^{(n + 1)}$  depends only linearly on  $x^{(n + 1)}$  and  $y_k^{(n + 1)}$ . Finally,  $y_0^{(n + 1)} = 0$  is linear in  $x$ , so the conclusion follows by induction.
With these two facts in mind, we verify that for each  $k$ , there exists a  $\theta_{k} \in \mathbb{R}^{d}$ , such that
$$
g (x, y, k) = g (x, 0, k) + y = \langle \theta_ {k}, x \rangle + y
$$
for all  $x, y$ . It follows from definition that  $g(x, y, 0) = y$ , so that  $\langle \theta_0, x \rangle = g(x, y, 0) - y = 0$ , so that  $\theta_0 = 0$ .
We now turn our attention to the third crucial fact: for all  $i$ ,
$$
g (x ^ {(i)}, y ^ {(i)}, k) = y _ {k} ^ {(i)} = \left\langle \theta_ {k}, x ^ {(i)} \right\rangle + y ^ {(i)}
$$
To see this, suppose that we let  $x^{(n + 1)} \coloneqq x^{(i)}$  for some  $i \in 1, \ldots, n$ . Then
$$
y _ {k + 1} ^ {(i)} = y _ {k} ^ {(i)} - \frac {1}{n} \sum_ {j = 1} ^ {n} x ^ {(i) T} A _ {k} x ^ {(j)} y _ {k} ^ {(j)}
$$
thus  $y_{k+1}^{(i)} = y_{k+1}^{(n+1)}$  if  $y_k^{(i)} = y_k^{(n+1)}$ , and the induction proof is completed by noting that  $y_0^{(i)} = y_0^{(n+1)}$  by definition. Let  $\bar{X} \in R^{d \times n}$  be the matrix whose columns are  $x^{(1)}, \ldots, x^{(n)}$ , leaving out  $x^{(n+1)}$ . Let  $\bar{Y}_k \in \mathbb{R}^{1 \times n}$  denote the vector of  $y_k^{(1)}, \ldots, y_k^{(n)}$ . Then it follows that
$$
\bar {Y} _ {k} = \bar {Y} _ {0} + \theta_ {k} ^ {T} \bar {X}.
$$
Using the above fact, the update formula for  $y_{k}^{(n + 1)}$  can be written as
$$
y _ {k + 1} ^ {(n + 1)} = y _ {k} ^ {(n + 1)} - \frac {1}{n} \left\langle A _ {k} X ^ {\top} Y _ {k}, x ^ {(n + 1)} \right\rangle
$$
$$
\begin{array}{l} \Rightarrow \quad \left\langle \theta_ {k + 1}, x ^ {(n + 1)} \right\rangle = \left\langle \theta_ {k}, x ^ {(n + 1)} \right\rangle - \frac {1}{n} \left\langle A _ {k} \bar {X} \left(\bar {X} ^ {T} \theta_ {k} + \bar {Y} _ {0}\right), x ^ {(n + 1)} \right\rangle \\ = \left\langle \theta_ {k}, x ^ {(n + 1)} \right\rangle - \frac {1}{n} \left\langle A _ {k} \bar {X} \left(\bar {X} ^ {T} \left(\theta_ {k} + w _ {\star}\right)\right), x ^ {(n + 1)} \right\rangle \\ \end{array}
$$
Since the choice of  $x^{(n + 1)}$  is arbitrary, we get the more general update formula
$$
\theta_ {k + 1} = \theta_ {k} - \frac {1}{n} A _ {k} \bar {X} \bar {X} ^ {T} \left(\theta_ {k} + w _ {\star}\right).
$$
We can treat  $A_{k}$  as a preconditioner. Let  $f(\theta)\coloneqq \frac{1}{2n} (\theta +w_{\star})^{T}\bar{X}\bar{X}^{T}(\theta +w_{\star})$  , then
$$
\theta_ {k + 1} = \theta_ {k} - \frac {1}{n} A _ {k} \nabla f (\theta).
$$
Finally, let  $w_{k}^{\mathrm{gd}} \coloneqq -\theta_{k}$ . We verify that  $f(-w) = R_{w_{*}}(w)$ , so that
$$
w _ {k + 1} ^ {\mathrm {g d}} = w _ {k} ^ {\mathrm {g d}} - \frac {1}{n} A _ {k} \nabla R _ {w _ {\star}} (w _ {k} ^ {\mathrm {g d}}).
$$
We also verify that for any  $x^{(n + 1)}$ , the prediction of  $y_{k}^{(n + 1)}$  is
$$
g \left(x ^ {(n + 1)}, y ^ {(n + 1)}, k\right) = y ^ {(n + 1)} - \left\langle \theta , x ^ {(n + 1)} \right\rangle = y ^ {(n + 1)} + \left\langle w _ {k} ^ {\mathbf {g d}}, x ^ {(n + 1)} \right\rangle .
$$
C.2 Reformulating the in-context loss
In this section, we will develop a re-formulation in-context loss, defined in (5), in a more convenient form (see Lemma 5).
For the entirety of this section, we assume that the transformer parameters  $\{P_i,Q_i\}_{i = 0}^k$  are of the form defined in (11), which we reproduce below for ease of reference:
$$
P _ {i} = \left[ \begin{array}{c c} B _ {i} & 0 \\ 0 & 1 \end{array} \right], \quad Q _ {i} = \left[ \begin{array}{c c} A _ {i} & 0 \\ 0 & 0 \end{array} \right].
$$
Recall the update dynamics in (4), which we reproduce below:
$$
Z _ {i + 1} = Z _ {i} + \frac {1}{n} P Z _ {i} M Z _ {i} ^ {\top} Q Z _ {i}, \tag {56}
$$
where  $M$  is a mask matrix given by  $M \coloneqq \left[ \begin{array}{cc} I_{n\times n} & 0 \\ 0 & 0 \end{array} \right]$ . Let  $X_{k} \in \mathbb{R}^{d\times n + 1}$  denote the first  $d$  rows of  $Z_{k}$  and let  $Y_{k} \in \mathbb{R}^{1\times n + 1}$  denote the  $(d + 1)^{th}$  (last) row of  $Z_{k}$ . Then the dynamics in (56) is equivalent to
$$
X _ {i + 1} = X _ {i} + \frac {1}{n} B _ {i} X _ {i} M X _ {i} ^ {T} A _ {i} X _ {i}
$$
$$
Y _ {i + 1} = Y _ {i} + \frac {1}{n} Y _ {i} M X _ {i} ^ {T} A _ {i} X _ {i}. \tag {57}
$$
We present below an equivalent form for the in-context loss from (5):
Lemma 5. Let  $p_x$  and  $p_w$  denote distributions over  $\mathbb{R}^d$ . Let  $x^{(1)}, \ldots, x^{(n+1)} \stackrel{iid}{\sim} p_x$  and  $w_\star \sim p_w$ . Let  $Z_0 \in \mathbb{R}^{d+1 \times n+1}$  be as defined in (1):
$$
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)} & 0 \end{array} \right] \in \mathbb {R} ^ {(d + 1) \times (n + 1)}.
$$
Let  $Z_{k}$  denote the output of the  $(k - 1)^{th}$  layer of the linear transformer (as defined in (56), initialized at  $Z_{0}$ ). Let  $f\left(\{P_{i}, Q_{i}\}_{i = 0}^{k}\right)$  denote the in-context loss defined in (5), i.e.
$$
f \left(\left\{P _ {i}, Q _ {i} \right\} _ {i = 0} ^ {k}\right) = \mathbb {E} _ {\left(Z _ {0}, w _ {\star}\right)} \left[ \left(\left[ Z _ {k} \right] _ {(d + 1), (n + 1)} + w _ {\star} ^ {\top} x ^ {(n + 1)}\right) ^ {2} \right]. \tag {58}
$$
Let  $\overline{Z}_0$  be defined as
where  $y^{(n + 1)} = \langle w_{\star},x^{(n + 1)}\rangle$ . Let  $\overline{Z}_k$  denote the output of the  $(k - 1)^{th}$  layer of the linear transformer (as defined in (56), initialized at  $\overline{Z}_0$ ). Assume  $\{P_i,Q_i\}_{i = 0}^k$  be of the form in (11). Then the loss in (5) has the equivalent form
$$
f \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}\right) := f \left(\left\{P _ {i}, Q _ {i} \right\} _ {i = 0} ^ {k}\right) = \mathbb {E} _ {\left(\bar {Z} _ {0}, w _ {\star}\right)} \left[ \operatorname {T r} \left(\left(I - M\right) \bar {Y} _ {k + 1} ^ {\top} \bar {Y} _ {k + 1} (I - M)\right) \right],
$$
where  $\overline{Y}_{k + 1}\in \mathbb{R}^{1\times n + 1}$  is the  $(d + 1)^{th}$  row of  $\overline{Z}_k$
Before proving Lemma 5, we first establish an intermediate result (Lemma 6 below). To facilitate discussion, let us define a function  $F_{X}\left(\{A_{i},B_{i}\}_{i = 0}^{k},X_{0},Y_{0}\right)$  and  $F_{Y}\left(\{A_{i},B_{i}\}_{i = 0}^{k},X_{0},Y_{0}\right)$  to be the outputs, after  $k$  layers of linear transformers respectively. I.e.
$$
F _ {X} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0}\right) = X _ {k + 1}
$$
$$
F _ {Y} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0}\right) = Y _ {k + 1},
$$
as defined in (57), given initialization  $X_0, Y_0$ .
We now prove a useful lemma showing that  $[Y_0]_{n + 1} = y^{(n + 1)}$  influences  $X_{i},Y_{i}$  in a very simple manner:
Lemma 6. Let  $X_{i}, Y_{i}$  follow the dynamics in (57). Then
1.  $[X_i]$  is are independent of  $[Y_0]_{n + 1}$ .
2. For  $j \neq n + 1$ ,  $[Y_i]_j$  is independent of  $[Y_0]_{n+1}$ .
3.  $[Y_i]_{n+1}$  depends additively on  $[Y_0]_{n+1}$ .
In other words, for  $C \coloneqq [0, 0, 0, \dots, 0, c] \in \mathbb{R}^{1 \times (n + 1)}$ ,
$$
\begin{array}{l} 1: F _ {X} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0} + C\right) = F _ {X} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0}\right) \\ 2 + 3: F _ {Y} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0} + C\right) = F _ {Y} \left(\left\{A _ {i}, B _ {i} \right\} _ {i = 0} ^ {k}, X _ {0}, Y _ {0}\right) + C \\ \end{array}
$$
Proof of Lemma 6. The first and second items follow directly from observing that the dynamics for  $X_{i}$  and  $Y_{i}$  in (57) do not involve  $[Y_i]_{n + 1}$ , due to the effect of  $M$ .
The third item again uses the fact that  $[Y_{i + 1} - Y_i]_{n + 1}$  does not depend on  $[Y_i]_{n + 1}$ .
We are now ready to prove Lemma 5
Proof of Lemma 5. Let  $Z_0, Z_k, \overline{Z}_0, \overline{Z}_k$  be as defined in the lemma statement. Let  $\overline{X}_k$  and  $\overline{Y}_k$  denote first  $d$  rows and last row of  $\overline{Z}_k$ . Then by Lemma 6,  $\overline{X}_{k+1} = X_{k+1}$  and  $\overline{Y}_{k+1} = Y_{k+1} + \left[ \begin{array}{cccc} 0 & 0 & \dots & 0 \end{array} \langle w_\star, x^{(n+1)} \rangle \right]$ . Therefore, (58) is equivalent to
$$
\begin{array}{l} \mathbb {E} _ {\left(\bar {Z} _ {0}, w _ {\star}\right)} \left[ \left(\left[ \bar {Z} _ {k + 1} \right] _ {(d + 1), (n + 1)}\right) ^ {2} \right] \\ = \mathbb {E} _ {\left(\bar {Z} _ {0}, w _ {\star}\right)} \left[ \left(\left[ \bar {Y} _ {k + 1} \right] _ {(n + 1)}\right) ^ {2} \right] \\ = \mathbb {E} _ {\left(\bar {Z} _ {0}, w _ {\star}\right)} \left[ \left\| (I - M) \bar {Y} _ {k + 1} ^ {\top} \right\| ^ {2} \right] \\ = \mathbb {E} _ {\left(\bar {Z} _ {0}, w _ {\star}\right)} \left[ \operatorname {T r} \left((I - M) \bar {Y} _ {k + 1} ^ {\top} \bar {Y} _ {k + 1} (I - M)\right) \right]. \\ \end{array}
$$
In this section, we present a few addition experimental results. We first present in Figure 5 a visualization of learned weights  $A_0$ ,  $A_1$ ,  $A_2$  for the setting of Theorem 4. One can see that the weight pattern matches the stationary point analyzed in Theorem 4; hence, combining Figure 4 and Figure 5, we corroborate our results from Theorem 4. Interestingly, it appears that the transformer implements a tiny gradient step using  $X_0$  (as  $A_0$  is small), and a large gradient step using  $X_2$  (as  $A_2$  is large). We believe that this is due to  $X_2$  being better-conditioned than  $X_1$ , due to the effects of  $B_0$ ,  $B_1$ .
Figure 5: Visualization of learned weights for the setting of Theorem 4. One can see that the weight pattern matches the stationary point analyzed in Theorem 4.
We next present some additional experiments that investigates the properties of the learned predictors of various algorithms. First, we plot the test losses against the number of examples provided in the prompt ("the number of ICL examples"). We compare four different algorithms: (i) the predictor learned by a three-layered of linear transformer, (ii) three steps of GD, (iii) three steps of preconditioned GD, and (iv) the ordinary least-squared solution (OLS). For GD and preconditioned GD, the optimal step sizes are found by gridsearch. For preconditioned GD, preconditioner is fixed to be  $\Sigma^{-1}$  for comparison. In all cases, the dimension  $d = 5$ , and for each  $N$ , the linear Transformer is trained using Adam. The result is presented in Figure 6.
Figure 6: Test loss comparison between (i) the predictor learned by a three-layered of linear transformer, (ii) three steps of GD, (iii) three steps of preconditioned GD, and (iv) the ordinary least-squared solution (OLS).
Lastly, in Figure 7, we plot the test losses against the number of layer  $L$  (or the number of steps in the case of gradient-based algorithms). For  $L = 1,2,3,4$ , we compare between (i) the predictor learned by  $L$ -linear transformer and (ii)  $L$ -steps of GD, (iii)  $L$ -steps of preconditioned GD. Again, the optimal stepsize is found by gridsearch, and for preconditioned GD, the preconditioner is fixed to be  $\Sigma^{-1}$ . In all cases, the dimension  $d = 5$ , and context length  $N = 20$ . The linear transformer is trained with Adam.
Figure 7: Test loss comparison between (i) the predictor learned by a  $L$ -layered linear transformer and (ii)  $L$ -steps of GD, (iii)  $L$ -steps of preconditioned GD, for  $L = 1,2,3,4$ .