ContrastiveLoss: expose computation dtype for extreme-temperature backward precision #9209

Is your feature request related to a problem? Please describe.

After the log-space correction proposed in #9208, ContrastiveLoss can return a finite, correct loss while its float32 backward has a material rounding residual at very small temperatures. For identical nonzero embeddings, the exact gradient is zero. At temperature 1e-8, the tested CPU implementation produces a residual of 1.0 in the example below.

This issue tracks an explicit computation-precision option for that residual backward limitation. Issue #9207 and PR #9208 track the direct-exponential overflow and stable forward evaluation; they currently document manual float64 input conversion but do not expose an internal computation dtype.

Reproduction against the #9208 implementation at 49542ec865f11a7b54a44e16fc367633d1f77041:

import torch
from monai.losses import ContrastiveLoss

x = torch.ones((2, 3), requires_grad=True)
y = torch.ones_like(x, requires_grad=True)
loss = ContrastiveLoss(temperature=1e-8)(x, y)
loss.backward()
print(loss.item(), x.grad.abs().max().item())
# Tested CPU result: 1.0986123085021973, 1.0
# Mathematical expectation: log(3), 0.0

The temperature scale amplifies cosine-backward roundoff. The effect is also reproducible using a weighted PyTorch cosine similarity without importing MONAI. Correcting the loss value alone does not establish backward accuracy.

Describe the solution you'd like

Add a keyword-only compute_dtype option to ContrastiveLoss, allowing torch.float32, torch.float64, or None. None retains the behavior of #9208. An explicit dtype converts the connected embeddings before cosine similarity and temperature construction:

loss_fn = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)
loss = loss_fn(x, y)  # x and y can remain float32
loss.backward()

The option centralizes a loss-specific numerical policy while retaining autograd connectivity to the original inputs. The output uses the selected dtype; input tensors and their gradient-storage dtypes are preserved. Float64 is an opt-in mitigation with additional memory/computation costs, rather than a promise of exact gradients at every temperature.

Describe alternatives you've considered

• Manually passing x.double() and y.double() is the existing equivalent workaround. Exposing this policy in the loss constructor avoids repeating input casts in training code.
• Casting only the computed loss to float64 leaves the preceding arithmetic unchanged and did not reduce the measured error.
• Replacing the reduction with cross-entropy did not remove the residual in the tested cases.
• Normalizing once per embedding improved some cases but did not eliminate all residuals; substituting F.normalize changed backward semantics near the cosine epsilon.
• Automatic float64 promotion based on a temperature threshold introduces an unmeasured performance policy; explicit selection retains the default path.

Additional context

With the proposed option, 24 fixed-seed identical-vector cases at temperature 1e-8 had a maximum gradient residual of 3.725290298461914e-9, versus 2.0 under default computation. Each case repeats one nonzero vector four times and splits it into two inputs of batch size 2; feature sizes are 3, 5 and 17.

Validation: 91 CPU tests passed (39 inherited from #9208, 52 new), 97% line coverage of the loss module, and Ruff/Black/isort passed. The initial 49 API tests failed on the parent because the option was absent; three mixed-input-dtype controls were subsequently added. Tests cover analytical zero gradients, exact equivalence to manual casting for values and nonzero gradients, input preservation, mixed dtypes, CPU autocast, unsupported options, singleton batches and legacy positional arguments.

Environment: Windows, Python 3.14.1, PyTorch 2.12.1+cpu, NumPy 2.3.5. GPU, distributed execution, training convergence, and performance were not tested. Float64 does not recover information already quantized in embeddings, increase the range of gradients stored in lower-precision inputs, or remove all floating-point limits.

Draft PR #9210 implements this option, retains #9208 as its prerequisite, and isolates this precision-control increment for review. Open/closed issues and PRs were searched for ContrastiveLoss dtype, float64, precision, gradient and roundoff; the existing related discussion is cross-linked above.

AI assistance disclosure: analysis, implementation, tests and documentation were prepared with OpenAI Codex, including independent adversarial reviews. The dragon-themed infographic was generated with native OpenAI image generation. The submitting author remains responsible under MONAI's tool-assisted contribution policy.