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

Learn image features by predicting the embeddings of hidden image blocks, not their pixels.

I-JEPA shows a Vision Transformer one large block of an image and trains it to predict the representations that a second network computes for four other blocks of the same image. It uses no second augmented view and no pixel decoder.

A ViT-Huge/14 trained this way on ImageNet takes under 72 hours on 16 A100 GPUs, reaches 79.3% on the linear probe, and beats DINO and iBOT at estimating depth on Clevr by 19.0 and 9.6 points.

Self-supervised pretraining trains an image encoder on unlabeled photos so that its features are useful later, on tasks that do have labels. The usual test is a linear probe: freeze the pretrained encoder, train only a linear classifier on its output, and report ImageNet accuracy. A high linear-probe score means the frozen features already separate the classes without further training. Before I-JEPA, most image methods followed one of two recipes.

The first recipe makes two versions of each photo by random cropping, rescaling and color jittering, and trains the encoder to give both versions the same embedding (the output vector the encoder computes for an image). The paper calls these invariance-based or view-invariance methods; DINO, iBOT (DINO plus a masked-patch loss) and MSN (which adds masking as one more augmentation) are examples, and they led the linear-probe tables. The encoder learns to ignore whatever the augmentations change, which suits classification (a dog is a dog at any crop or hue) but builds in a bias that the paper argues can hurt tasks needing different invariances. The augmentations are also specific to photos, with no obvious counterpart for audio, and every step processes two or more views of each image.

The second recipe hides patches of the image and trains the network to reconstruct them, either as pixels (MAE) or as discrete tokens from a pretrained tokenizer (BEiT). It needs less prior knowledge and carries over to other modalities, as masked-token prediction did for text in BERT. Its features score lower on the linear probe: MAE's ViT-H/14 reaches 77.2% after 1600 epochs, and methods of this kind usually need end-to-end fine-tuning to catch up.

I-JEPA keeps the masking and drops the pixels. The network sees one large context block of the image and predicts, for four target blocks of the same image, the embeddings that a separate target encoder computes from the full image. The loss is a distance between embeddings, and the target encoder is a slowly updated copy of the network being trained. The rest of the page covers, in order: where this sits among self-supervised architectures, one training step with real tensor shapes, the masking sampler, the loss and the moving-average target encoder, the ablations that justify each choice, and what the results and compute numbers do and do not show.

Where the loss is measured: three architectures

The paper places self-supervised methods in the framework of energy-based models. A model assigns a scalar energy to a pair of inputs $(x, y)$: low when the two are compatible (two crops of one photo, or the visible and hidden parts of one image) and high when they are not. Training lowers the energy of compatible pairs. The three families in the paper's Figure 2 differ in where that energy, a distance $D$, is computed. Click through them:

In the Joint-Embedding picture $D$ sits between two embeddings, and nothing in the distance stops both encoders from mapping every image to the same vector. That constant output gives $D = 0$ for every pair, a failure called representation collapse: the energy is flat and the embedding carries no information about the input. Joint-embedding methods add machinery to prevent it: contrastive losses that push apart embeddings of different images, losses that reduce redundancy between embedding dimensions, clustering with an entropy term, or an asymmetry between the two encoders, such as a stop-gradient (no gradient flows into one of them) combined with a slowly updated copy.

In the Generative picture $D$ sits in pixel space. For masked images, $x$ is the image with patches removed, $y$ is the full image, and $z$ is a set of mask and position tokens telling the decoder which patches to fill in. A constant decoder output cannot match a real image, so collapse is not a concern as long as $z$ carries little information compared with $y$. In exchange, the loss counts an error at every pixel, including texture and noise that a classifier does not need.

The Joint-Embedding Predictive picture keeps the predictor and the conditioning input $z$ from the generative design and moves $D$ back into embedding space. The predictor outputs $\hat{s}_y$, a guess at the embedding of the hidden region, and the loss compares it with $s_y$, the embedding the target encoder produces. Pixels never enter the loss. Because both sides of $D$ are now learned embeddings, collapse is possible again, and I-JEPA uses the asymmetric fix: the target encoder receives no gradient and is an exponential moving average of the context encoder, covered in the loss section below.

How I-JEPA works: one training step

I-JEPA uses three Vision Transformers. A ViT cuts the image into square patches, embeds each patch as a vector, and runs a stack of self-attention layers over the sequence, producing one output vector per patch. The running example on this page is the paper's ViT-H/14 at 224 × 224 pixels: 14-pixel patches give a 16 × 16 grid of 256 patches, each 14 × 14 × 3 = 588 pixel values, and the encoder has 32 layers of width 1280.

