Connecting Joint-Embedding Predictive Architecture with Contrastive Self-supervised Learning

Authors: shentong  mo
shengbang  tong
  cmu
  nyu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 05:38:23

Abstract: In recent advancements in unsupervised visual representation learning, the Joint-Embedding Predictive Architecture (JEPA) has emerged as a significant method for extracting visual features from unlabeled imagery through an innovative masking strategy. Despite its success, two primary limitations have been identified: the inefficacy of Exponential Moving Average (EMA) from I-JEPA in preventing entire collapse and the inadequacy of I-JEPA prediction in accurately learning the mean of patch representations. Addressing these challenges, this study introduces a novel framework, namely C-JEPA (Contrastive-JEPA), which integrates the Image-based Joint-Embedding Predictive Architecture with the Variance-Invariance-Covariance Regularization (VICReg) strategy. This integration is designed to effectively learn the variance/covariance for preventing entire collapse and ensuring invariance in the mean of augmented views, thereby overcoming the identified limitations. Through empirical and theoretical evaluations, our work demonstrates that C-JEPA significantly enhances the stability and quality of visual representation learning. When pre-trained on the ImageNet-1K dataset, C-JEPA exhibits rapid and improved convergence in both linear probing and fine-tuning performance metrics.

Main paper: 
 introduction: 
 Unsupervised learning of visual representations has recently seen remarkable progress, primarily due to the development of innovative architectures and strategies that exploit unlabeled imagery. Among these advancements, the Joint-Embedding Predictive Architecture (JEPA) [1, 2, 3] has distinguished itself as a powerful approach. I-JEPA [2] leverages a masking strategy to extract visual features, facilitating significant strides in understanding and utilizing unlabeled visual data. However, despite its successes, certain limitations within the JEPA framework have become apparent, particularly concerning its components I-JEPA Exponential Moving Average (EMA) and I-JEPA prediction capabilities. Specifically, I-JEPA EMA has been found to be inadequate in preventing the issue of entire collapse [4, 5] , while the I-JEPA prediction mechanism struggles to accurately learn the mean of patch representations. These challenges not only hinder the performance of JEPA but also limit its applicability in broader contexts. 
