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

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

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

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

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

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

@@ -7,7 +7,7 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear
 
-from einops import repeat, rearrange
+from einops import repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
 
 from hyper_connections import get_init_and_expand_reduce_stream_functions
@@ -214,7 +214,12 @@ class MemoryAsContextTransformer(Module):
             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 = NeuralMemory(
+                    dim = dim,
+                    chunk_size = num_longterm_mem_tokens + segment_len,
+                    **neural_memory_kwargs
+                )
+
                 mem = init_hyper_conn(dim = dim, branch = mem)
 
             self.neural_mem_layers.append(mem)
@@ -266,7 +271,7 @@ class MemoryAsContextTransformer(Module):
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len)
 
         mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
-        x = cat((mems, x), dim = -2)
+        x, mem_ps = pack((x, mems), 'b * d')
 
         x = inverse_segment(x)
 
@@ -283,21 +288,7 @@ class MemoryAsContextTransformer(Module):
         for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
 
             if exists(maybe_neural_mem):
-                batch_streams = x.shape[0]
-
-                x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
-
-                longterm_mems, x = x[:, :num_longterm_mem_tokens], x[:, num_longterm_mem_tokens:]
-
-                longterm_mems = rearrange(longterm_mems, '(b w) n d -> b (w n) d', b = batch_streams)
-
-                longterm_mems = maybe_neural_mem(longterm_mems)
-
-                longterm_mems = rearrange(longterm_mems, 'b (w n) d -> (b w) n d', n = num_longterm_mem_tokens)
-
-                x = cat((longterm_mems, x), dim = -2)
-
-                x = inverse_segment(x)
+                mems = maybe_neural_mem(mems)
 
             x = attn(x)
 
@@ -309,7 +300,7 @@ class MemoryAsContextTransformer(Module):
 
         x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
 
-        x = x[:, num_longterm_mem_tokens:]
+        x, mem = unpack(x, mem_ps, 'b * d')
 
         x = inverse_segment(x)
 

{titans_pytorch-0.0.36 → titans_pytorch-0.0.37}/titans_pytorch/titans.py

@@ -132,7 +132,7 @@ class NeuralMemory(Module):
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         default_mlp_kwargs: dict = dict(
-            depth = 4
+            depth = 2
         )
     ):
         super().__init__()

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

@@ -24,13 +24,13 @@ SHOULD_GENERATE = False
 SEQ_LEN = 512
 
 PROJECT_NAME = 'titans-mac-transformer'
-WANDB_ONLINE = False # turn this on to pipe experiment to cloud
+WANDB_ONLINE = True # turn this on to pipe experiment to cloud
 NEURAL_MEMORY_DEPTH = 2
 NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
-NEURAL_MEM_LAYERS = (2, 4, 6)
+NEURAL_MEM_LAYERS = (4,)
 WINDOW_SIZE = 32
-RUN_NAME = 'mac - 4 longterm mems, layers (2, 4, 6)'
+RUN_NAME = 'mac - 4 longterm mems, layers (4,)'
 
 # wandb experiment tracker
 
@@ -63,10 +63,10 @@ model = MemoryAsContextTransformer(
     num_longterm_mem_tokens = NUM_LONGTERM_MEM,
     neural_memory_layers = NEURAL_MEM_LAYERS,
     neural_memory_kwargs = dict(
+        dim_head = 64,
+        heads = 4,
         default_mlp_kwargs = dict(
             depth = NEURAL_MEMORY_DEPTH,
-            dim_head = 64,
-            heads = 4
         )
     )
 ).cuda()