Skip to content

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

Open
saud5150 wants to merge 12 commits into
lightly-ai:masterfrom
saud5150:feat/dinov3-phase1-losses
Open

saud5150 wants to merge 12 commits into
lightly-ai:masterfrom
saud5150:feat/dinov3-phase1-losses

Conversation

@saud5150

@saud5150 saud5150 commented Jul 28, 2026 •

Copy link
Copy Markdown
Contributor

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 each review round is kept in its own commits on top so the incremental diff stays readable. Since the last round master is merged in, and three commits on top drop the Python 3.7 workarounds now that 3.8 is the minimum (#2028), correct a comment that said GatherLayer averages gradients (it sums them), and strengthen the distributed tests.

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.

30 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 test of the teacher probabilities that fails if the cast back to the input dtype is dropped, since torch.einsum does not promote dtypes. It calls _teacher_probabilities directly because torch 1.10, the minimum, has no half-precision CPU log_softmax.

KoLeoLoss gets 14 tests, including two real two-rank tests on the shared gloo pool from #1982, marked DDP. Local batches of 6 with groups of 4 make one group span both ranks, so the tests also pin the rank order of the gathered batch; they fail if it is reversed. Workers read their rank from the process group, since the pool does not pin tasks to ranks. 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.

The distributed path of sinkhorn_knopp also runs on the pool, against the non-distributed result on the concatenated batch, with equal and with unequal batch sizes per rank; the unequal case fails if the number of samples is derived from the world size.

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 static-checks          # lock-check, ruff check + format, mypy: no issues in 239 source files
make test-distributed       # 6 passed
pytest tests --runslow      # 1578 passed, 186 skipped, 3 failed (macOS-only, see below)

Run on Python 3.9 with the lockfile's dependencies. The touched test files also pass at the floors of the minimal CI job, Python 3.8 with torch 1.10.0 and torchvision 0.11.1 (298 passed, 30 skipped for CUDA and the pool; the 6 pool tests pass with the pool on). The 3 failures are macOS-specific and unrelated to this PR: two PyAV spawn pickling tests and the DCL Lightning-DDP numerics test, all known to fail on this machine 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.

Summary

  • Add sinkhorn_knopp centering to DINOLoss, IBOTPatchLoss, and IBOTPlusPlusPatchLoss.
  • Share Sinkhorn-Knopp through center.py and keep SwaV compatibility.
  • Extend KoLeoLoss with grouped search, topk, and distributed gathering.
  • Preserve existing defaults and checkpoint-compatible center buffers.
  • Add validation and distributed, mixed-precision, and reference-comparison tests.

Validation

  • Static checks passed.
  • Distributed tests: 6 passed.
  • Slow tests: 1578 passed, 186 skipped, with 3 unrelated macOS failures reported.

@saud5150

Copy link
Copy Markdown
Contributor Author

@liopeer

liopeer commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

/review

@liopeer liopeer left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread lightly/loss/dino_loss.py Outdated
Comment thread lightly/models/modules/center.py Outdated
DINOv2 and DINOv3 replace the mean centering of the teacher output with the
Sinkhorn-Knopp centering from SwAV. DINOLoss now accepts
center_mode="sinkhorn_knopp", which normalizes the teacher logits so that every
prototype receives the same total weight across the batch instead of subtracting
a running center.

Following the reference implementation, the normalization is applied jointly
over all teacher views and no center is tracked. The center buffer stays
registered so that the state dict does not depend on the center mode.

The number of samples is reduced across processes rather than derived from the
world size, so that the shared helper also works for the iBOT loss where the
number of masked tokens differs between processes.
The DINOv3 paper applies Sinkhorn-Knopp centering to both the DINO and the iBOT
objective, and the reference training code asserts that centering is set to
sinkhorn_knopp. IBOTPatchLoss now accepts the same center_mode as DINOLoss.

Teacher normalization and the center update move into helpers so that
IBOTPlusPlusPatchLoss picks up the new mode without duplicating the branch.
DINOv3 applies the KoLeo regularizer to small batches of 16 samples. The
released configs reach that by running with a per-GPU batch size of 16, which is
not a setting lightly users are likely to train with.

KoLeoLoss now splits the batch into consecutive groups of group_size and
searches nearest neighbors within each group, so the paper setting is
reproducible at any batch size. topk and gather_distributed are added for parity
with the reference implementation.

Defaults are unchanged: without group_size the whole batch forms one group and
the loss is identical to before.
swav_loss.sinkhorn and the Sinkhorn-Knopp centering added for DINOv3 ran the
same algorithm. The implementation now lives in center.py and swav_loss.sinkhorn
delegates to it, keeping its signature and output dtype. Verified identical to
the previous implementation over 3600 random input, epsilon and iteration
combinations single-process, and over two gloo ranks with gather_distributed
both True and False.

The shared version calculates the exponential in float32, which fixes a NaN in
the SwaV path for half-precision input, and all-reduces the number of samples
instead of deriving it from the world size, which the iBOT loss needs because
the number of masked tokens differs between processes.

The docstring also no longer claims that DINOv2 uses Sinkhorn-Knopp centering.
DINOv2 offers it as an option but keeps mean centering in its released configs,
where it reports no difference on ImageNet-1k. DINOv3 uses it for both
objectives.
KoLeoLoss now rejects gather_distributed when torch.distributed is unavailable,
matching the ten other losses that do this, and reports a topk that is too
large for the group size instead of failing inside torch.topk. The previous
check skipped groups of size one, where a topk above one reached torch.topk
with k larger than the dimension.

DINOLoss and IBOTPatchLoss reject a negative sinkhorn_iterations, matching
MSNLoss, instead of silently running no iterations.

The Sinkhorn-Knopp probabilities are cast back to the dtype of the teacher
output, which was already the case but is now covered by a half-precision test
and a comment, because torch.einsum does not promote dtypes and the loss fails
without it.
The gathered KoLeoLoss was only covered by a mocked world_size=2 test, which
bypasses GatherLayer and therefore does not exercise its backward. The shared
gloo pool from lightly-ai#1982 makes a real two-rank test possible.

The forward test asserts that 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, because every rank
computes the loss of the whole global batch and GatherLayer all_reduces the
identical gradients before slicing out the local one. DDP cancels the factor
when it averages parameter gradients across ranks.
@saud5150
saud5150 force-pushed the feat/dinov3-phase1-losses branch from 8ce532a to 5336c4b Compare August 10, 2026 19:29
Comment thread lightly/models/modules/center.py
Comment thread lightly/loss/ibot_loss.py Outdated
Comment thread lightly/models/modules/center.py Outdated
Comment thread lightly/loss/ibot_loss.py Outdated
Replaces the CENTER_MODE_SINKHORN_KNOPP/VALID_CENTER_MODES module-level
constants with a Literal["mean", "sinkhorn_knopp"] type on center_mode in
DINOLoss and IBOTPatchLoss, and an explicit if/elif/else with a final
ValueError instead of a membership check against the shared constant list.
CENTER_MODE_TO_FUNCTION is kept since Center itself still relies on it.

Literal is imported behind TYPE_CHECKING because it requires Python 3.8+ and
this package still supports 3.7; the guard is exercised by deleting
typing.Literal and reloading both modules, which succeeds because the
annotation is never evaluated at runtime.

sinkhorn_iterations must now be at least 1 rather than merely non-negative,
since 0 iterations does not produce a valid probability distribution for
DINOv3 centering. The bound only moves in DINOLoss and IBOTPatchLoss: the
shared sinkhorn_knopp() function in center.py is left unconstrained, because
SwaVLoss also calls it and has long-tested, pre-existing support for
sinkhorn_iterations=0 as a deliberate configuration.
Two failures reproduced on the Python 3.7 CI job, both invisible locally
because this machine resolves a newer torch than the one pinned there
(torch==1.13.1).

_nearest_neighbor_indices returned an expression assembled from tensor
arithmetic directly; under torch 1.13.1's weaker operator stubs mypy widens
that to Any, tripping no-any-return. Assigning to a Tensor-annotated variable
before returning fixes it, matching the same gotcha hit on the VISReg PR
(lightly-ai#1968).

test_sinkhorn_knopp__half_precision ran the full DINOLoss forward pass in
half precision, which calls F.log_softmax on the student output. That has no
CPU kernel for Half in torch 1.13.1. The dtype-cast behavior the test guards
against lives entirely in _teacher_probabilities, which never touches
log_softmax, so the test now calls that method directly instead of going
through forward(). Confirmed this still catches the regression it was written
for: reverting the .to(teacher_out.dtype) cast makes the new assertion fail.
@saud5150 saud5150 mentioned this pull request Oct 21, 2026
5 tasks
saud5150 and others added 4 commits October 21, 2026 19:43
Resolves the import conflict in ibot_loss.py, where master added the cast
import (lightly-ai#2016) next to this branch's TYPE_CHECKING import.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Python 3.8 is now the minimum supported version (lightly-ai#2028) and mypy only runs
against the current torch stubs, so two workarounds from the Python 3.7 CI
leg are no longer needed:

- Literal is imported from typing directly instead of behind TYPE_CHECKING
  in DINOLoss and IBOTPatchLoss.
- _nearest_neighbor_indices returns its expression directly again; the
  Tensor-annotated intermediate only existed for the old torch stubs.

The half-precision DINOLoss test still calls _teacher_probabilities directly,
because the oldest supported torch (1.10) has no half-precision CPU kernel for
log_softmax. Its docstring now says so instead of pointing at the removed CI
leg.

sinkhorn_knopp's temperature uses a PEP 604 union with postponed annotations,
following the repo's typing convention.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The comment said GatherLayer averages the gradients. It sums them: its
backward all-reduces the gradients of all processes without dividing, so the
local gradient is world_size times the non-distributed one, as the multi-rank
gradient test asserts. DDP's averaging of parameter gradients cancels the
factor.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The distributed KoLeoLoss tests used local batches that were exactly one
group, so no group spanned two ranks and the order in which the gathered
batches are concatenated was never checked: the tests still passed with the
gathered batches reversed. They now use local batches of 6 with groups of 4,
which makes one group span both ranks; with the reversed concatenation the
pool tests fail.

The pool workers now read their rank from the process group. The pool hands
every task to whichever worker is free, so the rank passed with a task does
not always match the worker's process group rank (5 in 200 calls in a local
run). That was harmless while every group stayed within one rank, but with
cross-rank groups it would make the tests flaky. The gradient worker returns
its rank so that results are matched by rank rather than task order.

sinkhorn_knopp's distributed path, which SwaV now shares, had no multi-process
test. A pool test checks that it matches the non-distributed result on the
concatenated batch, with equal and with unequal batch sizes per rank, as
happens with the masked tokens of the iBOT loss. Deriving the number of
samples from the local batch size and the world size instead of reducing it
fails the unequal case.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@saud5150

Copy link
Copy Markdown
Contributor Author

@liopeer the Aug 12 feedback is in (Literal center modes per loss, VALID_CENTER_MODES gone, explicit if/elif/else with a final ValueError, sinkhorn_iterations >= 1), and master is merged in. Three small commits on top since then: dropping the Python 3.7 workarounds now that 3.8 is the minimum, a corrected comment in KoLeoLoss (GatherLayer sums gradients, it doesn't average them), and stronger distributed tests (a KoLeo group now spans both ranks, and the distributed Sinkhorn-Knopp path runs on the pool). Ready for another look. #2062 touches the same centering code, so whichever lands second needs a rebase; happy to take that on my side.

@coderabbitai

coderabbitai Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: Organization UI
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: cd51a1b1-d77a-44a2-8535-13dfe3ea4064
📥 Commits

Reviewing files that changed from the base of the PR and between dad178c and 5aea4b1.

📒 Files selected for processing: 11
  • lightly/loss/dino_loss.py
  • lightly/loss/ibot_loss.py
  • lightly/loss/koleo_loss.py
  • lightly/loss/swav_loss.py
  • lightly/models/modules/center.py
  • tests/loss/test_dino_loss.py
  • tests/loss/test_ibot_loss.py
  • tests/loss/test_ibot_plusplus_loss.py
  • tests/loss/test_koleo_loss.py
  • tests/loss/test_swav_loss.py
  • tests/models/modules/test_center.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The changes add shared Sinkhorn-Knopp normalization for SwAV, DINO, and iBOT losses. They also extend KoLeoLoss with grouped top-k neighbor selection and optional distributed feature gathering.

Changes: Sinkhorn-Knopp centering

Layer / File(s) Summary
Shared normalization and SwAV integration
lightly/models/modules/center.py, lightly/loss/swav_loss.py, tests/models/modules/test_center.py, tests/loss/test_swav_loss.py
Adds distributed-aware sinkhorn_knopp normalization and routes SwAV through it. Tests cover numerical results, distributed batch sizes, output shape, and dtype.
DINO teacher centering
lightly/loss/dino_loss.py, tests/loss/test_dino_loss.py
Adds mean and Sinkhorn-Knopp teacher probability modes to DINOLoss. Tests cover validation, joint view normalization, dtype preservation, and center state.
iBOT teacher centering
lightly/loss/ibot_loss.py, tests/loss/test_ibot_loss.py, tests/loss/test_ibot_plusplus_loss.py
Adds mean and Sinkhorn-Knopp teacher probability modes to both iBOT loss implementations. Tests cover validation, loss calculation, and center state.

KoLeo neighbor selection

Layer / File(s) Summary
Grouped and distributed neighbor selection
lightly/loss/koleo_loss.py, tests/loss/test_koleo_loss.py
Adds grouped top-k cosine-neighbor selection and optional distributed feature gathering. Tests cover local behavior, invalid inputs, distributed loss, and gradients.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant DINOLoss
  participant sinkhorn_knopp
  participant DistributedProcessGroup
  DINOLoss->>sinkhorn_knopp: pass teacher logits and iteration count
  sinkhorn_knopp->>DistributedProcessGroup: reduce totals and sample counts when gathering is enabled
  DistributedProcessGroup-->>sinkhorn_knopp: return aggregated totals and sample counts
  sinkhorn_knopp-->>DINOLoss: return detached teacher probabilities
Loading
sequenceDiagram
  participant KoLeoLoss
  participant DistributedFeatureGather
  participant NeighborSearch
  KoLeoLoss->>DistributedFeatureGather: gather features when enabled
  DistributedFeatureGather-->>KoLeoLoss: return global feature batch
  KoLeoLoss->>NeighborSearch: select grouped top-k neighbors
  NeighborSearch-->>KoLeoLoss: return neighbor indices
Loading

Possibly related issues

  • lightly-ai/lightly#1982: The PR adds distributed tests for KoLeo and Sinkhorn-Knopp, matching the issue’s request for broader distributed loss coverage.

Suggested reviewers: guarin

Merge Risk: ⚪ Minimal · up to 5aea4

The change adds opt-in Sinkhorn-Knopp centering for the DINO and iBOT losses and grouped, top-k and distributed neighbor selection for KoLeo, with defaults unchanged. No concrete merge-blocking risk was identified, so it looks ready to merge after the usual CI checks.

Pre-merge checks: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 58 functions across 11 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 4 passed
Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly summarizes the main changes: Sinkhorn-Knopp centering and grouped KoLeo support for DINOv3 losses.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create a new PR
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants