Contrastive-Equivariant Self-Supervised Learning Improves Alignment with Primate Visual Area IT

Thank you for your constructive review of our submission. We respond to the questions and limitations point-by-point below:

• Performance on Downstream Classification Tasks (W1, Q3): Please see our general response (Point 3) for our more detailed thoughts on why equivariant training on a significantly large dataset offers little in the way of improvements on downstream classification tasks. In short, our results on the ImageNet-100 training vs. ImageNet-1000 training suggest that the benefits induced by equivariant training “overlap” with those that come from using a large dataset (Line 252 of the submission). With this result, we aim to draw attention to the underappreciated fact that equivariant training seems to be more valuable when the pretraining dataset is smaller.

• Lack of a tradeoff between invariant and equivariant terms (W2): We think there might have been some confusion surrounding our use of the word “tradeoff”; when we refer to a “tradeoff between invariance and structured variability” (Line 54) we are referencing qualities of the artificial representations, not necessarily proposing that IT neurons are explicitly striking some balance between invariance and equivariance. In the artificial representations such a tradeoff is clear, for example, as the importance of the equivariance loss increases the amount of linearly accessible information about transformations applied to the input (Figure 2E), but this comes at the cost of a modest decrease in performance on classification on the in-distribution ImageNet-1k dataset (Figure 4). Furthermore while the improvements in neural predictivity are small in absolute terms, this is common for linear-regression based comparisons, and our improvement does result in a substantial increase in terms of the relative ranking of models on the IT Brain–Score leaderboard (see our general response and Figure R2 for more details).

  If we have misinterpreted this we are more than happy to engage further during the discussion period! We also will remove the reference to a “tradeoff” in line 58 with the following modification, to help avoid confusion: “We explore the impact of including an equivariant loss for predicting neural activity in IT,”

• Dependence on the input transformations (W3, Q4): We fully agree that this is an important point and could be better addressed. Thank you in particular for suggesting we experiment with testing how well the learned equivariances can generalize to unseen (stronger, but similarly typed) transformations. As a step towards answering this question we trained models using the same suite of augmentations but with a uniformly reduced strength, and repeated the parameter decoding experiments from Figure 2. These preliminary results suggest there is a degree of generalization to stronger transformations, see our general response and Figure R1 of the attached PDF for more details. We are currently working on the experiment you described involving rotations, and hope to have additional results to report during the discussion period.

  Considering other more complicated sets of input transformations than those commonly used in SSL is also an interesting direction for further research that we can highlight more directly. We feel the most important and ecologically relevant set of input transformations are those that mimic the visual experience of animals in the world, and as we note in the discussion we believe this is a very promising direction for future work!

• Is SimCLR less suited to incorporating equivariances (Q1): Our results do suggest that SimCLR may indeed be less amenable to smoothly shaping the representation via an equivariance loss. However another hypothesis is that these nonsmooth curves arise from the fact that SimCLR is in general more sensitive to various hyperparameter choices, and the one introduced by our equivariance loss is not an exception. For example, it is well known that SimCLR is more sensitive than other self-supervised learning methods to batch size [1]. To definitively determine whether this is a fundamental difference between SimCLR and the other objective functions would require large sweeps over a variety of hyperparamters (batch size, learning rate, etc.), and compute limitations prevent us from conducting such experiments. We will expand Appendix A.7: Limitations, to include this discussion.

• Interpreting the impact of the equivariance loss on neural predictivity (Q2): We agree that gaining better understandings and intuitive interpretations of how different features of learned representations lead to different levels of neural predictivity is an important goal for both our work and the field at large. We attempt to draw conclusions by analyzing how the representational analyses from Section 3.2 correlate with neural predictivity in Table 1. However we agree that there is more work to be done in terms of interpretation. For example one could examining the structure of the residuals and mapping weights obtained in the Model-to-Brain regression problem [2]. We could also begin to answer some of these questions using new and open source datasets of neural measurements in Macaque IT [3] and will highlight the importance of such investigations in the revised manuscript.

• Typos (Q5): Thank you for drawing our attention to this mistake, we will be sure to correct the figure reference in the revised manuscript.

[1] Zbontar, Jure, et al. "Barlow twins: Self-supervised learning via redundancy reduction." International conference on machine learning. PMLR, 2021.

