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

Authors: thomas  yerxa
jenelle  feather
eero p simoncelli
sueyeon  chung
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:44:17

Abstract: Models trained with self-supervised learning objectives have recently matched or surpassed models trained with traditional supervised object recognition in their ability to predict neural responses of object-selective neurons in the primate visual system. A self-supervised learning objective is arguably a more biologically plausible organizing principle, as the optimization does not require a large number of labeled examples. However, typical self-supervised objectives may result in network representations that are overly invariant to changes in the input. Here, we show that a representation with structured variability to the input transformations is better aligned with known features of visual perception and neural computation. We introduce a novel framework for converting standard invariant SSL losses into "contrastive-equivariant" versions that encourage preserving aspects of the input transformation without supervised access to the transformation parameters. We further demonstrate that our proposed method systematically increases models' ability to predict responses in macaque inferior temporal cortex. Our results demonstrate the promise of incorporating known features of neural computation into task-optimization for building better models of visual cortex.

Main paper: 
 introduction: 
 In the past decade, task-optimized deep neural networks (DNNs) have been used to predict responses of object-selective neurons in primates to natural image stimuli [Yamins et al., 2014 , Schrimpf et al., 2020 , Willeke et al., 2023] . Such networks have a pronounced advantage over more traditional models for explaining responses in deeper areas with more abstract representations, such as inferior temporal cortex (IT). This observation naturally leads to the hypothesis that task optimization can provide a normative account for IT neuron tuning properties: visual representations are shaped by the need to perform ecologically relevant tasks.However, the task that initially led to these advances was that of supervised object classification, a specific task that relies on an implausibly large number of labeled examples [Lindsay, 2021] . More recently, computer vision has undergone a "self-supervised learning" (SSL) revolution. A variety of methods have been proposed to learn representations that match or surpass supervised training on multiple tasks by deriving sources of supervision from the data itself rather than relying on human annotations. For example, many popular SSL strategies aim to group representations of different synthetic transformations of the same images (commonly referred to as "views"), while enforcing diversity among representations of distinct images. Additionally, self-supervised representations can predict neural responses with fidelity comparable to supervised representations [Zhuang et al., 2021 , Konkle and Alvarez, 2022 , Parthasarathy et al., 2024] .Both of these training objectives are forms of invariance learning: in object classification model outputs become invariant across different objects from the same class, and in self-supervised learning they become invariant to the transformations used to generate different views. However biological visual representations are not fully invariant across these sets of inputs [DiCarlo and Cox, 2007, Kuoch et al., 2024] . Indeed it has been demonstrated that training according to either of these two objectives leads to artificial representation with "excess invariance" relative to human perception [Feather et al., 2023] . Additionally, even in Area IT, which is thought to subserve invariant object recognition, neural populations encode a significant amount of "category orthogonal" information (e.g., object pose or viewing conditions that are unrelated to semantic category) [Hong et al., 2016] . Furthermore, such selectivity for object-orthogonal attributes is meaningfully organized within Area IT [Hong et al., 2016] (i.e. object orthogonal attributes are linearly decodable from population responses). Whether such structured variability emerges in invariance-trained networks is likely determined by the uncontrolled inductive biases of the network architecture [Alleman et al., 2024] .Here, we develop an equivariant learning framework that encourages such structured variability in network representations. Our contributions are:• We propose a novel framework that converts standard invariance-based self-supervised learning methods into "contrastive-equivariant" versions that produce structured, transformationrelated variability. Unlike previous approaches, our method does not require supervised access to transformation parameters or costly modifications to the training procedure.• We examine the tradeoff between invariance and structured variability through a series of representational analyses. We find that, relative to networks trained with invariance only, our contrastive-equivariant network learns structured transformation variability that is shared across images and factorized with respect to variability related to changes in image content.• We explore the impact of including an equivariant loss for predicting neural activity in IT, showing for the first time that explicitly encouraging structured variability via optimization leads to an improved ability to predict cortical responses to natural images. 
