Merge pull request #121 from lucidrains/fix-lfq-distance

calculate full euclidean distance in lfq

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/lookup_free_quantization.py

@@ -48,6 +48,14 @@ def log(t, eps = 1e-5):
 def entropy(prob):
     return (-prob * log(prob)).sum(dim=-1)
 
+# distance
+
+def euclidean_distance_squared(x, y):
+    x2 = reduce(x ** 2, '... n d -> ... n 1', 'sum')
+    y2 = reduce(y ** 2, 'n d -> n', 'sum')
+    xy = einsum('... i d, j d -> ... i j', x, y) * -2
+    return x2 + xy + y2
+
 # class
 
 class LFQ(Module):
@@ -218,8 +226,7 @@ def forward(
         # entropy aux loss
 
         if self.training:
-            # the same as euclidean distance up to a constant
-            distance = -2 * einsum('... i d, j d -> ... i j', original_input, self.codebook)
+            distance = euclidean_distance_squared(original_input, self.codebook)
 
             prob = (-distance * inv_temperature).softmax(dim = -1)