[JoaoMarcos44]

Fixes #9209.

Dependency and review scope

This PR depends on PR #9208, which supplies log-space forward stabilization. The new capability here is explicit computation precision for the remaining backward-roundoff limitation. The default path is unchanged.

Because #9208 is still open, this branch includes its commits while targeting dev. Review this contribution's two commits in the incremental comparison. This PR is ready for review and remains dependent on #9208 until the parent is integrated; the branch can then incorporate updated dev so the review diff contains only this increment. This is a separate precision-control enhancement, with its own issue and acceptance tests.

Description

At temperature 1e-8, the parent implementation can return the correct loss ln(3) but a float32 gradient residual of 1.0 for identical nonzero embeddings whose mathematical gradient is zero. The new keyword-only compute_dtype option makes the validated higher-precision workaround part of the loss configuration:

loss_fn = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)
loss = loss_fn(embeddings_a, embeddings_b)
loss.backward()

The connected embeddings are cast before cosine similarity and temperature construction. The output uses the requested dtype; gradients still propagate to the original inputs and are stored in their original dtypes. Input tensors are not mutated.

• None keeps the parent default behavior, including mixed-input promotion and float16/bfloat16 computation in at least float32.
• Explicit torch.float32 or torch.float64 selects the computation/output dtype. Explicit float32 can downcast float64 inputs.
• Unsupported options raise ValueError. Existing positional temperature/batch-size arguments remain valid.
• No automatic temperature threshold, default float64 promotion, gradient clipping, or change to cosine epsilon semantics is introduced.

Validation

Source: 116bc2a409604fd05953634a8314ac2bb3169ee3. Windows, Python 3.14.1, PyTorch 2.12.1+cpu, NumPy 2.3.5.

Check	Executed result
Initial 49 API cases on parent 49542ec8	All failed because compute_dtype was unavailable; 39 inherited cases passed
Final suite, including three additional mixed-dtype controls	91 passed, no skips: 39 inherited and 52 new cases
Loss-module line coverage	97%, 34/35 executable lines; existing mismatched-shape error line not exercised
Ruff, Black, isort on changed Python files	All passed
Numerical libraries / cleanup	One thread; disposable test dependencies, caches and temporary files removed

The new cases cover analytical zero gradients for identical embeddings, exact forward/backward equivalence to manual input casting for nonzero gradients, input preservation, mixed input dtypes, explicit float32/float64 output, CPU autocast, invalid options, singleton batches, and legacy arguments. Default behavior is also covered by all 39 parent cases.

The validator records the checkout revision before the final source commit; its source/test SHA-256 values match the final committed files.

python -m unittest tests.losses.test_contrastive_loss -v

A separate fixed-seed measurement repeats each of 24 vectors four times, split into two inputs of batch size 2, with feature sizes 3/5/17. At temperature 1e-8, the mathematical gradient is zero:

Computation	Maximum absolute gradient residual across the same 24 cases
Default	2.0
compute_dtype=torch.float64	3.725290298461914e-9

These CPU observations demonstrate reduced roundoff for the evaluated inputs. They do not guarantee exact gradients at every temperature, recover quantized embedding information, or expand the range of gradients stored in lower-precision inputs. Float64 requires device support and has additional memory/computation costs. GPU, distributed execution, full integration, training convergence, documentation builds, and performance were not tested. The inherited test named test_with_cuda used its CPU fallback.

Dragon infographic

The native-generated image explains the architecture, backward-roundoff problem, explicit float64 computation, and measured results. It is hosted separately from the code diff. No training benchmark or universal precision guarantee is implied.

Complete documentation and reproducibility

The backward limitation is real: cosine-similarity derivatives contain terms that should cancel for identical embeddings, but floating-point rounding can leave a residual. Scaling by 1 / temperature can amplify that residual at small temperatures. Stabilizing the forward reduction does not by itself correct the preceding cosine backward arithmetic. Computing cosine similarity in float64 reduces this error in the measured controls; casting the final scalar loss is too late.

