Metadata-Version: 2.4
Name: jepa-science
Version: 0.1.0
Summary: Sketched Isotropic Gaussian Regularization (SIGReg) from LeJEPA — a clean, standalone PyTorch primitive.
Project-URL: Homepage, https://github.com/gabsalvo/jepa_science
Project-URL: Paper, https://arxiv.org/abs/2511.08544
Project-URL: Reference implementation, https://github.com/rbalestr-lab/lejepa
Author: jepa_science contributors
License-Expression: Apache-2.0
License-File: LICENSE
License-File: NOTICE
Keywords: jepa,lejepa,pytorch,representation-learning,self-supervised-learning,sigreg
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.9
Requires-Dist: numpy>=1.23
Requires-Dist: torch>=2.0
Provides-Extra: demo
Requires-Dist: matplotlib>=3.6; extra == 'demo'
Requires-Dist: torchvision>=0.15; extra == 'demo'
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: scipy>=1.10; extra == 'dev'
Description-Content-Type: text/markdown

# jepa_science

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

[![tests](https://img.shields.io/badge/tests-passing-brightgreen)](tests)
[![license](https://img.shields.io/badge/license-Apache--2.0-blue)](LICENSE)

## 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`](NOTICE).

## Install

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

From source:

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

## Quickstart

```python
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 `λ`:

```python
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`:

![SIGReg drives embeddings to an isotropic Gaussian](docs/not_collapse.png)

Run the full LeJEPA demo on CIFAR-10:

```bash
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

```bash
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:

```bibtex
@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`](CITATION.cff) to cite this implementation as well.

## License

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