related work: 
 Self-supervised Learning. In the self-supervised literature, researchers aim to exploit the internal characteristics of data and leverage pretext tasks to train a model. Recently, an unsupervised framework that learns effective views with data augmentation was proposed by Tian et al. [7] to reduce the mutual information between views. CMC [8] introduced a multi-view contrastive learning framework with any number of views to learn view-agnostic representations. Another pretext task of solving jigsaw puzzles was developed in PIRL [9] to improve the semantic quality of learned image representations, achieving better object detection results than supervised pre-training. More recently, Masked image modeling (MIM) has been explored in many previous works [10, 11, 12, 13, 14] to reconstruct the masked image patch given the unmasked counterpart as clues. Some MIM approaches [10, 11, 12, 15, 16] designed customized masking strategies (i.e., random, block-wise) as pre-text tasks during pre-training. For example, block-wise masking was introduced in BEiT [10] to learn transferrable visual representations by recovering discrete tokens of masked image patches. Given features extracted from the 25% unmasked patches, the seminal work, MAE [12] directly reconstructed missing pixels of 75% masked patches. SimMIM [14] randomly masked the input image with a large square patch size (i.e., 32 ) and used a one-layer prediction layer after the encoder to predict RGB values of raw pixels. Other researchers [15, 16, 17, 18] started to leverage a teacher network like CLIP [19] or adversarial learning to generate the mask and supervision target.Contrastive Learning. In the past years, contrastive learning has shown its effectiveness in selfsupervised learning, where various instance-wise contrastive learning frameworks [20, 21, 22, 23, 24, 25, 26, 27, 4] and prototype-level contrastive methods [28, 29, 30, 31] were proposed. The general idea of the instance-wise contrastive learning is to close the distance of the embedding of different views from the same instance while pushing embeddings of views from different instances away. One common way is to use a large batch size to accumulate positive and negative pairs in the same batch. For instance, Chen et al. [20] proposed a simple framework with a learnable nonlinear projection head and a large batch size to improve the quality of the pre-trained representations. To make the best use of a large amount of unlabelled data, they present a bigger unsupervised pre-training network and introduce distillation with unlabeled data in SimCLR v2 [21] to improve the performance in downstream tasks. Without involving negative instances, BOYL [22] trains the online network from an augmented view of an image to predict the target network representation of the same image under a different augmented view (positive instance). Another broadly-used approach [23] in the ICL literature is to apply a momentum encoder to update negative instances from a large and consistent dictionary on the fly. The dynamic dictionary was used with a moving-averaged encoder in MoCo series [24, 23] to build a dynamic dictionary to update negative instances in a queue of considerable size. The Variance-Invariance-Covariance (VICReg) [6] regularization strategy has been proposed to address shortcomings in self-supervised learning by enforcing stability through variance and covariance constraints. Joint-Embedding Predictive Architectures. The Joint-Embedding Predictive Architecture (JEPA) framework has shown considerable promise by enabling the prediction of embeddings from a compatible signal [1] . Building on this, the Image-based Joint-Embedding Predictive Architecture (I-JEPA) introduced by Assran et al. [2] utilizes a context encoder and a target encoder, together with a predictor, to manage representations under a masking regime. This approach focuses on learning embedding spaces directly, contrasting with other methods that operate on pixel or token spaces. While I-JEPA has been successful, it faces challenges such as the inefficacy of Exponential Moving Average (EMA) strategies in preventing model collapse and inaccuracies in predicting the mean of patch representations. These limitations hinder the stability and efficacy of the learned representations, prompting a need for enhanced methodologies. By maintaining sufficient variance in the embeddings and minimizing redundancy among features, VICReg supports robust feature learning. To address the limitations of I-JEPA, our work integrates VICReg into the JEPA framework, thereby aiming to prevent total model collapse and ensure invariance across different views of the same image.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: See our discussion in the limitation section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methodology: 
 In this section, we present a novel masked modeling framework designed for the joint-embedding predictive architecture to avoid entire collapsing and improve the mean of patch representations. Our key idea is to integrate VICReg into the JEPA framework for alleviating entire model collapse and improving the invariance across different views of the same image. We first provide preliminaries in Section 3.1, then provide the theoretical and empirical connection between I-JEPA and SimSiam in Section 3.2, and finally connect I-JEPA with VICReg in Section 3.3 to show the benefit of reducing entire collapsing and learning the mean of patch representations.In this section, we first describe the problem setup and notations, and then revisit I-JEPA [2] , SimSiam [4] , and VICReg [6] for self-supervised image modeling.Problem Setup and Notations. Given a dataset X = {x i } N i=1 with images x i ∈ R c×h×w , our goal is to learn a neural network f θ (•) to extract unsupervised representations from these visual samples.Revisit I-JEPA [2] . Taking the self-supervised modeling as a joint-embedding predictive architecture, I-JEPA [2] utilized f θ (•) as a context encoder, a pair of neural network f ′ θ (•) as a target encoder. A predictor g θ (•) is applied to predict the target representations from M masked block patches b y (1), ..., b y (M ). For a target block b yi corresponding to a target mask B i , the predictor g θ (•, •) takes as input the output of the context encoder b x and a mask token for each patch to predict {m j } j∈Bi , and outputs the patch-level prediction{ byj } j∈Bi , that is, { byj } j∈Bi = g θ ({m j } j∈Bi ).The masking objective is optimized by the average ℓ 2 distance between the predicted patch-level representations byj and the target patch-level representation b yj , which is formulated as:EQUATIONwhere |M | denotes the total number of target blocks M , and B i is the generated mask corresponding to the i-th target block during training.Revisit SimSiam [4] . For contrastive self-supervised learning, SimSiam [4] tried to obtain two augmented views x i and x ′ i . Then they fed two views into a pair of neural networks f θ (•) and f ′ θ (•). To learn the invariance of two different views from the same image, a prediction MLP head, denoted as g, transforms the output z of one view and matches it to representations p of the other view. The overall objective of SimSiam is to minimize the distance between z and p for all random patches r j in the same image, which is formulated as:EQUATIONwhere |V | denotes the total number of augmented views V , and P i is the all random patches corresponding to the i-th view during training.Revisit VICReg [6] . To prevent a collapse in which the encoders produce constant or non-informative vectors in Siamese net architecture, VICReg [6] introduced variance and covariance regularization terms based on the representations space in the invariance terms. Firstly, the variance regularization term v as a hinge function on the standard deviation of the embeddings z along the batch dimension n to prevent collapse with all the embeddings mapped on the same vector, which is denoted as:EQUATIONwhere γ = 1 denotes a constant value for the standard deviation in their experiments, and ϵ denotes a small numerical scalar to stabilize training. Secondly, the covariance regularization term c minimizes the sum of the squared off-diagonal coefficients to encourage the off-diagonal coefficients of correlation matrix C along the batch dimension to be close to 0, which is formulated as:EQUATIONwhere z = 1 n n i=1 z i . This term is helpful to prevent different dimensions of the embeddings from encoding trivial information. Finally, the invariance term is defined to minimize the mean-squared Euclidean distance between each pair of vectors z and z rj from different views.In the realm of self-supervised learning, both I-JEPA [2] and SimSiam [4] frameworks aim to extract robust representations from images, yet they approach the problem with distinct architectures and objectives. Here, we draw theoretical and empirical parallels between these methodologies, leveraging their unique approaches to enhance the understanding of joint-embedding architectures.Regarding theoretical connections, I-JEPA [2] uses a predictive model where the encoder and predictor work together to forecast masked parts of the input, relying on partial views. Conversely, SimSiam [4] operates without explicit masking, utilizing dual augmentations of the same image to enforce consistency between the independently processed views. Both models share the underlying principle of minimizing the distance between certain representations. As shown in Figure 2 , I-JEPA focuses on the distance between predicted and actual masked patch representations, whereas SimSiam minimizes the distance between the two augmented views of an image, thereby encouraging consistency across different transformations of the same data.Meanwhile, empirical studies such as Tian et al. [5] have shown that linear predictors in frameworks like BYOL [22] , which are closely related to SimSiam, tend to learn alignment with the correlation matrix of representations. This suggests that SimSiam's approach without a separate predictor could inherently align with the mean of augmented views, a concept central to I-JEPA's strategy of predicting masked patches. A study [32] in joint-embedding self-supervised learning also indicates that these methods primarily focus on capturing the co-occurrence and distribution of image patches, which aligns with I-JEPA's objective of learning from masked representations.VICReg [6] introduces variance and covariance regularization to prevent the collapse of representations in Siamese networks, ensuring that the model learns informative and diverse features across different dimensions. By integrating VICReg into I-JEPA, our objective is to tackle the common challenges of model collapse and enhance the mean representation learning of image patches.VICReg's variance regularization ensures that all dimensions of the embedding space contain meaningful variance, which is crucial for preventing the model from collapsing to trivial solutions. I-JEPA, which aims to learn diverse patch representations, can benefit from such a mechanism to ensure that each masked patch contributes informatively to the overall representation. By minimizing the off-diagonal elements of the covariance matrix, VICReg encourages the features to be uncorrelated, enhancing the diversity of the learned features. This aspect can be particularly beneficial for I-JEPA, where diverse patch predictions are essential for effective representation learning. Literature on non-contrastive self-supervised learning [33] suggests that implicit variance regularization, as seen in VICReg, can facilitate the learning dynamics in joint-embedding architectures. Such regularization helps maintain a balance between similarity and diversity in learned representations, which is crucial for effective self-supervised learning.In the following, we consider linear predictor W P ∈ R M ×M in I-JPEA, the masking objective with the average ℓ 2 distance between the predicted patch-level representations z yj and the target patch-level representations z a yj asEQUATIONwhere |M | denotes the total number of target blocks M , and B i is the generated mask corresponding to the i-th target block during training. By connecting I-JEPA with SimSiam in Section 3.2, we can can diagonalize correlation matrix of z yj over R: C zy j = U D C U T , where U is an orthogonal matrix whose columns are the eigenvectors of C zy j and D C is the real-valued diagonal matrix of the eigenvalues s k , k ∈ [1, K]. Given this eigendecomposition, the predictor is directly set to W P = U D α C U T , where α is a positive constant exponent applied element-wise to D C . The eigenvalues λ k of the predictor matrix W P are then λ k = s α k . Assume ẑyj the representations expressed in the predictor's eigenbasis, the asymmetric loss L can be formulated as:EQUATIONFollowing non-contrastive self-supervised learning [33] , we can use NTK [34, 35] to characterize the learning gradient dynamics of neural networks as ∇ zy j L = (Dz yj -z t yj )D, and the representational dynamics of each mode k independently follow gradient of the loss -∇ ẑy j L, and decouple as K independent differential equations:EQUATIONBy taking the expectation over the same masking blocks, we can have the dynamics for each mask patch y j asEQUATIONNote that when λ k < 1, dẑ yj ,k has the same sign as ẑyj,k and they have opposite sign at λ k > 1.These convergent dynamics will push an eigenvalue λ k of one to prevent the collapse of each mode k in representation dynamics. When removing the stop-grad operator in L, we can havedẑ y j ,k dt = -η(1 -λ k ) 2 ẑyj,k, where dẑ yj ,k and ẑyj,k always have opposite signs and the learned representations will become zero with exponentially decaying eigenvalues to collapse. Without the predictor W P , we will have dẑ y j ,k dt = 0, and the representations will not be updated.Incorporating VICReg's regularization strategies into I-JEPA could prevent the entire collapsing of the model, especially when learning from a large and diverse dataset. This integration could also improve the granularity and utility of the patch-level representations by ensuring that each dimension of the embedding space remains informative and independent. Overall, the integration of VICReg and SimSiam principles into the I-JEPA framework offers promising avenues for enhancing the robustness and efficacy of self-supervised learning models, particularly in tasks requiring nuanced understanding of complex visual content.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] • 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: See our discussions in the broader impact section 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 