method: 
 2.1 Transformation-Invariant Self-Supervised Learning (iSSL)The influential work of [Chen et al., 2020] showed that applying two random transformations (often called "augmentations") to a batch of images, then training a network to identify which pairs of transformed images originated from the same sample with a cross-entropy style loss (the InfoNCE loss, first formulated in [Gutmann and Hyvärinen, 2010] ) yields representations that are competitive with supervised training in terms of object classification. Many subsequent studies have developed alternative objective functions that produce similar results: Barlow Twins [Zbontar et al., 2021] ,VICReg [Bardes et al., 2021] , and W-MSE [Ermolov et al., 2021] enforce augmentation invariance along with a constraint that the global covariance matrix is the identity; SimSiam [Chen and He, 2021] and BYOL [Grill et al., 2020] employ architectural constraints that regularize towards uniform representations and simply optimize for transformation invariance. Other studies have formalized the problem in terms of maximizing information [Ozsoy et al., 2022] or capacity [Yerxa et al., 2024] , subject to an invariance constraint, which has enabled connections to normative theories of coding efficiency and manifold capacity [Barlow et al., 1961 , Chung et al., 2018] .To formalize the definition of iSSL, we denote a dataset of images (e.g., ImageNet) by X ∈ R N ×D , where N is the number of images and D is their dimensionality. Let τ (•; ρ) : R D → R D be a function parameterized by ρ that maps images to images (for example τ might be the random crop operation, in which case ρ could specify the region to be cropped). The goal of iSSL algorithms is to learn the parameters W of some function f (•; W ) : R D → R d such that variability over ρ is minimal while preserving variability over X (to avoid trivial solutions). Many methods achieve this by observing pairs of randomly augmented views of a batch of images: X A = τ (X; ρ 1 ), X B = τ (X; ρ 2 ), with ρ 1 , ρ 2 ∼ p(ρ) where p(ρ) is a pre-chosen probability distribution over augmentation parameters. Generally iSSL frameworks employ an objective function that operates on the outputs of f , Z A = f (X A ; W ), Z B = f (X B ; W ). One popular framework is Barlow Twins [Zbontar et al., 2021] , which uses the objective: Left Gray Panel: corresponding images in each subset are augmented using the same set of two random transformations (so the total number of random transformations is halved relative to a standard invariant SSL training scheme). Every view is passed through a representation network f (ResNet-50 in this work) and the outputs are projected into two embedding spaces by different projector networks, g inv and g equi . In the invariant embedding space a standard iSSL loss is applied, while in the equivariant embedding space the same iSSL loss is applied to the difference vectors between transformation-positive pairs (visualized to the right).L BT = Σ i (1 -C ii ) 2 + λΣ i,i̸ =j (C ij ) 2 where C isthe cross-correlation matrix between Z A and Z B . The first term encourages the outputs in response to the same image subject to different augmentations to be correlated, while the second encourages the outputs in response to distinct images to be uncorrelated.Because complete invariance to the transformations employed in iSSL is actually harmful for downstream tasks, most frameworks employ a learnable "projector network" that maps the outputs of the representation network to an embedding space before applying the loss. The nearly ubiquitous use of this "guillotine regularization" [Bordes et al., 2022] , means that most iSSL methods aim to learn a function from which an augmentation invariant subspace can be extracted. While this approach does permit some transformation-related-variability in the representation, there is no explicit control or encouragement of that variability, and no incentive for that variability to be organized in a meaningful way.To induce structured variability to input transformations, we require that an equivariant subspace can be extracted from f alongside the invariant subspace described above. A function is equivariant to a set of input transformations if there exists a corresponding set of output transformations that induce the same changes. In the self-supervised learning setting this property can be written as:EQUATIONwhere τ ρ = τ (•; ρ) and P is the set of possible values of transformation parameters. Note that invariance is a special case of equivariance, in which T ρ is an identity transformation for all ρ. To avoid this degenerate solution, we will require both that similarly transformed inputs be related by the same transformation in the output space, and that differently transformed inputs are related to each other by different transformations. Our approach, detailed below and summarized in Fig. 1 , follows the principle first proposed by [Gupta et al., 2023] : "Equivariance should be learned from pairs of data, as in invariant contrastive learning."First we split our dataset of images into two random non-overlapping equal-sized partitions X 1 , X 2 ∈ R N 2 ×D . Next we apply a randomly selected augmentation to both X 1 and X 2 , so that corresponding rows of X A/B 1 and X A/B 2 contain distinct images that have been subjected to the same augmentations.Note that this decreases the total number of random samples of ρ by a factor of two relative to the standard invariant learning setup. Finally the resulting representation vectors are each fed through two distinct projector networks, zA/B i = g inv (r A/B i ) and zA/B i = g equi (r A/B i). The resulting embeddings are optimized to be invariant to transformations and discriminative across base images, or invariant to base images and discriminative across transformations, respectively. The overall objective (loss) functions is:EQUATIONwhere both terms are written in terms of L, a self-supervised learning loss function that encourages invariance and uniformity [Wang and Isola, 2020 ] (e.g., L BT ) and λ is a hyperparameter that determines the relative importance of extracting an invariant or equivariant subspace from the shared representation. In the notation of Eq. ( 1), by designing L CE-SSL to encourage similar transformations to induce similar displacements in the output space, we are implicitly specifying that our output transformations are of the form T ρ (z) = z + z ρ . Thus we leverage the principles underpinning contrastive invariance learning to encourage representations that contain useful transformation-related information; this choice differentiates this formulation from previous equivariant self-supervised learning approaches. 
