I-JEPA implementation

A compact, readable implementation of the core mechanics behind I-JEPA, the Image Joint Embedding Predictive Architecture. The goal of this repo is to make the moving parts obvious and to back every claim with a test that runs on CPU in a couple of seconds. Nothing is downloaded and every tensor stays tiny.

I-JEPA learns image representations without pixel reconstruction. Instead of predicting raw pixels it predicts the representation of unseen target patches from the representation of a visible context patch. Three pieces make that work:

1. A masking scheme that carves the patch grid into a context region and several target blocks that never overlap the context.
2. A context encoder (a small Vision Transformer) that embeds only the visible context patches.
3. A predictor that takes the context tokens plus positional mask tokens at the target locations and predicts the target representations, which come from a separate target encoder updated as an exponential moving average of the context encoder.

What is in here

src/
  config.py    tiny dataclass config (32x32 image, 8px patches, 48 dim encoder)
  masking.py   block masking: target blocks + a disjoint context block
  models.py    Vision Transformer encoder, predictor, and the IJEPA bundle
  ema.py       exponential moving average update for the target encoder
  loss.py      smooth L1 loss between predictions and detached targets
tests/
  test_masking.py  context and targets are disjoint, blocks are valid rectangles
  test_models.py   encoder and predictor shapes, target encoder is frozen
  test_train.py    loss drops on one batch, EMA math is correct

How the masking works

Patches live on a grid_size x grid_size grid and are addressed by flat row major indices. For each image we sample a handful of rectangular target blocks and one larger context block, then subtract the union of the target patches from the context. That subtraction is what guarantees the context and the targets share no patch, which is the property test_context_and_targets_are_disjoint checks. The MaskCollator truncates the per image index lists to a common length so a whole batch stacks into one tensor for a clean batched forward pass.

How prediction works

The context encoder embeds only the kept context patches. The predictor projects those tokens into a narrower width, appends one learnable mask token per target patch (each carrying the positional embedding of its target location), runs a short transformer, then reads out the tokens at the target positions and maps them back to the encoder width. The supervision signal is the representation of the same target patches produced by the target encoder over the full image. The target side is detached, so no gradient flows into the target encoder. It is nudged toward the context encoder only through the EMA update.

Running the tests

python -m pytest tests/ -q

On this machine all 13 tests pass in about two seconds on CPU.

A real result

The overfitting test drives the predictor and context encoder on a single fixed batch of four random images for sixty Adam steps. In one run the smooth L1 loss went from 0.5268 on the first step to 0.2060 on the last, which is the drop the test asserts on. Numbers will vary a little with the random seed, but the loss falling well below half of its starting value is the behavior the test pins down. These are sanity figures from a tiny synthetic setup, not benchmark results.

Notes and scope

This is a teaching sized implementation. It keeps the parts that define I-JEPA (block masking, a context encoder, a representation space predictor, and an EMA target encoder) and leaves out the engineering you would add for a full training run such as large backbones, multi crop data loading, and a learning rate schedule. The positional embeddings are fixed 2D sin/cos tables, the config is deliberately small, and every test uses synthetic tensors so the suite stays fast and offline.