• The target encoder $f_{\bar\theta}$ reads the whole image and returns 256 patch vectors, written $s_y = \{s_{y_1}, \dots, s_{y_N}\}$, a 256 × 1280 tensor.
• The context encoder $f_\theta$ reads only the patches of the context block and returns one vector per context patch, $s_x$. It has the same architecture as the target encoder and starts with identical weights.
• The predictor $g_\phi$ is a narrower ViT: width 384 and 12 layers for the ViT-H models (6 layers for ViT-B, 16 for ViT-G), with as many attention heads as the encoder.

A step on one image goes like this. The sampler picks four target blocks, which share one size within a batch; in the running example each is 7 × 6 = 42 patches (the next section shows where that size comes from). The target encoder runs on all 256 patches, and the rows of its output at each block's 42 positions become that block's target $s_y(i)$, where $B_i$ is the set of patch indices in block $i$:

The context encoder runs on the context patches, about 65 of them on average once the targets are carved out and the batch is trimmed (the next section derives that number). The predictor maps the 65 context vectors from width 1280 down to 384, appends 42 mask tokens, one per patch of block $i$, runs its 12 layers over the 107 tokens, and projects the outputs at the 42 mask-token positions back to width 1280:

Every mask token is the same learned vector plus the position embedding of the patch it stands for (in the code, a fixed 2D sine-cosine embedding). The mask tokens are the only thing that tells the predictor which block to predict, so position is the conditioning variable $z$ of the JEPA diagram. The predictor runs four times per image, once per target block, and the four predictions of shape 42 × 1280 are compared with the four targets.

The target encoder masks its output, not its input. Every one of the 256 patch vectors is computed by self-attention over the whole image, so the target for a patch of a dog's ear already reflects the rest of the dog. The alternative, cropping each target block first and encoding only its patches, gives target vectors computed from the block alone. The paper tests both on a ViT-H/16 for 300 epochs (Table 11, where the input-masking variant runs the target encoder once per target block): masking the output reaches 67.3% on the 1% ImageNet linear probe, masking the input 56.1%.

I-JEPA has no [cls] token, the extra summary token many ViTs add to read out a whole-image vector. For evaluation the paper average-pools the target encoder's 256 output vectors into one image vector. The pseudocode below writes out the step with the shapes for one image:

# One I-JEPA step, shapes for one image (ViT-H/14 at 224 px, width 1280)
ctx_idx, tgt_idx = batch_masks(B=128)       # ~65 context ids; 4 blocks of 42 ids
with no_grad():
    s_y = target_encoder(image)             # (256, 1280): every patch, full image
    s_y = layer_norm(s_y)                   # in the code; Eq. (1) has no norm
s_x = context_encoder(image, ctx_idx)       # (65, 1280): context patches only
loss = 0
for B_i in tgt_idx:                         # M = 4 target blocks
    m = mask_token + pos_embed[B_i]         # (42, 384): one token per target patch
    s_hat = predictor(s_x, m)               # (42, 1280)
    loss += ((s_hat - s_y[B_i]) ** 2).sum() # Eq. (1): squared L2, summed over B_i
loss = loss / 4                             # average over the M blocks
loss.backward()                             # gradients reach context_encoder
opt.step()                                  #   and predictor only (AdamW)
ema(target_encoder, context_encoder, m_t)   # Eq. (2), m_t from 0.996 to 1.0

Only the context encoder $\theta$ and the predictor $\phi$ are trained by gradient descent, with AdamW at batch size 2048: the learning rate warms up linearly from 10⁻⁴ to 10⁻³ over the first 15 epochs and then follows a cosine decay to 10⁻⁶, and the weight decay rises linearly from 0.04 to 0.4. The target encoder runs without gradient tracking and changes only through the moving-average update on the last line.

Choosing the context and the targets

The paper's multi-block sampler draws, for each image, $M = 4$ target blocks, each covering a random 15% to 20% of the image with an aspect ratio (height over width) between 0.75 and 1.5, and one context block covering 85% to 100% of the image, square. The target blocks may overlap each other. Because the context is drawn independently of the targets, it usually covers some of them, so every context patch that falls inside any target block is removed. That removal keeps the task non-trivial: the predictor never sees, in its input, a patch whose embedding it has to predict.