experimental setup: 
 Datasets. Following previous methods [12, 2] , we use ImageNet-1K [36] for image classification, MS-COCO [37] for object detection and instance segmentation, and ADE20K [38, 39] for semantic segmentation. We closely follow previous work [40, 41] , and adopt the Mask R-CNN [42] as the detector. The ViTs [43] backbone weights are initialized with weights pre-trained on ImageNet-1K using our C-JEPA. Following the settings in [12] , we use the Semantic FPN and UPerNet approach [44] based on our ImageNet-1K pre-trained ViTs for evaluation. For a fair comparison, we fine-tune the detector with the same learning rate in [12] . For video object segmentation, we use DAVIS-2017 dataset containing 60 training, 30 validation, and 60 testing videos. For low-level tasks, we follow the previous work [2] and use Clevr/Count and Clevr/Dist on Clevr [45] dataset.Evaluation Metrics. For image classification, we follow previous masked image modeling methods [12, 2] to report the classification accuracy of linear probing and fine-tuning. For object detection and instance segmentation on MS-COCO, we apply AP box and AP mask as metrics for the bounding boxes and the instance masks. mIoU results are reported to evaluate semantic segmentation on ADE20K. For video object segmentation on DAVIS-2017, we use Jabri-based (J &F) m , J m , F m as metrics to evaluate the quality of frozen representations of image patches by segmenting scenes with the nearest neighbor between consecutive frames. For object counting and depth prediction tasks on Clevr, we use object counting and depth prediction to evaluate the linear probing performance of our model.Implementation. For input images, we resized the resolution to 224 × 224, i.e., H = W = 224. Following prior work [12, 2] , we apply a patch size of 16, i.e., P = 16. We use the tiny, small, base, and large models of ViT [43] architecture for experiments. We set the embedding dimension of the predictor to 384, and keep the number of self-attention heads the same as the backbone context-encoder. For the ViT-T/16, ViT-S/16, and ViT-B/16 context-encoder, we set the depth of the predictor as 6. For ViT-L/16 context-encoders, we set the depth of the predictor to 12. Following I-JEPA [2] , we use AdamW to optimize the context-encoder and predictor weights. We train our model using the default batch size of 2048, and the learning rate linearly increased from 1e-4 to 1e-3 during the first 15 epochs of pre-training, and decay to 1e-6 following a cosine schedule. The weight decay is linearly increased from 0.04 to 0.4, and the target-encoder weights are initialized the same as the context-encoder weights, and updated via an exponential moving average. We use a momentum value of 0.996, and linearly increase this value to 1.0. For masking, we use the same strategy and settings as I-JEPA [2] for 4 possibly overlapping target block masks.In this work, we propose a novel and effective framework for connecting Joint-Embedding Predictive Architecture with VICReg in non-contrastive self-supervised learning. In order to demonstrate the effectiveness of the proposed C-JEPA, we comprehensively compare it to previous baselines [41, 10, 12, 46, 47, 2] on non-contrastive self-supervised learning and mask image modeling.ImageNet-1K image classification. For image classification on the ImageNet-1K benchmark, we report the quantitative comparison results on linear evaluation and fine-tuning results in Table 1 . We can observe that our C-JEPA achieves the best results in terms of ViT-B/16 compared to previous masked image modeling approaches. Specifically, the proposed method outperforms the I-JEPA by 0.8 and 1.0 in terms of top-1 accuracy on linear evaluation and fine-tuning settings. Meanwhile, we enjoy the advantage of fewer pre-training epochs compared to other masked image model frameworks.COCO object detection and instance segmentation. For COCO object detection & instance segmentation benchmarks, we also report the quantitative comparison results in Table 1 . As can be seen, we achieve significant performance gains of 0.8@AP box and 0.8@AP mask on COCO object detection and instance segmentation compared to I-JEPA. We also achieve better results than DINO [41] and MAE [12] regarding both settings.ADE20K semantic segmentation. For the ADE20K semantic segmentation, we report the quantitative comparison results in Table 1 . As can be seen, the proposed C-JEPA outperforms I-JEPA [2] , the current image-based joint-embedding predictive architecture by 1.1@mIoU. Also, we observed that our C-JEPA can achieve better performance than other masked image modeling baselines, including MAE [12] and data2vec [47] . These improvements suggest the importance of leveraging the VICReg's regularization to capture better semantics for dense prediction tasks.DAVIS video object segmentation. In addition, we scale up our experiments from ViT-B to ViT-L and report the results in Table 2 . The proposed C-JEPA still achieves better performance than I-JEPA [2] in terms of linear evaluation and fine-tuning settings. For DAVIS video object segmentation, our method also shows significant and consistent gains as shown in Table 2 . Compared to I-JEPA, ours achieved the results gains of 1.7@(J &F) m on ViT-L/16.Clevr object counting and depth prediction. We further measure the abilities of object-counting and depth prediction on Clevr benchmarks for our pre-trained models. Table 2 shows linear evaluation results of C-JEPA on the Clevr counting and depth benchmarks. Compared to I-JEPA [2] , we achieve the results gains of 1.2@Clevr/Count and 0.4@Clevr/Dist using ViT-L/16. These results further indicate the benefit of the proposed method in learning better representations than our vanilla I-JEPA baseline without VICReg regularization.In this section, we performed ablation studies to demonstrate the benefit of the proposed Variance/Covariance and Invariance terms from VICReg. Here, we conducted extensive experiments on ImageNet-1k pre-trained ViT-B/16 to explore the impact of representation mean and collapse, and learned meaningful qualitative patch-level representations. Variance/Covariance and Invariance in VICReg. We first conducted an ablation study to determine the efficacy of incorporating Variance and Covariance terms, along with the Invariance term, in preventing model collapse and improving the fidelity of representation means. These experiments are crucial for understanding how each component contributes to the overall stability and effectiveness of the learned embeddings. The results, as presented in Tables 3 and 4 , provide clear insights into the performance enhancements driven by these modifications. The inclusion of the Variance and Covariance terms significantly accelerated the training process, as evidenced by the metrics in Table 3 . Table 4 highlights improved accuracy and stability in the embeddings, demonstrating the beneficial impact of the Invariance term in aligning the representation means across different mask blocks of the same images.Representation Mean and Collapse. Further analysis focused on how well the C-JEPA framework mitigates the risk of model collapse and accurately learns the mean of patch representations: The integration of VICReg's terms with C-JEPA's architecture addresses the previously noted deficiencies in I-JEPA's EMA component, which was prone to collapsing. The results (referenced in the last rows of Tables 3 and 4 ) underscore the robustness added by these regularization terms from VICReg.Qualitative Attention Visualizations. To complement our quantitative findings, we also examined qualitative aspects through attention visualization techniques. Figures 3 and ? ? showcase the attention maps generated by both the base I-JEPA and the enhanced C-JEPA models. These visualizations illustrate the more focused and contextually relevant attention in C-JEPA, which correlates with the theoretical improvements expected from incorporating VICReg's terms. The attention maps further validate our approach by visually demonstrating the enhanced capability of C-JEPA to maintain stable and meaningful patch-level representations across various image contexts, an improvement over the base I-JEPA model. This qualitative evidence supports the quantitative improvements and highlights the benefits of our approach in capturing invariance. 
