random projection quantizer needs to have input normalized to mean 0 std of 1

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.1.2',
+  version = '1.1.4',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/random_projection_quantizer.py

@@ -25,6 +25,11 @@ def __init__(
 
         self.register_buffer('rand_projs', rand_projs)
 
+        # in section 3 of https://arxiv.org/abs/2202.01855
+        # "The input data is normalized to have 0 mean and standard deviation of 1 ... to prevent collapse"
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine = False)
+
         self.vq = VectorQuantize(
             dim = codebook_dim * num_codebooks,
             heads = num_codebooks,
@@ -37,6 +42,8 @@ def __init__(
     @torch.no_grad()
     def forward(self, x):
 
+        x = self.norm(x)
+
         x = einsum('b n d, h d e -> b n h e', x, self.rand_projs)
         x, ps = pack([x], 'b n *')