titans-pytorch 0.1.34__tar.gz → 0.1.36__tar.gz

{titans_pytorch-0.1.34 → titans_pytorch-0.1.36}/PKG-INFO

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

{titans_pytorch-0.1.34 → titans_pytorch-0.1.36}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.34"
+version = "0.1.36"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.1.34 → titans_pytorch-0.1.36}/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-0.1.34 → titans_pytorch-0.1.36}/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
     ):
@@ -912,6 +921,9 @@ class NeuralMemory(Module):
 
             out = (out, next_store_state)
 
+            if return_values:
+                out = (*out, self.zero)
+
             if not return_aux_kv_loss:
                 return out
 
@@ -933,6 +945,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 → titans_pytorch-0.1.36}/train_mac.py

@@ -9,7 +9,8 @@ from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
 from adam_atan2_pytorch import AdoptAtan2
-from titans_pytorch import MemoryAsContextTransformer
+
+from titans_pytorch import MemoryAsContextTransformer, MemoryMLP
 
 # constants
 
@@ -29,13 +30,16 @@ SEQ_LEN = 512
 NEURAL_MEMORY_DEPTH = 2
 NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
-NEURAL_MEM_LAYERS = (2, 4)
-NEURAL_MEM_GATE_ATTN_OUTPUT = True
+NEURAL_MEM_LAYERS = (2, 4, 6)               # layers 2, 4, 6 have neural memory, can add more
+NEURAL_MEM_GATE_ATTN_OUTPUT = False
 NEURAL_MEM_MOMENTUM = True
+NEURAL_MEM_ADD_VALUE_RESIDUAL = True,
 WINDOW_SIZE = 32
 NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2 # set smaller for more granularity for learning rate / momentum etc
 SLIDING_WINDOWS = True
+WEIGHT_TIE_MEMORY_MODEL = True  # set to have memory MLP shared across layers
 STORE_ATTN_POOL_CHUNKS = True # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
+MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
 KV_RECON_LOSS_WEIGHT = 0.
 
 # experiment related
@@ -84,15 +88,19 @@ model = MemoryAsContextTransformer(
     aux_kv_recon_loss_weight = KV_RECON_LOSS_WEIGHT,
     use_flex_attn = USE_FLEX_ATTN,
     sliding_window_attn = SLIDING_WINDOWS,
+    weight_tie_memory_model = WEIGHT_TIE_MEMORY_MODEL,
+    neural_memory_add_value_residual = NEURAL_MEM_ADD_VALUE_RESIDUAL,
+    neural_memory_model = MemoryMLP(
+        dim = 64,
+        depth = NEURAL_MEMORY_DEPTH
+    ),
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,
         attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
         momentum = NEURAL_MEM_MOMENTUM,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
-        default_model_kwargs = dict(
-            depth = NEURAL_MEMORY_DEPTH,
-        )
+        per_parameter_lr_modulation = MEMORY_MODEL_PER_LAYER_LEARNED_LR
     )
 ).cuda()