diff --git a/lightly/loss/dino_loss.py b/lightly/loss/dino_loss.py
index fcebae018..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
@@ -10,6 +11,9 @@
 from lightly.models.modules import center
 from lightly.models.modules.center import CENTER_MODE_TO_FUNCTION
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 
 class DINOLoss(Module):
     """Implementation of the loss described in 'Emerging Properties in
@@ -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 `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
+            that is optional in DINOv2 and used by DINOv3.
+        sinkhorn_iterations:
+            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:
@@ -63,17 +74,27 @@ 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.
 
         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 that is optional in DINOv2 and used by DINOv3.
+            sinkhorn_iterations:
+                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:
                 Number of epochs for the warmup phase of the teacher temperature (for backward compatibility).
+
+        Raises:
+            ValueError: If an unknown center mode is provided.
+            ValueError: If sinkhorn_iterations is less than 1.
         """
         super().__init__()
 
@@ -82,12 +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 CENTER_MODE_TO_FUNCTION:
+        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 'mean' and "
+                "'sinkhorn_knopp'."
+            )
+        if sinkhorn_iterations < 1:
             raise ValueError(
-                f"Unknown mode '{center_mode}'. Valid modes are "
-                f"{sorted(CENTER_MODE_TO_FUNCTION.keys())}."
+                f"sinkhorn_iterations must be at least 1 but is {sinkhorn_iterations}."
             )
-        self._center_fn = CENTER_MODE_TO_FUNCTION[center_mode]
+        self.center_mode = center_mode
+        self.sinkhorn_iterations = sinkhorn_iterations
+        self._center_fn = center_fn
         self.center: Parameter
         self.register_buffer("center", torch.zeros(1, 1, output_dim))
         self.center_momentum = center_momentum
@@ -140,8 +173,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 +204,56 @@ 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. Sinkhorn-Knopp probabilities are
+            detached from the computation graph, following the reference
+            implementation.
+        """
+        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)
+            probabilities = center.sinkhorn_knopp(
+                x=teacher_out.flatten(0, 1),
+                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)
+
     @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/loss/ibot_loss.py b/lightly/loss/ibot_loss.py
index 4e32a6d01..02681dce2 100644
--- a/lightly/loss/ibot_loss.py
+++ b/lightly/loss/ibot_loss.py
@@ -1,12 +1,18 @@
 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
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 
 class IBOTPatchLoss(Module):
     """Implementation of the iBOT patch loss [0] as used in DINOv2 [1].
@@ -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
+            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="sinkhorn_knopp"`.
     """
 
     def __init__(
@@ -35,21 +46,93 @@ 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:
-        """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.
+            ValueError: If sinkhorn_iterations is less than 1.
+        """
         super().__init__()
 
         self.teacher_temp = teacher_temp
         self.student_temp = student_temp
 
+        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 'mean' and "
+                "'sinkhorn_knopp'."
+            )
+        if sinkhorn_iterations < 1:
+            raise ValueError(
+                f"sinkhorn_iterations must be at least 1 but is {sinkhorn_iterations}."
+            )
+        self.center_mode = center_mode
+        self.sinkhorn_iterations = sinkhorn_iterations
+
+        # 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,
+            mode=center_mode if tracks_center 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. Sinkhorn-Knopp probabilities are
+            detached from the computation graph, following the reference
+            implementation.
+        """
+        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(
+                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 == "sinkhorn_knopp":
+            return
+        self.center.update(teacher_out)
+
     def forward(
         self,
         teacher_out: Tensor,
@@ -85,8 +168,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 +186,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 +310,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 +332,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/lightly/loss/koleo_loss.py b/lightly/loss/koleo_loss.py
index 5f8751cb1..36f2df93c 100644
--- a/lightly/loss/koleo_loss.py
+++ b/lightly/loss/koleo_loss.py
@@ -1,7 +1,52 @@
+from __future__ import annotations
+
 import torch
+import torch.distributed as torch_dist
 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. 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
+    indices: Tensor = (nn_idx + offset.view(-1, 1, 1)).flatten()
+    return indices
+
 
 class KoLeoLoss(Module):
     """KoLeo loss based on [0].
@@ -10,22 +55,50 @@ 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. 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)
+        >>>
+        >>> # 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 +107,39 @@ 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. 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
+        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 +150,43 @@ def forward(self, x: Tensor) -> Tensor:
 
         Returns:
             Loss value.
+
+        Raises:
+            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)
 
-        # 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. 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"topk {self.topk} must not be larger than {max_topk} for group size "
+                f"{group_size}."
+            )
 
-        # 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/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 21be7b9bf..795869565 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,88 @@ 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,
+    gather_distributed: bool = True,
+) -> Tensor:
+    """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 logits.
+        temperature:
+            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 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. 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]
+
+    # 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 gather:
+        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 gather:
+            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_TO_FUNCTION = {
     "mean": center_mean,
 }
diff --git a/tests/loss/test_dino_loss.py b/tests/loss/test_dino_loss.py
index 164c6e3a3..49a05b07b 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,94 @@ 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")  # type: ignore[arg-type]
+
+    @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.
+
+        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__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 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 = torch.stack(_generate_output(n_views=2, output_dim=4)).half()
+
+        probabilities = loss_fn._teacher_probabilities(
+            teacher_out=teacher_out, teacher_temp=torch.tensor(0.04)
+        )
+
+        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."""
+        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/loss/test_ibot_loss.py b/tests/loss/test_ibot_loss.py
index 25a202eb2..4c08f69cc 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,52 @@ 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")  # type: ignore[arg-type]
+
+    @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.
+
+        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)
diff --git a/tests/loss/test_koleo_loss.py b/tests/loss/test_koleo_loss.py
index bfc1d8bff..2222e6888 100644
--- a/tests/loss/test_koleo_loss.py
+++ b/tests/loss/test_koleo_loss.py
@@ -4,9 +4,37 @@
 import torch
 import torch.nn as nn
 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
+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:
@@ -41,3 +69,192 @@ 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).item() == 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.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.
+
+        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).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:
+        """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.item() == 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))
+
+    @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:
+        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.item() == 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)
+
+    @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)
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 ccdb20458..9bbea28c7 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,54 @@ 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, 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:
+    @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