From 69200af70cb5197d101f4f7995ac3c35a129a22b Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 15 Aug 2026 09:41:59 -0300 Subject: [PATCH] refactor(loss): move the NT-Xent equation into lightly.functional --- lightly/functional/__init__.py | 8 +++ lightly/functional/ntxent.py | 95 +++++++++++++++++++++++++++++++++ lightly/loss/__init__.py | 1 + lightly/loss/distributed.py | 35 ++++++++++++ lightly/loss/ntx_ent_loss.py | 87 ++++++++---------------------- tests/functional/__init__.py | 0 tests/functional/test_ntxent.py | 69 ++++++++++++++++++++++++ 7 files changed, 229 insertions(+), 66 deletions(-) create mode 100644 lightly/functional/__init__.py create mode 100644 lightly/functional/ntxent.py create mode 100644 lightly/loss/distributed.py create mode 100644 tests/functional/__init__.py create mode 100644 tests/functional/test_ntxent.py diff --git a/lightly/functional/__init__.py b/lightly/functional/__init__.py new file mode 100644 index 000000000..4fa276fb8 --- /dev/null +++ b/lightly/functional/__init__.py @@ -0,0 +1,8 @@ +"""Pure functions over tensors: one equation each, no parameters, no state. + +A block that owns parameters, buffers or optimiser state does not belong here. +""" + +from lightly.functional.ntxent import ntxent + +__all__ = ["ntxent"] diff --git a/lightly/functional/ntxent.py b/lightly/functional/ntxent.py new file mode 100644 index 000000000..5b8472e74 --- /dev/null +++ b/lightly/functional/ntxent.py @@ -0,0 +1,95 @@ +"""The NT-Xent equation, with no state and no distributed communication.""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import Tensor +from torch.nn.functional import cross_entropy + +__all__ = ["ntxent"] + + +def ntxent( + z0: Tensor, + z1: Tensor, + *, + temperature: float, + negatives: Optional[Tensor] = None, + z0_all: Optional[Tensor] = None, + z1_all: Optional[Tensor] = None, + positives_at: int = 0, +) -> Tensor: + """Normalized temperature-scaled cross entropy, as in SimCLR. + + Every argument is a tensor or a number, so this function communicates with + no other rank and holds nothing between calls. Gathering the negatives is + :class:`~lightly.loss.ntx_ent_loss.NTXentLoss`'s job, and it passes the + result in through ``z0_all`` and ``z1_all``. + + Reference: + SimCLR, 2020, https://arxiv.org/abs/2002.05709 + + Args: + z0: + Projections of the first view, shape ``(B, D)``, L2-normalized. + z1: + Projections of the second view, shape ``(B, D)``, L2-normalized. + temperature: + Scales the logits by its inverse. + negatives: + Negatives from a memory bank, shape ``(D, K)``. When given, the other + samples in the batch are not used as negatives. + z0_all: + The first view's projections across all ranks, shape ``(B * W, D)``. + Defaults to ``z0``, which is the single-rank case. + z1_all: + The second view's projections across all ranks. Defaults to ``z1``. + positives_at: + The column at which this rank's rows start inside ``z0_all``, so + ``rank * B``. Zero on one rank. + + Returns: + The loss, as a scalar. + """ + device = z0.device + batch_size = z0.size(0) + + if negatives is not None: + # sim_pos[i] is the similarity of sample i to its positive pair, and + # sim_neg[i, j] its similarity to the j-th vector in the bank. + sim_pos = torch.einsum("nc,nc->n", z0, z1).unsqueeze(-1) + sim_neg = torch.einsum("nc,ck->nk", z0, negatives.to(device)) + logits = torch.cat([sim_pos, sim_neg], dim=1) / temperature + labels = torch.zeros(logits.size(0), device=device, dtype=torch.long) + return cross_entropy(logits, labels) + + z0_all = z0 if z0_all is None else z0_all + z1_all = z1 if z1_all is None else z1_all + + # The similarities of a view with itself, which carry no signal and are + # dropped below. On one rank this is torch.eye. + rows = torch.arange(batch_size, device=device, dtype=torch.long) + diagonal = torch.zeros(batch_size, z0_all.size(0), device=device, dtype=torch.bool) + diagonal[rows, rows + positives_at] = True + + # n is the local batch size and m the batch across ranks, so every block is + # (n, m). + logits_00 = torch.einsum("nc,mc->nm", z0, z0_all) / temperature + logits_01 = torch.einsum("nc,mc->nm", z0, z1_all) / temperature + logits_10 = torch.einsum("nc,mc->nm", z1, z0_all) / temperature + logits_11 = torch.einsum("nc,mc->nm", z1, z1_all) / temperature + + logits_00 = logits_00[~diagonal].view(batch_size, -1) + logits_11 = logits_11[~diagonal].view(batch_size, -1) + + logits = torch.cat( + [ + torch.cat([logits_01, logits_00], dim=1), + torch.cat([logits_10, logits_11], dim=1), + ], + dim=0, + ) + labels = (rows + positives_at).repeat(2) + return cross_entropy(logits, labels) diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py index 453349d33..fd568dead 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -9,6 +9,7 @@ from lightly.loss.detcon_loss import DetConBLoss, DetConSLoss from lightly.loss.dino_loss import DINOLoss from lightly.loss.directclr_loss import DirectCLRLoss +from lightly.loss.distributed import DistributedKind from lightly.loss.emp_ssl_loss import EMPSSLLoss from lightly.loss.frossl_loss import FroSSLLoss from lightly.loss.ibot_loss import IBOTPatchLoss, IBOTPlusPlusPatchLoss diff --git a/lightly/loss/distributed.py b/lightly/loss/distributed.py new file mode 100644 index 000000000..f5b132ef2 --- /dev/null +++ b/lightly/loss/distributed.py @@ -0,0 +1,35 @@ +"""How a loss behaves when the batch is split across ranks.""" + +from enum import Enum, auto + +__all__ = ["DistributedKind"] + + +class DistributedKind(Enum): + """What a loss needs from the ranks around it. + + DDP averages the per-rank loss and the per-rank gradient. Whether that is + correct depends on the loss, and the difference is not visible in a forward + pass: issue #1920 shipped a correct SIGReg forward with a gradient off by + exactly ``1 / world_size``, and came back as #1977 against + ``BarlowTwinsLoss``. Declaring the kind is what selects the test. + """ + + RANK_LOCAL = auto() + """The per-rank value is a valid estimate, so DDP averaging is correct. + + An MSE reconstruction loss. + """ + + GATHER_FOR_NEGATIVES = auto() + """Correct only if the features are gathered with gradient first. + + NT-Xent and CLIP: every sample on every rank is a negative for every other. + """ + + GLOBAL_STATISTIC = auto() + """A function of the global batch, so averaging per rank is wrong. + + SIGReg and the distribution-matching regularisers. The gap between the two + is the variance across ranks, which closes only once the ranks agree. + """ diff --git a/lightly/loss/ntx_ent_loss.py b/lightly/loss/ntx_ent_loss.py index 62f82621f..8e92fe313 100644 --- a/lightly/loss/ntx_ent_loss.py +++ b/lightly/loss/ntx_ent_loss.py @@ -9,6 +9,8 @@ from torch import Tensor, nn from torch import distributed as torch_dist +from lightly.functional.ntxent import ntxent +from lightly.loss.distributed import DistributedKind from lightly.models.modules.memory_bank import MemoryBankModule from lightly.utils import dist @@ -59,6 +61,8 @@ class NTXentLoss(nn.Module): """ + distributed_kind = DistributedKind.GATHER_FOR_NEGATIVES + def __init__( self, temperature: float = 0.5, @@ -85,7 +89,6 @@ def __init__( ) self.temperature = temperature self.gather_distributed = gather_distributed - self.cross_entropy = nn.CrossEntropyLoss(reduction="mean") self.eps = 1e-8 if abs(self.temperature) < self.eps: @@ -117,10 +120,8 @@ def forward(self, out0: Tensor, out1: Tensor) -> Tensor: Returns: Contrastive Cross Entropy Loss value. """ - device = out0.device - batch_size, _ = out0.shape - - # Normalize the output to length 1 + # Normalize the output to length 1. The memory bank stores unit vectors, + # and ntxent takes them, so this happens here and only here. out0 = nn.functional.normalize(out0, dim=1) out1 = nn.functional.normalize(out1, dim=1) @@ -132,65 +133,19 @@ def forward(self, out0: Tensor, out1: Tensor) -> Tensor: # negatives: shape: (embedding_size, memory_bank_size) out1, negatives = self.memory_bank.forward(out1, update=out0.requires_grad) - # Use cosine similarity (dot product) as all vectors are normalized to unit length - # Notation in einsum: n = batch_size, c = embedding_size and k = memory_bank_size. - if negatives is not None: - # Use negatives from memory bank - negatives = negatives.to(device) - - # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity - # of the i-th sample in the batch to its positive pair - sim_pos = torch.einsum("nc,nc->n", out0, out1).unsqueeze(-1) - - # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity - # of the i-th sample to the j-th negative sample - sim_neg = torch.einsum("nc,ck->nk", out0, negatives) - - # Set the labels to maximize sim_pos in relation to sim_neg - logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature - labels = torch.zeros(logits.shape[0], device=device, dtype=torch.long) - - else: - # Use other samples from batch as negatives - # and create diagonal mask that only selects similarities between - # views of the same image - if self.gather_distributed and dist.world_size() > 1: - # Gather hidden representations from other processes - out0_large = torch.cat(dist.gather(out0), 0) - out1_large = torch.cat(dist.gather(out1), 0) - diag_mask = dist.eye_rank(batch_size, device=out0.device) - else: - # Single process - out0_large = out0 - out1_large = out1 - diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool) - - # Calculate similiarities - # Here n = batch_size and m = batch_size * world_size - # The resulting vectors have shape (n, m) - logits_00 = torch.einsum("nc,mc->nm", out0, out0_large) / self.temperature - logits_01 = torch.einsum("nc,mc->nm", out0, out1_large) / self.temperature - logits_10 = torch.einsum("nc,mc->nm", out1, out0_large) / self.temperature - logits_11 = torch.einsum("nc,mc->nm", out1, out1_large) / self.temperature - - # Remove simliarities between same views of the same image - logits_00 = logits_00[~diag_mask].view(batch_size, -1) - logits_11 = logits_11[~diag_mask].view(batch_size, -1) - - # Concatenate logits - # The logits tensor in the end has shape (2*n, 2*m-1) - logits_0100 = torch.cat([logits_01, logits_00], dim=1) - logits_1011 = torch.cat([logits_10, logits_11], dim=1) - logits = torch.cat([logits_0100, logits_1011], dim=0) - - # Create labels - labels = torch.arange(batch_size, device=device, dtype=torch.long) - if self.gather_distributed: - labels = labels + dist.rank() * batch_size - labels = labels.repeat(2) - - # Calculate the cross-entropy loss - loss: Tensor = self.cross_entropy(logits, labels) - - return loss + return ntxent(out0, out1, temperature=self.temperature, negatives=negatives) + + if self.gather_distributed and dist.world_size() > 1: + # GATHER_FOR_NEGATIVES: every sample on every rank is a negative for + # every other, so the gather happens with gradient before the loss. + return ntxent( + out0, + out1, + temperature=self.temperature, + z0_all=torch.cat(dist.gather(out0), 0), + z1_all=torch.cat(dist.gather(out1), 0), + positives_at=dist.rank() * out0.size(0), + ) + + return ntxent(out0, out1, temperature=self.temperature) diff --git a/tests/functional/__init__.py b/tests/functional/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/functional/test_ntxent.py b/tests/functional/test_ntxent.py new file mode 100644 index 000000000..8f3cf904f --- /dev/null +++ b/tests/functional/test_ntxent.py @@ -0,0 +1,69 @@ +"""Golden values for the NT-Xent equation. + +The numbers below are the second statement of the equation, as data rather than +as code. Editing ``lightly/functional/ntxent.py`` moves them, and someone then +has to approve the new number or restore the old one. A pinned loss sees changes +above roughly 1e-5. +""" + +from typing import Tuple + +import pytest +import torch +from torch import Tensor +from torch.nn.functional import normalize + +from lightly.functional import ntxent +from lightly.loss import NTXentLoss + +GOLDEN_LOSS = 5.350811958313 # seed 11, batch 8, dim 16, temperature 0.1 +GOLDEN_LOSS_MEMORY_BANK = 9.919633865356 # the same batch, 8 negatives from a bank +TOLERANCE = 1e-5 + + +def batch(seed: int = 11) -> Tuple[Tensor, Tensor]: + torch.manual_seed(seed) + return ( + normalize(torch.randn(8, 16), dim=1), + normalize(torch.randn(8, 16), dim=1), + ) + + +def test_golden_loss() -> None: + z0, z1 = batch() + assert float(ntxent(z0, z1, temperature=0.1)) == pytest.approx( + GOLDEN_LOSS, abs=TOLERANCE + ) + + +def test_golden_loss_with_a_memory_bank() -> None: + z0, z1 = batch() + torch.manual_seed(11) + negatives = normalize(torch.randn(8, 16), dim=1).T + loss = ntxent(z0, z1, temperature=0.1, negatives=negatives) + assert float(loss) == pytest.approx(GOLDEN_LOSS_MEMORY_BANK, abs=TOLERANCE) + + +def test_the_module_computes_the_same_equation() -> None: + z0, z1 = batch() + module = NTXentLoss(temperature=0.1) + assert float(module(z0, z1)) == pytest.approx(GOLDEN_LOSS, abs=TOLERANCE) + + +def test_temperature_scales_the_logits() -> None: + z0, z1 = batch() + assert float(ntxent(z0, z1, temperature=0.1)) > float( + ntxent(z0, z1, temperature=0.5) + ) + + +def test_positives_at_zero_is_the_single_rank_case() -> None: + z0, z1 = batch() + with_all = ntxent(z0, z1, temperature=0.1, z0_all=z0, z1_all=z1, positives_at=0) + assert torch.equal(with_all, ntxent(z0, z1, temperature=0.1)) + + +def test_the_loss_declares_how_it_behaves_across_ranks() -> None: + from lightly.loss import DistributedKind + + assert NTXentLoss.distributed_kind is DistributedKind.GATHER_FOR_NEGATIVES