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

Mahmoud Assran; Quentin Duval; Ishan Misra; Piotr Bojanowski; Pascal Vincent; Michael Rabbat; Yann LeCun; Nicolas Ballas

Source-linked AI summary

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

Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, Nicolas Ballas

arXiv:2301.08243v3Jan 19, 2023cs.CVcs.AIcs.LGeess.IV

TL;DR

Self-supervised image learning seeks semantic representations without relying on hand-crafted augmentations, whose biases can limit transfer across tasks. I-JEPA predicts large target-block representations from an informative context block in representation space, and achieves strong performance across semantic and low-level tasks while scaling efficiently.

Problem
Hand-crafted view augmentations can impose task-specific biases, while pixel-reconstruction methods typically produce lower-semantic-level representations in limited-supervision evaluations.
Method
I-JEPA predicts the representations of multiple target blocks from a single context block in abstract representation space, using large targets and a spatially distributed context.
Results
I-JEPA is competitive with view-invariant pretraining on semantic tasks, performs better on object counting and depth prediction, and requires less than 1200 GPU hours for ViT-H/14 ImageNet pretraining.
Takeaways & Limitations
Predicting in representation space provides a path to strong general image representations without hand-crafted view augmentations and with reduced pretraining computation.

Abstract

from arXiv · show

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

I-JEPA addresses limitations of augmentation-based and pixel-reconstruction self-supervision by predicting abstract image representations from context. Its masking strategy supports semantic learning, while experiments report strong performance and efficient scaling.

Motivation: Invariance-based methods use hand-crafted image augmentations that can impose biases unsuitable for tasks requiring different invariances or modalities.These methods can produce highly semantic representations, but their image-specific transformations are not straightforward to generalize.
Motivation: Masked generative methods require less prior knowledge and generalize beyond images, but typically learn lower-semantic-level representations than invariance-based methods.They often underperform invariance-based pretraining in linear probing and limited-supervision semantic transfer settings.
I-JEPA: I-JEPA predicts target-block representations from a context block in abstract representation space, without relying on extra image transformations.Target representations are computed by a learned target encoder.
I-JEPA: I-JEPA uses large target blocks and an informative, spatially distributed context block to guide representations toward semantic content.Its abstract targets can eliminate unnecessary pixel-level details compared with pixel/token prediction.
Results: I-JEPA learns strong off-the-shelf representations without hand-crafted view augmentations and outperforms pixel-reconstruction methods on ImageNet-1K and semantic transfer evaluations.The reported evaluations include linear probing, 1% ImageNet-1K, and semantic transfer tasks.
Results: Less than 1200 GPU hours are required to pretrain a ViT-H/14 on ImageNet, making I-JEPA over 2.5× faster than iBOT and over 10× more efficient than MAE.The paper attributes the computational reduction to predicting in representation space.

2. Background

Self-supervised methods can be described through energy-based relationships between compatible and incompatible inputs. Joint-embedding, generative, and predictive architectures differ in whether they compare embeddings, reconstruct signals, or predict embeddings.

Energy-Based View: Self-supervised learning captures relationships between inputs by assigning low energy to compatible inputs and high energy to incompatible inputs.This objective can be framed using Energy-Based Models.
Joint-Embedding Architectures: Joint-Embedding Architectures learn similar embeddings for compatible inputs and dissimilar embeddings for incompatible inputs, commonly using augmented views of the same image.Their main challenge is representation collapse, where the encoder produces a constant output.
Generative Architectures: Generative Architectures reconstruct a signal from a compatible input using a decoder conditioned on additional variables such as mask and position tokens.Masked image reconstruction is a common computer-vision example, and low-capacity conditioning variables prevent representation collapse.
Joint-Embedding Predictive Architectures: Joint-Embedding Predictive Architectures predict a signal’s embedding from a compatible input, applying the loss in embedding space rather than input space.JEPAs use a predictor conditioned on additional information and require asymmetric encoder architectures to avoid collapse.

3. Method

I-JEPA predicts representations of multiple target blocks from one context block using Vision Transformers, with masking designed to preserve semantic targets and informative context. Predictions are made in representation space against target-encoder outputs updated by exponential moving average.

I-JEPA predicts representations of multiple target blocks from a single context block in the same image.The context encoder processes visible patches, while a predictor uses positional mask tokens to predict each target location.
The method uses a non-generative ViT architecture whose predictions are made in representation space rather than pixel or token space.The context encoder, target encoder, and predictor all use Vision Transformer architectures.
Targets: Target blocks are sampled from target-encoder patch representations, typically using 4 blocks with scale 0.15–0.2 and aspect ratio 0.75–1.5.Masking the target-encoder output rather than the input is described as crucial for high-semantic-level target representations.
Context: A single context block uses scale 0.85–1.0, removes overlapping target regions, and is processed as visible context patches.This makes the context informative while keeping it sparse for efficient processing.
Prediction: The predictor applies positional mask tokens to generate patch-level predictions for each target block.It runs once per target block, conditioning on the locations to be predicted.
Loss: Training minimizes the average L2 distance between predicted and target patch representations, while the target encoder follows an exponential moving average of the context encoder.The predictor and context encoder are optimized by gradient descent; the moving-average target encoder is reported as essential for training with Vision Transformers.

