diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index d8b2c33e62..4878cc28af 100644 --- a/monai/losses/contrastive.py +++ b/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) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index 21a9e76417..6ea9c99483 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/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)