PhysRel-JEPA checkpoints

Checkpoints for PhysRel-JEPA, an object-centric JEPA video world model that routes object interactions through a narrow gated edge bottleneck instead of the predictor's self-attention.

Files

file	what	notes
savi.pt	frozen SAVi encoder, 64x64, N=7 slots, dim 128	trained 30 epochs on 5000 CLEVRER videos (8x H100, DDP)
pred_ocjepa.pt	OC-JEPA predictor (no masking, no edges)	5000-clip, 15-epoch run
pred_cjepa.pt	C-JEPA predictor (object masking, no edges)	5000-clip, 15-epoch run
pred_physrel.pt	PhysRel-JEPA predictor (masking + edge bottleneck + property memory)	5000-clip, 15-epoch run

Each .pt is a torch.save dict with keys model (state_dict), opt (optimizer state), and metadata. Predictor checkpoints include optimizer state, so they are resumable.

Loading

import torch
from physrel_jepa.config import load_config
from physrel_jepa.models.world_model import PhysRelJEPA
from physrel_jepa.models.encoder_savi import SAViEncoder

enc = SAViEncoder(num_slots=7, slot_dim=128, resolution=64)
enc.load_state_dict(torch.load("savi.pt", map_location="cpu")["model"])

cfg = load_config("configs/predictor_physrel.yaml")
pred = PhysRelJEPA(cfg.model)
pred.load_state_dict(torch.load("pred_physrel.pt", map_location="cpu")["model"])

Results (held-out future latent NMSE, lower is better)

5000 clips, 15 epochs, 5 seeds: OC-JEPA 0.00045, C-JEPA 0.00064, PhysRel-JEPA 0.00053. The edge bottleneck beats masking-only (PhysRel < C-JEPA) at this scale, while no-masking OC-JEPA is best under this budget. At 1500 clips / 200 epochs PhysRel-JEPA is best overall (0.00042, wins C-JEPA 5/5 seeds). See the repo for the full writeup and the gate-collapse ablation.