diff --git a/CHANGELOG.md b/CHANGELOG.md index b76c9207ab..bc8d798286 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 as in sCM and EDM2. Existing single-time, single-prediction calls remain unchanged. Existing DSM and flow-matching losses continue to consume single tensors. +- Adds `TrigFlowNoiseScheduler` and `TrigFlowPreconditioner` to the diffusion + noise schedulers and preconditioners in `physicsnemo.diffusion`. Pair them + with the same `sigma_data` to train and sample models along a trigonometric + diffusion path, with the backbone predicting normalized velocity. Built-in + linear Gaussian schedulers also provide differentiable `snr_inv()` methods + to convert signal-to-noise amplitude ratios back to diffusion times. Use + these inverses to match noise levels across schedules. The amplitude ratio + is `alpha(t) / sigma(t)`. The power signal-to-noise ratio is its square. + Custom schedulers can supply their own inverse. These additions are opt-in. + Existing schedulers, preconditioners, and training workflows work as before. - Adds standalone FLARE++ attention and model APIs, with input-conditioned dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside GeoTransolver. diff --git a/docs/api/diffusion/noise_schedulers.rst b/docs/api/diffusion/noise_schedulers.rst index dda757e7fe..90f0d58675 100644 --- a/docs/api/diffusion/noise_schedulers.rst +++ b/docs/api/diffusion/noise_schedulers.rst @@ -23,6 +23,46 @@ the coupling between noise schedule and reverse process in one place and avoids the need for schedule-specific sampler code. +Amplitude ratio and inverse +--------------------------- + +The amplitude ratio :math:`\alpha(t)/\sigma(t)` measures signal strength +relative to noise strength. The method ``scheduler.snr(t)`` returns this +ratio. The usual power signal-to-noise ratio is its square. + +Use ``scheduler.snr_inv(q)`` to recover a diffusion time from an amplitude +ratio or to match noise levels between schedules. EDM, VE, VP, rectified-flow, +continuous iDDPM, and TrigFlow provide differentiable inverses. Use finite +positive ratios on the schedule's invertible domain. Inverses do not clip +to training bounds. Add ``snr_inv`` for custom schedulers that need this +conversion, and keep it differentiable for training losses such as sCM. + +.. code-block:: python + + import torch + from physicsnemo.diffusion.noise_schedulers import TrigFlowNoiseScheduler + + scheduler = TrigFlowNoiseScheduler(sigma_data=0.5) + t = torch.tensor([0.2, 0.7]) + recovered = scheduler.snr_inv(scheduler.snr(t)) + torch.testing.assert_close(recovered, t) + +TrigFlow interpolates between data and Gaussian noise with trigonometric +coefficients. Its diffusion time is an angle in radians, whereas EDM uses +the noise standard deviation. + +Pair ``TrigFlowNoiseScheduler`` with ``TrigFlowPreconditioner`` and the same +``sigma_data``. Training samples diffusion times uniformly. Inference uses +an evenly spaced grid ending at zero. The denoising weight +:math:`1/(\sigma_d\sin(t))^2` gives unit weight to the normalized-velocity +residual with this preconditioner. Keep training diffusion times above zero +for denoising losses and below :math:`\pi/2` for reverse-process sampling. + +.. autoclass:: TrigFlowNoiseScheduler + :show-inheritance: + :members: + + Role in Training ---------------- diff --git a/docs/api/diffusion/preconditioners.rst b/docs/api/diffusion/preconditioners.rst index 8a80c59f24..5f150a6c94 100644 --- a/docs/api/diffusion/preconditioners.rst +++ b/docs/api/diffusion/preconditioners.rst @@ -140,6 +140,14 @@ API Reference :members: :exclude-members: forward +:code:`TrigFlowPreconditioner` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: physicsnemo.diffusion.preconditioners.TrigFlowPreconditioner + :show-inheritance: + :members: + :exclude-members: forward + :code:`VEPreconditioner` ~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/physicsnemo/diffusion/noise_schedulers/__init__.py b/physicsnemo/diffusion/noise_schedulers/__init__.py index 21cbe392a6..5316af9fc3 100644 --- a/physicsnemo/diffusion/noise_schedulers/__init__.py +++ b/physicsnemo/diffusion/noise_schedulers/__init__.py @@ -21,6 +21,7 @@ from .linear_gaussian import LinearGaussianNoiseScheduler # noqa: F401 from .rectified_flow import RectifiedFlowNoiseScheduler # noqa: F401 from .student_t_edm import StudentTEDMNoiseScheduler # noqa: F401 +from .trigflow import TrigFlowNoiseScheduler # noqa: F401 from .ve import VENoiseScheduler # noqa: F401 from .vp import VPNoiseScheduler # noqa: F401 diff --git a/physicsnemo/diffusion/noise_schedulers/edm.py b/physicsnemo/diffusion/noise_schedulers/edm.py index 6b0174e103..5189965ea3 100644 --- a/physicsnemo/diffusion/noise_schedulers/edm.py +++ b/physicsnemo/diffusion/noise_schedulers/edm.py @@ -71,7 +71,9 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the EDM noise scheduler: + The first sequence covers the EDM training and sampling interfaces with one + scheduler. Its final call converts a clean-data predictor into the right-hand side + used by an ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -98,7 +100,8 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): >>> denoiser(xN, tN).shape # ODE RHS for sampling torch.Size([4, 3, 8, 8]) - Per-channel ``sigma_data`` for heterogeneous channels: + A channel-specific ``sigma_data`` assigns different data scales to each channel. + The resulting loss weight retains a separate value for each channel. >>> sigma_per_ch = torch.tensor([0.3, 0.5, 0.7]) >>> scheduler_ch = EDMNoiseScheduler(sigma_data=sigma_per_ch) @@ -143,6 +146,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/iddpm.py b/physicsnemo/diffusion/noise_schedulers/iddpm.py index ca33329988..126a5fc589 100644 --- a/physicsnemo/diffusion/noise_schedulers/iddpm.py +++ b/physicsnemo/diffusion/noise_schedulers/iddpm.py @@ -59,7 +59,8 @@ class IDDPMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the iDDPM noise scheduler: + The first sequence builds noisy training samples and a reverse-time sampling state + for iDDPM. It then wraps a clean-data predictor as an ODE right-hand side. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import IDDPMNoiseScheduler @@ -131,6 +132,10 @@ def sigma_inv( r"""For iDDPM, :math:`t = \sigma` (identity mapping).""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find continuous diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return q.reciprocal() + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py index ca5d7ae3e0..65196bbd97 100644 --- a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py +++ b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py @@ -130,6 +130,7 @@ class LinearGaussianNoiseScheduler(ABC, NoiseScheduler): ... smin_rho = self.sigma_min**(1/self.rho) ... frac = i/(num_steps-1) ... t = (smax_rho + frac * (smin_rho - smax_rho))**self.rho + ... # Include the clean endpoint after the reverse-time grid. ... return torch.cat([t, torch.zeros(1, device=device)]) ... ... def sample_time(self, N, *, device=None, dtype=None): @@ -456,6 +457,50 @@ def snr( """ return self.alpha(t) / self.sigma(t) + def snr_inv( + self, + q: Float[Tensor, " *shape"], + ) -> Float[Tensor, " *shape"]: + r"""Find the diffusion time for a signal-to-noise amplitude ratio. + + Use this to match noise levels between diffusion schedules. The + ratio compares signal and noise amplitudes. The usual power + signal-to-noise ratio is its square. + + .. math:: + + q = \alpha(t) / \sigma(t), \qquad t = \mathrm{snr\_inv}(q). + + Parameters + ---------- + q : Tensor + Amplitude ratio, not its square. Built-in continuous inverses + support finite positive ratios without clipping to training bounds. + + Returns + ------- + Tensor + Diffusion times with the shape of ``q``. + + Notes + ----- + Override this method if your custom scheduler needs an inverse. + Keep the calculation differentiable for training losses such as sCM. + Schedulers without an inverse raise ``NotImplementedError`` here. + + Examples + -------- + For EDM, the signal coefficient is one and the noise coefficient equals time. + Inverting a signal-to-noise ratio of two gives time one half. + + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler + >>> EDMNoiseScheduler().snr_inv(torch.tensor([2.0])) + tensor([0.5000]) + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement snr_inv()." + ) + def x0_to_score( self, x0: Float[Tensor, " B *dims"], @@ -491,6 +536,10 @@ def x0_to_score( Examples -------- + A caller can convert a clean-data predictor into a score predictor when it needs + score parameterization. The conversion uses the noisy state and time to recover + the implied Gaussian score. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -562,6 +611,10 @@ def score_to_x0( Examples -------- + A caller can convert a score predictor into a clean-data predictor for a loss or + denoiser that expects that parameterization. The conversion combines the score + with the noisy state at the same time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -620,6 +673,10 @@ def epsilon_to_score( Examples -------- + A caller can convert a noise predictor into a score predictor without changing + its backbone. The scheduler applies its time-dependent noise coefficient to the + predicted noise. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -677,6 +734,9 @@ def score_to_epsilon( Examples -------- + A score predictor can also supply a noise prediction for a noise-based consumer. + The scheduler scales the score by its time-dependent noise coefficient. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion: @@ -740,6 +800,10 @@ def epsilon_to_x0( Examples -------- + A noise predictor can supply a clean-data estimate for losses that train on + clean targets. The conversion subtracts predicted noise from the noisy state + using the schedule coefficients. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -801,6 +865,10 @@ def x0_to_epsilon( Examples -------- + A clean-data predictor can supply a noise estimate without retraining the + backbone. The conversion subtracts the predicted clean signal from the noisy + state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion: @@ -867,6 +935,10 @@ def x0_to_flow( Examples -------- + A clean-data predictor can supply a velocity target for a flow-matching loss. + The conversion differentiates the schedule coefficients at the same diffusion + time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -934,6 +1006,10 @@ def flow_to_x0( Examples -------- + A caller can convert a velocity predictor to a clean-data predictor when it + needs the original sample. The scheduler solves the linear-Gaussian velocity + relation using the noisy state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion: @@ -1003,6 +1079,10 @@ def score_to_flow( Examples -------- + A score predictor can supply the velocity expected by a flow-matching loss. The + scheduler first relates the score to the clean signal, then applies the schedule + derivatives. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -1073,6 +1153,9 @@ def flow_to_score( Examples -------- + A velocity predictor can supply the score needed by a denoiser. The scheduler + uses the noisy state and schedule derivatives to perform the conversion. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion @@ -1190,7 +1273,9 @@ def get_denoiser( Examples -------- - Generate ODE RHS from a score-predictor: + A denoiser returns the reverse-process right-hand side used by a solver. This + first example starts from a score predictor and requests the probability-flow + ODE. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -1204,7 +1289,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an x0-predictor (score conversion is done internally): + A clean-data predictor can use the same ODE interface without producing a + score explicitly. The scheduler converts its output to the required score. >>> x0_pred = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2) # Toy x0-predictor >>> denoiser = scheduler.get_denoiser( @@ -1213,7 +1299,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an epsilon-predictor: + A noise predictor can also drive the ODE through the same factory. + The scheduler converts its output before evaluating the right-hand side. >>> eps_pred = lambda x, t: x * 0.1 # Toy epsilon-predictor >>> denoiser = scheduler.get_denoiser( diff --git a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py index 9ec1935f1d..f21eba4539 100644 --- a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py +++ b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py @@ -86,8 +86,9 @@ class RectifiedFlowNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Construct noisy states for training, then initialize sampling with the same - flow parameterization: + Training interpolates clean data toward noise, while sampling starts near the noisy + endpoint. The flow predictor then supplies the ODE right-hand side under the same + time convention. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import ( @@ -154,6 +155,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/(1+q)` for amplitude ratios ``q > 0``.""" + return 1 / (1 + q) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/trigflow.py b/physicsnemo/diffusion/noise_schedulers/trigflow.py new file mode 100644 index 0000000000..69e96759c3 --- /dev/null +++ b/physicsnemo/diffusion/noise_schedulers/trigflow.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Trigonometric interpolation between data and Gaussian noise.""" + +import math + +import torch +from jaxtyping import Float +from torch import Tensor + +from .linear_gaussian import LinearGaussianNoiseScheduler + + +class TrigFlowNoiseScheduler(LinearGaussianNoiseScheduler): + r"""Interpolate between clean data and Gaussian noise along a trigonometric path. + + Use this scheduler for TrigFlow denoising or consistency training, paired + with ``TrigFlowPreconditioner`` and the same ``sigma_data``. Diffusion time + is an angle in radians, whereas EDM uses the noise standard deviation. + Training draws diffusion times uniformly. Inference uses evenly spaced + diffusion times ending at zero. + + .. math:: + + x_t = \cos(t)x_0 + \sigma_d \sin(t)\epsilon, + \qquad \epsilon \sim \mathcal{N}(0,I). + + Parameters + ---------- + sigma_data : float, optional, default=0.5 + Data scale and noise standard deviation at :math:`t=\pi/2`. + Must be finite and positive. + t_min : float, optional, default=0.001 + Lower training diffusion time in radians, with + :math:`0 \leq t_{\min} < t_{\max}`. Use a positive value for + denoising loss weights, which diverge at zero. + t_max : float, optional, default=1.55 + Upper training diffusion time and first inference diffusion time, + in radians. Must be below ``pi / 2``. The default stays below that + endpoint in reduced precision. + + Examples + -------- + Use this scheduler when the model treats diffusion time as an angle and predicts + TrigFlow velocity. The example checks noisy-state shape, reverse-time grid size, and + the inverse signal-to-noise mapping. + + >>> scheduler = TrigFlowNoiseScheduler() + >>> t = scheduler.sample_time(2) + >>> scheduler.add_noise(torch.zeros(2, 3), t).shape + torch.Size([2, 3]) + >>> scheduler.timesteps(4).shape + torch.Size([5]) + >>> # Use interior angles for the inverse SNR check. + >>> t = torch.tensor([0.2, 0.7]) + >>> torch.allclose(scheduler.snr_inv(scheduler.snr(t)), t) + True + + Note + ---- + Coefficient functions accept angular times without clipping. The training + bounds do not limit coefficient evaluation. At zero, score conversions + and denoising weights are singular. At ``pi / 2``, conversions that divide + by the signal coefficient are singular. + """ + + def __init__( + self, sigma_data: float = 0.5, t_min: float = 0.001, t_max: float = 1.55 + ) -> None: + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + if not 0 <= t_min < t_max < math.pi / 2: + raise ValueError( + f"Expected 0 <= t_min < t_max < pi/2, got {t_min=}, {t_max=}." + ) + self.sigma_data = sigma_data + self.t_min = t_min + self.t_max = t_max + + def alpha(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal coefficient: :math:`\alpha(t)=\cos(t)`.""" + return t.cos() + + def alpha_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal derivative: :math:`\dot\alpha(t)=-\sin(t)`.""" + return -t.sin() + + def sigma(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise coefficient: :math:`\sigma(t)=\sigma_d\sin(t)`.""" + return self.sigma_data * t.sin() + + def sigma_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise derivative: :math:`\dot\sigma(t)=\sigma_d\cos(t)`.""" + return self.sigma_data * t.cos() + + def sigma_inv(self, sigma: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arcsin(\sigma/\sigma_d)` for noise scales in ``[0, sigma_data]``.""" + return (sigma / self.sigma_data).asin() + + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arctan(1/(\sigma_d q))` for nonnegative amplitude ratios.""" + return torch.atan2(torch.ones_like(q), self.sigma_data * q) + + def sample_time( + self, + N: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N"]: + r"""Draw training diffusion times uniformly between ``t_min`` and ``t_max``. + + Parameters + ---------- + N : int + Number of training examples. + device : torch.device, optional + Device for the sampled times. + dtype : torch.dtype, optional + Floating-point type for the sampled times. + + Returns + ------- + Tensor + Training diffusion times in radians with shape :math:`(N,)`. + """ + return self.t_min + (self.t_max - self.t_min) * torch.rand( + N, device=device, dtype=dtype + ) + + def timesteps( + self, + num_steps: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N+1"]: + r"""Return evenly spaced inference diffusion times from ``t_max`` to zero. + + Parameters + ---------- + num_steps : int + Number of solver steps, at least one. + device : torch.device, optional + Device for the inference grid. + dtype : torch.dtype, optional + Floating-point type for the grid. + + Returns + ------- + Tensor + Decreasing diffusion times in radians with shape + :math:`(\mathrm{num\_steps}+1,)`. + """ + if num_steps < 1: + raise ValueError(f"Expected num_steps >= 1, got {num_steps}.") + return torch.linspace(self.t_max, 0, num_steps + 1, device=device, dtype=dtype) + + def loss_weight(self, t: Float[Tensor, " N"]) -> Float[Tensor, " N"]: + r"""Weight denoising residuals by :math:`1/(\sigma_d\sin(t))^2`. + + Use with ``TrigFlowPreconditioner`` to give unit weight to the + normalized-velocity residual. Use positive diffusion times to avoid + the singularity at zero. + + Parameters + ---------- + t : Tensor + Training diffusion times in radians with shape :math:`(N,)`. + + Returns + ------- + Tensor + Per-example denoising weights. + """ + return self.sigma(t).square().reciprocal() diff --git a/physicsnemo/diffusion/noise_schedulers/ve.py b/physicsnemo/diffusion/noise_schedulers/ve.py index 608bbdc1de..4351b8faea 100644 --- a/physicsnemo/diffusion/noise_schedulers/ve.py +++ b/physicsnemo/diffusion/noise_schedulers/ve.py @@ -55,7 +55,9 @@ class VENoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VE noise scheduler: + The first sequence uses a variance-exploding schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VENoiseScheduler @@ -105,6 +107,10 @@ def sigma_inv( r"""Inverse VE mapping: :math:`t = \sigma^2`.""" return sigma**2 + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = q^{-2}` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/vp.py b/physicsnemo/diffusion/noise_schedulers/vp.py index ee6582822d..010f4c5cfd 100644 --- a/physicsnemo/diffusion/noise_schedulers/vp.py +++ b/physicsnemo/diffusion/noise_schedulers/vp.py @@ -66,7 +66,9 @@ class VPNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VP noise scheduler: + The first sequence uses a variance-preserving schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VPNoiseScheduler @@ -158,6 +160,20 @@ def sigma_inv( discriminant = self.beta_min**2 + 2 * self.beta_d * a return 2 * a / (self.beta_min + torch.sqrt(discriminant)) + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time for finite amplitude ratios ``q > 0``. + + .. math:: + + A = \log(1+q^{-2}), \qquad + t = \frac{2A}{\beta_{\min}+\sqrt{\beta_{\min}^2+2\beta_d A}}. + + The rationalized root avoids cancellation and supports ``beta_d=0`` + when ``beta_min > 0``. See the base ``snr_inv`` contract. + """ + a = torch.logaddexp(torch.zeros_like(q), -2 * q.log()) + return 2 * a / (self.beta_min + (self.beta_min**2 + 2 * self.beta_d * a).sqrt()) + def timesteps( self, num_steps: int, diff --git a/physicsnemo/diffusion/preconditioners/__init__.py b/physicsnemo/diffusion/preconditioners/__init__.py index 848114c5f8..8dcdabe631 100644 --- a/physicsnemo/diffusion/preconditioners/__init__.py +++ b/physicsnemo/diffusion/preconditioners/__init__.py @@ -28,6 +28,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) diff --git a/physicsnemo/diffusion/preconditioners/preconditioners.py b/physicsnemo/diffusion/preconditioners/preconditioners.py index eeac45dd74..3a5392ea02 100644 --- a/physicsnemo/diffusion/preconditioners/preconditioners.py +++ b/physicsnemo/diffusion/preconditioners/preconditioners.py @@ -580,6 +580,94 @@ def forward( return c_skip * x + c_out * F_x +class TrigFlowPreconditioner(BaseAffinePreconditioner): + r"""Convert normalized-velocity predictions into clean-data predictions. + + Use with ``TrigFlowNoiseScheduler`` and the same ``sigma_data`` for + TrigFlow denoising or consistency training. The backbone predicts velocity + divided by ``sigma_data``, rather than clean data or noise. Diffusion times + are angles in radians. A backbone can return a prediction paired with + ``DiffusionExtraOutputs``; its learned loss weight passes through unchanged. + + .. math:: + + D_\theta(x,t) = \cos(t)x - \sigma_d\sin(t)F_\theta(x/\sigma_d,t). + + Parameters + ---------- + model : physicsnemo.Module + Backbone predicting velocity divided by ``sigma_data``. + sigma_data : float, optional, default=0.5 + Data scale, matching the noise scheduler. Must be finite and positive. + + Forward + ------- + x : Tensor + Noisy states with shape :math:`(B, \ldots)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time in radians. Pass ``(t, extra_times)`` to include a + destination diffusion time ``r``. + condition : Tensor or TensorDict, optional + Conditioning passed to the backbone. + **model_kwargs : Any + Model-specific keyword arguments. + + Outputs + ------- + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + Clean-data prediction, optionally paired with its unchanged loss weight. + + Examples + -------- + The backbone in this example predicts zero normalized velocity, so the wrapper + returns the clean-data skip term. At zero angular time, that term equals the input. + + >>> class Velocity(Module): + ... def forward(self, x, t, condition=None): + ... return torch.zeros_like(x) + >>> model = TrigFlowPreconditioner(Velocity()) + >>> model(torch.ones(2, 3), torch.zeros(2)) + tensor([[1., 1., 1.], + [1., 1., 1.]]) + """ + + sigma_data: Float[torch.Tensor, ""] + + def __init__(self, model: Module, sigma_data: float = 0.5) -> None: + super().__init__(model) + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + self.register_buffer("sigma_data", torch.tensor(float(sigma_data))) + + def compute_coefficients( + self, t: Float[torch.Tensor, " B *dims"], / + ) -> tuple[ + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + ]: + r"""Return input, time, output, and skip coefficients for diffusion time. + + Parameters + ---------- + t : Tensor + Diffusion times in radians with shape :math:`(B, 1, \ldots)` + for broadcasting over the noisy states. + + Returns + ------- + tuple of Tensor + Coefficients ``(1 / sigma_data, t, -sigma_data * sin(t), cos(t))``. + """ + return ( + torch.ones_like(t) / self.sigma_data, + t, + -self.sigma_data * t.sin(), + t.cos(), + ) + + class VPPreconditioner(BaseAffinePreconditioner): r""" Variance Preserving (VP) preconditioner. diff --git a/test/diffusion/data/trigflow_precond_conv.mdlus b/test/diffusion/data/trigflow_precond_conv.mdlus new file mode 100644 index 0000000000..d5ba176859 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_conv.mdlus differ diff --git a/test/diffusion/data/trigflow_precond_conv_coefficients.pth b/test/diffusion/data/trigflow_precond_conv_coefficients.pth new file mode 100644 index 0000000000..ae5502fa5c Binary files /dev/null and b/test/diffusion/data/trigflow_precond_conv_coefficients.pth differ diff --git a/test/diffusion/data/trigflow_precond_conv_forward.pth b/test/diffusion/data/trigflow_precond_conv_forward.pth new file mode 100644 index 0000000000..b38b788800 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_conv_forward.pth differ diff --git a/test/diffusion/data/trigflow_precond_conv_sigma.pth b/test/diffusion/data/trigflow_precond_conv_sigma.pth new file mode 100644 index 0000000000..777d486de2 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_conv_sigma.pth differ diff --git a/test/diffusion/data/trigflow_precond_linear.mdlus b/test/diffusion/data/trigflow_precond_linear.mdlus new file mode 100644 index 0000000000..7b0d582227 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_linear.mdlus differ diff --git a/test/diffusion/data/trigflow_precond_linear_coefficients.pth b/test/diffusion/data/trigflow_precond_linear_coefficients.pth new file mode 100644 index 0000000000..57b36a0b26 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_linear_coefficients.pth differ diff --git a/test/diffusion/data/trigflow_precond_linear_forward.pth b/test/diffusion/data/trigflow_precond_linear_forward.pth new file mode 100644 index 0000000000..d1a46ac1ab Binary files /dev/null and b/test/diffusion/data/trigflow_precond_linear_forward.pth differ diff --git a/test/diffusion/data/trigflow_precond_linear_sigma.pth b/test/diffusion/data/trigflow_precond_linear_sigma.pth new file mode 100644 index 0000000000..76f4b47c54 Binary files /dev/null and b/test/diffusion/data/trigflow_precond_linear_sigma.pth differ diff --git a/test/diffusion/test_model_io.py b/test/diffusion/test_model_io.py index 4e4681b124..b1afdf679a 100644 --- a/test/diffusion/test_model_io.py +++ b/test/diffusion/test_model_io.py @@ -24,6 +24,7 @@ from physicsnemo.diffusion.preconditioners import ( EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -36,8 +37,9 @@ (VEPreconditioner, {}), (VPPreconditioner, {}), (IDDPMPreconditioner, {"M": 32}), + (TrigFlowPreconditioner, {"sigma_data": 1.4}), ) -IDS = ("edm", "ve", "vp", "iddpm") +IDS = ("edm", "ve", "vp", "iddpm", "trigflow") TIME_FIELDS = (None, (), ("r",), ("r1",), ("r2",), ("r1", "r2")) TIME_IDS = ("tensor", "empty", "r", "r1", "r2", "r1_r2") diff --git a/test/diffusion/test_noise_schedulers.py b/test/diffusion/test_noise_schedulers.py index 13fb98c917..20dbf741f7 100644 --- a/test/diffusion/test_noise_schedulers.py +++ b/test/diffusion/test_noise_schedulers.py @@ -289,6 +289,10 @@ def test_concrete_snr(self, device): t = torch.tensor([0.5, 2.0], device=device) assert torch.allclose(s.snr(t), 1 / t, atol=1e-7) + def test_optional_snr_inverse(self): + with pytest.raises(NotImplementedError, match="snr_inv"): + _MinimalScheduler().snr_inv(torch.ones(2)) + def test_concrete_add_noise(self, device): s = _MinimalScheduler() x0 = make_input((2, 3, 8, 8), seed=3, device=device) diff --git a/test/diffusion/test_preconditioners.py b/test/diffusion/test_preconditioners.py index beaccfda5d..8bee845b29 100644 --- a/test/diffusion/test_preconditioners.py +++ b/test/diffusion/test_preconditioners.py @@ -27,6 +27,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -113,6 +114,7 @@ def forward( # Preconditioner configurations for parameterized tests PRECOND_CONFIGS = [ + (TrigFlowPreconditioner, {"sigma_data": 1.4}, "trigflow_precond"), ( VPPreconditioner, {"beta_d": 19.9, "beta_min": 0.1, "M": 2000}, @@ -403,7 +405,7 @@ def test_single_element_sequence_stored_as_scalar(self, simple_model): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestNonRegression: """Non-regression tests parameterized across all preconditioner types.""" @@ -611,7 +613,7 @@ def create_fn(): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestAllPreconditioners: """Tests that apply to all preconditioner types.""" diff --git a/test/diffusion/test_trigflow.py b/test/diffusion/test_trigflow.py new file mode 100644 index 0000000000..8a32f9acee --- /dev/null +++ b/test/diffusion/test_trigflow.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Check continuous inverse SNR and TrigFlow coordinates.""" + +import pytest +import torch + +from physicsnemo.diffusion import ( + DiffusionExtraTimes, + preconditioners, +) +from physicsnemo.diffusion import noise_schedulers as schedules + +from .helpers import StructuredDiffusionModel + +CONFIGS = ( + ("EDMNoiseScheduler", {}), + ("EDMLogUniformNoiseScheduler", {}), + ("VENoiseScheduler", {}), + ("VPNoiseScheduler", {}), + ("VPNoiseScheduler", {"beta_d": 0.0}), + ("IDDPMNoiseScheduler", {"M": 32}), + ("RectifiedFlowNoiseScheduler", {}), + ("TrigFlowNoiseScheduler", {"sigma_data": 0.5}), + ("TrigFlowNoiseScheduler", {"sigma_data": 1.4}), +) +IDS = ( + "edm", + "edm-uniform", + "ve", + "vp", + "vp-constant", + "iddpm", + "rf", + "trig", + "trig-scaled", +) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestConsistency: + def test_snr_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + recovered, tangent = torch.func.jvp( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + (t,), + (torch.ones_like(t),), + ) + torch.testing.assert_close(recovered, t) + torch.testing.assert_close(tangent, torch.ones_like(t)) + + def test_coefficient_derivatives(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + for coefficient, derivative in ( + (scheduler.alpha, scheduler.alpha_dot), + (scheduler.sigma, scheduler.sigma_dot), + ): + _, actual = torch.func.jvp(coefficient, (t,), (torch.ones_like(t),)) + torch.testing.assert_close(actual, derivative(t)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestGradientFlow: + def test_trigflow(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).double() + x = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + t = torch.tensor([0.4, 0.8], dtype=torch.float64, requires_grad=True) + + def forward(x, t): + prediction, extras = model( + x, (t, DiffusionExtraTimes(r=t / 2, batch_size=[2])) + ) + return ( + prediction, + extras.logvar, + scheduler.snr_inv(scheduler.snr(t)), + ) + + assert torch.autograd.gradcheck(forward, (x, t)) + + +class TestNumerics: + @pytest.mark.parametrize( + "dtype,ratio", + ( + (torch.float16, 0.003), + (torch.float32, 1e-20), + (torch.float64, 1e-200), + ), + ) + def test_vp_small_ratio(self, dtype, ratio, device): + scheduler = schedules.VPNoiseScheduler() + q = torch.tensor([ratio], dtype=dtype, device=device, requires_grad=True) + result = scheduler.snr_inv(q) + (gradient,) = torch.autograd.grad(result.sum(), (q,)) + q64 = q.detach().double() + a = torch.logaddexp(torch.zeros_like(q64), -2 * q64.log()) + root = (scheduler.beta_min**2 + 2 * scheduler.beta_d * a).sqrt() + expected = 2 * a / (scheduler.beta_min + root) + expected_gradient = -2 / (q64 * (1 + q64.square()) * root) + torch.testing.assert_close(result, expected.to(dtype)) + torch.testing.assert_close(gradient, expected_gradient.to(dtype)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestConstructor: + def test_trigflow(self, sigma_data, device): + scheduler = schedules.TrigFlowNoiseScheduler( + sigma_data=sigma_data, t_min=0.1, t_max=1.3 + ) + t = scheduler.sample_time(128, device=device, dtype=torch.float64) + assert torch.all((t >= 0.1) & (t <= 1.3)) + grid = scheduler.timesteps(8, device=device) + assert grid.shape == (9,) + torch.testing.assert_close(grid, torch.linspace(1.3, 0, 9, device=device)) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).to(device) + x = torch.randn(2, 3, device=device) + t = torch.tensor([0.4, 0.8], device=device) + result, extras = model(x, t) + raw, raw_extras = model.model(x / sigma_data, t) + expected = t.cos()[:, None] * x - sigma_data * t.sin()[:, None] * raw + torch.testing.assert_close(result, expected) + torch.testing.assert_close(extras.logvar, raw_extras.logvar) + torch.testing.assert_close( + scheduler.loss_weight(t), 1 / (sigma_data * t.sin()) ** 2 + ) + + def test_zero_time(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + t = torch.zeros(2) + torch.testing.assert_close(scheduler.alpha(t), torch.ones_like(t)) + torch.testing.assert_close(scheduler.sigma(t), t) + torch.testing.assert_close(scheduler.alpha_dot(t), t) + torch.testing.assert_close( + scheduler.sigma_dot(t), torch.full_like(t, sigma_data) + ) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestCompile: + def test_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.2, 0.7], device=device, requires_grad=True) + fn = torch.compile( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + backend="aot_eager", + fullgraph=True, + ) + result = fn(t) + with torch._dynamo.config.patch(error_on_recompile=True): + torch.testing.assert_close(fn(t), result) + torch.testing.assert_close(result, t) + (gradient,) = torch.autograd.grad(result.sum(), (t,)) + torch.testing.assert_close(gradient, torch.ones_like(t)) + + +class TestDiagnostics: + @pytest.mark.parametrize( + "kwargs", + ( + {"sigma_data": 0}, + {"sigma_data": float("nan")}, + {"t_min": -0.1}, + {"t_max": 1.6}, + {"t_min": 0.7, "t_max": 0.5}, + ), + ) + def test_invalid_scheduler(self, kwargs): + with pytest.raises(ValueError): + schedules.TrigFlowNoiseScheduler(**kwargs) + + @pytest.mark.parametrize("sigma_data", (0, -1, float("nan"))) + def test_invalid_preconditioner(self, sigma_data): + with pytest.raises(ValueError): + preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data + )