Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture — part 2/3

Cleaned-text code points [14882, 30562). Bibliography omitted here; source.md includes it.

The idea of image denoising has recently been revisited in the context of masked image modelling [ 36 , 71 , 9 ] , where a Vision Transformer [ 29 ] is used to reconstruct missing input patches. The work on Masked Autoencoders (MAE) [ 36 ] proposed an efficient architecture that only requires the encoder to process visible image patches. By reconstructing missing patches in pixels space, MAE achieves strong performance when fine-tuned end-to-end on large labeled datasets and exhibits good scaling properties. BEiT [ 9 ] predicts the value of missing patches in a tokenized space; specifically, tokenizing image patches using a frozen discreteVAE, which is trained on a dataset containing 250 million images [ 58 ] . Yet, pixel-level pre-training has been shown to outperform BEiT for fine-tuning [ 36 ] . Another work, SimMIM [ 71 ] , explores reconstruction targets based on the classic Histogram of Gradients [ 27 ] feature space, and demonstrates some advantage over pixel space reconstruction. Different from those works, our representation space is learned during training through a Joint-Embedding Predictive Architecture. Our goal is to learn semantic representations that do not require extensive fine-tuning on downstream tasks.

Closest to our work is data2vec [ 8 ] and Context Autoencoders [ 25 ] . The data2vec method learns to predict the representation of missing patches computed through an online target encoder; by avoiding handcrafted augmentations, the method can be applied to diverse modalities with promising results in vision, text and speech. Context Autoencoders use an encoder/decoder architecture optimized via the sum of a reconstruction loss and an alignment constraint, which enforces predictability of missing patches in representation space. Compared to these methods, I-JEPA exhibits significant improvements in computational efficiency and learns more semantic off-the-shelf representations. Concurrent to our work, data2vec-v2 [ 7 ] explores efficient architectures for learning with various modalities.

We also compare I-JEPA with various methods based on joint-embedding architectures; e.g., DINO [ 18 ] , MSN [ 4 ] and iBOT [ 79 ] . Theses methods rely on hand-crafted data augmentations during pretraining to learn semantic image representations. The work on MSN [ 4 ] , uses masking as an additional data-augmentation during pretraining, while iBOT combines a data2vec-style patch-level reconstruction loss with the DINO view-invariance loss. Common to these approaches is the need to process multiple user-generated views of each input image, thereby hindering scalability. By contrast, I-JEPA only requires processing a single view of each image. We find that a ViT-Huge/14 trained with I-JEPA requires less computational effort than a ViT-Small/16 trained with iBOT.

## 5 Image Classification

To demonstrate that I-JEPA learns high-level representations without relying on hand-crafted data-augmentations, we report results on various image classification tasks using the linear probing and partial fine-tuning protocols. In this section, we consider self-supervised models that have been pretrained on the ImageNet-1K dataset [ 60 ] . Pretraining and evaluation implementation details are described in the Appendix A . All I-JEPA models are trained at resolution 224 × 224 224\times 224 pixels, unless stated otherwise.

#### ImageNet-1K.

Table 1 shows performance on the common ImageNet-1K linear-evaluation benchmark. After self-supervised pretraining, the model weights are frozen and a linear classifier is trained on top using the full ImageNet-1K training set. Compared to popular methods such as Masked Autoencoders (MAE) [ 36 ] , Context Autoencoders (CAE) [ 22 ] , and data2vec [ 8 ] , which also do not rely on extensive hand-crafted data-augmentations during pretraining, we see that I-JEPA significantly improves linear probing performance, while using less computational effort (see section 7 ). By leveraging the improved efficiency of I-JEPA, we can train larger models that outperform the best CAE model while using a fraction of the compute. I-JEPA also benefits from scale; in particular, a ViT-H/16 trained at resolution 448 × 448 448\times 448 pixels matches the performance of view-invariant approaches such as iBOT [ 79 ] , despite avoiding the use of hand-crafted data-augmentations.

#### Low-Shot ImageNet-1K.

Table 2 shows performance on the 1% ImageNet benchmark. Here the idea is to adapt the pretrained models for ImageNet classification using only 1% of the available ImageNet labels, corresponding to roughly 12 or 13 images per class. Models are adapted via fine-tuning or linear-probing, depending on whichever works best for each respective method. I-JEPA outperforms MAE while requiring less pretraining epochs when using a similar encoder architecture. I-JEPA, using a ViT-H/14 architecture, matches the performance of a ViT-L/16 pretrained with data2vec [ 8 ] , while using significantly less computational effort (see Section 7 ). By increasing the image input resolution, I-JEPA outperforms previous methods including joint-embedding methods that do leverage extra hand-crafted data-augmentations during pretraining, such as MSN [ 4 ] , DINO [ 17 ] , and iBOT [ 79 ] .