implementation details: 
 Architecture and invariance ojective. For all experiments we use a ResNet-50 architecture [He et al., 2016] as the backbone representation network f . Our training scheme is compatible with any choice of invariance based self-supervised learning frameworks (which correspond to using different choices for L iSSL ). We experimented with three different base methods chosen to span the spectrum from "instance contrastive" to "dimension contrastive" [Garrido et al., 2023a] : SimCLR [Chen et al., 2020] , MMCR [Yerxa et al., 2024] , and Barlow Twins [Zbontar et al., 2021] . In each case we adopt the projector network architecture proposed in the original work for g inv . To retain the synergy between the normalization scheme, loss function, and projector architecture achieved by each framework we use the same architecture for both g inv and g equi .Pretraining dataset and augmentations. We train using the ImageNet-1k dataset and the standard set of augmentations first introduced in [Grill et al., 2020] (which includes random resized cropping, color jittering, Gaussian blurring, solarization, and horizontal flips). See the Appendix A.2 for exact training details.Invariance-equivariance tradeoff. For each of the 3 choices L iSSL we swept the hyperparameter λ over [0.0, 0.001, 0.1, 0.2, 0.3, 0.4, 0.5], yielding a total of 21 learned representations. λ = 0 corresponds to the invariant case. We did not explore beyond λ = 0.5 because classification performance begins to severely degrade beyond this point (see Appendix A.4).Bures metric comparisons. We conducted a series of experiments to determine the extent to which various sources of variability in our dataset were meaningfully organized. The first set of experiments utilized the Bures metric, which is the Wasserstein ("Earth Mover's") distance between mean-centered Gaussian distributions with covariance matrices C 1 and C 2 :EQUATION)When C 1 and C 2 are normalized to have a trace of 1, the maximal distance of 2.0 occurs when the variabilities lie in orthogonal subspaces (or are completely "factorized" from each other) and the minimum distance of 0.0 occurs when the covariances are equal. More generally, a large Bures distance indicates two sources of variability are factorized from each other and a low distance indicates shared structure. We first estimate the trace-normalized covariance of the outputs of some network f over two sources of variability and compute the Bures metric between the two.Because we are mainly interested in the impact of the equivariance loss relative to the invariant baseline, each analysis below is carefully controlled to expose any structural differences. In particular, we estimate C 1 and C 2 over identical inputs for an invariant network and an equivariant network trained using the same base objective but a non-zero value of λ. We then can directly compare the measured Bures distance for the invariant network and each equivariant network (λ ̸ = 0):∆D B = D B (C λ=0 1 , C λ=0 2 ) -D B (C λ̸ =0 1 , C λ̸ =0 2). In each of the panels in the bottom row of Fig. 2 we show how E[∆D B ] evolves as a function of λ when C 1 and C 2 are estimated over different sources of variability (axis labels indicate the ensembles over which the variabilities were estimated). We show the joint distribution of(D B (C λ=0 1 , C λ=0 2 ), D B (C λ̸ =0 1 , C λ̸ =0 2)) and summarize the sources of variability in each experiment described below in Appendix A.5.Augmentation-Augmentation alignment. First we determine the extent to which augmentation variability is shared across base images in each network (Fig. 2A ). For these experiments, both C 1 and C 2 are estimated over many random transformations of single images in the validation set (we will refer to the responses to such a group as a "augmentation manifold"). The expectation is then over randomly sampled pairs of augmentation manifolds. The positive expected difference of distance indicates the equivariant networks consistently produce lower distance between augmentation manifolds, indicating more shared augmentation variability across base images. This structure is closely related to what is encouraged by the equivariance loss term and the orderly increase as a function of λ suggests that we are optimizing effectively.Augmentation-Centroid factorization. We next investigate the extent to which variability over augmentations is factorized from variability over base images (Fig. 2B ). We use the "centroid manifold," [Yerxa et al., 2024] to characterize the variability over base images, by measuring the covariance over base images of the means over augmentations. That is, for these experiments (for each network respectively) C 1 is held fixed as the covariance of the centroids of all augmentation manifolds and C 2 is the covariance of a randomly selected augmentation manifold. We observe that equivariant networks generally exhibit a larger distance between centroid and augmentation manifolds indicating increased factorization of image-content variability and image-augmentation variability. This structure was not explicitly encouraged by the objective and can be considered an emergent property of the equivariant learning procedure.Spatial-Photometric factorization. Next we ask whether our equivariant training procedure induced increased factorization of variability to different types of input transformations (Fig. 2C ). The standard augmentation procedure involves first taking a random crop (spatial variability) of a given image and then applying a series of pixel-level transformations (color-jittering, gaussian blurring, etc.) (photometric variability). To assess the impact of these two distinct classes of image transformations we first chose 20 random crops a given image, then applied the same set of 20 random photometric transformations to each individual crop, yielding 400 different views of each base image. C 1 and C 2 are then estimated over network responses that are averaged over different crops or different photometric transformations respectively, and the expectation is taken over different (single) base images. We observe the equivariant networks consistently exhibit increased factorization (i.e. larger Bures distances relative to the invariant trained network). This again is an emergent property of the equivariant learning procedure, and is particularly interesting in light of recent work that discovered that this form of transformation-factorization is more correlated with neural predictivity than transformation invariance [Lindsey and Issa, 2024] .Class-Class factorization. Finally we asked whether within-class variability was more or less shared between distinct classes in equivariant networks by estimating C 1 and C 2 over responses to all images in distinct classes in the validation set (the expectation is then taken over different random pairs of classes) (Fig. 2D ). Increased sharing of variability between class manifolds has been demonstrated to increase manifold capacity, and can make representations better suited for multi-task evaluations [Wakhloo et al., 2023 [Wakhloo et al., , 2024]] . We observe lower expected pairwise Bures distances in the equivariant networks indicating that the "class manifolds" are better aligned in equivariant networks relative to the invariant networks.Linear embedding of augmentation-related information. While the above experiments demonstrate that equivariant training induces increased alignment of transformation-related variability between images, this does not necessarily imply that this variability is coherently organized. To assess this more directly, we measure the extent to which augmentation parameters can be linearly decoded from the networks' representations. Specificlly, we regress the concatenated outputs of a clean and transformed image onto the parameters of the applied augmentation. We report the resulting coefficient of determination (R 2 ) on a heldout set of validation images (Fig. 2 ). The equivariant training is seen to increase the amount of linearly accessible augmentation information relative to invariant training (the leftmost points plotted in Fig 2E ). We further analyzed a set of equivariant models trained with weaker augmentation parameters (see A.6 for details). In these networks, we again observe that equivariant training increased the amount of linearly accessible augmentation information compared to invariant training. This holds not only for augmentation parameters within the training range (left panel 7) but also for parameter values beyond the training range (right panel 7). Thus, the equivariance properties of the models generalize beyond the training distribution. Future work could examine generalization to unseen types of augmentations.We utilized the BrainScore evaluation pipeline Schrimpf et al. [2018] to measure the extent to which each learned representation can linearly predict neural responses measured in macaque area IT, for four different experimental datasets. At the time of testing, our highest performing model (Barlow Twins objective, λ = 0.2) had the 10th highest average predictivity for area IT out of approximately 250 publicly available models on the Brain-Score leaderboard. Across a reasonably large range of values of λ, the equivariant models improved the neural predictivity relative to the invariant baseline (λ = 0) for all four datasets (Fig. 3 ). Many previous publications have noted that changes in training objective function have a small effect on neural predictivity, relative to other factors such as training dataset [Tuckute et al., 2022 , Conwell et al., 2023 , Yerxa et al., 2024] . In contrast, encouraging equivariance produced much larger gains than choosing between different base invariant objectives: the range of predictivities over the sweep of λ was around 4 times larger than the range of predictivities over objective functions for the invariant baseline. We further contextualize the scale of predictivity improvements in Fig 4 by comparing models to all public submissions on the BrainScore leaderboard; our equivariant training procedure improves performance of the already-strong invariant models to nearly state-of-the-art levels of IT predictivity. By training the most predictive model for 1000 epochs (rather than 100), we achieved 0.5355 mean fraction of explained variance, which makes this the top IT brain prediction model. To ensure that the observed alignment increases are not architecture specific, we trained a smaller set of models using different backbone architectures and observed similar trends when using both smaller and larger networks (see Appendix A.7 for details).We quantified the correlation between our various representational measurements and the neural predictivity for each of the four electrophysiology datasets in Table 1 . We observed that the only representational metric with a correlation greater than 0.4 across all four neural datasets was the Spatial-Photometric distance, which is the metric most closely related to the factorization score described in [Lindsey and Issa, 2024] . While this previous study described a correlation between structured variability and neural predictivity measured from a large set of pre-trained models, our results demonstrate that explicitly encouraging such structures can improve alignment between artificial and biological representations. In addition to the previously described representational measurements, we also looked at the linear decoding of the hue modulation parameter in isolation. Hue modulation is one of 12 augmentation parameters that are linearly decoded in the parameter regression measures described in Section 3.2. We observed a strong correlation between neural predictivity and hue modulation, particularly with the Sanghavi-Jozwik dataset, which is the only response dataset that included color image stimuli (last column of Table 1 ).Several previous studies that aim to reduce augmentation-invariance of self-supervised features have reported that the resulting representations generalize better to out-of-distribution classification datasets [Gupta et al., 2023 , Xiao et al., 2020 , Suau et al., 2023 , Chavhan et al., 2022] . However most of these studies focused on using the smaller ImageNet-100 dataset for training, in one case reporting that the transfer learning gains diminish or disappear when using ImageNet-1k [Chavhan et al., 2022] . We tested our set of networks on 6 different downstream tasks and found limited evidence that the equivariant features confer an advantage in terms of out-of-distribution generalization when training on a sufficiently large and diverse dataset (see Table 2 ). To address this discrepancy with the literature we conducted additional experiments on networks trained using the ImageNet-100 dataset, and we do recover the improvement in generalization to diverse downstream tasks.It is also worth noting that CE-SSL trained networks do not outperform their invariant counterparts on in-distribution generalization (see A.2.1 and Fig. 5 ). This is not necessarily surprising in light of the fact that the suite of augmentations and architectures employed in SSL have been in some sense optimized by the community in order to improve performance on this task (by aligning the transformation invariance task with the standard in-distribution classification task). However, for out-of-distribution classification tasks where the task-alignment is worse, the equivariance task could mitigate this mis-match. A concrete example is the Flowers-102 dataset, where the color of petals is a much stronger predictor of class than color is in, say, the ImageNet-1k dataset (so the color insensitivity induced by the standard augmentations could be detrimental). For this dataset we do see marginal improvements, but note that the improvements are much more pronounced when pretraining on smaller datasets (ImageNet-100). There are at least 2 possible explanations for this:(1) for ImageNet-1k pretraining the performance of the networks is already quite high, and the task is saturated, or (2) there is a more fundamental reason that the improvements in transfer learning induced by equivariance decrease as the size and diversity of the pretraining dataset grows. Future work could explicitly disambiguate between these hypotheses to determine why the benefits of transformation-related variability for out of distribution generalization are outweighed by the gains of scaling the dataset. Furthermore this result shows that the increased neural predictivity we observe in ImageNet-1k trained networks cannot be explained by a need to perform better on a variety of invariant-classification tasks. 
