We are going to take a short look at the joint-embedding predictive architecture (JEPA). JEPA is a form of self-supervised learning (SSL) whose basic idea is the prediction between latent space representations of related objects [LeCun, 2022]. In particular, the representation is required to be predictable from unobserved parts of the input, giving rise to a latent space that ignores unpredictable details.
A concrete JEPA for images was introduced by [Assran et al., 2023]. This so-called I-JEPA combines ideas from both reconstruction-based and teacher-student SSL approaches. An image is embedded such that the representation of an image part can be predicted from the representation of another contextual part.
A more detailed discussion of I-JEPA is provided below.
I-JEPA
I-JEPA is an SSL strategy for vision transformers (ViTs) that is based on reconstructing latent space representations. As in the masked autoencoder (MAE), a patch-based masking strategy is used to hide parts of an image. In contrast to the MAE, the missing information is predicted inside latent space, not inside image space. This avoids the pixel-level objective of generative SSL and encourages the embedding to ignore noise and to focus on relevant features only.
On the other hand, this formulation suffers from representation collapse, i.e. the learned representation may, at worst, collapse into a constant. This is a well-known problem in other SSL approaches that also match latent space representations. As a remedy, I-JEPA employs two different ViT encoders (same architecture, different weights) together with an asymmetric teacher-student-style learning scheme.
The following figure visualizes I-JEPA:
First of all, an input image is decomposed into non-overlapping patches. These patches are then processed by three different ViTs. The target encoder computes a global representation of the full image. From the ensemble of patch embeddings obtained this way, one generates random, possibly overlapping blocks. At the same time, a single context block is randomly created by removing patches from the input. This block is not allowed to overlap with the target blocks. While the removed parts are simply ignored, the context encoder embeds the visible context patches.
After the target and context embeddings are computed, another ViT predicts the target from the context. To this end, the predictor relies on context embeddings and, in addition, a mask token for each target patch. These mask tokens are based on a shared learnable parameter vector and the standard encodings of the patch positions. The resulting patch-level embeddings are then compared against the target ones. The original paper simply uses an \(L_2\) loss function here.
In order to train the combined system of the target encoder, context encoder and predictor, I-JEPA employs an updating scheme that is reminiscent of teacher-student methods. Context encoder and predictor are learned directly through gradient-based optimization. The parameters of the target encoder, however, are updated more slowly as the exponential moving average (EMA) of the context encoder’s weights. This way, the target encoder provides a stabilized target to the context encoder. After training, the final embedding is the average-pooled output of the target encoder.