SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = - \eta \mathbb {E} \left[ \frac {\partial \mathcal {L} _ {\mathrm {n o S G}} ^ {\cos}}{\partial \hat {z} _ {m}} \right] \\ = \eta \lambda_ {m} \gamma_ {m} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {k} \cdot \mathbb {E} \left[ \chi_ {k} ^ {2} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi_ {m} \right] - \lambda_ {m} \cdot \mathbb {E} \left[ \chi_ {m} \chi_ {k} \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi_ {k} \right]\right) \\ + \eta \lambda_ {m} \gamma_ {m} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {m} \lambda_ {k} \mathbb {E} \left[ \chi_ {k} ^ {3} \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \mathbb {E} \left[ \chi_ {k} \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi_ {m} \chi_ {k} \| \hat {\boldsymbol {z}} \| \right]\right) \\ + \eta \lambda_ {m} ^ {2} \gamma_ {m} \left(\lambda_ {m} ^ {2} \mathbb {E} \left[ \chi_ {m} ^ {3} \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \mathbb {E} \left[ \chi_ {m} \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi_ {m} ^ {2} \| \hat {\boldsymbol {z}} \| \right]\right). \\ \end{array}
$$
In the asymptotic regime with dominant eigenvalue  $\lambda_{1}$ , we get the dynamics:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z} _ {1}}{\mathrm {d} t} = \eta \lambda_ {1} \gamma_ {1} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {k} \epsilon^ {2} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] - \lambda_ {m} \epsilon^ {2} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right]\right) \\ + \eta \lambda_ {1} \gamma_ {1} \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {1} \lambda_ {m} \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \epsilon^ {2} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \| \hat {\boldsymbol {z}} \| \right]\right) \\ + \eta \lambda_ {1} ^ {2} \gamma_ {1} \left(\lambda_ {1} ^ {2} \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \| \hat {\boldsymbol {z}} \| ]\right) \\ \approx \eta \lambda_ {1} ^ {4} \gamma_ {1} \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] \\ \end{array}
$$
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z} _ {m \neq 1}}{\mathrm {d} t} = \eta \lambda_ {m} \gamma_ {m} \sum_ {k \notin \{m, 1 \}} \lambda_ {k} \left(\lambda_ {k} \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] - \lambda_ {m} \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {2}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right]\right) \\ + \eta \lambda_ {m} \gamma_ {m} \lambda_ {1} \left(\lambda_ {1} \epsilon \mathbb {E} \left[ \frac {\| \hat {z} \| ^ {2}}{\| D \hat {z} \| ^ {3}} \right] - \lambda_ {m} \epsilon \mathbb {E} \left[ \frac {\| \hat {z} \| ^ {2}}{\| D \hat {z} \| ^ {3}} \right]\right) \\ + \eta \lambda_ {m} \gamma_ {m} \sum_ {k \notin \{m, 1 \}} \lambda_ {k} \left(\lambda_ {m} \lambda_ {k} \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \| \hat {\boldsymbol {z}} \| ]\right) \\ + \eta \lambda_ {m} \gamma_ {m} \lambda_ {1} \left(\lambda_ {m} \lambda_ {1} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \epsilon \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \| \hat {\boldsymbol {z}} \| \right]\right) \\ + \eta \lambda_ {m} ^ {2} \gamma_ {m} \left(\lambda_ {m} ^ {2} \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] - \epsilon^ {3} \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} [ \| \hat {\boldsymbol {z}} \| ]\right) \\ \approx \eta \lambda_ {m} ^ {2} \gamma_ {m} \lambda_ {1} ^ {2} \cdot \mathbb {E} \left[ \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right], \\ \end{array}
$$
Thus, all eigenmodes diverge because  $\gamma_{m}\frac{\mathrm{d}\hat{z}_{m}}{\mathrm{d}t} >0$
Similarly, we find divergent dynamics when starting in the near-uniform regime:
$$
\begin{array}{l} \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \lambda_ {m} \gamma_ {m} \mathbb {E} \left[ \chi^ {2} \frac {\| \hat {\pmb {z}} \| ^ {2}}{\| D \hat {\pmb {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi \right] \sum_ {k \neq m} \lambda_ {k} \left(\lambda_ {k} - \lambda_ {m}\right) \\ + \eta \lambda_ {m} ^ {2} \gamma_ {m} \mathbb {E} \left[ \chi^ {3} \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] \sum_ {k} \lambda_ {k} ^ {2} \\ - \eta \lambda_ {m} \gamma_ {m} \mathbb {E} \left[ \chi \frac {\| \hat {\boldsymbol {z}} \|}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \chi^ {2} \| \hat {\boldsymbol {z}} \| \right] \sum_ {k} \lambda_ {k} \\ \approx \eta \lambda_ {m} ^ {2} \gamma_ {m} \mathbb {E} \left[ \chi^ {3} \frac {\| \hat {\boldsymbol {z}} \| ^ {3}}{\| D \hat {\boldsymbol {z}} \| ^ {3}} \right] \cdot \mathbb {E} \left[ \frac {1}{\| \hat {\boldsymbol {z}} \|} \right] \sum_ {k} \lambda_ {k} ^ {2}, \\ \end{array}
$$
selecting the terms with the highest power in the eigenvalues.
Thus, omission of stop-grad precludes successful representation learning for both the Euclidean and the cosine loss, but due to different mechanisms. Euclidean loss yields collapse, whereas the cosine loss succumbs to run-away activity.
C.3 Removing the predictor from the Euclidean loss  $\mathcal{L}^{\mathrm{euc}}$
To analyze the representational dynamics in the absence of the predictor network, we consider  $\mathcal{L}_{\mathrm{noPred}}^{\mathrm{euc}}$ :
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {n o P r e d}} ^ {\mathrm {e u c}} = \frac {1}{2} \| \boldsymbol {z} ^ {(1)} - \operatorname {S G} (\boldsymbol {z} ^ {(2)}) \| ^ {2} \\ = \frac {1}{2} \sum_ {m} ^ {M} | \hat {z} _ {m} ^ {(1)} - \mathrm {S G} (\hat {z} _ {m} ^ {(2)}) | ^ {2}. \\ \end{array}
$$
The dynamics resulting from this loss function are a special case of the dynamics derived in Theorem 1 with all the eigenvalues equal to one  $(\lambda_{k} = 1)$ . In particular Eq. (8) becomes:
$$
\frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = - \eta \frac {\partial \mathcal {L} _ {\mathrm {n o P r e d}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(1)}} (t) = \eta (\hat {z} _ {m} ^ {(2)} - \hat {z} _ {m} ^ {(1)}),
$$
which evaluates to 0 under expectation over augmentations. Hence there is no learning without the predictor.
C.4 Removing the predictor from the cosine loss  $\mathcal{L}^{\mathrm{cos}}$
Similarly, when we remove the predictor from  $\mathcal{L}^{\mathrm{cos}}$  it yields:
$$
\mathcal {L} _ {\mathrm {n o P r e d}} ^ {\mathrm {c o s}} = - \sum_ {m} ^ {M} \frac {\hat {z} _ {m} ^ {(1)} \mathrm {S G} (\hat {z} _ {m} ^ {(2)})}{\| \hat {\boldsymbol {z}} ^ {(1)} \| \| \mathrm {S G} (\hat {\boldsymbol {z}} ^ {(2)}) \|},
$$
so that Eq. (10) becomes:
$$
\frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = \frac {\eta}{\| \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \sum_ {k \neq m} \left(\left(\hat {z} _ {k} ^ {(1)}\right) ^ {2} \hat {z} _ {m} ^ {(2)} - \hat {z} _ {m} ^ {(1)} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right). \tag {24}
$$
Here, the near-uniform approximation (Eq. (11)) of ignoring the differences in  $\chi$  between different eigenmodes is not valid. This is because the  $\lambda$ -terms are no longer present, and the effects of the  $\chi$ -terms on the dynamics cannot be treated as negligible. In particular, setting  $W_{\mathrm{P}} = I$  in the dynamics derived in Theorem 2 would yield zero dynamics. However taking the expectation of Eq. (24) over augmentations yields the non-zero dynamics:
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \sum_ {k \neq m} \left(\mathbb {E} \left[ \frac {\hat {z} _ {k} ^ {2}}{\| \hat {\boldsymbol {z}} \| ^ {3}} \right] \mathbb {E} \left[ \frac {\hat {z} _ {m}}{\| \hat {\boldsymbol {z}} \|} \right] - \mathbb {E} \left[ \frac {\hat {z} _ {m} \hat {z} _ {k}}{\| \hat {\boldsymbol {z}} \| ^ {3}} \right] \mathbb {E} \left[ \frac {\hat {z} _ {k}}{\| \hat {\boldsymbol {z}} \|} \right]\right), \tag {25}
$$
which consists of terms of order  $\mathcal{O}\left(\frac{1}{\hat{z}_m}\right)$  and  $\mathcal{O}(1)$  in  $\hat{z}_m$ , hinting at slower dynamics compared to Eq. (10). To study these granular effects, we would need to explicitly model the effect of the augmentations, for which we do not have a good statistical model. In lieu of deriving these dynamics analytically, we make an observation which restricts the possible dynamical behavior. Specifically, the sum of the eigenvalues remains constant throughout training. To show this, we begin by writing out the expression for the derivative over time of the sum of all the eigenvalues:
$$
\begin{array}{l} \frac {\mathrm {d}}{\mathrm {d} t} \sum_ {m} \lambda_ {m} = \sum_ {m} \frac {\mathrm {d} \lambda_ {m}}{\mathrm {d} t} = \sum_ {m} \frac {\mathrm {d}}{\mathrm {d} t} \mathbb {E} _ {\mathrm {d a t a}} \left[ \hat {z} _ {m} ^ {2} \right] \\ = \mathbb {E} _ {\mathrm {d a t a}} \left[ \sum_ {m} \frac {\mathrm {d} \hat {z} _ {m} ^ {2}}{\mathrm {d} t} \right] \\ = \mathbb {E} _ {\mathrm {d a t a}} \left[ \sum_ {m} \hat {z} _ {m} \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} \right]. \\ \end{array}
$$
We can derive the term inside the expectation by adding up the dynamics given by Eq. (25) for all the different eigenmodes:
$$
\begin{array}{l} \hat {z} _ {m} ^ {(1)} \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = \frac {\eta}{\| \hat {z} ^ {(1)} \| ^ {3} \| \hat {z} ^ {(2)} \|} \sum_ {k \neq m} \left((\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(1)} \hat {z} _ {m} ^ {(2)} - (\hat {z} _ {m} ^ {(1)}) ^ {2} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right) \\ \Rightarrow \sum_ {m} \hat {z} _ {m} ^ {(1)} \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} = \frac {\eta}{\| \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \sum_ {m} \sum_ {k \neq m} \left((\hat {z} _ {k} ^ {(1)}) ^ {2} \hat {z} _ {m} ^ {(1)} \hat {z} _ {m} ^ {(2)} - (\hat {z} _ {m} ^ {(1)}) ^ {2} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}\right) \\ = 0 \\ \Longrightarrow \mathbb {E} _ {\mathrm {d a t a}} \left[ \mathbb {E} _ {\mathrm {a u g}} \left[ \sum_ {m} \hat {z} _ {m} ^ {(1)} \frac {\mathrm {d} \hat {z} _ {m} ^ {(1)}}{\mathrm {d} t} \right] \right] = \mathbb {E} _ {\mathrm {d a t a}} \left[ \sum_ {m} \hat {z} _ {m} \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} \right] = 0, \\ \end{array}
$$
proving that  $\frac{\mathrm{d}}{\mathrm{d}t}\sum_{m}\lambda_{m} = 0$ , i.e., the sum of the eigenvalues is conserved. This precludes collapsing dynamics where all eigenvalues go to zero as well as diverging dynamics where at least one eigenvalue diverges.
C.5 Isotropic losses for equalized convergence rates
In Expressions (9) and (11) we see that the overall learning dynamics have a quadratic dependence on the eigenvalues with a root near collapsed solutions, which causes these modes to learn slower. We reasoned that this anisotropy could be detrimental for learning. To address this issue, we sought to derive alternative loss functions that encourage isotropic learning dynamics for all modes.
C.5.1 Euclidean IsoLoss.
We start by deriving an IsoLoss function for the Euclidean case  $\mathcal{L}^{\mathrm{euc}}$ . To avoid the unwanted quadratic dependence, we first note that we would like to arrive at the following expression for the dynamics:
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta (1 - \lambda_ {m}) \hat {z} _ {m}.
$$
By recalling the Euclidean loss and corresponding dynamics:
$$
\mathcal {L} ^ {\mathrm {e u c}} = \frac {1}{2} \sum_ {m} ^ {M} | \lambda_ {m} \hat {z} _ {m} ^ {(1)} - \mathrm {S G} (\hat {z} _ {m} ^ {(2)}) | ^ {2} \Rightarrow \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \lambda_ {m} (1 - \lambda_ {m}) \hat {z} _ {m},
$$
we note that the leading  $\lambda_{m}$  term has no influence on the overall sign of the dynamics, and is introduced by the second step in the chain rule:
$$
\frac {\partial \mathcal {L} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(1)}} = (\lambda_ {m} \hat {z} ^ {(1)} - \hat {z} ^ {(2)}) \cdot \frac {\partial}{\partial \hat {z} _ {m} ^ {(1)}} (\lambda_ {m} \hat {z} ^ {(1)} - \hat {z} ^ {(2)}) .
$$
Based on this realization we see that this second step needs to be modified. To that end, we start with the desired derivative:
$$
\frac {\partial \mathcal {L} _ {\mathrm {i s o}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(1)}} = (\lambda_ {m} \hat {z} ^ {(1)} - \hat {z} ^ {(2)}) \cdot \frac {\partial}{\partial \hat {z} _ {m} ^ {(1)}} (\hat {z} ^ {(1)} - \hat {z} ^ {(2)}) ,
$$
and see that several loss functions are possible. The one we have reported in Eq. (15) we derived by applying an appropriate stop-grad while integrating:
$$
\frac {\partial \mathcal {L} _ {\mathrm {i s o}} ^ {\mathrm {e u c}}}{\partial \hat {z} _ {m} ^ {(1)}} = (\hat {z} _ {m} ^ {(1)} + \lambda_ {m} \hat {z} ^ {(1)} - \hat {z} ^ {(2)} - \hat {z} _ {m} ^ {(1)}) \cdot \frac {\partial}{\partial \hat {z} _ {m} ^ {(1)}} (\hat {z} ^ {(1)} - \hat {z} ^ {(2)}) .
$$
to give:
$$
\mathcal {L} _ {\mathrm {i s o}} ^ {\mathrm {e u c}} = \frac {1}{2} \sum_ {m} ^ {M} | \hat {z} _ {m} ^ {(1)} - \mathrm {S G} (\hat {z} _ {m} ^ {(2)} + \hat {z} _ {m} ^ {(1)} - \lambda_ {m} \hat {z} _ {m} ^ {(1)}) | ^ {2}
$$
Another alternative loss with the same desired isotropic learning dynamics, but using a different placement of the stop-gradient operators, is given by:
$$
\mathcal {L} _ {\mathrm {i s o}} ^ {\mathrm {e u c}} = \sum_ {m} ^ {M} \operatorname {S G} \left(\lambda_ {m} \hat {z} _ {m} ^ {(1)} - \hat {z} _ {m} ^ {(2)}\right) \cdot \left(\hat {z} _ {m} ^ {(1)} - \operatorname {S G} (\hat {z} _ {m} ^ {(2)})\right)
$$
C.5.2 Cosine Similarity IsoLoss.
Since most practical SSL approaches rely on cosine similarity, which suffers from a similar anisotropy of the learning dynamics, we sought to find IsoLosses in this setting. With the same goal as above, we would like to arrive at the dynamics:
$$
\frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \frac {\hat {z} _ {m} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \hat {\boldsymbol {z}} ^ {(2)} \|} - \eta \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} \hat {z} _ {m} ^ {(1)}
$$
starting from the cosine loss and corresponding dynamics:
$$
\begin{array}{l} \mathcal {L} ^ {\cos} = - \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) \|} (26) \\ \Rightarrow \frac {\mathrm {d} \hat {z} _ {m}}{\mathrm {d} t} = \eta \frac {\lambda_ {m} \hat {z} _ {m} ^ {(2)}}{\| D \hat {z} ^ {(1)} \| \| \hat {z} ^ {(2)} \|} - \eta \frac {\sum_ {k} \lambda_ {k} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}}{\| D \hat {z} ^ {(1)} \| ^ {3} \| \hat {z} ^ {(2)} \|} \cdot \lambda_ {m} ^ {2} \hat {z} _ {m} ^ {(1)}. (27) \\ \end{array}
$$
The IsoLoss in this case can be derived by noting how  $\lambda_{m}$  arises in each of the two terms in Eq. (27), and engineering an alternative loss function corresponding to each term separately.
In the first term,  $\lambda_{m}$  arises from the partial derivative of the numerator  $\lambda_{m}\hat{z}_{m}^{(1)}\mathrm{SG}(\hat{z}_{m}^{(2)})$  in the original loss (Eq. (26)). This can be remediated by using  $\hat{z}_m^{(1)}\mathrm{SG}(\hat{z}_m^{(2)})$  as the numerator instead.
In the second term in Eq. (27),  $\lambda_{m}^{2}$  arises from the partial derivative of  $\| D\hat{z}^{(1)}\| = \sqrt{\sum_{k}(\lambda_{k}\hat{z}_{m}^{(1)})^{2}}$  in the denominator. We can reduce  $\lambda_{m}^{2}$  to  $\lambda_{m}$  by instead taking the partial derivative of  $\| D^{1 / 2}\hat{z}^{(1)}\| = \sqrt{\sum_{k}(\lambda_{k}^{1 / 2}\hat{z}_{m}^{(1)})^{2}}$ .
Putting these insights together, we arrive at the desired partial derivative:
$$
\begin{array}{l} \frac {\partial \mathcal {L} _ {\mathrm {i s o}} ^ {\mathrm {c o s}}}{\partial \hat {z} _ {m} ^ {(1)}} = \frac {- 1}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \hat {\boldsymbol {z}} ^ {(2)} \|} \cdot \frac {\partial \hat {z} _ {m} ^ {(1)} \hat {z} _ {m} ^ {(2)}}{\partial \hat {z} _ {m} ^ {(1)}} + \frac {\sum_ {k} \lambda_ {k} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \cdot \frac {1}{2} \frac {\partial \lambda_ {m} (\hat {z} _ {m} ^ {(1)}) ^ {2}}{\partial \hat {z} _ {m} ^ {(1)}} \\ = \frac {- 1}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \hat {\boldsymbol {z}} ^ {(2)} \|} \cdot \frac {\partial (\hat {\boldsymbol {z}} ^ {(1)}) ^ {\top} \hat {\boldsymbol {z}} ^ {(2)}}{\partial \hat {z} _ {m} ^ {(1)}} + \frac {\sum_ {k} \lambda_ {k} \hat {z} _ {k} ^ {(1)} \hat {z} _ {k} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|} \cdot \frac {1}{2} \frac {\partial \| D ^ {1 / 2} \hat {\boldsymbol {z}} ^ {(1)} \| ^ {2}}{\partial \hat {z} _ {m} ^ {(1)}} , \\ \end{array}
$$
and the integrated IsoLoss in eigenspace:
$$
\mathcal {L} _ {\mathrm {i s o}} ^ {\cos} = - (\hat {\boldsymbol {z}} ^ {(1)}) ^ {\top} \operatorname {S G} \left(\frac {\hat {\boldsymbol {z}} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| \| \hat {\boldsymbol {z}} ^ {(2)} \|}\right) + \frac {1}{2} \operatorname {S G} \left(\frac {(D \hat {\boldsymbol {z}} ^ {(1)}) ^ {\top} \hat {\boldsymbol {z}} ^ {(2)}}{\| D \hat {\boldsymbol {z}} ^ {(1)} \| ^ {3} \| \hat {\boldsymbol {z}} ^ {(2)} \|}\right) \| D ^ {1 / 2} \hat {\boldsymbol {z}} ^ {(1)} \| ^ {2}.
$$
Rotating all terms back to the original space gives the desired IsoLoss for cosine similarity as reported (Eq. (17)):
$$
\mathcal {L} _ {\mathrm {i s o}} = - \left(\boldsymbol {z} ^ {(1)}\right) ^ {\top} \operatorname {S G} \left(\frac {\boldsymbol {z} ^ {(2)}}{\| W _ {\mathrm {P}} \boldsymbol {z} ^ {(1)} \| \| \boldsymbol {z} ^ {(2)} \|}\right) + \frac {1}{2} \operatorname {S G} \left(\frac {\left(W _ {\mathrm {P}} \boldsymbol {z} ^ {(1)}\right) ^ {\top} \boldsymbol {z} ^ {(2)}}{\| W _ {\mathrm {P}} \boldsymbol {z} ^ {(1)} \| ^ {3} \| \boldsymbol {z} ^ {(2)} \|}\right) \| W _ {\mathrm {P}} ^ {1 / 2} \boldsymbol {z} ^ {(1)} \| ^ {2}.
$$
D Experimental methods
Self-supervised pretraining. We used the CIFAR-10, CIFAR-100 [3], STL-10 [4], and TinyImageNet [5] datasets for self-supervised learning with a ResNet-18 [6] encoder and the SimCLR set of transformations [7]. We also adopted several modifications of ResNet-18 which have been proposed to deal with the low resolution of the images in these datasets [7]. The ResNet modifications comprise using  $3 \times 3$  convolutional kernels instead of  $7 \times 7$  kernels and skipping the first max-pooling operation. The Solo-learn library [8] also provides specialized sets of augmentations that work well for these datasets, which we adopted as well. The configurations we used for each dataset are summarized in Table 4. We used BatchNorm in the backbone and the projector multi-layer perceptron (MLP) in the hidden layer for all methods. For BYOL, we included BatchNorm also in the hidden layer of the predictor MLP.
We used a projection dimension of 256 for the projection MLP using one hidden layer with 4096 units, and the same architecture for the nonlinear predictor for the BYOL baseline. For networks using EMA target networks, we used the LARS optimizer with learning rate 1.0 whereas for networks without the EMA, we used stochastic gradient descent with momentum 0.9 and learning rate 0.1. Furthermore, we used a warmup period of 10 epochs for the learning rate followed by a cosine decay schedule and a batch size of 256. We also used a weight decay  $4 \times 10^{-4}$  for the closed-form predictor models and  $10^{-5}$  for the nonlinear predictor models. For evaluation, we removed the projection MLP and used the embeddings at the pooled output of the ResNet convolutional layers following standard practice. For the EMA, we started with  $\tau_{\mathrm{base}} = 0.99$  and increased  $\tau_{\mathrm{EMA}}$  to 1 with a cosine schedule exactly following the configuration reported in [9]. For DirectPred, we used  $\alpha = 0.5$  and  $\tau = 0.998$  for the moving average estimate of the correlation matrix updated at every step.
Table 4:
CIFAR-10	CIFAR-100	STL-10	TinyImageNet
Resolution	32 × 32	32 × 32	96 × 96	64 × 64
Kernel size	3 × 3	3 × 3	7 × 7	3 × 3
First max-pool	No	No	Yes	Yes
Blur	No	No	Yes	No
Linear evaluation protocol. We reported the held-out classification accuracy on the test sets for CIFAR-10/100 and STL-10, and the validation set for TinyImageNet, after online training of the gradient-isolated linear classifier on each labeled example in the training set during pretraining.
Compute resources All simulations were run on an in-house cluster consisting of 5 nodes with 4 V100 NVIDIA GPUs each, one node with 4 A100 NVIDIA GPUs, and one node with 8 A40 NVIDIA GPUs. Runs on CIFAR-10/100 took about 8 hours each, and the STL-10 and TinyImagenet runs took about 24 hours each.