feat(diffusion): Added Flow Matching

[Dibyajyoti-Chakraborty]

Add rectified flow matching to physicsnemo.diffusion

Adds rectified flow matching (conditional optimal transport) support to the diffusion framework, following the same design as the existing noise scheduler / loss / denoiser stack.

Added

• physicsnemo.diffusion.noise_schedulers.FlowMatchingNoiseScheduler — linear interpolation path (α(t) = 1-t, σ(t) = t) with a closed-form get_denoiser RHS.
• physicsnemo.diffusion.metrics.losses.FlowMatchingLoss — velocity-objective training loss; also supports x0/epsilon/score prediction types on any LinearGaussianNoiseScheduler.
• LinearGaussianNoiseScheduler.x0_to_velocity / .velocity_to_x0 conversions, and velocity_predictor support in get_denoiser (shared by all linear-Gaussian schedulers, not just flow matching).
• ConcatConditionWrapper.time_scale — rescales [0,1]-valued flow-matching times into a backbone's expected embedding range (e.g. 999.0). Backward compatible (defaults to 1.0, not a registered buffer, so old checkpoints load unchanged).

Docs

• New API reference entries in noise_schedulers.rst / metrics.rst.
• CHANGELOG.md updated.

Tests

• test/diffusion/test_flow_matching.py (new) — scheduler constructor, coefficients, timesteps/sampling, get_denoiser (velocity/x0/epsilon/score/SDE), sampling round-trip.
• FlowMatchingLoss coverage added to test_losses.py (constructor, non-regression, gradient flow, compile).
• velocity_predictor/x0_to_velocity/velocity_to_x0 coverage added to test_noise_schedulers.py.
• time_scale coverage added to test_wrappers.py.

Verified on a CPU node: full test/diffusion/ suite.

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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit 61e5816bda7a. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @CharlelieLrt — 21 file(s)

• docs/api/diffusion/metrics.rst
• docs/api/diffusion/noise_schedulers.rst
• physicsnemo/diffusion/metrics/init.py
• physicsnemo/diffusion/metrics/losses.py
• physicsnemo/diffusion/noise_schedulers/init.py
• physicsnemo/diffusion/noise_schedulers/flow_matching.py
• physicsnemo/diffusion/noise_schedulers/linear_gaussian.py
• physicsnemo/diffusion/noise_schedulers/noise_schedulers.py
• physicsnemo/diffusion/utils/model_wrappers.py
• test/diffusion/data/test_losses_fm_1d_epsilon.pth
• test/diffusion/data/test_losses_fm_1d_score.pth
• test/diffusion/data/test_losses_fm_1d_velocity.pth
• test/diffusion/data/test_losses_fm_1d_x0.pth
• test/diffusion/data/test_losses_fm_2d_epsilon.pth
• test/diffusion/data/test_losses_fm_2d_score.pth
• test/diffusion/data/test_losses_fm_2d_velocity.pth
• test/diffusion/data/test_losses_fm_2d_x0.pth
• test/diffusion/test_flow_matching.py
• test/diffusion/test_losses.py
• test/diffusion/test_noise_schedulers.py
• test/diffusion/test_wrappers.py

⏳ @megnvidia — 2 file(s)

• docs/api/diffusion/metrics.rst
• docs/api/diffusion/noise_schedulers.rst

No CODEOWNER

• CHANGELOG.md

Comment /codeowners-info to refresh.

[greptile-apps]

Greptile Summary

The PR adds rectified flow matching to the diffusion stack, including a scheduler, training loss, velocity conversions, and timestep scaling support.

• Adds FlowMatchingNoiseScheduler with linear interpolation coefficients and predictor-specific denoisers.
• Adds FlowMatchingLoss supporting velocity, x0, epsilon, and score prediction.
• Extends linear-Gaussian schedulers with velocity conversion and denoiser support.
• Adds documentation, compatibility exports, regression fixtures, and focused tests.

Important Files Changed

Filename	Overview
physicsnemo/diffusion/noise_schedulers/flow_matching.py	Adds the flow-matching schedule and closed-form ODE/SDE denoisers; the SDE path should reject its singular default t=1 endpoint.
physicsnemo/diffusion/metrics/losses.py	Adds a flow-matching training loss with shared-noise velocity targets and prediction conversions.
physicsnemo/diffusion/noise_schedulers/linear_gaussian.py	Adds x0/velocity conversions and velocity-predictor support to generic linear-Gaussian denoisers.
physicsnemo/diffusion/utils/model_wrappers.py	Adds backward-compatible timestep scaling before forwarding times to supported backbones.
test/diffusion/test_flow_matching.py	Covers scheduler coefficients, denoiser variants, and ODE sampling, but avoids rather than validates rejection of the singular SDE endpoint.

