feat(loss): add Sinkhorn-Knopp centering and KoLeo group sizes for DINOv3

[saud5150]

Part of #1881 — phase 1 of 3, so this does not close the issue.

Description

• My change is breaking

The three loss-level pieces for DINOv3's first training phase, scoped in #1881 (comment). Paper and reference evidence for each is in that thread; this is what changed.

DINOLoss(output_dim=65536, center_mode="sinkhorn_knopp")       # next to "mean"
IBOTPatchLoss(output_dim=65536, center_mode="sinkhorn_knopp")
KoLeoLoss(group_size=16, topk=1, gather_distributed=False)

center_mode mirrors DINOv2/v3's own train.centering key. The shared sinkhorn_knopp() helper sits in lightly/models/modules/center.py beside Center; in that mode the center buffer stays registered but unused, so checkpoints load across modes, and IBOTPlusPlusPatchLoss inherits the mode via the extracted _teacher_probabilities(). KoLeoLoss searches nearest neighbors within consecutive groups of group_size; topk and gather_distributed complete the reference KoLeoLossDistributed surface in ~15 lines.

Defaults are unchanged and no existing test was adapted.

Three things not in the issue thread:

• exp(logit / 0.04) overflows fp32 at |logit| >= ~3.6. DINO heads L2-normalize before the weight-normed prototype layer, so logits stay in [-1, 1] and this cannot trigger; the reference has the identical exposure, and stabilizing under DDP needs a globally reduced max, i.e. an extra collective per step. I matched the reference.
• sinkhorn_knopp does not reuse swav_loss.py::sinkhorn, which normalizes by local_batch_size * world_size — wrong for iBOT, where the masked-token count differs per rank.
• DINOLoss already matches DINOv3's loss composition: upstream's separately-weighted global and local crop terms are algebraically the uniform mean over view pairs it already computes (12.2663984 vs 12.2663975 on identical inputs, for the pretrain default reweight_dino_local_loss: false).

Three commits, independently reviewable; the KoLeo commit shares no file with the other two.

Tests

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

sinkhorn_knopp is verified against a copy of the DINOv3 reference across 9 temperature x iteration combinations, following the OriginalDINOLoss pattern already in test_dino_loss.py, plus row-sum, prototype-marginal and detachment tests. The loss-level tests cover wiring rather than repeating the algorithm: Sinkhorn-Knopp applied jointly over all teacher views as upstream does, the center staying zero while remaining in the state dict, and IBOTPlusPlusPatchLoss inheriting the mode. KoLeoLoss gets 11 new tests, including a mocker-simulated world_size=2 gather whose loss and gradient match the non-distributed run on the concatenated global batch — a correct gathered forward can still have a wrong backward, as #1977 showed. GatherLayer itself is mocked out, so that test covers grouping and neighbor indexing keeping gradients on the right features, not GatherLayer's own backward.

make format-check   # All checks passed / 560 files already formatted
make lint           # All checks passed
make type-check     # Success: no issues found in 541 source files
make -C docs html-noplot                           # build succeeded
python -m pytest tests --runslow                   # 1751 passed, 217 skipped

Run on Python 3.9 with current torch; the 3.7 and 3.12 CI legs are unexercised locally. The full run also reports 4 macOS-specific failures (assertWarns dict iteration, two PyAV spawn pickling, one DCL gloo mismatch) that reproduce on a clean upstream/master.

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.

Every new parameter is documented on both the class Attributes and the __init__ Args of the affected loss. lightly.loss.html is the only rendered page that changes, attached as a PDF printout.

Improvements put into another issue:

• Multi-rank test for KoLeoLoss(gather_distributed=True) once the shared gloo process pool from Proper distributed tests #1982 lands; currently covered by a mocker-simulated world_size=2.
• DINOLoss cannot express the scheduled local-crop reweighting (reweight_dino_local_loss: true) used by the DINOv3 gram-anchoring config. Not needed for phase 1; I will raise it with the refinement-phase work.

Issues covering the breaking change:

• None. Not a breaking change: all defaults are unchanged and no existing test was adapted.

[saud5150]

DINOv3_losses_docs_page.pdf