feat(losses): add ContrastiveLoss computation dtype

monai/losses/contrastive.py

@@ -28,20 +28,38 @@ 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. By default, 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())``, or set ``compute_dtype=torch.float64``.
+    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:
+    def __init__(
+        self, temperature: float = 0.5, batch_size: int = -1, *, compute_dtype: torch.dtype | None = None
+    ) -> None:
         """
         Args:
             temperature: Can be scaled between 0 and 1 for learning from negative samples, ideally set to 0.5.
+            batch_size: Legacy parameter. The batch size is estimated dynamically in ``forward``.
+            compute_dtype: Floating-point dtype for similarity and loss computations, either ``torch.float32``
+                or ``torch.float64``. Defaults to None, preserving float64 inputs and promoting float16/bfloat16
+                inputs to float32. Selecting float64 can reduce backward roundoff at very small temperatures
+                on devices supporting float64, with additional computation and memory costs. Gradients are
+                still stored in the original input dtype; existing input quantization cannot be recovered.
 
         Raises:
-            ValueError: When an input of dimension length > 2 is passed
-            ValueError: When input and target are of different shapes
+            ValueError: If ``compute_dtype`` is neither None nor a supported computation dtype.
 
         """
         super().__init__()
+        if compute_dtype not in (None, torch.float32, torch.float64):
+            raise ValueError("compute_dtype must be None, torch.float32, or torch.float64.")
         self.temperature = temperature
+        self.compute_dtype = compute_dtype
 
         if batch_size != -1:
             warn("batch_size is no longer required to be set. It will be estimated dynamically in the forward call")
@@ -51,6 +69,13 @@ 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 in ``compute_dtype`` when specified. Otherwise, float16/bfloat16
+            inputs use float32, and float32/float64 inputs retain their dtype.
+
+        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 +86,25 @@ 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)
+        compute_dtype = self.compute_dtype
+        if compute_dtype is not None:
+            repr = repr.to(dtype=compute_dtype)
+        elif 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)
@@ -82,5 +202,109 @@ def check_warning_rasied(self):
             ContrastiveLoss(temperature=0.5, batch_size=1)
 
 