_{Reviews (1): Last reviewed commit: "Added Flow Matching" | Re-trigger Greptile}

[laserkelvin]

Just two minor general comments.

I think the broader one that I have, which affects the interface more is if we wanted to mask loss contributions before the reduction. For example in a chemistry case where we have a padded batch of atoms, we would want to apply a padding mask

[megnvidia]

LGTM

[CharlelieLrt]

Just two minor general comments.

I think the broader one that I have, which affects the interface more is if we wanted to mask loss contributions before the reduction. For example in a chemistry case where we have a padded batch of atoms, we would want to apply a padding mask

For parity with the denoising score matching loss, we should have a WeightedFlowMatchingLoss, which supports arbitrary weights (including a binary mask).

@laserkelvin would this work for the chemistry case you mentioned?

[Dibyajyoti-Chakraborty]

Updated done based on the comments:

• Named rectified flow matching: RectifiedFlowNoiseScheduler, (@CharlelieLrt)
• Name changes - "flow" (named to avoid clashing with v-prediction), flow_predictor/x0_to_flow/flow_to_x0 support. (@CharlelieLrt)
• PredictorType (name changed @CharlelieLrt ) .
• Added WeightedFlowMatchingLoss (support masking etc. @laserkelvin)
• Simplified the loss to delegate to the scheduler's add_noise/x0_to_flow instead of duplicating the noise-path formula. (@CharlelieLrt )

[laserkelvin]

Just two minor general comments.
I think the broader one that I have, which affects the interface more is if we wanted to mask loss contributions before the reduction. For example in a chemistry case where we have a padded batch of atoms, we would want to apply a padding mask

For parity with the denoising score matching loss, we should have a WeightedFlowMatchingLoss, which supports arbitrary weights (including a binary mask).

@laserkelvin would this work for the chemistry case you mentioned?

Yeah weight argument would work for this, thanks!

[Dibyajyoti-Chakraborty]

/ok to test cc0b84d

[CharlelieLrt]

IMO that's not a good pattern because it deviates too much from the existing MSEDSMLoss. Nothing wrong in the implementation in itself, but it would be preferrable to have API parity with the existing loss.
For example, the [optional] conversion methods should be provided as callbacks rather than inferred from the scheduler object. This can be useful when one needs to change just the prediction type without defining a new scheduler class. So, there is redundancy, but it improves API flexibility

[CharlelieLrt]

One detail though: the denoising score matching losses expose conversion to x0 (score-to-x0 and noise-to-x0). That is because the loss formulation is based on x0, and any conversion is just an adaptation of this. IMO, the basis prediction for the FlowMatching loss should be the flow prediction, so it should expose callbacks to convert to x0 (x0-to-flow, epsilon-to-flow, score-to-flow)

[Dibyajyoti-Chakraborty]

Done, both points. The noise_scheduler parameter is documented as the NoiseScheduler protocol only, and the implementation no longer inspects the scheduler type.

The flow matching losses now expose x0_to_flow_fn, score_to_flow_fn, and epsilon_to_flow_fn callbacks with the same pattern as MSEDSMLoss. Flow is the base prediction: every prediction type converts to flow. One note: x0_to_flow_fn is required for all prediction types because it also computes the regression target from clean data.

[CharlelieLrt]

The new tests should be reorganized for clarity:

1. All pytest configs at the top
2. Constructor tests together
3. Non regression tests for the non-weighted losses (both DSM and flow mathcing)
4. Same for the weighted version

[Dibyajyoti-Chakraborty]

Done. The file now follows this layout, and the FM losses run through the same classes as the DSM losses via LOSS_CONFIGS/WEIGHTED_LOSS_CONFIGS rows (loss_cls, sched_cls, kwargs, name) — analogous to SCHEDULER_CONFIGS. Same for the multi-diffusion loss/predictor test files.

[CharlelieLrt]

Two major concerns:

• The PR claims to add support for "flow" prediction, but this support is only very partial with missing conversion callbacks at multiple place in the module
• New tests are a not great, with massive amounts of code duplication, new class-specific configs and tests, while it would be much better to extend the existing tests, etc...