📝 (BetaSchedules) rework references in docstrings + add equations

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

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)

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)

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

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)

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 ⎷α̅