| name | lejepa-loss |
| description | Emits the SIGReg (Sketched Isotropic Gaussian Regularization) loss module from LeJEPA in PyTorch or equinox+JAX. SIGReg is a sliced characteristic- function test for standard Gaussianity. Use when the user wants only the SIGReg loss as a drop-in module, references LeJEPA's regularizer, or needs to add SIGReg to an existing pretraining pipeline. For the full LeJEPA recipe (backbone + projector + invariance), use `lejepa-model` or `lejepa-encoder` instead. |
LeJEPA: SIGReg loss
SIGReg is a loss module that pushes a batch of projected embeddings toward an isotropic standard Gaussian. The contract is sigreg(proj) -> scalar where proj has shape (B, V, D). For equinox the signature is sigreg(proj, key) -> scalar because the random slicing matrix is sampled per call.
Defaults
| Knob | Default |
|---|
knots | 17 |
t_max | 3.0 |
num_slices | 256 |
These reproduce MINIMAL.md. The sample count in the chi-squared scaling is B * V, not proj.shape[-2], so the layout (B, V, D) is correct regardless of internal transposes.
PyTorch skeleton
class SIGReg(nn.Module):
def __init__(self, knots=17, t_max=3.0, num_slices=256):
super().__init__()
def forward(self, proj):
...
Equinox skeleton
class SIGReg(eqx.Module):
t: jax.Array
phi: jax.Array
weights: jax.Array
num_slices: int = eqx.field(static=True)
def __call__(self, proj, key):
...
For full implementations of both, the quadrature derivation, the layout fix vs MINIMAL.md, and the testing patterns, see references/REFERENCE.md. The script at scripts/sanity_check.py verifies Gaussian/non-Gaussian discrimination, the shape contract, and the gradient finite-difference check.