conclusion: 
 In this work, we introduced C-JEPA, a novel enhancement to the Joint-Embedding Predictive Architecture that integrates the robust features of Variance-Invariance-Covariance Regularization (VICReg) to address critical limitations in the existing I-JEPA model. By refining the EMA and prediction components of I-JEPA, C-JEPA successfully prevents entire model collapse and more effectively learns the mean of patch representations, thereby advancing the state of unsupervised visual representation learning. Our theoretical analysis and empirical results have demonstrated the effectiveness of C-JEPA, particularly when pre-trained on the ImageNet-1K dataset. The framework shows marked improvements in convergence rates and performance in both linear probing and finetuning scenarios, outperforming existing methods. This underscores C-JEPA's capability to not only address the weaknesses of its predecessor but also to provide substantial improvements in learning quality and stability.  
 
 Appendix: Our C-JEPA's adoption of VICReg principles enhances its versatility and applicability across various unsupervised learning contexts, making it a valuable tool for tasks requiring robust and reliable visual representations. The implications of this research are significant, offering a pathway for future studies to explore and expand upon the integration of contrastive and joint-embedding techniques. As we move forward, further refinements and explorations into the scalability of C-JEPA in more diverse and challenging datasets will be crucial. Additionally, investigating the adaptability of C-JEPA to different domains and its effectiveness in real-world applications will be essential to fully realize its potential. C-JEPA represents a significant step forward in the field of machine learning, particularly in the unsupervised learning of visual representations. It sets a new benchmark for future research and development in this area, promising enhanced performance and broader applicability in tackling complex visual understanding tasks.to the empirical NTKΘ(Y, Y) in the original basis. Each block of the full empirical NTK is denoted as Θt (y j , y q ) = (y T j y q )I M , where y j , y q denote the high-dimensional inputs from mask patches j and q in the training samples. I M ∈ R M ×M is the identity. Based on this and the central limit theorem for an i.i.d standard Gaussian distribution, we can formulate the representational dynamics for each mask patch y j under the L asEQUATIONwhere D t is a diagonal matrix composed of the eigenvalues λ k . For kth entry of dẑ y j ,k dt , we can formulate the representational dynamics under the L asEQUATIONSince the expectation over the same masking blocks is ẑyj,k , we can have the dynamics for each mask patch y j asEQUATIONB Representation learning dynamics without the stop-grad and predictor B.1 Learning dynamics of the loss L without the stop-gradRecall we consider linear predictor W P ∈ R M ×M in I-JPEA, the masking objective with the average ℓ 2 distance between the predicted patch-level representations z yj and the target patch-level representations z a yj asEQUATIONWhen removing the stop-grad operator in L, we can have the masking objective asEQUATION)The representational dynamics for each mask patch y j under the L is formulated asEQUATIONSince the expectation over the same masking blocks is ẑyj,k , we can have the dynamics for each mask patch y j asEQUATIONwhere dẑ yj ,k and ẑyj,k always have opposite signs and the learned representations will become zero with exponentially decaying eigenvalues to collapse.Without the predictor W P , we can have the masking objective asEQUATIONIn this case, we have λ k = 1 for all independent eigenvalues in the eigenbasis. Based on the representation dynamics for each mask patch y j as dẑ yj ,k dt = ηλ k (1 -λ k )ẑ yj ,k , we will have dẑ yj ,k dt = 0, and thus the representations will not be updated in the end.In this section, we provide detailed algorithmic and implementation aspects of the Contrastive-JEPA (C-JEPA) framework, elaborating on the loss functions, pseudo-code, and experimental setup used throughout our studies. These details aim to enable reproducibility and deeper understanding of the modifications introduced in the C-JEPA compared to its predecessors.The loss function for C-JEPA is a combination of the original JEPA loss and the VICReg-inspired regularization terms, structured as follows:EQUATION)where βvicreg = 0.001 provides a scaling factor to balance the original JEPA loss with the new regularization terms introduced from VICReg. The VICReg loss component is defined as:EQUATION)with β sim = 25, β std = 25, and β cov = 1, which scale the similarity, standard deviation, and covariance regularization terms respectively, integral to maintaining the diversity and stability of the learned representations.The following pseudo-code outlines the batch-wise application of VICReg within the C-JEPA framework, providing clarity on the computational steps involved during training: loss ← 0 10:num_masks ← SHAPE(zc) [1] 11:for i ← 0 to num_masks -1 do 12:for j ← 0 to num_masks -1 do 13:if i ̸ = j then 14:zc i ← zc[:, i, :] 15:zc j ← zc[:, j, :]if 'vicreg' in args then 17:zc_gather ← FULLGATHERLAYER(zc)18:zc_gather i ← zc_gather[:, i, :] 19:zc_gather j ← zc_gather[:, j, :]loss ← loss + CRITERION(zc i , zc j , zc_gather i , zc_gather j )loss ← loss -CRITERION(zc i , zc j )end forend for 27:loss ← loss/(num_masks × (num_masks -1))28:return loss 29: end function For a comprehensive implementation, we detail the specific settings used during our experiments:• Image Preprocessing: All input images were resized to a resolution of 224 × 224 pixels.• Patch Size: Following prior work, a patch size of 16 × 16 pixels was used.Table 9 : ImageNet-1K image classification. We perform a linear evaluation on pre-trained ViT-T/16 and ViT-S/16 models for image classification on ImageNet-1K benchmark. We report the Top-1 accuracies using knn, linprob and fine-tune settings to evaluate the quality of pre-trained representations. The best results are indicated in bold. • Vision Transformer Configurations: We utilized tiny, small, base, and large models of the ViT architecture, with specific settings for each model scale provided in Tables 5,• Optimizer: The AdamW optimizer was used for training, with specific learning rate schedules and weight decay settings described.• Batch Size and Training Schedule: A default batch size of 2048 was used, with a learning rate and momentum adjustment schedule to optimize convergence.To validate the efficacy of C-JEPA, we employed several benchmark datasets and tasks:• Image Classification on ImageNet-1K [36] • Object Detection and Instance Segmentation on MS-COCO [37] • Semantic Segmentation on ADE20K [38, 39] • Video Object Segmentation on DAVIS-2017 [50] • Low-Level Tasks on Clevr Dataset [45] For each task, we adopted established methodologies and frameworks, such as Mask R-CNN [42] for object detection and Semantic FPN and UPerNet [44] for semantic segmentation, ensuring that our findings were comparable with state-of-the-art results.In this part, we provide additional experimental analyses to validate the effectiveness of the proposed C-JEPA framework further. These experiments encompass various downstream tasks, ablation studies, and the impact of different components and coefficients in the model, specifically focusing on the Variance, Covariance, and Invariance terms from VICReg across different configurations of Vision Transformers.Downstream Results on ViT-T and ViT-S. We evaluated the performance of pre-trained ViT-T/16 and ViT-S/16 models across a series of downstream tasks to assess the quality of the representations learned by the C-JEPA framework. For image classification, we performed linear evaluations using knn, linear probing (linprob), and fine-tuning methods: The results of top-1 accuracies, presented in Table 9 , indicate the robustness and quality of the features extracted by our pre-trained models under different evaluation settings. For object detection and instance segmentation, we fine-tuned the pre-trained models to perform tasks on the COCO dataset. The average precision for bounding box detection (AP box ) and mask segmentation (AP mask ) are reported in Table 10 , demonstrating the applicability of C-JEPA pre-trained models in complex vision tasks. The models were also finetuned for semantic segmentation tasks, and metrics such as mean Intersection over Union (mIoU), average class accuracy (aAcc), and mean accuracy (mAcc) are detailed in Table 11 , highlighting the effectiveness of our approach in handling pixel-level classification tasks.Variance/Covariance and Invariance in VICReg. To dissect the contributions of individual VICReg components within the C-JEPA, we conducted focused ablation studies using the ViT-T/16 model. For faster convergence, initial ablations in Table 12 analyze the impact of the Variance/Covariance and Invariance modules on speeding up the model training. Further studies on better convergence, shown in Table 13 , investigate how these modules contribute to achieving better training dynamics and model performance.VICReg and Invariance Coefficients. For VICReg coefficients β vicreg , we varied the coefficients to understand their influence on model training and final performance, with results compiled in Table 14 .The impact of varying the invariance coefficient β sim is explored to optimize the balance between similarity and variance, with findings in Table 15 .These experiments underline the nuanced contributions of the VICReg components to the overall efficacy of the C-JEPA framework. By systematically varying key parameters and assessing their impacts across a broad spectrum of tasks, we establish a comprehensive understanding of the model's behavior and performance characteristics. This rigorous empirical foundation supports the theoretical enhancements proposed in C-JEPA, affirming its potential to advance the state-of-the-art in unsupervised visual representation learning.In order to provide a comprehensive understanding of the advancements offered by the C-JEPA framework, we include a detailed visual analysis alongside our quantitative evaluations. This section of the appendix is dedicated to showcasing attention visualizations that illustrate the practical impacts of the theoretical improvements achieved through the integration of VICReg components within the C-JEPA framework.To better understand how the C-JEPA model processes visual information, we utilize attention visualization techniques. These techniques highlight the model's ability to focus on relevant features within an image, which is crucial for tasks involving fine-grained recognition and localization. The following figures 4, 5, 6, 7, and 8 provide a visual comparison between the base I-JEPA and the Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: Summary: The paper presents an improvement to I-JEPA to improve representation learning quality. The authors observe that I-JEPA has limitations in terms of model collapse and learns mean patch representations. The authors connect I-JEPA to VICReg and SimSiam to address both these drawbacks respectively. The method is tested on ImageNet.

