Add VISRegLoss (VISReg paper)#1968
Merged
Merged
Conversation
Implements the variance-invariance-sketching regularization loss from https://arxiv.org/abs/2606.02572 per Algorithm 1, with per-component losses, per-device random slices for distributed training, and the Eq. 8 prediction term over both global and local views. Part 1 of lightly-ai#1960.
Contributor
Author
Contributor
|
Thank you so much for the nice contribution!!! |
Contributor
|
/review |
Contributor
|
/codex review |
The tensor subtraction resolves to Any with the pinned torch stubs used in CI, tripping mypy's no-any-return. Annotate the return explicitly, matching visreg_scale_loss.
gabrielfruet
reviewed
Jul 8, 2026
gabrielfruet
left a comment
gabrielfruet
left a comment
Contributor
There was a problem hiding this comment.
Implementation is looking good to me :). Very nice and clean PR. I just think it's missing the test_forward__cuda. Thank you for the nice PR!
Contributor
|
The test failure is not related to your changes FYI. |
Other loss tests include a skip-guarded CUDA forward test; add the matching one for VISRegLoss to verify the tensors created inside the shape loss (slices, quantiles) land on the input's device. Addresses review feedback on lightly-ai#1968.
Contributor
Author
|
Thanks for the review! Added test_forward__cuda in 03d90e3, following the other loss tests. |
gabrielfruet
added a commit
that referenced
this pull request
Jul 9, 2026
… noise (#1970) * test(models): loosen tolerance in test_singular_mask for FP noise pool_masked uses scatter_reduce_(reduce="mean"), whose CPU summation order differs from torch.mean by ~1e-7 in float32. With the default torch.allclose(atol=1e-8, rtol=1e-5) tolerance this fails intermittently (~4% of random seeds), surfacing in CI on PR #1968. The result is mathematically equivalent; loosen the tolerance to atol=1e-5, rtol=1e-3 which is still ~100x tighter than typical float32 noise. * test(models): apply ruff formatting --------- Co-authored-by: gabrielfruet <gabrielfruet@users.noreply.github.com> Co-authored-by: Lionel Peer <lionel@lightly.ai>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1968 +/- ##
==========================================
+ Coverage 86.40% 86.54% +0.13%
==========================================
Files 173 174 +1
Lines 7386 7467 +81
==========================================
+ Hits 6382 6462 +80
- Misses 1004 1005 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
gabrielfruet
approved these changes
Jul 9, 2026
gabrielfruet
left a comment
gabrielfruet
left a comment
Contributor
There was a problem hiding this comment.
LGTM! Thank you for the nice contribution :)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related to #1960 (PR 1 of 2).
Description
Adds
lightly/loss/visreg_loss.pyimplementing VISReg (Wu et al., 2026) per Algorithm 1, structured likelejepa_loss.py. Covers the open points from #1960:broadcast/all_reduce(Section 3.2, avoiding [BUG] SIGReg withgather_distributed=Trueproduces gradient which is off by1/world_sizeunder DDP #1920).gather_distributed=Trueoptionally gathers embeddings for the statistics via the same autograd-aware path asVICRegLoss.visreg_center/scale/shape_losshelpers plusforward_components()returning a(total, pred, scale, shape, center)NamedTuple (per Feature Request: Return all losses with multi-loss methods e.g., VICReg #1422).forward()returns a scalar.lejepa_invariance_loss; intended forLeJEPAProjectionHead.lambda_param=0.9,num_slices=4096, equal weights).Implemented from the paper's math only (official repo is CC BY-NC). Three choices worth review:
forward(*, local_proj, global_proj)rather than the issue'sforward(z)sketch (the prediction term needs the global views); prediction and regularization run over all views per Eq. 8 / Section 3.1 (differs fromLeJEPALoss's local-only, pinned by a test); and Algorithm 1's(N,K)-(N,)broadcasting bug is fixed with anunsqueeze.Tests
37 tests in
tests/loss/test_visreg_loss.py(patterns fromtest_lejepa_loss.py): shape loss pinned against a naive Algorithm 1 transcription, non-vanishing gradient under collapse, stop-gradient behavior, all-views regularization scope, no-collectives on the default distributed path, and gather matching the global-batch loss.make format,make static-checks, andmake testpass (1693 passed; 4 failures are pre-existing macOS/Python-3.9 env issues, not hit by CI).Documentation
.rstfiles).VISRegLossautoclass entry added todocs/source/lightly.loss.rst(with:special-members: __init__). PDF of the rendered page attached.Implications / comments / further issues
examples/+docs/source/examples/visreg.rst, building on this loss.