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

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

# Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture

###### Abstract

This paper demonstrates an approach for learning highly semantic image representations without relying on hand-crafted data-augmentations. We introduce the Image-based Joint-Embedding Predictive Architecture (I-JEPA), a non-generative approach for self-supervised learning from images. The idea behind I-JEPA is simple: from a single context block, predict the representations of various target blocks in the same image. A core design choice to guide I-JEPA towards producing semantic representations is the masking strategy; specifically, it is crucial to (a) sample target blocks with sufficiently large scale (semantic), and to (b) use a sufficiently informative (spatially distributed) context block. Empirically, when combined with Vision Transformers, we find I-JEPA to be highly scalable. For instance, we train a ViT-Huge/14 on ImageNet using 16 A100 GPUs in under 72 hours to achieve strong downstream performance across a wide range of tasks, from linear classification to object counting and depth prediction.

## 1 Introduction

In computer vision, there are two common families of approaches for self-supervised learning from images: invariance-based methods [ 4 , 37 , 17 , 24 , 35 , 18 , 74 , 10 , 1 ] and generative methods [ 28 , 8 , 57 , 36 ] .

Invariance-based pretraining methods optimize an encoder to produce similar embeddings for two or more views of the same image [ 15 , 20 ] , with image views typically constructed using a set of hand-crafted data augmentations, such as random scaling, cropping, and color jittering [ 20 ] , amongst others [ 35 ] . These pretraining methods can produce representations of a high semantic level [ 18 , 4 ] , but they also introduce strong biases that may be detrimental for certain downstream tasks or even for pretraining tasks with different data distributions [ 2 ] . Often, it is unclear how to generalize these biases for tasks requiring different levels of abstraction. For example, image classification and instance segmentation do not require the same invariances [ 11 ] . Additionally, it is not straightforward to generalize these image-specific augmentations to other modalities such as audio.

Cognitive learning theories have suggested that a driving mechanism behind representation learning in biological systems is the adaptation of an internal model to predict sensory input responses [ 59 , 31 ] . This idea is at the core of self-supervised generative methods, which remove or corrupt portions of the input and learn to predict the corrupted content [ 67 , 57 , 36 , 9 , 71 , 68 ] . In particular, mask-denoising approaches learn representations by reconstructing randomly masked patches from an input, either at the pixel or token level. Masked pretraining tasks require less prior knowledge than view-invariance approaches and easily generalize beyond the image modality [ 8 ] . However, the resulting representations are typically of a lower semantic level and underperform invariance-based pretraining in off-the-shelf evaluations (e.g., linear-probing) and in transfer settings with limited supervision for semantic classification tasks [ 4 ] . Consequently, a more involved adaptation mechanism (e.g., end-to-end fine-tuning) is required to reap the full advantage of these methods.

In this work, we explore how to improve the semantic level of self-supervised representations without using extra prior knowledge encoded through image transformations. To that end, we introduce a joint-embedding predictive architecture [ 48 ] for images (I-JEPA). An illustration of the method is provided in Figure 3 . The idea behind I-JEPA is to predict missing information in an abstract representation space; e.g., given a single context block, predict the representations of various target blocks in the same image, where target representations are computed by a learned target-encoder network.

Compared to generative methods that predict in pixel/token space, I-JEPA makes use of abstract prediction targets for which unnecessary pixel-level details are potentially eliminated, thereby leading the model to learn more semantic features. Another core design choice to guide I-JEPA towards producing semantic representations is the proposed multi-block masking strategy. Specifically, we demonstrate the importance of predicting sufficiently large target blocks in the image, using an informative (spatially distributed) context block.

Through an extensive empirical evaluation, we demonstrate that: • I-JEPA learns strong off-the-shelf representations without the use of hand-crafted view augmentations (cf. Fig. 1 ). I-JEPA outperforms pixel-reconstruction methods such as MAE [ 36 ] on ImageNet-1K linear probing, semi-supervised 1% ImageNet-1K, and semantic transfer tasks.

• I-JEPA is competitive with view-invariant pretraining approaches on semantic tasks and achieves better performance on low-level visions tasks such as object counting and depth prediction (Sections 5 and 6 ). By using a simpler model with less rigid inductive bias, I-JEPA is applicable to a wider set of tasks.

• I-JEPA is also scalable and efficient (Section 7 ). Pre-training a ViT-H/14 on ImageNet requires less than 1200 GPU hours, which is over 2.5 × 2.5\times faster than a ViT-S/16 pretrained with iBOT [ 79 ] and over 10 × 10\times more efficient than a ViT-H/14 pretrained with MAE. Predicting in representation space significantly reduces the total computation needed for self-supervised pretraining.

## 2 Background

Self-supervised learning is an approach to representation learning in which a system learns to capture the relationships between its inputs. This objective can be readily described using the framework of Energy-Based Models (EBMs) [ 49 ] in which the self-supervised objective is to assign a high energy to incompatible inputs, and to assign a low energy to compatible inputs. Many existing generative and non-generative approaches to self-supervised learning can indeed be cast in this framework; see Figure 2 .