Positives:
1. Clearly written paper with a clear motivation and method section. It was easy to follow the problem that the authors were trying to solve and how they solved it.
2. Strong ablation studies strengthen the main claims in this work
3. The final method improves over the recent I-JEPA work and the authors demonstrate this improvement across multiple model sizes, and downstream tasks.
4. The method holds gains even after longer training (800 and 1200 epochs) which shows this technique isn't merely accelerating convergence. Please add this to the paper appendix.

Negatives:
1. Missing comparisons to stronger baselines like DinoV2. This would admittedly require training on much larger pretraining datasets.

Reviewers:
The reviewers are positive about this work and maintain the positive ratings after the author rebuttal.

Decision:
Given the reviewer recommendations and results of this paper, it should be accepted at NeurIPS. 
Further, given that it shows strong empirical fixes to I-JEPA this work should be made a spotlight.
However, since the proposed improvements are limited to a single method (rather than all non-contrastive methods like BYOL, DINO), the AC does not recommend it to be an oral.

Reviews: review: paper_summary: paper_summary: The paper presents C-JEPA, a novel framework integrating VICReg into the Image Joint-Embedding Predictive Architecture (I-JEPA) to address its limitations in preventing model collapse and learning mean patch representations. Empirical and theoretical evaluations demonstrate that C-JEPA enhances the stability and quality of visual representation learning, showing superior performance across multiple benchmarks.
strength_weakness: strength_weakness: The additional regularization terms may lead to increased computational overhead.
Further testing on larger and more diverse datasets is needed to confirm scalability and generalization.
While the results are promising, testing across more varied domains would strengthen the conclusions.d 
The new combination of established SSL methods (I-JEPA and VICReg) comprises an incremental development in hte fiels.
questions: questions: Can you provide more details on the computational overhead introduced by the VICReg integration?
Have you considered testing C-JEPA on more diverse datasets beyond ImageNet-1K and more diverse downstream tasks?
limitations: limitations: The authors have acknowledged the potential computational complexity and the need for further testing on diverse datasets. They have made significant efforts to address the limitations of C-JEPA, and their proposed future work aims to explore these aspects further.
paper_summary: paper_summary: The paper introduces C-JEPA, an enhancement to the Joint-Embedding Predictive Architecture incorporating Variance-Invariance-Covariance Regularization (VICReg) for non-contrastive self-supervised learning. This approach addresses limitations such as model collapse and inaccurate mean patch representations, enhancing stability and learning quality. C-JEPA shows improved performance across diverse tasks, including image classification, object detection, semantic segmentation, and video object segmentation.
strength_weakness: strength_weakness: 1. The framework would benefit from comparisons with latest leading methods like Dinov2[1], Mocov3[2], and IWM[3] to benchmark against current advancements.
2. The paper claims that the inclusion of Variance and Covariance terms enhances training speed and stability (lines 291–295). However, the basis for these conclusions is not clearly articulated, necessitating further explanation to validate these claims.
3. While the paper presents results from pre-training on image dataset like ImageNet, its applicability to video-related tasks can be explored. An adaptation on video JEPA methods, such as V-JEPA[4] for video domains, could significantly extend the framework's utility.
4. There are typographical and formatting errors, such as the unexplained markers in line 302.

