titans-pytorch 0.0.39__py3-none-any.whl → 0.0.41__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -94,6 +94,7 @@ class SegmentedAttention(Module):
         num_longterm_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
+        accept_value_residual = False,
         attend_kwargs: dict = dict()
     ):
         super().__init__()
@@ -108,6 +109,12 @@ class SegmentedAttention(Module):
         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
 
@@ -118,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
 
@@ -136,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)
@@ -159,7 +180,7 @@ class SegmentedAttention(Module):
 
         out = inverse_segment(out)
 
-        return out
+        return out, orig_v
 
 # MAC transformer
 
@@ -171,6 +192,7 @@ class MemoryAsContextTransformer(Module):
         dim,
         depth,
         segment_len,
+        neural_memory_segment_len = None,
         num_longterm_mem_tokens = 0,
         num_persist_mem_tokens = 0,
         dim_head = 64,
@@ -200,7 +222,9 @@ class MemoryAsContextTransformer(Module):
         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([])
+
         self.neural_mem_layers = ModuleList([])
+        neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
 
         layers = tuple(range(1, depth + 1))
 
@@ -210,6 +234,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
 
@@ -220,7 +245,7 @@ class MemoryAsContextTransformer(Module):
 
                 mem = NeuralMemory(
                     dim = dim,
-                    chunk_size = num_longterm_mem_tokens + segment_len,
+                    chunk_size = neural_memory_segment_len,
                     **neural_memory_kwargs
                 )
 
@@ -235,6 +260,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
             )
@@ -285,6 +311,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)
@@ -295,7 +325,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.39.dist-info → titans_pytorch-0.0.41.dist-info}/METADATA

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

{titans_pytorch-0.0.39.dist-info → titans_pytorch-0.0.41.dist-info}/RECORD

@@ -1,9 +1,9 @@
 titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=h58sHfufxMnSXZXyWuW-KBwzq8xwBYmFjU2XtOjUixk,8512
+titans_pytorch/mac_transformer.py,sha256=szHg8m97ew7OtipVih3pkOe1jsvhBnqvohJVJBrU5ks,9452
 titans_pytorch/titans.py,sha256=bv2Ceq-_4nNb5FNx4hLd2inC93m5MmJxO2-Mbf6PKK4,14378
 titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.39.dist-info/METADATA,sha256=3KD2hmJ-uOyQ87Z3VB6JfaKtDcLBnoKA8037DpzJuPE,3968
-titans_pytorch-0.0.39.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.39.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.39.dist-info/RECORD,,
+titans_pytorch-0.0.41.dist-info/METADATA,sha256=RXjWshmGlzoA-OYaiI6GwieXBwbsq5ZR_hk-PM66aeg,3968
+titans_pytorch-0.0.41.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.41.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.41.dist-info/RECORD,,