The code fixes a few details the paper leaves out. One uniform random number sets both a block's scale and its aspect ratio, so larger target blocks are also taller. Block sizes are drawn once per batch and shared by every image on the GPU; only locations vary per image. For ViT-H/14 the arithmetic runs as in the listing below: a draw of 0.5 gives scale 0.175, int(256 × 0.175) = 44 patches, aspect 1.125, and a block of 7 rows by 6 columns (42 patches). The same draw for the context gives scale 0.925, 236 patches, and a 15 × 15 block of 225.

# Multi-block masks, after src/masks/multiblock.py (ViT-H/14 at 224 px: G = 16)
def block_size(scale, aspect, G=16):
    r = rand()                                    # one draw sets scale AND aspect
    s = scale[0] + r * (scale[1] - scale[0])
    ar = aspect[0] + r * (aspect[1] - aspect[0])
    keep = int(G * G * s)                         # int(256 * 0.175) = 44 at r = 0.5
    return round(sqrt(keep * ar)), round(sqrt(keep / ar))   # (7, 6) at r = 0.5

def batch_masks(B, M=4):
    th, tw = block_size((0.15, 0.2), (0.75, 1.5))  # target size, shared by batch
    ch, cw = block_size((0.85, 1.0), (1.0, 1.0))   # context size, shared by batch
    ctx, tgt = [], []
    for b in range(B):                              # locations differ per image
        t = [random_block(th, tw) for _ in range(M)]
        c = random_block(ch, cw) - union(t)         # carve the targets out
        ctx.append(sorted(c)); tgt.append(t)
    k = min(len(c) for c in ctx)                    # equal length across the batch
    return [c[:k] for c in ctx], tgt                # first k indices, raster order

The last two lines of batch_masks set the context size the encoder actually sees. After the carve-out, a single image's context averages about 42% of its patches, because four overlapping targets cover about 45% of the image. To stack the batch into one tensor, the collator cuts every image's context list to the length of the shortest one in the batch, keeping the first patches in raster order (row by row, left to right). With 128 images per GPU the shortest context averages about 25% of the patches, which matches the 0.25 the paper reports as the multi-block context's “Avg. Ratio” in Table 6. On the 256-patch grid, 25% is about 65 patches, the number used in the training step above. (The 42%, 45% and 25% come from running a port of the official sampler; the paper states only the 0.25.)

The figure runs that port on a 14 × 14 grid, the grid of the /16 models (ViT-B/16, ViT-L/16) that the masking ablations use. Each resample draws a batch of 128 and shows its first image. Turn on the batch trim to see which context patches the equal-length cut removes, and move the sliders to see how the averages respond:

The trim keeps the first patches in raster order, so for every image except the one with the shortest context it removes the bottom rows of the carved context block.

The loss and the slowly moving target encoder

The loss is the squared L2 distance between predicted and target patch embeddings, summed over the patches of each target block and averaged over the $M$ blocks:

For the running example, each $\hat{s}_{y_j}$ and $s_{y_j}$ is a 1280-dimensional vector, each block has 42 patches, and $M = 4$, so one image contributes 4 × 42 = 168 squared distances, divided by 4. The released code differs in two ways: it layer-normalizes each target vector over its 1280 features before comparing (the operation described in the layer normalization explainer), and it uses PyTorch's smooth L1 loss averaged over every element instead of a summed squared error. Smooth L1 equals half the squared error when a coordinate is off by less than 1 and grows linearly beyond that, so the two losses agree up to scale on small errors and the code penalizes large errors less.

Gradients of (1) flow into the predictor $\phi$ and, through $s_x$, into the context encoder $\theta$. The target encoder $\bar\theta$ gets none. After every optimizer step its weights move a small fraction of the way toward the context encoder's:

The momentum $m$ starts at 0.996 and rises linearly to 1.0 by the end of pretraining. At 0.996 each update moves the target encoder 0.4% of the way toward the context encoder, which makes it an exponential average over roughly the last $1/(1-m)$, or 250, steps; the weight on a step from 173 updates ago has halved, since $0.996^{173} \approx 0.5$. At batch 2048, one epoch of ImageNet's 1.28 million training images is about 626 steps, and 300 epochs about 188,000, so the averaging window starts near 0.13% of the run and grows without bound as $m$ reaches 1.0, when the target encoder stops changing.

Why this prevents collapse is not settled. With a stop-gradient, the target is held fixed during each step, so gradient descent cannot lower the loss by moving the target toward the prediction; it can only move the prediction toward the target. The paper groups asymmetric designs with the heuristic defences against collapse, notes that the EMA target encoder “has proven essential” for joint-embedding methods with Vision Transformers, and reports the same for I-JEPA. It does not run an ablation with a frozen target ($m = 1$ throughout) or with a target that copies the context encoder every step ($m = 0$).

