SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We used a standard linear evaluation technique, freezing the parameters of the encoding network and training a linear classifier with supervision (14), to verify that our method extracts semantically relevant features from the data. We also perform semi-supervised evaluation, where all model parameters are fine tuned using a small number of labelled examples, and also check whether the learned features generalize to three out-of-distribution datasets: Flowers-102, Food-101, and the Describable Textures Dataset (57; 11; 19). The results are summarized in Table 1 and details of the training of the linear classifier are provided in Appendix G. Finally we also evaluate our best model in and each of the baselines on object detection on the VOC07 dataset. We followed (41; 77), fine tuning the representation networks for detection with a Faster R-CNN head and C-4 backbone using the 1x schedule. MMCR achieved a mean average precision (mAP) of 54.6 and the baseline method performance ranged from 53.1 to 56.0, demonstrating that though our framework was inspired by a theory of classification the learned features do generalize to other vision tasks. See Appendix M for detailed results.
Note that all included models were trained using the same backbone architecture (ResNet-50), dataset (ImageNet-1k), and number of pretraining epochs (100; we briefly explore the impact of longer pretraining in Appendix P). The results for networks pretrained on smaller datasets can be found in the Appendix I.
3.2 Analyses of Learned Representations
We next conduct a series of experiments to elucidate the differences between representations learned with different SSL procedures, and suggest a mechanism by which augmentation manifold compress
sion gives rise to class separability. To reduce the computational requirements, these analyses are carried out on models trained on the CIFAR-10 dataset.
Figure 2: Mean field Manifold Capacity analysis. Manifold dimensionality (top row), radius (middle row), and capacity (bottom row), as a function of stage in the representational hierarchy (from pixel inputs to output of the encoder/learned representation). Shaded regions indicate  $95\%$  confidence intervals around the mean (analysis was conducted with 5 different random samples from the dataset, see Appendix E).
Mean Field Theory Manifold Capacity. In Fig. 2 we show that our representation, which is optimized using an objective that assumes elliptical manifold geometry, nevertheless yields representations with high values of the more general mean field manifold capacity (relative to baseline methods). For completeness we also analyzed the geometries of class manifolds, whose points are the representations of different examples from the same class. This analysis provided further evidence that learning to maximize augmentation manifold capacity compresses and separates class manifolds, leading to a useful representation. Interestingly MMCRs seem to use a different strategy than the baseline methods to increase the capacity, yieldingclass/augmentation manifolds with larger radii, but lower dimensionality (Fig. 2) These geometrical differences do not emerge until the tail end of the hierarchy, suggesting early layers of each network are carrying out stereotyped transformations and loss function induced specialization does not emerge until later layers.
Emergence of neural manifolds via gradient coherence. We hypothesize that class separability in MMCRs arises because augmentation manifolds corresponding to examples from the same class are optimally compressed by more similar transformations than those stemming from distinct classes. To investigate this empirically, we evaluate the gradient of the objective function for inputs belonging to the same class. We can then check whether gradients obtained from (distinct) batches of the same class are more similar to each other than those obtained from different classes, which would suggest that the strategy for compressing augmentation manifolds from the same class are relatively similar to each other. Fig. 3 demonstrates that this is the case: within class gradient coherence, as measured by cosine similarity, is consistently higher than across class coherence across both training epochs and model hierarchy.
Manifold subspace alignment. Within-class gradient coherence constitutes a plausible mechanistic explanation for the emergence of class separability, but it does not explain why members of the same class share similar compression strategies. To explore this question we examine the geometric properties of augmentation manifolds in the pixel domain. Here we observe small but measurable differences between the distributions of within-class similarity and across-class similarity, as demonstrated in the top row of Fig. 4. The subtle difference in the geometric properties of augmentation manifolds in the pixel domain in turn leads to the increased gradient coherence observed above, which leads to a representation that rearranges and reshapes augmentation manifolds from the same class in a similar fashion (bottom row of Fig. 4), thus allowing better linear separation of classes. Not only are centroids
Figure 3: Cosine similarity of gradients for pairs of single-class batches. We plot the mean pairwise similarity for pairs of gradients for different subsets of the model parameters (all parameters, and the first and last linear operators) obtained from single-class-batches coming from the same or distinct classes over the course of training. Because a large number of bathces were used  $95\%$  confidence intervals about these means are too small to be visible. To the left is a visualization of the fact that single-class gradients flow backward through the model in more similar directions.
of same-class-manifolds in more similar regions of the representation space than those coming from distinct classes (Fig. 4 third column bottom row) but additionally same-class-manifolds have more similar shapes to each other (Fig. 4 bottom row columns 1 and 2 show same-class-manifolds occupy subspaces with lower relative angles and share more variance).
We next ask how the representation learned with the MMCR objective differs from those optimized for other self-supervised loss functions. While MMCR encourages centroids to be orthogonal, the InfoNCE loss (14) encourages negative pairs to be as dissimilar as possible, which is achieved when they lie in opposite regions of the same subspace. The Barlow Twins (77) loss is not an explicit function of feature vector similarities, but instead encourages individual features to be correlated and distinct features to be uncorrelated, across the batch dimension. Fig. 5 shows that these intuitions are borne out empirically: the MMCR representation produces augmentation manifold centroids that are significantly more orthogonal to each other than the two baseline methods.
Figure 4: The distributions of various similarity metrics for augmentation manifolds from the same (orange) and distinct (blue) classes. These are shown for both the input images (top row), and the learned representation (bottom row). Left: schematic illustration of the exemplar-augmentation manifold-class manifold structure of the learned representation.
3.3 Biological relevance
Neuroscience has provided motivation for many of the developments in artificial neural networks, and it is of interest to ask whether SSL networks can characterize the measured behaviors of neurons in biological visual systems. As a simple test, Table 2 shows performance of our model compared with five other SSL models on the BrainScore repository (64; 30; 53; 51). We find that MMCR achieves the highest performance in explaining neural data from primate visual areas V2 and V4, second-highest for V1 and IT, and is the most or second most predictive for 8 out of the 11 individual datasets (see Appendix K).
In addition, we examine general response spectral characteristics that have been recently described for neural populations. In particular, Stringer et al. (67) reported that the eigenspectrum of the
Figure 5: Cosine similarities of centroids for models trained according to different SSL objectives. The left panel shows the distribution of cosine similarities for centroids of augmentation manifolds for examples of the same class, while the right shows the same distribution for examples from distinct classes. Note that because we are analyzing the outputs of the ResNet backbone (which are rectified), the minimum possible cosine similarity is 0.
Table 2: Neural prediction and spectral properties of self-supervised models. First four columns provide average BrainScore correlation values (64) for neurons recorded in four areas of primate ventral stream. Last two provide participation ratio (PR) and spectral decay coefficients  $(\alpha)$ , estimated for 10 bootstrapped samples of the features for the ImageNet validation set. When estimating the spectral decay coefficient we omitted the tails (where the power decays more rapidly), see L for exact details. Entries indicate mean and standard error of the mean, and boldface indicates best performance within standard error. MMCR results are for a model trained with 8 views.
Model	V1	V2	V4	IT	PR	α
MMCR	.494 ± .006	.311 ± .005	.481 ± .005	.416 ± .003	279.2 ± .3	1.04 ± 1e-4
SimCLR	.500 ± .007	.288 ± .007	.475 ± .004	.420 ± .003	124.6 ± .1	1.34 ± 3e-4
BYOL	.500 ± .006	.291 ± .007	.477 ± .005	.404 ± .003	248.0 ± .4	1.35 ± 1e-4
MoCo	.499 ± .007	.293 ± .006	.477 ± .005	.417 ± .003	147.7 ± .2	1.44 ± 3e-4
Barlow	.498 ± .007	.293 ± .008	.477 ± .005	.404 ± .003	252.2 ± .2	1.21 ± 3e-4
SwAV	.488 ± .007	.296 ± .009	.463 ± .004	.398 ± .003	163.9 ± .2	1.15 ± 3e-4
covariance matrix of population activity in visual area V1 follows a power law decay with a decay coefficient of approximately 1 ( $\lambda_{n} \propto n^{-\alpha}$  with  $\alpha \approx 1$  where  $\lambda_{n}$  is the  $n^{th}$  eigenvalues). Subsequent studies in artificial networks have found that such a decay spectrum is associated with increased robustness to adversarial perturbations and favorable generalization properties (55; 2). Additionally, several recent works have investigated the connection between representational dimensionality and neural predictivity (70; 63). In particular, Elmoznino and Bonner (27) report that high intrinsic dimensionality (as measured with the participation ratio of the representation covariance) is correlated with stronger ability to predict neural activity. Table 2 provides values for the participation ratio of each representation over the ImageNet validation set, as well as the decay coefficient of the covariance spectrum (see Appendix L for more details on each experiment). We see that the MMCR has the highest participation ratio (dimensionality) - note that this differs from the quantity optimized in the objective function, which lies in the embedding space. In addition, MMCR also also yields features with a decay coefficient that is nearest to one.
We find that in this controlled setting each model explains a very similar fraction of neural variance through linear regression. This is consistent with recent and concurrent works (21; 62) which have identified this lack of ability to discriminate between alternative models as a weakness of the dominant paradigm used for model-to-brain comparisons. However, our results demonstrate that different SSL algorithms produce representations with meaningfully different geometries (as evidenced by the large spread in the spectral properties such as the participation ratio and decay coefficient). This suggests the need for the development of new metrics for comparing models to data, such as geometrical measures, that capture these important differences between candidate models.
We have presented a novel self-supervised learning algorithm inspired by manifold capacity theory. Many existing SSL methods can be categorized as either "contrastive" or "non-contrastive" depending on whether they avoid collapse by imposing constraints on the embedding gram or covariance matrix, respectively. Our framework strikes a compromise, optimizing the singular values of the embedding matrix itself. By directly optimizing this population level feature (the spectrum), we are able to encourage alignment and uniformity (71) simultaneously with a single-term objective. Additionally, this formulation circumvents the need for making a large number of pairwise comparisons, either between instances or dimensions. As a result learning MMCRs is efficient, requiring neither large batch size nor large embedding dimension. Finally, our method extends naturally to the multi-view case, offering improved performance with minimal increases in computational cost.
Our formulation approximates manifold geometries as elliptical, reducing computational requirements while still yielding a useful learning signal and networks with high manifold capacity. Specifically, we were able to leverage manifold capacity analysis in its full generality to gain insight into the geometry of MMCR networks after training. Further research could explore objectives based on more modest reductions of mean field manifold capacity that capture non-elliptical structure. Intriguingly, our method produces augmentation and class manifolds with lower dimensionality but larger radius than either Barlow Twins or SimCLR (Fig. 2). We do not understand why this is the case, but the differences indicate that capacity analysis can provide a useful tool for elucidating the different encoding strategies encouraged by various SSL paradigms.
Finally we investigated two recently proposed theories on ideal spectral properties for neural representations. For our considered set of models a spectral decay coefficient near 1 was associated with better performance on the within-distribution task and generalization to unseen datasets (MMCR, SwAV, and Barlow were the top three models for each of the out-of-distribution classification tasks), a finding which is broadly aligned with both the empirical and theoretical findings of Agrawal et al. (2). However, we also found that high dimensionality did not always correspond to strong neural predictivity: Despite having the lowest dimensionality, SimCLR performed strongly in terms of neural predictivity. This implies, perhaps unsurprisingly, that global dimensionality alone is not sufficient to explain the response properties of neurons in the primate ventral stream. In fact our experiments add to a growing body of work on the need for complementary approaches to linear predictivity for discriminating between candidate models of the visual system (21). Happily the field is already moving to address this limitation, for instance concurrent work (12) finds that decomposing neural predictivity error into distinct modes can yield insights into how different models fit different aspects of the neural data (even if they yield similar overall predictivity).
One promising direction for improving the quality of artificial networks as models of neural computations is to incorporate the constraints associated with biologically plausible learning (i.e. the need for local learning rules, Illing et al. (44)). A complementary direction is to better align training data diets with ecological inputs. For example, temporal invariance rather than augmentation invariance seems a more plausible objective for a neural system to optimize (73; 3; 61). We speculate that a variant of the MMCR objective that operates over time may be well-suited to a neural circuit implementation as its computation would only require mechanisms for tracking (1) short timescale temporal means (to form centroids) and (2) the singular values of the population activity over longer timescales.
Neuroscience is brimming with newly collected datasets recorded from ever larger populations of neurons, and there is a long history of methods that aim to make sense of these measurements through a normative lens. Here we have demonstrated that one such technique that has proven useful for gleaning insights from neural data (60; 75; 31) can be reformulated for use as an objective function to learn a useful abstract representation of images. Future work should aim to close the loop between modelling and analysis by using these learned models to generate experimentally testable predictions and constrain new experimental designs.
[71] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929-9939. PMLR, 2020.
[72] Yu Wang, Jingyang Lin, Qi Cai, Yingwei Pan, Ting Yao, Hongyang Chao, and Tao Mei. A low rank promoting prior for unsupervised contrastive learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[73] Laurenz Wiskott and Terrence J Sejnowski. Slow feature analysis: Unsupervised learning of invariances. Neural computation, 14(4):715-770, 2002.
[74] Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3733-3742, 2018.
[75] Justin D Yao, Klavdia O Zemlianova, David L Hocker, Cristina Savin, Christine M Constantinople, SueYeon Chung, and Dan H Sanes. Transformation of acoustic information to sensory decision variables in the parietal cortex. Proceedings of the National Academy of Sciences, 120 (2):e2212120120, 2023.
[76] Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning diverse and discriminative representations via the principle of maximal coding rate reduction. Advances in Neural Information Processing Systems, 33:9422-9434, 2020.
[77] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pages 12310-12320. PMLR, 2021.
A Optimal Embeddings
Recall the setting of self-supervised learning as described in Balestriero and LeCun (6): given a dataset  $\mathbf{X}' = [\pmb{x}_1, \dots, \pmb{x}_N]^T \in \mathbb{R}^{N \times D'}$  we construct a new dataset by creating  $k$  randomly augmented views of the original data,  $\mathbf{X} = [\mathrm{view}_1(\mathbf{X}')', \dots, \mathrm{view}_k(\mathbf{X}')'] \in \mathbb{R}^{Nk \times D}$ . The advantage of doing so is that we can now leverage the knowledge that different views of the same underlying datapoint are semantically related. We can express this notion of similarity in the symmetric matrix  $\mathbf{G} \in \{0, 1\}^{Nk \times Nk}$  with  $G_{ij} = 1$  if augmented datapoints  $i$  and  $j$  are semantically related (and  $G_{ii} = 1$  as any datapoint is related to itself). We can normalize  $\mathbf{G}$  such that its rows and columns sum to 1 (so rows of  $\mathbf{G}$  are  $k$ -sparse with nonzero entries equal to  $1 / k$ ).
$$
\begin{array}{l} \mathcal {L} = - \left| \left| G Z \right| \right| _ {*} \\ = - \left\| \boldsymbol {Q} \boldsymbol {\Lambda} \boldsymbol {Q} ^ {T} \boldsymbol {U} \boldsymbol {S} \boldsymbol {V} ^ {T} \right\| _ {*} \tag {7} \\ = - \left\| \boldsymbol {\Lambda} \boldsymbol {Q} ^ {T} \boldsymbol {U} \boldsymbol {S} \right\| _ {*} \\ \end{array}
$$
Where we have taken the eigendecomposition of  $G$  which is real and symmetric and the SVD of  $Z$ , and then used the fact that the singular value spectrum is invariant under left or right orthogonal transformations. We now show that a global optima of this objective is achieved when the left singular vectors of  $Z$  are the eigenvectors of  $G$  and the singular values of  $Z$  are proportional to the eigenvalues of  $G$ . Throughout we will assume that the size of the dataset is greater than the dimensionality of the embeddings,  $N > d$ , as is the case in practical applications. First we prove a simple lemma about the spectrum of matrices who are extended by zeros (i.e. embedded in a higher dimensional space).
Lemma A.1: For  $\mathbf{A} \in \mathbb{R}^{N \times N}$ ,  $\mathbf{B} \in \mathbb{R}^{N \times d}$  with  $d < N$ ,  $||\mathbf{A}\mathbf{B}||_{*} = ||\mathbf{A}\tilde{\mathbf{B}}||_{*}$  where  $\tilde{\mathbf{B}} = [\mathbf{B},\mathbf{0}] \in \mathbb{R}^{N \times N}$ .
Proof: First note that  $\mathbf{A}\tilde{\mathbf{B}} = [\mathbf{A}\mathbf{B},\mathbf{0}]$  so it suffices to show that for arbitrary  $\mathbf{X}$  that  $\sigma (X) = \sigma ([X,0])$ . Taking the SVD of  $\mathbf{X}$ ,
$$
\boldsymbol {X} = \left[ \begin{array}{l l} \boldsymbol {U} & \widetilde {\boldsymbol {U}} \end{array} \right] \left[ \begin{array}{l} \boldsymbol {\Sigma} \\ \boldsymbol {0} \end{array} \right] \left[ \begin{array}{l} \boldsymbol {V} ^ {T} \end{array} \right] = \boldsymbol {U} \boldsymbol {\Sigma} \boldsymbol {V} ^ {T}
$$
Then a valid singular value decomposition for  $\tilde{X}$  is
$$
\tilde {\boldsymbol {X}} = \left[ \begin{array}{l l} \boldsymbol {U} & \tilde {\boldsymbol {U}} \end{array} \right] \left[ \begin{array}{l l} \boldsymbol {\Sigma} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {0} \end{array} \right] \left[ \begin{array}{l l} \boldsymbol {V} ^ {T} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {I} \end{array} \right]
$$
Clearly then,  $||X||_{*} = ||\tilde{X} ||_{*}$
Theorem: The proposed loss achieves a global minimum when the left singular vectors of  $Z$  are the eigenvectors of  $G$ , and the singular values of  $Z$  are proportional to the top  $d$  eigenvalues of  $G$ .
Proof: Let  $\tilde{\pmb{Z}} = [\pmb {Z},\mathbf{0}]\in \mathbb{R}^{N\times N}$ . By Lemma A.1 we have  $||\pmb {G}\pmb {Z}||_{*} = ||\pmb {G}\tilde{\pmb{Z}} ||_{*}$ . Von Neumann's trace inequality can be used to show  $||\pmb {G}\tilde{\pmb{Z}} ||_{*}\leq \sum_{i = 1}^{Nk}\sigma_{i}(\pmb {G})\sigma_{i}(\tilde{\pmb{Z}})$  (see Marshall et al. (54) for proof). Examining (4) it is clear that this bound is achieved when  $\pmb {U} = \pmb{Q}$ . The problem can therefore be reduced to the constrained optimization problem,
$$
\min  _ {\boldsymbol {\sigma} _ {i} (\tilde {\boldsymbol {Z}})} \sum_ {i = 1} ^ {N k} \boldsymbol {\sigma} _ {i} (\boldsymbol {G}) \boldsymbol {\sigma} _ {i} (\tilde {\boldsymbol {Z}})
$$
$$
\text {s u b j e c t} \sum_ {i = 1} ^ {N k} \sigma_ {i} (\tilde {\mathbf {Z}}) ^ {2} = N k
$$
where the constraint comes from the fact that columns of  $\mathbf{Z}$  are unit vectors. Intuitively, we are maximizing the inner product between a fixed vector  $\sigma(\mathbf{G})$  and a vector with fixed L2 norm. The
solution of course is to align the two vectors as closely as possible, i.e. when  $\sigma_{i}(\tilde{\mathbf{Z}})\propto \sigma_{i}(\mathbf{G})$  for  $i = 1,\ldots ,d$ . It is worth noting that by construction  $\sigma_{i}(\tilde{\mathbf{Z}}) = 0$  for  $i > d$  and the columns of  $\mathbf{U}$  associated with these zero valued singular values are unconstrained.
B Pytorch Style Pseudocode for MMCR
1 # h: encoder
2 # g: projection head
3 # T: momentum temperature
4 # B: batch size
5 # K: number of augmentations
6 # D: projector output dimensionality
7 # 
8 # lmbda: trade-off parameter
9 
10 f_o, g_o = ResNet50(), MLP() # online networks
11 
# initialize momentum network with identical params
12 f_m, g_m = f_o.copy, g_o.copy()
13 # momentum networks are not updated via gradient descent
14 
f_mrequires_grad = False
15 # momentum networks are not updated via gradient descent
16 f_mrequires_grad = False
17 g_mrequires_grad = False
18 
for x in loader:
20 # K randomly augmented views
21 x = multi_augment(x) # B x K x H x W
22 
# push through encoder and projector
23 z_o = g_o(h_o(x)) # B x K x D
24 z_m = g_m(h_m(x)) # B x K x D
25 z = concatenate(z_o, z_m, dim=1) # append outputs
26 
# project onto unit sphere
27 
z = normalize(z, dim=-1)
28 
# calculate centroids (mean over augmentation axis)
29 c = z.mean(dim=1) # B x D
30 
# calculate singular values
31 U_z, S_z, V_z = svd(z) # batch svd
32 U_c, S_c, V_c = svd(c)
33 
# calculate loss
34 
loss = -1.0 * sum(S_c) + lambda * sum(S_z) / B
35 
# backward pass and optimization step
36 
loss-backward()
37 
# perform momentum update
38 with torch.no_grad():
40 f_m.params() = (1 - T) * f_o.params() + T * f_m.
41 parameters()
42 g_m.params() = (1 - T) * g_o.params() + T * g_m.
43 parameters()