From 131fca54eaa474e2cd52f0ce443d53b2dc5857db Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Tue, 28 Jul 2026 18:05:30 +0500 Subject: [PATCH 1/8] feat(loss): add Sinkhorn-Knopp centering to DINOLoss 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. --- lightly/loss/dino_loss.py | 77 +++++++++++++++++++++++--- lightly/models/modules/center.py | 86 ++++++++++++++++++++++++++++- tests/loss/test_dino_loss.py | 66 +++++++++++++++++++++- tests/models/modules/test_center.py | 82 ++++++++++++++++++++++++++- 4 files changed, 300 insertions(+), 11 deletions(-) diff --git a/lightly/loss/dino_loss.py b/lightly/loss/dino_loss.py index fcebae018..607c8f0aa 100644 --- a/lightly/loss/dino_loss.py +++ b/lightly/loss/dino_loss.py @@ -8,7 +8,11 @@ from torch.nn import Module, Parameter from lightly.models.modules import center -from lightly.models.modules.center import CENTER_MODE_TO_FUNCTION +from lightly.models.modules.center import ( + CENTER_MODE_SINKHORN_KNOPP, + CENTER_MODE_TO_FUNCTION, + VALID_CENTER_MODES, +) class DINOLoss(Module): @@ -32,9 +36,16 @@ class DINOLoss(Module): Temperature parameter for the student network. center: Center used for the teacher output. It is updated with a moving average - during training. + during training. Unused if 'sinkhorn_knopp' centering is selected. center_momentum: Momentum term for the center calculation. + center_mode: + Mode used to normalize the teacher output. Either 'mean' for the mean + centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp centering + used by DINOv2 and DINOv3. + sinkhorn_iterations: + Number of Sinkhorn-Knopp iterations. Only used if center_mode is + 'sinkhorn_knopp'. warmup_teacher_temp_epochs: Number of epochs for the warmup phase of the teacher temperature (for backward compatibility). teacher_temp_schedule: @@ -64,16 +75,25 @@ def __init__( student_temp: float = 0.1, center_momentum: float = 0.9, center_mode: str = "mean", + sinkhorn_iterations: int = 3, ) -> None: """Initializes the DINOLoss Module. Args: center_mode: - Mode for center calculation. Only 'mean' is supported. + Mode used to normalize the teacher output. Either 'mean' for the mean + centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp + centering used by DINOv2 and DINOv3. + sinkhorn_iterations: + Number of Sinkhorn-Knopp iterations. Only used if center_mode is + 'sinkhorn_knopp'. warmup_teacher_temp: Initial temperature for the teacher network (for backward compatibility). warmup_teacher_temp_epochs: Number of epochs for the warmup phase of the teacher temperature (for backward compatibility). + + Raises: + ValueError: If an unknown center mode is provided. """ super().__init__() @@ -82,12 +102,16 @@ def __init__( # TODO(Guarin, 08/24): Refactor this to use the Center module directly once # we do a breaking change. - if center_mode not in CENTER_MODE_TO_FUNCTION: + if center_mode not in VALID_CENTER_MODES: raise ValueError( f"Unknown mode '{center_mode}'. Valid modes are " - f"{sorted(CENTER_MODE_TO_FUNCTION.keys())}." + f"{sorted(VALID_CENTER_MODES)}." ) - self._center_fn = CENTER_MODE_TO_FUNCTION[center_mode] + self.center_mode = center_mode + self.sinkhorn_iterations = sinkhorn_iterations + # No center is tracked with Sinkhorn-Knopp centering. The center buffer is + # still registered to keep the state dict independent of the center mode. + self._center_fn = CENTER_MODE_TO_FUNCTION.get(center_mode) self.center: Parameter self.register_buffer("center", torch.zeros(1, 1, output_dim)) self.center_momentum = center_momentum @@ -140,8 +164,8 @@ def forward( # Calculate cross-entropy loss. teacher_out_stacked = torch.stack(teacher_out) - t_out: Tensor = F.softmax( - (teacher_out_stacked - self.center) / teacher_temperature, dim=-1 + t_out = self._teacher_probabilities( + teacher_out=teacher_out_stacked, teacher_temp=teacher_temperature ) student_out_stacked = torch.stack(student_out) s_out = F.log_softmax(student_out_stacked / self.student_temp, dim=-1) @@ -171,15 +195,52 @@ def forward( return loss + def _teacher_probabilities( + self, teacher_out: Tensor, teacher_temp: Tensor + ) -> Tensor: + """Returns the sharpened teacher probabilities for the given teacher output. + + Applies either mean centering followed by a softmax, or Sinkhorn-Knopp + centering, depending on center_mode. + + Args: + teacher_out: + Tensor with shape (num_views, batch_size, output_dim) containing + features from the teacher model. + teacher_temp: + The temperature used for the teacher output. + + Returns: + Tensor with the same shape as teacher_out containing probabilities that + sum to one along the last dimension. + """ + if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + # Sinkhorn-Knopp is applied jointly over all views, following the reference + # implementation. + # (num_views, batch_size, output_dim) -> (num_views * batch_size, output_dim) + probabilities = center.sinkhorn_knopp( + x=teacher_out.flatten(0, 1), + temperature=teacher_temp, + num_iterations=self.sinkhorn_iterations, + ) + return probabilities.reshape(teacher_out.shape).to(teacher_out.dtype) + return F.softmax((teacher_out - self.center) / teacher_temp, dim=-1) + @torch.no_grad() def update_center(self, teacher_out: Tensor) -> None: """Moving average update of the center used for the teacher output. + Does nothing if Sinkhorn-Knopp centering is used, as it does not track a + center. + Args: teacher_out: Tensor with shape (num_views, batch_size, output_dim) containing features from the teacher model. """ + if self._center_fn is None: + return + # Calculate the batch center using the specified center function batch_center = self._center_fn(x=teacher_out, dim=(0, 1)) diff --git a/lightly/models/modules/center.py b/lightly/models/modules/center.py index 21be7b9bf..6d9d466f0 100644 --- a/lightly/models/modules/center.py +++ b/lightly/models/modules/center.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Tuple, Union import torch import torch.distributed as dist @@ -106,6 +106,90 @@ def center_momentum(center: Tensor, batch_center: Tensor, momentum: float) -> Te return center * momentum + batch_center * (1 - momentum) +@torch.no_grad() +def sinkhorn_knopp( + x: Tensor, + temperature: Union[float, Tensor] = 0.04, + num_iterations: int = 3, +) -> Tensor: + """Returns teacher probabilities with Sinkhorn-Knopp centering as used in DINOv2 [0] + and DINOv3 [1]. + + Sinkhorn-Knopp centering originates from SwAV [2] and replaces the mean centering + followed by a softmax used in DINO [3]. Instead of subtracting a running center, the + sharpened teacher logits are normalized such that every prototype receives the same + total weight across the batch, which avoids collapse without tracking any state. + + Implementation is based on [4]. In distributed settings the normalization is computed + over the batches of all processes. + + This differs from lightly.loss.swav_loss.sinkhorn, which normalizes by + local_batch_size * world_size and only reduces across processes if requested. Here + the number of samples is reduced across processes, which is required for the iBOT + loss where the number of masked tokens differs between processes. + + - [0]: DINOv2, 2023, https://arxiv.org/abs/2304.07193 + - [1]: DINOv3, 2025, https://arxiv.org/abs/2508.10104 + - [2]: SwAV, 2020, https://arxiv.org/abs/2006.09882 + - [3]: DINO, 2021, https://arxiv.org/abs/2104.14294 + - [4]: https://github.com/facebookresearch/dinov3/blob/main/dinov3/loss/dino_clstoken_loss.py + + Args: + x: + Tensor with shape (batch_size, num_prototypes) containing the teacher + logits. + temperature: + Temperature used to sharpen the teacher logits. + num_iterations: + Number of Sinkhorn-Knopp iterations. + + Returns: + Tensor with shape (batch_size, num_prototypes) containing the teacher + probabilities in float32. Every row sums to one. The probabilities are detached + from the computation graph, following the reference implementation. + """ + # (batch_size, num_prototypes) -> (num_prototypes, batch_size) following the + # notation of the reference implementation. + Q = torch.exp(x.float() / temperature).t() + num_prototypes = Q.shape[0] + + # The number of samples is reduced together with the sum over Q to save a + # synchronization point. The actual number of samples is reduced instead of + # multiplying the local batch size with the world size because processes can have + # different batch sizes. This happens for example in the iBOT loss where the number + # of masked tokens differs between processes. + local_num_samples = torch.tensor(Q.shape[1], device=Q.device, dtype=Q.dtype) + sums = torch.stack([Q.sum(), local_num_samples]) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(sums) + sum_Q, num_samples = sums[0], sums[1] + + # Make the matrix sum to 1. + Q /= sum_Q + + for _ in range(num_iterations): + # Normalize rows: the total weight per prototype must be 1 / num_prototypes. + sum_of_rows = torch.sum(Q, dim=1, keepdim=True) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(sum_of_rows) + Q /= sum_of_rows + Q /= num_prototypes + + # Normalize columns: the total weight per sample must be 1 / num_samples. + Q /= torch.sum(Q, dim=0, keepdim=True) + Q /= num_samples + + # Scale the columns to sum to one, such that Q is an assignment. + Q *= num_samples + return Q.t() + + +CENTER_MODE_SINKHORN_KNOPP = "sinkhorn_knopp" + CENTER_MODE_TO_FUNCTION = { "mean": center_mean, } + +# Modes accepted by the losses. The Center module only supports the modes in +# CENTER_MODE_TO_FUNCTION, as Sinkhorn-Knopp does not track a center. +VALID_CENTER_MODES = [*CENTER_MODE_TO_FUNCTION, CENTER_MODE_SINKHORN_KNOPP] diff --git a/tests/loss/test_dino_loss.py b/tests/loss/test_dino_loss.py index 164c6e3a3..cf1319563 100644 --- a/tests/loss/test_dino_loss.py +++ b/tests/loss/test_dino_loss.py @@ -9,7 +9,7 @@ from torch import Tensor, nn from lightly.loss import DINOLoss -from lightly.models.modules.center import Center +from lightly.models.modules.center import Center, sinkhorn_knopp from lightly.models.utils import deactivate_requires_grad @@ -131,6 +131,70 @@ def test_other_parameters( center_momentum=center_momentum, ) + def test__init__invalid_center_mode(self) -> None: + with pytest.raises(ValueError, match="Unknown mode"): + DINOLoss(output_dim=4, center_mode="invalid") + + def test_sinkhorn_knopp(self) -> None: + """Sinkhorn-Knopp centering is applied jointly over all teacher views. + + The reference implementation centers the teacher output of all views at once + and then averages the cross-entropy over all view pairs with a different view + index: + https://github.com/facebookresearch/dinov3/blob/main/dinov3/loss/dino_clstoken_loss.py + """ + batch_size, output_dim = 3, 4 + teacher_temp, student_temp = 0.04, 0.1 + n_global, n_local = 2, 6 + teacher_out = _generate_output( + batch_size=batch_size, n_views=n_global, output_dim=output_dim, seed=0 + ) + student_out = _generate_output( + batch_size=batch_size, + n_views=n_global + n_local, + output_dim=output_dim, + seed=1, + ) + + loss_fn = DINOLoss( + output_dim=output_dim, + teacher_temp=teacher_temp, + student_temp=student_temp, + center_mode="sinkhorn_knopp", + ) + loss = loss_fn(teacher_out=teacher_out, student_out=student_out) + + teacher_probs = sinkhorn_knopp( + x=torch.cat(teacher_out), temperature=teacher_temp + ).reshape(n_global, batch_size, output_dim) + student_log_probs = F.log_softmax( + torch.stack(student_out) / student_temp, dim=-1 + ) + expected = torch.tensor(0.0) + n_terms = 0 + for t, teacher_view in enumerate(teacher_probs): + for s, student_view in enumerate(student_log_probs): + if s == t: + continue + expected = expected - (teacher_view * student_view).sum() + n_terms += batch_size + expected = expected / n_terms + + assert torch.allclose(loss, expected) + + def test_sinkhorn_knopp__center_not_updated(self) -> None: + """Sinkhorn-Knopp does not track a center, but keeps the buffer registered.""" + loss_fn = DINOLoss(output_dim=4, center_mode="sinkhorn_knopp") + teacher_out = _generate_output(n_views=2, output_dim=4, seed=0) + student_out = _generate_output(n_views=4, output_dim=4, seed=1) + + loss_fn(teacher_out=teacher_out, student_out=student_out) + assert "center" in loss_fn.state_dict() + assert torch.all(loss_fn.center == 0) + + loss_fn.update_center(teacher_out=torch.stack(teacher_out)) + assert torch.all(loss_fn.center == 0) + def test_single_view_raises(self) -> None: # A single teacher view and a single student view leave no cross-view # terms (the diagonal is excluded), which previously produced a silent diff --git a/tests/models/modules/test_center.py b/tests/models/modules/test_center.py index ccdb20458..d11554c04 100644 --- a/tests/models/modules/test_center.py +++ b/tests/models/modules/test_center.py @@ -1,8 +1,43 @@ +import typing + import pytest import torch from torch import Tensor +from torch.nn import functional as F + +from lightly.models.modules.center import Center, sinkhorn_knopp + + +@typing.no_type_check +@torch.no_grad() +def original_sinkhorn_knopp_teacher(teacher_output, teacher_temp, n_iterations=3): + """Copy paste from the original DINOv3 implementation. We use this to verify our + implementation. + + The only change from the original code is that distributed training is no longer + assumed. + + Source: https://github.com/facebookresearch/dinov3/blob/6876159a11b4df116f30f667f8c9888617df0751/dinov3/loss/dino_clstoken_loss.py#L43 + """ + teacher_output = teacher_output.float() + world_size = 1 + Q = torch.exp(teacher_output / teacher_temp).t() + B = Q.shape[1] * world_size + K = Q.shape[0] + + sum_Q = torch.sum(Q) + Q /= sum_Q + + for _ in range(n_iterations): + sum_of_rows = torch.sum(Q, dim=1, keepdim=True) + Q /= sum_of_rows + Q /= K + + Q /= torch.sum(Q, dim=0, keepdim=True) + Q /= B -from lightly.models.modules.center import Center + Q *= B + return Q.t() class TestCenter: @@ -39,3 +74,48 @@ def test_update__momentum(self, momentum: float, expected: Tensor) -> None: center = Center(size=(1, 2), mode="mean", momentum=momentum) center.update(torch.tensor([[1.0, 2.0]])) assert torch.all(center.value == expected) + + +def _teacher_logits(batch_size: int, num_prototypes: int) -> Tensor: + """Returns logits in [-1, 1], as produced by a projection head with weight norm.""" + return F.normalize(torch.randn(batch_size, num_prototypes), dim=-1) + + +class TestSinkhornKnopp: + @pytest.mark.parametrize("temperature", [0.04, 0.07, 0.1]) + @pytest.mark.parametrize("num_iterations", [1, 3, 5]) + def test__matches_original_implementation( + self, temperature: float, num_iterations: int + ) -> None: + torch.manual_seed(0) + x = _teacher_logits(batch_size=16, num_prototypes=8) + + expected = original_sinkhorn_knopp_teacher( + teacher_output=x, teacher_temp=temperature, n_iterations=num_iterations + ) + probabilities = sinkhorn_knopp( + x=x, temperature=temperature, num_iterations=num_iterations + ) + assert torch.allclose(probabilities, expected, atol=1e-6) + + def test__rows_sum_to_one(self) -> None: + torch.manual_seed(0) + x = _teacher_logits(batch_size=16, num_prototypes=8) + probabilities = sinkhorn_knopp(x=x, temperature=0.04) + assert torch.allclose(probabilities.sum(dim=1), torch.ones(16), atol=1e-5) + + def test__prototypes_are_equally_used(self) -> None: + """Every prototype gets the same total weight across the batch.""" + torch.manual_seed(0) + batch_size, num_prototypes = 32, 8 + probabilities = sinkhorn_knopp( + x=_teacher_logits(batch_size=batch_size, num_prototypes=num_prototypes), + temperature=0.1, + num_iterations=100, + ) + expected = torch.full((num_prototypes,), batch_size / num_prototypes) + assert torch.allclose(probabilities.sum(dim=0), expected, atol=1e-2) + + def test__no_grad(self) -> None: + x = _teacher_logits(batch_size=4, num_prototypes=8).requires_grad_() + assert not sinkhorn_knopp(x=x, temperature=0.04).requires_grad From bd26189486368f630e107d2557a56251b3181627 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Tue, 28 Jul 2026 18:05:41 +0500 Subject: [PATCH 2/8] feat(loss): add Sinkhorn-Knopp centering to IBOTPatchLoss 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. --- lightly/loss/ibot_loss.py | 89 ++++++++++++++++++++++++--- tests/loss/test_ibot_loss.py | 46 ++++++++++++++ tests/loss/test_ibot_plusplus_loss.py | 13 ++++ 3 files changed, 138 insertions(+), 10 deletions(-) diff --git a/lightly/loss/ibot_loss.py b/lightly/loss/ibot_loss.py index 4e32a6d01..50d6a3758 100644 --- a/lightly/loss/ibot_loss.py +++ b/lightly/loss/ibot_loss.py @@ -5,7 +5,13 @@ from torch.nn import Module from torch.nn import functional as F -from lightly.models.modules.center import Center +from lightly.models.modules import center as center_module +from lightly.models.modules.center import ( + CENTER_MODE_SINKHORN_KNOPP, + CENTER_MODE_TO_FUNCTION, + VALID_CENTER_MODES, + Center, +) class IBOTPatchLoss(Module): @@ -25,9 +31,14 @@ class IBOTPatchLoss(Module): student_temp: Temperature for the student output. center_mode: - Mode for center calculation. Only 'mean' is supported. + Mode used to normalize the teacher output. Either 'mean' for the mean + centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp centering + used by DINOv2 and DINOv3. center_momentum: Momentum term for the center update. + sinkhorn_iterations: + Number of Sinkhorn-Knopp iterations. Only used if center_mode is + 'sinkhorn_knopp'. """ def __init__( @@ -37,19 +48,77 @@ def __init__( student_temp: float = 0.1, center_mode: str = "mean", center_momentum: float = 0.9, + sinkhorn_iterations: int = 3, ) -> None: - """Initializes the iBOTPatchLoss module with the specified parameters.""" + """Initializes the iBOTPatchLoss module with the specified parameters. + + Raises: + ValueError: If an unknown center mode is provided. + """ super().__init__() self.teacher_temp = teacher_temp self.student_temp = student_temp + if center_mode not in VALID_CENTER_MODES: + raise ValueError( + f"Unknown mode '{center_mode}'. Valid modes are " + f"{sorted(VALID_CENTER_MODES)}." + ) + self.center_mode = center_mode + self.sinkhorn_iterations = sinkhorn_iterations + + # Sinkhorn-Knopp centering does not track a center. The Center module is still + # created to keep the state dict independent of the center mode. self.center = Center( size=(1, output_dim), - mode=center_mode, + mode=center_mode if center_mode in CENTER_MODE_TO_FUNCTION else "mean", momentum=center_momentum, ) + def _teacher_probabilities( + self, teacher_out: Tensor, teacher_temp: Tensor + ) -> Tensor: + """Returns the sharpened teacher probabilities for the given teacher output. + + Applies either mean centering followed by a softmax, or Sinkhorn-Knopp + centering, depending on center_mode. + + Args: + teacher_out: + Tensor with shape (num_tokens, output_dim) containing the teacher + output. + teacher_temp: + The temperature used for the teacher output. + + Returns: + Tensor with the same shape as teacher_out containing probabilities that + sum to one along the last dimension. + """ + if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + probabilities = center_module.sinkhorn_knopp( + x=teacher_out, + temperature=teacher_temp, + num_iterations=self.sinkhorn_iterations, + ) + return probabilities.to(teacher_out.dtype) + return F.softmax((teacher_out - self.center.value) / teacher_temp, dim=-1) + + def _update_center(self, teacher_out: Tensor) -> None: + """Updates the center with the given teacher output. + + Does nothing if no center is tracked, which is the case for Sinkhorn-Knopp + centering. + + Args: + teacher_out: + Tensor with shape (num_tokens, output_dim) containing the teacher + output. + """ + if self.center_mode not in CENTER_MODE_TO_FUNCTION: + return + self.center.update(teacher_out) + def forward( self, teacher_out: Tensor, @@ -85,8 +154,8 @@ def forward( ) # Calculate cross-entropy loss. - teacher_softmax = F.softmax( - (teacher_out - self.center.value) / teacher_temperature, dim=-1 + teacher_softmax = self._teacher_probabilities( + teacher_out=teacher_out, teacher_temp=teacher_temperature ) student_log_softmax = F.log_softmax(student_out / self.student_temp, dim=-1) @@ -103,7 +172,7 @@ def forward( B = mask.shape[0] loss = (loss * weight).sum() / B - self.center.update(teacher_out) + self._update_center(teacher_out) return loss @@ -227,8 +296,8 @@ def forward( ) # (B * N, K) - teacher_softmax = F.softmax( - (teacher_flat - self.center.value) / teacher_temperature, dim=-1 + teacher_softmax = self._teacher_probabilities( + teacher_out=teacher_flat, teacher_temp=teacher_temperature ) student_log_softmax = F.log_softmax(student_flat / self.student_temp, dim=-1) @@ -249,6 +318,6 @@ def forward( visible_loss = (ce * (1.0 - mask_flat)).sum(dim=1) / n_visible loss = (masked_loss + visible_loss_weight * visible_loss).mean() - self.center.update(teacher_flat) + self._update_center(teacher_flat) return loss diff --git a/tests/loss/test_ibot_loss.py b/tests/loss/test_ibot_loss.py index 25a202eb2..348963b85 100644 --- a/tests/loss/test_ibot_loss.py +++ b/tests/loss/test_ibot_loss.py @@ -1,7 +1,9 @@ import pytest import torch +from torch.nn import functional as F from lightly.loss.ibot_loss import IBOTPatchLoss +from lightly.models.modules.center import sinkhorn_knopp class TestIBOTPatchLoss: @@ -44,3 +46,47 @@ def test_forward(self, device: str) -> None: # teacher_patch_tokens_masked=orig_t_center, # student_masks_flat=mask.flatten(start_dim=1), # ) + + def test__init__invalid_center_mode(self) -> None: + with pytest.raises(ValueError, match="Unknown mode"): + IBOTPatchLoss(output_dim=2, center_mode="invalid") + + def test_sinkhorn_knopp(self) -> None: + """Sinkhorn-Knopp centering replaces the mean centering of the teacher output. + + DINOv3 applies Sinkhorn-Knopp to the masked teacher tokens, see + https://github.com/facebookresearch/dinov3/blob/main/dinov3/loss/ibot_patch_loss.py + """ + teacher_temp, student_temp = 0.04, 0.2 + criterion = IBOTPatchLoss( + output_dim=2, + teacher_temp=teacher_temp, + student_temp=student_temp, + center_mode="sinkhorn_knopp", + ) + teacher_out = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + student_out = torch.tensor([[0.7, 0.8], [0.9, 1.0], [1.1, 1.2]]) + mask = torch.tensor( + [ + [[True, False], [True, False]], + [[False, False], [False, True]], + [[False, False], [False, False]], + ] + ) + + loss = criterion.forward( + teacher_out=teacher_out, student_out=student_out, mask=mask + ) + + teacher_probs = sinkhorn_knopp(x=teacher_out, temperature=teacher_temp) + cross_entropy = -( + teacher_probs * F.log_softmax(student_out / student_temp, dim=-1) + ).sum(dim=-1) + weight = (1.0 / mask.sum(dim=(1, 2), keepdim=True).clamp(min=1.0)).expand_as( + mask + )[mask] + expected = (cross_entropy * weight).sum() / mask.shape[0] + + assert loss == pytest.approx(expected.item(), rel=1e-6) + # No center is tracked with Sinkhorn-Knopp centering. + assert torch.all(criterion.center.value == 0) diff --git a/tests/loss/test_ibot_plusplus_loss.py b/tests/loss/test_ibot_plusplus_loss.py index a12d210c9..69c8c0541 100644 --- a/tests/loss/test_ibot_plusplus_loss.py +++ b/tests/loss/test_ibot_plusplus_loss.py @@ -187,6 +187,19 @@ def test_invalid_shapes_raise(self) -> None: mask=torch.zeros(3, 2, dtype=torch.bool), ) + def test_forward__sinkhorn_knopp(self) -> None: + """The Sinkhorn-Knopp centering of the parent class is inherited.""" + torch.manual_seed(0) + criterion = IBOTPlusPlusPatchLoss(output_dim=8, center_mode="sinkhorn_knopp") + teacher_out = F.normalize(torch.randn(4, 16, 8), dim=-1) + student_out = torch.randn(4, 16, 8) + + loss = criterion(teacher_out=teacher_out, student_out=student_out) + + assert loss.isfinite() + # No center is tracked with Sinkhorn-Knopp centering. + assert torch.all(criterion.center.value == 0) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="No cuda") def test_cuda_forward(self) -> None: torch.manual_seed(0) From 19cf75ec0f89b05719910935a5c6aca48310429e Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Tue, 28 Jul 2026 18:05:50 +0500 Subject: [PATCH 3/8] feat(loss): add group size, topk and distributed gather to KoLeoLoss 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. --- lightly/loss/koleo_loss.py | 123 +++++++++++++++++++++++++++++-- tests/loss/test_koleo_loss.py | 133 ++++++++++++++++++++++++++++++++++ 2 files changed, 249 insertions(+), 7 deletions(-) diff --git a/lightly/loss/koleo_loss.py b/lightly/loss/koleo_loss.py index 5f8751cb1..c6450eeb7 100644 --- a/lightly/loss/koleo_loss.py +++ b/lightly/loss/koleo_loss.py @@ -1,7 +1,48 @@ +from __future__ import annotations + import torch from torch import Tensor from torch.nn import Module, PairwiseDistance, functional +from lightly.utils import dist as lightly_dist + + +@torch.no_grad() +def _nearest_neighbor_indices(x: Tensor, group_size: int, topk: int) -> Tensor: + """Returns the indices of the topk nearest neighbors within every group. + + The batch is split into consecutive groups of group_size features and neighbors + are searched by cosine similarity within each group. A feature is never its own + neighbor unless the group holds no other feature. + + Args: + x: + Tensor with shape (batch_size, embedding_size) containing L2-normalized + features. The batch size must be divisible by group_size. + group_size: + Number of features per group. + topk: + Number of neighbors per feature. + + Returns: + Tensor with shape (batch_size * topk,) containing indices into x, ordered by + feature and then by neighbor. + """ + # (batch_size, embedding_size) -> (num_groups, group_size, embedding_size) + x_grouped = x.view(-1, group_size, x.shape[1]) + num_groups = x_grouped.shape[0] + + # Cosine similarity within every group, with self-similarity masked out. + cos_sim = torch.bmm(x_grouped, x_grouped.transpose(1, 2)) + cos_sim.diagonal(dim1=-2, dim2=-1).fill_(-2) + + # (num_groups, group_size, topk) + nn_idx = cos_sim.topk(k=topk, dim=-1).indices + + # Shift the group-local indices so that they index into the flat batch. + offset = torch.arange(num_groups, device=x.device) * group_size + return (nn_idx + offset.view(-1, 1, 1)).flatten() + class KoLeoLoss(Module): """KoLeo loss based on [0]. @@ -10,22 +51,47 @@ class KoLeoLoss(Module): batch by penalizing the distance between the features and their nearest neighbors. - Implementation is based on [1]. + Implementation is based on [1]. Nearest neighbors are searched within groups of + group_size features, following the distributed KoLeo loss used by DINOv3 [2][3]. - [0]: Spreading vectors for similarity search, 2019, https://arxiv.org/abs/1806.03198 - [1]: https://github.com/facebookresearch/dinov2/blob/main/dinov2/loss/koleo_loss.py + - [2]: DINOv3, 2025, https://arxiv.org/abs/2508.10104 + - [3]: https://github.com/facebookresearch/dinov3/blob/main/dinov3/loss/koleo_loss.py Attributes: p: The norm degree for pairwise distance calculation. eps: Small value to avoid division by zero. + topk: + Number of nearest neighbors per feature that are penalized. + group_size: + Number of features within which nearest neighbors are searched. The batch + is split into consecutive groups of this size. If None, the whole batch is + used as a single group. DINOv3 uses a group size of 16. + gather_distributed: + If True, features from all GPUs are gathered before the batch is split into + groups. Requires that the distributed process group is initialized. + + Examples: + >>> # initialize loss function + >>> loss_fn = KoLeoLoss() + >>> + >>> # generate the features of a batch of images + >>> features = model(images) + >>> + >>> # calculate loss + >>> loss = loss_fn(features) """ def __init__( self, p: float = 2, eps: float = 1e-8, + topk: int = 1, + group_size: int | None = None, + gather_distributed: bool = False, ): """Initializes the KoLeoLoss module with the specified parameters. @@ -34,10 +100,30 @@ def __init__( The norm degree for pairwise distance calculation. eps: Small value to avoid division by zero. + topk: + Number of nearest neighbors per feature that are penalized. + group_size: + Number of features within which nearest neighbors are searched. The + batch is split into consecutive groups of this size. If None, the whole + batch is used as a single group. DINOv3 uses a group size of 16. + gather_distributed: + If True, features from all GPUs are gathered before the batch is split + into groups. Requires that the distributed process group is initialized. + + Raises: + ValueError: If topk or group_size are not positive. """ super().__init__() + if topk < 1: + raise ValueError(f"topk must be positive but is {topk}.") + if group_size is not None and group_size < 1: + raise ValueError(f"group_size must be positive but is {group_size}.") + self.p = p self.eps = eps + self.topk = topk + self.group_size = group_size + self.gather_distributed = gather_distributed self.pairwise_distance = PairwiseDistance(p=p, eps=eps) def forward(self, x: Tensor) -> Tensor: @@ -48,17 +134,40 @@ def forward(self, x: Tensor) -> Tensor: Returns: Loss value. + + Raises: + ValueError: If the batch size is not divisible by group_size or if + group_size is not larger than topk. """ # Normalize the input tensor x = functional.normalize(x, p=2, dim=-1, eps=self.eps) - # Calculate cosine similarity. - cos_sim = torch.mm(x, x.t()) - cos_sim.fill_diagonal_(-2) + # Gather features from all GPUs. The loss is calculated over the global batch + # on every process, gradients are averaged by the GatherLayer. + if self.gather_distributed and lightly_dist.world_size() > 1: + x = torch.cat(lightly_dist.gather(x), dim=0) + + batch_size = x.shape[0] + if batch_size == 0: + raise ValueError("KoLeoLoss requires a non-empty batch.") + + group_size = self.group_size if self.group_size is not None else batch_size + if batch_size % group_size != 0: + raise ValueError( + f"Batch size {batch_size} must be divisible by group size {group_size}." + ) + # A group must hold topk neighbors besides the feature itself, except for the + # degenerate group of size one, which is allowed for backwards compatibility. + if group_size > 1 and self.topk >= group_size: + raise ValueError( + f"Group size {group_size} must be larger than topk {self.topk}." + ) - # Get nearest neighbors. - nn_idx = cos_sim.argmax(dim=1) - nn_dist: Tensor = self.pairwise_distance(x, x[nn_idx]) + # Get the nearest neighbors and their distances. + nn_idx = _nearest_neighbor_indices(x=x, group_size=group_size, topk=self.topk) + nn_dist: Tensor = self.pairwise_distance( + x.repeat_interleave(self.topk, dim=0), x[nn_idx] + ) # Compute the loss loss = -(nn_dist + self.eps).log().mean() diff --git a/tests/loss/test_koleo_loss.py b/tests/loss/test_koleo_loss.py index bfc1d8bff..5ff18dae9 100644 --- a/tests/loss/test_koleo_loss.py +++ b/tests/loss/test_koleo_loss.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from pytest_mock import MockerFixture from torch import Tensor from lightly.loss.koleo_loss import KoLeoLoss @@ -41,3 +42,135 @@ def test_forward(self, x: Tensor, expected_loss: float, device: str) -> None: x = x.to(device) loss = KoLeoLoss().to(device) assert loss(x).item() == pytest.approx(expected_loss, rel=1e-4) + + def test_forward__group_size_none_is_full_batch(self) -> None: + torch.manual_seed(0) + x = torch.randn(8, 4) + assert KoLeoLoss(group_size=8)(x) == pytest.approx(KoLeoLoss()(x).item()) + + def test_forward__group_size(self) -> None: + """Groups are consecutive chunks of the batch and are averaged over.""" + torch.manual_seed(0) + x = torch.randn(8, 4) + + loss = KoLeoLoss(group_size=4)(x) + + loss_fn = KoLeoLoss() + expected = 0.5 * (loss_fn(x[:4]) + loss_fn(x[4:])) + assert loss == pytest.approx(expected.item(), rel=1e-5) + + def test_forward__group_size_finds_neighbors_within_group(self) -> None: + """A neighbor in another group is ignored. + + The batch holds two copies of the same feature in different groups. Without + grouping the nearest neighbor is the identical feature at distance zero, which + gives a much larger loss than the within-group neighbor. + """ + x = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]]) + assert KoLeoLoss(group_size=2)(x) == pytest.approx(-math.log(2**0.5), rel=1e-4) + assert KoLeoLoss()(x) > 10.0 + + @pytest.mark.parametrize("topk", [1, 2, 3]) + def test_forward__topk(self, topk: int) -> None: + """The k nearest neighbors of every feature are penalized.""" + torch.manual_seed(0) + x = F.normalize(torch.randn(6, 4), dim=-1) + + loss = KoLeoLoss(topk=topk)(x) + + cos_sim = x @ x.t() + cos_sim.fill_diagonal_(-2) + nn_idx = cos_sim.topk(k=topk, dim=-1).indices + distances = torch.stack( + [(x - x[idx]).norm(dim=-1) for idx in nn_idx.unbind(dim=-1)] + ) + expected = -(distances + 1e-8).log().mean() + assert loss == pytest.approx(expected.item(), rel=1e-4) + + def test_forward__empty_batch(self) -> None: + with pytest.raises(ValueError, match="non-empty batch"): + KoLeoLoss()(torch.randn(0, 4)) + + def test_forward__batch_size_not_divisible_by_group_size(self) -> None: + with pytest.raises(ValueError, match="must be divisible by group size"): + KoLeoLoss(group_size=3)(torch.randn(8, 4)) + + def test_forward__topk_not_smaller_than_group_size(self) -> None: + with pytest.raises(ValueError, match="must be larger than topk"): + KoLeoLoss(topk=4, group_size=4)(torch.randn(8, 4)) + + @pytest.mark.parametrize("topk, group_size", [(0, None), (1, 0)]) + def test__init__invalid_parameters(self, topk: int, group_size: int) -> None: + with pytest.raises(ValueError, match="must be positive"): + KoLeoLoss(topk=topk, group_size=group_size) + + def test_gather_distributed_world_size_one_does_not_gather( + self, mocker: MockerFixture + ) -> None: + mock_gather = mocker.patch("lightly.loss.koleo_loss.lightly_dist.gather") + + torch.manual_seed(0) + loss = KoLeoLoss(gather_distributed=True)(torch.randn(4, 8)) + + assert loss.isfinite() + mock_gather.assert_not_called() + + def test_gather_distributed_matches_non_distributed( + self, mocker: MockerFixture + ) -> None: + """Gathered forward equals non-distributed forward on the global batch. + + Simulates ``world_size=2`` with different data on both ranks, so the global + batch is the concatenation of the two local batches. + """ + torch.manual_seed(0) + rank_0, rank_1 = torch.randn(4, 8), torch.randn(4, 8) + + # Non-distributed truth: loss on the concatenated global batch. + expected = KoLeoLoss(group_size=4)(torch.cat([rank_0, rank_1])) + + mocker.patch("lightly.loss.koleo_loss.lightly_dist.world_size", return_value=2) + # gather is called with the already normalized features of the local rank. + mocker.patch( + "lightly.loss.koleo_loss.lightly_dist.gather", + side_effect=lambda tensor: ( + tensor, + F.normalize(rank_1, p=2, dim=-1, eps=1e-8), + ), + ) + loss = KoLeoLoss(group_size=4, gather_distributed=True)(rank_0) + + assert loss == pytest.approx(expected.item(), rel=1e-5) + + def test_gather_distributed_gradient_matches_non_distributed( + self, mocker: MockerFixture + ) -> None: + """The gradient w.r.t. the local features matches the non-distributed one. + + A gathered forward that is correct can still produce a wrong backward, see + #1977. GatherLayer itself is mocked here, so this covers that grouping and + neighbor indexing keep the gradients attached to the right features. + """ + torch.manual_seed(0) + rank_0 = torch.randn(4, 8) + rank_1 = torch.randn(4, 8) + + # Non-distributed truth: gradient of the loss on the concatenated global batch. + global_batch = torch.cat([rank_0, rank_1]).requires_grad_() + KoLeoLoss(group_size=4)(global_batch).backward() + assert global_batch.grad is not None + expected_grad = global_batch.grad[: len(rank_0)] + + mocker.patch("lightly.loss.koleo_loss.lightly_dist.world_size", return_value=2) + mocker.patch( + "lightly.loss.koleo_loss.lightly_dist.gather", + side_effect=lambda tensor: ( + tensor, + F.normalize(rank_1, p=2, dim=-1, eps=1e-8), + ), + ) + local_batch = rank_0.clone().requires_grad_() + KoLeoLoss(group_size=4, gather_distributed=True)(local_batch).backward() + + assert local_batch.grad is not None + assert torch.allclose(local_batch.grad, expected_grad, atol=1e-6) From 335d3736889aaa335975b3e54c72957c90170bd3 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Mon, 10 Aug 2026 23:56:23 +0500 Subject: [PATCH 4/8] refactor(loss): share the Sinkhorn-Knopp implementation with SwaV 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. --- lightly/loss/swav_loss.py | 38 ++++++------------ lightly/models/modules/center.py | 62 +++++++++++++++-------------- tests/loss/test_swav_loss.py | 17 +++++++- tests/models/modules/test_center.py | 12 ++++-- 4 files changed, 70 insertions(+), 59 deletions(-) diff --git a/lightly/loss/swav_loss.py b/lightly/loss/swav_loss.py index 625081f1a..c58aedb43 100644 --- a/lightly/loss/swav_loss.py +++ b/lightly/loss/swav_loss.py @@ -6,6 +6,8 @@ import torch.nn.functional as F from torch import Tensor +from lightly.models.modules import center + @torch.no_grad() def sinkhorn( @@ -16,7 +18,9 @@ def sinkhorn( ) -> Tensor: """Distributed sinkhorn algorithm. - As outlined in [0] and implemented in [1]. + As outlined in [0] and implemented in [1]. Thin wrapper around + lightly.models.modules.center.sinkhorn_knopp, which holds the shared + implementation. - [0]: SwaV, 2020, https://arxiv.org/abs/2006.09882 - [1]: https://github.com/facebookresearch/swav/ @@ -35,31 +39,13 @@ def sinkhorn( Returns: Soft codes Q assigning each feature to a prototype. """ - world_size = 1 - if gather_distributed and dist.is_initialized(): - world_size = dist.get_world_size() - - # Get the exponential matrix and make it sum to 1 - Q = torch.exp(out / epsilon).t() - sum_Q = torch.sum(Q) - if world_size > 1: - dist.all_reduce(sum_Q) - Q /= sum_Q - - B = Q.shape[1] * world_size - - for _ in range(iterations): - # Normalize rows - sum_of_rows = torch.sum(Q, dim=1, keepdim=True) - if world_size > 1: - dist.all_reduce(sum_of_rows) - Q /= sum_of_rows - # Normalize columns - Q /= torch.sum(Q, dim=0, keepdim=True) - Q /= B - - Q *= B - return Q.t() + codes = center.sinkhorn_knopp( + x=out, + temperature=epsilon, + num_iterations=iterations, + gather_distributed=gather_distributed, + ) + return codes.to(out.dtype) class SwaVLoss(nn.Module): diff --git a/lightly/models/modules/center.py b/lightly/models/modules/center.py index 6d9d466f0..10fa66eb7 100644 --- a/lightly/models/modules/center.py +++ b/lightly/models/modules/center.py @@ -111,45 +111,49 @@ def sinkhorn_knopp( x: Tensor, temperature: Union[float, Tensor] = 0.04, num_iterations: int = 3, + gather_distributed: bool = True, ) -> Tensor: - """Returns teacher probabilities with Sinkhorn-Knopp centering as used in DINOv2 [0] - and DINOv3 [1]. - - Sinkhorn-Knopp centering originates from SwAV [2] and replaces the mean centering - followed by a softmax used in DINO [3]. Instead of subtracting a running center, the - sharpened teacher logits are normalized such that every prototype receives the same - total weight across the batch, which avoids collapse without tracking any state. - - Implementation is based on [4]. In distributed settings the normalization is computed - over the batches of all processes. - - This differs from lightly.loss.swav_loss.sinkhorn, which normalizes by - local_batch_size * world_size and only reduces across processes if requested. Here - the number of samples is reduced across processes, which is required for the iBOT - loss where the number of masked tokens differs between processes. - - - [0]: DINOv2, 2023, https://arxiv.org/abs/2304.07193 - - [1]: DINOv3, 2025, https://arxiv.org/abs/2508.10104 - - [2]: SwAV, 2020, https://arxiv.org/abs/2006.09882 - - [3]: DINO, 2021, https://arxiv.org/abs/2104.14294 + """Returns Sinkhorn-Knopp normalized probabilities as introduced in SwAV [0]. + + Instead of subtracting a running center, the sharpened logits are normalized such + that every prototype receives the same total weight across the batch, which avoids + collapse without tracking any state. DINOv2 [1] offers this as an alternative to the + mean centering of DINO [2], but keeps mean centering in its released configs, where + it reports no difference on ImageNet-1k. DINOv3 [3] uses it for both the DINO and + the iBOT objective. + + Implementation is based on [4] and shared with lightly.loss.swav_loss.sinkhorn. + The number of samples is reduced across processes instead of being derived from the + world size, because processes can hold a different number of samples. This is the + case for the iBOT loss, where the number of masked tokens differs between processes. + + - [0]: SwAV, 2020, https://arxiv.org/abs/2006.09882 + - [1]: DINOv2, 2023, https://arxiv.org/abs/2304.07193 + - [2]: DINO, 2021, https://arxiv.org/abs/2104.14294 + - [3]: DINOv3, 2025, https://arxiv.org/abs/2508.10104 - [4]: https://github.com/facebookresearch/dinov3/blob/main/dinov3/loss/dino_clstoken_loss.py Args: x: - Tensor with shape (batch_size, num_prototypes) containing the teacher - logits. + Tensor with shape (batch_size, num_prototypes) containing the logits. temperature: - Temperature used to sharpen the teacher logits. + Temperature used to sharpen the logits. num_iterations: Number of Sinkhorn-Knopp iterations. + gather_distributed: + If True, the normalization is computed over the batches of all processes. Returns: - Tensor with shape (batch_size, num_prototypes) containing the teacher - probabilities in float32. Every row sums to one. The probabilities are detached - from the computation graph, following the reference implementation. + Tensor with shape (batch_size, num_prototypes) containing the probabilities in + float32. Every row sums to one if num_iterations is at least one. The + probabilities are detached from the computation graph, following the reference + implementation. """ + gather = gather_distributed and dist.is_available() and dist.is_initialized() + # (batch_size, num_prototypes) -> (num_prototypes, batch_size) following the - # notation of the reference implementation. + # notation of the reference implementation. The exponential is computed in float32 + # because it overflows in half precision for the low temperatures used by DINO. Q = torch.exp(x.float() / temperature).t() num_prototypes = Q.shape[0] @@ -160,7 +164,7 @@ def sinkhorn_knopp( # of masked tokens differs between processes. local_num_samples = torch.tensor(Q.shape[1], device=Q.device, dtype=Q.dtype) sums = torch.stack([Q.sum(), local_num_samples]) - if dist.is_available() and dist.is_initialized(): + if gather: dist.all_reduce(sums) sum_Q, num_samples = sums[0], sums[1] @@ -170,7 +174,7 @@ def sinkhorn_knopp( for _ in range(num_iterations): # Normalize rows: the total weight per prototype must be 1 / num_prototypes. sum_of_rows = torch.sum(Q, dim=1, keepdim=True) - if dist.is_available() and dist.is_initialized(): + if gather: dist.all_reduce(sum_of_rows) Q /= sum_of_rows Q /= num_prototypes diff --git a/tests/loss/test_swav_loss.py b/tests/loss/test_swav_loss.py index df3ebc92b..c105a2973 100644 --- a/tests/loss/test_swav_loss.py +++ b/tests/loss/test_swav_loss.py @@ -3,7 +3,22 @@ from pytest_mock import MockerFixture from torch import distributed as dist -from lightly.loss import SwaVLoss +from lightly.loss import SwaVLoss, swav_loss + + +class TestSinkhorn: + def test__preserves_dtype(self) -> None: + # The shared implementation calculates in float32, the codes must be returned + # in the dtype of the input. + out = torch.rand(4, 8, dtype=torch.float16) + assert swav_loss.sinkhorn(out).dtype == torch.float16 + + @pytest.mark.parametrize("iterations", range(1, 4)) + def test__soft_codes(self, iterations: int) -> None: + out = torch.rand(4, 8) + codes = swav_loss.sinkhorn(out, iterations=iterations) + assert codes.shape == out.shape + assert torch.allclose(codes.sum(dim=1), torch.ones(4)) class TestSwaVLoss: diff --git a/tests/models/modules/test_center.py b/tests/models/modules/test_center.py index d11554c04..9bbea28c7 100644 --- a/tests/models/modules/test_center.py +++ b/tests/models/modules/test_center.py @@ -76,9 +76,15 @@ def test_update__momentum(self, momentum: float, expected: Tensor) -> None: assert torch.all(center.value == expected) -def _teacher_logits(batch_size: int, num_prototypes: int) -> Tensor: - """Returns logits in [-1, 1], as produced by a projection head with weight norm.""" - return F.normalize(torch.randn(batch_size, num_prototypes), dim=-1) +def _teacher_logits(batch_size: int, num_prototypes: int, dim: int = 16) -> Tensor: + """Returns logits as produced by a DINO projection head. + + The head L2-normalizes the features before the weight-normed prototype layer, so + the logits are cosine similarities and lie in [-1, 1]. + """ + features = F.normalize(torch.randn(batch_size, dim), dim=-1) + prototypes = F.normalize(torch.randn(dim, num_prototypes), dim=0) + return features @ prototypes class TestSinkhornKnopp: From f00cf527613db91e6a855291a44b4ffb42d1af27 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Mon, 10 Aug 2026 23:56:31 +0500 Subject: [PATCH 5/8] fix(loss): tighten validation of the new loss options 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. --- lightly/loss/dino_loss.py | 26 ++++++++++++++++++-------- lightly/loss/ibot_loss.py | 27 +++++++++++++++++++-------- lightly/loss/koleo_loss.py | 32 ++++++++++++++++++++++++-------- tests/loss/test_dino_loss.py | 19 +++++++++++++++++++ tests/loss/test_ibot_loss.py | 5 +++++ tests/loss/test_koleo_loss.py | 34 +++++++++++++++++++++++++--------- 6 files changed, 110 insertions(+), 33 deletions(-) diff --git a/lightly/loss/dino_loss.py b/lightly/loss/dino_loss.py index 607c8f0aa..30eb80e4f 100644 --- a/lightly/loss/dino_loss.py +++ b/lightly/loss/dino_loss.py @@ -36,16 +36,16 @@ class DINOLoss(Module): Temperature parameter for the student network. center: Center used for the teacher output. It is updated with a moving average - during training. Unused if 'sinkhorn_knopp' centering is selected. + during training. Unused if `center_mode="sinkhorn_knopp"` is selected. center_momentum: Momentum term for the center calculation. center_mode: Mode used to normalize the teacher output. Either 'mean' for the mean centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp centering - used by DINOv2 and DINOv3. + that is optional in DINOv2 and used by DINOv3. sinkhorn_iterations: - Number of Sinkhorn-Knopp iterations. Only used if center_mode is - 'sinkhorn_knopp'. + Number of Sinkhorn-Knopp iterations. Only used if + `center_mode="sinkhorn_knopp"`. warmup_teacher_temp_epochs: Number of epochs for the warmup phase of the teacher temperature (for backward compatibility). teacher_temp_schedule: @@ -83,10 +83,10 @@ def __init__( center_mode: Mode used to normalize the teacher output. Either 'mean' for the mean centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp - centering used by DINOv2 and DINOv3. + centering that is optional in DINOv2 and used by DINOv3. sinkhorn_iterations: - Number of Sinkhorn-Knopp iterations. Only used if center_mode is - 'sinkhorn_knopp'. + Number of Sinkhorn-Knopp iterations. Only used if + `center_mode="sinkhorn_knopp"`. warmup_teacher_temp: Initial temperature for the teacher network (for backward compatibility). warmup_teacher_temp_epochs: @@ -94,6 +94,7 @@ def __init__( Raises: ValueError: If an unknown center mode is provided. + ValueError: If sinkhorn_iterations is negative. """ super().__init__() @@ -107,6 +108,11 @@ def __init__( f"Unknown mode '{center_mode}'. Valid modes are " f"{sorted(VALID_CENTER_MODES)}." ) + if sinkhorn_iterations < 0: + raise ValueError( + f"sinkhorn_iterations must not be negative but is " + f"{sinkhorn_iterations}." + ) self.center_mode = center_mode self.sinkhorn_iterations = sinkhorn_iterations # No center is tracked with Sinkhorn-Knopp centering. The center buffer is @@ -212,7 +218,9 @@ def _teacher_probabilities( Returns: Tensor with the same shape as teacher_out containing probabilities that - sum to one along the last dimension. + sum to one along the last dimension. Sinkhorn-Knopp probabilities are + detached from the computation graph, following the reference + implementation. """ if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: # Sinkhorn-Knopp is applied jointly over all views, following the reference @@ -223,6 +231,8 @@ def _teacher_probabilities( temperature=teacher_temp, num_iterations=self.sinkhorn_iterations, ) + # Sinkhorn-Knopp calculates in float32. The probabilities are cast back + # because the einsum with the student output does not promote dtypes. return probabilities.reshape(teacher_out.shape).to(teacher_out.dtype) return F.softmax((teacher_out - self.center) / teacher_temp, dim=-1) diff --git a/lightly/loss/ibot_loss.py b/lightly/loss/ibot_loss.py index 50d6a3758..eb2e769ba 100644 --- a/lightly/loss/ibot_loss.py +++ b/lightly/loss/ibot_loss.py @@ -8,7 +8,6 @@ from lightly.models.modules import center as center_module from lightly.models.modules.center import ( CENTER_MODE_SINKHORN_KNOPP, - CENTER_MODE_TO_FUNCTION, VALID_CENTER_MODES, Center, ) @@ -33,12 +32,12 @@ class IBOTPatchLoss(Module): center_mode: Mode used to normalize the teacher output. Either 'mean' for the mean centering from DINO or 'sinkhorn_knopp' for the Sinkhorn-Knopp centering - used by DINOv2 and DINOv3. + that is optional in DINOv2 and used by DINOv3. center_momentum: Momentum term for the center update. sinkhorn_iterations: - Number of Sinkhorn-Knopp iterations. Only used if center_mode is - 'sinkhorn_knopp'. + Number of Sinkhorn-Knopp iterations. Only used if + `center_mode="sinkhorn_knopp"`. """ def __init__( @@ -54,6 +53,7 @@ def __init__( Raises: ValueError: If an unknown center mode is provided. + ValueError: If sinkhorn_iterations is negative. """ super().__init__() @@ -65,14 +65,21 @@ def __init__( f"Unknown mode '{center_mode}'. Valid modes are " f"{sorted(VALID_CENTER_MODES)}." ) + if sinkhorn_iterations < 0: + raise ValueError( + f"sinkhorn_iterations must not be negative but is " + f"{sinkhorn_iterations}." + ) self.center_mode = center_mode self.sinkhorn_iterations = sinkhorn_iterations # Sinkhorn-Knopp centering does not track a center. The Center module is still - # created to keep the state dict independent of the center mode. + # created, with the default mode, to keep the state dict independent of the + # center mode. + tracks_center = center_mode != CENTER_MODE_SINKHORN_KNOPP self.center = Center( size=(1, output_dim), - mode=center_mode if center_mode in CENTER_MODE_TO_FUNCTION else "mean", + mode=center_mode if tracks_center else "mean", momentum=center_momentum, ) @@ -93,9 +100,13 @@ def _teacher_probabilities( Returns: Tensor with the same shape as teacher_out containing probabilities that - sum to one along the last dimension. + sum to one along the last dimension. Sinkhorn-Knopp probabilities are + detached from the computation graph, following the reference + implementation. """ if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + # Sinkhorn-Knopp calculates in float32, the probabilities are cast back to + # keep the dtype of the loss independent of the center mode. probabilities = center_module.sinkhorn_knopp( x=teacher_out, temperature=teacher_temp, @@ -115,7 +126,7 @@ def _update_center(self, teacher_out: Tensor) -> None: Tensor with shape (num_tokens, output_dim) containing the teacher output. """ - if self.center_mode not in CENTER_MODE_TO_FUNCTION: + if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: return self.center.update(teacher_out) diff --git a/lightly/loss/koleo_loss.py b/lightly/loss/koleo_loss.py index c6450eeb7..612ef90ad 100644 --- a/lightly/loss/koleo_loss.py +++ b/lightly/loss/koleo_loss.py @@ -1,6 +1,7 @@ from __future__ import annotations import torch +import torch.distributed as torch_dist from torch import Tensor from torch.nn import Module, PairwiseDistance, functional @@ -72,12 +73,15 @@ class KoLeoLoss(Module): used as a single group. DINOv3 uses a group size of 16. gather_distributed: If True, features from all GPUs are gathered before the batch is split into - groups. Requires that the distributed process group is initialized. + groups. Has no effect if the distributed process group is not initialized. Examples: >>> # initialize loss function >>> loss_fn = KoLeoLoss() >>> + >>> # or with the settings used by DINOv3 + >>> loss_fn = KoLeoLoss(group_size=16, gather_distributed=True) + >>> >>> # generate the features of a batch of images >>> features = model(images) >>> @@ -108,16 +112,25 @@ def __init__( batch is used as a single group. DINOv3 uses a group size of 16. gather_distributed: If True, features from all GPUs are gathered before the batch is split - into groups. Requires that the distributed process group is initialized. + into groups. Has no effect if the distributed process group is not + initialized. Raises: ValueError: If topk or group_size are not positive. + ValueError: If gather_distributed is True but torch.distributed is not + available. """ super().__init__() if topk < 1: raise ValueError(f"topk must be positive but is {topk}.") if group_size is not None and group_size < 1: raise ValueError(f"group_size must be positive but is {group_size}.") + if gather_distributed and not torch_dist.is_available(): + raise ValueError( + "gather_distributed is True but torch.distributed is not available. " + "Please set gather_distributed=False or install a torch version with " + "distributed support." + ) self.p = p self.eps = eps @@ -136,8 +149,8 @@ def forward(self, x: Tensor) -> Tensor: Loss value. Raises: - ValueError: If the batch size is not divisible by group_size or if - group_size is not larger than topk. + ValueError: If the batch is empty, if the batch size is not divisible by + group_size, or if topk is too large for the group size. """ # Normalize the input tensor x = functional.normalize(x, p=2, dim=-1, eps=self.eps) @@ -156,11 +169,14 @@ def forward(self, x: Tensor) -> Tensor: raise ValueError( f"Batch size {batch_size} must be divisible by group size {group_size}." ) - # A group must hold topk neighbors besides the feature itself, except for the - # degenerate group of size one, which is allowed for backwards compatibility. - if group_size > 1 and self.topk >= group_size: + # A group must hold topk neighbors besides the feature itself. Groups of size + # one are the exception: there the feature is its own neighbor, which keeps + # a batch size of one working as it did before groups existed. + max_topk = max(group_size - 1, 1) + if self.topk > max_topk: raise ValueError( - f"Group size {group_size} must be larger than topk {self.topk}." + f"topk {self.topk} must not be larger than {max_topk} for group size " + f"{group_size}." ) # Get the nearest neighbors and their distances. diff --git a/tests/loss/test_dino_loss.py b/tests/loss/test_dino_loss.py index cf1319563..12d86e2b3 100644 --- a/tests/loss/test_dino_loss.py +++ b/tests/loss/test_dino_loss.py @@ -135,6 +135,11 @@ def test__init__invalid_center_mode(self) -> None: with pytest.raises(ValueError, match="Unknown mode"): DINOLoss(output_dim=4, center_mode="invalid") + def test__init__negative_sinkhorn_iterations(self) -> None: + DINOLoss(output_dim=4, sinkhorn_iterations=0) + with pytest.raises(ValueError, match="must not be negative"): + DINOLoss(output_dim=4, sinkhorn_iterations=-1) + def test_sinkhorn_knopp(self) -> None: """Sinkhorn-Knopp centering is applied jointly over all teacher views. @@ -182,6 +187,20 @@ def test_sinkhorn_knopp(self) -> None: assert torch.allclose(loss, expected) + def test_sinkhorn_knopp__half_precision(self) -> None: + """The teacher probabilities keep the dtype of the teacher output. + + Sinkhorn-Knopp calculates in float32 and the einsum with the student output + does not promote dtypes, so the loss fails if the cast back is dropped. + """ + loss_fn = DINOLoss(output_dim=4, center_mode="sinkhorn_knopp") + teacher_out = [t.half() for t in _generate_output(n_views=2, output_dim=4)] + student_out = [s.half() for s in _generate_output(n_views=4, output_dim=4)] + + loss = loss_fn(teacher_out=teacher_out, student_out=student_out) + + assert loss.isfinite() + def test_sinkhorn_knopp__center_not_updated(self) -> None: """Sinkhorn-Knopp does not track a center, but keeps the buffer registered.""" loss_fn = DINOLoss(output_dim=4, center_mode="sinkhorn_knopp") diff --git a/tests/loss/test_ibot_loss.py b/tests/loss/test_ibot_loss.py index 348963b85..3b6e4fa40 100644 --- a/tests/loss/test_ibot_loss.py +++ b/tests/loss/test_ibot_loss.py @@ -51,6 +51,11 @@ def test__init__invalid_center_mode(self) -> None: with pytest.raises(ValueError, match="Unknown mode"): IBOTPatchLoss(output_dim=2, center_mode="invalid") + def test__init__negative_sinkhorn_iterations(self) -> None: + IBOTPatchLoss(output_dim=2, sinkhorn_iterations=0) + with pytest.raises(ValueError, match="must not be negative"): + IBOTPatchLoss(output_dim=2, sinkhorn_iterations=-1) + def test_sinkhorn_knopp(self) -> None: """Sinkhorn-Knopp centering replaces the mean centering of the teacher output. diff --git a/tests/loss/test_koleo_loss.py b/tests/loss/test_koleo_loss.py index 5ff18dae9..a7d18f0fc 100644 --- a/tests/loss/test_koleo_loss.py +++ b/tests/loss/test_koleo_loss.py @@ -6,6 +6,7 @@ import torch.nn.functional as F from pytest_mock import MockerFixture from torch import Tensor +from torch import distributed as torch_dist from lightly.loss.koleo_loss import KoLeoLoss @@ -46,7 +47,7 @@ def test_forward(self, x: Tensor, expected_loss: float, device: str) -> None: def test_forward__group_size_none_is_full_batch(self) -> None: torch.manual_seed(0) x = torch.randn(8, 4) - assert KoLeoLoss(group_size=8)(x) == pytest.approx(KoLeoLoss()(x).item()) + assert KoLeoLoss(group_size=8)(x).item() == pytest.approx(KoLeoLoss()(x).item()) def test_forward__group_size(self) -> None: """Groups are consecutive chunks of the batch and are averaged over.""" @@ -57,7 +58,7 @@ def test_forward__group_size(self) -> None: loss_fn = KoLeoLoss() expected = 0.5 * (loss_fn(x[:4]) + loss_fn(x[4:])) - assert loss == pytest.approx(expected.item(), rel=1e-5) + assert loss.item() == pytest.approx(expected.item(), rel=1e-5) def test_forward__group_size_finds_neighbors_within_group(self) -> None: """A neighbor in another group is ignored. @@ -67,8 +68,10 @@ def test_forward__group_size_finds_neighbors_within_group(self) -> None: gives a much larger loss than the within-group neighbor. """ x = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]]) - assert KoLeoLoss(group_size=2)(x) == pytest.approx(-math.log(2**0.5), rel=1e-4) - assert KoLeoLoss()(x) > 10.0 + assert KoLeoLoss(group_size=2)(x).item() == pytest.approx( + -math.log(2**0.5), rel=1e-4 + ) + assert KoLeoLoss()(x).item() > 10.0 @pytest.mark.parametrize("topk", [1, 2, 3]) def test_forward__topk(self, topk: int) -> None: @@ -85,7 +88,7 @@ def test_forward__topk(self, topk: int) -> None: [(x - x[idx]).norm(dim=-1) for idx in nn_idx.unbind(dim=-1)] ) expected = -(distances + 1e-8).log().mean() - assert loss == pytest.approx(expected.item(), rel=1e-4) + assert loss.item() == pytest.approx(expected.item(), rel=1e-4) def test_forward__empty_batch(self) -> None: with pytest.raises(ValueError, match="non-empty batch"): @@ -95,15 +98,28 @@ def test_forward__batch_size_not_divisible_by_group_size(self) -> None: with pytest.raises(ValueError, match="must be divisible by group size"): KoLeoLoss(group_size=3)(torch.randn(8, 4)) - def test_forward__topk_not_smaller_than_group_size(self) -> None: - with pytest.raises(ValueError, match="must be larger than topk"): - KoLeoLoss(topk=4, group_size=4)(torch.randn(8, 4)) + @pytest.mark.parametrize("topk, group_size", [(4, 4), (5, 4), (2, 1)]) + def test_forward__topk_too_large_for_group( + self, topk: int, group_size: int + ) -> None: + with pytest.raises(ValueError, match="must not be larger than"): + KoLeoLoss(topk=topk, group_size=group_size)(torch.randn(8, 4)) @pytest.mark.parametrize("topk, group_size", [(0, None), (1, 0)]) def test__init__invalid_parameters(self, topk: int, group_size: int) -> None: with pytest.raises(ValueError, match="must be positive"): KoLeoLoss(topk=topk, group_size=group_size) + def test__init__gather_distributed_dist_not_available( + self, mocker: MockerFixture + ) -> None: + mock_is_available = mocker.patch.object( + torch_dist, "is_available", return_value=False + ) + with pytest.raises(ValueError, match="torch.distributed is not available"): + KoLeoLoss(gather_distributed=True) + mock_is_available.assert_called_once() + def test_gather_distributed_world_size_one_does_not_gather( self, mocker: MockerFixture ) -> None: @@ -140,7 +156,7 @@ def test_gather_distributed_matches_non_distributed( ) loss = KoLeoLoss(group_size=4, gather_distributed=True)(rank_0) - assert loss == pytest.approx(expected.item(), rel=1e-5) + assert loss.item() == pytest.approx(expected.item(), rel=1e-5) def test_gather_distributed_gradient_matches_non_distributed( self, mocker: MockerFixture From 5336c4bba69effea49e847431ffe021e4c326c03 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Tue, 11 Aug 2026 00:07:53 +0500 Subject: [PATCH 6/8] test(loss): add multi-rank KoLeoLoss gather tests 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 #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. --- tests/loss/test_koleo_loss.py | 68 +++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/tests/loss/test_koleo_loss.py b/tests/loss/test_koleo_loss.py index a7d18f0fc..2222e6888 100644 --- a/tests/loss/test_koleo_loss.py +++ b/tests/loss/test_koleo_loss.py @@ -9,6 +9,32 @@ from torch import distributed as torch_dist from lightly.loss.koleo_loss import KoLeoLoss +from tests.ddp_helpers import NUM_PROCESSES, USE_PYTEST_POOL + +# Group size used by the distributed tests. The global batch is the concatenation of +# the local batches, so a group spans the features of both ranks. +DDP_GROUP_SIZE = 4 + + +def _local_batch(batch: Tensor, rank: int, world_size: int) -> Tensor: + """Returns the chunk of the global batch that belongs to the given rank.""" + return batch.chunk(world_size)[rank] + + +def _forward_worker(rank: int, world_size: int, batch: Tensor) -> Tensor: + # Pool worker: distributed forward on one rank. See #1982. + loss: Tensor = KoLeoLoss(group_size=DDP_GROUP_SIZE, gather_distributed=True)( + _local_batch(batch, rank, world_size) + ) + return loss.detach() + + +def _gradient_worker(rank: int, world_size: int, batch: Tensor) -> Tensor: + # Pool worker: distributed backward on one rank, returns the input gradient. + x = _local_batch(batch, rank, world_size).clone().requires_grad_(True) + KoLeoLoss(group_size=DDP_GROUP_SIZE, gather_distributed=True)(x).backward() + assert x.grad is not None + return x.grad class TestKoLeoLoss: @@ -190,3 +216,45 @@ def test_gather_distributed_gradient_matches_non_distributed( assert local_batch.grad is not None assert torch.allclose(local_batch.grad, expected_grad, atol=1e-6) + + @pytest.mark.DDP + @pytest.mark.skipif(not USE_PYTEST_POOL, reason="DDP pool is not available") + def test__gather_distributed_forward_matches_non_distributed(self) -> None: + # Every rank gathers the same global batch, so every rank must see the loss + # of the non-distributed run on the concatenated batch. + torch.manual_seed(0) + batch = torch.randn(NUM_PROCESSES * DDP_GROUP_SIZE, 8) + + losses = pytest.pool.starmap( # type: ignore[attr-defined] + _forward_worker, + [(rank, NUM_PROCESSES, batch) for rank in range(NUM_PROCESSES)], + ) + loss_truth = KoLeoLoss(group_size=DDP_GROUP_SIZE)(batch) + + assert all(torch.allclose(loss, loss_truth, atol=1e-5) for loss in losses) + + @pytest.mark.DDP + @pytest.mark.skipif(not USE_PYTEST_POOL, reason="DDP pool is not available") + def test__gather_distributed_gradient_matches_non_distributed(self) -> None: + # Complements the mocked world_size=2 test above, which bypasses GatherLayer + # and therefore does not cover its backward. + # + # Every rank computes the loss of the whole global batch, so GatherLayer's + # backward all_reduces NUM_PROCESSES identical gradients before slicing out + # the local one. The local gradient is therefore NUM_PROCESSES times the + # non-distributed one, which DDP cancels when it averages parameter + # gradients across ranks. + torch.manual_seed(0) + batch = torch.randn(NUM_PROCESSES * DDP_GROUP_SIZE, 8) + + grads = pytest.pool.starmap( # type: ignore[attr-defined] + _gradient_worker, + [(rank, NUM_PROCESSES, batch) for rank in range(NUM_PROCESSES)], + ) + batch_ref = batch.clone().requires_grad_(True) + KoLeoLoss(group_size=DDP_GROUP_SIZE)(batch_ref).backward() + assert batch_ref.grad is not None + + for rank, grad in enumerate(grads): + expected = NUM_PROCESSES * _local_batch(batch_ref.grad, rank, NUM_PROCESSES) + assert torch.allclose(grad, expected, atol=1e-5) From 71088bfaffffeecfb491984762150918cd3b8bd0 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Thu, 13 Aug 2026 00:26:50 +0500 Subject: [PATCH 7/8] fix(loss): address review feedback on center_mode typing and validation 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. --- lightly/loss/dino_loss.py | 37 +++++++++++++++------------- lightly/loss/ibot_loss.py | 41 +++++++++++++++++--------------- lightly/models/modules/center.py | 6 ----- tests/loss/test_dino_loss.py | 10 ++++---- tests/loss/test_ibot_loss.py | 10 ++++---- 5 files changed, 52 insertions(+), 52 deletions(-) diff --git a/lightly/loss/dino_loss.py b/lightly/loss/dino_loss.py index 30eb80e4f..acd581c2b 100644 --- a/lightly/loss/dino_loss.py +++ b/lightly/loss/dino_loss.py @@ -1,6 +1,7 @@ from __future__ import annotations import warnings +from typing import TYPE_CHECKING import torch import torch.nn.functional as F @@ -8,11 +9,10 @@ from torch.nn import Module, Parameter from lightly.models.modules import center -from lightly.models.modules.center import ( - CENTER_MODE_SINKHORN_KNOPP, - CENTER_MODE_TO_FUNCTION, - VALID_CENTER_MODES, -) +from lightly.models.modules.center import CENTER_MODE_TO_FUNCTION + +if TYPE_CHECKING: + from typing import Literal class DINOLoss(Module): @@ -74,7 +74,7 @@ def __init__( warmup_teacher_temp_epochs: int = 30, student_temp: float = 0.1, center_momentum: float = 0.9, - center_mode: str = "mean", + center_mode: Literal["mean", "sinkhorn_knopp"] = "mean", sinkhorn_iterations: int = 3, ) -> None: """Initializes the DINOLoss Module. @@ -94,7 +94,7 @@ def __init__( Raises: ValueError: If an unknown center mode is provided. - ValueError: If sinkhorn_iterations is negative. + ValueError: If sinkhorn_iterations is less than 1. """ super().__init__() @@ -103,21 +103,24 @@ def __init__( # TODO(Guarin, 08/24): Refactor this to use the Center module directly once # we do a breaking change. - if center_mode not in VALID_CENTER_MODES: + if center_mode == "mean": + center_fn = CENTER_MODE_TO_FUNCTION["mean"] + elif center_mode == "sinkhorn_knopp": + # No center is tracked with Sinkhorn-Knopp centering. The center buffer is + # still registered to keep the state dict independent of the center mode. + center_fn = None + else: raise ValueError( - f"Unknown mode '{center_mode}'. Valid modes are " - f"{sorted(VALID_CENTER_MODES)}." + f"Unknown mode '{center_mode}'. Valid modes are 'mean' and " + "'sinkhorn_knopp'." ) - if sinkhorn_iterations < 0: + if sinkhorn_iterations < 1: raise ValueError( - f"sinkhorn_iterations must not be negative but is " - f"{sinkhorn_iterations}." + f"sinkhorn_iterations must be at least 1 but is {sinkhorn_iterations}." ) self.center_mode = center_mode self.sinkhorn_iterations = sinkhorn_iterations - # No center is tracked with Sinkhorn-Knopp centering. The center buffer is - # still registered to keep the state dict independent of the center mode. - self._center_fn = CENTER_MODE_TO_FUNCTION.get(center_mode) + self._center_fn = center_fn self.center: Parameter self.register_buffer("center", torch.zeros(1, 1, output_dim)) self.center_momentum = center_momentum @@ -222,7 +225,7 @@ def _teacher_probabilities( detached from the computation graph, following the reference implementation. """ - if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + if self.center_mode == "sinkhorn_knopp": # Sinkhorn-Knopp is applied jointly over all views, following the reference # implementation. # (num_views, batch_size, output_dim) -> (num_views * batch_size, output_dim) diff --git a/lightly/loss/ibot_loss.py b/lightly/loss/ibot_loss.py index eb2e769ba..02681dce2 100644 --- a/lightly/loss/ibot_loss.py +++ b/lightly/loss/ibot_loss.py @@ -1,16 +1,17 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import torch from torch import Tensor from torch.nn import Module from torch.nn import functional as F from lightly.models.modules import center as center_module -from lightly.models.modules.center import ( - CENTER_MODE_SINKHORN_KNOPP, - VALID_CENTER_MODES, - Center, -) +from lightly.models.modules.center import Center + +if TYPE_CHECKING: + from typing import Literal class IBOTPatchLoss(Module): @@ -45,7 +46,7 @@ def __init__( output_dim: int = 65536, teacher_temp: float = 0.04, student_temp: float = 0.1, - center_mode: str = "mean", + center_mode: Literal["mean", "sinkhorn_knopp"] = "mean", center_momentum: float = 0.9, sinkhorn_iterations: int = 3, ) -> None: @@ -53,30 +54,32 @@ def __init__( Raises: ValueError: If an unknown center mode is provided. - ValueError: If sinkhorn_iterations is negative. + ValueError: If sinkhorn_iterations is less than 1. """ super().__init__() self.teacher_temp = teacher_temp self.student_temp = student_temp - if center_mode not in VALID_CENTER_MODES: + if center_mode == "mean": + tracks_center = True + elif center_mode == "sinkhorn_knopp": + # Sinkhorn-Knopp centering does not track a center. + tracks_center = False + else: raise ValueError( - f"Unknown mode '{center_mode}'. Valid modes are " - f"{sorted(VALID_CENTER_MODES)}." + f"Unknown mode '{center_mode}'. Valid modes are 'mean' and " + "'sinkhorn_knopp'." ) - if sinkhorn_iterations < 0: + if sinkhorn_iterations < 1: raise ValueError( - f"sinkhorn_iterations must not be negative but is " - f"{sinkhorn_iterations}." + f"sinkhorn_iterations must be at least 1 but is {sinkhorn_iterations}." ) self.center_mode = center_mode self.sinkhorn_iterations = sinkhorn_iterations - # Sinkhorn-Knopp centering does not track a center. The Center module is still - # created, with the default mode, to keep the state dict independent of the - # center mode. - tracks_center = center_mode != CENTER_MODE_SINKHORN_KNOPP + # The Center module is still created, with the default mode, to keep the state + # dict independent of the center mode. self.center = Center( size=(1, output_dim), mode=center_mode if tracks_center else "mean", @@ -104,7 +107,7 @@ def _teacher_probabilities( detached from the computation graph, following the reference implementation. """ - if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + if self.center_mode == "sinkhorn_knopp": # Sinkhorn-Knopp calculates in float32, the probabilities are cast back to # keep the dtype of the loss independent of the center mode. probabilities = center_module.sinkhorn_knopp( @@ -126,7 +129,7 @@ def _update_center(self, teacher_out: Tensor) -> None: Tensor with shape (num_tokens, output_dim) containing the teacher output. """ - if self.center_mode == CENTER_MODE_SINKHORN_KNOPP: + if self.center_mode == "sinkhorn_knopp": return self.center.update(teacher_out) diff --git a/lightly/models/modules/center.py b/lightly/models/modules/center.py index 10fa66eb7..795869565 100644 --- a/lightly/models/modules/center.py +++ b/lightly/models/modules/center.py @@ -188,12 +188,6 @@ def sinkhorn_knopp( return Q.t() -CENTER_MODE_SINKHORN_KNOPP = "sinkhorn_knopp" - CENTER_MODE_TO_FUNCTION = { "mean": center_mean, } - -# Modes accepted by the losses. The Center module only supports the modes in -# CENTER_MODE_TO_FUNCTION, as Sinkhorn-Knopp does not track a center. -VALID_CENTER_MODES = [*CENTER_MODE_TO_FUNCTION, CENTER_MODE_SINKHORN_KNOPP] diff --git a/tests/loss/test_dino_loss.py b/tests/loss/test_dino_loss.py index 12d86e2b3..ced1d4515 100644 --- a/tests/loss/test_dino_loss.py +++ b/tests/loss/test_dino_loss.py @@ -133,12 +133,12 @@ def test_other_parameters( def test__init__invalid_center_mode(self) -> None: with pytest.raises(ValueError, match="Unknown mode"): - DINOLoss(output_dim=4, center_mode="invalid") + DINOLoss(output_dim=4, center_mode="invalid") # type: ignore[arg-type] - def test__init__negative_sinkhorn_iterations(self) -> None: - DINOLoss(output_dim=4, sinkhorn_iterations=0) - with pytest.raises(ValueError, match="must not be negative"): - DINOLoss(output_dim=4, sinkhorn_iterations=-1) + @pytest.mark.parametrize("sinkhorn_iterations", [0, -1]) + def test__init__invalid_sinkhorn_iterations(self, sinkhorn_iterations: int) -> None: + with pytest.raises(ValueError, match="must be at least 1"): + DINOLoss(output_dim=4, sinkhorn_iterations=sinkhorn_iterations) def test_sinkhorn_knopp(self) -> None: """Sinkhorn-Knopp centering is applied jointly over all teacher views. diff --git a/tests/loss/test_ibot_loss.py b/tests/loss/test_ibot_loss.py index 3b6e4fa40..4c08f69cc 100644 --- a/tests/loss/test_ibot_loss.py +++ b/tests/loss/test_ibot_loss.py @@ -49,12 +49,12 @@ def test_forward(self, device: str) -> None: def test__init__invalid_center_mode(self) -> None: with pytest.raises(ValueError, match="Unknown mode"): - IBOTPatchLoss(output_dim=2, center_mode="invalid") + IBOTPatchLoss(output_dim=2, center_mode="invalid") # type: ignore[arg-type] - def test__init__negative_sinkhorn_iterations(self) -> None: - IBOTPatchLoss(output_dim=2, sinkhorn_iterations=0) - with pytest.raises(ValueError, match="must not be negative"): - IBOTPatchLoss(output_dim=2, sinkhorn_iterations=-1) + @pytest.mark.parametrize("sinkhorn_iterations", [0, -1]) + def test__init__invalid_sinkhorn_iterations(self, sinkhorn_iterations: int) -> None: + with pytest.raises(ValueError, match="must be at least 1"): + IBOTPatchLoss(output_dim=2, sinkhorn_iterations=sinkhorn_iterations) def test_sinkhorn_knopp(self) -> None: """Sinkhorn-Knopp centering replaces the mean centering of the teacher output. From 732a5b06becb3076aeae8cd044c83caa547e86f4 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Thu, 13 Aug 2026 00:27:01 +0500 Subject: [PATCH 8/8] fix(loss): fix CI failures on the pinned Python 3.7 leg 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 (#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. --- lightly/loss/koleo_loss.py | 7 +++++-- tests/loss/test_dino_loss.py | 17 +++++++++++------ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/lightly/loss/koleo_loss.py b/lightly/loss/koleo_loss.py index 612ef90ad..36f2df93c 100644 --- a/lightly/loss/koleo_loss.py +++ b/lightly/loss/koleo_loss.py @@ -40,9 +40,12 @@ def _nearest_neighbor_indices(x: Tensor, group_size: int, topk: int) -> Tensor: # (num_groups, group_size, topk) nn_idx = cos_sim.topk(k=topk, dim=-1).indices - # Shift the group-local indices so that they index into the flat batch. + # Shift the group-local indices so that they index into the flat batch. Assigned + # to a Tensor-annotated variable because mypy infers Any from the addition when + # checked against the older torch stubs used by the pinned CI environment. offset = torch.arange(num_groups, device=x.device) * group_size - return (nn_idx + offset.view(-1, 1, 1)).flatten() + indices: Tensor = (nn_idx + offset.view(-1, 1, 1)).flatten() + return indices class KoLeoLoss(Module): diff --git a/tests/loss/test_dino_loss.py b/tests/loss/test_dino_loss.py index ced1d4515..49a05b07b 100644 --- a/tests/loss/test_dino_loss.py +++ b/tests/loss/test_dino_loss.py @@ -190,16 +190,21 @@ def test_sinkhorn_knopp(self) -> None: def test_sinkhorn_knopp__half_precision(self) -> None: """The teacher probabilities keep the dtype of the teacher output. - Sinkhorn-Knopp calculates in float32 and the einsum with the student output - does not promote dtypes, so the loss fails if the cast back is dropped. + Sinkhorn-Knopp calculates in float32 and the einsum with the student output in + forward() does not promote dtypes, so the loss fails if the cast back is + dropped. Tested on _teacher_probabilities directly, rather than through + forward(), because F.log_softmax on the student output has no half-precision + CPU kernel in the older torch version pinned for the Python 3.7 CI leg. """ loss_fn = DINOLoss(output_dim=4, center_mode="sinkhorn_knopp") - teacher_out = [t.half() for t in _generate_output(n_views=2, output_dim=4)] - student_out = [s.half() for s in _generate_output(n_views=4, output_dim=4)] + teacher_out = torch.stack(_generate_output(n_views=2, output_dim=4)).half() - loss = loss_fn(teacher_out=teacher_out, student_out=student_out) + probabilities = loss_fn._teacher_probabilities( + teacher_out=teacher_out, teacher_temp=torch.tensor(0.04) + ) - assert loss.isfinite() + assert probabilities.dtype == torch.float16 + assert probabilities.isfinite().all() def test_sinkhorn_knopp__center_not_updated(self) -> None: """Sinkhorn-Knopp does not track a center, but keeps the buffer registered."""