[1] Oquab, Maxime, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193 (2023).
[2] Xie, Zhenda, et al. self-supervised learning with swin transformers. arXiv preprint arXiv:2105.04553 (2021).
[3] Garrido, Quentin, et.al. Learning and Leveraging World Models in Visual Representation Learning. arXiv:2403.00504 (2024)
[4] Bardes, Adrien, et.al. Revisiting Feature Prediction for Learning Visual Representations from Video. arXiv:2404.08471 (2024)
questions: questions: Please refer to the weakness.
limitations: limitations: Yes
paper_summary: paper_summary: This paper proposes a novel contrastive self-supervised learning framework based on JEPA, namely C-JEPA. The main idea of C-JEPA is to address the limitations of the I-JEPA, especially the limited prevention of collapse with EMA, by incorporating the principles of VICReg. The authors demonstrate the effectiveness of C-JEPA through rigorous empirical and theoretical evaluations. They show that C-JEPA achieves superior performance metrics compared to existing frameworks with faster and better convergence.
strength_weakness: strength_weakness: **[W1]** Though the paper shows a strong performance, I think the contrastive learning on JEPA lacks novelty. The proposed approach is a simple combination of I-JEPA and VICReg’s regularization strategy.

**[W2]** Does C-JEPA perform better than I-JEPA with more pre-training, e.g., 800 or 1200 epochs? In figure 1, it seems that the slope of I-JEPA is much larger than C-JEPA.

**[W3]** Invalid citation: Line 302.
questions: questions: Please address my concerns in the Weaknesses.
limitations: limitations: They did not address the limitations and potential negative societal impact of their work in the paper.

They must include the limitations in the final manuscript, e.g., the proposed approach requires a large computation and a large network capacity, therefore raising environmental concerns, e.g., carbon generation [1].

[1] Schwartz, Roy, et al. "Green ai." Communications of the ACM 63.12 (2020): 54-63.

Comments: title: Response to Reviewer 1VXy
Response to Reviewer Rspm
comment: Dear Reviewer 1VXy,

Thank you for your continued engagement and support. We will add those clarifications and experiments to the final version. Thank you once again for your insightful comments.
Dear Reviewer Rspm,

Thank you for your continued engagement and support. We will add those comparisons and clarifications to the final version. Thank you once again for your insightful comments.
Thank you for the clarifications provided. The response has effectively addressed most of my concerns. I have increased the review score accordingly. 
Regarding the training speed and stability of C-JEPA, it is encouraging to see empirical evidence, such as loss curves and performance metrics, which support faster convergence and improved stability compared to I-JEPA.
Thank you for your response. My concerns are well addressed in the rebuttal. I raise the rating to 7 and keep recommending acceptance.