titans-pytorch 0.0.55__py3-none-any.whl → 0.0.57__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -285,6 +285,50 @@ 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
+    ):
+        super().__init__()
+        self.attn = attn
+        self.neural_mem = neural_mem
+        self.to_gates = nn.Linear(dim, dim) if exists(neural_mem) else None
+
+    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, first_kv_aux_loss = mem(seq, return_aux_kv_loss = True)
+
+        seq = seq + retrieved
+
+        # attention
+
+        attn_out, values = self.attn(seq, *args, **kwargs)
+
+        # another retrieve, but this time gate the attention output
+
+        retrieved, second_kv_aux_loss = mem(attn_out, return_aux_kv_loss = True)
+
+        attn_out = attn_out * self.to_gates(retrieved).sigmoid()
+
+        return (attn_out, values), (first_kv_aux_loss + second_kv_aux_loss)
+
 # MAC transformer
 
 class MemoryAsContextTransformer(Module):
@@ -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,10 @@ 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,
             )
 
             ff = FeedForward(dim = dim, mult = ff_mult)
@@ -445,8 +488,7 @@ class MemoryAsContextTransformer(Module):
         flex_attn_fn = None
 
         if use_flex_attn:
-            block_mask = create_mac_block_mask(seq_len, self.segment_len, self.num_persist_mem_tokens)
-
+            block_mask = create_mac_block_mask(seq_len_with_mem, self.segment_len, self.num_persist_mem_tokens)
             flex_attn_fn = partial(flex_attention, block_mask = block_mask)
 
         # value residual
@@ -461,13 +503,16 @@ class MemoryAsContextTransformer(Module):
 
         x = self.expand_streams(x)
 
-        for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
+        for attn, ff in self.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
+            (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
+            )
 
-            x, values = 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)
 

{titans_pytorch-0.0.55.dist-info → titans_pytorch-0.0.57.dist-info}/METADATA

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

titans_pytorch-0.0.57.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=-WS_SI--_5f1whMaJOH-mYCU37EjYU_iZTurGfs8zgI,15331
+titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
+titans_pytorch-0.0.57.dist-info/METADATA,sha256=rwLIRndtBo22oJt0Xm9xK9zqOYV50Jfo6g7oVrKq7CQ,4457
+titans_pytorch-0.0.57.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.57.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.57.dist-info/RECORD,,

titans_pytorch-0.0.55.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=-VN8bURUaqHXH_96UqGYDhWcfgCaFdHGdM6faVuYDgQ,14159
-titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
-titans_pytorch-0.0.55.dist-info/METADATA,sha256=VYP1B5d9tejIXr7u6ML4cSjvgIDlWYyp5KTyydlUqV8,4457
-titans_pytorch-0.0.55.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.55.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.55.dist-info/RECORD,,