• Complete report with architecture, API usage, results, limitations, dependency and Dragon infographic.
• Published validation evidence and reproduction scripts: initial failing tests, final test/coverage logs, source hashes and all 24 paired numerical measurements.
• To reproduce the recorded validation, check out code commit 116bc2a409604fd05953634a8314ac2bb3169ee3, copy the three published Python helper scripts into qa/gradient_precision/ at that checkout's root, then run python -B qa/gradient_precision/run_validation.py --phase green and python -B qa/gradient_precision/observe_precision.py. The helpers use disposable sandboxes and one thread per numerical library; the validator installs its small test-tool dependencies into the sandbox. PyTorch and NumPy must already be installed. These helpers use a minimal module bootstrap, so the recorded result is scoped CPU loss validation rather than a full MONAI integration run.

At the publication update, all six GitHub checks displayed success: DCO, pre-commit.ci, DeepSource Python/Shell/Docker, and CodeRabbit. This status does not imply GPU or training validation; the local validation limits above still apply. This PR is open and ready for review; merge dependency on parent #9208 remains in effect.

Duplicate check

Open and closed issues/PRs were searched for ContrastiveLoss dtype, compute_dtype, float64, precision, gradient, and roundoff. Related #9207/#9208 document the limitation and manual casting, but do not implement this explicit control. This increment leaves their forward stabilization in the parent contribution.

Checklist

• Optional keyword-only API; default behavior and legacy positional arguments preserved.
• Regression/acceptance tests and inline documentation added.
• CPU validation, source/test hash verification, and cleanup completed.
• Parent fix(losses): stabilize ContrastiveLoss at small temperatures #9208 integrated.
• GPU/distributed/full integration or training validation executed locally.
• Documentation built locally.

AI assistance disclosure: investigation, implementation, tests and documentation were prepared with OpenAI Codex, including independent adversarial reviews. The Dragon infographic was generated with native OpenAI image generation. All commits carry DCO sign-offs and AI assistance trailers; the submitting author remains responsible under MONAI's tool-assisted contribution policy.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: Project-MONAI/MONAI/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: cd92fe48-2928-4de0-b76b-ffb23e3a5c3d

📥 Commits

Reviewing files that changed from the base of the PR and between bbebca6 and 116bc2a.

📒 Files selected for processing: 2

• monai/losses/contrastive.py
• tests/losses/test_contrastive_loss.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 5 remain after this review.

📝 Walkthrough

ContrastiveLoss adds a keyword-only compute_dtype option and validates supported values. The forward pass promotes reduced-precision inputs to float32 by default or uses the selected computation dtype. It replaces explicit exponential reductions with shifted logits, self-pair masking, and logsumexp. Tests cover loss values, gradients, dtype behavior, autocast, and edge cases.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 116bc

No identified issue blocks merging; the untested GPU and performance cases remain outside the stated validation scope.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 9.52% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the main change: adding a computation dtype option to ContrastiveLoss.
Description check		The description is complete and directly related to the change. It includes the issue reference, implementation details, compatibility behavior, tests, validation results, limitations, dependency stat…
Linked Issues check		Issue #9209 requires a keyword-only compute_dtype accepting None, torch.float32, or torch.float64. ContrastiveLoss.__init__ validates these values and raises ValueError otherwise. `forward…
Out of Scope Changes check		The reviewed changes are limited to monai/losses/contrastive.py and its focused test file. The log-space forward path is the stated #9208 prerequisite for the None behavior and is directly relevan…

Full details: Description check

Explanation

The description is complete and directly related to the change. It includes the issue reference, implementation details, compatibility behavior, tests, validation results, limitations, dependency status, and checklist information.

Full details: Linked Issues check

Explanation

Issue #9209 requires a keyword-only compute_dtype accepting None, torch.float32, or torch.float64. ContrastiveLoss.__init__ validates these values and raises ValueError otherwise. forward casts the connected representation before cosine similarity and constructs temperature with the computation dtype. The selected dtype controls the output, while autograd keeps gradients connected to unchanged inputs and their original gradient dtypes. The default reduced-precision promotion and legacy positional arguments remain supported. Tests cover extreme-temperature gradients, casting equivalence, mixed input dtypes, output dtypes, input preservation, autocast, invalid options, singleton batches, and legacy arguments.

Full details: Out of Scope Changes check

Explanation

The reviewed changes are limited to monai/losses/contrastive.py and its focused test file. The log-space forward path is the stated #9208 prerequisite for the None behavior and is directly relevant to this loss precision objective. The added tests support the requested dtype behavior and regression coverage. No unrelated product behavior or unrelated files are shown.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}