microJEPA — JEPA From Scratch

Work in progress.

Yann LeCun has been talking about JEPA for years. Now his new lab AMI raised over $1B to build world models based on it. The whole field is moving that way — World Labs, Genie 3, all of them.

So I sat down and implemented the thing from scratch. No PyTorch, no NumPy. Just math and random. I wanted to understand what JEPA actually does, not just read about it.

The result is microjepa — a single Python file that implements I-JEPA with a tensor engine, autograd, Vision Transformer, attention, Adam optimizer, all built from nothing. About 375 lines.

What JEPA does differently

Current models (LLMs, diffusion) are generative. They predict the next token or reconstruct every pixel. LeCun says this is wasteful. The world is full of details you cannot predict — the exact texture of a wall, the position of every leaf. A good model should not try to predict all that.

JEPA predicts representations instead of raw data:

MAE: mask image patches, reconstruct the pixels. Loss in pixel space.
JEPA: mask image patches, predict the representations. Loss in embedding space.

This means JEPA can ignore low-level noise and focus on what matters — structure, semantics, meaning. That is the bet for building something that understands physics and common sense.

Here is an interactive comparison of the three main self-supervised approaches — MAE, JEPA, and contrastive learning. Click the tabs to see how each one works:

The full code

375 lines. Tensor engine, autograd, ViT, predictor, EMA, training loop — everything in one file, zero dependencies.

Show full implementation (microjepa.py)

"""
microjepa.py — JEPA from scratch. Zero dependencies. Pure Python.
Implements I-JEPA (LeCun et al.): ViT encoder, predictor, EMA, L1 loss,
on a from-scratch tensor + autograd engine. Run: python microjepa.py
"""
import math, random

# --- Tensor Engine: minimal nd-array with autograd ---

def _shape(d):
    s = []
    while isinstance(d, list): s.append(len(d)); d = d[0] if d else []
    return tuple(s)

def _fill(shape, val): return val if not shape else [_fill(shape[1:], val) for _ in range(shape[0])]
def _fillfn(shape, fn): return fn() if not shape else [_fillfn(shape[1:], fn) for _ in range(shape[0])]
def _ewise(a, b, fn): return fn(a, b) if not isinstance(a, list) else [_ewise(ai, bi, fn) for ai, bi in zip(a, b)]

def _acc(tgt, src):
    for i in range(len(tgt)):
        if isinstance(tgt[i], list): _acc(tgt[i], src[i])
        else: tgt[i] += src[i]

def _flat(d): return [d] if not isinstance(d, list) else [x for sub in d for x in _flat(sub)]

def _unflat(f, s):
    if not s: return f[0]
    sz = 1
    for x in s[1:]: sz *= x
    return [_unflat(f[i*sz:(i+1)*sz], s[1:]) for i in range(s[0])]

class Tensor:
    def __init__(self, data, _ch=(), _op='', requires_grad=False):
        self.data, self.shape = data, _shape(data)
        self.grad, self.requires_grad, self._backward, self._prev = None, requires_grad, lambda: None, set(_ch)

    def item(self):
        d = self.data
        while isinstance(d, list): d = d[0]
        return d

    @staticmethod
    def zeros(s): return Tensor(_fill(s, 0.0))
    @staticmethod
    def randn(s, std=0.02): return Tensor(_fillfn(s, lambda: random.gauss(0, std)))

    def _ensure_grad(self):
        if self.grad is None: self.grad = Tensor.zeros(self.shape)

    def __add__(self, other):
        o = other if isinstance(other, Tensor) else Tensor(_fill(self.shape, other))
        out = Tensor(_ewise(self.data, o.data, lambda a, b: a+b), (self, o), '+')
        def _bw():
            for t in (self, o):
                if t.requires_grad: t._ensure_grad(); _acc(t.grad.data, out.grad.data)
        out._backward, out.requires_grad = _bw, self.requires_grad or o.requires_grad
        return out

    def __mul__(self, other):
        o = other if isinstance(other, Tensor) else Tensor(_fill(self.shape, other))
        out = Tensor(_ewise(self.data, o.data, lambda a, b: a*b), (self, o), '*')
        def _bw():
            if self.requires_grad: self._ensure_grad(); _acc(self.grad.data, _ewise(out.grad.data, o.data, lambda g,v: g*v))
            if o.requires_grad: o._ensure_grad(); _acc(o.grad.data, _ewise(out.grad.data, self.data, lambda g,v: g*v))
        out._backward, out.requires_grad = _bw, self.requires_grad or o.requires_grad
        return out

    def __neg__(self): return self * Tensor(_fill(self.shape, -1.0))
    def __sub__(self, o): return self + (-o)
    def __radd__(self, o): return self + o
    def __rmul__(self, o): return self * o

    def backward(self):
        topo, vis = [], set()
        def build(v):
            if v not in vis: vis.add(v); [build(c) for c in v._prev]; topo.append(v)
        build(self)
        self.grad = Tensor(_fill(self.shape, 1.0))
        for v in reversed(topo): v._backward()

