jepa-science 0.1.0

Sketched Isotropic Gaussian Regularization (SIGReg) from LeJEPA — a clean, standalone PyTorch primitive.

jepa_science

A clean, standalone PyTorch implementation of SIGReg — the Sketched Isotropic Gaussian Regularization loss from LeJEPA (Balestriero & LeCun, 2025, arXiv:2511.08544) — plus a minimal, heavily annotated LeJEPA training demo on CIFAR-10.

What is JEPA?

A Joint-Embedding Predictive Architecture (JEPA) learns representations by predicting the embedding of one view of an input from the embedding of another, in latent space rather than in pixel space. The danger is collapse: the encoder can trivially satisfy the prediction objective by mapping everything to a constant. Existing methods fight collapse with heuristics — stop-gradients, teacher/student networks, EMA targets, schedulers.

What is SIGReg?

LeJEPA proves that the optimal embedding distribution for minimizing downstream risk is the standard isotropic Gaussian N(0, I), and introduces SIGReg to push embeddings toward it — the single, principled anti-collapse term. SIGReg works by a sketch:

1. Cramér–Wold (paper Lemma 3): two distributions are equal iff all their 1-D projections are equal. So matching N(0, I) reduces to making every projection aᵀz look like N(0, 1).
2. Slice: sample K random unit directions, project the batch onto each.
3. Test: on each 1-D projection apply the Epps–Pulley statistic — the weighted L² distance between the sample's empirical characteristic function and that of N(0, 1), computed by trapezoidal quadrature.

The loss is the mean of that statistic over the slices. It is minimized exactly when the embeddings have zero mean and identity covariance and are Gaussian along every direction — which is precisely why it prevents collapse. It has linear time and memory cost and bounded, stable gradients (paper Theorem 4).

Why this exists

The reference repo bundles SIGReg inside a full training framework. This package pulls it out as a clean, dependency-light primitive (torch + numpy) — a SciPy-style building block you can drop into any JEPA/SSL pipeline — and pairs it with a runnable, readable demo so you can watch embeddings Gaussianize.

Provenance: this is a clean-room implementation written from the paper's mathematics. It does not copy code from the reference repository (which ships under a NonCommercial CC BY-NC 4.0 license, incompatible with this package's Apache-2.0). Methods/algorithms aren't copyrightable; the code here is original. See NOTICE.

Install

pip install jepa-science            # the primitive (torch + numpy)
pip install "jepa-science[demo]"    # + torchvision & matplotlib for the CIFAR-10 demo

From source:

git clone https://github.com/gabsalvo/jepa_science && cd jepa_science
pip install -e ".[demo,dev]"

Quickstart

import torch
from jepa_science import sigreg_loss, SIGReg

z = torch.randn(256, 64, requires_grad=True)   # (batch, embedding_dim)

# Functional form:
loss = sigreg_loss(z, num_slices=1024)         # scalar; minimized at N(0, I)

# Module form (precomputes quadrature buffers, moves with .to(device)):
criterion = SIGReg(num_slices=1024)
loss = criterion(z)
loss.backward()

Use it in a LeJEPA-style objective with the single trade-off hyperparameter λ:

embedding_loss = lam * sigreg_loss(z) + (1 - lam) * prediction_loss   # λ ≈ 0.05

sigreg_loss returns only the SIGReg statistic — λ and the JEPA prediction term live in your training loop (see lejepa_lite/), keeping this a standalone primitive.

Watch it not collapse

Minimizing only SIGReg turns an anisotropic, partially-collapsed batch into N(0, I): the covariance condition number → 1, the smallest variance is pushed up away from zero (anti-collapse), and a random 1-D projection matches the N(0,1) density. Reproduce with python docs/make_figure.py:

Run the full LeJEPA demo on CIFAR-10:

python -m lejepa_lite.train                 # downloads CIFAR-10, trains, saves plots
python -m lejepa_lite.train --smoke         # tiny CPU run, no download (also CI path)

Tests

pytest            # CPU-only, deterministic

The suite proves correctness against an independent numpy oracle and a hand-verified golden value (cross-checked against a high-accuracy scipy.quad integral — see SPEC.md §6.2), asserts the Theorem 4 gradient bound numerically, and demonstrates collapse avoidance on a toy tensor.

Fidelity notes (this implementation vs. the paper/reference)

We follow the authors' runnable reference numerics where the paper text and code diverge, and flag every divergence here so nothing is silent:

Knob	Paper text	Reference code	This package (default)
Quadrature window	w(t)=e^{-t²/σ²}, σ=1 (e^{-t²})	e^{-t²/2} (σ=√2)	e^{-t²/2}
Integration domain	[-5, 5]	[0, 3] folded (=±3)	[0, 3] folded
Quadrature knots	17	17	17
Slices K	1024	256 (MINIMAL.md)	1024
λ (trade-off)	0.05	0.02 (MINIMAL.md)	0.05 (in the demo)

Other deliberate choices: SIGReg does not internally center/whiten embeddings (doing so would mask collapse — any full-rank covariance would pass); the demo uses CIFAR-10 and a small CNN (the reference used Imagenette + ViT-small, so this is not a benchmark-parity claim); and DDP/all_reduce is out of scope for v1.

Citation

If you use this software, please cite the LeJEPA paper:

@misc{balestriero2025lejepa,
  title  = {LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics},
  author = {Randall Balestriero and Yann LeCun},
  year   = {2025},
  eprint = {2511.08544},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url    = {https://arxiv.org/abs/2511.08544}
}

See CITATION.cff to cite this implementation as well.

License

Apache-2.0. The implemented method is from the LeJEPA paper; this code is an independent expression of it (see NOTICE).