I-JEPA on histology images

PathMNIST tiles with I-JEPA context and target block masks — I-JEPA masking on PathMNIST tiles. Left: image. Middle: the context the encoder sees (132 patches). Right: four target blocks the predictor must predict in latent space.

Why

Ball-in-cup taught the action-conditioned side of JEPA. The other big use case is a self-supervised visual backbone for domains with many images and few labels, like pathology. I wanted to see it work at a realistic resolution on my own hardware: PathMNIST at 224 px, ViT encoder, block masking, no EMA target.

Engineering on a 24 GB Mac

The 224 px split ships as a compressed .npz, and loading it pulls ~13.5 GB into RAM, which put training straight into swap. Decompressing each split once to an uncompressed .npy and memory-mapping it brought data RAM to roughly zero. It also let multiple data-loader workers share pages instead of copying.

Probing the frozen encoder

Frozen features, 9 tissue classes, 7,180 test images. The same probes run on a randomly initialised encoder as a baseline:

Random-init encoderlinear probe

74.3%

I-JEPA, epoch 1linear probe

77.2%

I-JEPA, epoch 4linear probe

78.6%

Random-init encoder5-NN

63.1%

I-JEPA, epoch 45-NN

78.1%

The linear-probe gain over random init is real but modest. The nearest-neighbour gain (+15 points) is the clearer sign that the representation organises tissue types. The features' effective rank also roughly doubled between epochs 1 and 4.

Lesson A random-init baseline is essential. A linear probe on random ViT features already reaches 74%, so an absolute number on its own says little about what pretraining did. Longer runs weren't automatically better either: probe accuracy has to be tracked across checkpoints, not just read at the end.

2-D embedding of mean-pooled encoder features coloured by tissue class — Mean-pooled features from a later checkpoint, projected to 2-D. Background, adipose and debris separate cleanly. The epithelial and stromal classes stay entangled.