torch_world.losses¶

SIGReg ¶

SIGReg(
    num_slices: int = 1024,
    knots: int = 17,
    t_max: float = 3.0,
)

Bases: Module

Sketched Isotropic Gaussian Regularization.

Projects embeddings on num_slices random unit directions and averages the Epps-Pulley statistic, which measures each 1D projection's distance from a standard normal by comparing characteristic functions on knots trapezoid points in [0, t_max].

Balestriero & LeCun, "LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics", 2025. Paper: https://arxiv.org/abs/2511.08544 Code: https://github.com/galilai-group/lejepa

forward ¶

forward(
    x: Tensor, generator: Generator | None = None
) -> Tensor

x is (..., N, D) with samples on dim -2; returns the statistic averaged over slices and leading dims.

info_nce ¶

info_nce(
    a: Tensor,
    b: Tensor,
    temperature: float | Tensor = 0.07,
    symmetric: bool = True,
) -> Tensor

InfoNCE between paired rows of a and b, both (B, D); averaged over both directions when symmetric.