dew.diffusion.schedules

The noise schedules, one module per family.

Name	Summary
ContinuousNoiseScheduler	A schedule whose time is a fraction of the trajectory, not an index.
CosineContinuousNoiseScheduler	alpha = cos(pi t / 2), sigma = sin(pi t / 2), weighted by sigma^2, which is 1 / (1 + SNR).
CosineGeneralNoiseScheduler	Sigmas placed so that log-SNR runs along a cosine in t, on the variance exploding form.
CosineNoiseScheduler	The cosine beta table of Nichol and Dhariwal 2021.
DiscreteNoiseScheduler	A variance preserving schedule tabulated from betas, DDPM style.
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).
ExpNoiseScheduler	A beta table whose cumulative alpha decays as exp(-12 t).
FlowMatchingScheduler	Rectified flow / conditional flow matching on the linear path.
GeneralizedNoiseScheduler	The variance exploding family of Karras et al.
KarrasVENoiseScheduler	Sigmas placed along t with the rho spacing of Karras et al.
LinearNoiseScheduler	The linear beta table of Ho et al.
NoiseScheduler	The forward process on [0, T], with t = T the fully noised end.
SqrtContinuousNoiseScheduler	Square-root schedule from Diffusion-LM (Li et al.
compute_resolution_shift	Flux-style resolution dependent timestep shift.
cosine_beta_schedule	Nichol and Dhariwal 2021, Eq.
exp_beta_schedule	Betas whose cumulative alpha decays as exp(-12 t), each clipped at beta_end.
expand	A per-example coefficient [B] shaped to broadcast against x [B, ...].
linear_beta_schedule	Ho et al.

ContinuousNoiseScheduler

class source

class ContinuousNoiseScheduler(NoiseScheduler)

A schedule whose time is a fraction of the trajectory, not an index.

T is 1.0, so t = 1 is fully noised, and training draws t uniformly. A subclass gives the rates and the weight of its parameterization.

ContinuousNoiseScheduler.sample_t

def sample_t(key, n)

CosineContinuousNoiseScheduler

class source

class CosineContinuousNoiseScheduler(ContinuousNoiseScheduler)

alpha = cos(pi t / 2), sigma = sin(pi t / 2), weighted by sigma^2, which is 1 / (1 + SNR).

CosineContinuousNoiseScheduler.rates

def rates(t)

CosineContinuousNoiseScheduler.weight

def weight(t)

CosineGeneralNoiseScheduler

class source

class CosineGeneralNoiseScheduler(
    sigma_min: float = 0.02,
    sigma_max: float = 80.0,
    kappa: float = 1.0,
    sigma_data: float = 0.5,
)

Sigmas placed so that log-SNR runs along a cosine in t, on the variance exploding form.

CosineGeneralNoiseScheduler.sigmas

def sigmas(t)

CosineGeneralNoiseScheduler.t_of_sigma

def t_of_sigma(sigma)

CosineNoiseScheduler

class source

class CosineNoiseScheduler(
    timesteps: int,
    beta_start: float = 0.008,
    beta_end: float = 0.999,
    p2_loss_weight_k: float = 1,
    p2_loss_weight_gamma: float = 1,
)

The cosine beta table of Nichol and Dhariwal 2021.

DiscreteNoiseScheduler

class source

class DiscreteNoiseScheduler(
    betas: np.ndarray,
    p2_loss_weight_k: float = 1,
    p2_loss_weight_gamma: float = 1,
)

A variance preserving schedule tabulated from betas, DDPM style.

signal_rate^2 + noise_rate^2 = 1 at every index, and t is the index into the table, so T is the number of entries. The loss weight is the P2 weight of Choi et al. 2022, (k + SNR)^-gamma. At the defaults k = 1, gamma = 1 it is 1 / (1 + SNR), which on a v-prediction loss (whose error is 1 + SNR times the x_0 error) is exactly an unweighted x_0 loss.

DiscreteNoiseScheduler.index

def index(t) -> jax.Array

t as a table index; a time grid may reach T itself, which is the last entry.

DiscreteNoiseScheduler.rates

def rates(t)

DiscreteNoiseScheduler.sample_t

def sample_t(key, n)

DiscreteNoiseScheduler.weight

def weight(t)

EDMNoiseScheduler

class source

class EDMNoiseScheduler(
    sigma_min: float = 0.002,
    sigma_max: float = 80.0,
    sigma_data: float = 0.5,
    P_mean: float = -0.4,
    P_std: float = 1.0,
)

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).

Defaults are EDM2’s (Karras et al. 2024); EDM1’s -1.2/1.2 concentrated too much mass on low noise levels for larger models. Pass them explicitly to reproduce an EDM1 run.

EDMNoiseScheduler.sigmas

def sigmas(t)

EDMNoiseScheduler.t_of_sigma

def t_of_sigma(sigma)

EDMNoiseScheduler.sample_t

def sample_t(key, n)

ExpNoiseScheduler

class source

class ExpNoiseScheduler(
    timesteps: int,
    beta_end: float = 0.999,
    p2_loss_weight_k: float = 1,
    p2_loss_weight_gamma: float = 1,
)

A beta table whose cumulative alpha decays as exp(-12 t).

FlowMatchingScheduler

class source

class FlowMatchingScheduler(
    shift: float = 1.0,
    logit_mean: float = 0.0,
    logit_std: float = 1.0,
)

Rectified flow / conditional flow matching on the linear path.

x_t = (1 - t) * x_0 + t * epsilon for t in [0, 1], so alpha + sigma = 1 and the model input needs no scaling. Timesteps are drawn logit-normal as in SD3, which concentrates training on the middle of the trajectory where the velocity is hardest to predict.

FlowMatchingScheduler.shift_timesteps

