From d9275a744c31abd656dd18bea98cfcbc6ace2933 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Wed, 17 Aug 2022 11:42:54 -0700
Subject: [PATCH] add weight standardization prior to groupnorm, lessen cosine
 sim attention scale to 10 for fp16

---
 README.md                                     | 10 ++++++++++
 .../denoising_diffusion_pytorch.py            | 20 ++++++++++++++++---
 setup.py                                      |  2 +-
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index f910378..b499b0e 100644
--- a/README.md
+++ b/README.md
@@ -181,3 +181,13 @@ $ accelerate launch train.py
     primaryClass = {cs.CV}
 }
 ```
+
+```bibtex
+@article{Qiao2019WeightS,
+    title   = {Weight Standardization},
+    author  = {Siyuan Qiao and Huiyu Wang and Chenxi Liu and Wei Shen and Alan Loddon Yuille},
+    journal = {ArXiv},
+    year    = {2019},
+    volume  = {abs/1903.10520}
+}
+```
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index be0eeb4..5a3e096 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -88,6 +88,21 @@ def Upsample(dim, dim_out = None):
 def Downsample(dim, dim_out = None):
     return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
 
+class WeightStandardizedConv2d(nn.Conv2d):
+    """
+    https://arxiv.org/abs/1903.10520
+    weight standardization purportedly works synergistically with group normalization
+    """
+    def forward(self, x):
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
+
+        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()
+
+        return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
 class LayerNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -147,7 +162,7 @@ class LearnedSinusoidalPosEmb(nn.Module):
 class Block(nn.Module):
     def __init__(self, dim, dim_out, groups = 8):
         super().__init__()
-        self.proj = nn.Conv2d(dim, dim_out, 3, padding = 1)
+        self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding = 1)
         self.norm = nn.GroupNorm(groups, dim_out)
         self.act = nn.SiLU()
 
@@ -219,7 +234,7 @@ class LinearAttention(nn.Module):
         return self.to_out(out)
 
 class Attention(nn.Module):
-    def __init__(self, dim, heads = 4, dim_head = 32, scale = 16):
+    def __init__(self, dim, heads = 4, dim_head = 32, scale = 10):
         super().__init__()
         self.scale = scale
         self.heads = heads
@@ -236,7 +251,6 @@ class Attention(nn.Module):
 
         sim = einsum('b h d i, b h d j -> b h i j', q, k) * self.scale
         attn = sim.softmax(dim = -1)
-
         out = einsum('b h i j, b h d j -> b h i d', attn, v)
         out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
         return self.to_out(out)
diff --git a/setup.py b/setup.py
index bc7fbf8..1ab23d5 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.27.2',
+  version = '0.27.4',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',