I-JEPA: Predicting Image Representations Behind a Mask

3 minute read

Published:

TL;DR: I-JEPA learns image features by predicting the target encoder's representations of hidden regions. Large target blocks and informative visible context encourage a meaningful prediction task. The target encoder follows the context encoder through a moving average; downstream evaluation checks what information the resulting features retain.

A hidden region can have a predictable role

Intuition first. Cover one wheel of a bicycle. The remaining frame and the other wheel provide evidence about the missing region, even if its precise texture is uncertain. This motivates a target that can express useful structure without requiring every pixel to be reproduced.

Assran et al. implement this with a context encoder, a predictor, and a target encoder. The target encoder processes the full image; target patch features are selected from its output. The context encoder receives visible patches with target regions removed. Positional information tells the predictor which hidden locations to predict.

Where the training signal comes from

Let \(B_m\) be target block \(m\), \(z_j\) its target feature at position \(j\), and \(\hat z_j\) the corresponding prediction. The paper’s squared-distance objective can be written as

\[ \mathcal L=\frac{1}{M}\sum_{m=1}^{M} \sum_{j\in B_m}\|\hat z_j-\operatorname{sg}(z_j)\|_2^2. \]

The stop-gradient operator \(\operatorname{sg}\) blocks backpropagation into the target branch. Its parameters instead follow an exponential moving average:

\[ \xi\leftarrow\tau\xi+(1-\tau)\theta. \]

The paper explains the loss and target updates; the official implementation provides the training configuration. A slowly changing target helps stabilise learning, but the equation alone is not a general proof against collapse.

The two views have different access to the imageA visible context grid contains a missing central region. A full target grid contains that region. The predictor is scored on features at the hidden positions. Visible contextFull target viewCompare hiddenpatch featuresThe target branch supplies training labels in feature space.
A schematic mask, not the paper's actual mask dimensions. The target branch sees information withheld from the context branch.

Worked example: one patch, two coordinates

Use an illustrative target \(z=(1,2)\) and prediction \(\hat z=(0.8,2.3)\). For one block containing this single patch,

\[ \mathcal L=(0.8-1)^2+(2.3-2)^2=0.04+0.09=0.13. \]

The gradient with respect to the prediction is \((-0.4,0.6)\). A gradient step of size \(0.1\) applied directly to this illustrative prediction gives \((0.84,2.24)\) and reduces the error to \(0.0832\). In the network, that gradient instead flows through the predictor and context encoder parameters.

The two coordinates are abstract features; they are not automatically object position or colour. Assigning physical meaning requires an additional measurement or decoder.

Why the mask matters

The paper emphasises sufficiently large target regions and spatially informative context. A tiny missing patch can often be inferred from immediate neighbours; a larger missing region demands broader evidence. Removing too much context can make the target underdetermined. Mask design therefore changes the problem the model must solve, not just its computation cost. I-JEPA masking analysis.

For a concrete diagnostic, imagine two datasets with identical objects. In one, the background always identifies the class; in the other, backgrounds vary independently. Good performance on the first does not establish that the representation encodes the object. Evaluate controlled changes in the nuisance factors you want the model to ignore.

Key Insight: Both sides of the prediction problem matter: the visible context determines what can be inferred, and the target encoder determines what counts as a correct inference.

Continue with LeJEPA for collapse and feature geometry, or V-JEPA for prediction across space and time.

References

  1. Assran, M., et al. Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023.
  2. Meta FAIR. Official I-JEPA code and models.