more feedback; authors used uncentered covariance matrix, detail in appendix C

Author: lucidrains

perfusion_pytorch/perfusion.py

@@ -36,6 +36,8 @@ def calculate_input_covariance(
 
     all_embeds = []
 
+    length = len(texts)
+
     for batch_ind in range(num_batches):
         start_index = batch_ind * batch_size
         batch_texts = texts[start_index:(start_index + batch_size)]
@@ -44,9 +46,8 @@ def calculate_input_covariance(
         all_embeds.append(embeds[mask])
 
     all_embeds = torch.cat((all_embeds), dim = 0)
-    all_embeds = rearrange(all_embeds, 'n d -> d n')
 
-    return torch.cov(all_embeds, correction = 0, **cov_kwargs)
+    return einsum('n d, n e -> d e', all_embeds, all_embeds) / length
 
 # a module that wraps the keys and values projection of the cross attentions to text encodings
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.16',
+  version = '0.0.17',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',