dew.diffusion

Forward processes, noise schedules and the parameterizations over them.

Process pairs a schedule with a prediction transform, which is what a run’s objective and every sampler read. presets holds the named combinations, and discrete is the masked-token process.

Name	Summary
ConsistencyBoundary	The boundary parameterization of a latent consistency model, as Diffusers’ LCMScheduler reads it (Luo et al.
ContinuousNoiseScheduler	A schedule whose time is a fraction of the trajectory, not an index. Documented in dew.diffusion.schedules.
CosineContinuousNoiseScheduler	alpha = cos(pi t / 2), sigma = sin(pi t / 2), weighted by sigma^2, which is 1 / (1 + SNR). Documented in dew.diffusion.schedules.
CosineGeneralNoiseScheduler	Sigmas placed so that log-SNR runs along a cosine in t, on the variance exploding form. Documented in dew.diffusion.schedules.
CosineNoiseScheduler	The cosine beta table of Nichol and Dhariwal 2021. Documented in dew.diffusion.schedules.
Denoiser	Denoises with one model, its parameters and its conditions. Documented in dew.diffusion.process.
DirectPredictionTransform
DiscreteNoiseScheduler	A variance preserving schedule tabulated from betas, DDPM style. Documented in dew.diffusion.schedules.
EDMNoiseScheduler	Training sigmas drawn from exp(N(P_mean, P_std^2)): t is the standard normal draw and sigma(t) = exp(P_mean + P_std t). Documented in dew.diffusion.schedules.
EpsilonPredictionTransform
ExpNoiseScheduler	A beta table whose cumulative alpha decays as exp(-12 t). Documented in dew.diffusion.schedules.
FlowMatchPredictionTransform	The model predicts the rectified flow velocity u = epsilon - x_0.
FlowMatchingScheduler	Rectified flow / conditional flow matching on the linear path. Documented in dew.diffusion.schedules.
GeneralizedNoiseScheduler	The variance exploding family of Karras et al. Documented in dew.diffusion.schedules.
KarrasPredictionTransform	The EDM preconditioning of Karras et al.
KarrasVENoiseScheduler	Sigmas placed along t with the rho spacing of Karras et al. Documented in dew.diffusion.schedules.
LinearNoiseScheduler	The linear beta table of Ho et al. Documented in dew.diffusion.schedules.
MinSNR	Weights the loss with min-SNR-gamma (Hang et al.
NoiseScheduler	The forward process on [0, T], with t = T the fully noised end. Documented in dew.diffusion.schedules.
PredictionTransform	What the model predicts, and how x_0 and epsilon are read back out.
Process	Pairs a schedule with what the model predicts on it and how the loss is weighted. Documented in dew.diffusion.process.
ScheduleWeighting	Weights the loss with the schedule’s own weight.
SourceLimitedPrediction	Limits x_0 the way a published scheduler’s step does.
SqrtContinuousNoiseScheduler	Square-root schedule from Diffusion-LM (Li et al. Documented in dew.diffusion.schedules.
VPredictionTransform	v = alpha eps - sigma x_0, normalized by the total variance.
Weighting	Weights a per-example loss, given the schedule and what it predicts.
broadcast_rates	The schedule’s rates at t, shaped to broadcast against x.
compute_resolution_shift	Flux-style resolution dependent timestep shift. Documented in dew.diffusion.schedules.
cosine_beta_schedule	Nichol and Dhariwal 2021, Eq. Documented in dew.diffusion.schedules.
discrete	Masked (absorbing-state) discrete diffusion, shaped like the Gaussian one.
exp_beta_schedule	Betas whose cumulative alpha decays as exp(-12 t), each clipped at beta_end. Documented in dew.diffusion.schedules.
expand	A per-example coefficient [B] shaped to broadcast against x [B, ...]. Documented in dew.diffusion.schedules.
linear_beta_schedule	Ho et al. Documented in dew.diffusion.schedules.
presets	Documented in dew.registry.

ConsistencyBoundary

class source

class ConsistencyBoundary(
    inner: PredictionTransform,
    timestep_scaling: float = 10.0,
    sigma_data: float = 0.5,
)

The boundary parameterization of a latent consistency model, as Diffusers’ LCMScheduler reads it (Luo et al. 2023, arXiv 2310.04378).

The model predicts x_0 in inner’s space, and the consistency function is f = c_skip x_t + c_out x_0, with c_skip = sigma_data^2 / (s^2 + sigma_data^2) and c_out = s / sqrt(s^2 + sigma_data^2) at the scaled time s = timestep_scaling t. f is therefore x_t itself at t = 0, and x_0 and epsilon are read out of f.

ConsistencyBoundary.pred_transform

def pred_transform(x_t, prediction, rates, t)

ConsistencyBoundary.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

ConsistencyBoundary.get_input_scale

def get_input_scale(rates)

DirectPredictionTransform

class source

class DirectPredictionTransform(PredictionTransform)

DirectPredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

EpsilonPredictionTransform

class source

class EpsilonPredictionTransform(PredictionTransform)

EpsilonPredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

EpsilonPredictionTransform.get_target

def get_target(x_0, epsilon, rates)

EpsilonPredictionTransform.target_error_scale

def target_error_scale(snr)

FlowMatchPredictionTransform

class source

class FlowMatchPredictionTransform(PredictionTransform)

The model predicts the rectified flow velocity u = epsilon - x_0.

That is the constant velocity of the linear path, so both endpoints are one step away.

FlowMatchPredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

FlowMatchPredictionTransform.get_target

def get_target(x_0, epsilon, rates)

FlowMatchPredictionTransform.target_error_scale

def target_error_scale(snr)

KarrasPredictionTransform

class source

class KarrasPredictionTransform(sigma_data: float = 0.5, *, velocity: bool = False)

The EDM preconditioning of Karras et al. 2022, Table 1.

The model sees c_in x_t and its raw output F is read as x_0 = c_skip x_t + c_out F. Every denominator is at least sigma_data, so none needs a guard.

velocity is Diffusers 0.34.0’s EDM prediction_type="v_prediction", whose precondition_outputs negates c_out, so the model’s output is the velocity of the preconditioned path rather than its endpoint offset.

KarrasPredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

KarrasPredictionTransform.pred_transform

def pred_transform(x_t, prediction, rates, t)

KarrasPredictionTransform.get_input_scale

def get_input_scale(rates)

KarrasPredictionTransform.target_error_scale

def target_error_scale(snr)

MinSNR

dataclass source

class MinSNR(gamma: float)

Weights the loss with min-SNR-gamma (Hang et al. 2023).

min(SNR, gamma) on the x_0 loss, converted into the space the model trains in. It replaces the schedule’s own weight.

PredictionTransform

class source

class PredictionTransform(*, normalize_input: bool = False)

What the model predicts, and how x_0 and epsilon are read back out.

The base supplies the x_0 target and the identity output transform. A subclass gives backward_diffusion, without which the parameterization is incomplete.

PredictionTransform.pred_transform

def pred_transform(x_t, prediction, rates, t) -> jax.Array

The model’s raw output at (x_t, t) as a prediction in target space.

PredictionTransform.forward_diffusion

def forward_diffusion(x_0, epsilon, rates) -> tuple[jax.Array, ArrayLike, jax.Array]

(x_t, c_in, target): the noised sample, the model input scale, and what the model should output for it.

PredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates) -> tuple[jax.Array, jax.Array]

(x_0, epsilon) read out of a prediction in target space.

PredictionTransform.get_target

def get_target(x_0, epsilon, rates) -> jax.Array

PredictionTransform.get_input_scale

def get_input_scale(rates) -> ArrayLike

PredictionTransform.target_error_scale

def target_error_scale(snr) -> ArrayLike

||target error||^2 / ||x_0 error||^2 at the given SNR.

min-SNR-gamma and the other loss weights are defined on the x_0 loss. Dividing by this converts them into the space the model trains in.

ScheduleWeighting

dataclass source

class ScheduleWeighting()

Weights the loss with the schedule’s own weight.

SourceLimitedPrediction

class source

class SourceLimitedPrediction(
    inner: PredictionTransform,
    *,
    clip: float | None = None,
    threshold: tuple[float, float] | None = None,
    recompute_epsilon: bool = True,
)

Limits x_0 the way a published scheduler’s step does.

inner reads x_0 out of the model’s output, and then either dynamic thresholding or a plain clamp to clip limits it. Thresholding is Saharia et al. 2022: clamp each sample to its own ratio quantile of |x_0|, never below 1 and never above maximum, then divide by it. Thresholding wins where a source declares both, the way its step tests them.

recompute_epsilon is whether the source re-derives epsilon from the limited x_0. DDPM’s posterior, DEIS and the noise-prediction DPM-Solver algorithms do, so their update carries the limit. DDIM keeps the model’s own output as its epsilon, and only its x_0 term is limited.

The limit is not linear in the model’s output, so it belongs to the conversion a guided walk runs once on the combined output rather than to each guidance branch.

SourceLimitedPrediction.pred_transform

def pred_transform(x_t, prediction, rates, t)

SourceLimitedPrediction.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

SourceLimitedPrediction.get_target

def get_target(x_0, epsilon, rates)

SourceLimitedPrediction.get_input_scale

def get_input_scale(rates)

SourceLimitedPrediction.target_error_scale

def target_error_scale(snr)

VPredictionTransform

class source

class VPredictionTransform(PredictionTransform)

v = alpha eps - sigma x_0, normalized by the total variance.

VPredictionTransform.backward_diffusion

def backward_diffusion(x_t, prediction, rates)

VPredictionTransform.get_target

def get_target(x_0, epsilon, rates)

VPredictionTransform.target_error_scale

def target_error_scale(snr)

Weighting

class source

class Weighting(Protocol)

Weights a per-example loss, given the schedule and what it predicts.

broadcast_rates

function source

def broadcast_rates(schedule: NoiseScheduler, t, x) -> tuple[jax.Array, jax.Array]

The schedule’s rates at t, shaped to broadcast against x.

discrete

module

Masked (absorbing-state) discrete diffusion, shaped like the Gaussian one.

The forward process replaces each token by a mask id independently, with a probability that grows along t in [0, 1]; MaskingSchedule.alpha(t) is the fraction of tokens still visible. Training is the continuous-time negative ELBO of MDLM (Sahoo et al. 2024, “Simple and Effective Masked Diffusion Language Models”): the cross entropy of the model’s prediction at the masked positions, weighted by -alpha’(t) / (1 - alpha(t)). Sampling reverses the process one interval at a time: a masked token is revealed with probability (alpha(s) - alpha(t)) / (1 - alpha(t)) and, when revealed, drawn from the model’s categorical, which is MDLM’s _ddpm_update.

DiscreteProcess has the surface dew.sampling.sample walks: a time grid, an initial state, and a denoiser whose two outputs are the model’s argmax fill of the masked positions and the log-probabilities the solver draws from, in the slots a Gaussian denoiser puts x_0 and epsilon.