# 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.

Pick for self-supervised pretraining without a hand-built augmentation pipeline. The stop-gradient on the target branch is what prevents representation collapse.

- Category: Computer Vision
- Paper: Assran et al. (2023). Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023 (https://arxiv.org/abs/2301.08243)
- Layers: 11
- Parameters: 40.74M
- Input shape (batchless): 3 × 224 × 224
- Output shape: 196 × 768
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/ijepa/model.json
- Open on the canvas: https://neurarch.com/?template=ijepa

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | Image 224x224 | Input | shape=[3, 224, 224] | 3 × 224 × 224 |
| 2 | Context Patch Embed | Patch Embed | embedDim=768, patchSize=16 | 196 × 768 |
| 3 | Context Encoder Block 1 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 196 × 768 |
| 4 | Context Encoder Block 2 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 196 × 768 |
| 5 | Predictor (narrow, 384) | JEPA Predictor | embedDim=768, numHeads=12 | 196 × 768 |
| 6 | predicted target reps | Output |  | 196 × 768 |
| 7 | Target Patch Embed | Patch Embed | embedDim=768, patchSize=16 | 196 × 768 |
| 8 | Target Encoder Block 1 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 196 × 768 |
| 9 | Target Encoder Block 2 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 196 × 768 |
| 10 | EMA Target (stop-gradient) | EMA Target / Stop-Gradient |  | 196 × 768 |
| 11 | target reps (no grad) | Output |  | 196 × 768 |

## Verifier findings

- **info** `deep-no-norm`: 8 layers with no BatchNorm, LayerNorm, or GroupNorm. Without normalization, activations can explode or vanish across layers, causing slow or unstable training. Fix: Add BatchNorm after Conv2d (CV tasks), LayerNorm after attention/FFN (NLP/LLM), or GroupNorm for small batch sizes.

## Exported PyTorch (first 46 lines)

```python
# 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()
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md