I-JEPA: predict embeddings of masked blocks, not their pixels

How a person fills in a dog whose head is covered

The paper, in order — every section

Open arXiv 2301.08243 (v3, 13 Apr 2023) next to this. Headings below match the PDF.

Abstract

Learn semantic image representations without hand-crafted view augmentations. I-JEPA: from a single context block, predict representations of various target blocks in the same image. Masking must (a) sample large-scale (semantic) targets and (b) keep a spatially distributed informative context. ViT-H/14 on ImageNet, 16 A100s, under 72 hours; transfers to linear classification, object counting, depth.

Figure 1 — ImageNet linear vs GPU hours

The poster plot: I-JEPA sits above MAE/data2vec at fewer GPU hours because the loss is already in embedding space, so you do not train a pixel decoder. Read the axes before you memorise a single top-1.

Figure 2 — three energy cartoons

(a) Joint-Embedding Architecture: pull $E(x)$ and $E(y)$ together for two augmented views. Collapse is the failure mode. (b) Generative: decode $y$ from $E(x)$ and mask tokens $z$. Collapse is unlikely if $z$ is small. (c) JEPA: predict $E(y)$ from $E(x)$ given $z$. Looks like (b), loses like (a). This figure is LeCun 2022 drawn for CVPR. If you cannot redraw it, do not start §3.

§1 Introduction

Two SSL families in vision. Invariance methods (SimCLR, BYOL, DINO, VICReg, iBOT): strong semantics, but the augmentation list is a prior that does not travel to audio and is wrong for some tasks (classification vs instance segmentation want different invariances). Generative / mask-denoising (MAE, BEiT): less prior, travels across modalities, but representations are lower-level — you need full fine-tuning to catch up on linear probe. I-JEPA wants MAE’s lack of augmentation prior and DINO’s off-the-shelf semantics.

Three empirical claims to score later: (1) beats MAE-like methods on ImageNet linear / 1% / semantic transfer; (2) competitive with view-invariant methods on semantics and better on counting and depth; (3) cheaper: ViT-H/14 < 1200 GPU hours, $2.5\times$ faster than iBOT ViT-S/16, $10\times$ more efficient than MAE ViT-H/14.

§2 Background

EBMs: low energy on compatible pairs. JEA collapse fixes: contrastive negatives, VICReg redundancy, clustering entropy, or asymmetric EMA. Generative collapse is rare because you still have to paint $y$. JEPA collapse is real; they steal BYOL’s EMA teacher plus a predictor. Conditioning on $z$ is what makes JEPA not “make two views identical.”

§3 Method — Figure 3

Targets. Run the target encoder on the full image (or all patches), then keep only the tokens whose patches fall in a sampled target block. Targets are representations, not pixels, not discrete visual tokens.

Context. Sample a large block, remove any overlap with the targets, feed the rest to the context encoder. The encoder never sees mask tokens.

Predictor. Narrow ViT. Inputs: context tokens + one learned positional token per target patch location. Output: one vector per target patch, same width as the encoder.

Loss. Average L2 (paper: average of squared error in representation space) over predicted vs target tokens. Stop-gradient through the target encoder. EMA update of target weights.

Figure 4 — masking

§4 Related work

Places I-JEPA against MAE, data2vec, CAE, MSN, DINO, iBOT. data2vec is the closest cousin (predict EMA representations) but uses a different mask and is not sold as “no view augmentations.” Cite the PDF, not a blog, for the exact differences.

§5 Image classification — Table 1

Protocol: freeze encoder, linear probe on full ImageNet-1K train, average-pooled tokens (no [CLS]). LARS, batch 16384, 50 epochs (appendix). Numbers from Table 1:

View-aug methods still sit higher at similar size (DINO ViT-B/8 80.1, iBOT ViT-L/16 81.0). The paper’s claim is not “beats DINO.” It is “matches that ballpark without the augmentation prior, and does it with fewer epochs than MAE.”

1% ImageNet semi-supervised (≈12 labels/class): ViT-L/16 69.4, ViT-H/14 73.3, ViT-H/16₄₄₈ 77.3 (Meta blog quotes the 12-shot story). Full fine-tune, appendix Table 15: I-JEPA ViT-H/16₄₄₈ 87.1 vs MAE ViT-H/14₄₄₈ 87.8 after 1600 MAE epochs — I-JEPA used 300. Within 1%, 5.3× fewer epochs. Do not say “beats MAE at fine-tune.”

§6 Local prediction tasks

This is the section that justifies dropping SimCLR. Counting (CLEVR) and depth: view-invariance methods throw away localisation that those tasks need. I-JEPA keeps it because the predictor is still spatially addressed. Quote the PDF’s tables here rather than a tweet. The qualitative claim in the Meta blog: a decoder fitted on predictor outputs sketches a dog head / bird leg in the right pose — semantics without discarding position.