relationship to existing augmentation-sensitive ssl methods: 
 A key feature that differentiates our approach is that it encourages equivariant structure without explicit access to augmentation parameters. This is enabled by the "paired augmentation" data generation procedure, and to the best of our knowledge CARE [Gupta et al., 2023] is the only Describable Textures Database [Cimpoi et al., 2014] , Flowers-102 [Nilsback and Zisserman, 2008] , and Food-101 [Bossard et al., 2014] . We closely follow the evaluation procedure from [Lee et al., 2021] (see Appendix A.4 for details) and report top1 accuracy for each objective/dataset. In all cases we report the mean over 5 runs of the evaluation procedure, we observed very little variability (maximum of .2%, over all evaluations, we report the standard deviation over runs in Appendix A.4. The equivariant networks are denoted by prepending a "CE" before the objective and were trained using λ = 0.1, which enabled a substantial amount of structure variability without significantly impacting frozen-linear classification on the SSL training dataset (see Appendix A.2). For ImageNet-1k trained networks out of distribution performance decreased for most evaluations, while for ImageNet-100 trained networks performance was improved in 15 of 18 cases.Objective existing work that shares this feature. Our method has two advantages over CARE: (1) in CARE the equivariance loss is applied in the same space as the invariance constraint, and because there are no "negative equivariant pairs" in the CARE framework, learning an invariant representation would perfectly satisfy the equivariant constraint; and (2) in CARE the standard augmentation pipeline is used to optimize the base L iSSL loss and paired augmentation are used in parallel to optimize the equivariance term (so an increased number of passes through the network is necessary relative to standard training).In EquiMod [Devillers and Lefort, 2023] the projector network is conditioned on the augmentation parameters by appending the parameters to the output of f . The authors theorize that knowledge of the augmentation parameters could allow the projector network to better extract an invariant subspace tailored to each transformation, thus allowing for more structured variability in the representation space. Alternatively, the projector network could simply ignore the augmentation parameters, resulting in a structure that is identical to invariant SSL. In practice [Garrido et al., 2023b] have found this to be the case. Split-Invariant-Equivariant (SIE) and Amortised Invariance (AI) learning [Garrido et al., 2023b ] each improve on this principle by using a hypernetwork approach: a separate network takes as inputs the augmentation parameters and outputs the parameters of either g or both f and g respectively. While the collapse issue of EquiMod is avoided, this comes at the expense of significantly complicating the network computation, and in the case of AI introduces new parameters that need to be tuned for every downstream task (when augmentation information is not available).Still other methods supplement the standard invariant SSL loss with an auxillary term that involves predicting the parameters of the input transformation [Lee et al., 2021 , Dangovski et al., 2021] .The relationship of our method to these is analogous to the relationship of transformation-invariant self-supervised learning to supervised classification. 
discussion: 
 We've developed a new objective for explicitly encouraging structured variability in SSL-trained networks, and demonstrated that it can produce increased alignment with neural responses in IT. While we are not the first to incorporate a notion of equivariance to self-supervised learning, our method improves on existing work in several ways: it require no extra passes through the network relative to invariance-based learning, it encourages diversity in the representation of transformation-related information by leveraging advances in invariance-based learning, and it does not rely on supervised access to transformation parameters. The parsimony of our approach (applying the same objective to both outputs of individual images and to displacements between similarly transformed images) allows our technique to be easily adapted to other settings such as temporal self-supervised learning (discussed below). 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.Our approach induced several interesting features in the learned representations: transformation variability is shared across base images and factorized with respect to variability over base images, the variability induced by distinct types of transformations are factorized from each other, there is increased alignment between class manifolds, and transformation related information is linearly encoded. Some of these properties are closely related to the imposed objective and some are emergent. We also confirmed that several of these representational properties are correlated with increased neural predictivity. Future work can go beyond these correlative observations to better understand how increasing transformation sensitivity improves neural alignment. For example, one could analyze the residuals of predicted neural firing rates of distinct models to determine how "overlapping" the variance predicted by each is (or alternatively, attempt to fit the residual variance of one model with another). Such analyses are quickly becoming more feasible as more large scale datasets of neural responses to natural images are collected and released (i.e. [Madan et al., 2024] ). We view this result as demonstrating the promise of incorporating knowledge gained from experimental observations and large scale comparative studies into optimization procedures to produce better models.Implicitly enforcing an additive equivariant structure with a contrastive objective does however lead to fewer guarantees regarding the learned structure. For example in schemes where the output transformations (T ρ 's) are explicitly represented or learned the resulting representation is "steerable" by default. It is of interest to investigate whether the output transformations could be reliably recovered from our learned representations. Additionally it would be interesting to consider other types of output transformations (CARE [Gupta et al., 2023] focuses on orthogonal transformations, and in the case where output transformations are learned they can be computed with nonlinear neural networks).Complementary work could explore the impact of training using more ecologically relevant sources of data, e.g., by replacing synthetically transformed views of images with temporally adjacent frames of natural videos. This approach is particularly appealing from the perspective of biological plausibility, as the pairing of training examples is readily available from natural visual experience. Several recent publications have shown that such a strategy can produce representations with competitive neural predictivity and performance on computer vision tasks [Zhuang et al., 2021 , Parthasarathy et al., 2023 , Venkataramanan et al., 2023] . In this context, the typical invariance loss can be thought of as incentivizing representational slowness [Földiák, 1991, Wiskott and Sejnowski, 2002] .The equivariance mechanism described in this work could be implemented by applying the same invariance-based loss function to the first temporal derivative of the responses, i.e. by encouraging the displacement between successive pairs of frames to be constant. Such a temporal-equivariance objective would instead incentivize representational straightness. Temporal straightness been used to describe features of both human perception and neural activity in the ventral stream [Hénaff et al., 2021 [Hénaff et al., , 2019]] , and straightness in artificial representations has been found to be correlated with both neural predictivity and adversarial robustness [Lindsey and Issa, 2024, Harrington et al., 2022] .Exploring this connection is a promising direction for future work.All code used for pretraining, evaluation, and analyses, will be made available in a public github repository upon publication.Here we report some additional hyperparameters not included in the main text.Optimization: For all experiments we trained for 100 epochs using a batch size of 2048 and used the LARS optimizer [You et al., 2017] with weight decay of 1e-6 and momentum of 0.9. Note that the L CE-SSL loss is evaluated on pairs of augmented views and thus had an effective batch size of 1024.We use a base learning rate of 4.8 and a learning rate schedule consisting of linear warm-up for the first 10 epochs followed by cosine decay throughout training.Projector Architectures Each trained network uses two projectors with matching architectures. For Barlow Twins we used the architecture proposed in the original work [Zbontar et al., 2021] (3 layer MLP with hidden layer and output layer widths of 8192). For MMCR we also used a 3 layer MLP with 8192 hidden width but 512 output units (also in line with the original work [Yerxa et al., 2024] ).For SimCLR we used the same projector architecture as MMCR, which is larger than the MLP described originally because subsequent work [Garrido et al., 2023a] has found that SimCLR benefits from a more expressive projector.For Barlow Twins we set the λ BT , which balances the on and off diagonal loss terms, hyperparameter to 5e -3. For SimCLR we used a temperature of τ = 0.15.Besides the change of dataset, the only hyperparameter change in this setting is that we increased the number of pretraining epochs from 100 to 200 to be more in line with previous work.Because frozen-linear evaluation on large datasets is computationally intensive we instead opt for online-linear classification. During pretraining the representation network outputs are detached from the gradient propogation graph and fed through a linear layer that is optimized with the standard supervised cross entropy loss. Previous work has shown that online-evaluation is very strongly correlated with frozen-linear evaluation and incurs only a minimal cost on top of self-supervised pretraining. We report the accuracies for ImageNet-1k trained networks in Fig. 5 and the smaller set of ImageNet-100 trained models in Table 3 . We closely follow the evaluation procedure from [Lee et al., 2021] , we repeat the details here for completeness. First images are resized such that the shortest edge is 224 pixels, then center cropped to 224x224 resolution. Then features are extracted from train, validation, and test splits of each dataset. L-BFGS is used to optimize the standard cross entropy loss with L 2 regularization, the value of the ridge parameter is swept over selected via performance on the validation set. Subsequently the linear classifier is retrained using both the train and validation sets, and we report the final accuracy on the held out test set. This classification procedure was run 5 times with different random initializations, we reported the mean performance in 3.4, and report the standard deviation over runs below.Table 4 : Standard deviation of top 1 accuracies over 5 independent runs of the transfer learning evaluation procedure.Cifar-10 Cifar-100 Pets DTD Flowers-102 Food-101MMCR 1e-2 7e-3 4e-2 1e-5 7e-2 2e-1 Barlow 1e-4 2e-4 2e-2 3e-1 3e-2 9e-3 SimCLR 1e-2 7e-3 1e-2 4e-2 7e-2 5e-3 CE-MMCR 1e-2 2e-2 4e-2 3e-2 1e-1 1e-2 CE-Barlow 5e-2 2e-2 3e-2 3e-2 8e-2 1e-2 CE-SimCLR 5e-3 2e-2 2e-2 2e-2 5e-2 1e-2Objective Cifar-10 Cifar-100 Pets DTD Flowers-102 Food-101 MMCR 1e-2 2e-2 8e-2 4e-2 2e-2 7e-3 Barlow 1e-2 1e-1 1e-1 4e-2 3e-2 3e-3 SimCLR 2e-2 7e-3 1e-4 4e-2 2e-2 8e-3 CE-MMCR 9e-3 2e-2 2e-1 2e-1 6e-2 2e-2 CE-Barlow 1e-5 1e-2 9e-2 1e-5 5e-2 9e-2 CE-SimCLR 1e-2 1e-2 2e-2 2e-2 7e-2 8e-3Below we depict the joint distributions of the distances described in 3.2. For each setting (unique objective function and value of λ there are 800 unique measurements (i.e. points of a unique hue on an individual plot). Meaning, for example, there are 800 random pairs of augmentation manifolds compared in both invariant representation space and equivariant representation space for each equivariant network in the left most column. We summarize describe the sources of variability over which covariance matrices are estimated and the variables over which the expected Bures distance is calculated for the experiments in Fig. 2 A-D .We aim to test whether learning equivariances using weak augmentations induces structured variability in response to stronger augmentations (the extent to which learned equivariances generalize beyond the range of transformations seen during training). We trained models using the Barlow Twins objective and the same sweep over values of λ using "weak" augmentations of (1) double the minimum crop size, (2) half the maximum value of color jittering, and (3) half the maximum size of Gaussian blurring kernel. We then repeat the parameter decoding experiments from the main paper Fig. 2E on these weak augmentations (left panel Fig. 7 ), and on the non-overlapping part of the parameter space between the weak and strong augmentation distributions, i.e only for augmentations whose parameters are in distribution for the models trained as in the main text but out of distribution for the new models trained using weaker augmentations (right panel Fig. 7 ). In the first panel, 7 : Augmentation parameter decoding performance on held-out test images for networks trained using either strong or weak transformations. In the left panel we plot the decoding performance for weak transformations only, and in the right panel the performance for strong transformations only (which are not seen by the weak-trained networks during pretraining).we can see that the best parameter decoding performance occurs when the pretraining distribution of transformations is matched to the evaluation transformations (i.e. weak-trained models slightly outperform the strong-trained models at decoding the parameters of weak transformations). In the right panel we see that models trained with strong augmentations have higher decoding performance, but models trained only on weak augmentations still demonstrate significantly increased ability to linearly decode strong augmentation parameters relative to the invariant trained models (λ = 0 models), indicating a degree of generalization in the learned equivariances.To verify that the effect of equivariance neural predictivity observed in the main text is not limited to a specific choice of architecture, we trained invariant (λ = 0.0) and equivariant (λ ∈ [0.1, 0.2]) networks using L BT with smaller (ResNet-34) and larger (ResNet-50) backbones. As shown in Fig. 8 , we observe the same trend across different architectures: contrastive-equivariant training increases alignment to area IT as measured by linear predictivity.Figure 8 : Neural predictivity results for models trained using the MMCR base loss and a ResNet-34 backbone (left panel), and the Barlow Twins base loss and a ResNet-101 backbone (left panel). We see similar trends in terms of increased neural predictivity for equivariant trained models as observed using the ResNet-50 backbone in the main text.All pretraining runs used 8 A100 Nvidia GPUs with 40GB of memory each. In our setting pretraining run times were around 15 hours, and we note that CE-SSL training generally increased training time by approximately 10% relative to standard self-supervised training. Subsequent evaluations ran on a single A100.We discuss some limitations not addressed in the discussion here. Our current training setup requires selection of the hyperparameter λ 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 λ. It is worth noting that some invariant SSL methods may be more or less sensitive to this additional hyperparameter. For example, some of the non-monotonicity of SimCLR curves in Fig. 2 may suggest that SimCLR representations are more difficult to smoothly shape within the CE-SSL paradigm.Additionally, compute limitations prevent us from doing an extensive architecture search over the two projector networks employed in contrastive equivariant training. In addition to varying the depth and width of each projector, it would also be interesting to "split" the represetnation space and have each projector only see some dimensions as input (as in SIE [Garrido et al., 2023b] ). Additonally compute limitations prevented us from extensively evaluating the variability in neural predicitivity over independent runs of contrastive equivariant training (the BrainScore framework recently stopped providing estimates of the error of neural predicitivity though in previous studies the reported error for IT predicitivity was ≈ 3e-3, which is small relative to the variability we observed across the parameter of interest λ). Pilot experiments indicated to us that the variability over training runs was small relative to the variability over values of lambda (we trained two invariant models and two with λ = 0.1 to get a rough estimate of this, in both cases we kept the more predictive model for inclusion in all analyses that appear in this paper).In this work we propose one strategy for inducing increased alignment between artificial and biological visual representations. Better understanding the computational principles underlying visual representations has the potential to benefit both the quality of computer vision applications and our ability to clinically treat disorders related to visual perception.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: Data sources are clearly documented and source code will be released in the form of a public repository upon publication. Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The broader impacts of the work are described in an appendix. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper describes a novel SSL approach incorporating constraints for learning less invariant and more equivariant representations. 

The reviewers unanimously found the proposal novel and worthy. One main limitation of the work is the seemingly limited improvements in OOD classification compared to current methods. Even in the case of brain predictability, one could argue that the improvements are marginal. Nonetheless, the work appears to be carried well, and there is unanimous agreement that the approach is moving the field in the right direction.

Hence, the AC recommends the paper to be accepted.

Reviews: review: paper_summary: paper_summary: The paper introduces a novel framework, CE-SSL, to address the limitations of traditional self-supervised learning (SSL) objectives, which often result in overly invariant network representations, with a goal to improve the neuronal plausibility of resulting representations. The authors propose a method that incorporates structured variability in response to input transformations, aligning the representations more closely with known features of visual perception and neural computation. The CE-SSL framework converts standard invariant SSL losses into contrastive-equivariant versions, encouraging the preservation of aspects of the input transformation without supervised access to transformation parameters. The method was validated through representational analyses, neural predictivity evaluations using the BrainScore pipeline, and various downstream tasks, demonstrating improved alignment with neural responses and increased structured variability in the representations.
strength_weakness: strength_weakness: Sensitivity to Hyperparameters: The performance of the method is sensitive to the choice of the hyperparameter λ, which controls the balance between invariant and equivariant loss terms, requiring careful tuning.
Limited Generalization: The method showed limited improvement in generalization to out-of-distribution tasks when trained on large datasets like ImageNet-1k, which could limit its applicability in diverse real-world scenarios.
Limited broader impact: The method improves biological plausibility of learned representations, but the broader impact to representation learning in AI/ML is not clear.
Presentation: Some concepts are insufficiently described, for example the projection network architecture is not descibed in sufficient detail. Overall, the organization of the paper is a bit convoluted and more work is needed to streamline the narrative.
questions: questions: How does the method perform with different backbone architectures other than ResNet-50? Have the authors considered testing with more recent architectures?
How robust is the method to changes in the choice of augmentations? Have the authors tested with different sets of augmentations?
What are the potential applications of the CE-SSL framework in other domains beyond visual perception, such as audio or text processing?
Have the authors considered integrating additional forms of self-supervision, such as clustering or reconstruction, to further enhance the learned representations?
limitations: limitations: The authors have adequately addressed the limitations of their method, including the computational overhead, sensitivity to hyperparameters, and limited out-of-distribution generalization. They have also provided constructive suggestions for future work to address these limitations, such as exploring temporal self-supervised learning and more ecologically relevant data sources. The authors' transparency in discussing these limitations is commendable and aligns with the best practices for responsible machine learning research.
paper_summary: paper_summary: In their paper ‘Contrastive-Equivariant Self-Supervised Learning Improves Alignment with Primate Visual Area IT’ the authors construct a new kind of contrastive-equivariant loss to optimise ResNet-50 architectures. Their new loss is based on the idea of using both self-supervised learning via object representations invariant across transformations of an image, and equivariant representations of transformation across images. They show that their new loss results in improved alignment to brain data recorded in IT. In a series of additional investigations, they show which part of the achieved representations is most likely causing the alignment to data. They also show that their method shows similar generalization advantages like prior alignment techniques but that this improvement is actually easily compensated for by training on a larger dataset.
strength_weakness: strength_weakness: I believe that the two weaknesses, one of experimental nature and one on the writing:

Experimental – 
The main goal of authors was to achieve higher alignment to brain data, but I found the comparison to other methods a bit thin. Authors say their method currently ranks 10th but it would be helpful to at least have a feeling for how far off the method is from 1st place. Many readers will not know whether the differences in the top 10 are in the magnitude of full percentage points or more / less. As such, I would suggest to perhaps at least provide the performance of the 11th and 9th, and 1st. This should be easy to address. 
I also wonder how ‘unique’ the achieved explainability of their method is, i.e. authors contrast their method with mostly task-optimised networks and would they expect that adding their method to task optimised networks would result in an overall best performant network or is the explainability captured by their method also the variance captured by task optimisation so that combining them does not promise any further improvements? I can see that from a biological plausibility perspective the newly proposed technique is still nicer, but if the classification loss achieves learning even more brain-like presentations which also contain a similar trade-off between invariance and equivariance, then perhaps there is multiple ways for the brain to learn such representations and the loss presented here would not achieve any unique advancements. 

Writing –
I found the section ‘3.2 Representational Analyses’ somewhat difficult to follow – I appreciate that authors with Figure 2 tried to provide visualisations of manifolds for different scenarios but perhaps it would be more helpful to visualise which cases are actually compared for which distance? Perhaps this could be done in the form of table in the appendix, if otherwise authors run out of space? From the text and current figures I struggled to exactly get which images / augmentations etc go into which set to construct the manifolds. Also, I believe the line plots at the bottom of Figure 2 would be significantly easier to understand if authors would use something like ‘Strength of equivariance loss (Lambda)’.
questions: questions: Apart from the suggestions mentioned in weaknesses, I have some minor comments: 
-	Some of the in-line references seem to be lacking the brackets in the proper locations, I noticed this in the references in line 86 and 227, where no overall bracket is around the reference but only around the year
-	In line 161 & 162 authors say ‘for each equivariant network’ and I am somewhat confused why there is a plural here? Does this mean the projection relating to each specific transformation?
-	In line 113 you mention a grey inset in Figure 2 but I am not sure what you are referring to with this?
limitations: limitations: The limitations discussed in the discussion and appendix seem appropriate.
paper_summary: paper_summary: Summary:

The authors propose a novel approach to self-supervised learning (SSL) by incorporating contrastive equivariant training into existing successful SSL methods based on creating invariance to input transformations. The authors observe that such an incorporation of equivariant contrastive learning produces the following downstream benefits: 1) Improved representational similarity to biological visual representations, and 2) Enhanced transfer learning performance on downstream tasks. Both these findings have been evaluated via experiments performed using the BrainScore benchmark and by computing linear probe classification accuracy of various SSL training methods on downstream image classification tasks.
strength_weakness: strength_weakness: - Overall, the improvements from using CE-SSL in the BrainScore performance is only quite marginal, hence raising the question of how impactful the current work will be in the broader scheme of methods to enhance neural predictivity.
- The writing quality of this work could be further improved. Particularly, Section 2 is currently written in a convoluted manner in my opinion and could be further improved to enhance readability. There are other minor writing issues, for e.g.:
    - Lines 145-146 and equation 2 are in conflict with each other. Throughout the paper, it looks like $\lambda=0$ refers to the invariant case. In that case, the invariant term of Eqn. 2 should be weighted by $(1-\lambda)$. Is that correct? Currently the invariant term is weighted by $\lambda$ in which case $\lambda=0$ would correspond to training only with CE-SSL.
    - Line 113, I believe the authors intended to refer to Figure 1 but they incorrectly refer to Figure 2 in the submission
    - Citation issue in Line 86, Zbontar et al. [2021].
    - It would help to use different icons to represent $f$ and $g_{inv}$ or  $g_{equi}$ as there are significant structural differences between the backbone and projector networks.