+class TestContrastiveComputeDtype(unittest.TestCase):
+
+    @parameterized.expand([(features, seed) for features in (3, 5, 17) for seed in range(8)])
+    def test_identical_extreme_temperature_gradients(self, features, seed):
+        generator = torch.Generator().manual_seed(seed)
+        vector = torch.randn((1, features), generator=generator)
+        inputs = vector.repeat(2, 1).requires_grad_()
+        targets = vector.repeat(2, 1).requires_grad_()
+        result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(inputs, targets)
+        self.assertEqual(result.dtype, torch.float64)
+        torch.testing.assert_close(result, result.new_tensor(math.log(3)))
+        result.backward()
+        # Each cosine is stationary at identical nonzero vectors; its mathematical gradient is zero.
+        torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs), atol=1e-7, rtol=0)
+        torch.testing.assert_close(targets.grad, torch.zeros_like(targets), atol=1e-7, rtol=0)
+
+    @parameterized.expand(
+        [
+            (dtype, temperature)
+            for dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64)
+            for temperature in (0.5, 0.001)
+        ]
+        + [(torch.float32, 1e-8), (torch.float64, 1e-8)]
+    )
+    def test_equivalent_to_casting_inputs(self, dtype, temperature):
+        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_()
+        original_inputs, original_targets = inputs.detach().clone(), targets.detach().clone()
+        ref_inputs = inputs.detach().clone().requires_grad_()
+        ref_targets = targets.detach().clone().requires_grad_()
+        result = ContrastiveLoss(temperature=temperature, compute_dtype=torch.float64)(inputs, targets)
+        expected = ContrastiveLoss(temperature=temperature)(ref_inputs.double(), ref_targets.double())
+        torch.testing.assert_close(result, expected, atol=0, rtol=0)
+        result.backward()
+        expected.backward()
+        torch.testing.assert_close(inputs.grad, ref_inputs.grad, atol=0, rtol=0)
+        torch.testing.assert_close(targets.grad, ref_targets.grad, atol=0, rtol=0)
+        torch.testing.assert_close(inputs, original_inputs, atol=0, rtol=0)
+        torch.testing.assert_close(targets, original_targets, atol=0, rtol=0)
+        self.assertTrue(torch.isfinite(inputs.grad).all())
+        self.assertTrue(torch.isfinite(targets.grad).all())
+
+    @parameterized.expand(
+        [
+            (input_dtype, compute_dtype)
+            for input_dtype in (torch.float32, torch.float64)
+            for compute_dtype in (None, torch.float32, torch.float64)
+        ]
+    )
+    def test_output_dtype(self, input_dtype, compute_dtype):
+        inputs = torch.ones((2, 3), dtype=input_dtype)
+        result = ContrastiveLoss(compute_dtype=compute_dtype)(inputs, inputs)
+        self.assertEqual(result.dtype, input_dtype if compute_dtype is None else compute_dtype)
+
+    @parameterized.expand([(None,), (torch.float32,), (torch.float64,)])
+    def test_mixed_input_dtype_backward(self, compute_dtype):
+        inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], requires_grad=True)
+        targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=torch.float64, requires_grad=True)
+        ref_inputs = inputs.detach().clone().requires_grad_()
+        ref_targets = targets.detach().clone().requires_grad_()
+        result = ContrastiveLoss(temperature=0.001, compute_dtype=compute_dtype)(inputs, targets)
+        expected_dtype = torch.float64 if compute_dtype is None else compute_dtype
+        expected = ContrastiveLoss(temperature=0.001)(ref_inputs.to(expected_dtype), ref_targets.to(expected_dtype))
+        torch.testing.assert_close(result, expected, atol=0, rtol=0)
+        result.backward()
+        expected.backward()
+        torch.testing.assert_close(inputs.grad, ref_inputs.grad, atol=0, rtol=0)
+        torch.testing.assert_close(targets.grad, ref_targets.grad, atol=0, rtol=0)
+        self.assertEqual(inputs.grad.dtype, torch.float32)
+        self.assertEqual(targets.grad.dtype, torch.float64)
+
+    @parameterized.expand([(torch.float16,), (torch.bfloat16,), (torch.int64,), (torch.bool,), ("float64",)])
+    def test_invalid_compute_dtype(self, compute_dtype):
+        with self.assertRaisesRegex(ValueError, "compute_dtype"):
+            ContrastiveLoss(compute_dtype=compute_dtype)
+
+    @parameterized.expand([(torch.float16,), (torch.bfloat16,)])
+    def test_cpu_autocast_respects_compute_dtype(self, autocast_dtype):
+        inputs = torch.ones((2, 3), requires_grad=True)
+        targets = torch.ones_like(inputs, requires_grad=True)
+        with torch.autocast(device_type="cpu", dtype=autocast_dtype):
+            result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(inputs, targets)
+        self.assertEqual(result.dtype, torch.float64)
+        result.backward()
+        torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs), atol=1e-7, rtol=0)
+        torch.testing.assert_close(targets.grad, torch.zeros_like(targets), atol=1e-7, rtol=0)
+
+    def test_singleton_high_precision(self):
+        inputs = torch.tensor([[0.3, -0.7, 0.2]], requires_grad=True)
+        targets = torch.tensor([[-0.4, 0.1, 0.8]], requires_grad=True)
+        result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(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))
+
+    def test_legacy_arguments_preserved(self):
+        with self.assertWarnsRegex(Warning, "batch_size"):
+            loss = ContrastiveLoss(0.5, 2, compute_dtype=torch.float64)
+        inputs = torch.ones((2, 3))
+        self.assertEqual(loss(inputs, inputs).dtype, torch.float64)
+
+
 if __name__ == "__main__":
     unittest.main()