Irie et al. notices that the original Oord implementation of VQ sets cluster sizes of 0 initially, leading to worse convergence. not an issue if kmeans init is turned on

Author: lucidrains

README.md

@@ -679,3 +679,12 @@ assert loss.item() >= 0
     primaryClass = {cs.LG}
 }
 ```
+
+```bibtex
+@inproceedings{Irie2023SelfOrganisingND,
+    title   = {Self-Organising Neural Discrete Representation Learning \`a la Kohonen},
+    author  = {Kazuki Irie and R'obert Csord'as and J{\"u}rgen Schmidhuber},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:256901024}
+}
+```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.15.2"
+version = "1.15.3"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -312,7 +312,7 @@ def __init__(
         self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
-        self.register_buffer('cluster_size', torch.zeros(num_codebooks, codebook_size))
+        self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
         self.register_buffer('embed_avg', embed.clone())
 
         self.learnable_codebook = learnable_codebook
@@ -582,7 +582,7 @@ def __init__(
         self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
-        self.register_buffer('cluster_size', torch.zeros(num_codebooks, codebook_size))
+        self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
         self.register_buffer('embed_avg', embed.clone())
 
         self.learnable_codebook = learnable_codebook