titans-pytorch 0.1.34__py3-none-any.whl → 0.1.35__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -479,7 +479,8 @@ class MemoryAsContextTransformer(Module):
         depth,
         segment_len,
         neural_memory_segment_len = None,
-        neural_mem_gate_attn_output = True,
+        neural_mem_gate_attn_output = False,
+        neural_memory_add_value_residual = False,
         num_longterm_mem_tokens = 0,
         num_persist_mem_tokens = 0,
         dim_head = 64,
@@ -535,6 +536,11 @@ class MemoryAsContextTransformer(Module):
 
         self.weight_tie_memory_model = weight_tie_memory_model
 
+        # value residual learning for neural memory
+
+        is_first_mem = True
+        self.mem_add_value_residual = neural_memory_add_value_residual
+
         # mem, attn, and feedforward layers
 
         for layer in layers:
@@ -564,9 +570,11 @@ class MemoryAsContextTransformer(Module):
                     dim = dim,
                     chunk_size = self.neural_memory_segment_len,
                     model = maybe_copy(neural_memory_model),
+                    accept_value_residual = not is_first_mem and neural_memory_add_value_residual,
                     **neural_memory_kwargs
                 )
 
+                is_first_mem = False
 
             ff = FeedForward(dim = dim, mult = ff_mult)
 
@@ -757,6 +765,8 @@ class MemoryAsContextTransformer(Module):
 
         value_residual = None
 
+        mem_value_residual = None
+
         # aux losses
 
         kv_recon_losses = self.zero
@@ -784,21 +794,28 @@ class MemoryAsContextTransformer(Module):
                 mem_input, add_residual = mem_hyper_conn(x)
 
                 if not is_inferencing:
-                    (retrieved, next_neural_mem_cache), mem_kv_aux_loss = mem(
+                    (retrieved, next_neural_mem_cache, next_mem_value_residual), mem_kv_aux_loss = mem(
                         mem_input,
                         return_aux_kv_loss = True,
+                        return_values = True,
+                        value_residual = mem_value_residual,
                         prev_layer_updates = neural_memory_updates
                     )
 
                     kv_recon_losses = kv_recon_losses + mem_kv_aux_loss
 
                 else:
-                    retrieved, next_neural_mem_cache = mem.forward_inference(
+                    (retrieved, next_neural_mem_cache, next_mem_value_residual) = mem.forward_inference(
                         mem_input,
                         state = next(neural_mem_caches, None),
+                        return_values = True,
+                        value_residual = mem_value_residual,
                         prev_layer_updates = neural_memory_updates
                     )
 
+                if self.mem_add_value_residual:
+                    mem_value_residual = next_mem_value_residual
+
                 if weight_tie_memory_model:
                     neural_memory_updates = next_neural_mem_cache.updates
 

titans_pytorch/neural_memory.py

@@ -822,7 +822,9 @@ class NeuralMemory(Module):
         self,
         token: Tensor,
         state = None,
-        prev_layer_updates: dict[str, Tensor] | None = None
+        prev_layer_updates: dict[str, Tensor] | None = None,
+        return_values = False,
+        value_residual = None,
     ):
 
         # unpack previous state
@@ -870,11 +872,12 @@ class NeuralMemory(Module):
 
         if store_seq_cache_len == self.chunk_size:
 
-            next_updates, next_states, _ = self.store_memories(
+            next_updates, next_states, values = self.store_memories(
                 cache_store_seq,
                 weights,
                 past_state = past_states,
                 prev_layer_updates = prev_layer_updates,
+                value_residual = value_residual
             )
 
             updates = next_updates
@@ -888,7 +891,12 @@ class NeuralMemory(Module):
 
         next_state = NeuralMemCache(curr_seq_len, cache_store_seq, next_states, updates)
 
-        return retrieved, next_state
+        output = (retrieved, next_state)
+
+        if return_values:
+            output = (*output, values)
+
+        return output
 
     def forward(
         self,
@@ -900,6 +908,7 @@ class NeuralMemory(Module):
         chunk_size = None,
         store_chunk_size = None,
         return_values = False,
+        value_residual = None,
         return_next_state = False,
         prev_layer_updates: dict[str, Tensor] | None = None
     ):
@@ -933,6 +942,7 @@ class NeuralMemory(Module):
             mem_model_weights,
             chunk_size = store_chunk_size,
             prev_layer_updates = prev_layer_updates,
+            value_residual = value_residual,
             return_aux_kv_loss = True
         )
 

{titans_pytorch-0.1.34.dist-info → titans_pytorch-0.1.35.dist-info}/METADATA

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

titans_pytorch-0.1.35.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=rMT99CPQFH4Gudp0FmVPWGKfhBf6xksGEaFEcOVdqjs,230
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=KBwo-Fr_fDzVaAa7xg1ggEpNlE4vRUoGMEjB-I2ZWTU,26463
+titans_pytorch/neural_memory.py,sha256=wFOLFe3ViXiQfNvUiAGJ6BfiaDhr0BYDRDnLNMHWQhU,28938
+titans_pytorch-0.1.35.dist-info/METADATA,sha256=5e5qPt4hAOhxhDWqdjutjJuUmht44zYq_KqgagKjqxE,6826
+titans_pytorch-0.1.35.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.35.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.35.dist-info/RECORD,,

titans_pytorch-0.1.34.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=rMT99CPQFH4Gudp0FmVPWGKfhBf6xksGEaFEcOVdqjs,230
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=JvA4mhQaW9LD4j6boRUfLfjyzDCtjqybIr4Ajeio8n8,25708
-titans_pytorch/neural_memory.py,sha256=nNAxhkubuHCGs3bty_eA_yBhWqepPZJgKKvkWXO6IK4,28653
-titans_pytorch-0.1.34.dist-info/METADATA,sha256=pVgjCX_YTT9_5WPcFfXpoaBvzrg1-esvwS0kPpeJAYU,6826
-titans_pytorch-0.1.34.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.34.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.34.dist-info/RECORD,,