diff --git a/src/BetaSchedulers.jl b/src/BetaSchedulers.jl index 8d98477..7571c3e 100644 --- a/src/BetaSchedulers.jl +++ b/src/BetaSchedulers.jl @@ -1,13 +1,29 @@ import NNlib: sigmoid +""" +Linear beta schedule. + +cf. https://arxiv.org/abs/2006.11239 +""" function linear_beta_schedule(num_timesteps::Int, β_start=0.0001f0, β_end=0.02f0) range(β_start, β_end, length=num_timesteps) end +""" +Scaled linear beta schedule. +(very specific to latent diffusion models) + +cf. https://arxiv.org/abs/2006.11239 +""" function scaled_linear_beta_schedule(num_timesteps::Int, β_start=0.0001f0, β_end=0.02f0) range(β_start^0.5, β_end^0.5, length=num_timesteps) .^ 2 end +""" +Cosine beta schedule. + +cf. https://arxiv.org/abs/2102.09672 +""" function cosine_beta_schedule(num_timesteps::Int, max_beta=0.999f0, ϵ=1e-3f0) α_bar(t) = cos((t + ϵ) / (1 + ϵ) * π / 2)^2 @@ -21,6 +37,12 @@ function cosine_beta_schedule(num_timesteps::Int, max_beta=0.999f0, ϵ=1e-3f0) return βs end +""" +Sigmoid beta schedule. + +cf. https://arxiv.org/abs/2203.02923 +and https://github.com/MinkaiXu/GeoDiff/blob/main/models/epsnet/diffusion.py#L34 +""" function sigmoid_beta_schedule(num_timesteps::Int, β_start=0.0001f0, β_end=0.02f0) x = range(-6, 6, length=num_timesteps) sigmoid(x) * (β_end - β_start) + β_start