fix(losses): stabilize ContrastiveLoss at small temperatures

monai/losses/contrastive.py

@@ -28,16 +28,21 @@ class ContrastiveLoss(_Loss):
     Adapted from:
         https://github.com/Sara-Ahmed/SiT/blob/1aacd6adcd39b71efc903d16b4e9095b97dda76f/losses.py#L5
 
+    The loss is evaluated using positive-shifted logits in log space. Float16 and
+    bfloat16 inputs use float32 arithmetic to preserve precision and accumulation range.
+
+    Very small temperatures can amplify roundoff in cosine-similarity gradients.
+    For higher precision, pass both embeddings as float64 before evaluating the loss,
+    for example ``loss(input.double(), target.double())``. Casting the computed loss
+    to float64 does not improve the precision of its backward computation.
+
     """
 
     def __init__(self, temperature: float = 0.5, batch_size: int = -1) -> None:
         """
         Args:
             temperature: Can be scaled between 0 and 1 for learning from negative samples, ideally set to 0.5.
-
-        Raises:
-            ValueError: When an input of dimension length > 2 is passed
-            ValueError: When input and target are of different shapes
+            batch_size: Legacy parameter. The batch size is estimated dynamically in ``forward``.
 
         """
         super().__init__()
@@ -51,6 +56,12 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Args:
             input: the shape should be B[F].
             target: the shape should be B[F].
+
+        Returns:
+            A scalar loss tensor, using float32 for float16 and bfloat16 inputs.
+
+        Raises:
+            ValueError: If either input has more than two dimensions or their shapes differ.
         """
         if len(target.shape) > 2 or len(input.shape) > 2:
             raise ValueError(
@@ -61,21 +72,22 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if target.shape != input.shape:
             raise ValueError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})")
 
-        temperature_tensor = torch.as_tensor(self.temperature).to(input.device)
         batch_size = input.shape[0]
 
-        negatives_mask = ~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool)
-        negatives_mask = torch.clone(negatives_mask.type(torch.float)).to(input.device)
+        self_mask = torch.eye(batch_size * 2, dtype=torch.bool, device=input.device)
 
         repr = torch.cat([input, target], dim=0)
+        if repr.dtype in (torch.float16, torch.bfloat16):
+            repr = repr.float()
+        temperature_tensor = torch.as_tensor(self.temperature, dtype=repr.dtype, device=input.device)
         sim_matrix = F.cosine_similarity(repr.unsqueeze(1), repr.unsqueeze(0), dim=2)
         sim_ij = torch.diag(sim_matrix, batch_size)
         sim_ji = torch.diag(sim_matrix, -batch_size)
 
         positives = torch.cat([sim_ij, sim_ji], dim=0)
-        nominator = torch.exp(positives / temperature_tensor)
-        denominator = negatives_mask * torch.exp(sim_matrix / temperature_tensor)
+        # Shift before scaling so the positive logit is zero, avoiding cancellation after logsumexp.
+        logits = ((sim_matrix - positives.unsqueeze(1)) / temperature_tensor).masked_fill(self_mask, -torch.inf)
 
-        loss_partial = -torch.log(nominator / torch.sum(denominator, dim=1))
+        loss_partial = torch.logsumexp(logits, dim=1)
 
         return torch.sum(loss_partial) / (2 * batch_size)

tests/losses/test_contrastive_loss.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import math
 import unittest
 
 import numpy as np
@@ -56,6 +57,125 @@
 
 class TestContrastiveLoss(unittest.TestCase):
 
