set new defaults by recommendation from @AranKomat

README.md

@@ -18,9 +18,9 @@ from product_key_memory import PKM
 
 pkm = PKM(
     dim = 512,
-    heads = 8,
-    num_keys = 512,       # number of subkeys, # values will be num_keys ^ 2
-    topk = 10,            # the top number of subkeys to select
+    heads = 4,
+    num_keys = 256,       # number of subkeys, # values will be num_keys ^ 2
+    topk = 32,            # the top number of subkeys to select
     use_evonorm = True    # usually PKM requires decent batch sizes with batchnorm to work well. this is an experimental feature using the new evonorm-s0 for batch-independent normalization
 )
 

product_key_memory/product_key_memory.py

@@ -28,7 +28,7 @@ def forward(self, x):
         return x.reshape(*shape)
 
 class PKM(nn.Module):
-    def __init__(self, dim, heads = 8, num_keys = 128, topk = 10, input_dropout = 0., query_dropout = 0., value_dropout = 0., use_evonorm = False):
+    def __init__(self, dim, heads = 4, num_keys = 256, topk = 32, input_dropout = 0., query_dropout = 0., value_dropout = 0., use_evonorm = False):
         super().__init__()
         assert (dim % heads == 0), 'dimension must be divisible by number of heads'
         self.topk = topk

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'product_key_memory',
   packages = find_packages(),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Product Key Memory',
   author = 'Phil Wang',