titans-pytorch 0.0.34__tar.gz → 0.0.36__tar.gz

{titans_pytorch-0.0.34 → titans_pytorch-0.0.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.34
+Version: 0.0.36
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch

{titans_pytorch-0.0.34 → titans_pytorch-0.0.36}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.34"
+version = "0.0.36"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.0.34 → titans_pytorch-0.0.36}/titans_pytorch/__init__.py

@@ -2,3 +2,5 @@ from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
 )
+
+from titans_pytorch.mac_transformer import MemoryAsContextTransformer

{titans_pytorch-0.0.34 → titans_pytorch-0.0.36}/titans_pytorch/mac_transformer.py

@@ -185,7 +185,9 @@ class MemoryAsContextTransformer(Module):
         # long term mem tokens
 
         self.segment_len = segment_len
+
         self.num_longterm_mem_tokens = num_longterm_mem_tokens
+        has_longterm_mems = num_longterm_mem_tokens > 0
 
         self.longterm_mems = nn.Parameter(torch.randn(num_longterm_mem_tokens, dim) * 0.02)
 
@@ -197,7 +199,11 @@ class MemoryAsContextTransformer(Module):
         self.neural_mem_layers = ModuleList([])
 
         layers = tuple(range(1, depth + 1))
-        neural_memory_layers = set(default(neural_memory_layers, layers))
+
+        if not exists(neural_memory_layers):
+            neural_memory_layers = layers if has_longterm_mems else ()
+
+        assert not (num_longterm_mem_tokens > 0 and len(neural_memory_layers) == 0), 'empty `neural_memory_layers` when longterm memory tokens are present'
 
         for layer in layers:
 
@@ -205,7 +211,9 @@ class MemoryAsContextTransformer(Module):
 
             mem = None
 
-            if num_longterm_mem_tokens > 0 and layer in neural_memory_layers:
+            if layer in neural_memory_layers:
+                assert has_longterm_mems, '`num_longterm_mem_tokens` must be greater than 0'
+
                 mem = NeuralMemory(dim = dim, chunk_size = num_longterm_mem_tokens)
                 mem = init_hyper_conn(dim = dim, branch = mem)
 

{titans_pytorch-0.0.34 → titans_pytorch-0.0.36}/train_mac.py

@@ -25,10 +25,12 @@ SEQ_LEN = 512
 
 PROJECT_NAME = 'titans-mac-transformer'
 WANDB_ONLINE = False # turn this on to pipe experiment to cloud
-GLOBAL_LAYERS = (2, 4)
 NEURAL_MEMORY_DEPTH = 2
-WINDOW_SIZE = 64
-RUN_NAME = 'mac'
+NUM_PERSIST_MEM = 4
+NUM_LONGTERM_MEM = 4
+NEURAL_MEM_LAYERS = (2, 4, 6)
+WINDOW_SIZE = 32
+RUN_NAME = 'mac - 4 longterm mems, layers (2, 4, 6)'
 
 # wandb experiment tracker
 
@@ -57,12 +59,14 @@ model = MemoryAsContextTransformer(
     dim = 384,
     depth = 8,
     segment_len = WINDOW_SIZE,
-    num_persist_mem_tokens = 16,
-    num_longterm_mem_tokens = 16,
-    neural_memory_layers = (3, 4),
+    num_persist_mem_tokens = NUM_PERSIST_MEM,
+    num_longterm_mem_tokens = NUM_LONGTERM_MEM,
+    neural_memory_layers = NEURAL_MEM_LAYERS,
     neural_memory_kwargs = dict(
         default_mlp_kwargs = dict(
-            depth = NEURAL_MEMORY_DEPTH
+            depth = NEURAL_MEMORY_DEPTH,
+            dim_head = 64,
+            heads = 4
         )
     )
 ).cuda()