pocket-jepa

Pocket JEPA

Pocket JEPA is a small-compute joint-embedding predictive architecture for 8x8 visual worlds. The online network sees a partially hidden image and predicts the latent representation produced by an exponential-moving-average target encoder on the complete image. It never reconstructs pixels and never uses class labels during pretraining.

The experiment measures whether the learned J-space becomes useful under a strict ten-labels-per-class linear-probe budget. An architecture-matched random encoder and a raw-pixel probe are retained as controls. A separate masked-view evaluation tests whether predictive pretraining creates occlusion-tolerant features.

This is a compact JEPA-inspired experiment, not an implementation or reproduction claim for any specific large-scale I-JEPA result.

Verified results

After 260 unlabeled pretraining epochs on 1,347 images, a linear classifier received only ten labels per class:

The deployed encoder has 14,848 parameters. Its training-only predictor has 4,192 parameters. Pocket JEPA improved the masked-view result by 30.22 points over the random encoder and 13.56 points over raw pixels.

Reproduce

uv run python projects/pocket-jepa/train.py