titans-pytorch 0.0.27__py3-none-any.whl → 0.0.30__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import math
+from math import ceil
 from functools import partial
 
 import torch
@@ -7,11 +7,16 @@ 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
+
+from axial_positional_embedding import ContinuousAxialPositionalEmbedding
+from rotary_embedding_torch import RotaryEmbedding
+
 # constants
 
 LinearNoBias = partial(Linear, bias = False)
@@ -25,7 +30,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
 
@@ -49,7 +54,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,
     ):
@@ -58,29 +64,37 @@ 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)
 
         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.segment_seq = Rearrange('b (n w) d -> (b n) w d', n = total_segment_len)
+        self.merge_seq_back = Rearrange('(b n) w d -> b (n w) d', n = total_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
+        need_segment = seq_len >= total_segment_len
 
         if need_segment:
-            next_seq_len = round_up_multiple(seq_len, self.segment_len)
+            next_seq_len = round_up_multiple(seq_len, total_segment_len)
             padding = next_seq_len - seq_len
 
             if padding > 0:
@@ -99,6 +113,12 @@ class SegmentedAttention(Module):
 
         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)
 
@@ -125,7 +145,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,
@@ -135,6 +156,17 @@ class MemoryAsContextTransformer(Module):
 
         self.token_emb = nn.Embedding(num_tokens, dim)
 
+        self.axial_pos_emb = ContinuousAxialPositionalEmbedding(dim = dim, num_axial_dims = 2)
+
+        # 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)
 
         self.layers = ModuleList([])
@@ -145,6 +177,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
             )
 
@@ -161,8 +194,45 @@ class MemoryAsContextTransformer(Module):
 
     def forward(self, x):
 
+        # 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
+
+        need_segment = seq_len >= segment_len
+
+        if need_segment:
+            next_seq_len = round_up_multiple(seq_len, segment_len)
+            padding = next_seq_len - seq_len
+
+            if padding > 0:
+                x = F.pad(x, (0, 0, 0, padding))
+
+            x = rearrange(x, 'b (w n) d -> (b w) n d', n = segment_len)
+
+        mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
+        x = torch.cat((mems, x), dim = -2)
+
+        if need_segment:
+            x = rearrange(x, '(b w) n d -> b (w n) d', b = batch)
+            x = x[:, :seq_len]
+
+        # 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[:seq_len]
+
+        # expand and reduce streams for hyper connections
+
         x = self.expand_streams(x)
 
         for attn, ff in self.layers:
@@ -171,21 +241,8 @@ class MemoryAsContextTransformer(Module):
 
         x = self.reduce_streams(x)
 
+        # to logits
+
         x = self.norm(x)
 
         return self.to_logits(x)
-
-# main
-
-if __name__ == '__main__':
-    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)

{titans_pytorch-0.0.27.dist-info → titans_pytorch-0.0.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.27
+Version: 0.0.30
 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.27.dist-info → titans_pytorch-0.0.30.dist-info}/RECORD

@@ -1,9 +1,9 @@
 titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=4xUSfGDdVsR-WmeXX7yXoFfybROvNCjOxL_EHDJ_Wlk,4681
+titans_pytorch/mac_transformer.py,sha256=XDLc9NWXpVMza03XjU0lkw5lRvtJ25ReTPKNoGslOOk,6773
 titans_pytorch/titans.py,sha256=Kx_tl_QkeDccvkMwPZ0xQ_saYjZfbKzDNPTTSHNWYcA,14304
 titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.27.dist-info/METADATA,sha256=qJp1IDbphEUfW7EyNvQ7RfmHuvB7SH5h_tlnCVwV4EY,3851
-titans_pytorch-0.0.27.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.27.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.27.dist-info/RECORD,,
+titans_pytorch-0.0.30.dist-info/METADATA,sha256=7H6WPsgfBE9ByEUP7r6C-cfBX6K13yNCaVhZMzwUvf8,3938
+titans_pytorch-0.0.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.30.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.30.dist-info/RECORD,,