From 2b89319b120bf8b8ebba22dd007413a3af070ed8 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 16:46:31 -0300 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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/8] 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/8] 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: From 3a3d40934e4d9c548b5b2e0a616f64609abe57ae Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:32:25 -0300 Subject: [PATCH 7/8] test(losses): specify contrastive computation dtype controls On parent PR #9208 head49542ec8, all 49 new API cases fail because compute_dtype is unavailable; all 39 inherited cases pass. Cases specify float64 extreme-temperature zero gradients, equivalence to input casting, input/output dtypes, autocast, invalid options, singleton loss and legacy arguments. Red validation completed in a cleaned sandbox. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- tests/losses/test_contrastive_loss.py | 80 +++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index 6ea9c994836..11e3efb7844 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -202,5 +202,85 @@ 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([(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() From 116bc2a409604fd05953634a8314ac2bb3169ee3 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 20 Oct 2026 17:38:07 -0300 Subject: [PATCH 8/8] feat(losses): add explicit contrastive computation dtype Expose keyword-only compute_dtype=None/float32/float64. Cast connected embeddings before cosine and temperature construction, preserving the parent default path. All 91 CPU tests pass, including 52 new cases and mixed-dtype backward controls; coverage97 percent and Ruff/Black/isort pass. In 24 identical-vector controls at temperature1e-8, explicit float64 reduces maximum gradient residual from2.0 to3.725290298461914e-9. This is an opt-in precision policy, not a universal accuracy guarantee or automatic dtype change. Depends on upstream PR#9208. Assisted-by: OpenAI Codex Signed-off-by: joaomarcos --- monai/losses/contrastive.py | 31 ++++++++++++++++++++------ tests/losses/test_contrastive_loss.py | 32 +++++++++++++++++++++++---- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index 4878cc28afd..91225fc61af 100644 --- a/monai/losses/contrastive.py +++ b/monai/losses/contrastive.py @@ -28,25 +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. Float16 and - bfloat16 inputs use float32 arithmetic to preserve precision and accumulation range. + 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())``. Casting the computed loss - to float64 does not improve the precision of its backward computation. + 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: 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") @@ -58,7 +71,8 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: target: the shape should be B[F]. Returns: - A scalar loss tensor, using float32 for float16 and bfloat16 inputs. + 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. @@ -77,7 +91,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: 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): + 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) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index 11e3efb7844..325d9cd3111 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -219,8 +219,12 @@ def test_identical_extreme_temperature_gradients(self, features, seed): 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)] + [ + (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) @@ -242,14 +246,34 @@ def test_equivalent_to_casting_inputs(self, dtype, temperature): 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)] + [ + (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"):