The figure is a toy: one weight of the context encoder, rising noisily as it would under gradient updates, and the target-encoder copy of that weight under (2). Drag the momentum from 0.5 (the copy follows within a few steps) to 1.0 (it never leaves its starting value), and turn on the schedule to raise $m$ linearly to 1.0 across the plotted run:

Because the two encoders start with identical weights, the target at step 0 is the context encoder's own output on the full image.

Why predict embeddings instead of pixels

The paper's Table 7 isolates the target. Two ViT-L/16 models use the same masking and architecture; one predicts the target encoder's output, the other predicts the raw pixels of the target blocks, with the target encoder removed from the loss. Both are evaluated with a linear probe on 1% of the ImageNet labels:

The representation target wins by 26.2 points while training for 300 fewer epochs. The paper's explanation is stated as a conjecture: the target encoder can produce abstract targets from which irrelevant pixel-level detail has been removed. A squared error on pixels scores every pixel equally, so a model that gets an object right and its fur texture wrong is penalized for the texture. A squared error on target-encoder outputs scores only what the target encoder represents, and that changes during training as the target encoder follows the context encoder.

The same comparison appears in the paper's results tables, with the architecture and epochs no longer matched: MAE predicts pixels with an encoder that also sees only visible patches, and its ViT-H/14 reaches 77.2% on the full linear probe after 1600 epochs against I-JEPA's 79.3% after 300.

Which masking choices matter

Table 6 compares multi-block masking with three other ways of splitting an image into context and targets, on a ViT-B/16 trained for 300 epochs and probed on 1% of ImageNet. Rasterized cuts the image into four quadrants and predicts three from the fourth. Block predicts one block covering 60% of the image from the rest. Random predicts a random 60% of the patches from the rest. In every strategy the context and the targets are disjoint, and the average context is 25% to 40% of the image:

Rasterized masking has the same 25% average context as multi-block and scores 15.5%. Its context is one quadrant, a view of one corner of the scene, while multi-block's 25% is cut from a block covering most of the image. Block and random give the model more context (40%) and still score 20.2% and 17.6%; their target is one large region or a scatter of single patches rather than several object-sized blocks. The paper's summary of these results is that the targets should be several fairly large blocks and the context should be informative and spread across the image.

Appendix C varies one setting at a time from the multi-block default. The figure puts each sweep on its own tab, with the paper's default in bright teal. Click a bar to read its value against the default; on the three masking tabs, the thumbnail draws one sample mask at that setting:

Target scale (Table 8) peaks at the default range of 0.15 to 0.2 (54.2%). On a 224-pixel image a block of 17.5% is about 94 × 94 pixels, while the lowest setting allows blocks down to 7.5%, about 61 × 61 pixels, and scores 19.2%. Raising the range to 0.2 to 0.3 lowers the score to 33.6%, since four such blocks cover much more of the image and leave less context.

Context scale (Table 9) rises steadily as the context block grows: 31.2% when it may be as small as 40% of the image, 47.1% at 65%, 49.3% at 75%, 54.2% at the default 85%.

Number of targets (Table 10) gives the steepest curve: 9.0% with one target block, 22.0% with two, 48.5% with three and 54.2% with four. Each extra target block adds 42 predicted patches per image in the running example at the cost of one more predictor pass; the target encoder still runs once per image. The paper stops at four.

Predictor size (Tables 12 and 14): depth helps and width does not. A 12-layer predictor beats a 6-layer one on ViT-L/16 (66.9% against 64.0%). A narrow predictor beats a wide one: width 384 scores 70.7% and width 1024, the ViT-L encoder's own width, scores 68.4%. The paper reports that the width bottleneck helps without offering a mechanism. Table 14 was run for 600 epochs and its caption describes fine-tuning, so its numbers are not comparable with Table 12's.

The paper also looks at what the predictor has learned (Sec. 8). It freezes the context encoder and predictor, trains a separate generative model (RCDM) to turn the average-pooled predictor output back into an image, and samples several decodings. The samples agree on the pose and identity of high-level object parts, such as the back of a bird or the top of a car, and vary in the fine detail, which the paper reads as the predictor capturing positional uncertainty rather than exact pixels.

Results: where I-JEPA wins and where it trails