#### Joint-Embedding Architectures.

Invariance-based pretraining can be cast in the framework of EBMs using a Joint-Embedding Architecture (JEA), which learns to output similar embeddings for compatible inputs, 𝒙 , 𝒚 {\bm{x}},{\bm{y}} , and dissimilar embeddings for incompatible inputs; see Figure 2(a) . In the context of image-based pretraining, compatible 𝒙 , 𝒚 {\bm{x}},{\bm{y}} pairs are typically constructed by randomly applying hand-crafted data augmentations to the same input image [ 20 ] .

The main challenge with JEAs is representation collapse, wherein the energy landscape is flat (i.e., the encoder produces a constant output regardless of the input). During the past few years, several approaches have been investigated to prevent representation collapse, such as contrastive losses that explicitly push apart embeddings of negative examples [ 15 , 37 , 24 ] , non-contrastive losses that minimize the informational redundancy across embeddings [ 74 , 10 ] , and clustering-based approaches that maximize the entropy of the average embedding [ 18 , 5 , 4 ] . There are also heuristic approaches that leverage an asymmetric architectural design between the x x -encoder and y y -encoder to avoid collapse [ 24 , 35 , 8 ] .

#### Generative Architectures.

Reconstruction-based methods for self-supervised learning can also be cast in the framework of EBMs using Generative Architectures; see Figure 2(b) . Generative Architectures learn to directly reconstruct a signal 𝒚 {\bm{y}} from a compatible signal 𝒙 {\bm{x}} , using a decoder network that is conditioned on an additional (possibly latent) variable 𝒛 {\bm{z}} to facilitate reconstruction. In the context of image-based pretraining, one common approach in computer vision is to produce compatible 𝒙 , 𝒚 {\bm{x}},{\bm{y}} pairs using masking [ 38 , 9 ] where 𝒙 {\bm{x}} is a copy of the image 𝒚 {\bm{y}} , but with some of the patches masked. The conditioning variable 𝒛 {\bm{z}} then corresponds to a set of (possibly learnable) mask and position tokens, that specifies to the decoder which image patches to reconstruct. Representation collapse is not a concern with these architectures as long as the informational capacity of 𝒛 {\bm{z}} is low compared to the signal 𝒚 {\bm{y}} .

#### Joint-Embedding Predictive Architectures.

As shown in Figure 2(c) , Joint-Embedding Predictive Architectures [ 48 ] are conceptually similar to Generative Architectures; however, a key difference is that the loss function is applied in embedding space, not input space. JEPAs learn to predict the embeddings of a signal 𝒚 {\bm{y}} from a compatible signal 𝒙 {\bm{x}} , using a predictor network that is conditioned on an additional (possibly latent) variable 𝒛 {\bm{z}} to facilitate prediction. Our proposed I-JEPA provides an instantiation of this architecture in the context of images using masking; see Figure 3 .

In contrast to Joint-Embedding Architectures, JEPAs do not seek representations invariant to a set of hand-crafted data augmentations, but instead seek representations that are predictive of each other when conditioned on additional information 𝒛 {\bm{z}} . However, as with Joint-Embedding Architectures, representation collapse is also a concern with JEPAs; we leverage an asymmetric architecture between the 𝒙 {\bm{x}} - and 𝒚 {\bm{y}} -encoders to avoid representation collapse.

## 3 Method

We now describe the proposed Image-based Joint-Embedding Predictive Architecture (I-JEPA), illustrated in Figure 3 . The overall objective is as follows: given a context block, predict the representations of various target blocks in the same image. We use a Vision Transformer [ 63 , 29 ] (ViT) architecture for the context-encoder, target-encoder, and predictor. A ViT is composed of a stack of transformer layers, each consisting of a self-attention [ 66 ] operation followed by a fully-connected MLP. Our encoder/predictor architecture is reminiscent of the generative masked autoencoders (MAE) [ 36 ] method. However, one key difference is that the I-JEPA method is non-generative and the predictions are made in representation space.

#### Targets.

We first describe how we produce the targets in the I-JEPA framework: in I-JEPA, the targets correspond to the representations of image blocks. Given an input image 𝒚 {\bm{y}} , we convert it into a sequence of N N non-overlapping patches, and feed this through the target-encoder f θ ¯ f_{\bar{\theta}} to obtain a corresponding patch-level representation 𝒔 y = { 𝒔 y 1 , … , 𝒔 y N } {\bm{s}}_{y}=\{{\bm{s}}_{y_{1}},\dots,{\bm{s}}_{y_{N}}\} where 𝒔 y k {\bm{s}}_{y_{k}} is the representation associated with the k th k^{\text{th}} patch. To obtain the targets for our loss, we randomly sample M M (possibly overlapping) blocks from the target representations 𝒔 y {\bm{s}}_{y} . We denote by B i B_{i} the mask corresponding of the i th i^{\text{th}} block and by 𝒔 y ​ ( i ) = { 𝒔 y j } j ∈ B i {\bm{s}}_{y}(i)=\{{\bm{s}}_{y_{j}}\}_{j\in B_{i}} its patch-level representation. Typically, we set M M equal to 4, and sample the blocks with a random aspect ratio in the range ( 0.75 , 1.5 ) (0.75,1.5) and random scale in the range ( 0.15 , 0.2 ) (0.15,0.2) . Note that the target blocks are obtained by masking the output of the target-encoder, not the input. This distinction is crucial to ensure target representations of a high semantic level; see, e.g., [ 8 ] .