§7 Scalability

ViT-H/14 < 1200 GPU hours; 16×A100 < 72 h. Faster than iBOT ViT-S/16 by $2.5\times$, cheaper than MAE ViT-H/14 by $10\times$. Predicting in representation space removes the decoder FLOPs and converges in ~5× fewer iterations than MAE (paper’s comparison).

§8 Predictor visualisations

Already above. The predictor is a restricted world model of spatial uncertainty in a still image. It is not a physics engine.

§9 Ablations

Load-bearing knobs, from the PDF: large target blocks beat small/random MAE masks; context must remain informative (do not mask 90% like MAE); predicting encoder outputs beats predicting pixels (Table 11 flavour: masking the target-encoder output vs input). Predictor width is a bottleneck on purpose — too wide and it copies texture.

§10 Conclusion

Semantic features, no hand-crafted views, cheaper than pixel reconstruction. Path for joint-embedding methods that are not SimCLR. Next sentence in history is V-JEPA (replace image blocks with video tubes).

What the paper does not claim

Not a world model of time. Not an agent. Not “beats DINO.” Not hierarchical JEPA. No language. No robot.

The engineering tension

Prerequisites

1. 01

   Vision Transformers

   Patch tokens, positional embeddings, no [CLS] required (I-JEPA pools). If ViT is fuzzy, Figure 3 is unreadable.

2. 02

   MAE-style masking

   I-JEPA reuses the “hide patches, predict them” game. The change is the output space.

3. 03

   BYOL / EMA teacher

   Stop-gradient + slow copy of the encoder is the collapse brake. Same idea as Grill et al. 2020.

4. 04

   The JEPA cartoon

   Read LeCun 2022 until Figure 2 of this paper is obvious.

Knowledge graph

Architecture

Multi-block masking

Default recipe from Figure 4 / §3: $M=4$ targets, scale $U(0.15,0.2)$, aspect $U(0.75,1.5)$; context scale $U(0.85,1.0)$ then subtract overlaps. This is not MAE’s random 75%. A random sprinkle of patches can be inpainted from local texture; a 20% block cannot, unless you know what object was there.

Results, honestly

Protocol reminder: linear numbers are frozen-encoder, average pool, LARS probe. Do not compare them to a paper that fine-tunes the backbone unless you say so.

Worked intuition

ViT-B/16, $224\times 224$ → $14\times 14 = 196$ patches, $d=768$. Sample a target block of ~20% area ≈ 40 patches. Context keeps ~$196-40=156$ patches (more, because four smaller blocks). Context encoder: 156 tokens → 156 vectors in $\mathbb{R}^{768}$. Predictor (narrower, say 384-d inside, project back) + 40 position tokens → 40 vectors in $\mathbb{R}^{768}$. Target encoder: 196 tokens, index the 40. Loss: mean of $\| \hat{v}_i - \mathrm{sg}(\bar{v}_i) \|_2^2$. EMA $\bar{\theta} \leftarrow m\bar{\theta} + (1-m)\theta$. Nobody decoded $16\times 16\times 3$ RGB cells.

Caveats

• Still ImageNet-centric pretraining. “No augmentations” still uses random resized crop to 224 in the usual ViT pipeline — they mean no colour jitter / blur / multi-crop views.
• Linear-probe wins versus MAE are the headline; fine-tune is a near-tie.
• Predictor visualisations use a separate generative probe. Pretty pictures are not the loss.
• Not a temporal world model. That is V-JEPA.

Study plan

80 minutes: case (10) → Figure 2 cartoons + Figure 3 (20) → Table 1 protocol and numbers (15) → Figure 4 masking + §9 (15) → §6 counting/depth vs DINO (10) → open V-JEPA (10). Hold the PDF on Figure 3 the whole time.

Primary sources

• Paper: arXiv:2301.08243 · arXiv HTML v3
• CVPR page: Assran et al., CVPR 2023, pp. 15619–15629
• Code / weights: github.com/facebookresearch/ijepa
• Official blog: Meta AI — I-JEPA
• Blueprint: LeCun 2022 JEPA · next: V-JEPA

@InProceedings{Assran_2023_CVPR,
  author    = {Assran, Mahmoud and Duval, Quentin and Misra, Ishan and Bojanowski, Piotr and Vincent, Pascal and Rabbat, Michael and LeCun, Yann and Ballas, Nicolas},
  title     = {Self-Supervised Learning From Images With a Joint-Embedding Predictive Architecture},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  month     = {June},
  year      = {2023},
  pages     = {15619--15629}
}