titans-pytorch 0.0.38__tar.gz → 0.0.40__tar.gz

{titans_pytorch-0.0.38 → titans_pytorch-0.0.40}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.38
+Version: 0.0.40
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -43,6 +43,7 @@ Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict
 Requires-Dist: torch>=2.2
+Requires-Dist: x-transformers
 Provides-Extra: examples
 Requires-Dist: local-attention>=1.10.1; extra == 'examples'
 Requires-Dist: taylor-series-linear-attention; extra == 'examples'

{titans_pytorch-0.0.38 → titans_pytorch-0.0.40}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.38"
+version = "0.0.40"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -34,6 +34,7 @@ dependencies = [
     "rotary-embedding-torch",
     "tensordict",
     "torch>=2.2",
+    "x-transformers"
 ]
 
 [project.urls]

{titans_pytorch-0.0.38 → titans_pytorch-0.0.40}/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, pack, unpack
+from einops import einsum, repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
 
 from hyper_connections import get_init_and_expand_reduce_stream_functions
@@ -16,6 +16,7 @@ from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 from axial_positional_embedding import ContinuousAxialPositionalEmbedding
 from rotary_embedding_torch import RotaryEmbedding
+from x_transformers.attend import Attend
 
 # proposed neural memory
 
@@ -93,6 +94,8 @@ class SegmentedAttention(Module):
         num_longterm_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
+        accept_value_residual = False,
+        attend_kwargs: dict = dict()
     ):
         super().__init__()
         self.norm = nn.RMSNorm(dim)
@@ -101,9 +104,17 @@ class SegmentedAttention(Module):
 
         self.rotary_emb = RotaryEmbedding(dim_head)
 
+        self.attend = Attend(causal = True, **attend_kwargs)
+
         self.to_qkv = LinearNoBias(dim, dim_inner * 3)
         self.to_out = LinearNoBias(dim_inner, dim)
 
+        self.to_learned_v_mix = nn.Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b n h -> b h n 1'),
+            nn.Sigmoid()
+        ) if accept_value_residual else None
+
         self.segment_len = segment_len
         self.num_longterm_mem_tokens = num_longterm_mem_tokens
 
@@ -114,7 +125,13 @@ class SegmentedAttention(Module):
 
         self.persistent_memory = nn.Parameter(torch.zeros(2, heads, num_persist_mem_tokens, dim_head))
 
-    def forward(self, seq):
+    def forward(
+        self,
+        seq,
+        value_residual = None
+    ):
+        assert not (exists(value_residual) ^ exists(self.to_learned_v_mix))
+
         segment_len, num_longterm_mem_tokens = self.segment_len, self.num_longterm_mem_tokens
         total_segment_len = segment_len + num_longterm_mem_tokens
 
@@ -132,6 +149,14 @@ class SegmentedAttention(Module):
         q, k, v = self.to_qkv(seq).chunk(3, dim = -1)
         q, k, v = map(self.split_heads, (q, k, v))
 
+        # value residual
+
+        orig_v = v
+
+        if exists(self.to_learned_v_mix):
+            mix = self.to_learned_v_mix(seq)
+            v = v.lerp(value_residual, mix)
+
         # 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)
@@ -145,9 +170,9 @@ class SegmentedAttention(Module):
         k = cat((pmk, k), dim = -2)
         v = cat((pmv, v), dim = -2)
 
-        # sdpa
+        # attention
 
-        out = F.scaled_dot_product_attention(q, k, v, is_causal = True)
+        out, _ = self.attend(q, k, v)
 
         out = self.merge_heads(out)
 
@@ -155,7 +180,7 @@ class SegmentedAttention(Module):
 
         out = inverse_segment(out)
 
-        return out
+        return out, orig_v
 
 # MAC transformer
 
@@ -206,6 +231,7 @@ class MemoryAsContextTransformer(Module):
         assert not (num_longterm_mem_tokens > 0 and len(neural_memory_layers) == 0), 'empty `neural_memory_layers` when longterm memory tokens are present'
 
         for layer in layers:
+            is_first = layer == 1
 
             # neural memory
 
@@ -231,6 +257,7 @@ class MemoryAsContextTransformer(Module):
                 dim_head = dim_head,
                 heads = heads,
                 segment_len = segment_len,
+                accept_value_residual = not is_first,
                 num_longterm_mem_tokens = num_longterm_mem_tokens,
                 num_persist_mem_tokens = num_persist_mem_tokens
             )
@@ -281,6 +308,10 @@ class MemoryAsContextTransformer(Module):
         pos_emb = self.axial_pos_emb((windows, total_segment_len), flatten = True)
         x = x + pos_emb[:x.shape[-2]]
 
+        # value residual
+
+        value_residual = None
+
         # expand and reduce streams for hyper connections
 
         x = self.expand_streams(x)
@@ -291,7 +322,9 @@ class MemoryAsContextTransformer(Module):
                 x = maybe_neural_mem(x)
 
 
-            x = attn(x)
+            x, values = attn(x, value_residual = value_residual)
+
+            value_residual = default(value_residual, values)
 
             x = ff(x)
 

{titans_pytorch-0.0.38 → titans_pytorch-0.0.40}/train_mac.py

@@ -24,7 +24,7 @@ SHOULD_GENERATE = False
 SEQ_LEN = 512
 
 PROJECT_NAME = 'titans-mac-transformer'
-WANDB_ONLINE = True # turn this on to pipe experiment to cloud
+WANDB_ONLINE = False # turn this on to pipe experiment to cloud
 NEURAL_MEMORY_DEPTH = 2
 NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4