titans-pytorch 0.0.42__tar.gz → 0.0.44__tar.gz

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.42
+Version: 0.0.44
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -83,6 +83,26 @@ retrieved = mem(seq)
 assert seq.shape == retrieved.shape
 ```
 
+A transformer with the `MAC` configuration can be used as
+
+```python
+import torch
+from titans_pytorch import MemoryAsContextTransformer
+
+transformer = MemoryAsContextTransformer(
+    num_tokens = 256,
+    dim = 256,
+    depth = 2,
+    segment_len = 128,              # local attention window size
+    num_persist_mem_tokens = 4,
+    num_longterm_mem_tokens = 16,
+)
+
+token_ids = torch.randint(0, 256, (1, 1023))
+
+logits = transformer(token_ids) # (1, 1023, 256)
+```
+
 ## Experiments
 
 ```bash

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/README.md

@@ -30,6 +30,26 @@ retrieved = mem(seq)
 assert seq.shape == retrieved.shape
 ```
 
+A transformer with the `MAC` configuration can be used as
+
+```python
+import torch
+from titans_pytorch import MemoryAsContextTransformer
+
+transformer = MemoryAsContextTransformer(
+    num_tokens = 256,
+    dim = 256,
+    depth = 2,
+    segment_len = 128,              # local attention window size
+    num_persist_mem_tokens = 4,
+    num_longterm_mem_tokens = 16,
+)
+
+token_ids = torch.randint(0, 256, (1, 1023))
+
+logits = transformer(token_ids) # (1, 1023, 256)
+```
+
 ## Experiments
 
 ```bash

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/pyproject.toml

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

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/tests/test_titans.py

@@ -1,5 +1,6 @@
 import torch
 import pytest
+from titans_pytorch import NeuralMemory
 
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
@@ -7,9 +8,6 @@ def test_titans(
     seq_len,
     max_grad_norm
 ):
-
-    from titans_pytorch import NeuralMemory
-
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
@@ -22,11 +20,14 @@ def test_titans(
     assert seq.shape == retrieved.shape
 
 def test_titans_attn_memory():
-    from titans_pytorch.titans_attn_memory import NeuralMemory
+    from titans_pytorch.titans import MemoryAttention
 
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
+        model = MemoryAttention(
+            dim = 384
+        )
     )
 
     seq = torch.randn(2, 1024, 384)

titans_pytorch-0.0.44/titans_pytorch/__init__.py

@@ -0,0 +1,8 @@
+from titans_pytorch.titans import (
+    NeuralMemory,
+    MemoryMLP,
+)
+
+from titans_pytorch.mac_transformer import (
+    MemoryAsContextTransformer
+)

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/titans_pytorch/mac_transformer.py

@@ -224,7 +224,7 @@ class MemoryAsContextTransformer(Module):
         self.layers = ModuleList([])
 
         self.neural_mem_layers = ModuleList([])
-        neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
+        self.neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
 
         layers = tuple(range(1, depth + 1))
 
@@ -245,7 +245,7 @@ class MemoryAsContextTransformer(Module):
 
                 mem = NeuralMemory(
                     dim = dim,
-                    chunk_size = neural_memory_segment_len,
+                    chunk_size = self.neural_memory_segment_len,
                     **neural_memory_kwargs
                 )
 
@@ -287,10 +287,7 @@ class MemoryAsContextTransformer(Module):
 
         # math
 
-        batch, seq_len, segment_len, num_longterm_mem_tokens= *x.shape, self.segment_len, self.num_longterm_mem_tokens
-
-        windows = ceil(seq_len / segment_len)
-        total_segment_len = segment_len + num_longterm_mem_tokens
+        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens
 
         # token embedding
 
@@ -305,11 +302,16 @@ class MemoryAsContextTransformer(Module):
 
         x = inverse_segment(x)
 
+        seq_len_with_mem = x.shape[-2]
+
         # apply axial positional embedding
         # so intra and inter segment can be more easily discerned by the network
 
-        pos_emb = self.axial_pos_emb((windows, total_segment_len), flatten = True)
-        x = x + pos_emb[:x.shape[-2]]
+        neural_mem_windows = ceil(seq_len_with_mem / neural_mem_segment_len)
+
+        pos_emb = self.axial_pos_emb((neural_mem_windows, neural_mem_segment_len), flatten = True)
+
+        x = x + pos_emb[:seq_len_with_mem]
 
         # value residual
 
@@ -334,7 +336,7 @@ class MemoryAsContextTransformer(Module):
 
         # excise out the memories
 
-        x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
+        x, inverse_segment = pad_and_segment_with_inverse(x, segment_len + num_longterm_mem_tokens)
 
         x, _ = unpack(x, mem_ps, 'b * d')
 

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/titans_pytorch/titans.py

@@ -425,11 +425,7 @@ class NeuralMemory(Module):
         next_seq_len = round_up_multiple(curtailed_seq_len, chunk_size)
 
         padding = next_seq_len - curtailed_seq_len
-
-        needs_pad = padding > 0
-
-        if needs_pad:
-            seq = pad_at_dim(seq, (0, padding), dim = 1)
+        seq = pad_at_dim(seq, (0, padding), dim = 1)
 
         # the parameters of the memory model stores the memories of the key / values
         # when the MLP has only 1 weight matrix, it is equivalent to `kv` fast weight memories from linear attention literature (recall fetching of memories is q @ (kv)) / schmidhuber's paper
@@ -481,10 +477,7 @@ class NeuralMemory(Module):
         empty_memory_embeds = self.init_empty_memory_embed(values.shape[0], chunk_size - 1)
         values = torch.cat((empty_memory_embeds, values), dim = -2)
 
-        if needs_pad:
-            values = values[:, :-padding]
-
-        return values
+        return values[:, :seq_len]
 
     def forward(
         self,

{titans_pytorch-0.0.42 → titans_pytorch-0.0.44}/train_mac.py

@@ -62,6 +62,7 @@ model = MemoryAsContextTransformer(
     num_persist_mem_tokens = NUM_PERSIST_MEM,
     num_longterm_mem_tokens = NUM_LONGTERM_MEM,
     neural_memory_layers = NEURAL_MEM_LAYERS,
+    neural_memory_segment_len = WINDOW_SIZE // 2,
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,

titans_pytorch-0.0.42/titans_pytorch/__init__.py

@@ -1,6 +0,0 @@
-from titans_pytorch.titans import (
-    NeuralMemory,
-    MemoryMLP,
-)
-
-from titans_pytorch.mac_transformer import MemoryAsContextTransformer