4. Related Work

Related work spans denoising and masked image modeling, joint-embedding methods, and approaches using handcrafted views. I-JEPA differs by predicting abstract representations from a single image view, supporting efficiency and broader modality applicability.

Earlier denoising approaches reconstruct missing or corrupted image content, including pixels or input patches.Examples include denoising autoencoders, context encoders, colorization, and masked autoencoders.
Representation-prediction methods such as data2vec avoid handcrafted augmentations and can apply across vision, text, and speech.Context Autoencoders also predict missing patches in representation space using reconstruction and alignment losses.
I-JEPA improves linear probing performance over other augmentation-free methods while using less computational effort.Its larger models can outperform the best Context Autoencoder model with a fraction of the compute.
Joint-embedding methods such as DINO, MSN, and iBOT rely on handcrafted augmentations and typically process multiple views of each image.I-JEPA instead requires processing a single view and predicts in representation space.
On ImageNet-1%, I-JEPA outperforms MAE and benefits from scaling to larger models and higher resolution.A ViT-H/16 trained at resolution 448 surpasses previous methods, including methods using handcrafted augmentations.

5. Image Classification

Across ImageNet classification evaluations, I-JEPA learns strong representations without handcrafted view augmentations. It improves over augmentation-free methods and narrows or closes gaps with view-invariant approaches across linear, low-shot, and transfer settings.

I-JEPA significantly improves ImageNet-1K linear probing over MAE, CAE, and data2vec while using less computational effort.Larger I-JEPA models outperform the best CAE model using a fraction of the compute.
At resolution 448 × 448, a ViT-H/16 I-JEPA model matches view-invariant approaches without handcrafted data augmentations.The result is reported as evidence of scaling benefits in ImageNet linear evaluation.
I-JEPA learns strong off-the-shelf representations without hand-crafted view augmentations across ImageNet linear probing, semi-supervised classification, and semantic transfer.This broad evaluation pattern is stated as a central result of the classification experiments.
Low-Shot ImageNet-1K: In the 1% ImageNet setting, I-JEPA outperforms MAE with a similar encoder while requiring fewer pretraining epochs.A ViT-H/14 I-JEPA model matches a ViT-L/16 data2vec model, and the higher-resolution ViT-H/16 surpasses previous methods.
Transfer learning: For downstream image classification transfer, I-JEPA outperforms augmentation-free MAE and data2vec and surpasses DINO on CIFAR100 and Place205 with a linear probe.It also decreases the gap with the strongest view-invariance-based methods.

6. Local Prediction Tasks

I-JEPA is reported to capture both semantic and local image features. On linear-probe object counting and depth prediction, it outperforms view-invariance methods, especially on depth prediction.

The results indicate that I-JEPA captures local image features during pretraining in addition to improving downstream semantic classification.This conclusion connects the local-task evaluation with the classification findings reported in the preceding section.
I-JEPA surpasses view-invariance-based methods on object counting and depth prediction.These are evaluated as low-level downstream tasks using frozen encoders and linear probes.
On Clevr, I-JEPA outperforms DINO and iBOT for object counting and by a large margin for depth prediction.The tasks are Clevr/Count and Clevr/Dist, respectively.

7. Scalability

I-JEPA scales efficiently with both compute and pretraining data, while larger models improve semantic transfer but may not help local prediction tasks. Its compute advantage comes from faster convergence despite modest per-iteration overhead.

Compute efficiency: I-JEPA requires less compute than previous methods while achieving strong semi-supervised ImageNet-1K performance.Compared with MAE and data2vec, it uses fewer pretraining epochs; compared with iBOT, ViT-H/14 uses less compute than ViT-S/16.
Data scaling: Larger and more diverse pretraining datasets improve transfer-learning performance across conceptually different semantic and low-level tasks.The comparison is between ImageNet-1K and ImageNet-22K pretraining.
Model scaling: ViT-G/16 improves semantic classification transfer over ViT-H/14, but does not improve low-level tasks because its larger input patches can harm local prediction.The reported semantic tasks include Place205 and INat18.

8. Predictor Visualizations

The predictor visualization conditions generated samples on positional mask tokens, allowing qualitative inspection of the representations predicted for target regions.

Visualization setup: The visualization compares each original image with its context image processed by a pretrained I-JEPA ViT-H/14 encoder.Subsequent columns show samples decoded from the predictor output.
Visualization setup: Green bounding boxes identify target regions whose predictor representations are decoded by a generative model.The predictor is conditioned on positional mask tokens specifying each green-box location.
Interpretation: The visualization probes whether positional conditioning lets the predictor capture uncertainty about the target representation.This is presented as a qualitative investigation of the predictor's behavior.

