ContrastiveLoss returns NaN loss and gradients at small positive temperatures #9207

Describe the bug

ContrastiveLoss returns a non-finite loss and non-finite gradients for finite embeddings at small positive temperatures. For example, temperature=0.01 is within the documented range but produces NaN in float32 even when all four embeddings are identical and the exact loss is ln(3).

This prevents an optimizer from receiving usable gradients for an otherwise valid loss evaluation.

To Reproduce

import math
import torch
from monai.losses import ContrastiveLoss

x = torch.ones((2, 4), requires_grad=True)
loss = ContrastiveLoss(temperature=0.01)(x, x)
print(loss.item())  # nan; expected math.log(3) == 1.0986122886681098
loss.backward()
print(torch.isfinite(x.grad).all().item())  # False

Reproduced against unmodified dev at bbebca6f17cabf98ea5fb8ec5776eeacd10aa60d.

Expected behavior

Each of the four anchors has three equally similar candidates after excluding itself; the probability of its positive is 1/3. The loss must therefore be -log(1/3) = ln(3), with finite gradients. A singleton batch has one candidate per anchor and should have zero loss, including at small positive temperatures.

Environment

• Windows, Python 3.14.1
• PyTorch 2.12.1+cpu; NumPy 2.3.5
• CPU float32 and float64 regressions tested, including temperatures 0.01 and 0.001
• GPU and distributed execution were not tested

Root cause and proposed correction

forward() exponentiates the cosine similarities divided by temperature before masking out the diagonal. A float32 self-similarity of 1 at temperature 0.01 evaluates exp(100), which overflows. Multiplying this by the zero diagonal mask creates 0 * inf = NaN; the positive-to-denominator ratio can also overflow or underflow independently.

Promote float16/bfloat16 embeddings to float32 and evaluate logsumexp((similarity - positive_similarity) / temperature) with self comparisons excluded. Construct the temperature scalar in the computation dtype. This is algebraically the same contrastive objective and avoids the direct exponentials and probability ratio. No epsilon or loss clamping is needed.

Validation and duplicate check

The initial reproducer added 18 parameterized cases alongside seven existing tests; adversarial review expanded the final suite to 39 cases. On the unmodified implementation, 11 of the 25 cases fail with incorrect non-finite results. The final positive-shifted correction passes 39 cases, including analytic values, singleton batches, low-precision accumulation, CPU autocast, gradcheck, and independent double-precision cross-entropy loss/gradient comparisons. Ruff, Black, and isort pass on both changed files. Tests use disposable sandboxes with cleanup and one numerical-library thread.

Searched open and closed issues/PRs for ContrastiveLoss, contrastive nan, contrastive temperature, contrastive overflow, contrastive underflow, contrastive stability, NT-Xent, and logsumexp. The existing contrastive-loss items concern introduction, batch sizing, warnings, reduction, and redundant normalization; no matching numerical-stability fix was found.

AI assistance disclosure: investigation, regression tests, fix, and this report were prepared with OpenAI Codex. The infographic was generated with the native OpenAI image tool. The submitting author remains responsible for the contribution under MONAI's tool-assisted contribution policy.

The root-cause correction, regression tests, inline documentation, and explanatory infographic are published in PR #9208. Final instrumented validation at 49542ec passed all 39 tests and measured 93% line coverage of the changed loss module. Adversarial review additionally corrected lost float32 promotion for FP16/BF16, cancellation after logsumexp, and temperature-scalar underflow in float64. Extremely small temperatures can still amplify gradient rounding; universal gradient accuracy is not claimed. GPU and distributed testing are not claimed.

Extreme-temperature precision follow-up: CPU diagnostics reproduced cosine-similarity backward roundoff without MONAI and evaluated 1,363 configurations with seven explicit control checks. At temperature 1e-8, 24 identical-vector cases had maximum gradient residual 2.0 in float32 versus approximately 3.7e-9 when the same class received float64 embeddings. Convert both embeddings before loss evaluation (loss_fn(x.double(), y.double())); casting the already-computed loss does not improve backward precision. Normalization replacements left residuals or changed near-epsilon gradient semantics, so the algorithm and default dtype policy remain unchanged. This is a measured mitigation, not a universal accuracy guarantee. Guidance is included in the docstring and PR #9208; the 39-case suite and style checks pass. GPU and performance were not tested.