feat(benchmarks): add VISReg to the ImageNet vitb16 benchmark

[saud5150]

Related to #1960 — adds the ViT benchmark requested in this maintainer comment, building on the VISRegLoss from #1968.

Description

• My change is breaking

Registers VISReg as a selectable method in the imagenet/vitb16 runner: a new benchmarks/imagenet/vitb16/visreg.py (VISReg LightningModule + multi-crop transform) plus a two-line entry in main.py's METHODS dict. The module mirrors lejepa.py so the two run head-to-head on the same backbone, optimizer, and eval, differing only where VISReg's paper says they should.

Per the maintainer note ("use the best setting in the paper as much as possible, like LeJEPA"), the settings are VISReg's own paper-best, each cross-checked against the paper and the official repo (not just the paper prose, which diverges from the code in two places):

Setting	Value	Source
backbone	vit_small_patch16_224	ViT-S/16 — smaller than the paper's ViT-B/16 per the maintainer ("we train smaller models")
views	4 global + 6 local	paper App. A.1 / repo n_global: 4 (needs a custom MultiViewTransform; DINOTransform is fixed at 2 global)
augmentation	one uniform recipe on every crop (blur 0.5, solarize 0.2)	reference _shared_aug
lr	9e-4	paper Table 3 (9e-4 > 5e-4 at this scale); repo lr: 9e-4
projector	3-layer MLP embed→2048→2048→256, BatchNorm + GELU, last layer bare	reference ViTEncoder.proj (proj_gelu: true); projects the last-layer CLS only
loss	VISRegLoss(num_slices=2048, gather_distributed=True)	K per paper App. A.1; gather_distributed=True matches the reference's autograd-aware all_gather before the reg statistics, and the sibling LeJEPA
λ / weight decay	0.9 / 5e-2	paper + repo

Two spots where the reference code overrides the paper's own text, worth noting since they look like bugs otherwise: the projector runs on the last-layer CLS token only (self.proj(last_cls)), not the "concatenated last-two-layer CLS" the paper describes — that concat is the eval-probe embedding — so input_dim=embed_dim, not 2×; and the augmentation is a single distribution shared by all crops, with no DINO-style global asymmetry.

Purely additive: no lightly/ source is touched, VISRegLoss is only imported, and the new METHODS["visreg"] key affects a run only when --methods visreg is passed.

Tests

• My change is covered by existing tests.
• My change needs new tests.
• I have added/adapted the tests accordingly.
• I have manually tested the change.

Benchmark method modules aren't unit-tested in this repo (same as dino.py, lejepa.py, etc.) and sit outside the pytest/mypy targets, so I verified with a smoke test plus the standard checks. Before: --methods visreg is not a registered method. After: it constructs and trains a step end-to-end.

smoke test: transform emits 4 global + 6 local, projector is the paper's GELU

384->2048->2048->256, and one training step runs end-to-end

python - <<'PY'
import sys; sys.path.insert(0, "benchmarks/imagenet/vitb16")
import torch, visreg
from PIL import Image
v = visreg.transform(Image.new("RGB", (256, 256)))
print(len(v), "views;", sum(x.shape[-1] == 224 for x in v), "global,", sum(x.shape[-1] == 96 for x in v), "local")
m = visreg.VISReg(batch_size_per_device=4, num_classes=10)
m.log = m.log_dict = lambda *a, **k: None
b = [torch.randn(4, 3, 224, 224)] * 4 + [torch.randn(4, 3, 96, 96)] * 6
print("loss", float(m.training_step((b, torch.randint(0, 10, (4,)), [""] * 4), 0).detach()))
PY

-> 10 views; 4 global, 6 local

-> loss ~3 (finite scalar; exact value varies with random init)

make format-check # All checks passed!
make lint # All checks passed!
make type-check # Success: no issues found in 542 source files
make test-fast # 1687 passed, 249 skipped; 4 failures are pre-existing macOS/Py3.9 env issues, not hit by CI and unrelated to this diff
ruff check benchmarks/imagenet/vitb16/visreg.py benchmarks/imagenet/vitb16/main.py # All checks passed!

The full ImageNet run (ViT-S/16, `--methods visreg`) that produces the comparison numbers + checkpoint is left to the maintainers' cluster.

## Documentation
- [ ] I have added docstrings to all changed/added public functions/methods.
- [ ] My change requires a change to the documentation ( `.rst` files).
- [ ] I have updated the documentation accordingly.
- [ ] The autodocs update the documentation accordingly.

Benchmark method modules carry no docstrings by convention (matching the sibling files); the non-obvious paper deviations are noted in short module comments. The benchmark results table needs real ImageNet numbers from a cluster run, so the `docs/.../benchmarks.rst` row is a follow-up once those exist — no `.rst` change here.

## Implications / comments / further issues
- `benchmarks/imagenet/resnet50/visreg.py` (the maintainer's stated follow-up) and the docs results row are separate PRs.
- The lightly runner imposes a couple of schedule differences from the paper the maintainer already OK'd ("changing the schedule and batch size is ok"): ~1-epoch warmup (vs 5) and no grad clipping (vs the repo's `clip_grad_norm: 1.0`). The cosine floor already matches (`end_value=0.001` == repo `final_lr_div: 1000`).
- Only integration surface is the `METHODS` registry; existing methods are unaffected.

[gabrielfruet]

LGTM! Thank you