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

titans_pytorch/mac_transformer.py

@@ -292,12 +292,13 @@ class NeuralMemoryGatingWrapper(Module):
         self,
         dim,
         attn: SegmentedAttention,
-        neural_mem: NeuralMemory | None = None
+        neural_mem: NeuralMemory | None = None,
+        gate_attn_output = True
     ):
         super().__init__()
         self.attn = attn
         self.neural_mem = neural_mem
-        self.to_gates = nn.Linear(dim, dim) if exists(neural_mem) else None
+        self.gate_attn_output = gate_attn_output
 
     def forward(
         self,
@@ -313,21 +314,19 @@ class NeuralMemoryGatingWrapper(Module):
 
         # 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)
+        retrieved, kv_aux_loss = mem(seq, return_aux_kv_loss = True)
 
-        seq = seq + retrieved
+        if not self.gate_attn_output:
+            seq = seq + retrieved
 
         # attention
 
         attn_out, values = self.attn(seq, *args, **kwargs)
 
-        # another retrieve, but this time gate the attention output
+        if self.gate_attn_output:
+            attn_out = attn_out * retrieved.sigmoid()
 
-        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)
+        return (attn_out, values), kv_aux_loss
 
 # MAC transformer
 
@@ -340,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,
@@ -414,6 +414,7 @@ class MemoryAsContextTransformer(Module):
                 dim,
                 attn = attn,
                 neural_mem = mem,
+                gate_attn_output = neural_mem_gate_attn_output
             )
 
             ff = FeedForward(dim = dim, mult = ff_mult)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.57
+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.57.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=-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,,