titans-pytorch 0.0.29__tar.gz → 0.0.31__tar.gz

{titans_pytorch-0.0.29 → titans_pytorch-0.0.31}/PKG-INFO

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

{titans_pytorch-0.0.29 → titans_pytorch-0.0.31}/pyproject.toml

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

{titans_pytorch-0.0.29 → titans_pytorch-0.0.31}/tests/test_titans.py

@@ -34,14 +34,20 @@ def test_titans_attn_memory():
 
     assert seq.shape == retrieved.shape
 
-def test_mac():
+@pytest.mark.parametrize('num_persist_mem_tokens', (0, 16))
+@pytest.mark.parametrize('num_longterm_mem_tokens', (0, 16))
+def test_mac(
+    num_persist_mem_tokens,
+    num_longterm_mem_tokens
+):
     from titans_pytorch.mac_transformer import MemoryAsContextTransformer
 
     transformer = MemoryAsContextTransformer(
         num_tokens = 256,
         dim = 256,
         depth = 2,
-        num_persist_mem_tokens = 16,
+        num_persist_mem_tokens = num_persist_mem_tokens,
+        num_longterm_mem_tokens = num_longterm_mem_tokens,
         segment_len = 128,
     )
 

{titans_pytorch-0.0.29 → titans_pytorch-0.0.31}/titans_pytorch/mac_transformer.py

@@ -7,10 +7,9 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear
 
-from einops import repeat
+from einops import repeat, rearrange
 from einops.layers.torch import Rearrange
 
-
 from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 # absolute and relative positions
@@ -30,9 +29,33 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def identity(t):
+    return t
+
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
+def pad_and_segment_with_inverse(seq, segment_len):
+    batch, seq_len = seq.shape[:2]
+
+    need_segment = seq_len >= segment_len
+
+    if not need_segment:
+        return seq, identity
+
+    next_seq_len_mult = round_up_multiple(seq_len, segment_len)
+
+    padding = next_seq_len_mult - seq_len
+    seq = F.pad(seq, (0, 0, 0, padding))
+
+    seq = rearrange(seq, 'b (w n) d -> (b w) n d', n = segment_len)
+
+    def inverse(out):
+        out = rearrange(out, '(b w) n d -> b (w n) d', b = batch)
+        return out[:, :-padding]
+
+    return seq, inverse
+
 # feedforward and attention
 
 class GEGLU(Module):
@@ -55,7 +78,8 @@ class SegmentedAttention(Module):
         self,
         dim,
         segment_len,
-        num_persist_mem_tokens,
+        num_persist_mem_tokens = 0,
+        num_longterm_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
     ):
@@ -70,31 +94,25 @@ class SegmentedAttention(Module):
         self.to_out = LinearNoBias(dim_inner, dim)
 
         self.segment_len = segment_len
+        self.num_longterm_mem_tokens = num_longterm_mem_tokens
+
+        total_segment_len = segment_len + num_longterm_mem_tokens
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
 
-        self.segment_seq = Rearrange('b (n w) d -> (b n) w d', n = segment_len)
-        self.merge_seq_back = Rearrange('(b n) w d -> b (n w) d', n = segment_len)
-
         self.persistent_memory = nn.Parameter(torch.zeros(2, heads, num_persist_mem_tokens, dim_head))
 
     def forward(self, seq):
+        segment_len, num_longterm_mem_tokens = self.segment_len, self.num_longterm_mem_tokens
+        total_segment_len = segment_len + num_longterm_mem_tokens
+
         batch, seq_len = seq.shape[:2]
 
         # auto pad to multiple
         # todo - get rid of logic with flex attention
 
-        need_segment = seq_len >= self.segment_len
-
-        if need_segment:
-            next_seq_len = round_up_multiple(seq_len, self.segment_len)
-            padding = next_seq_len - seq_len
-
-            if padding > 0:
-                seq = F.pad(seq, (0, 0, 0, padding))
-
-            seq = self.segment_seq(seq)
+        seq, inverse_segment = pad_and_segment_with_inverse(seq, total_segment_len)
 
         # attention
 
@@ -124,10 +142,9 @@ class SegmentedAttention(Module):
 
         out = self.to_out(out)
 
-        if need_segment:
-            out = self.merge_seq_back(out)
+        out = inverse_segment(out)
 
-        return out[:, :seq_len]
+        return out
 
 # MAC transformer
 
@@ -139,7 +156,8 @@ class MemoryAsContextTransformer(Module):
         dim,
         depth,
         segment_len,
-        num_persist_mem_tokens,
+        num_longterm_mem_tokens = 0,
+        num_persist_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
         ff_mult = 4,
@@ -147,10 +165,18 @@ class MemoryAsContextTransformer(Module):
     ):
         super().__init__()
 
-        self.segment_len = segment_len
+        self.token_emb = nn.Embedding(num_tokens, dim)
+
         self.axial_pos_emb = ContinuousAxialPositionalEmbedding(dim = dim, num_axial_dims = 2)
 
-        self.token_emb = nn.Embedding(num_tokens, dim)
+        # long term mem tokens
+
+        self.segment_len = segment_len
+        self.num_longterm_mem_tokens = num_longterm_mem_tokens
+
+        self.longterm_mems = nn.Parameter(torch.randn(num_longterm_mem_tokens, dim) * 0.02)
+
+        # hyper conection
 
         init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
 
@@ -162,6 +188,7 @@ class MemoryAsContextTransformer(Module):
                 dim_head = dim_head,
                 heads = heads,
                 segment_len = segment_len,
+                num_longterm_mem_tokens = num_longterm_mem_tokens,
                 num_persist_mem_tokens = num_persist_mem_tokens
             )
 
@@ -177,16 +204,32 @@ class MemoryAsContextTransformer(Module):
         self.to_logits = LinearNoBias(dim, num_tokens)
 
     def forward(self, x):
-        seq_len, segment_len = x.shape[-1], self.segment_len
+
+        # 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
+
+        # token embedding
 
         x = self.token_emb(x)
 
+        # intersperse longterm memory
+
+        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 = torch.cat((mems, x), dim = -2)
+
+        x = inverse_segment(x)
+
         # apply axial positional embedding
         # so intra and inter segment can be more easily discerned by the network
 
-        pos_emb = self.axial_pos_emb((windows, segment_len), flatten = True)
-        x = x + pos_emb[:seq_len]
+        pos_emb = self.axial_pos_emb((windows, total_segment_len), flatten = True)
+        x = x + pos_emb[:x.shape[-2]]
 
         # expand and reduce streams for hyper connections
 
@@ -198,21 +241,16 @@ class MemoryAsContextTransformer(Module):
 
         x = self.reduce_streams(x)
 
-        x = self.norm(x)
+        # excise out the memories
 
-        return self.to_logits(x)
+        x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
 
-# main
+        x = x[:, self.num_longterm_mem_tokens:]
 
-if __name__ == '__main__':
-    transformer = MemoryAsContextTransformer(
-        num_tokens = 256,
-        dim = 256,
-        depth = 2,
-        num_persist_mem_tokens = 16,
-        segment_len = 128,
-    )
+        x = inverse_segment(x)
+
+        # to logits
 
-    x = torch.randint(0, 256, (1, 1023))
+        x = self.norm(x)
 
-    logits = transformer(x)
+        return self.to_logits(x)