just move towards using rmsnorm, given success of llama. and to avoid using pytorch layernorm, which has issues

Author: lucidrains

denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -104,26 +104,23 @@ def forward(self, x):
         weight = self.weight
         mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
         var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
-        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+        normalized_weight = (weight - mean) * var.clamp(min = eps).rsqrt()
 
         return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
 
     def forward(self, x):
-        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) * (var + eps).rsqrt() * self.g
+        return F.normalize(x, dim = 1) * self.g * (x.shape[1] ** 0.5)
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
         self.fn = fn
-        self.norm = LayerNorm(dim)
+        self.norm = RMSNorm(dim)
 
     def forward(self, x):
         x = self.norm(x)
@@ -220,7 +217,7 @@ def __init__(self, dim, heads = 4, dim_head = 32):
 
         self.to_out = nn.Sequential(
             nn.Conv2d(hidden_dim, dim, 1),
-            LayerNorm(dim)
+            RMSNorm(dim)
         )
 
     def forward(self, x):

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -109,26 +109,23 @@ def forward(self, x):
         weight = self.weight
         mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
         var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
-        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+        normalized_weight = (weight - mean) * var.clamp(min = eps).rsqrt()
 
         return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
 
     def forward(self, x):
-        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) * (var + eps).rsqrt() * self.g
+        return F.normalize(x, dim = 1) * self.g * (x.shape[1] ** 0.5)
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
         self.fn = fn
-        self.norm = LayerNorm(dim)
+        self.norm = RMSNorm(dim)
 
     def forward(self, x):
         x = self.norm(x)
@@ -223,7 +220,7 @@ def __init__(self, dim, heads = 4, dim_head = 32):
 
         self.to_out = nn.Sequential(
             nn.Conv2d(hidden_dim, dim, 1),
-            LayerNorm(dim)
+            RMSNorm(dim)
         )
 
     def forward(self, x):

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -96,26 +96,23 @@ def forward(self, x):
         weight = self.weight
         mean = reduce(weight, 'o ... -> o 1 1', 'mean')
         var = reduce(weight, 'o ... -> o 1 1', partial(torch.var, unbiased = False))
-        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+        normalized_weight = (weight - mean) * var.clamp(min = eps).rsqrt()
 
         return F.conv1d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.g = nn.Parameter(torch.ones(1, dim, 1))
 
     def forward(self, x):
-        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) * (var + eps).rsqrt() * self.g
+        return F.normalize(x, dim = 1) * self.g * (x.shape[1] ** 0.5)
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
         self.fn = fn
-        self.norm = LayerNorm(dim)
+        self.norm = RMSNorm(dim)
 
     def forward(self, x):
         x = self.norm(x)
@@ -210,7 +207,7 @@ def __init__(self, dim, heads = 4, dim_head = 32):
 
         self.to_out = nn.Sequential(
             nn.Conv1d(hidden_dim, dim, 1),
-            LayerNorm(dim)
+            RMSNorm(dim)
         )
 
     def forward(self, x):
@@ -868,9 +865,9 @@ def train(self):
                             milestone = self.step // self.save_and_sample_every
                             batches = num_to_groups(self.num_samples, self.batch_size)
                             all_samples_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))
-                        #
+
                         all_samples = torch.cat(all_samples_list, dim = 0)
-                        #
+
                         torch.save(all_samples, str(self.results_folder / f'sample-{milestone}.png'))
                         self.save(milestone)
 

denoising_diffusion_pytorch/guided_diffusion.py

@@ -104,26 +104,23 @@ def forward(self, x):
         weight = self.weight
         mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
         var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
-        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+        normalized_weight = (weight - mean) * var.clamp(min = eps).rsqrt()
 
         return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
 
     def forward(self, x):
-        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) * (var + eps).rsqrt() * self.g
+        return F.normalize(x, dim = 1) * self.g * (x.shape[-1] ** 0.5)
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
         self.fn = fn
-        self.norm = LayerNorm(dim)
+        self.norm = RMSNorm(dim)
 
     def forward(self, x):
         x = self.norm(x)
@@ -218,7 +215,7 @@ def __init__(self, dim, heads = 4, dim_head = 32):
 
         self.to_out = nn.Sequential(
             nn.Conv2d(hidden_dim, dim, 1),
-            LayerNorm(dim)
+            RMSNorm(dim)
         )
 
     def forward(self, x):

denoising_diffusion_pytorch/simple_diffusion.py

@@ -83,7 +83,7 @@ def Downsample(
         nn.Conv2d(dim * (factor ** 2), default(dim_out, dim), 1)
     )
 
-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, dim, scale = True, normalize_dim = 2):
         super().__init__()
         self.g = nn.Parameter(torch.ones(dim)) if scale else 1
@@ -94,11 +94,7 @@ def __init__(self, dim, scale = True, normalize_dim = 2):
     def forward(self, x):
         normalize_dim = self.normalize_dim
         scale = append_dims(self.g, x.ndim - self.normalize_dim - 1) if self.scale else 1
-
-        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
-        var = torch.var(x, dim = normalize_dim, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = normalize_dim, keepdim = True)
-        return (x - mean) * var.clamp(min = eps).rsqrt() * scale
+        return F.normalize(x, dim = normalize_dim) * scale * (x.shape[normalize_dim] ** 0.5)
 
 # sinusoidal positional embeds
 
@@ -169,12 +165,12 @@ def __init__(self, dim, heads = 4, dim_head = 32):
         self.heads = heads
         hidden_dim = dim_head * heads
 
-        self.norm = LayerNorm(dim, normalize_dim = 1)
+        self.norm = RMSNorm(dim, normalize_dim = 1)
         self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
 
         self.to_out = nn.Sequential(
             nn.Conv2d(hidden_dim, dim, 1),
-            LayerNorm(dim, normalize_dim = 1)
+            RMSNorm(dim, normalize_dim = 1)
         )
 
     def forward(self, x):
@@ -207,7 +203,7 @@ def __init__(self, dim, heads = 4, dim_head = 32, scale = 8, dropout = 0.):
         self.heads = heads
         hidden_dim = dim_head * heads
 
-        self.norm = LayerNorm(dim)
+        self.norm = RMSNorm(dim)
 
         self.attn_dropout = nn.Dropout(dropout)
         self.to_qkv = nn.Linear(dim, hidden_dim * 3, bias = False)
@@ -247,7 +243,7 @@ def __init__(
         dropout = 0.
     ):
         super().__init__()
-        self.norm = LayerNorm(dim, scale = False)
+        self.norm = RMSNorm(dim, scale = False)
         dim_hidden = dim * mult
 
         self.to_scale_shift = nn.Sequential(
@@ -359,10 +355,11 @@ def __init__(
                 self.init_conv = nn.Conv2d(channels, init_dim, patch_size, stride = patch_size)
             else:
                 self.init_conv = nn.Sequential(
-                    Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1 = patch_size, p2 = patch_size),
-                    LayerNorm(input_channels, normalize_dim = 1),
-                    nn.Conv2d(input_channels, init_dim, 1),
-                    LayerNorm(init_dim, normalize_dim = 1)
+                    Rearrange('b c (h p1) (w p2) -> b h w (c p1 p2)', p1 = patch_size, p2 = patch_size),
+                    nn.LayerNorm(input_channels),
+                    nn.Linear(input_channels, init_dim),
+                    nn.LayerNorm(init_dim),
+                    Rearrange('b h w c -> b c h w')
                 )
 
             self.unpatchify = nn.ConvTranspose2d(input_channels, channels, patch_size, stride = patch_size)

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.6.4'
+__version__ = '1.7.1'