SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {L} = - \frac {\left(W _ {\mathrm {P}} z ^ {(1)}\right) ^ {\top} \operatorname {S G} (z ^ {(2)})}{\| W _ {\mathrm {P}} z ^ {(1)} \| \| \operatorname {S G} (z ^ {(2)}) \|} \\ = - \frac {\left(\boldsymbol {z} ^ {(1)}\right) ^ {\top} U D U ^ {\top} \operatorname {S G} \left(\boldsymbol {z} ^ {(2)}\right)}{\| U D U ^ {\top} \boldsymbol {z} ^ {(1)} \| \| \operatorname {S G} \left(U U ^ {\top} \boldsymbol {z} ^ {(2)}\right) \|} \\ = - \frac {\left(\hat {\boldsymbol {z}} ^ {(1)}\right) ^ {\top} D \operatorname {S G} \left(\hat {\boldsymbol {z}} ^ {(2)}\right)}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \operatorname {S G} \left(\hat {\boldsymbol {z}} ^ {(2)}\right) \|} \\ = - \sum_ {m} ^ {M} \frac {\lambda_ {m} \hat {z} _ {m} ^ {(1)} \operatorname {S G} \left(\hat {z} _ {m} ^ {(2)}\right)}{\| D \hat {z} ^ {(1)} \| \| \operatorname {S G} \left(\hat {z} ^ {(2)}\right) \|}, \\ \end{array}
$$
where we used the fact that  $U$  is orthogonal and therefore does not change the Euclidean norm.  $\hat{z} = U^{\top}z$  is the representation rotated into the eigenbasis.
Lemma 2. (Learning dynamics in a rotated basis) Assuming that a given loss  $\mathcal{L}$  is optimized by gradient descent on the parameters of a neural network with network outputs  $z$ , a given orthogonal transformation  $\hat{z} = U^{\top}z$  and learning rate  $\eta$ , then the rotated representations  $\hat{z}$  evolve according to the dynamics:
$$
\frac {\mathrm {d} \hat {\pmb {z}}}{\mathrm {d} t} = - \eta \hat {\Theta} _ {t} (\pmb {x}, \mathcal {X}) \nabla_ {\hat {\mathcal {Z}}} \mathcal {L} ,
$$
where  $\hat{\Theta}_t(\mathcal{X},\mathcal{X}) = \nabla_\theta \hat{\mathcal{Z}}\nabla_\theta \hat{\mathcal{Z}}^\top$  is the empirical NTK expressed in the rotated basis.
Proof. Let  $\theta$  be the parameters of the neural network. Then we obtain the representational dynamics using the chain rule in the continuous-time gradient-flow setting [2]:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z}}{\mathrm {d} t} = \nabla_ {\theta} \hat {z} \frac {\mathrm {d} \theta}{\mathrm {d} t} \\ = \nabla_ {\boldsymbol {\theta}} \hat {z} (- \eta \nabla_ {\boldsymbol {\theta}} \mathcal {L}) \\ = \nabla_ {\boldsymbol {\theta}} \hat {z} \left(- \eta \nabla_ {\boldsymbol {\theta}} \hat {\mathcal {Z}} ^ {\top} \nabla_ {\hat {\mathcal {Z}}} \mathcal {L}\right) \\ = - \eta \hat {\Theta} _ {t} (x, \mathcal {X}) \nabla_ {\hat {Z}} \mathcal {L}. \\ \end{array}
$$
The above is a reiteration of the derivation of Eq. (2) given by Lee et al. [2], with an additional orthogonal transformation on the network outputs.
We proceed by proving the following Lemma which we will use in our proofs of Theorems 1 and 2.
Lemma 3. The NTK for a linear network is invariant under orthogonal transformations of the network output.
Proof. We first note that for a linear network, the parameters  $\theta$  are just the feedforward weights  $W$ . Therefore, for any orthogonal transformation  $U$  of the network output:
$$
\begin{array}{l} \hat {z} = U ^ {\top} f (\boldsymbol {x}) = U ^ {\top} W \boldsymbol {x} \\ \Rightarrow \nabla_ {\theta} \hat {\boldsymbol {z}} = \nabla_ {W} \hat {\boldsymbol {z}} = \nabla_ {W} \left(U ^ {\top} W \boldsymbol {x}\right) = \boldsymbol {x} ^ {\top} \otimes U ^ {\top}, \tag {18} \\ \end{array}
$$
where  $\otimes$  is the Kronecker product resulting from the fact that every input vector component appears in the update once for each output component.
We now study  $\hat{\Theta}_t(\mathcal{X},\mathcal{X})$ , the transformed empirical NTK (cf. Lemma 2). The  $(M\times M)$  diagonal blocks in the full  $(M|\mathbb{D}|\times M|\mathbb{D}|)$  empirical NTK  $\hat{\Theta}_t(\mathcal{X},\mathcal{X})$  correspond to single samples and the off-diagonal blocks are cross-terms between samples, where  $|\mathbb{D}|$  denotes the size of the training dataset and  $M$  the dimension of the outputs. We can develop a generic expression for each  $(M\times M)$  block  $\hat{\Theta}_t(\pmb {x}_i,\pmb {x}_j)$  corresponding to the interactions between samples  $i$  and  $j$  as:
$$
\begin{array}{l} \hat {\Theta} _ {t} \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}\right) = \nabla_ {W} \hat {\boldsymbol {z}} _ {i} \nabla_ {W} \hat {\boldsymbol {z}} _ {j} ^ {\top} \\ = \left(\boldsymbol {x} _ {i} ^ {\top} \otimes U ^ {\top}\right) \left(\boldsymbol {x} _ {j} ^ {\top} \otimes U ^ {\top}\right) ^ {\top} \\ = \left(\boldsymbol {x} _ {i} ^ {\top} \otimes U ^ {\top}\right) \left(\boldsymbol {x} _ {j} \otimes U\right) \\ = \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {j}\right) \otimes \left(U ^ {\top} U\right) \\ = \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {j}\right) \otimes I _ {\mathrm {M}} \\ = \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {j}\right) I _ {\mathrm {M}}. \tag {19} \\ \end{array}
$$
where we have used the fact that  $(A\otimes B)^{\top} = A^{\top}\otimes B^{\top}$  and  $(A\otimes B)(C\otimes D) = AC\otimes BD$ . Here,  $I_{\mathrm{M}}$  is the identity matrix of size  $M$ . Noting that Eq. (19) is unchanged when  $U$  is just the identity matrix completes the proof.
Proof. For a linear network with weights  $W \in \mathbb{R}^{M \times N}$ , we have from Lemma 3 that the empirical NTK  $\hat{\Theta}(\mathcal{X}, \mathcal{X})$  in the orthogonal eigenbasis is equal to the empirical NTK  $\Theta(\mathcal{X}, \mathcal{X})$  in the original basis. Furthermore from the proof for the lemma (see Eq. (19) above), each  $(M \times M)$  block of the full  $(M|\mathsf{D}| \times M|\mathsf{D}|)$  empirical NTK is given by:
$$
\hat {\Theta} _ {t} \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}\right) = \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {j}\right) I _ {\mathrm {M}}. \tag {20}
$$
where  $I_M \in \mathbb{R}^{M \times M}$  is the identity. Eq. (20) gives the total effective interaction between the samples  $i$  and  $j$  from the dataset. For high-dimensional inputs  $\pmb{x}$  drawn from an i.i.d standard Gaussian distribution, we have  $\pmb{x}_i^\top \pmb{x}_j \approx \delta_{ij}$  by the central limit theorem. Therefore, in the special case of a linear network with Gaussian i.i.d inputs, the representational dynamics (Lemma 2) simplify as follows:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {\boldsymbol {z}} _ {i} ^ {(1)}}{\mathrm {d} t} = - \eta \hat {\Theta} _ {t} (\boldsymbol {x} _ {i}, \mathcal {X}) \nabla_ {\hat {\mathcal {Z}}} \mathcal {L} \\ = - \eta \hat {\Theta} _ {t} \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {i}\right) \nabla_ {\hat {z} _ {i}} \mathcal {L} - \eta \sum_ {j \neq i} \hat {\Theta} _ {t} \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}\right) \nabla_ {\hat {z} _ {j}} \mathcal {L} \\ = - \eta \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {i}\right) \nabla_ {\hat {\boldsymbol {z}} _ {i}} \mathcal {L} - \eta \sum_ {j \neq i} \left(\boldsymbol {x} _ {i} ^ {\top} \boldsymbol {x} _ {j}\right) \nabla_ {\hat {\boldsymbol {z}} _ {j}} \mathcal {L} \\ = - \eta \nabla_ {\hat {z} _ {i}} \mathcal {L}. \tag {21} \\ \end{array}
$$
While the assumption of Gaussian i.i.d inputs is quite restrictive, we offer a generalizing interpretation here. Specifically, the above argument also holds when the inputs  $\pmb{x}$  are not all mutually orthogonal, but fall into  $P$  orthogonal clusters in the input dataset. Then, we would have  $\pmb{x}_i^\top \pmb{x}_j \approx \delta_{p_i = p_j}$  where  $p_i$  is the "label" of the cluster corresponding to sample  $i$ . If  $\mathcal{P}_i$  is the number of all the samples with the same label  $p_i$ , then Eq. (21) would simply be scaled to give  $\frac{\mathrm{d}\hat{z}_i^{(1)}}{\mathrm{d}t} = -\eta \mathcal{P}_i\nabla_{\hat{z}_i}\mathcal{L}$ .
For brevity, we proceed with the simplest case Eq. (21) in which every input is orthogonal. For  $\mathcal{L}^{\mathrm{euc}}$ , the representational gradient  $\nabla_{\hat{z}_i}\mathcal{L}$  is then given by:
$$
\nabla_ {\hat {z} _ {i}} \mathcal {L} ^ {\mathrm {e u c}} = \left(D _ {t} \hat {\boldsymbol {z}} _ {i} ^ {(1)} - \hat {\boldsymbol {z}} _ {i} ^ {(2)}\right) D _ {t}
$$
Noting that  $D_{t}$  is just a diagonal matrix containing the eigenvalues  $\lambda_{m}$  and dropping the sample subscript  $i$  for notational ease, we obtain for the  $m$ -th component of  $\nabla_{\hat{z}_i}\mathcal{L}^{\mathrm{euc}}$ :
$$
\frac {\partial \mathcal {L} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m}} = \lambda_ {m} (\lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)}) .
$$
Substituting this result in Eq. (21) gives us Eq. (8), the expression we were looking for. Finally, introducing  $\hat{z}_m \equiv \mathbb{E}[\hat{z}_m^{(1)}] = \mathbb{E}[\hat{z}_m^{(2)}]$  as the expectation over augmentations, we find that each eigenmode evolves independently in expectation value as:
$$
\begin{array}{l} \mathbb {E} \left[ \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} \right] = \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \lambda_ {m} \left(\mathbb {E} [ \hat {z} _ {m} ^ {(2)} ] - \lambda_ {m} \mathbb {E} [ \hat {z} _ {m} ^ {(1)} ]\right) \\ = \eta \lambda_ {m} (1 - \lambda_ {m}) \hat {z} _ {m}. \\ \end{array}
$$
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} \approx \eta \lambda_ {m} \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m} ^ {2}}{\| D \hat {z} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m}}{\| \hat {z} \|} \right] \cdot \sum_ {k \neq m} \lambda_ {k} (\lambda_ {k} - \lambda_ {m}), \tag {11}
$$
Proof. We can retrace the steps from the proof for Theorem 1 until Eq. (21):
$$
\frac {\mathrm {d} \hat {\boldsymbol {z}} _ {i} ^ {(1)}}{\mathrm {d} t} = - \eta \nabla_ {\hat {\boldsymbol {z}} _ {i}} \mathcal {L}.
$$
$\nabla_{\hat{z}_i}\mathcal{L}$  is a vector of dimension  $M$ . Ignoring the sample subscript  $i$  for simplicity, and focusing on the  $m$ -th component of  $\nabla_{\hat{z}_i}\mathcal{L}$ , we get:
$$
\begin{array}{l} \mathcal {L} ^ {\mathrm {c o s}} = - \sum_ {m} ^ {M} \frac {\lambda_ {m} \hat {z} _ {m} ^ {(1)} \mathrm {S G} (\hat {z} _ {m} ^ {(2)})}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \mathrm {S G} (\hat {\boldsymbol {z}} ^ {(2)}) \|} \\ \Rightarrow \frac {\partial \mathcal {L}}{\partial \hat {z} _ {m} ^ {(1)}} = - \frac {\lambda_ {m} \hat {z} _ {m} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \hat {\boldsymbol {z}} ^ {(2)} \|} + \frac {\sum_ {k} \lambda_ {k} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \cdot \lambda_ {m} ^ {2} \hat {z} _ {m} ^ {(1)} \\ = - \frac {\lambda_ {m}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \left[ \| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {2} \hat {z} _ {m} ^ {(2)} - \lambda_ {m} \hat {z} _ {m} ^ {(1)} \left(\sum_ {k} \lambda_ {k} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right) \right] \\ = - \frac {\lambda_ {m}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \left[ \left(\sum_ {k} \lambda_ {k} ^ {2} (\hat {\boldsymbol {z}} _ {k} ^ {(1)}) ^ {2}\right) \hat {\boldsymbol {z}} _ {m} ^ {(2)} - \lambda_ {m} \hat {\boldsymbol {z}} _ {m} ^ {(1)} \left(\sum_ {k} \lambda_ {k} \hat {\boldsymbol {z}} _ {k} ^ {(1)} \hat {\boldsymbol {z}} _ {k} ^ {(2)}\right) \right] \\ = - \frac {\lambda_ {m}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {k} (\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(2)} - \lambda_ {m} \hat {z} _ {m} ^ {(1)} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right) \\ \end{array}
$$
$$
\begin{array}{l} \Rightarrow \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = - \eta \frac {\partial \mathcal {L}}{\partial \hat {z} _ {m} ^ {(1)}} \\ = \frac {\eta \lambda_ {m}}{\| D \hat {\pmb {z}} ^ {(1)} \| ^ {3} \| \hat {\pmb {z}} ^ {(2)} \|} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {k} (\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(2)} - \lambda_ {m} \hat {z} _ {m} ^ {(1)} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right), \\ \end{array}
$$
proving Eq. (10). Assuming sufficiently small augmentations,  $\hat{z}_k^{(1)}$  and  $\hat{z}_k^{(2)}$  carry the same sign, and the net sign of both terms inside the parenthesis is fully determined by  $\gamma_m \equiv \mathrm{sign}(\hat{z}_m^{(1)})$ . Hence, we may write:
$$
\frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = \frac {\eta \lambda_ {m} \gamma_ {m}}{\| D \hat {\mathbf {z}} ^ {(1)} \| ^ {3} \| \hat {\mathbf {z}} ^ {(2)} \|} \sum_ {k \neq m} \left(\lambda_ {k} ^ {2} (\hat {z} _ {k} ^ {(1)}) ^ {2} | \hat {z} _ {m} ^ {(2)} | - \lambda_ {m} \lambda_ {k} | \hat {z} _ {m} ^ {(1)} | | \hat {z} _ {k} ^ {(1)} | | \hat {z} _ {k} ^ {(2)} |\right) .
$$
It is useful to separate out  $\gamma_{m}$  in this manner because every other term in the expression is now non-negative. Then  $\mathrm{sign}(\gamma_m\cdot \frac{\mathrm{d}\hat{z}}{\mathrm{d}t}) = \mathrm{sign}(\hat{z}_m\cdot \frac{\mathrm{d}\hat{z}_m}{\mathrm{d}t})$  tells us whether  $\hat{z}_{m}$  tends to increase or decrease in magnitude, as we have argued in the main text.
Asymptotic analysis. To get a handle on how the different eigenvalues influence each other, we consider two important limiting cases. First, we consider the asymptotic regime dominated by one eigenvalue, and show that it tends towards a more symmetric solution in which the gap between different eigenvalues decreases. Second, we derive asymptotic expressions for the near-uniform regime in which all eigenvalues are comparable in size and show that this solution tends toward the uniform solution (cf. Eq. (11)).
To facilitate our analysis, we define each mode's relative contribution  $\chi_{m}\equiv \frac{|\hat{z}_{m}|}{\|\hat{z}\|}$  and evaluate Eq. (10) taking the expectation value over augmentations:
$$
\begin{array}{l} \mathbb {E} \left[ \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} \right] = \eta \lambda_ {m} \sum_ {k \neq m} \left(\lambda_ {k} ^ {2} \cdot \mathbb {E} \left[ \frac {(\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(2)}}{\| D \hat {z} ^ {(1)} \| ^ {3} \| \hat {z} ^ {(2)} \|} \right] - \lambda_ {m} \lambda_ {k} \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m} ^ {(1)} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}}{\| D \hat {z} ^ {(1)} \| ^ {3} \| \hat {z} ^ {(2)} \|} \right]\right) \\ \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \lambda_ {m} \sum_ {k \neq m} \left(\lambda_ {k} ^ {2} \cdot \mathbb {E} \left[ \frac {\hat {z} _ {k} ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m}}{\| \hat {\boldsymbol {z}} \|} \right] - \lambda_ {m} \lambda_ {k} \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m} \hat {z} _ {k}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {\hat {z} _ {k}}{\| \hat {\boldsymbol {z}} \|} \right]\right) \\ = \eta \lambda_ {m} \gamma_ {m} \sum_ {k \neq m} \left(\lambda_ {k} ^ {2} \cdot \mathbb {E} \left[ \chi_ {k} ^ {2} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \chi_ {m} ] - \lambda_ {m} \lambda_ {k} \cdot \mathbb {E} \left[ \chi_ {m} \chi_ {k} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \chi_ {k} ]\right). \tag {22} \\ \end{array}
$$
In the second equality we used the fact that the expectation value taken over augmentations is conditioned on the input sample, which makes them conditionally independent.
One dominant eigenvalue. First, we consider the low-rank regime in which one eigenvalue dominates. Without loss of generality, we assume  $\lambda_1 \gg \lambda_k \forall k \neq 1$ . We then have:
$$
\chi_ {1} \sim 1
$$
$$
\chi_ {k} \sim \epsilon \quad (0 <   \epsilon \ll 1) \quad \forall \quad k \neq 1
$$
Plugging these values into Eq. (22) gives the following dynamics for the dominant eigenmode:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z} _ {1}}{\mathrm {d} t} \approx \eta \lambda_ {1} \gamma_ {1} \sum_ {k \neq 1} \left(\lambda_ {k} ^ {2} \cdot \mathbb {E} \left[ \epsilon^ {2} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ 1 ] - \lambda_ {1} \lambda_ {k} \cdot \mathbb {E} \left[ \epsilon \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \mathbb {E} [ \epsilon ]\right) \\ = \eta \lambda_ {1} \gamma_ {1} \sum_ {k \neq 1} \left(\lambda_ {k} ^ {2} \epsilon^ {2} \cdot \mathbb {E} \left[ \frac {\| \hat {z} \| ^ {2}}{\| D \hat {z} \| ^ {3}} \right] \cdot \mathbb {E} [ 1 ] - \lambda_ {1} \lambda_ {k} \epsilon^ {2} \cdot \mathbb {E} \left[ \frac {\| \hat {z} \| ^ {2}}{\| D \hat {z} \| ^ {3}} \right]\right) \\ = \eta \lambda_ {1} \gamma_ {1} \epsilon^ {2} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \sum_ {k \neq 1} \lambda_ {k} (\lambda_ {k} - \lambda_ {1}) \quad . \\ \end{array}
$$
These updates are always opposite in sign to the representation component, which corresponds to decaying dynamics for the leading eigenmode because  $\gamma_{1}\frac{\mathrm{d}\hat{z}_{1}}{\mathrm{d}t} < 0$ .
For all other modes we have:
$$
\begin{array}{l} \frac{\mathrm{d}\hat{z}_{m\neq 1}}{\mathrm{d}t}\approx \eta \lambda_{m}\gamma_{m}\sum_{k\notin \{m,1\}}\left(\lambda_{k}^{2}\epsilon^{2}\cdot \mathbb{E}\left[\frac{\| \hat{\boldsymbol{z}}\|^{2}}{\|D\hat{\boldsymbol{z}}\|^{3}}\right]\cdot \mathbb{E}\left[\epsilon \right] - \lambda_{m}\lambda_{k}\epsilon^{2}\cdot \mathbb{E}\left[\frac{\|\hat{\boldsymbol{z}}\|^{2}}{\|D\hat{\boldsymbol{z}}\|^{3}}\right]\cdot \mathbb{E}\left[\epsilon \right]\right) \\ + \eta \lambda_ {m} \gamma_ {m} \left(\lambda_ {1} ^ {2} \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \epsilon ] - \lambda_ {m} \lambda_ {1} \epsilon \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ 1 ]\right) \\ = \eta \lambda_ {m} \gamma_ {m} \epsilon \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \left(\lambda_ {1} (\lambda_ {1} - \lambda_ {m}) + \epsilon^ {2} \sum_ {k \notin \{m, 1 \}} \lambda_ {k} (\lambda_ {k} - \lambda_ {m})\right) \\ \approx \eta \lambda_ {m} \gamma_ {m} \lambda_ {1} \epsilon \cdot \mathbb {E} \left[ \frac {\| \hat {\pmb {z}} \| ^ {2}}{\| D \hat {\pmb {z}} \| ^ {3}} \right] (\lambda_ {1} - \lambda_ {m}) , \\ \end{array}
$$
so that  $\gamma_{m}\frac{\mathrm{d}\hat{z}_{m}}{\mathrm{d}t} >0$ , i.e, the updates have the same sign as the representation component, which corresponds to growth dynamics. In other words: The dominant eigenvalue "pulls all the other eigenvalues up," a form of implicit cooperation between the eigenmodes. We also note that the non-dominant eigenmodes increase at a rate proportional to  $\epsilon$ , whereas the dominant eigenmode decreases at a slower rate proportional to  $\epsilon^2$ . Thus, for sensible initializations with at least one large and many small eigenvalues, the modes will tend toward an equilibrium at some non-zero intermediate value, without a dominant mode. Next we study this other limiting case in which all eigenvalues are of similar size.
Near-uniform regime. To study the dynamics in a near-uniform regime, we note that all  $\chi_{m}$  are of order  $\mathcal{O}(1)$  in  $\hat{z}_{m}$ , whereas the eigenvalues  $\lambda_{m}$  are of order  $\mathcal{O}(\hat{z}_{m}^{2})$ . In this setting, the effect of the eigenvalue terms  $\lambda_{m}$  on the dynamics is stronger than the  $\chi_{m}$  terms which are bounded between 0 and 1. With a sufficiently high-dimensional representation, all  $\chi_{m}$  terms will be centered around  $1 / \sqrt{M}$ . Based on these observations, we may make the simplifying assumption that the contributions are all approximately equal, i.e.,  $\chi_{i} = \chi$  for all  $i$ . Substituting this value in Eq (22) gives:
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \lambda_ {m} \gamma_ {m} \cdot \mathbb {E} \left[ \chi^ {2} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \chi ] \cdot \sum_ {k \neq m} \lambda_ {k} (\lambda_ {k} - \lambda_ {m}). \tag {23}
$$
Finally, substituting for  $\chi$ , which by assumption are all approximately equal:
$$
\chi \approx \chi_ {m} = \frac {\left| \hat {z} _ {m} \right|}{\left\| \hat {\boldsymbol {z}} \right\|},
$$
and absorbing back the sign from  $\gamma_{m}$ , we obtain the approximate dynamics in Eq (11):
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} \approx \eta \lambda_ {m} \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m} ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {\hat {z} _ {m}}{\| \hat {\boldsymbol {z}} \|} \right] \cdot \sum_ {k \neq m} \lambda_ {k} (\lambda_ {k} - \lambda_ {m})
$$
C Derivation of idealized learning dynamics for different loss variations
C.1 Removing the stop-grad from the Euclidean loss  $\mathcal{L}^{\mathrm{euc}}$
Omitting the stop-grad operator from  $\mathcal{L}^{\mathrm{euc}}$  gives:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {e u c}} = \frac {1}{2} \| W _ {\mathrm {P}} z ^ {(1)} - z ^ {(2)} \| ^ {2} \\ = \frac {1}{2} \sum_ {m} ^ {M} | \lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)} | ^ {2}. \\ \end{array}
$$
Tracing the steps to prove Theorem 1 and assuming Gaussian i.i.d inputs for a linear network, we write:
$$
\begin{array}{l} \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m}} = \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(1)}} + \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(2)}} \\ = \left(\lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)}\right) \lambda_ {m} - \left(\lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)}\right) \\ = \left(\lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)}\right) (\lambda_ {m} - 1) \\ \Rightarrow \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = - \eta \mathbb {E} \left[ \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m}} \right] \\ = - \eta \left(\lambda_ {m} \mathbb {E} [ \hat {z} _ {m} ^ {(1)} ] - \mathbb {E} [ \hat {z} _ {m} ^ {(2)} ]\right) (\lambda_ {m} - 1) \\ = - \eta (1 - \lambda_ {m}) ^ {2} \hat {z} _ {m}, \\ \end{array}
$$
which results in decaying representations and thus collapse.
C.2 Removing the stop-grad from the cosine loss  $\mathcal{L}^{\mathrm{cos}}$
Following the same arguments as above, omitting the stop-grad operator from  $\mathcal{L}^{\mathrm{cos}}$  gives:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {c o s}} = - \frac {\left(W _ {\mathrm {P}} z ^ {(1)}\right) ^ {\top} z ^ {(2)}}{\| W _ {\mathrm {P}} z ^ {(1)} \| \| z ^ {(2)} \|} \\ \Rightarrow \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\mathrm {c o s}}}{\partial \hat {z} _ {m}} = \frac {- \lambda_ {m}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \sum_ {k \neq m} \left(\lambda_ {k} ^ {2} (\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(2)} - \lambda_ {m} \lambda_ {k} \hat {z} _ {m} ^ {(1)} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)} + \lambda_ {k} ^ {2} \lambda_ {m} (\hat {z} _ {k} ^ {(1)}) ^ {3} - \lambda_ {k} \hat {z} _ {m} ^ {(2)} \hat {z} _ {k} ^ {(2)} \hat {z} _ {k} ^ {(1)}\right) \\ + \frac {- \lambda_ {m}}{\| D \hat {\pmb {z}} ^ {(1)} \| ^ {3} \| \hat {\pmb {z}} ^ {(2)} \|} \left(\lambda_ {m} ^ {3} (\hat {z} _ {m} ^ {(1)}) ^ {3} - \lambda_ {m} (\hat {z} _ {m} ^ {(2)}) ^ {2} \hat {z} _ {m} ^ {(1)}\right) , \\ \end{array}
$$
so that, when taking the expectation value over augmentations, the dynamics follow: