the magnitude preserving unet works best with inverse square root decay learning schedule

Author: lucidrains

denoising_diffusion_pytorch/__init__.py

@@ -8,4 +8,4 @@
 
 from denoising_diffusion_pytorch.denoising_diffusion_pytorch_1d import GaussianDiffusion1D, Unet1D, Trainer1D, Dataset1D
 
-from denoising_diffusion_pytorch.karras_unet import KarrasUnet
+from denoising_diffusion_pytorch.karras_unet import KarrasUnet, InvSqrtDecayLRSched

denoising_diffusion_pytorch/karras_unet.py

@@ -9,6 +9,7 @@
 import torch
 from torch import nn, einsum
 from torch.nn import Module, ModuleList
+from torch.optim.lr_scheduler import LambdaLR
 import torch.nn.functional as F
 
 from einops import rearrange, repeat, pack, unpack
@@ -680,6 +681,21 @@ def forward(self, x):
 
         return x
 
+# works best with inverse square root decay schedule
+
+def InvSqrtDecayLRSched(
+    optimizer,
+    t_ref = 70000,
+    sigma_ref = 0.01
+):
+    """
+    refer to equation 67 and Table1
+    """
+    def inv_sqrt_decay_fn(step: int):
+        return sigma_ref / sqrt(max(t / t_ref, 1.))
+
+    return LambdaLR(optimizer, lr_lambda = inv_sqrt_decay_fn)
+
 # example
 
 if __name__ == '__main__':

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.10.5'
+__version__ = '1.10.6'