#### Transfer learning.

Table 3 shows performance on various downstream image classification tasks using a linear probe. I-JEPA significantly outperforms previous methods that do not use augmentations (MAE and data2vec), and decreases the gap with the best view-invariance-based methods, which leverage hand-crafted data augmentations during pretraining, even surpassing the popular DINO [ 18 ] on CIFAR100 and Place205 with a linear probe.

## 6 Local Prediction Tasks

As demonstrated in Section 5 , I-JEPA learns semantic image representations that significantly improve the downstream image classification performance of previous methods, such as MAE and data2vec. Additionally, I-JEPA benefits from scale and can close the gap, and even surpass, view-invariance based methods that leverage extra hand-crafted data augmentations. In this section, we find that I-JEPA also learns local image features and surpasses view-invariance based methods on low-level and dense prediction tasks, such as object counting and depth prediction.

Table 4 shows performance on various low-level tasks using a linear probe. After pretraining, the encoder weights are frozen and a linear model is trained on top to perform object-counting and depth prediction on the Clevr dataset [ 43 ] . Compared to view-invariance methods such as DINO and iBOT, the I-JEPA method effectively captures low-level image features during pretraining and outperforms them in object counting (Clevr/Count) and (by a large margin) depth prediction (Clevr/Dist).

## 7 Scalability

#### Model Efficiency.

I-JEPA is highly scalable compared to previous approaches. Figure 5 shows semi-supervised evaluation on 1% ImageNet-1K as a function of GPU hours. I-JEPA requires less compute than previous methods and achieves strong performance without relying on hand-crafted data-augmentations. Compared to reconstruction-based methods, such as MAE, which directly use pixels as targets, I-JEPA introduces extra overhead by computing targets in representation space (about 7% slower time per iteration). However, since I-JEPA converges in roughly 5 × 5\times fewer iterations, we still see significant compute savings in practice. Compared to view-invariance based methods, such as iBOT, which rely on hand-crafted data augmentations to create and process multiple views of each image, I-JEPA also runs significantly faster. In particular, a huge I-JEPA model (ViT-H/14) requires less compute than a small iBOT model (ViT-S/16).

#### Scaling data size.

We also find I-JEPA to benefit from pretraining with larger datasets. Table 5 shows transfer learning performance on semantic and low level tasks when increasing the size of the pretraining dataset (IN1K versus IN22K). Transfer learning performance on these conceptually different tasks improves when pretraining on a larger more diverse dataset.

#### Scaling model size.

Table 5 also shows that I-JEPA benefit from larger model size when pretraining on IN22K. Pretraining a ViT-G/16 significantly improves the downstream performances on image classification tasks such as Place205 and INat18 compared to a ViT-H/14 model, but does not improve performance on low-level downstream tasks — the ViT-G/16 uses larger input patches, which can be detrimental for the local prediction tasks.

## 8 Predictor Visualizations

The role of the predictor in I-JEPA is to take the output of the context encoder and, conditioned on positional mask tokens, to predict the representations of a target black at the location specified by the mask tokens. One natural question is whether the predictor conditioned on the positional mask tokens is learning to correctly capture positional uncertainty in the target. To qualitatively investigate this question, we visualize the outputs of the predictor. We use the following visualization approach to enable the research community to independently reproduce our findings. After pretraining, we freeze the context-encoder and predictor weights, and train a decoder following the RCDM framework [ 13 ] to map the average-pool of the predictor outputs back to pixel space. Figure 6 shows decoder outputs for various random seeds. Qualities that are common across samples represent information that is contained in the average-pooled predictor representation. The I-JEPA predictor correctly captures positional uncertainty and produces high-level object parts with the correct pose (e.g., back of the bird and top of the car).

## 9 Ablations

#### Predicting in representation space.

Table 7 compares low-shot performance on 1% ImageNet-1K using a linear probe when the loss is computed in pixel-space versus representation space. We conjecture that a crucial component of I-JEPA is that the loss is computed entirely in representation space, thereby giving the target encoder the ability to produce abstract prediction targets, for which irrelevant pixel-level details are eliminated. From Table 7 , it is clear that predicting in pixel-space leads to a significant degradation in the linear probing performance.

#### Masking strategy.

