Fix LocalNormalizedCrossCorrelationLoss exceeding its expected range in low-variance regions

monai/losses/image_dissimilarity.py

@@ -72,7 +72,8 @@ class LocalNormalizedCrossCorrelationLoss(_Loss):
             - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
             - ``"sum"``: the output will be summed.
         smooth_nr: a small constant added to the numerator to avoid nan.
-        smooth_dr: a small constant added to the denominator to avoid nan.
+        smooth_dr: strictly positive minimum informative local variance. Windows where either
+            input has variance less than or equal to this threshold contribute zero to the loss.
 
     Returns:
         torch.Tensor: The computed loss value. The output range is approximately [-1, 0], where:
@@ -84,6 +85,7 @@ class LocalNormalizedCrossCorrelationLoss(_Loss):
         The implementation computes the squared normalized cross-correlation coefficient
         and then negates it, transforming the correlation maximization problem into a
         loss minimization problem suitable for standard PyTorch optimizers.
+        Squared correlation is capped at 1 to guard against numerical overshoots.
 
         Interpretation:
             - Loss ≈ -1: Perfect correlation between images
@@ -118,6 +120,8 @@ def __init__(
 
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
+        if not self.smooth_dr > 0:
+            raise ValueError(f"smooth_dr must be strictly positive, got {self.smooth_dr}")
 
     def get_kernel_vol(self) -> torch.Tensor:
         assert self.kernel is not None
@@ -165,13 +169,14 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         #     = sum[t*p] - sum[t] * mean[p] = cross
         # the following is actually squared ncc
         cross = tp_sum - p_avg * t_sum
-        t_var = torch.max(
-            t2_sum - t_avg * t_sum, torch.as_tensor(self.smooth_dr, dtype=t2_sum.dtype, device=t2_sum.device)
-        )
-        p_var = torch.max(
-            p2_sum - p_avg * p_sum, torch.as_tensor(self.smooth_dr, dtype=p2_sum.dtype, device=p2_sum.device)
-        )
+        t_var = t2_sum - t_avg * t_sum
+        p_var = p2_sum - p_avg * p_sum
+        valid = (t_var > self.smooth_dr) & (p_var > self.smooth_dr)
+        # Use a safe denominator before masking to keep invalid windows' gradients finite.
+        t_var = torch.where(valid, t_var, torch.ones_like(t_var))
+        p_var = torch.where(valid, p_var, torch.ones_like(p_var))
         ncc: torch.Tensor = (cross * cross + self.smooth_nr) / (t_var * p_var)
+        ncc = torch.where(valid, ncc.clamp(max=1.0), torch.zeros_like(ncc))
 
         if self.reduction == LossReduction.SUM.value:
             return torch.sum(ncc).neg()  # sum over the batch, channel and spatial ndims

tests/losses/image_dissimilarity/test_local_normalized_cross_correlation_loss.py

@@ -137,6 +137,48 @@
 
 class TestLocalNormalizedCrossCorrelationLoss(unittest.TestCase):
 
+    @parameterized.expand([1e-5, 1e-6])
+    def test_low_variance_range(self, smooth_dr):
+        # gh-9178: cancellation in float32 moments can produce squared NCC > 1.
+        generator = torch.Generator().manual_seed(9178)
+        pred = (1.0 + 1e-4 * torch.randn(1, 1, 32, 32, generator=generator)).requires_grad_()
+        target = 1.0 + 1e-4 * torch.randn(1, 1, 32, 32, generator=generator)
+        loss = LocalNormalizedCrossCorrelationLoss(spatial_dims=2, kernel_size=9, smooth_dr=smooth_dr, reduction="none")
+        result = loss(pred, target)
+        self.assertTrue(torch.isfinite(result).all())
+        self.assertGreaterEqual(result.min().item(), -1.0)
+        self.assertLessEqual(result.max().item(), 0.0)
+        result.mean().backward()
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+    @parameterized.expand([(False,), (True,)])
+    def test_uninformative_windows(self, swap):
+        # Only one input needs to have insufficient variance to mask the window.
+        pred = torch.tensor([[[0.0, 1e-4, 2e-4]]], dtype=torch.float64)
+        target = torch.tensor([[[0.0, 1.0, 2.0]]], dtype=torch.float64)
+        if swap:
+            pred, target = target, pred
+        pred.requires_grad_()
+        loss = LocalNormalizedCrossCorrelationLoss(spatial_dims=1, reduction="none", smooth_nr=1e-8)
+        result = loss(pred, target)
+        torch.testing.assert_close(result, torch.zeros_like(result))
+        result.sum().backward()
+        torch.testing.assert_close(pred.grad, torch.zeros_like(pred))
+
+    @parameterized.expand([0.0, -1e-5, float("nan")])
+    def test_nonpositive_smooth_dr(self, smooth_dr):
+        with self.assertRaisesRegex(ValueError, "smooth_dr must be strictly positive"):
+            LocalNormalizedCrossCorrelationLoss(smooth_dr=smooth_dr)
+
+    def test_variance_at_threshold(self):
+        pred = torch.tensor([[[0.0, 1.0, 2.0]]], requires_grad=True)
+        loss = LocalNormalizedCrossCorrelationLoss(spatial_dims=1, smooth_dr=2.0, reduction="none")
+        # The center window has variance exactly equal to the threshold.
+        result = loss(pred, pred)
+        self.assertEqual(result[0, 0, 1].item(), 0.0)
+        result[0, 0, 1].backward()
+        torch.testing.assert_close(pred.grad, torch.zeros_like(pred))
+
     @parameterized.expand(TEST_CASES)
     def test_shape(self, input_param, input_data, expected_val):
         result = LocalNormalizedCrossCorrelationLoss(**input_param).forward(**input_data)