I-JEPA From Scratch

A compact, readable, from-scratch implementation of I-JEPA — the Image-based Joint-Embedding Predictive Architecture (Assran et al., Meta AI, CVPR 2023) — trained self-supervised on CIFAR-10, with a linear-probe evaluation and a one-click Google Colab notebook.

The defaults are sized to run in a few minutes on a free Colab GPU while keeping every real I-JEPA mechanism intact: multi-block masking, an EMA target encoder, latent-space prediction, and a stop-gradient.

---

The core idea

Most self-supervised vision methods either:

• reconstruct pixels of masked regions (MAE) — wastes capacity on low-level detail, or
• contrast augmentations of the same image (SimCLR, DINO) — needs carefully hand-tuned augmentations and negative pairs.

I-JEPA does neither. It predicts, in representation space, what a target network would have encoded for a masked region — given only a single visible context region. Learning happens in latent space, so the model is free to ignore unpredictable pixel-level noise and focus on semantics.

flowchart TD
    IMG([Full image])

    IMG -->|context block only| CE["Context Encoder<br/>(ViT, trained)"]
    IMG -->|full image, no grad| TE["Target Encoder<br/>(ViT, EMA copy)"]

    CE -->|context tokens| PRED["Predictor<br/>(narrow ViT)"]
    POS["Target positions<br/>(+ mask tokens)"] --> PRED

    TE -->|LayerNorm + stop-grad| TGT([Target representations])

    PRED -->|predicted reps| LOSS{{"smooth-L1 loss<br/>(latent space)"}}
    TGT --> LOSS

    LOSS -. gradients .-> CE
    LOSS -. gradients .-> PRED
    CE -. EMA update .-> TE

    classDef trained fill:#dbeafe,stroke:#2563eb,color:#1e3a8a;
    classDef frozen fill:#fee2e2,stroke:#dc2626,color:#7f1d1d;
    classDef io fill:#f1f5f9,stroke:#64748b,color:#0f172a;
    class CE,PRED trained;
    class TE frozen;
    class IMG,POS,TGT,LOSS io;

Loading

Blue = trained by gradient descent · Red = frozen, EMA-updated only · the dotted arrows are the gradient flow and the EMA update.

The three networks

Network	Sees	Trained?	Role
Context encoder	only context patches	Yes, gradient descent	builds the representation we keep
Target encoder	the whole image	No, EMA of context encoder	produces the prediction targets
Predictor	context tokens + target positions	Yes, gradient descent	predicts target reps in latent space

Why it doesn't collapse

A model that maps everything to a constant would trivially minimise the loss. Four design choices prevent that:

1. EMA target encoder — targets change slowly, so the context encoder can't instantly chase a degenerate solution.
2. Stop-gradient — no gradient flows into the target encoder.
3. LayerNorm on the targets — removes the trivial "all-zero" minimum.
4. Asymmetry — the context encoder sees less than the target encoder, so there is real information to predict.

---

Repository layout

JEPA_From_Scratch/
├── src/
│   ├── config.py      # IJEPAConfig: every hyperparameter in one dataclass
│   ├── data.py        # CIFAR-10 loaders (light aug for pretrain, none for probe)
│   ├── masking.py     # MultiBlockMaskCollator — context/target sampling
│   ├── vit.py         # ViT encoder (patch embed, 2D sin-cos pos, blocks)
│   ├── predictor.py   # narrow ViT predictor with learnable mask tokens
│   ├── ijepa.py       # wires the 3 nets + EMA update + latent loss
│   └── train.py       # pretraining loop (AdamW, cosine LR, EMA schedule)
├── eval_probe.py      # linear probe on frozen encoder features
├── smoke_test.py      # fast CPU sanity checks (masking, loss, EMA, no-grad)
├── notebooks/
│   └── ijepa_colab.ipynb   # end-to-end Colab demo
└── requirements.txt

How the pieces connect

1. src/masking.py samples, per batch, one large context block and four small target blocks, and removes target patches from the context. It returns patch-index tensors enc_mask and pred_masks.
2. src/ijepa.py runs the target encoder on the full image (no grad) → LayerNorm → these are the targets. The context encoder runs on the context patches. The predictor predicts each target block from the context tokens + target positions. Loss = smooth-L1 over the blocks.
3. After each optimiser step, src/train.py EMA-updates the target encoder with momentum ramped 0.996 → 1.0.
4. eval_probe.py freezes the target encoder, mean-pools its patch tokens, and trains a single linear layer on labels.

---

Quick start (Colab — recommended)

1. Push this repo to GitHub.
2. Open notebooks/ijepa_colab.ipynb in Colab (or File ▸ Open notebook ▸ GitHub).
3. Runtime ▸ Change runtime type ▸ T4 GPU.
4. Runtime ▸ Run all. It clones the repo, pretrains, plots the loss, and runs the linear probe.

Quick start (local)

python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt          # or the CPU wheel from pytorch.org

python -m src.train                       # self-supervised pretraining
python eval_probe.py                      # linear probe on the saved encoder

python smoke_test.py                      # fast correctness checks (CPU, seconds)

python -m src.train saves the target encoder to ./ijepa_ckpt.pt, which eval_probe.py then loads.

---

Tuning it

Everything is in src/config.py. Useful knobs:

Field	Effect
epochs, embed_dim, depth	bigger/longer → better probe accuracy, slower
num_target_blocks, target_scale	harder prediction task
context_scale	how much the context encoder gets to see
ema_start / ema_end	target-encoder update speed (collapse control)
patch_size, image_size	granularity / resolution

What good looks like: the latent loss trends steadily down (not crashing to 0 instantly — that would signal collapse), and the linear probe lands well above the 10% random baseline (a short run already clears it; longer/bigger runs go much higher).

---

Notes & gotchas

• Predict in latent space — never pixels. There is no decoder anywhere.
• The target encoder is never in the optimiser — it only moves via EMA.
• Masking must remove context/target overlap, or the task becomes trivial.
• Masks are sampled once per batch and shared across images (as in the official repo) — efficient and keeps index tensors rectangular.

Reference

Mahmoud Assran et al., "Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture", CVPR 2023. arXiv:2301.08243 · official code

This repository is an educational reimplementation, not the official codebase.