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 of 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. Sinkhorn-Knopp already existed in the repo as swav_loss.sinkhorn, so the shared implementation now lives in lightly/models/modules/center.py beside Center and swav_loss.sinkhorn delegates to it, keeping its signature and output dtype. In Sinkhorn-Knopp mode the center buffer stays registered but unused, so checkpoints load across modes, and IBOTPlusPlusPatchLoss inherits the mode through 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.

Sharing the implementation changes two things for SwaV, both improvements:

• The exponential is computed in float32. sinkhorn previously returned NaN in fp16 for |out| > ~0.55 at the default epsilon=0.05, which is reachable under AMP with cosine-similarity logits. The codes are cast back, so the returned dtype is unchanged.
• The number of samples is all-reduced instead of derived from the world size. Identical whenever ranks hold equally many samples, which is always the case for SwaV, and it is what lets iBOT reuse the helper, where the masked-token count differs per rank. It rides in the existing sum_Q all-reduce, so the number of collectives is unchanged.

Two 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 it under DDP needs a globally reduced max, i.e. an extra collective per step. I matched the reference.
• DINOLoss already matches DINOv3's loss composition. Upstream averages the global-crop and local-crop terms separately and recombines them with dino_global_scale / dino_local_scale, which are proportional to the term counts; for the pretrain default reweight_dino_local_loss: false that is algebraically the uniform mean over view pairs DINOLoss already computes. Checked numerically on identical inputs with 2 global and 6 local views: 5.9371023 vs 5.9371023 for mean centering, 5.6693587 vs 5.6693592 for Sinkhorn-Knopp.

Commits are independently reviewable; the KoLeo commit shares no file with the others, and the review round is kept in its own commits on top so the incremental diff stays readable.

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.

29 new tests. 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, IBOTPlusPlusPatchLoss inheriting the mode, and a half-precision forward that fails if the cast back to the input dtype is dropped, since torch.einsum does not promote dtypes.

KoLeoLoss gets 14 tests, including two real two-rank tests on the shared gloo pool from #1982, marked DDP. The forward test asserts every rank sees the loss of the non-distributed run on the concatenated global batch; the gradient test asserts the local gradient is NUM_PROCESSES times the non-distributed one, since every rank computes the loss of the whole global batch and GatherLayer all-reduces the identical gradients before slicing out the local one — the factor DDP cancels when it averages parameter gradients. A correct gathered forward can still have a wrong backward, as #1977 showed. The mocker-simulated world_size=2 test is kept alongside them because it isolates grouping and neighbor indexing from GatherLayer.

Manually verified that the shared helper is a no-op for SwaV: bit-identical to the previous swav_loss.sinkhorn over 3600 random (input, epsilon, iterations) combinations single-process, and over two gloo ranks with gather_distributed both True and False.

Also ran a toy convergence check, since the tests above are all unit-level. A student/teacher pair with a DINOProjectionHead overfits 300 steps on a fixed batch of 256 synthetic views with 64 prototypes. Entropy is that of the mean teacher assignment as a fraction of its maximum, so 1.0 is perfectly balanced prototype usage and 0.0 is collapse:

                                first    last    teacher entropy
DINOLoss      mean              2.537    0.028    0.820
DINOLoss      sinkhorn_knopp    3.034    0.063    0.992
IBOTPatchLoss mean              2.402    0.211    0.802
IBOTPatchLoss sinkhorn_knopp    2.875    0.319    0.998

Both modes train in both losses. Sinkhorn-Knopp holds prototype usage essentially uniform where mean centering drifts to 0.80, which is the property it is chosen for; its slightly higher final loss is the expected cost of enforcing balance per batch rather than nudging a running center. This is a synthetic overfit check, not a benchmark, and says nothing about downstream quality. KoLeoLoss decreases in the same setup with and without grouping (0.076 to -0.232 ungrouped, -0.125 to -0.412 with group_size=8), and grouping leaves the global nearest-neighbor cosine similarity higher (0.56 vs 0.21) because only within-group neighbors are pushed apart.

make format-check   # All checks passed / 576 files already formatted
make lint           # All checks passed
make type-check     # Success: no issues found in 551 source files
make test-distributed                              # 4 passed, 1 skipped
make -C docs html-noplot                           # build succeeded
python -m pytest tests --runslow                   # 1832 passed, 224 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:

• 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.
• DINOLoss(center_mode="mean") fails on half-precision input, because the float32 center buffer promotes the teacher output while the student output stays fp16 and torch.einsum does not promote. Pre-existing on master and untouched here; happy to open an issue.

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

[liopeer]

/review

[liopeer]

It looks correct, I would mainly try to re-use the existing SwAV code though. I will have again a detailed look again later.