[2] Canatar, Abdulkadir, et al. "A spectral theory of neural prediction and alignment." Advances in Neural Information Processing Systems 36 (2024).

[3] Madan, Spandan, et al. "Benchmarking Out-of-Distribution Generalization Capabilities of DNN-based Encoding Models for the Ventral Visual Cortex." arXiv preprint arXiv:2406.16935 (2024).

Thank you for your thoughtful assessment of our paper. Find below is our point-by-point response to the weaknesses and questions raised:

• Marginal Effect on Neural Predictivity: Please see our general response Point 1 and Figure R2 in the rebuttal pdf which directly address this concern about the marginal effect on IT predictions. Specifically, the changes observed are substantial in terms of the relative ranking compared to other models on this benchmark.

• Writing/Clarity Issues: You are correct, there is a misleading typo in Eqn. 2 and the weightings of the two terms are reversed. Thank you for pointing out this mistake, we will correct it in the revised manuscript to $L_{\rm overall} = (1 - \lambda) L_{iSSL} + \lambda L_{CE-SSL}$. Similarly the reference on line 113 should indeed be to Figure 1. Finally, we will also correct the citation error on line 86 and change the icon for projector networks in Figure 1 to better indicate the transition from a convolutional backbone to an MLP projector.

• Performance on ImageNet-1k Classification: Please see our general response Point 3, which provides reasoning for why downstream accuracy on ImageNet-1k degrades as we increase the importance of the equivariance loss. We agree that this issue could be better highlighted in the main text and propose to include language similar to that in our general response in the final version of the paper.

Thank you for your constructive assessment of this submission. Below is our point-by-point response to the weaknesses and questions raised:

• Lack of baselines/context for brain data: We agree that the contribution is clarified by providing more context for the predictivities for a range of models. See our general response Point 1 and the Figure R2 in the attached PDF for our plan to do so.

• Uniqueness of Improvements: This is indeed a relevant consideration. We suspect that our method induces representations that are explaining disparate portions of the neural response variance from task-trained models for three reasons. . First, self-supervised trained models provide neural predictivity that is on par with task trained networks (i.e. supervised recognition networks) [1, 2]. Second, above a certain threshold (in the regime that modern computer vision networks occupy) task performance (top-1 accuracy on ImageNet-1k) is actually negatively correlated with neural predictivity [3]. Finally it is worth noting that after retraining our most predictive model for an increased number of epochs, we now obtain SOTA performance on IT predictivity. However, it would certainly be interesting to more directly answer the question of how “shared” the explained variance is between disparate sets of models, for example by examining the structure of the residuals obtained in the Model-to-Brain regression problem [4]. We could begin to answer some of these questions using new and open source datasets of neural measurements in Macaque IT [5] and will highlight the importance of such investigations in the revised manuscript.

• Writing in Section 3.2: We agree that these experiments could be described more clearly. The idea for a table that describes the sources of variability being compared in each panel is a good one and we will include an Appendix with additional details in the revised manuscript. We will also make the suggested updates to axis labels.

• Thank you for pointing out these typographical issues! We will correct these during the revision process.

• Confusion regarding lines 161-162: Thank you for pointing out the ambiguity in our word choice. There need not be a plural in this sentence, the intended meaning was that we repeated this pairwise comparison (between one invariant and one equivariant network) for each of the equivariant networks we trained (different base objective functions and different values of lambda). We propose to change this sentence to:

  "In particular, we estimate C1 and C2 over identical inputs for an invariant network and an equivariant network trained using the same base objective but a non-zero value of $\lambda$."

• Grey Inset: This was referring to the gray square in Figure 1 and is another typographical error. We apologize for the confusion and will correct this error.

[1] Zhuang, Chengxu, et al. "Unsupervised neural network models of the ventral visual stream." Proceedings of the National Academy of Sciences 118.3 (2021): e2014196118.

[2] Conwell, Colin, et al. "What can 1.8 billion regressions tell us about the pressures shaping high-level visual representation in brains and machines?." BioRxiv (2022): 2022-03.

[3] Schrimpf, Martin, et al. "Brain-score: Which artificial neural network for object recognition is most brain-like?." BioRxiv (2018): 407007.

