titans-pytorch 0.0.56__py3-none-any.whl → 0.0.58__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -285,6 +285,49 @@ class SegmentedAttention(Module):
 
         return out, orig_v
 
+# Attention + Neural Memory gating configuration, as depicted in Figure 2
+
+class NeuralMemoryGatingWrapper(Module):
+    def __init__(
+        self,
+        dim,
+        attn: SegmentedAttention,
+        neural_mem: NeuralMemory | None = None,
+        gate_attn_output = True
+    ):
+        super().__init__()
+        self.attn = attn
+        self.neural_mem = neural_mem
+        self.gate_attn_output = gate_attn_output
+
+    def forward(
+        self,
+        seq,
+        *args,
+        **kwargs
+    ):
+        batch, seq_len = seq.shape[:2]
+        mem = self.neural_mem
+
+        if not exists(mem):
+            return self.attn(seq, *args, **kwargs), 0.
+
+        # initial retrieve, still should store first, it doesn't make sense not to, unless if all layers share the same neural memory
+
+        retrieved, kv_aux_loss = mem(seq, return_aux_kv_loss = True)
+
+        if not self.gate_attn_output:
+            seq = seq + retrieved
+
+        # attention
+
+        attn_out, values = self.attn(seq, *args, **kwargs)
+
+        if self.gate_attn_output:
+            attn_out = attn_out * retrieved.sigmoid()
+
+        return (attn_out, values), kv_aux_loss
+
 # MAC transformer
 
 class MemoryAsContextTransformer(Module):
@@ -296,6 +339,7 @@ class MemoryAsContextTransformer(Module):
         depth,
         segment_len,
         neural_memory_segment_len = None,
+        neural_mem_gate_attn_output = True,
         num_longterm_mem_tokens = 0,
         num_persist_mem_tokens = 0,
         dim_head = 64,
@@ -328,7 +372,6 @@ class MemoryAsContextTransformer(Module):
 
         self.layers = ModuleList([])
 
-        self.neural_mem_layers = ModuleList([])
         self.neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
 
         layers = tuple(range(1, depth + 1))
@@ -343,7 +386,18 @@ class MemoryAsContextTransformer(Module):
         for layer in layers:
             is_first = layer == 1
 
-            # neural memory
+            # attention and feedforward
+
+            attn = SegmentedAttention(
+                dim = dim,
+                dim_head = dim_head,
+                heads = heads,
+                segment_len = segment_len,
+                use_flex_attn = use_flex_attn,
+                accept_value_residual = not is_first,
+                num_longterm_mem_tokens = num_longterm_mem_tokens,
+                num_persist_mem_tokens = num_persist_mem_tokens
+            )
 
             mem = None
 
@@ -356,21 +410,11 @@ class MemoryAsContextTransformer(Module):
                     **neural_memory_kwargs
                 )
 
-                mem = init_hyper_conn(dim = dim, branch = mem)
-
-            self.neural_mem_layers.append(mem)
-
-            # attention and feedforward
-
-            attn = SegmentedAttention(
-                dim = dim,
-                dim_head = dim_head,
-                heads = heads,
-                segment_len = segment_len,
-                use_flex_attn = use_flex_attn,
-                accept_value_residual = not is_first,
-                num_longterm_mem_tokens = num_longterm_mem_tokens,
-                num_persist_mem_tokens = num_persist_mem_tokens
+            attn = NeuralMemoryGatingWrapper(
+                dim,
+                attn = attn,
+                neural_mem = mem,
+                gate_attn_output = neural_mem_gate_attn_output
             )
 
             ff = FeedForward(dim = dim, mult = ff_mult)
@@ -460,19 +504,17 @@ class MemoryAsContextTransformer(Module):
 
         x = self.expand_streams(x)
 
-        for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
-
-            if exists(maybe_neural_mem):
-                x, aux_kv_loss = maybe_neural_mem(x, return_aux_kv_loss = True)
-                kv_recon_losses = kv_recon_losses + aux_kv_loss
+        for attn, ff in self.layers:
 
-            x, values = attn(
+            (x, values), maybe_mem_kv_aux_loss = attn(
                 x,
                 value_residual = value_residual,
                 disable_flex_attn = disable_flex_attn,
                 flex_attn_fn = flex_attn_fn
             )
 
+            kv_recon_losses = kv_recon_losses + maybe_mem_kv_aux_loss
+
             value_residual = default(value_residual, values)
 
             x = ff(x)

{titans_pytorch-0.0.56.dist-info → titans_pytorch-0.0.58.dist-info}/METADATA

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

titans_pytorch-0.0.58.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=Q0MQA3RS8vqzs-KzSGZkvLR7afQ6ZW9uMOq1MeNuFoY,170
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=kk8s8Q2WmbJxCVi8PcqSUyJBc8-CDAHrVjt6M0d_kFs,15323
+titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
+titans_pytorch-0.0.58.dist-info/METADATA,sha256=a-Y6MV_89D44HlB7eKpurh-sw5DDiS-pIVei3Uw_uGE,4457
+titans_pytorch-0.0.58.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.58.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.58.dist-info/RECORD,,

titans_pytorch-0.0.56.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=Q0MQA3RS8vqzs-KzSGZkvLR7afQ6ZW9uMOq1MeNuFoY,170
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=_Vsco5YuR6uxouWcjFj-s-zPhrBcaapIzqoyi7qqY0Q,14245
-titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
-titans_pytorch-0.0.56.dist-info/METADATA,sha256=QlCmHqajHiaZTps0W9gKXIHE6dShZER3PqPoYi2zRe4,4457
-titans_pytorch-0.0.56.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.56.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.56.dist-info/RECORD,,