File size: 3,651 Bytes
714d1ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | ---
library_name: pytorch
tags:
- self-supervised-learning
- vision-transformer
- jepa
- representation-learning
- pytorch
datasets:
- stl10
- cifar10
pipeline_tag: image-feature-extraction
---
# Self-Supervised to JEPA — pretrained ViT-Tiny checkpoints
Pre-training checkpoints for **[self-supervised-to-jepa](https://github.com/hugoplanchat/SNU-Internship)**: 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_dict`s: you need the model classes from the [GitHub repo](https://github.com/hugoplanchat/SNU-Internship) to instantiate the architecture, then pour the weights in.
```python
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
|