From cfd090b6e267f66580ed00634f9e33e84596adec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Laure=CE=B7t?= <laurent@fainsin.bzh>
Date: Tue, 25 Jul 2023 20:07:26 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20add=20docstrings=20to=20BetaSche?=
 =?UTF-8?q?dulers.jl?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/BetaSchedulers.jl | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

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