The paper evaluates on four kinds of task: the ImageNet linear probe, ImageNet with 1% of the labels (about 12 or 13 images per class), linear-probe transfer to three other classification datasets, and two low-level tasks on Clevr, a dataset of rendered 3D scenes of simple objects: counting the objects and estimating depth. Two baselines are close relatives, and the paper names both as its closest prior work: data2vec also predicts representations of masked patches computed by a target encoder, and CAE (Context Autoencoders) adds an alignment term in representation space to a reconstruction loss. The figure shows every row of each table. Rows are grouped as in the paper, by whether the method uses view augmentations, and the readout compares I-JEPA's ViT-H/14 with the best other row in each group:

ImageNet linear probe (Table 1). The ViT-H/14 reaches 79.3% and the ViT-H/16 at 448-pixel resolution 81.1%, both after 300 epochs. Among methods without view augmentations, the strongest other row is CAE's ViT-L/16 at 78.1%, so the H/14 leads by 1.2 points (and data2vec's 77.3% by 2.0). The view-augmentation methods still lead at comparable size: iBOT's ViT-L/16 scores 81.0% and DINO's ViT-B/8 80.1%. The 448-pixel I-JEPA matches iBOT (81.1 against 81.0).

1% of ImageNet labels (Table 2). Each method is adapted by fine-tuning or linear probing, whichever works better for it. The ViT-H/14 scores 73.3%, level with data2vec's ViT-L/16 and 1.8 points above MAE's ViT-H/14, but below MSN's ViT-B/4 at 75.7%. Only the 448-pixel model, at 77.3%, is above every row in the table.

Transfer (Table 3). On CIFAR100, Places205 and iNaturalist 2018, the ViT-H/14 scores 87.5, 58.4 and 47.6. That is above MAE's ViT-H/14 on all three (by 10.2, 3.4 and 14.7 points) and above data2vec's ViT-L/16 on all three. It beats DINO's ViT-B/8 on CIFAR100 and Places205 but trails it by 8.3 points on iNaturalist, and it trails iBOT's ViT-L/16 on all three datasets (by 0.8, 2.0 and 9.7). The paper's own summary is that I-JEPA “decreases the gap” with the best view-invariance methods.

Counting and depth (Table 4). On Clevr/Count the ViT-H/14 scores 86.7, just above DINO (86.6) and iBOT (85.7). On Clevr/Dist, depth estimation, it scores 72.4 against DINO's 53.4 and iBOT's 62.8, leads of 19.0 and 9.6 points. The paper cites these as evidence that I-JEPA keeps low-level information that view-invariance methods lose, and scopes the claim to them: MAE's ViT-H/14, which also trains without view augmentations, scores 90.5 on counting and ties I-JEPA at 72.4 on depth.

Table 5 adds scale. Pretraining the same ViT-H/14 on ImageNet-22K (about 14 million images in about 22,000 classes) instead of ImageNet-1K raises four of the five transfer scores (CIFAR100 to 89.5, iNaturalist to 50.5, Clevr/Count to 88.6, Clevr/Dist to 75.0); Places205 moves from 58.4 to 57.8. A ViT-G/16 on ImageNet-22K improves Places205 and iNaturalist further (59.1 and 55.3) but not the Clevr tasks, which the authors attribute to its larger 16-pixel patches.

Compute: a ViT-H/14 in under 1200 GPU-hours

The abstract's headline is a ViT-Huge/14 trained on ImageNet with 16 A100 GPUs in under 72 hours of wall clock, which is 16 × 72 = 1,152 GPU-hours; the paper rounds it to “less than 1200 GPU hours”. The released config for this model runs 300 epochs at batch 2048, which is 300 × 1.28 million, or about 384 million, image passes. Finishing that in 72 hours takes at least 1,480 images per second across the 16 GPUs, about 93 per GPU.

The paper compares that budget two ways. Against view-invariance methods, pretraining the ViT-H/14 is over 2.5 times faster than pretraining a ViT-S/16 with iBOT, although the ViT-S has about 22 million parameters and the ViT-H about 630 million; the paper attributes this to I-JEPA processing a single view of each image where iBOT processes several. Against MAE, the ViT-H/14 is over 10 times more efficient than an MAE ViT-H/14. Each I-JEPA iteration is about 7% slower than an MAE iteration, since computing targets means a forward pass of the target encoder over all 256 patches, but I-JEPA converges in roughly 5 times fewer iterations.

Within one step, the context encoder, the only encoder that needs a backward pass, processes about 65 of the 256 patches. The target encoder processes all 256 but only forward, with no stored activations for backpropagation. The predictor is 384 wide instead of 1280 and runs once per target block on about 107 tokens.