# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# 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.
from abc import ABC, abstractmethod
from typing import Optional, Tuple, Type
import numpy as np
import torch
from nemo.collections.common.parts.utils import mask_sequence_tensor
from nemo.core.classes import NeuralModule, typecheck
from nemo.core.neural_types import FloatType, LengthsType, NeuralType, SpectrogramType, VoidType
from nemo.utils import logging
[docs]
class StochasticDifferentialEquation(NeuralModule, ABC):
"""Base class for stochastic differential equations."""
def __init__(self, time_min: float, time_max: float, num_steps: int):
super().__init__()
# min and max time
if time_min <= 0:
raise ValueError(f'time_min should be positive, current value {time_min}')
if time_max <= time_min:
raise ValueError(f'time_max should be larger than time_min, current max {time_max} and min {time_min}')
self.time_min = time_min
self.time_max = time_max
# number of steps
if num_steps <= 0:
raise ValueError(f'num_steps needs to be positive: current value {num_steps}')
self.num_steps = num_steps
@property
def dt(self) -> float:
"""Time step for this SDE.
This denotes the step size between `0` and `self.time_max` when using `self.num_steps`.
"""
return self.time_max / self.num_steps
@property
def time_delta(self) -> float:
"""Time range for this SDE."""
return self.time_max - self.time_min
[docs]
def generate_time(self, size: int, device: torch.device) -> torch.Tensor:
"""Generate random time steps in the valid range.
Time steps are generated between `self.time_min` and `self.time_max`.
Args:
size: number of samples
device: device to use
Returns:
A tensor of floats with shape (size,)
"""
time = torch.rand(size, device=device) * self.time_delta + self.time_min
return time
[docs]
@abstractmethod
def coefficients(self, state: torch.Tensor, time: torch.Tensor, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Args:
state: tensor of shape (B, C, D, T)
time: tensor of shape (B,)
Returns:
Tuple with drift and diffusion coefficients.
"""
pass
[docs]
@typecheck(
input_types={
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
},
output_types={
"sample": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
},
)
@abstractmethod
def prior_sampling(self, prior_mean: torch.Tensor) -> torch.Tensor:
"""Generate a sample from the prior distribution p_T.
Args:
prior_mean: Mean of the prior distribution
Returns:
A sample from the prior distribution.
"""
pass
[docs]
def discretize(
self, *, state: torch.Tensor, time: torch.Tensor, state_length: Optional[torch.Tensor] = None, **kwargs
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Assume we have the following SDE:
dx = drift(x, t) * dt + diffusion(x, t) * dwt
where `wt` is the standard Wiener process.
We assume the following discretization:
new_state = current_state + total_drift + total_diffusion * z_norm
where `z_norm` is sampled from normal distribution with zero mean and unit variance.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
state_length: length of the valid time steps for each example in the batch, shape (B,)
**kwargs: other parameters
Returns:
Drift and diffusion.
"""
# Get coefficients
drift_coefficient, diffusion_coefficient = self.coefficients(
state=state, time=time, state_length=state_length, **kwargs
)
# Discretized drift
drift = drift_coefficient * self.dt
# Note:
# Scale with sqrt(dt) because z_norm is sampled from a normal distribution with zero mean and
# unit variance and dwt is normally distributed with zero mean and variance dt
diffusion = diffusion_coefficient * np.sqrt(self.dt)
return drift, diffusion
[docs]
@abstractmethod
def copy(self):
"""Create a copy of this SDE."""
pass
def __repr__(self):
desc = f'{self.__class__.__name__}(time_min={self.time_min}, time_max={self.time_max}, num_steps={self.num_steps})'
desc += f'\n\tdt: {self.dt}'
desc += f'\n\ttime_delta: {self.time_delta}'
return desc
[docs]
class OrnsteinUhlenbeckVarianceExplodingSDE(StochasticDifferentialEquation):
"""This class implements the Ornstein-Uhlenbeck SDE with variance exploding noise schedule.
The SDE is given by:
dx = theta * (y - x) dt + g(t) dw
where `theta` is the stiffness parameter and `g(t)` is the diffusion coefficient:
g(t) = std_min * (std_max/std_min)^t * sqrt(2 * log(std_max/std_min))
References:
Richter et al., Speech Enhancement and Dereverberation with Diffusion-based Generative Models, Tr. ASLP 2023
"""
def __init__(
self,
stiffness: float,
std_min: float,
std_max: float,
num_steps: int = 100,
time_min: float = 3e-2,
time_max: float = 1.0,
eps: float = 1e-8,
):
super().__init__(time_min=time_min, time_max=time_max, num_steps=num_steps)
# Small regularization
if eps <= 0:
raise ValueError(f'eps should be positive, current value {eps}')
self.eps = eps
# stifness
self.stiffness = stiffness
# noise schedule
if std_min <= 0:
raise ValueError(f'std_min should be positive, current value {std_min}')
if std_max <= std_min:
raise ValueError(f'std_max should be larger than std_min, current max {std_max} and min {std_min}')
self.std_min = std_min
self.std_max = std_max
logging.debug('Initialized %s with', self.__class__.__name__)
logging.debug('\tstiffness: %s', self.stiffness)
logging.debug('\tstd_min: %s', self.std_min)
logging.debug('\tstd_max: %s', self.std_max)
logging.debug('\tnum_steps: %s', self.num_steps)
logging.debug('\ttime_min: %s', self.time_min)
logging.debug('\ttime_max: %s', self.time_max)
logging.debug('\teps: %s', self.eps)
@property
def std_ratio(self) -> float:
return self.std_max / (self.std_min + self.eps)
@property
def log_std_ratio(self) -> float:
return np.log(self.std_ratio + self.eps)
[docs]
@typecheck(
input_types={
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"time": NeuralType(tuple('B'), FloatType()),
},
output_types={
"mean": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
},
)
def perturb_kernel_mean(self, state: torch.Tensor, prior_mean: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
"""Return the mean of the perturbation kernel for this SDE.
Args:
state: current state of the process, shape (B, C, D, T)
prior_mean: mean of the prior distribution
time: current time of the process, shape (B,)
Returns:
A tensor of shape (B, C, D, T)
"""
# exponential weighting
weight = torch.exp(-self.stiffness * time)
# view as [B, C, D, T]
weight = weight.view(-1, 1, 1, 1)
# closed-form mean
mean = weight * state + (1 - weight) * prior_mean
return mean
[docs]
@typecheck(
input_types={
"time": NeuralType(tuple('B'), FloatType()),
},
output_types={
"std": NeuralType(tuple('B'), FloatType()),
},
)
def perturb_kernel_std(self, time: torch.Tensor) -> torch.Tensor:
"""Return the standard deviation of the perturbation kernel for this SDE.
Note that the standard deviation depends on the time and the noise schedule,
which is parametrized using `self.stiffness`, `self.std_min` and `self.std_max`.
Args:
time: current time of the process, shape (B,)
Returns:
A tensor of shape (B,)
"""
var = (self.std_min**2) * self.log_std_ratio
var *= torch.pow(self.std_ratio, 2 * time) - torch.exp(-2 * self.stiffness * time)
var /= self.stiffness + self.log_std_ratio
std = torch.sqrt(var)
return std
[docs]
@typecheck(
input_types={
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"time": NeuralType(tuple('B'), FloatType()),
},
output_types={
"mean": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
"std": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
},
)
def perturb_kernel_params(self, state: torch.Tensor, prior_mean: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
"""Return the mean and standard deviation of the perturbation kernel for this SDE.
Args:
state: current state of the process, shape (B, C, D, T)
prior_mean: mean of the prior distribution
time: current time of the process, shape (B,)
"""
assert torch.all(time <= self.time_max)
assert torch.all(time >= self.time_min)
# compute the mean
mean = self.perturb_kernel_mean(state=state, prior_mean=prior_mean, time=time)
# compute the standard deviation
std = self.perturb_kernel_std(time=time)
# view as [B, C, D, T]
std = std.view(-1, 1, 1, 1)
return mean, std
[docs]
@typecheck(
input_types={
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"time": NeuralType(tuple('B'), VoidType()),
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
},
output_types={
"drift_coefficient": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
"diffusion_coefficient": NeuralType(('B', 'C', 'D', 'T'), FloatType()),
},
)
def coefficients(
self,
state: torch.Tensor,
time: torch.Tensor,
prior_mean: torch.Tensor,
state_length: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute drift and diffusion coefficients for this SDE.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
prior_mean: mean of the prior distribution
state_length: length of the valid time steps for each example in the batch
Returns:
Drift and diffusion coefficients.
"""
# Drift coefficient
drift_coefficient = self.stiffness * (prior_mean - state)
# Diffusion coefficient
diffusion_coefficient = self.std_min * torch.pow(self.std_ratio, time) * np.sqrt(2 * self.log_std_ratio)
# View in the same shape as the state
diffusion_coefficient = diffusion_coefficient.view(-1, *([1] * (state.dim() - 1)))
if state_length is not None:
drift_coefficient = mask_sequence_tensor(drift_coefficient, state_length)
diffusion_coefficient = mask_sequence_tensor(diffusion_coefficient, state_length)
return drift_coefficient, diffusion_coefficient
[docs]
def prior_sampling(self, prior_mean: torch.Tensor) -> torch.Tensor:
"""Generate a sample from the prior distribution p_T.
Args:
prior_mean: Mean of the prior distribution
"""
# Final time step for all samples in the batch
time = self.time_max * torch.ones(prior_mean.shape[0], device=prior_mean.device)
# Compute the std of the prior distribution
std = self.perturb_kernel_std(time=time)
# view as [B, C, D, T]
std = std.view(-1, 1, 1, 1)
# Generate a sample from a normal distribution centered at prior_mean
sample = prior_mean + torch.randn_like(prior_mean) * std
return sample
[docs]
def copy(self):
return OrnsteinUhlenbeckVarianceExplodingSDE(
stiffness=self.stiffness,
std_min=self.std_min,
std_max=self.std_max,
num_steps=self.num_steps,
time_min=self.time_min,
time_max=self.time_max,
eps=self.eps,
)
def __repr__(self):
desc = f'{self.__class__.__name__}(stiffness={self.stiffness}, std_min={self.std_min}, std_max={self.std_max}, num_steps={self.num_steps}, time_min={self.time_min}, time_max={self.time_max}, eps={self.eps})'
desc += f'\n\tdt: {self.dt}'
desc += f'\n\ttime_delta: {self.time_delta}'
desc += f'\n\tstd_ratio: {self.std_ratio}'
desc += f'\n\tlog_std_ratio: {self.log_std_ratio}'
return desc
[docs]
class ReverseStochasticDifferentialEquation(StochasticDifferentialEquation):
def __init__(self, *, sde: Type[StochasticDifferentialEquation], score_estimator: Type[NeuralModule]):
"""Use the forward SDE and a score estimator to define the reverse SDE.
Args:
sde: forward SDE
score_estimator: neural score estimator
"""
super().__init__(time_min=sde.time_min, time_max=sde.time_max, num_steps=sde.num_steps)
self.score_estimator = score_estimator
self.forward_sde = sde
logging.debug('Initialized %s', self.__class__.__name__)
[docs]
def coefficients(
self,
state: torch.Tensor,
time: torch.Tensor,
score_condition: Optional[torch.Tensor] = None,
state_length: Optional[torch.Tensor] = None,
**kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute drift and diffusion coefficients for the reverse SDE.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
"""
raise NotImplementedError('Coefficients not necessary for the reverse SDE.')
[docs]
def prior_sampling(self, shape: torch.Size, device: torch.device) -> torch.Tensor:
"""Prior sampling is not necessary for the reverse SDE."""
raise NotImplementedError('Prior sampling not necessary for the reverse SDE.')
[docs]
def discretize(
self,
*,
state: torch.Tensor,
time: torch.Tensor,
score_condition: Optional[torch.Tensor] = None,
state_length: Optional[torch.Tensor] = None,
**kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Discretize the reverse SDE.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
score_condition: condition for the score estimator
state_length: length of the valid time steps for each example in the batch
**kwargs: other parameters for discretization of the forward SDE
"""
# Drift and diffusion from the forward SDE
forward_drift, forward_diffusion = self.forward_sde.discretize(state=state, time=time, **kwargs)
# For input for the score estimator:
# - if no condition is provided, use the state
# - if a condition is provided, concatenate the state and the condition along the channel dimension
score_input = state if score_condition is None else torch.cat([state, score_condition], dim=1)
# Estimate score
score, _ = self.score_estimator(input=score_input, input_length=state_length, condition=time)
# Adjust drift
drift = forward_drift - forward_diffusion.pow(2) * score
# Adjust diffusion
diffusion = forward_diffusion
if state_length is not None:
drift = mask_sequence_tensor(drift, state_length)
diffusion = mask_sequence_tensor(diffusion, state_length)
return drift, diffusion
[docs]
def copy(self):
return ReverseStochasticDifferentialEquation(sde=self.forward_sde.copy(), score_estimator=self.score_estimator)
def __repr__(self):
desc = f'{self.__class__.__name__}(sde={self.forward_sde}, score_estimator={self.score_estimator})'
return desc
[docs]
class PredictorCorrectorSampler(NeuralModule):
"""Predictor-Corrector sampler for the reverse SDE.
Args:
sde: forward SDE
score_estimator: neural score estimator
predictor: predictor for the reverse process
corrector: corrector for the reverse process
num_steps: number of time steps for the reverse process
num_corrector_steps: number of corrector steps
time_max: maximum time
time_min: minimum time
snr: SNR for Annealed Langevin Dynamics
output_type: type of the output ('state' for the final state, or 'mean' for the mean of the final state)
References:
- Song et al., Score-based generative modeling through stochastic differential equations, 2021
"""
def __init__(
self,
sde,
score_estimator,
predictor: str = 'reverse_diffusion',
corrector: str = 'annealed_langevin_dynamics',
num_steps: int = 50,
num_corrector_steps: int = 1,
time_max: Optional[float] = None,
time_min: Optional[float] = None,
snr: float = 0.5,
output_type: str = 'mean',
):
super().__init__()
# Create a copy of SDE
self.sde = sde.copy()
# Update SDE parameters for sampling
if time_max is not None:
self.sde.time_max = time_max
logging.info('sde.time_max set to: %s', self.sde.time_max)
if time_min is not None:
self.sde.time_min = time_min
logging.info('sde.time_min set to: %s', self.sde.time_min)
self.sde.num_steps = num_steps
logging.info('sde.num_steps set to: %s', self.sde.num_steps)
# Update local values
self.time_max = self.sde.time_max
self.time_min = self.sde.time_min
self.num_steps = self.sde.num_steps
# Predictor setup
if predictor == 'reverse_diffusion':
self.predictor = ReverseDiffusionPredictor(sde=self.sde, score_estimator=score_estimator)
else:
raise RuntimeError(f'Unexpected predictor: {predictor}')
# Corrector setup
if corrector == 'annealed_langevin_dynamics':
self.corrector = AnnealedLangevinDynamics(
sde=self.sde, score_estimator=score_estimator, snr=snr, num_steps=num_corrector_steps
)
else:
raise RuntimeError(f'Unexpected corrector: {corrector}')
if output_type not in ['mean', 'state']:
raise ValueError(f'Unexpected output type: {output_type}')
self.output_type = output_type
logging.debug('Initialized %s with', self.__class__.__name__)
logging.debug('\tpredictor: %s', predictor)
logging.debug('\tcorrector: %s', corrector)
logging.debug('\tnum_steps: %s', self.num_steps)
logging.debug('\ttime_min: %s', self.time_min)
logging.debug('\ttime_max: %s', self.time_max)
logging.debug('\tnum_corrector_steps: %s', num_corrector_steps)
logging.debug('\tsnr: %s', snr)
logging.debug('\toutput_type: %s', self.output_type)
[docs]
@typecheck(
input_types={
"prior_mean": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType()),
"score_condition": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType(), optional=True),
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
},
output_types={
"sample": NeuralType(('B', 'C', 'D', 'T'), SpectrogramType()),
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
},
)
@torch.inference_mode()
def forward(
self, prior_mean: torch.Tensor, score_condition: torch.Tensor, state_length: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""Takes prior (noisy) mean and generates a sample by solving the reverse SDE.
Args:
prior_mean: mean for the prior distribution, e.g., noisy observation
score_condition: conditioning for the score estimator
state_length: length of the valid time steps for each example in the batch
Returns:
Generated `sample` and the corresponding `sample_length`.
"""
# Sample from the prior distribution
state = self.sde.prior_sampling(prior_mean=prior_mean)
if state_length is not None:
state = mask_sequence_tensor(state, state_length)
# Time steps for evaluation
time_steps = torch.linspace(self.time_max, self.time_min, self.num_steps, device=state.device)
# Sampling
for t in time_steps:
# time steps for the whole batch
time = t * torch.ones(state.shape[0], device=state.device)
# corrector step
state, _ = self.corrector(
state=state, time=time, score_condition=score_condition, state_length=state_length
)
# predictor step
state, state_mean = self.predictor(
state=state,
time=time,
score_condition=score_condition,
prior_mean=prior_mean,
state_length=state_length,
)
# Final output
if self.output_type == 'state':
sample = state
elif self.output_type == 'mean':
sample = state_mean
else:
raise RuntimeError(f'Unexpected output type: {self.output_type}')
if state_length is not None:
sample = mask_sequence_tensor(sample, state_length)
return sample, state_length
[docs]
class Predictor(torch.nn.Module, ABC):
"""Predictor for the reverse process.
Args:
sde: forward SDE
score_estimator: neural score estimator
"""
def __init__(self, sde, score_estimator):
super().__init__()
self.reverse_sde = ReverseStochasticDifferentialEquation(sde=sde, score_estimator=score_estimator)
[docs]
@abstractmethod
@torch.inference_mode()
def forward(
self,
*,
state: torch.Tensor,
time: torch.Tensor,
score_condition: Optional[torch.Tensor] = None,
state_length: Optional[torch.Tensor] = None,
**kwargs,
):
"""Predict the next state of the reverse process.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
score_condition: conditioning for the score estimator
state_length: length of the valid time steps for each example in the batch
Returns:
New state and mean.
"""
pass
[docs]
class ReverseDiffusionPredictor(Predictor):
"""Predict the next state of the reverse process using the reverse diffusion process.
Args:
sde: forward SDE
score_estimator: neural score estimator
"""
def __init__(self, sde, score_estimator):
super().__init__(sde=sde, score_estimator=score_estimator)
[docs]
@torch.inference_mode()
def forward(self, *, state, time, score_condition=None, state_length=None, **kwargs):
"""Predict the next state of the reverse process using the reverse diffusion process.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
score_condition: conditioning for the score estimator
state_length: length of the valid time steps for each example in the batch
Returns:
New state and mean of the diffusion process.
"""
drift, diffusion = self.reverse_sde.discretize(
state=state, time=time, score_condition=score_condition, state_length=state_length, **kwargs
)
# Generate a random sample from a standard normal distribution
z_norm = torch.randn_like(state)
# Compute the mean of the next state
mean = state - drift
# Compute new state by sampling
new_state = mean + diffusion * z_norm
if state_length is not None:
new_state = mask_sequence_tensor(new_state, state_length)
mean = mask_sequence_tensor(mean, state_length)
return new_state, mean
[docs]
class Corrector(NeuralModule, ABC):
"""Corrector for the reverse process.
Args:
sde: forward SDE
score_estimator: neural score estimator
snr: SNR for Annealed Langevin Dynamics
num_steps: number of steps for the corrector
"""
def __init__(
self,
sde: Type[StochasticDifferentialEquation],
score_estimator: Type[NeuralModule],
snr: float,
num_steps: int,
):
super().__init__()
self.sde = sde
self.score_estimator = score_estimator
self.snr = snr
self.num_steps = num_steps
logging.debug('Initialized %s with', self.__class__.__name__)
logging.debug('\tsnr: %s', snr)
logging.debug('\tnum_steps: %s', num_steps)
[docs]
@abstractmethod
@typecheck(
input_types={
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
"time": NeuralType(tuple('B'), FloatType()),
"score_condition": NeuralType(('B', 'C', 'D', 'T'), VoidType(), optional=True),
"state_length": NeuralType(tuple('B'), LengthsType(), optional=True),
},
output_types={
"state": NeuralType(('B', 'C', 'D', 'T'), VoidType()),
},
)
@torch.inference_mode()
def forward(self, state, time, score_condition=None, state_length=None):
"""
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
score_condition: conditioning for the score estimator
state_length: length of the valid time steps for each example in the batch
Returns:
New state and mean.
"""
pass
[docs]
class AnnealedLangevinDynamics(Corrector):
"""Annealed Langevin Dynamics for the reverse process.
References:
- Song et al., Score-based generative modeling through stochastic differential equations, 2021
"""
def __init__(self, sde, **kwargs):
if not isinstance(sde, OrnsteinUhlenbeckVarianceExplodingSDE):
raise ValueError(f'Expected an instance of OrnsteinUhlenbeckVarianceExplodingSDE, got {type(sde)}')
super().__init__(sde=sde, **kwargs)
[docs]
@torch.inference_mode()
def forward(self, state, time, score_condition=None, state_length=None):
"""Correct the state using Annealed Langevin Dynamics.
Args:
state: current state of the process, shape (B, C, D, T)
time: current time of the process, shape (B,)
score_condition: conditioning for the score estimator
state_length: length of the valid time steps for each example in the batch
Returns:
New state and mean of the diffusion process.
References:
Alg. 4 in http://arxiv.org/abs/2011.13456
"""
# Compute the standard deviation of the diffusion process
std = self.sde.perturb_kernel_std(time=time)
# View as [B, 1, 1, 1]
std = std.view(-1, *([1] * (state.dim() - 1)))
for i in range(self.num_steps):
# prepare input for the score estimator, concatenate conditioning along the channel dimension
score_input = state if score_condition is None else torch.cat([state, score_condition], dim=1)
# calculate the score
score, _ = self.score_estimator(input=score_input, input_length=state_length, condition=time)
# generate a sample from a standard normal distribution
z_norm = torch.randn_like(state)
# compute the step size
# note: this is slightly different than in the paper, where std = ||z_norm||_2 / ||score||_2
step_size = 2 * (self.snr * std).pow(2)
# update the mean
mean = state + step_size * score
# update the state
state = mean + z_norm * torch.sqrt(step_size * 2)
if state_length is not None:
state = mask_sequence_tensor(state, state_length)
mean = mask_sequence_tensor(mean, state_length)
return state, mean
