# 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.
"""Second-order stochastic Heun sampler from the EDM paper."""
import math
from typing import Callable
import torch
from jaxtyping import Float
from torch import Tensor
from physicsnemo.diffusion.base import Denoiser
from .base import Solver
[docs]
class EDMStochasticHeunSolver(Solver):
r"""
Second-order stochastic Heun sampler from the EDM paper.
Implements stochastic sampling with configurable noise injection
controlled by the "churn" parameters, using a second-order Heun
correction step.
.. important::
This is **not** a true SDE solver. It performs ad-hoc noise injection
("churn") at each step to improve sample diversity, but the underlying
integration is still an ODE step. Therefore, the denoiser should return
the right-hand side of the **ODE**, not the SDE.
By default, noise injection is performed directly in time-step space.
For linear-Gaussian noise schedules where diffusion time and noise level
are not equal (e.g., VP schedule), provide ``sigma_fn`` and
``sigma_inv_fn`` to apply churn in noise-level space rather than
time-step space. Optionally provide ``diffusion_fn`` to control the
time-dependent magnitude of the injected noise.
.. code-block:: python
def sigma_fn(
t: Tensor, # shape: (B,) or broadcastable
) -> Tensor: ... # noise level, same shape as t
def sigma_inv_fn(
sigma: Tensor, # shape: (B,) or broadcastable
) -> Tensor: ... # diffusion time, same shape as sigma
def diffusion_fn(
x: Tensor, # shape: (B, *dims)
t: Tensor, # shape: (B,)
) -> Tensor: ... # g^2(x, t), broadcastable to shape of x
Parameters
----------
denoiser : Denoiser
A callable implementing the
:class:`~physicsnemo.diffusion.Denoiser` interface. Should
return the right-hand side of the **ODE** (not the SDE, since the
stochastic noise injection is handled internally by this solver).
Typically obtained via
:meth:`~physicsnemo.diffusion.noise_schedulers.NoiseScheduler.get_denoiser`
with ``denoising_type="ode"``.
alpha : float, optional
Interpolation parameter for the corrector step, must be in (0, 1].
``alpha=1`` gives the standard Heun method (trapezoidal rule),
``alpha=0.5`` gives the midpoint method. By default 1.
S_churn : float, optional
Controls the amount of noise added at each step. Higher values add
more stochasticity. By default 0 (deterministic), in which case this
solver is equivalent to the deterministic :class:`HeunSolver`.
S_min : float, optional
Minimum diffusion time (or noise level if ``sigma_fn`` and
``sigma_inv_fn`` are provided) for applying churn. By default 0.
S_max : float, optional
Maximum diffusion time (or noise level if ``sigma_fn`` and
``sigma_inv_fn`` are provided) for applying churn. By default
``float("inf")``.
S_noise : float, optional
Noise scaling factor. Large values add more noise to the latent state.
By default 1.
num_steps : int, optional
Total number of sampling steps, used to scale churn. By default 18.
sigma_fn : Callable[[Tensor], Tensor] | None, optional
Maps time to noise level :math:`\sigma(t)`. Useful for linear-Gaussian
schedules where :math:`\sigma(t) \neq t`. Typically
:meth:`~physicsnemo.diffusion.noise_schedulers.LinearGaussianNoiseScheduler.sigma`.
If provided, ``sigma_inv_fn`` must also be provided.
By default ``None`` (identity mapping).
sigma_inv_fn : Callable[[Tensor], Tensor] | None, optional
Maps noise level back to time. Typically
:meth:`~physicsnemo.diffusion.noise_schedulers.LinearGaussianNoiseScheduler.sigma_inv`.
If provided, ``sigma_fn`` must also be provided.
By default ``None`` (identity mapping).
diffusion_fn : Callable[[Tensor, Tensor], Tensor] | None, optional
Controls the time-dependent magnitude of the injected
noise, in addition of the ``S_noise`` scaling factor. Typically the
squared diffusion coefficient :math:`g^2(\mathbf{x}, t)` from the
reverse SDE, obtained from
:meth:`~physicsnemo.diffusion.noise_schedulers.LinearGaussianNoiseScheduler.diffusion`.
By default ``None`` (:math:`g^2 = 2t`), which corresponds to an
EDM-like noise schedule.
Note
----
Reference: `Elucidating the Design Space of Diffusion-Based
Generative Models <https://arxiv.org/abs/2206.00364>`_
Examples
--------
Basic usage with default parameters (noise injection in time-step space):
>>> import torch
>>> from physicsnemo.diffusion.samplers import (
... EDMStochasticHeunSolver,
... )
>>> denoiser = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2) # Toy denoiser
>>> solver = EDMStochasticHeunSolver(denoiser, S_churn=40, num_steps=18)
>>> x_t = torch.randn(1, 3, 8, 8)
>>> t_cur = torch.tensor([1.0])
>>> t_next = torch.tensor([0.5])
>>> x_tm1 = solver.step(x_t, t_cur, t_next)
>>> x_tm1.shape
torch.Size([1, 3, 8, 8])
Using noise scheduler methods for linear-Gaussian schedules where
:math:`\sigma(t) \neq t` (e.g., VP schedule). The callbacks map between
time and noise level, allowing the churn to be applied in noise-level
space before converting back to time-step space:
>>> from physicsnemo.diffusion.noise_schedulers import VPNoiseScheduler
>>> scheduler = VPNoiseScheduler()
>>> num_steps = 10
>>> solver = EDMStochasticHeunSolver(
... denoiser,
... S_churn=40,
... num_steps=num_steps,
... sigma_fn=scheduler.sigma,
... sigma_inv_fn=scheduler.sigma_inv,
... diffusion_fn=scheduler.diffusion,
... )
>>> x_tm1 = solver.step(x_t, t_cur, t_next)
>>> x_tm1.shape
torch.Size([1, 3, 8, 8])
"""
def __init__(
self,
denoiser: Denoiser,
alpha: float = 1.0,
S_churn: float = 0,
S_min: float = 0,
S_max: float = float("inf"),
S_noise: float = 1,
num_steps: int = 18,
sigma_fn: Callable[[Float[Tensor, " *shape"]], Float[Tensor, " *shape"]]
| None = None,
sigma_inv_fn: Callable[[Float[Tensor, " *shape"]], Float[Tensor, " *shape"]]
| None = None,
diffusion_fn: Callable[
[Float[Tensor, " B *dims"], Float[Tensor, " B"]], Float[Tensor, " B *_"]
]
| None = None,
) -> None:
self.denoiser = denoiser
if not 0 < alpha <= 1:
raise ValueError(f"alpha must be in (0, 1], got {alpha}")
self.alpha = alpha
self.S_churn = S_churn
self.S_min = S_min
self.S_max = S_max
self.S_noise = S_noise
self.num_steps = num_steps
# Validate sigma_fn and sigma_inv_fn
if (sigma_fn is None) != (sigma_inv_fn is None):
raise ValueError(
"sigma_fn and sigma_inv_fn must both be provided or both None."
)
if sigma_fn is None and sigma_inv_fn is None:
self.sigma_fn = lambda t: t
self.sigma_inv_fn = lambda sigma: sigma
self._use_noise_level_space = False
else:
self.sigma_fn = sigma_fn
self.sigma_inv_fn = sigma_inv_fn
self._use_noise_level_space = True
if diffusion_fn is None:
self.diffusion_fn = lambda x, t: 2 * t.reshape(-1, *([1] * (x.ndim - 1)))
else:
self.diffusion_fn = diffusion_fn
[docs]
def step(
self,
x: Float[Tensor, " B *dims"],
t_cur: Float[Tensor, " B"],
t_next: Float[Tensor, " B"],
) -> Float[Tensor, " B *dims"]:
r"""
Perform one stochastic Heun sampling step.
Parameters
----------
x : Tensor
Current noisy latent state :math:`\mathbf{x}_n` of shape
:math:`(B, *)` where :math:`B` is the batch size.
t_cur : Tensor
Current diffusion time :math:`t_n` of shape :math:`(B,)`.
t_next : Tensor
Target diffusion time :math:`t_{n-1}` of shape :math:`(B,)`.
Returns
-------
Tensor
Updated latent state :math:`\mathbf{x}_{n-1}` at time
``t_next``, same shape as ``x``.
"""
# Ensure contiguous strides so successive denoiser calls (across
# sampling steps) present the same stride layout to torch.compile,
# avoiding spurious recompilations / silently divergent traces.
t_cur = t_cur.contiguous()
t_next = t_next.contiguous()
# Reshape t for broadcasting: (B,) -> (B, 1, ..., 1)
expected_shape = (-1,) + (1,) * (x.ndim - 1)
t_cur_bc = t_cur.reshape(expected_shape)
t_next_bc = t_next.reshape(expected_shape)
gamma_base = min(self.S_churn / self.num_steps, math.sqrt(2) - 1)
# Compute perturbed time t_hat with increased noise
# NOTE: sigma_fn and sigma_inv_fn are identity if not provided (stays
# in time-step space). diffusion_fn defaults to g^2 = 2t (EDM-like
# noise schedule).
sigma_cur_bc = self.sigma_fn(t_cur_bc)
# Mask: apply churn only where S_min <= sigma <= S_max
churn_mask = (sigma_cur_bc >= self.S_min) & (sigma_cur_bc <= self.S_max)
gamma_bc = torch.where(churn_mask, gamma_base, 0.0)
sigma_hat_bc = sigma_cur_bc + gamma_bc * sigma_cur_bc
t_hat_bc = self.sigma_inv_fn(sigma_hat_bc)
# Noise scale: sqrt(sigma_hat^2 - sigma_cur^2) * S_noise * g(x,t) / sqrt(2*t)
g_sq_bc = self.diffusion_fn(x, t_cur)
safe_t_cur_bc = torch.where(t_cur_bc == 0, torch.ones_like(t_cur_bc), t_cur_bc)
noise_scale_bc = (
(sigma_hat_bc**2 - sigma_cur_bc**2).clamp(min=0).sqrt()
* self.S_noise
* (g_sq_bc / (2 * safe_t_cur_bc)).sqrt()
)
noise_scale_bc = torch.where(
t_cur_bc == 0, torch.zeros_like(noise_scale_bc), noise_scale_bc
)
# Perturb latent with noise
x_hat = x + noise_scale_bc * torch.randn_like(x)
# Euler step from t_hat to intermediate point (predictor)
t_hat = t_hat_bc.reshape(x.shape[0])
h_bc = t_next_bc - t_hat_bc
d_cur = self.denoiser(x_hat, t_hat)
t_prime_bc = t_hat_bc + self.alpha * h_bc
x_prime = x_hat + self.alpha * h_bc * d_cur
# Mask for elements where t_next != 0 (need 2nd order correction)
mask_bc = (t_next_bc != 0).float()
# Second RHS evaluation (compute everywhere, masked later)
t_prime = t_prime_bc.reshape(x.shape[0])
# Avoid issues by using t_hat where t_prime would be 0
t_prime_safe = torch.where(t_prime == 0, t_hat, t_prime)
d_prime = self.denoiser(x_prime, t_prime_safe)
# Apply 2nd order correction only where t_next != 0
w_cur = 1 - 1 / (2 * self.alpha)
w_prime = 1 / (2 * self.alpha)
x_euler = x_hat + h_bc * d_cur
x_heun = x_hat + h_bc * (w_cur * d_cur + w_prime * d_prime)
x_next = mask_bc * x_heun + (1 - mask_bc) * x_euler
return x_next
