| name | lejepa |
| description | Entry point for LeJEPA (Lean Joint-Embedding Predictive Architecture) and SIGReg (Sketched Isotropic Gaussian Regularization). Explains the paper's recipe (backbone, projector, SIGReg loss, invariance loss with a single trade-off lambda) and the four child skills that instantiate it at different scopes in either PyTorch or equinox+JAX. Use when the user mentions LeJEPA broadly, asks about the paper, is unsure which child skill applies, or wants background before picking a scope. |
LeJEPA
LeJEPA is a self-supervised learning recipe with one trade-off hyperparameter and no heuristics. No stop-gradient, no teacher-student, no register tokens. The core innovation is SIGReg, a loss that pushes the projected embeddings toward an isotropic standard Gaussian by measuring the squared characteristic-function distance along random slicing directions. Paper: https://arxiv.org/abs/2511.08544. Official repo: https://github.com/galilai-group/lejepa (PyTorch only).
The full loss is L = lam * SIGReg(proj) + (1 - lam) * inv(proj), where inv is the mean squared distance between each view's projection and the per-image mean projection.
Architecture
images (B, V, C, H, W)
-> backbone -> emb (B, V, D_emb) used for downstream eval
-> projector -> proj (B, V, D_proj) used for SIGReg + invariance
-> loss = lam * SIGReg(proj) + (1 - lam) * inv(proj)
The backbone is interchangeable. The paper validates the recipe across 60+ architectures. The projector is fixed at MLP(emb_dim, [2048, 2048, proj_dim]) with BatchNorm1d.
Child skills
| Skill | Scope |
|---|
lejepa-loss | SIGReg loss module only. Bring your own encoder. |
lejepa-encoder | Backbone (default or BYO) + projector + SIGReg as separate objects, MINIMAL.md style. |
lejepa-model | Bundled LeJEPAModel class with embed(x) and loss(views) methods. |
lejepa-pretrain | Single-GPU reference pretraining script with optimizer, schedule, bf16, and optional online probe. |
If the user's intent maps unambiguously to one of these, invoke that child skill directly via the Skill tool. Otherwise ask which scope they want.
Paper-default knobs
These are validated by the paper. Do not change without an explicit reason.
| Knob | Value |
|---|
lam | 0.02 |
| projector hidden | [2048, 2048, proj_dim] with BatchNorm1d (load-bearing) |
proj_dim | 128 |
SIGReg knots | 17 |
SIGReg t_max | 3.0 |
SIGReg num_slices | 256 |
| optimizer | AdamW |
| weight decay | 5e-2 for ViT, 5e-4 for ResNet |
| schedule | linear warmup of one epoch then cosine annealing to lr/1000 |
| precision | bf16 |
Framework note
PyTorch uses timm for backbones. The JAX path is equinox only (no flax, no haiku) per the user's standing instruction. SIGReg and the bundled model carry an explicit key argument in equinox because the random slicing matrix is sampled per forward.
For derivations, the connection to the paper's normality-test family, and pointers to the alternative SIGReg variants in lejepa.univariate and lejepa.multivariate, see references/REFERENCE.md. The script at scripts/sanity_check.py verifies that SIGReg discriminates Gaussian from non-Gaussian samples in whichever framework is installed.
