lejepa-loss

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__()
        # build t, phi, weights buffers ... (full impl in REFERENCE.md)

    def forward(self, proj):  # (B, V, D) -> scalar
        ...

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):  # (B, V, D), key -> scalar
        ...

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.