| name | lejepa-encoder |
| description | Emits a LeJEPA encoder (backbone, 3-layer MLP projector with BatchNorm, and SIGReg loss as separate objects) in PyTorch or equinox+JAX, mirroring the MINIMAL.md assembly. Bring your own backbone is the primary API; a default backbone is included (timm ViT-S/8 for PyTorch, a minimal self-contained equinox ViT for JAX). Use when the user wants the three pieces as separate objects and intends to write their own training loop. For a single bundled object, use `lejepa-model` instead. |
LeJEPA: encoder
The encoder produces both raw backbone embeddings (for downstream eval) and projector outputs (for the SSL loss). The contract is
encoder(x: (B, V, C, H, W)) -> (emb: (B, V, D_emb), proj: (B, V, D_proj))
The projector is MLP(D_emb, [2048, 2048, D_proj]) with BatchNorm1d. BatchNorm is load-bearing. Substituting LayerNorm degrades performance materially per the paper's ablations.
SIGReg is constructed separately and composed by the user at training time. The encoder does not own the loss.
Default backbones
| Framework | Default backbone |
|---|
| PyTorch | timm.create_model("vit_small_patch8_224", pretrained=False, num_classes=D_emb, drop_path_rate=0.1, img_size=128) |
| equinox+JAX | a minimal self-contained ViT defined inside the skill's REFERENCE.md (patch embed, transformer blocks, cls-token head, no register tokens, no pretrained weights) |
To plug in a different backbone, pass it to Encoder.__init__ as the backbone argument. Any module producing (B*V, D_emb) from (B*V, C, H, W) will work.
Defaults
| Knob | Default |
|---|
emb_dim | 512 |
proj_dim | 128 |
| projector hidden | [2048, 2048] |
| projector norm | BatchNorm1d (do not substitute) |
PRNG note
The PyTorch encoder uses the module-global RNG (no key arg). The equinox encoder is deterministic in __call__ because the random slicing is owned by SIGReg, not the encoder. When composing them, the SIGReg __call__ takes a key explicitly.
For full implementations of both, including the minimal equinox ViT and the projector wiring, see references/REFERENCE.md. The script at scripts/sanity_check.py verifies output shapes, BatchNorm presence in the projector, and that a stub backbone substitutes cleanly.
