DSMLoss implementation + minor bugfixes in noise_schedulers.py and pr…

[CharlelieLrt]

…econditioners.py

PhysicsNeMo Pull Request

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[CharlelieLrt]

/blossom-ci

[greptile-apps]

Greptile Summary

This PR introduces DSMLoss (Denoising Score Matching loss) for training diffusion models, along with supporting infrastructure changes to noise schedulers and preconditioners.

Major additions:

• New DSMLoss class in physicsnemo/diffusion/metrics/losses.py with comprehensive documentation, examples, and support for both x0-prediction and score-prediction training modes
• Added loss_weight() method to all noise scheduler implementations to compute per-sample training weights
• Added score_to_x0() method to LinearGaussianNoiseScheduler for converting score predictions to x0 estimates

Improvements to noise schedulers:

• Changed training time sampling in EDMNoiseScheduler and StudentTEDMNoiseScheduler from log-uniform to log-normal distribution
• Added new parameters (sigma_data, P_mean, P_std) to EDM-based schedulers for better control
• Fixed _sample_student_t signature to use *shape instead of shape: Tuple[int, ...] for cleaner syntax

Bug fixes:

• Removed unused _exponent helper method from VPNoiseScheduler and inlined its logic into alpha() method
• Fixed Student-t sampling to properly handle device and dtype parameters

Documentation updates:

• Enhanced preconditioner docstrings to clarify terminology ("noise schedule" → "time-to-noise-level mapping")
• Added usage guidance linking each preconditioner to its corresponding noise scheduler and loss function
• All new code follows NumPy docstring conventions with proper LaTeX math notation

Important Files Changed

Filename	Overview
physicsnemo/diffusion/metrics/losses.py	New DSMLoss class implementation with comprehensive docstrings and examples, following NumPy-style conventions
physicsnemo/diffusion/noise_schedulers/noise_schedulers.py	Added loss_weight and score_to_x0 methods to NoiseScheduler protocol and implementations; updated EDM/StudentT time sampling to log-normal; refactored VPNoiseScheduler to inline _exponent; fixed Student-t sampling signature

_{Last reviewed commit: 0dd567e}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[pzharrington]

LGTM!

[CharlelieLrt]

/blossom-ci