| name | lejepa-model |
| description | Emits a bundled `LeJEPAModel` class in PyTorch or equinox+JAX with `embed(x)` for downstream evaluation and `loss(views)` for training. The class owns backbone, projector, and SIGReg internally, and exposes lambda as a constructor argument. The online linear probe lives outside the model. Use when the user wants a single-object LeJEPA API. For separately wired pieces (backbone, projector, SIGReg as distinct objects), use `lejepa-encoder` instead. |
LeJEPA: bundled model
LeJEPAModel bundles the encoder and SIGReg into one object with a two-method API. The probe is not part of the model. The training loop owns the probe because the probe consumes detached embeddings and trains on a separate label stream.
API contract
model.embed(x: (B, C, H, W)) -> (B, D_emb)
model.loss(views: (B, V, C, H, W)) -> (scalar, {"sigreg": scalar, "inv": scalar})
For equinox the signature is model.loss(views, key) -> (scalar, aux) because SIGReg owns the random slicing.
Defaults
| Knob | Default |
|---|
lam | 0.02 |
emb_dim | 512 |
proj_dim | 128 |
sigreg_knots | 17 |
sigreg_t_max | 3.0 |
sigreg_num_slices | 256 |
Lambda is a constructor argument. Mutating it after construction is unsupported; build a new model.
Loss decomposition
inv = mean_b mean_v || proj_{b,v} - mean_v proj_{b,.} ||^2 # invariance
sig = SIGReg(proj) # isotropy
total = lam * sig + (1 - lam) * inv
The aux dict returned by loss contains sigreg and inv as detached scalars (PyTorch) or stop-gradient scalars (JAX), suitable for logging without leaking gradient.
JIT placement
For equinox, model.loss is not JIT-decorated. JIT at the training-step level instead, so the optimizer step, gradient computation, and loss are compiled as one graph. The PyTorch model also avoids @torch.compile decoration internally for the same reason.
For full implementations of both, the lambda edge cases, the BatchNorm-state threading for equinox, and the testing patterns, see references/REFERENCE.md. The script at scripts/sanity_check.py verifies the API contract, the lambda edge cases, and gradient flow to backbone parameters.