9. Ablations

The ablations identify representation-space prediction and multi-block masking as central to I-JEPA's semantic representations. The masking comparison evaluates alternative context-target arrangements under linear evaluation.

Prediction targets: Predicting in pixel space significantly degrades 1% ImageNet-1K linear-probing performance compared with predicting in representation space.The representation-space loss permits abstract targets in which irrelevant pixel-level details are eliminated.
Masking strategy: Rasterized masking uses one quadrant to predict the other three, whereas block masking predicts one block from the image complement.Random masking predicts a set of random patches from the image complement; context and target regions do not overlap.
Predictor visualizations: The predictor visualization decodes target-region outputs conditioned on positional mask tokens, providing a qualitative probe of predicted representations.The visualization uses samples from a generative model to display predictor outputs.
Masking strategy: Multi-block masking helps guide I-JEPA toward learning semantic representations.It is compared with rasterized, block, and random masking strategies.

10. Conclusion

I-JEPA learns semantic image representations efficiently without hand-crafted data augmentations by predicting in representation space. The conclusion contrasts this joint-embedding approach with both pixel reconstruction and view-invariance methods.

Conclusion: I-JEPA learns high-level semantic representations without relying on hand-crafted data augmentations.The method is described as simple and efficient.
Conclusion: Predicting in representation space lets I-JEPA converge faster than pixel-reconstruction methods.The target encoder supplies abstract prediction targets rather than pixel-level targets.
Architecture: I-JEPA uses Vision Transformer architectures for its context encoder, target encoder, and predictor.The predictor is a lightweight narrow ViT with fixed embedding dimension 384.
Optimization: The target encoder is initialized from the context encoder and updated thereafter using an exponential moving average.The momentum starts at 0.996 and increases linearly to 1.0 during pretraining.
Masking: The default masking design samples four target blocks and one context block, removing context regions that overlap targets.Target blocks are sampled with scales from 0.15 to 0.2, while the context block scale ranges from 0.85 to 1.0.

A.2. Downstream Tasks

The supplied passages describe evaluation protocols and masking ablations for I-JEPA, emphasizing how target-block scale, context informativeness, and target count affect low-shot performance.

Evaluation protocols: Linear evaluation freezes the encoder and trains a classifier across iNaturalist18, CIFAR100, Clevr/Count, Clevr/Dist, and Places205.For I-JEPA, average-pooled patch representations replace the [cls] token.
Evaluation protocols: ImageNet linear evaluation uses average-pooled representations, LARS, batch-size 16384, and a 50-epoch probe.The learning rate is decayed by a factor of 10 every 15 epochs while sweeping three reference learning rates and two weight-decay values.
Evaluation protocols: Low-shot evaluation fine-tunes ViT-L/H models for 50 epochs on ImageNet-1% using AdamW and a cosine learning-rate schedule.The protocol uses batch size 512, layer decay 0.75, label smoothing 0.1, and MAE-style randaugment.
Related methods: Joint-embedding methods commonly prevent collapse through explicit regularization or architectural constraints such as gradient stopping, momentum encoders, or asymmetric prediction heads.The passage places these approaches within active research on self-supervised visual representation learning.
Multiblock masking: Larger target blocks improve low-shot accuracy when the context is sufficiently informative, while reducing context size degrades performance.The ablations vary target-block scale and context scale during ViT-B/16 pretraining for 300 epochs.
Multiblock masking: Increasing the number of target blocks improves low-shot accuracy on the 1% ImageNet-1K benchmark.This ablation uses I-JEPA with a ViT-B/16 pretrained for 300 epochs and evaluates with a linear probe.

C. Additional Ablations

Additional ablations show that I-JEPA’s semantic representations depend on target-encoder output masking, predictor design, and multiblock masking, while full-ImageNet fine-tuning remains competitive.

Masking design: Masking the target-encoder output produces more semantic prediction targets and improves linear probing performance.Masking the input instead requires a separate target-encoder forward pass for each target region.
Predictor design: Increasing predictor depth from 6 to 12 layers significantly improves downstream low-shot performance.The comparison uses a ViT-L/16 pretrained for 500 epochs.
Optimization: A smaller pretraining weight decay improves ImageNet-1% fine-tuning but degrades linear evaluation.The comparison is between a fixed weight decay of 0.05 and a schedule increasing from 0.04 to 0.4.
Predictor design: A predictor-width bottleneck improves ImageNet-1% downstream performance relative to matching the ViT-L encoder’s 1024 channels.The compared predictor widths are 384 and 1024 channels.
Representation visualizations: RCDM visualizations decode representations with a diffusion-model decoder and random noise, revealing captured information through features common across generated samples.I-JEPA target-encoder visualizations capture high-level object information and poses while discarding low-level details; MSN samples show higher variability.