# --- Higher-level ops ---

def gelu(x):
    fx, C = _flat(x.data), math.sqrt(2.0/math.pi)
    fo = [0.5*v*(1.0+math.tanh(C*(v+0.044715*v**3))) for v in fx]
    out = Tensor(_unflat(fo, x.shape), (x,), 'gelu')
    def _bw():
        if not x.requires_grad: return
        x._ensure_grad()
        gs = []
        for v, g in zip(fx, _flat(out.grad.data)):
            t = math.tanh(C*(v+0.044715*v**3))
            gs.append(g*(0.5*(1+t)+0.5*v*(1-t*t)*C*(1+3*0.044715*v*v)))
        _acc(x.grad.data, _unflat(gs, x.shape))
    out._backward, out.requires_grad = _bw, x.requires_grad
    return out

def softmax(x):
    flat, N = _flat(x.data), x.shape[-1]
    out_f = []
    for g in range(len(flat)//N):
        grp = flat[g*N:(g+1)*N]; mx = max(grp); e = [math.exp(v-mx) for v in grp]; s = sum(e)
        out_f.extend([v/s for v in e])
    return Tensor(_unflat(out_f, x.shape))

def layernorm(x, gamma, beta, eps=1e-5):
    B, N, D = x.shape
    od = _fill((B, N, D), 0.0)
    for b in range(B):
        for n in range(N):
            row = x.data[b][n]; mu = sum(row)/D; var = sum((v-mu)**2 for v in row)/D; std = math.sqrt(var+eps)
            for d in range(D): od[b][n][d] = gamma.data[d]*(row[d]-mu)/std + beta.data[d]
    out = Tensor(od, (x, gamma, beta), 'ln')
    def _bw():
        if x.requires_grad: x._ensure_grad()
        if gamma.requires_grad: gamma._ensure_grad()
        if beta.requires_grad: beta._ensure_grad()
        for b in range(B):
            for n in range(N):
                row = x.data[b][n]; mu = sum(row)/D; var = sum((v-mu)**2 for v in row)/D; std = math.sqrt(var+eps)
                normed = [(v-mu)/std for v in row]; dout = out.grad.data[b][n]
                for d in range(D):
                    if gamma.requires_grad: gamma.grad.data[d] += dout[d]*normed[d]
                    if beta.requires_grad: beta.grad.data[d] += dout[d]
                if x.requires_grad:
                    dn = [dout[d]*gamma.data[d] for d in range(D)]
                    dm = -sum(dn)/(D*std); dv = sum(-dn[d]*(row[d]-mu)/(2*D*(var+eps)) for d in range(D))
                    for d in range(D): x.grad.data[b][n][d] += dn[d]/std + dm + 2*dv*(row[d]-mu)
    out._backward = _bw
    out.requires_grad = x.requires_grad or gamma.requires_grad or beta.requires_grad
    return out

def smooth_l1(pred, target):
    fp, ft, n = _flat(pred.data), _flat(target.data), len(_flat(pred.data))
    total = sum(0.5*d*d if abs(d:=p-t)<1 else abs(d)-0.5 for p, t in zip(fp, ft))
    out = Tensor([[total/n]], (pred,), 'loss'); out.shape = (1,1)
    def _bw():
        if not pred.requires_grad: return
        pred._ensure_grad()
        gs = [(d/n if abs(d:=p-t)<1 else (1.0 if d>0 else -1.0)/n) for p, t in zip(fp, ft)]
        _acc(pred.grad.data, _unflat(gs, pred.shape))
    out._backward, out.requires_grad = _bw, pred.requires_grad
    return out

# --- Neural network layers ---

class Linear:
    def __init__(self, din, dout):
        s = 1.0/math.sqrt(din)
        self.w = Tensor(_fillfn((din, dout), lambda: random.gauss(0, s)), requires_grad=True)
        self.b = Tensor([0.0]*dout, requires_grad=True); self.b.shape = (dout,)
    def params(self): return [self.w, self.b]
    def __call__(self, x):
        B, N, Di = x.shape; Do = self.w.shape[1]
        od = _fill((B, N, Do), 0.0)
        for b in range(B):
            for n in range(N):
                for j in range(Do):
                    s = self.b.data[j]
                    for i in range(Di): s += x.data[b][n][i]*self.w.data[i][j]
                    od[b][n][j] = s
        out = Tensor(od, (x, self.w, self.b), 'linear')
        def _bw():
            if x.requires_grad: x._ensure_grad()
            if self.w.requires_grad: self.w._ensure_grad()
            if self.b.requires_grad: self.b._ensure_grad()
            for b in range(B):
                for n in range(N):
                    for j in range(Do):
                        g = out.grad.data[b][n][j]
                        if self.b.requires_grad: self.b.grad.data[j] += g
                        for i in range(Di):
                            if self.w.requires_grad: self.w.grad.data[i][j] += x.data[b][n][i]*g
                            if x.requires_grad: x.grad.data[b][n][i] += self.w.data[i][j]*g
        out._backward, out.requires_grad = _bw, True
        return out

class LN:
    def __init__(self, d):
        self.g = Tensor([1.0]*d, requires_grad=True); self.g.shape = (d,)
        self.b = Tensor([0.0]*d, requires_grad=True); self.b.shape = (d,)
    def params(self): return [self.g, self.b]
    def __call__(self, x): return layernorm(x, self.g, self.b)

# --- Transformer components ---

class Attention:
    def __init__(self, dim):
        self.qkv, self.proj, self.dim = Linear(dim, dim*3), Linear(dim, dim), dim
    def params(self): return self.qkv.params() + self.proj.params()
    def __call__(self, x):
        B, N, D = x.shape; qkv = self.qkv(x); scale = 1.0/math.sqrt(D)
        Q = [[[qkv.data[b][n][d] for d in range(D)] for n in range(N)] for b in range(B)]
        K = [[[qkv.data[b][n][D+d] for d in range(D)] for n in range(N)] for b in range(B)]
        V = [[[qkv.data[b][n][2*D+d] for d in range(D)] for n in range(N)] for b in range(B)]
        sc = _fill((B,N,N), 0.0)
        for b in range(B):
            for i in range(N):
                for j in range(N): sc[b][i][j] = sum(Q[b][i][d]*K[b][j][d] for d in range(D))*scale
        sw = softmax(Tensor(sc)).data
        od = _fill((B,N,D), 0.0)
        for b in range(B):
            for i in range(N):
                for d in range(D): od[b][i][d] = sum(sw[b][i][j]*V[b][j][d] for j in range(N))
        ao = Tensor(od, (qkv,), 'attn'); ao.requires_grad = True
        def _bw():
            if not qkv.requires_grad: return
            qkv._ensure_grad()
            for b in range(B):
                for i in range(N):
                    for d in range(D):
                        g = ao.grad.data[b][i][d] if ao.grad else 0
                        for j in range(N):
                            qkv.grad.data[b][j][2*D+d] += sw[b][i][j]*g
                        for d2 in range(D):
                            for j in range(N):
                                qkv.grad.data[b][i][d2] += (ao.grad.data[b][i][d] if ao.grad else 0)*K[b][j][d2]*V[b][j][d]*scale/N
        ao._backward = _bw
        return self.proj(ao)

class Block:
    def __init__(self, dim):
        self.ln1, self.attn, self.ln2 = LN(dim), Attention(dim), LN(dim)
        self.fc1, self.fc2 = Linear(dim, dim*2), Linear(dim*2, dim)
    def params(self): return self.ln1.params()+self.attn.params()+self.ln2.params()+self.fc1.params()+self.fc2.params()
    def __call__(self, x): return (x:=x+self.attn(self.ln1(x))) + self.fc2(gelu(self.fc1(self.ln2(x))))

# --- Vision Transformer ---

class ViT:
    def __init__(self, dim, n_layers, n_patches):
        self.proj = Linear(PATCH_DIM, dim)
        self.pos = Tensor(_fillfn((1, n_patches, dim), lambda: random.gauss(0, 0.02)), requires_grad=True)
        self.blocks, self.norm = [Block(dim) for _ in range(n_layers)], LN(dim)
    def params(self):
        p = self.proj.params()+[self.pos]
        for b in self.blocks: p += b.params()
        return p + self.norm.params()
    def __call__(self, patches, indices=None):
        tok = self.proj(patches); B, N, D = tok.shape
        pos = Tensor([[[self.pos.data[0][idx][d] for d in range(D)] for idx in (indices or range(N))] for _ in range(B)])
        tok = tok + pos
        for b in self.blocks: tok = b(tok)
        return self.norm(tok)

# --- Predictor: narrow transformer bridging context -> target ---

class Predictor:
    def __init__(self, enc_dim, pred_dim, n_layers, n_patches):
        self.inp = Linear(enc_dim, pred_dim)
        self.mask_tok = Tensor(_fillfn((1,1,pred_dim), lambda: random.gauss(0, 0.02)), requires_grad=True)
        self.pos = Tensor(_fillfn((1, n_patches, pred_dim), lambda: random.gauss(0, 0.02)), requires_grad=True)
        self.blocks, self.norm, self.out_proj = [Block(pred_dim) for _ in range(n_layers)], LN(pred_dim), Linear(pred_dim, enc_dim)
    def params(self):
        p = self.inp.params()+[self.mask_tok, self.pos]
        for b in self.blocks: p += b.params()
        return p + self.norm.params() + self.out_proj.params()
    def __call__(self, ctx_tok, ctx_idx, tgt_idx):
        B, Nc, _ = ctx_tok.shape; D = self.inp.w.shape[1]; Nt = len(tgt_idx)
        ctx = self.inp(ctx_tok)
        for b in range(B):
            for i, idx in enumerate(ctx_idx):
                for d in range(D): ctx.data[b][i][d] += self.pos.data[0][idx][d]
        tgt = [[[self.mask_tok.data[0][0][d]+self.pos.data[0][idx][d] for d in range(D)] for idx in tgt_idx] for _ in range(B)]
        cat = _fill((B, Nc+Nt, D), 0.0)
        for b in range(B):
            for n in range(Nc): cat[b][n] = ctx.data[b][n][:]
            for n in range(Nt): cat[b][Nc+n] = tgt[b][n][:]
        tokens = Tensor(cat, (ctx,), 'cat'); tokens.requires_grad = True
        def _bw():
            if ctx.requires_grad:
                ctx._ensure_grad()
                for b in range(B):
                    for n in range(Nc):
                        for d in range(D): ctx.grad.data[b][n][d] += (tokens.grad.data[b][n][d] if tokens.grad else 0)
        tokens._backward = _bw
        for blk in self.blocks: tokens = blk(tokens)
        tokens = self.norm(tokens)
        pred = Tensor([[[tokens.data[b][Nc+n][d] for d in range(D)] for n in range(Nt)] for b in range(B)], (tokens,), 'ext')
        pred.requires_grad = True
        def _bw2():
            if tokens.requires_grad:
                tokens._ensure_grad()
                for b in range(B):
                    for n in range(Nt):
                        for d in range(D): tokens.grad.data[b][Nc+n][d] += (pred.grad.data[b][n][d] if pred.grad else 0)
        pred._backward = _bw2
        return self.out_proj(pred)

# --- EMA update: the collapse prevention trick ---

def ema_update(online, target, decay):
    for po, pt in zip(online, target):
        fo, ft = _flat(po.data), _flat(pt.data)
        pt.data = _unflat([decay*t+(1-decay)*o for o, t in zip(fo, ft)], pt.shape)

# --- Adam optimizer from scratch ---

class Adam:
    def __init__(self, params, lr=3e-4):
        self.params, self.lr, self.t = params, lr, 0
        self.m = [_fill(p.shape, 0.0) for p in params]
        self.v = [_fill(p.shape, 0.0) for p in params]
    def zero_grad(self):
        for p in self.params: p.grad = None
    def step(self):
        self.t += 1; b1, b2, eps = 0.9, 0.999, 1e-8
        for i, p in enumerate(self.params):
            if p.grad is None: continue
            fg, fm, fv, fp = _flat(p.grad.data), _flat(self.m[i]), _flat(self.v[i]), _flat(p.data)
            for j in range(len(fp)):
                fm[j] = b1*fm[j]+(1-b1)*fg[j]; fv[j] = b2*fv[j]+(1-b2)*fg[j]**2
                fp[j] -= self.lr*(fm[j]/(1-b1**self.t))/(math.sqrt(fv[j]/(1-b2**self.t))+eps)
            p.data, self.m[i], self.v[i] = _unflat(fp, p.shape), _unflat(fm, p.shape), _unflat(fv, p.shape)

# --- Config & synthetic data ---

IMG, PATCH = 6, 2           # 6x6 images, 2x2 patches -> 3x3 = 9 patches
GRID, N_PATCHES = IMG//PATCH, (IMG//PATCH)**2
PATCH_DIM, ENC_DIM, PRED_DIM = PATCH*PATCH*3, 16, 8

def make_image():
    bg, fg = [random.random()*0.3 for _ in range(3)], [random.random()*0.7+0.3 for _ in range(3)]
    r1, r2 = sorted(random.sample(range(IMG), 2)); c1, c2 = sorted(random.sample(range(IMG), 2))
    img = [[[bg[c] for c in range(3)] for _ in range(IMG)] for _ in range(IMG)]
    for r in range(r1, r2+1):
        for c in range(c1, c2+1): img[r][c] = fg[:]
    return img

def to_patches(img):
    patches = []
    for gr in range(GRID):
        for gc in range(GRID):
            patches.append([img[gr*PATCH+dr][gc*PATCH+dc][ch] for dr in range(PATCH) for dc in range(PATCH) for ch in range(3)])
    return patches

def make_batch(bs): return Tensor([to_patches(make_image()) for _ in range(bs)])

def sample_masks():
    idx = list(range(N_PATCHES)); random.shuffle(idx); k = min(3, N_PATCHES//3)
    return sorted(idx[k:]), sorted(idx[:k])

def select(t, idx):
    B, D = t.shape[0], t.shape[2]
    return Tensor([[[t.data[b][i][d] for d in range(D)] for i in idx] for b in range(B)])

# --- Training loop: the complete JEPA algorithm ---

def train():
    print(f"microjepa | zero deps | {IMG}x{IMG} imgs | {N_PATCHES} patches | enc={ENC_DIM} pred={PRED_DIM}")
    ctx_enc = ViT(ENC_DIM, 1, N_PATCHES); tgt_enc = ViT(ENC_DIM, 1, N_PATCHES)
    pred = Predictor(ENC_DIM, PRED_DIM, 1, N_PATCHES)
    for pc, pt in zip(ctx_enc.params(), tgt_enc.params()):
        pt.data = _unflat(_flat(pc.data)[:], pc.shape); pt.requires_grad = False
    train_p = ctx_enc.params() + pred.params()
    for p in train_p: p.requires_grad = True
    opt = Adam(train_p)
    print(f"trainable params: {sum(len(_flat(p.data)) for p in train_p):,}")
    for epoch in range(20):
        total = 0.0
        for _ in range(10):
            batch = make_batch(4); batch.requires_grad = True
            ci, ti = sample_masks()
            target_repr = select(tgt_enc(batch), ti)              # target encoder: full image
            context_repr = ctx_enc(select(batch, ci), ci)         # context encoder: masked
            pred_repr = pred(context_repr, ci, ti)                # predictor: context -> target
            loss = smooth_l1(pred_repr, target_repr)              # L1 loss in repr space
            opt.zero_grad(); loss.backward(); opt.step()          # backprop
            ema_update(ctx_enc.params(), tgt_enc.params(), 0.99)  # EMA
            total += loss.item()
        print(f"epoch {epoch+1:3d}/20 | loss {total/10:.4f}")
    print("\ndone. the representations are learned. everything else is just efficiency.")

if __name__ == '__main__':
    train()

Here is an interactive diagram of the full I-JEPA architecture. Click the patches in the grid to toggle between context (blue) and target (red):

How it works

The code above is the whole thing. Let me walk through the key parts.

Tensors and autograd. We need tensors that track their operations so we can compute gradients. Every + or * creates a new tensor with a _backward function that knows how to propagate gradients back. Backpropagation is just a topological sort of this graph, then walking backwards. Same thing PyTorch does, but in 6 lines.

Here is what that computation graph looks like in practice. Click "Run backward" to watch gradients flow from the loss back through each operation:

Here is an interactive visualization of that pipeline. Hover over the image grid to select a patch, or click pixels to change their color:

Two encoders. There is a context encoder that sees only the visible patches and gets gradients. And a target encoder that sees all patches but gets no gradients — it is updated as a slow moving average of the context encoder. This prevents collapse. Without it, the model could just map everything to the same vector and get zero loss.

The predictor. A small transformer (half the width of the encoder) that takes context representations and predicts what the target representations should be. It is small on purpose — we want the encoder to do the heavy lifting, not the predictor.

The loss. Smooth L1 in representation space. Not pixel space. The model learns to predict what matters, not every visual detail.

Why this matters

I-JEPA on images is step one. The roadmap goes:

I-JEPA — predict masked image patches in representation space
V-JEPA — same thing but with video frames
V-JEPA 2 — add actions, so you can predict what happens when you do something. That is a primitive world model.
H-JEPA — hierarchical predictions at different time scales. The full world model.

The core idea stays the same at every step: predict representations, not raw data. That is what LeCun is building at AMI Labs.

Here is an interactive timeline showing how each stage builds toward the full world model:

References

LeCun, A Path Towards Autonomous Machine Intelligence, 2022
Assran et al., Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture, 2023
Bardes et al., V-JEPA, 2024

More coming soon.

[WIP] microJEPA — JEPA From Scratch in Pure Python