| name | lejepa-pretrain |
| description | Emits a single-GPU LeJEPA pretraining script in PyTorch or equinox+JAX, matching MINIMAL.md. Includes AdamW with separate weight decay groups, linear warmup plus cosine annealing schedule, bf16 mixed precision, the lightweight V-augmented-views augmentation pipeline, and an optional online linear probe controlled by a flag. Use when the user wants a runnable pretraining script with paper-default hyperparameters. Does not ship DDP, checkpointing, or full DINO multi-crop (those are documented as upgrade paths in the reference). |
LeJEPA: pretraining
A reference single-GPU pretraining script in each framework. The PyTorch version is a faithful adaptation of MINIMAL.md with the probe optional and W&B optional. The equinox version is an optax-based port. Both follow the lightweight V-views augmentation, not DINO multi-crop.
Annotated knobs
Paper defaults. Do not change without a stated reason.
| Knob | Value |
|---|
lam | 0.02 |
proj_dim | 128 |
| projector hidden | [2048, 2048] with BatchNorm1d |
SIGReg knots / t_max / num_slices | 17 / 3.0 / 256 |
| optimizer | AdamW |
| weight decay | 5e-2 for ViT backbones, 5e-4 for ResNet |
| schedule | linear warmup of one epoch then cosine annealing to lr / 1000 |
| precision | bf16 (torch.amp.autocast or jnp.bfloat16 cast) |
| probe optimizer | AdamW, lr=1e-3, wd=1e-7 |
User-tunable. Set for your run.
| Knob | Notes |
|---|
dataset | Any image dataset producing (C, H, W) after transform |
backbone | timm name (PyTorch) or equinox module (JAX) |
batch_size | MINIMAL.md uses 256 |
lr | MINIMAL.md uses 2e-3 |
epochs | MINIMAL.md case study uses 800 |
V | Number of augmented views per image; MINIMAL.md uses 4 |
probe_classes | None to disable the probe, else number of classes |
Augmentation
Lightweight V augmented copies per image, all at the same resolution (default 128x128), with RandomResizedCrop(scale=(0.08, 1.0)), ColorJitter, RandomGrayscale, GaussianBlur, RandomSolarize, and RandomHorizontalFlip. This is MINIMAL.md's pipeline. The DINO-style multi-crop (2 global at 224, 6 local at 98) is documented as an upgrade path in the reference but not shipped.
Excluded by design
- No DDP or
torch.distributed setup. Wrap the model yourself if multi-GPU.
- No checkpoint manager. Save with the framework's standard pattern at the boundaries you choose.
- No DINO-style multi-crop.
Online probe
The probe is nn.Sequential(nn.LayerNorm(emb_dim), nn.Linear(emb_dim, probe_classes)) (PyTorch) or the equinox equivalent. It trains on model.embed(view_0).detach() (or jax.lax.stop_gradient in JAX) with its own AdamW param group. Set probe_classes=None to skip the probe entirely.
For full PyTorch and equinox pretraining scripts, the optimizer param-group construction, the optax schedule recipe, the bf16 patterns in each framework, and the testing patterns, see references/REFERENCE.md. The script at scripts/sanity_check.py runs a single training step on synthetic data and verifies the probe-disable flag.
