HugoPlanchat/self-supervised-to-jepa

Self-Supervised to JEPA — pretrained ViT-Tiny checkpoints

Pre-training checkpoints for self-supervised-to-jepa: a controlled benchmark of ten self-supervised methods re-implemented from scratch on a single shared ViT-Tiny (patch 4, d=128, 6 layers, 4 heads, ≈1.2M params), pre-trained on STL-10 (unlabeled, 100 epochs) and evaluated on CIFAR-10. These are small research models meant to reproduce this benchmark, not general-purpose backbones — useful to regenerate every figure and downstream number without re-training on a GPU.

What is in this repo (for now)

The 10 pre-training checkpoints, one per method. Each *_pretrain.pt holds, besides the SSL weights, the full monitoring history recorded during training:

key	content
model	full SSL model state_dict (all branches)
encoder	the shared ViT-Tiny trunk state_dict (the reusable part)
losses (or hist["loss"])	per-epoch training loss
knn_history	list of (epoch, weighted k-NN accuracy) pairs on CIFAR-10
rank_history	list of (epoch, effective rank) pairs (MAE stores lr_history instead)
method-specific loss terms	the objective's component breakdown: SimCLR pos/neg, MoCo lookup, BYOL feature std, CPC row_acc, VICReg inv/var/cov, LeJEPA pred/sig; DINOv3 and VISReg bundle theirs in a hist dict

Linear-probe and fine-tune checkpoints will be added later.

Downstream results (CIFAR-10, %)

Method	k-NN	Linear probe	Fine-tune
MAE	46.77	49.05	83.57
CPC	39.93	36.50	73.67
SimCLR	74.17	72.92	86.75
MoCo	74.13	69.68	84.95
BYOL	73.07	73.76	86.23
DINOv3	59.13	54.98	81.90
VICReg	73.04	73.08	86.65
I-JEPA	57.71	56.27	83.25
LeJEPA	65.06	62.00	81.83
VISReg	71.57	68.87	85.48

Supervised ViT-Tiny (from scratch, reference): 78.69 fine-tune.

Files

Checkpoints mirror the repo layout, so you can drop them straight into ssl/checkpoints/:

04_mae/mae_pretrain.pt        09_dinov3/dinov3_pretrain.pt
05_cpc/cpc_pretrain.pt        10_VICReg/vicreg_pretrain.pt
06_simclr/simclr_pretrain.pt  11_IJEPA/ijepa_pretrain.pt
07_moco/moco_pretrain.pt      12_LEJEPA/lejepa_pretrain.pt
08_byol/byol_pretrain.pt      13_VISReg/visreg_pretrain.pt

Load a checkpoint

The .pt files are state_dicts: you need the model classes from the GitHub repo to instantiate the architecture, then pour the weights in.

import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download("HugoPlanchat/self-supervised-to-jepa", "08_byol/byol_pretrain.pt")
ck = torch.load(path, map_location="cpu", weights_only=False)

# ck["encoder"] is the shared ViT-Tiny trunk; rebuild ViTEncoder() from the notebook, then:
# encoder.load_state_dict(ck["encoder"])

print(ck.keys())                 # weights + monitoring histories
print(ck["knn_history"][-1])     # e.g. (100, 0.7307)

To make the notebooks restart-and-run-all instant, download each file into ssl/checkpoints/<method>/ and the load-or-train guard will reload instead of training.

Links

Code, report and full write-up: https://github.com/hugoplanchat/SNU-Internship