- In Figure 4, as the CE-SSL loss contribution increases, the ImageNet classification accuracy seems to be dropping. This is an issue that needs to be explicitly discussed more clearly.
questions: questions: Please refer to my review above.
limitations: limitations: Limitations are discussed adequately in this submission.
paper_summary: paper_summary: This paper presents an equivariant learning framework that modifies standard invariant-based self-supervised learning methods by integrating differences between object classes (Loss_iSSL) and the changes observed in images before and after identical transformations (Loss_CE-SSL). This adaptation enables the model to effectively incorporate transformation-related variability without necessitating additional transformation parameters or extensive modifications to the training protocol. The authors then demonstrate that this approach systematically enhances the model's capacity to predict neural activity in the inferotemporal (IT) cortex, revealing that optimizing for structured variability significantly boosts the accuracy of predicting cortical responses to natural images.
strength_weakness: strength_weakness: 1. The model's performance in transfer learning tasks, as shown in Table 2, does not demonstrate a significant improvement; in many cases, it merely maintains accuracy levels similar to those before implementation.
2. While the paper analyzes representations based on different values of $\lamda$, these analyses seem to offer little help in predicting neuronal activity or improving transfer learning. The proposed "tradeoff between invariance and structured variability" does not appear to manifest significantly.
3. The equivariance learned by the method may be limited by the types and varieties of transformations applied to the images, an aspect the paper does not analyze or discuss.
questions: questions: 1. The SimCLR curve in the lowest part of Figure 2 appears more jagged compared to the other two models. Could this indicate that SimCLR is less suitable for incorporating equivariance?
2. When predicting IT neuron activity, how can we interpret the variations caused by different $\lamda$ values across datasets?
3. In the transfer learning results, why does the addition of an equivariance loss seem to have minimal impact on the outcomes?
4. Is the equivariance learned by this method limited by the types and varieties of transformations applied to the images? For instance, if the model only learns transformations involving rotations between 0-30 degrees, can it represent equivariance for a 60-degree rotation? Furthermore, can the method effectively learn more complex, non-linear transformations, and can learning multiple transformations enable the model to represent combinations of these transformations equivariantly?
5. There appears to be a typographical error in line 225: “($\lamda$ = 0) for all four datasets (Fig. 3.3),” as the figure number seems incorrect.
limitations: limitations: Yes, the paper concludes with a discussion on some limitations of the model and directions for future research.

Comments: title: Response
Reply
comment: Thank you for taking the time to consider our rebuttal experiments. We are hard at work on the rotation case! Preliminary results suggest a trend similar to that of Figure R1.
Thank you we will improve our score.
Thank you for your response. I really like the results presented in Figure R1 and am looking forward to seeing the attempts involving rotations. These results have convinced me well, and I will increase my score.