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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.26
+Version: 0.0.29
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -35,10 +35,12 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
+Requires-Dist: axial-positional-embedding>=0.3.5
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8
 Requires-Dist: ninja
+Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict
 Requires-Dist: torch>=2.2
 Provides-Extra: examples

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

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.26"
+version = "0.0.29"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -26,10 +26,12 @@ classifiers=[
 
 dependencies = [
     "accelerated-scan>=0.2.0",
+    "axial_positional_embedding>=0.3.5",
     "einx>=0.3.0",
     "einops>=0.8.0",
     "hyper-connections>=0.1.8",
     "Ninja",
+    "rotary-embedding-torch",
     "tensordict",
     "torch>=2.2",
 ]

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

@@ -33,3 +33,19 @@ def test_titans_attn_memory():
     retrieved = mem(seq)
 
     assert seq.shape == retrieved.shape
+
+def test_mac():
+    from titans_pytorch.mac_transformer import MemoryAsContextTransformer
+
+    transformer = MemoryAsContextTransformer(
+        num_tokens = 256,
+        dim = 256,
+        depth = 2,
+        num_persist_mem_tokens = 16,
+        segment_len = 128,
+    )
+
+    x = torch.randint(0, 256, (1, 1023))
+
+    logits = transformer(x)
+    assert logits.shape == (1, 1023, 256)

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

@@ -1,16 +1,23 @@
 from __future__ import annotations
-import math
+from math import ceil
 from functools import partial
 
 import torch
-from torch import nn
+from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear
 
+from einops import repeat
 from einops.layers.torch import Rearrange
 
+
 from hyper_connections import get_init_and_expand_reduce_stream_functions
 
+# absolute and relative positions
+
+from axial_positional_embedding import ContinuousAxialPositionalEmbedding
+from rotary_embedding_torch import RotaryEmbedding
+
 # constants
 
 LinearNoBias = partial(Linear, bias = False)
@@ -24,7 +31,7 @@ def default(v, d):
     return v if exists(v) else d
 
 def round_up_multiple(seq, mult):
-    return math.ceil(seq / mult) * mult
+    return ceil(seq / mult) * mult
 
 # feedforward and attention
 
@@ -48,6 +55,7 @@ class SegmentedAttention(Module):
         self,
         dim,
         segment_len,
+        num_persist_mem_tokens,
         dim_head = 64,
         heads = 8,
     ):
@@ -56,6 +64,8 @@ class SegmentedAttention(Module):
 
         dim_inner = dim_head * heads
 
+        self.rotary_emb = RotaryEmbedding(dim_head)
+
         self.to_qkv = LinearNoBias(dim, dim_inner * 3)
         self.to_out = LinearNoBias(dim_inner, dim)
 
@@ -67,6 +77,7 @@ class SegmentedAttention(Module):
         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):
         batch, seq_len = seq.shape[:2]
@@ -92,6 +103,21 @@ class SegmentedAttention(Module):
         q, k, v = self.to_qkv(seq).chunk(3, dim = -1)
         q, k, v = map(self.split_heads, (q, k, v))
 
+        # take care of persistent memory key / values
+
+        pmk, pmv = tuple(repeat(t, 'h n d -> b h n d', b = seq.shape[0]) for t in self.persistent_memory)
+
+        # relative positions
+
+        q, k = self.rotary_emb.rotate_queries_with_cached_keys(q, k)
+
+        # persistent memory
+
+        k = cat((pmk, k), dim = -2)
+        v = cat((pmv, v), dim = -2)
+
+        # sdpa
+
         out = F.scaled_dot_product_attention(q, k, v, is_causal = True)
 
         out = self.merge_heads(out)
@@ -113,6 +139,7 @@ class MemoryAsContextTransformer(Module):
         dim,
         depth,
         segment_len,
+        num_persist_mem_tokens,
         dim_head = 64,
         heads = 8,
         ff_mult = 4,
@@ -120,6 +147,9 @@ class MemoryAsContextTransformer(Module):
     ):
         super().__init__()
 
+        self.segment_len = segment_len
+        self.axial_pos_emb = ContinuousAxialPositionalEmbedding(dim = dim, num_axial_dims = 2)
+
         self.token_emb = nn.Embedding(num_tokens, dim)
 
         init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
@@ -127,7 +157,14 @@ class MemoryAsContextTransformer(Module):
         self.layers = ModuleList([])
 
         for _ in range(depth):
-            attn = SegmentedAttention(dim = dim, dim_head = dim_head, heads = heads, segment_len = segment_len)
+            attn = SegmentedAttention(
+                dim = dim,
+                dim_head = dim_head,
+                heads = heads,
+                segment_len = segment_len,
+                num_persist_mem_tokens = num_persist_mem_tokens
+            )
+
             ff = FeedForward(dim = dim, mult = ff_mult)
 
             self.layers.append(ModuleList([
@@ -140,9 +177,19 @@ class MemoryAsContextTransformer(Module):
         self.to_logits = LinearNoBias(dim, num_tokens)
 
     def forward(self, x):
+        seq_len, segment_len = x.shape[-1], self.segment_len
+        windows = ceil(seq_len / segment_len)
 
         x = self.token_emb(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]
+
+        # expand and reduce streams for hyper connections
+
         x = self.expand_streams(x)
 
         for attn, ff in self.layers:
@@ -162,6 +209,7 @@ if __name__ == '__main__':
         num_tokens = 256,
         dim = 256,
         depth = 2,
+        num_persist_mem_tokens = 16,
         segment_len = 128,
     )