Implement VISReg (Variance-Invariance-Sketching Regularization) loss + examples

[gabrielfruet]

Add VISReg loss + examples

Paper: https://arxiv.org/abs/2606.02572 (Wu, Balestriero, Levine, Jun 2026)
Official repo: https://github.com/HaiyuWu/visreg

Sits between VICRegLoss and SIGReg/LeJEPALoss: keeps the variance term from VICReg, replaces covariance with a sliced-Wasserstein sketching term against an isotropic Gaussian. O(NDK), and unlike SIGReg the gradient doesn't vanish under collapse (their Fig. 2).

Note: the official repo is CC BY-NC 4.0, non-commercial. Fine as a reference for the math (Algorithm 1 is unambiguous), but we shouldn't port code directly given our license.

Two PRs:

PR 1: Loss

• lightly/loss/visreg_loss.py, structured like LeJEPALoss
  • Don't forget to reference the paper and authors.
• register in lightly/loss/__init__.py
• update docs/source/lightly.loss.rst (autodoc entry, same pattern as VICRegLoss/LeJEPALoss)

class VISRegLoss(nn.Module):
    def __init__(
        self,
        lambda_param: float = 0.9,
        num_slices: int = 4096,
        lambda_scale: float = 1.0,
        lambda_shape: float = 1.0,
        lambda_center: float = 1.0,
        gather_distributed: bool = False,
        eps: float = 1e-4,
    ):
        ...

    def forward(self, z: Tensor) -> Tensor:
        ...

Open points:

1. Distributed: paper generates K slices per GPU, no batch gather (Sec 3.2, Fig. 6). Use that instead of the SIGReg all_reduce pattern from [BUG] SIGReg with gather_distributed=True produces gradient which is off by 1/world_size under DDP #1920.
2. Return component losses (scale/shape/center/pred) separately — same as Feature Request: Return all losses with multi-loss methods e.g., VICReg #1422 for VICReg. Table 12 shows shape-loss weight is the main tuning lever for long-tailed/low-rank data.
3. Reuse LeJEPAProjectionHead, no new head needed (no changes to lightly/models/modules/heads.py expected).
4. Use the best hyperparameters combinations defined in the paper.

Tests:

• Try to match same testing patterns for the other losses.

PR 2: Examples (depends on PR 1)

• examples/pytorch/visreg.py
• examples/pytorch_lightning/visreg.py
• examples/pytorch_lightning_distributed/visreg.py
• docs/source/examples/visreg.rst, same format as vicreg.rst, plus adding it to the examples toctree/index
• add VISReg to the model list in the main README and docs landing page

Follow the VICReg examples structure: backbone + LeJEPAProjectionHead + VISRegLoss, multi-crop augmentation per paper's A.1 setup.

[gabrielfruet]

@RecreationalMath @saud5150 @MohamedFarag21 If any of you find this interesting and would like to contribute, please let me know :).

We would love your contributions.

[RecreationalMath]

This looks interesting.
I'm occupied till July 10. If not needed immediately, then I'm happy to pick it up in the week of July 13th.

[RecreationalMath]

FYI, I found BarlowTwins to have the same bug as the one pointed out for SIGReg in #1920.
I will open an issue for it and fix that as well.

[saud5150]

Thanks for the ping — this is right up my alley. I'll pick up PR 1 (loss implementation): distributed K-slices per GPU (no all_reduce), component loss returns, and reusing LeJEPAProjectionHead. I'll open a draft PR with tests before July 10, so @RecreationalMath has something concrete to build the examples on top of when they're back.

[gabrielfruet]

@RecreationalMath @saud5150 Sounds like a plan to me!