+    @parameterized.expand(
+        [(temperature, dtype) for temperature in (0.5, 0.01, 0.001) for dtype in (torch.float16, torch.bfloat16)]
+    )
+    def test_reduced_precision_loss_and_gradients(self, temperature, dtype):
+        # All nonself similarities are 0.5, but the embedding gradients are nonzero.
+        vectors = torch.cat([torch.ones(4, 1), torch.eye(4)], dim=1).to(dtype)
+        inputs = vectors[:2].clone().requires_grad_()
+        targets = vectors[2:].clone().requires_grad_()
+        ref_inputs = inputs.detach().double().requires_grad_()
+        ref_targets = targets.detach().double().requires_grad_()
+        representations = torch.nn.functional.normalize(torch.cat([ref_inputs, ref_targets]), dim=1)
+        logits = (representations @ representations.T / temperature).masked_fill(
+            torch.eye(4, dtype=torch.bool), -torch.inf
+        )
+        expected = torch.nn.functional.cross_entropy(logits, torch.tensor([2, 3, 0, 1]))
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        self.assertEqual(result.dtype, torch.float32)
+        torch.testing.assert_close(result.double(), expected, atol=1e-5, rtol=1e-5)
+        result.backward()
+        expected.backward()
+        torch.testing.assert_close(inputs.grad, ref_inputs.grad.to(dtype), atol=1e-4, rtol=1e-2)
+        torch.testing.assert_close(targets.grad, ref_targets.grad.to(dtype), atol=1e-4, rtol=1e-2)
+
+    @parameterized.expand(
+        [(1e-7, torch.float32), (1e-8, torch.float32), (1e-12, torch.float64), (1e-50, torch.float64)]
+    )
+    def test_small_temperature_cancellation(self, temperature, dtype):
+        inputs = torch.ones((2, 4), dtype=dtype, requires_grad=True)
+        targets = torch.ones_like(inputs, requires_grad=True)
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        torch.testing.assert_close(result, result.new_tensor(math.log(3)), atol=1e-7, rtol=1e-7)
+        result.backward()
+        self.assertTrue(torch.isfinite(inputs.grad).all())
+        self.assertTrue(torch.isfinite(targets.grad).all())
+
+    def test_float16_batch_reduction(self):
+        # The mean is finite although a float16 sum of the per-anchor losses overflows.
+        batch_size, temperature = 1280, 0.1
+        inputs = torch.tensor([[1.0, 0.0]], dtype=torch.float16).repeat(batch_size, 1)
+        result = ContrastiveLoss(temperature=temperature)(inputs, -inputs)
+        expected = 2 / temperature + math.log(batch_size - 1 + batch_size * math.exp(-2 / temperature))
+        self.assertEqual(result.dtype, torch.float32)
+        torch.testing.assert_close(result, result.new_tensor(expected))
+
+    @parameterized.expand([(torch.float16,), (torch.bfloat16,)])
+    def test_cpu_autocast(self, dtype):
+        vectors = torch.cat([torch.ones(4, 1), torch.eye(4)], dim=1).to(dtype)
+        inputs = vectors[:2].clone().requires_grad_()
+        targets = vectors[2:].clone().requires_grad_()
+        with torch.autocast(device_type="cpu", dtype=dtype):
+            result = ContrastiveLoss(temperature=0.001)(inputs, targets)
+        self.assertEqual(result.dtype, torch.float32)
+        torch.testing.assert_close(result, result.new_tensor(math.log(3)))
+        result.backward()
+        self.assertTrue(torch.isfinite(inputs.grad).all())
+        self.assertTrue(torch.isfinite(targets.grad).all())
+
+    def test_gradcheck(self):
+        inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], dtype=torch.float64, requires_grad=True)
+        targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=torch.float64, requires_grad=True)
+        self.assertTrue(torch.autograd.gradcheck(ContrastiveLoss(temperature=0.1), (inputs, targets)))
+
+    @parameterized.expand(
+        [(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)]
+    )
+    def test_identical_low_temperature(self, temperature, dtype):
+        inputs = torch.ones((2, 4), dtype=dtype, requires_grad=True)
+        targets = torch.ones((2, 4), dtype=dtype, requires_grad=True)
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        # Every anchor has three equally similar candidates, including its positive.
+        torch.testing.assert_close(result, result.new_tensor(math.log(3)), atol=1e-4, rtol=1e-4)
+        result.backward()
+        self.assertTrue(torch.isfinite(inputs.grad).all())
+        self.assertTrue(torch.isfinite(targets.grad).all())
+
+    @parameterized.expand(
+        [(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)]
+    )
+    def test_antipodal_low_temperature(self, temperature, dtype):
+        inputs = torch.eye(2, dtype=dtype, requires_grad=True)
+        targets = (-torch.eye(2, dtype=dtype)).requires_grad_()
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        # Positive similarity is -1, and both other candidates have similarity 0.
+        expected = 1 / temperature + math.log(2 + math.exp(-1 / temperature))
+        torch.testing.assert_close(result, result.new_tensor(expected))
+        result.backward()
+        self.assertTrue(torch.isfinite(inputs.grad).all())
+        self.assertTrue(torch.isfinite(targets.grad).all())
+
+    @parameterized.expand(
+        [(temperature, dtype) for temperature in (0.5, 0.05, 0.01) for dtype in (torch.float32, torch.float64)]
+    )
+    def test_reference_loss_and_gradients(self, temperature, dtype):
+        inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], dtype=dtype, requires_grad=True)
+        targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=dtype, requires_grad=True)
+        reference_inputs = inputs.detach().double().requires_grad_()
+        reference_targets = targets.detach().double().requires_grad_()
+        representations = torch.nn.functional.normalize(torch.cat([reference_inputs, reference_targets]), dim=1)
+        logits = representations @ representations.T / temperature
+        logits = logits.masked_fill(torch.eye(4, dtype=torch.bool), -torch.inf)
+        # The positive for each anchor is the corresponding sample in the other view.
+        expected = torch.nn.functional.cross_entropy(logits, torch.tensor([2, 3, 0, 1]))
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        torch.testing.assert_close(result.double(), expected, atol=1e-4, rtol=1e-4)
+        result.backward()
+        expected.backward()
+        torch.testing.assert_close(inputs.grad.double(), reference_inputs.grad, atol=1e-4, rtol=1e-4)
+        torch.testing.assert_close(targets.grad.double(), reference_targets.grad, atol=1e-4, rtol=1e-4)
+
+    @parameterized.expand([(temperature, sign) for temperature in (0.01, 0.001) for sign in (1, -1)])
+    def test_single_pair_low_temperature(self, temperature, sign):
+        inputs = torch.tensor([[1.0, 0.0]], requires_grad=True)
+        targets = (sign * inputs.detach()).requires_grad_()
+        result = ContrastiveLoss(temperature=temperature)(inputs, targets)
+        torch.testing.assert_close(result, torch.zeros_like(result))
+        result.backward()
+        torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs))
+        torch.testing.assert_close(targets.grad, torch.zeros_like(targets))
+
     @parameterized.expand(TEST_CASES)
     def test_result(self, input_param, input_data, expected_val):
         contrastiveloss = ContrastiveLoss(**input_param)