[4] Canatar, Abdulkadir, et al. "A spectral theory of neural prediction and alignment." Advances in Neural Information Processing Systems 36 (2024).

[5] Madan, Spandan, et al. "Benchmarking Out-of-Distribution Generalization Capabilities of DNN-based Encoding Models for the Ventral Visual Cortex." arXiv preprint arXiv:2406.16935 (2024).

Thank you for your careful review of our submission. We respond to the listed weaknesses and questions below:

• Introducing a new hyperparameter to tune: we agree that this is a fundamental weakness of our framework. It would be preferable to balance this tradeoff without using two loss functions and a Lagrange multiplier. We will add the following to the limitations section (Line 524):

  “Our current training setup requires selection of the hyperparameter $\lambda$ to balance between the equivariant and invariant loss functions. Future work could investigate methods to balance the two losses without explicitly training an individual network for each choice of $\lambda$.”

• Performance on OOD Classification tasks: Please see our general response for a more detailed explanation for why equivariance leads to little improvement in OOD classification performance (relative to invariant SSL) when trained on a sufficiently large and diverse set of natural images. While we agree that this in some sense limits the applicability of our method in AI/ML representation learning, we feel that this observation is in and of itself a contribution to the community. A similar effect was observed in some past work on Equivariant SSL (see Table 2 vs Table 3 in [1]), but we feel that it can only benefit the community to draw attention to the interplay between inductive biases (enforced via the loss) and the effect of dataset scaling on learned representations.

• Other backbone architectures: We agree that it is important to confirm that the effects we report are not limited to a specific choice of architecture. In response to this concern, we trained a limited set of models using either ResNet-34 or ResNet-101 backbones in place of the ResNet-50 considered in the text. We found similar effects in terms of the neural predictivity of the equivariant networks relative to their invariant counterparts (see Figure R3 in the rebuttal pdf). We acknowledge that this is not a very radical departure from our initial architecture, but time prevented us from considering more modern backbones such as ResNext and ViT. We will endeavor to include experiments with these backbones in the final version of our paper, and will certainly include the experiments from Figure R3 of the Rebuttal PDF as a new Appendix.

• Robustness to choice of augmentations: This is an interesting question to consider. To begin answering this we conducted experiments involving pretraining models using weaker versions of the same transformations (see general response and Figure R1 of the PDF for more details). Of course these experiments do not entirely answer the question of how robust our method and observations are to the choice of input transformations, though they do suggest a form of generalization (from representing weak-to-strong transformations) is possible. We feel the most important and ecologically relevant set of input transformations are those that mimic the visual experience of animals in the world, and as we note in the discussion we believe this is a very promising direction for future work!

• Other input modalities: Great question - we have been thinking about how this technique might be applied to learning representations of audio signals! In this set-up we envision using a speech-in-background setting: where “positive” pairs would consist of the same speech signal with different choices of background noise. Networks can then be trained to be invariant to background noise, or equivariant to background noise using the method described in this work. While these experiments are beyond the scope of this paper, we plan on following up in this direction with subsequent work. We had not considered applying our method to text, though in view of our analogy between invariance/equivariance and slowness/straightness and the recent observations in [2], this may also be an interesting possibility. We will add the following sentence to the discussion line 296 to capture this idea:

  “Although in this work we focused on the visual domain, similar equivariant and invariant objectives could be investigated for other domains such as audio and langauge representation learning.”

• Other forms of Self-Supervision: In terms of clustering, we believe that existing methods that use clustering based contrastive losses (i.e. SwAV) would likely show similar results to the base objective functions considered in this work. A reconstruction loss would certainly also “fight-against” the invariance term, as it explicitly encourages that information about the input pixels be preserved. It would be interesting to see whether such a scheme also produced more predictive representations of IT, though this may be out of the scope of the current work.

• Presentation: We will also add additional details on projector network architectures to Appendix A.2: Additional pretraining details.

[1] Chavhan, Ruchika, et al. "Amortised Invariance Learning for Contrastive Self-Supervision." The Eleventh International Conference on Learning Representations, 2023.

[2] Hosseini, Eghbal, and Evelina Fedorenko. "Large language models implicitly learn to straighten neural sentence trajectories to construct a predictive representation of natural language." Advances in Neural Information Processing Systems 36 (2024).