diff --git a/src/BetaSchedules/Cosine.jl b/src/BetaSchedules/Cosine.jl index 637511a..cf62260 100644 --- a/src/BetaSchedules/Cosine.jl +++ b/src/BetaSchedules/Cosine.jl @@ -1,6 +1,10 @@ """ Cosine beta schedule. +```math +\\overline{\\alpha}_t = \\cos \\left( \\frac{t / T + \\epsilon}{1 + \\epsilon} \\frac{\\pi}{2} \\right) +``` + ## Input * `T::Int`: number of timesteps * `βₘₐₓ::Real=0.999f0`: maximum value of β @@ -10,8 +14,8 @@ Cosine beta schedule. * `β::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) +* [nichol2021improved; Improved Denoising Diffusion Probabilistic Models](@cite) +* [github:openai/improved-diffusion/improved_diffusion/gaussian_diffusion.py](https://github.com/openai/improved-diffusion/blob/783b6740edb79fdb7d063250db2c51cc9545dcd1/improved_diffusion/gaussian_diffusion.py#L36) """ function cosine_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0, ϵ::Real=1.0f-3) α̅(t) = cos((t / T + ϵ) / (1 + ϵ) * π / 2)^2 diff --git a/src/BetaSchedules/Exponential.jl b/src/BetaSchedules/Exponential.jl index 3d32f4e..4f68080 100644 --- a/src/BetaSchedules/Exponential.jl +++ b/src/BetaSchedules/Exponential.jl @@ -1,14 +1,16 @@ """ Exponential beta schedule. +```math +\\overline{\\alpha}_t = \\exp \\left( \\frac{-12 t}{T} \\right) +``` + ## Input * `T::Int`: number of timesteps * `βₘₐₓ::Real=0.999f0`: maximum value of β ## Output * `β::Vector{Real}`: βₜ values at each timestep t - -## References """ function exponential_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0) α̅(t) = exp(-12 * t / T) diff --git a/src/BetaSchedules/Linear.jl b/src/BetaSchedules/Linear.jl index a77ed13..b42710d 100644 --- a/src/BetaSchedules/Linear.jl +++ b/src/BetaSchedules/Linear.jl @@ -1,6 +1,10 @@ """ Linear beta schedule. +```math +\\beta_t = \\beta_1 + \\frac{t - 1}{T - 1} (\\beta_{-1} - \\beta_1) +``` + ## Input * `T::Integer`: number of timesteps * `β₁::Real=1.0f-4`: initial (t=1) value of β @@ -10,7 +14,7 @@ Linear beta schedule. * `β::Vector{Real}`: βₜ values at each timestep t ## References - * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) + * [ho2020denoising; Denoising Diffusion Probabilistic Models](@cite) """ function linear_beta_schedule(T::Integer, β₁::Real=1.0f-4, β₋₁::Real=2.0f-2) return range(start=β₁, stop=β₋₁, length=T) diff --git a/src/BetaSchedules/ScaledLinear.jl b/src/BetaSchedules/ScaledLinear.jl index 7be4130..03b3442 100644 --- a/src/BetaSchedules/ScaledLinear.jl +++ b/src/BetaSchedules/ScaledLinear.jl @@ -1,6 +1,10 @@ """ Scaled linear beta schedule. +```math +\\beta_t = \\left( \\sqrt{\\beta_1} + \\frac{t - 1}{T - 1} \\left( \\sqrt{\\beta_{-1}} - \\sqrt{\\beta_1} \\right) \\right)^2 +``` + ## Input * `T::Int`: number of timesteps * `β₁::Real=1.0f-4`: initial value of β @@ -10,7 +14,7 @@ Scaled linear beta schedule. * `β::Vector{Real}`: βₜ values at each timestep t ## References - * [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) + * [ho2020denoising; Denoising Diffusion Probabilistic Models](@cite) """ function scaled_linear_beta_schedule(T::Integer, β₁::Real=1.0f-4, β₋₁::Real=2.0f-2) return range(start=√β₁, stop=√β₋₁, length=T) .^ 2 diff --git a/src/BetaSchedules/Sigmoid.jl b/src/BetaSchedules/Sigmoid.jl index 9e22fd3..84a92ce 100644 --- a/src/BetaSchedules/Sigmoid.jl +++ b/src/BetaSchedules/Sigmoid.jl @@ -3,6 +3,10 @@ import NNlib: sigmoid """ Sigmoid beta schedule. +```math +\\beta_t = \\sigma \\left( 12 \\frac{t - 1}{T - 1} - 6 \\right) ( \\beta_{-1} - \\beta_1 ) + \\beta_1 +``` + ## Input * `T::Int`: number of timesteps * `β₁::Real=1.0f-4`: initial value of β @@ -12,8 +16,8 @@ Sigmoid beta schedule. * `β::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) + * [xu2022geodiff; GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](@cite) + * [github.com:MinkaiXu/GeoDiff/models/epsnet/diffusion.py](https://github.com/MinkaiXu/GeoDiff/blob/ea0ca48045a2f7abfccd7f0df449e45eb6eae638/models/epsnet/diffusion.py#L57) """ function sigmoid_beta_schedule(T::Integer, β₁::Real=1.0f-4, β₋₁::Real=2.0f-2) x = range(start=-6, stop=6, length=T) diff --git a/src/BetaSchedules/ZeroSNR.jl b/src/BetaSchedules/ZeroSNR.jl index bae0ba3..81dbdb2 100644 --- a/src/BetaSchedules/ZeroSNR.jl +++ b/src/BetaSchedules/ZeroSNR.jl @@ -1,5 +1,5 @@ """ -Rescale betas to have zero terminal Signal to Noise Ratio (SNR). +Rescale β to have zero terminal Signal to Noise Ratio (SNR). ## Input * `β::AbstractArray`: βₜ values at each timestep t @@ -8,7 +8,7 @@ Rescale betas to have zero terminal Signal to Noise Ratio (SNR). * `β::Vector{Real}`: rescaled βₜ values at each timestep t ## References - * [[2305.08891] Rescaling Diffusion Models](https://arxiv.org/abs/2305.08891) (Alg. 1) + * [lin2023common; Rescaling Diffusion Models (Alg. 1)](@cite) """ function rescale_zero_terminal_snr(β::AbstractArray) # convert β to ⎷α̅