#### Context.

Recall, the goal behind I-JEPA is to predict the target block representations from a single context block. To obtain the context in I-JEPA, we first sample a single block 𝒙 {\bm{x}} from the image with a random scale in the range ( 0.85 , 1.0 ) (0.85,1.0) and unit aspect ratio. We denote by B x B_{x} the mask associated with the context block 𝒙 {\bm{x}} . Since the target blocks are sampled independently from the context block, there may be significant overlap. To ensure a non-trivial prediction task, we remove any overlapping regions from the context block. Figure 4 shows examples of various context and target blocks in practice. Next, the masked context block, 𝒙 {\bm{x}} , is fed through the context encoder f θ f_{\theta} to obtain a corresponding patch-level representation 𝒔 x = { 𝒔 x j } j ∈ B x {\bm{s}}_{x}=\{{\bm{s}}_{x_{j}}\}_{j\in B_{x}} .

#### Prediction.

Given the output of the context encoder, 𝒔 x {\bm{s}}_{x} , we wish to predict the M M target block representations 𝒔 y ​ ( 1 ) , … , 𝒔 y ​ ( M ) {\bm{s}}_{y}(1),\ldots,{\bm{s}}_{y}(M) . To that end, for a given target block 𝒔 y ​ ( i ) {\bm{s}}_{y}(i) corresponding to a target mask B i B_{i} , the predictor g ϕ ​ ( ⋅ , ⋅ ) g_{\phi}(\cdot,\cdot) takes as input the output of the context encoder 𝒔 x {\bm{s}}_{x} and a mask token for each patch we wish to predict, { 𝒎 j } j ∈ B i \{{\bm{m}}_{j}\}_{j\in B_{i}} , and outputs a patch-level prediction 𝒔 ^ y ​ ( i ) = { 𝒔 ^ y j } j ∈ B i = g ϕ ​ ( 𝒔 x , { 𝒎 j } j ∈ B i ) \hat{{\bm{s}}}_{y}(i)=\{\hat{{\bm{s}}}_{y_{j}}\}_{j\in B_{i}}=g_{\phi}({\bm{s}}_{x},\{{\bm{m}}_{j}\}_{j\in B_{i}}) . The mask tokens are parameterized by a shared learnable vector with an added positional embedding. Since we wish to make predictions for M M target blocks, we apply our predictor M M times, each time conditioning on the mask tokens corresponding to the target-block locations we wish to predict, and obtain predictions 𝒔 ^ y ​ ( 1 ) , … , 𝒔 ^ y ​ ( M ) \hat{{\bm{s}}}_{y}(1),\ldots,\hat{{\bm{s}}}_{y}(M) .

#### Loss.

The loss is simply the average L 2 L_{2} distance between the predicted patch-level representations 𝒔 ^ y ​ ( i ) \hat{{\bm{s}}}_{y}(i) and the target patch-level representation 𝒔 y ​ ( i ) {\bm{s}}_{y}(i) ; i.e., 1 M ​ ∑ i = 1 M D ⁡ ( 𝒔 ^ y ​ ( i ) , 𝒔 y ​ ( i ) ) = 1 M ​ ∑ i = 1 M ∑ j ∈ B i ∥ 𝒔 ^ y j − 𝒔 y j ∥ 2 2 . \frac{1}{M}\sum^{M}_{i=1}{D}\left(\hat{{\bm{s}}}_{y}(i),{\bm{s}}_{y}(i)\right)=\frac{1}{M}\sum^{M}_{i=1}\sum_{j\in B_{i}}\lVert\hat{{\bm{s}}}_{y_{j}}-{\bm{s}}_{y_{j}}\rVert^{2}_{2}. The parameters of the predictor, ϕ \phi , and context encoder, θ \theta , are learned through gradient-based optimization, while the parameters of the target encoder θ ¯ \bar{\theta} are updated via an exponential moving average of the context-encoder parameters. The use of an exponential moving average target-encoder has proven essential for training JEAs with Vision Transformers [ 25 , 79 , 18 ] , we find the same to be true for I-JEPA.

## 4 Related Work

A long line of work has explored visual representation learning by predicting the values of missing or corrupted sensory inputs. Denoising autoencoders use random noise as input corruption [ 67 ] . Context encoders regress an entire image region based on its surrounding [ 57 ] . Other works cast image colorization as a denoising task [ 77 , 46 , 47 ] .