Architectures / Computer Vision

🧩 I-JEPA (Joint-Embedding Predictive Architecture)

CVPR 2023 - prediction in representation space rather than pixel space. The target encoder is an EMA copy that carries no gradient.

From Assran et al. (2023). Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

11 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 43 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class JEPAPredictor(nn.Module):
    """I-JEPA's predictor, narrow on purpose: a predictor as wide as the encoder
    is a second encoder. Prediction happens in representation space."""

    def __init__(self, embed_dim: int = 768, predictor_dim: int = 384,
                 depth: int = 6, num_heads: int = 12):
        super().__init__()
        heads = num_heads if predictor_dim % num_heads == 0 else 1
        self.inp = nn.Linear(embed_dim, predictor_dim)
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=predictor_dim, nhead=heads,
                dim_feedforward=4 * predictor_dim, batch_first=True)
            for _ in range(max(1, depth))
        ])
        self.out = nn.Linear(predictor_dim, embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.inp(x)
        for blk in self.blocks:
            h = blk(h)
        return self.out(h)


class EMATarget(nn.Module):
    """Not a computation: a declaration that this branch is an exponential
    moving average of its upstream and carries NO gradient. Dropping the detach
    is how BYOL / SimSiam / I-JEPA collapse to a constant."""

    def __init__(self, momentum: float = 0.996):
        super().__init__()
        self.momentum = momentum

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.detach()

For agents

Also in Computer Vision