Table 6 compare our multi-block masking with other masking strategies such as rasterized masking, where the image is split into four large quadrants, and the goal is to use one quadrant as a context to predict the other three quadrants, and the traditional block and random masking typically used in reconstruction-based methods. In block masking, the target is a single image block and the context is the image complement. In random masking, the target is a set of random patches and the context is the image complement. Note that there is no overlap between the context and target blocks in all considered strategies. We find multi-block masking helpful for guiding I-JEPA to learning semantic representations. Additional ablations on multi-block masking can be found in Appendix C .

## 10 Conclusion

We proposed I-JEPA, a simple and efficient method for learning semantic image representations without relying on hand-crafted data augmentations. We show that by predicting in representation space, I-JEPA converges faster than pixel reconstruction methods and learns representations of a high semantic level. In contrast to view-invariance based methods, I-JEPA highlights a path for learning general representations with joint-embedding architectures, without relying on hand-crafted view augmentations.

## Appendix A Implementation Details

### A.1 Pretraining

#### Architectures.

For I-JEPA pretraining, we use Vision Transformer [ 29 ] (ViT) architectures for the context-encoder, target-encoder, and the predictor. While the context-encoders and target-encoders correspond to standard ViT architectures, the predictor is designed as a light-weight (narrow) ViT architecture. Specifically, we fix the embedding dimension of the predictor to 384, while keeping the number of self-attention heads equal to that of the backbone context-encoder. For the smaller ViT-B/16 context-encoder, we set the depth of the predictor to 6. For ViT-L/16, ViT-H/16, and ViT-H/14 context-encoders, we set the depth of the predictor to 12. Finally, the ViT-G/16 uses a predictor of depth 16. I-JEPA is pretrained without a [cls] token. We use the target-encoder for evaluation and average pool its output to produce a global image representation.

#### Optimization.

We use AdamW [ 51 ] to optimize the context-encoder and predictor weights. Our default batch-size is 2048, and the learning rate is linearly increased from 10 − 4 10^{-4} to 10 − 3 10^{-3} during the first 15 epochs of pretraining, and decayed to 10 − 6 10^{-6} following a cosine schedule thereafter. Following [ 18 , 4 ] , the weight-decay is linearly increased from 0.04 0.04 to 0.4 0.4 throughout pretraining. The target-encoder weights are identical to the context-encoder weights at initialization, and updated via an exponential moving average thereafter [ 61 , 37 , 23 , 35 , 18 , 4 ] . We use a momentum value of 0.996 0.996 , and linearly increase this value to 1.0 1.0 throughout pretraining, following [ 18 , 4 ] .

#### Masking.

By default, we sample 4 4 possibly overlapping target blocks masks with random scale in the range ( 015 , 0.2 ) (015,0.2) and aspect ratio in the range ( 0.75 , 1.5 ) (0.75,1.5) . We sample 1 1 context block mask with random scale in the range ( 0.85 , 1.0 ) (0.85,1.0) and unit aspect ratio. We subsequently eliminate any regions in the context block mask that overlap with any of the 4 4 target block masks. The context-block mask and target-block masks are sampled independently for each image in the mini-batch. To ensure efficient batch processing, we restrict the size of all context masks on a co-located GPU to be identical. Similarly, we restrict the size of all target masks on a co-located GPU to be identical. The mask-sampler is efficiently implemented in only a few lines of code in PyTorch [ 56 ] using a batch-collator function, which runs in the data loader processes. In short, in each iteration, the data loader returns a mini-batch of images and a set of context and target masks for each image, identifying the patch indices to keep for the context and target views.

### A.2 Downstream Tasks

#### Linear evaluation.

When evaluating methods such as iBOT [ 79 ] , DINO [ 18 ] or MAE [ 36 ] , which leverage Vision Transformers [ 29 ] with an additional [cls] token, we use the default configurations of VISSL [ 34 ] to evaluate all the models on iNaturalist18 [ 65 ] , CIFAR100 [ 45 ] , Clevr/Count [ 42 , 75 ] , Clevr/Dist [ 42 , 75 ] , and Places205 [ 78 ] . We freeze the encoder and return the best number among the following representations: 1) the [cls] token representation of the last layer, 2) the concatenation of the last 4 4 layers of the [cls] token. For each representation, we try two different heads: 1) a linear head, or 2) a linear head preceded by a batch normalization, and return the best number. We use the default data augmentations of VISSL [ 34 ] : random resize cropping and horizontal flipping, with the exception of Clevr/Count and Clevr/Dist, where we only use center crop and horizontal flipping, as random cropping interferes with the capability of counting objects and estimating distance, removing critical objects from the scene. For CIFAR100, we resize the images to 224 × 224 224\times 224 pixels, so as to keep the number of patches equal to that used during pretraining.