From 4dfa3c92dc7e4c7bb0271d2eee332cff35cfe13f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Laure=CE=B7t?= Date: Sat, 29 Jul 2023 15:40:43 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=9A=20(BetaSchedulers)=20move=20beta?= =?UTF-8?q?=20variance=20schedules=20in=20their=20own=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 1 - examples/swissroll.jl | 5 +- src/BetaSchedulers.jl | 125 ----------------------------- src/BetaSchedules/BetaSchedules.jl | 21 +++++ src/BetaSchedules/Cosine.jl | 30 +++++++ src/BetaSchedules/Linear.jl | 17 ++++ src/BetaSchedules/ScaledLinear.jl | 17 ++++ src/BetaSchedules/Sigmoid.jl | 21 +++++ src/BetaSchedules/ZeroSNR.jl | 36 +++++++++ src/Diffusers.jl | 3 +- 10 files changed, 147 insertions(+), 129 deletions(-) delete mode 100644 src/BetaSchedulers.jl create mode 100644 src/BetaSchedules/BetaSchedules.jl create mode 100644 src/BetaSchedules/Cosine.jl create mode 100644 src/BetaSchedules/Linear.jl create mode 100644 src/BetaSchedules/ScaledLinear.jl create mode 100644 src/BetaSchedules/Sigmoid.jl create mode 100644 src/BetaSchedules/ZeroSNR.jl diff --git a/README.md b/README.md index b58772d..e7c46e7 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,6 @@ [🤗 diffusers](https://github.com/huggingface/diffusers/), but in Julia (and with worse code) ! - ## Credits - [@huggingface](https://github.com/huggingface) for their amazing [🤗 diffusers](https://github.com/huggingface/diffusers/) library. diff --git a/examples/swissroll.jl b/examples/swissroll.jl index 0e8d7ee..046c8a7 100644 --- a/examples/swissroll.jl +++ b/examples/swissroll.jl @@ -1,4 +1,5 @@ import Diffusers +import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr using Flux using Random using Plots @@ -42,8 +43,8 @@ scatter(data[1, :], data[2, :], num_timesteps = 100 scheduler = Diffusers.DDPM( Vector{Float64}, - Diffusers.rescale_zero_terminal_snr( - Diffusers.cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0) + rescale_zero_terminal_snr( + cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0) ) ) diff --git a/src/BetaSchedulers.jl b/src/BetaSchedulers.jl deleted file mode 100644 index e47a150..0000000 --- a/src/BetaSchedulers.jl +++ /dev/null @@ -1,125 +0,0 @@ -import NNlib: sigmoid - -""" -Linear beta schedule. - -## Input - * `T::Integer`: number of timesteps - * `β₁::Real=0.0001f0`: initial (t=1) value of β - * `β₋₁::Real=0.02f0`: final (t=T) value of β - -## Output - * `β::Vector{Real}`: βₜ values at each timestep t - -## References - * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) -""" -function linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) - return range(start=β₁, stop=β₋₁, length=T) -end - -""" -Scaled linear beta schedule. - -## Input - * `T::Int`: number of timesteps - * `β₁::Real=0.0001f0`: initial value of β - * `β₋₁::Real=0.02f0`: final value of β - -## Output - * `β::Vector{Real}`: βₜ values at each timestep t - -## References - * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) -""" -function scaled_linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) - return range(start=β₁^0.5, stop=β₋₁^0.5, length=T) .^ 2 -end - -""" -Sigmoid beta schedule. - -## Input - * `T::Int`: number of timesteps - * `β₁::Real=0.0001f0`: initial value of β - * `β₋₁::Real=0.02f0`: final value of β - -## Output - * `β::Vector{Real}`: βₜ values at each timestep t - -## References - * [[2203.02923] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://arxiv.org/abs/2203.02923) - * [github.com:MinkaiXu/GeoDiff](https://github.com/MinkaiXu/GeoDiff/blob/ea0ca48045a2f7abfccd7f0df449e45eb6eae638/models/epsnet/diffusion.py#L57) -""" -function sigmoid_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) - x = range(start=-6, stop=6, length=T) - return sigmoid(x) .* (β₋₁ - β₁) .+ β₁ -end - -""" -Cosine beta schedule. - -## Input - * `T::Int`: number of timesteps - * `βₘₐₓ::Real=0.999f0`: maximum value of β - * `ϵ::Real=1e-3f0`: small value used to avoid division by zero - -## Output - * `β::Vector{Real}`: βₜ values at each timestep t - -## References - * [[2102.09672] Improved Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2102.09672) - * [github:openai/improved-diffusion](https://github.com/openai/improved-diffusion/blob/783b6740edb79fdb7d063250db2c51cc9545dcd1/improved_diffusion/gaussian_diffusion.py#L36) -""" -function cosine_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0, ϵ::Real=0.001f0) - α̅(t) = cos((t / T + ϵ) / (1 + ϵ) * π / 2)^2 - - β = Vector{Real}(undef, T) - for t in 1:T - αₜ = α̅(t) / α̅(t - 1) - - βₜ = 1 - αₜ - βₜ = min(βₘₐₓ, βₜ) - - β[t] = βₜ - end - - return β -end - -""" -Rescale betas to have zero terminal Signal to Noise Ratio (SNR). - -## Input - * `β::AbstractArray`: βₜ values at each timestep t - -## Output - * `β::Vector{Real}`: rescaled βₜ values at each timestep t - -## References - * [[2305.08891] Rescaling Diffusion Models](https://arxiv.org/abs/2305.08891) (Alg. 1) -""" -function rescale_zero_terminal_snr(β::AbstractArray) - # convert β to ⎷α̅ - α = 1 .- β - α̅ = cumprod(α) - ⎷α̅ = sqrt.(α̅) - - # store old extrema values - ⎷α̅₁ = ⎷α̅[1] - ⎷α̅₋₁ = ⎷α̅[end] - - # shift last timestep to zero - ⎷α̅ .-= ⎷α̅₋₁ - - # scale so that first timestep reaches old values - ⎷α̅ *= ⎷α̅₁ / (⎷α̅₁ - ⎷α̅₋₁) - - # convert back ⎷α̅ to β - α̅ = ⎷α̅ .^ 2 - α = α̅[2:end] ./ α̅[1:end-1] - α = vcat(α̅[1], α) - β = 1 .- α - - return β -end diff --git a/src/BetaSchedules/BetaSchedules.jl b/src/BetaSchedules/BetaSchedules.jl new file mode 100644 index 0000000..976f21d --- /dev/null +++ b/src/BetaSchedules/BetaSchedules.jl @@ -0,0 +1,21 @@ +module BetaSchedules + +# variance schedulers +include("Linear.jl") +include("ScaledLinear.jl") +include("Cosine.jl") +include("Sigmoid.jl") + +export + linear_beta_schedule, + scaled_linear_beta_schedule, + cosine_beta_schedule, + sigmoid_beta_schedule + +# utils +include("ZeroSNR.jl") + +export + rescale_zero_terminal_snr + +end # module BetaSchedules diff --git a/src/BetaSchedules/Cosine.jl b/src/BetaSchedules/Cosine.jl new file mode 100644 index 0000000..0e571b7 --- /dev/null +++ b/src/BetaSchedules/Cosine.jl @@ -0,0 +1,30 @@ +""" +Cosine beta schedule. + +## Input + * `T::Int`: number of timesteps + * `βₘₐₓ::Real=0.999f0`: maximum value of β + * `ϵ::Real=1e-3f0`: small value used to avoid division by zero + +## Output + * `β::Vector{Real}`: βₜ values at each timestep t + +## References + * [[2102.09672] Improved Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2102.09672) + * [github:openai/improved-diffusion](https://github.com/openai/improved-diffusion/blob/783b6740edb79fdb7d063250db2c51cc9545dcd1/improved_diffusion/gaussian_diffusion.py#L36) +""" +function cosine_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0, ϵ::Real=0.001f0) + α̅(t) = cos((t / T + ϵ) / (1 + ϵ) * π / 2)^2 + + β = Vector{Real}(undef, T) + for t in 1:T + αₜ = α̅(t) / α̅(t - 1) + + βₜ = 1 - αₜ + βₜ = min(βₘₐₓ, βₜ) + + β[t] = βₜ + end + + return β +end diff --git a/src/BetaSchedules/Linear.jl b/src/BetaSchedules/Linear.jl new file mode 100644 index 0000000..6c19e4a --- /dev/null +++ b/src/BetaSchedules/Linear.jl @@ -0,0 +1,17 @@ +""" +Linear beta schedule. + +## Input + * `T::Integer`: number of timesteps + * `β₁::Real=0.0001f0`: initial (t=1) value of β + * `β₋₁::Real=0.02f0`: final (t=T) value of β + +## Output + * `β::Vector{Real}`: βₜ values at each timestep t + +## References + * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) +""" +function linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) + return range(start=β₁, stop=β₋₁, length=T) +end diff --git a/src/BetaSchedules/ScaledLinear.jl b/src/BetaSchedules/ScaledLinear.jl new file mode 100644 index 0000000..4472cfc --- /dev/null +++ b/src/BetaSchedules/ScaledLinear.jl @@ -0,0 +1,17 @@ +""" +Scaled linear beta schedule. + +## Input + * `T::Int`: number of timesteps + * `β₁::Real=0.0001f0`: initial value of β + * `β₋₁::Real=0.02f0`: final value of β + +## Output + * `β::Vector{Real}`: βₜ values at each timestep t + +## References + * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) +""" +function scaled_linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) + return range(start=β₁^0.5, stop=β₋₁^0.5, length=T) .^ 2 +end diff --git a/src/BetaSchedules/Sigmoid.jl b/src/BetaSchedules/Sigmoid.jl new file mode 100644 index 0000000..277a9aa --- /dev/null +++ b/src/BetaSchedules/Sigmoid.jl @@ -0,0 +1,21 @@ +import NNlib: sigmoid + +""" +Sigmoid beta schedule. + +## Input + * `T::Int`: number of timesteps + * `β₁::Real=0.0001f0`: initial value of β + * `β₋₁::Real=0.02f0`: final value of β + +## Output + * `β::Vector{Real}`: βₜ values at each timestep t + +## References + * [[2203.02923] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://arxiv.org/abs/2203.02923) + * [github.com:MinkaiXu/GeoDiff](https://github.com/MinkaiXu/GeoDiff/blob/ea0ca48045a2f7abfccd7f0df449e45eb6eae638/models/epsnet/diffusion.py#L57) +""" +function sigmoid_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0) + x = range(start=-6, stop=6, length=T) + return sigmoid(x) .* (β₋₁ - β₁) .+ β₁ +end diff --git a/src/BetaSchedules/ZeroSNR.jl b/src/BetaSchedules/ZeroSNR.jl new file mode 100644 index 0000000..bae0ba3 --- /dev/null +++ b/src/BetaSchedules/ZeroSNR.jl @@ -0,0 +1,36 @@ +""" +Rescale betas to have zero terminal Signal to Noise Ratio (SNR). + +## Input + * `β::AbstractArray`: βₜ values at each timestep t + +## Output + * `β::Vector{Real}`: rescaled βₜ values at each timestep t + +## References + * [[2305.08891] Rescaling Diffusion Models](https://arxiv.org/abs/2305.08891) (Alg. 1) +""" +function rescale_zero_terminal_snr(β::AbstractArray) + # convert β to ⎷α̅ + α = 1 .- β + α̅ = cumprod(α) + ⎷α̅ = sqrt.(α̅) + + # store old extrema values + ⎷α̅₁ = ⎷α̅[1] + ⎷α̅₋₁ = ⎷α̅[end] + + # shift last timestep to zero + ⎷α̅ .-= ⎷α̅₋₁ + + # scale so that first timestep reaches old values + ⎷α̅ *= ⎷α̅₁ / (⎷α̅₁ - ⎷α̅₋₁) + + # convert back ⎷α̅ to β + α̅ = ⎷α̅ .^ 2 + α = α̅[2:end] ./ α̅[1:end-1] + α = vcat(α̅[1], α) + β = 1 .- α + + return β +end diff --git a/src/Diffusers.jl b/src/Diffusers.jl index 785e4fd..f4d683a 100644 --- a/src/Diffusers.jl +++ b/src/Diffusers.jl @@ -1,8 +1,9 @@ module Diffusers +include("BetaSchedules/BetaSchedules.jl") + # abtract types include("Schedulers.jl") -include("BetaSchedulers.jl") # concrete types include("DDPM.jl")