def shift_timesteps(t) -> jax.Array

FlowMatchingScheduler.sample_t

def sample_t(key, n)

FlowMatchingScheduler.rates

def rates(t)

FlowMatchingScheduler.weight

def weight(t)

FlowMatchingScheduler.model_time

def model_time(t)

GeneralizedNoiseScheduler

class source

class GeneralizedNoiseScheduler(
    sigma_min: float = 0.002,
    sigma_max: float = 80.0,
    sigma_data: float = 0.5,
)

The variance exploding family of Karras et al. 2022 (“Elucidating the Design Space of Diffusion-Based Generative Models”).

alpha is 1, and the paired preconditioning scales the model input. Every member conditions the model on c_noise = log(sigma) / 4 and weights the loss with lambda(sigma) = (sigma^2 + sigma_data^2) / (sigma sigma_data)^2 (Eq. 8 of the paper), written in a form that needs no epsilon guard; a subclass places the sigmas along t and inverts that placement for the solvers that step in sigma.

GeneralizedNoiseScheduler.sigmas

def sigmas(t) -> jax.Array

The noise level at t.

GeneralizedNoiseScheduler.t_of_sigma

def t_of_sigma(sigma) -> jax.Array

The inverse of sigmas.

GeneralizedNoiseScheduler.rates

def rates(t)

GeneralizedNoiseScheduler.sample_t

def sample_t(key, n)

GeneralizedNoiseScheduler.weight

def weight(t)

GeneralizedNoiseScheduler.model_time

def model_time(t)

KarrasVENoiseScheduler

class source

class KarrasVENoiseScheduler(
    sigma_min: float = 0.002,
    sigma_max: float = 80.0,
    rho: float = 7.0,
    sigma_data: float = 0.5,
)

Sigmas placed along t with the rho spacing of Karras et al. 2022 (Eq. 5): sigma(t) = (sigma_max^(1/rho) + (1 - t) (sigma_min^(1/rho) - sigma_max^(1/rho)))^rho, so a uniform grid in t is the paper’s sampling grid in sigma.

KarrasVENoiseScheduler.sigmas

def sigmas(t)

KarrasVENoiseScheduler.t_of_sigma

def t_of_sigma(sigma)

LinearNoiseScheduler

class source

class LinearNoiseScheduler(
    timesteps: int,
    beta_start: float = 0.0001,
    beta_end: float = 0.02,
    p2_loss_weight_k: float = 1,
    p2_loss_weight_gamma: float = 1,
)

The linear beta table of Ho et al. 2020, scaled to the step count.

NoiseScheduler

class source

class NoiseScheduler(ABC)

The forward process on [0, T], with t = T the fully noised end.

NoiseScheduler.rates

def rates(t) -> tuple[jax.Array, jax.Array]

(alpha, sigma) at t, shaped like t.

NoiseScheduler.sample_t

def sample_t(key, n: int) -> jax.Array

n training times drawn the way this schedule trains.

NoiseScheduler.weight

def weight(t) -> jax.Array

The schedule’s own loss weight at t, in the space its paired parameterization computes the loss in.

NoiseScheduler.model_time

def model_time(t) -> jax.Array

What the model is conditioned on at t; the time itself unless the schedule says otherwise.

NoiseScheduler.prior_scale

def prior_scale() -> jax.Array

Standard deviation of the initial Gaussian draw.

NoiseScheduler.step_interval

def step_interval(t, t_next) -> jax.Array

The transfer interval; ordinary grids advance to their next point.

NoiseScheduler.half_interval

def half_interval(t, t_next) -> jax.Array

Half a grid interval, as a Runge-Kutta stage places its intermediate points.

NoiseScheduler.snr

def snr(t) -> jax.Array

SqrtContinuousNoiseScheduler

class source

class SqrtContinuousNoiseScheduler(ContinuousNoiseScheduler)

Square-root schedule from Diffusion-LM (Li et al. 2022).

alpha(t) = sqrt(1 - t) and sigma(t) = sqrt(t) for t in [0, 1], so it is variance preserving (alpha^2 + sigma^2 = 1) with SNR(t) = (1 - t) / t. Noise ramps up much faster near t = 0 than in the cosine schedule. The low-noise end carries little signal about the token identity for discrete/embedding data, so the schedule spends fewer steps there. The paper trains the plain x_0 loss, so the weight is one.

SqrtContinuousNoiseScheduler.rates

def rates(t)

SqrtContinuousNoiseScheduler.weight

def weight(t)

compute_resolution_shift

function source

def compute_resolution_shift(
    sequence_length,
    base_seq_len=256,
    max_seq_len=4096,
    base_shift=0.5,
    max_shift=1.15,
) -> float

Flux-style resolution dependent timestep shift.

Longer token sequences carry more redundancy, so the trajectory has to spend more of its budget at high noise for the global structure to settle. mu is interpolated linearly in sequence length and the shift is exp(mu).

cosine_beta_schedule

function source

def cosine_beta_schedule(timesteps, start_angle=0.008, end_angle=0.999)

Nichol and Dhariwal 2021, Eq. 17: the cumulative alpha follows cos^2((t / T + s) / (1 + s) * pi / 2), s = start_angle, and each beta is clipped at end_angle.

exp_beta_schedule

function source

def exp_beta_schedule(timesteps, beta_end=0.999)

Betas whose cumulative alpha decays as exp(-12 t), each clipped at beta_end.

expand

function source

def expand(coefficient, x)

A per-example coefficient [B] shaped to broadcast against x [B, ...].

linear_beta_schedule

function source

def linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Ho et al. 2020’s betas, linear from beta_start to beta_end over 1000 steps and scaled so another step count keeps the same cumulative alpha.