From 2b89319b120bf8b8ebba22dd007413a3af070ed8 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 16:46:31 -0300 Subject: [PATCH 1/6] test(losses): reproduce contrastive loss numerical instability On unmodified dev bbebca6f, 11 of 25 focused tests fail with NaN or infinite loss at temperatures 0.01 and 0.001. Covers analytic values, singleton batches, and independent cross-entropy loss/gradient reference. Disposable sandbox removed after validation. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- tests/losses/test_contrastive_loss.py | 52 +++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index 21a9e76417c..17ee7edd04f 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -12,6 +12,7 @@ from __future__ import annotations import unittest +import math import numpy as np import torch @@ -56,6 +57,57 @@ class TestContrastiveLoss(unittest.TestCase): + @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) From d9b792753aa46a6edf706552db7f4359865e7ec2 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 16:47:41 -0300 Subject: [PATCH 2/6] fix(losses): compute contrastive loss in log space Exclude diagonal similarities before logsumexp and subtract the positive logit instead of exponentiating a ratio. All 25 focused tests pass, including 18 new analytic and independent loss/gradient regression cases; Ruff, Black, and isort pass. CPU validation with isolated dependencies and sandbox cleanup. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- monai/losses/contrastive.py | 11 ++++++----- tests/losses/test_contrastive_loss.py | 14 ++++++++++---- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index d8b2c33e62e..b6211a18fb8 100644 --- a/monai/losses/contrastive.py +++ b/monai/losses/contrastive.py @@ -28,6 +28,9 @@ class ContrastiveLoss(_Loss): Adapted from: https://github.com/Sara-Ahmed/SiT/blob/1aacd6adcd39b71efc903d16b4e9095b97dda76f/losses.py#L5 + The loss is evaluated in log space to avoid exponential overflow or underflow + at small positive temperatures. + """ def __init__(self, temperature: float = 0.5, batch_size: int = -1) -> None: @@ -64,8 +67,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: 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) sim_matrix = F.cosine_similarity(repr.unsqueeze(1), repr.unsqueeze(0), dim=2) @@ -73,9 +75,8 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: 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) + logits = (sim_matrix / 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) - positives / temperature_tensor 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 17ee7edd04f..a2da20548c4 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -11,8 +11,8 @@ from __future__ import annotations -import unittest import math +import unittest import numpy as np import torch @@ -57,7 +57,9 @@ class TestContrastiveLoss(unittest.TestCase): - @parameterized.expand([(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)]) + @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) @@ -68,7 +70,9 @@ def test_identical_low_temperature(self, temperature, dtype): 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)]) + @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_() @@ -80,7 +84,9 @@ def test_antipodal_low_temperature(self, temperature, dtype): 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)]) + @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) From 046cc2618d1e5634366b8bd1e98f4c007d1043a3 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:01:41 -0300 Subject: [PATCH 3/6] test(losses): expose precision regressions in contrastive fix Adversarial review reproduced 10 failing cases on d9b79275 among 35 tests: float16/bfloat16 output promotion and nonzero gradient checks, float16 batch reduction overflow, and cancellation for tiny temperatures. Both independent reviewers rejected the initial fix. Red validation completed in a removed disposable sandbox. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- tests/losses/test_contrastive_loss.py | 42 +++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index a2da20548c4..d3abe3e9b31 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -57,6 +57,48 @@ 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)]) + 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( [(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)] ) From 70edaa823c3c9cfae422ebad94b18c42f4f0afea Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:04:44 -0300 Subject: [PATCH 4/6] test(losses): cover autocast and float64 temperature precision CPU autocast and numerical gradcheck pass on the intermediate correction. One additional regression reproduces float64 temperature 1e-50 becoming zero through default float32 scalar construction; 1 of 39 focused cases fails until the scalar dtype is preserved. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- tests/losses/test_contrastive_loss.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index d3abe3e9b31..6ea9c994836 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -80,7 +80,9 @@ def test_reduced_precision_loss_and_gradients(self, temperature, dtype): 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)]) + @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) @@ -99,6 +101,24 @@ def test_float16_batch_reduction(self): 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)] ) From 87f243cdf5466c1edee5c251e0061b39da63159e Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:06:01 -0300 Subject: [PATCH 5/6] fix(losses): preserve contrastive precision before logsumexp Adversarial review found a new float16 sum overflow from lost float32 promotion and residual cancellation after logsumexp. Compute reduced-precision inputs in float32, shift similarities by the positive before temperature scaling, and preserve the temperature dtype for float64. All 39 focused tests, finite-difference gradcheck, CPU autocast, and independent gradient references pass. Changed loss module line coverage is 93%; Ruff, Black, and isort pass. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- monai/losses/contrastive.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index b6211a18fb8..464c9fb47f3 100644 --- a/monai/losses/contrastive.py +++ b/monai/losses/contrastive.py @@ -28,8 +28,8 @@ class ContrastiveLoss(_Loss): Adapted from: https://github.com/Sara-Ahmed/SiT/blob/1aacd6adcd39b71efc903d16b4e9095b97dda76f/losses.py#L5 - The loss is evaluated in log space to avoid exponential overflow or underflow - at small positive temperatures. + The loss is evaluated using positive-shifted logits in log space. Float16 and + bfloat16 inputs use float32 arithmetic to preserve precision and accumulation range. """ @@ -37,10 +37,7 @@ 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__() @@ -54,6 +51,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( @@ -64,19 +67,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] 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) - logits = (sim_matrix / temperature_tensor).masked_fill(self_mask, -torch.inf) + # 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.logsumexp(logits, dim=1) - positives / temperature_tensor + loss_partial = torch.logsumexp(logits, dim=1) return torch.sum(loss_partial) / (2 * batch_size) From 49542ec865f11a7b54a44e16fc367633d1f77041 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:25:12 -0300 Subject: [PATCH 6/6] docs(losses): explain extreme-temperature gradient precision Isolated cosine backward roundoff and compared seven evaluation strategies across 1363 diagnostic configurations. Passing float64 embeddings before the loss reduces the measured gradient residual; casting the computed loss does not. Normalization prototypes leave residuals or alter near-epsilon backward semantics, so the algorithm stays unchanged. All 39 focused tests, Ruff, Black, and isort pass; module line coverage remains 93 percent. CPU evidence only; no universal gradient-accuracy claim or automatic dtype policy. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- monai/losses/contrastive.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index 464c9fb47f3..4878cc28afd 100644 --- a/monai/losses/contrastive.py +++ b/monai/losses/contrastive.py @@ -31,6 +31,11 @@ class ContrastiveLoss(_Loss): 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: