titans-pytorch 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -217,7 +217,8 @@ class SegmentedAttention(Module):
         self,
         seq,
         value_residual = None,
-        flex_attn_fn: Callable | None = None
+        flex_attn_fn: Callable | None = None,
+        output_gating = None
     ):
 
         assert not (exists(value_residual) ^ exists(self.to_learned_v_mix))
@@ -267,6 +268,9 @@ class SegmentedAttention(Module):
 
         out = self.to_out(out)
 
+        if exists(output_gating):
+            out = out * output_gating
+
         return out, orig_v
 
     def forward(
@@ -274,10 +278,11 @@ class SegmentedAttention(Module):
         seq,
         value_residual = None,
         flex_attn_fn: Callable | None = None,
-        disable_flex_attn = False
+        disable_flex_attn = False,
+        output_gating = None
     ):
         if seq.is_cuda and self.use_flex_attn and not disable_flex_attn:
-            return self.forward_flex(seq, value_residual, flex_attn_fn)
+            return self.forward_flex(seq, value_residual, flex_attn_fn, output_gating = output_gating)
 
         assert not (exists(value_residual) ^ exists(self.to_learned_v_mix))
 
@@ -361,50 +366,10 @@ class SegmentedAttention(Module):
 
         out = inverse_segment(out)
 
-        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.
+        if exists(output_gating):
+            out = out * output_gating
 
-        # 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
+        return out, orig_v
 
 # MAC transformer
 
@@ -494,16 +459,10 @@ class MemoryAsContextTransformer(Module):
                     **neural_memory_kwargs
                 )
 
-            attn = NeuralMemoryGatingWrapper(
-                dim,
-                attn = attn,
-                neural_mem = mem,
-                gate_attn_output = neural_mem_gate_attn_output
-            )
-
             ff = FeedForward(dim = dim, mult = ff_mult)
 
             self.layers.append(ModuleList([
+                init_hyper_conn(dim = dim, branch = mem, add_branch_out_to_residual = not neural_mem_gate_attn_output) if exists(mem) else None,
                 init_hyper_conn(dim = dim, branch = attn),
                 init_hyper_conn(dim = dim, branch = ff)
             ]))
@@ -512,6 +471,10 @@ class MemoryAsContextTransformer(Module):
 
         self.to_logits = LinearNoBias(dim, num_tokens)
 
+        # whether to gate the attention output with the retrieved memories
+
+        self.gate_attn_output = neural_mem_gate_attn_output
+
         # auxiliary loss on kv recon
 
         self.has_aux_kv_recon_loss = aux_kv_recon_loss_weight > 0.
@@ -652,19 +615,34 @@ class MemoryAsContextTransformer(Module):
 
         x = self.expand_streams(x)
 
-        for attn, ff in self.layers:
+        for mem, attn, ff in self.layers:
+
+            retrieved = None
+            attn_out_gates = None
+
+            if exists(mem):
+                retrieved, mem_kv_aux_loss = mem(x, return_aux_kv_loss = True)
+                kv_recon_losses = kv_recon_losses + mem_kv_aux_loss
 
-            (x, values), maybe_mem_kv_aux_loss = attn(
+                if self.gate_attn_output:
+                    attn_out_gates = retrieved.sigmoid()
+                else:
+                    seq = retrieved
+
+            # attention
+
+            x, values = attn(
                 x,
                 value_residual = value_residual,
                 disable_flex_attn = disable_flex_attn,
-                flex_attn_fn = flex_attn_fn
+                flex_attn_fn = flex_attn_fn,
+                output_gating = attn_out_gates
             )
 
-            kv_recon_losses = kv_recon_losses + maybe_mem_kv_aux_loss
-
             value_residual = default(value_residual, values)
 
+            # feedforward
+
             x = ff(x)
 
         x = self.reduce_streams(x)

{titans_pytorch-0.1.17.dist-info → titans_pytorch-0.1.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.17
+Version: 0.1.18
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -38,7 +38,7 @@ Requires-Dist: accelerated-scan>=0.2.0
 Requires-Dist: axial-positional-embedding>=0.3.9
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
-Requires-Dist: hyper-connections>=0.1.8
+Requires-Dist: hyper-connections>=0.1.9
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict

titans_pytorch-0.1.18.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=9HHZ-KyWrnIXMG68WdevnuX_yigVoxKeikT0J5yDFY8,19957
+titans_pytorch/titans.py,sha256=R0e25ly2uTHkHSZEb-9Eqb0DqtFq8wFBB8iH1T6bYVg,22240
+titans_pytorch-0.1.18.dist-info/METADATA,sha256=w5G49IHPqC1eqCRhAgJYPLHYVmaO8zucI-GGVPemoKg,6340
+titans_pytorch-0.1.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.18.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.18.dist-info/RECORD,,

titans_pytorch-0.1.17.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=ajsX-djEUzstig5n99yF_NimRzKNfv0MSz-EIV-Fe1A,20393
-titans_pytorch/titans.py,sha256=R0e25ly2uTHkHSZEb-9Eqb0DqtFq8wFBB8iH1T6bYVg,22240
-titans_pytorch-0.1.17.dist-info/METADATA,sha256=E9nwWCKZLSqT9Mr85nrJQzinYpKZnkLeexeaYyOIqrU,6340
-titans_pytorch-0.1.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.17.dist-info/RECORD,,