titans-pytorch 0.1.33__tar.gz → 0.1.35__tar.gz

{titans_pytorch-0.1.33 → titans_pytorch-0.1.35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.33
+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.33 → titans_pytorch-0.1.35}/pyproject.toml

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

{titans_pytorch-0.1.33 → titans_pytorch-0.1.35}/tests/test_titans.py

@@ -86,25 +86,32 @@ def test_retrieve_store_diff_seq():
     assert retrieve_seq.shape == retrieved.shape
 
 def test_weight_tied_mlp_neural_mem():
+    from titans_pytorch import MemoryMLP
+
+    mlp = MemoryMLP(64, depth = 2)
+
     mem  = NeuralMemory(
         dim = 384,
         dim_head = 64,
         heads = 2,
-        chunk_size = 2
+        chunk_size = 2,
+        model = mlp
     )
 
     mem2 = NeuralMemory(
         dim = 384,
         dim_head = 64,
         heads = 2,
-        chunk_size = 2
+        chunk_size = 2,
+        model = mlp
     )
 
     mem3 = NeuralMemory(
         dim = 384,
         dim_head = 64,
         heads = 2,
-        chunk_size = 2
+        chunk_size = 2,
+        model = mlp
     )
 
     seq = torch.randn(2, 128, 384)
@@ -113,6 +120,31 @@ def test_weight_tied_mlp_neural_mem():
     seq, cache2 = mem2(seq, prev_layer_updates = cache.updates)
     seq, cache3 = mem3(seq, prev_layer_updates = cache2.updates)
 
+def test_mac_with_weight_tied_neural_mem():
+    from titans_pytorch import MemoryMLP, MemoryAsContextTransformer
+
+    transformer = MemoryAsContextTransformer(
+        num_tokens = 256,
+        dim = 256,
+        depth = 2,
+        segment_len = 2,
+        num_persist_mem_tokens = 0,
+        num_longterm_mem_tokens = 2,
+        neural_memory_segment_len = 2,
+        sliding_window_attn = True,
+        neural_memory_layers = (1, 2),
+        neural_memory_model = MemoryMLP(256, depth = 1),
+        num_residual_streams = 4,
+        weight_tie_memory_model = True,
+        neural_mem_gate_attn_output = True,
+    )
+
+
+    ids = torch.randint(0, 256, (1, 1023))
+    logits = transformer(ids)
+
+    assert logits.shape == (1, 1023, 256)
+
 def test_overriding_chunk_size():
     mem = NeuralMemory(
         dim = 384,

{titans_pytorch-0.1.33 → titans_pytorch-0.1.35}/titans_pytorch/mac_transformer.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from typing import Callable
 
 from math import ceil
+from copy import deepcopy
 from functools import partial
 from collections import namedtuple
 
@@ -478,18 +479,21 @@ 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,
         heads = 8,
         ff_mult = 4,
         num_residual_streams = 4,
+        neural_memory_model: Module | None = None,
         neural_memory_kwargs: dict = dict(),
         neural_memory_layers: tuple[int, ...] | None = None,
         aux_kv_recon_loss_weight = 0.,
         use_flex_attn = False,
-        sliding_window_attn = False
+        sliding_window_attn = False,
+        weight_tie_memory_model = False
     ):
         super().__init__()
 
@@ -523,6 +527,20 @@ class MemoryAsContextTransformer(Module):
 
         neural_memory_layers = default(neural_memory_layers, layers)
 
+        # weight tying neural memory model
+
+        maybe_copy = deepcopy if not weight_tie_memory_model else identity
+
+        if weight_tie_memory_model:
+            assert exists(neural_memory_model), '`neural_memory_model` must be explicitly set'
+
+        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:
@@ -551,9 +569,12 @@ class MemoryAsContextTransformer(Module):
                 mem = NeuralMemory(
                     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)
 
@@ -683,7 +704,7 @@ class MemoryAsContextTransformer(Module):
 
         # math
 
-        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens, attn_window_size = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens, self.attn_window_size
+        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens, attn_window_size, weight_tie_memory_model = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens, self.attn_window_size, self.weight_tie_memory_model
 
         seq_len_with_mem = self.seq_len_with_longterm_mem(seq_len)
 
@@ -736,10 +757,16 @@ class MemoryAsContextTransformer(Module):
         next_kv_caches = []
         next_neural_mem_caches = []
 
+        # weight tied neural memory
+
+        neural_memory_updates = None
+
         # value residual
 
         value_residual = None
 
+        mem_value_residual = None
+
         # aux losses
 
         kv_recon_losses = self.zero
@@ -767,19 +794,31 @@ 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_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)
+                        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
+
                 if self.gate_attn_output:
                     attn_out_gates = retrieved.sigmoid()
                 else:

{titans_pytorch-0.1.33 → titans_pytorch-0.1.35}/titans_pytorch/neural_memory.py

@@ -822,6 +822,9 @@ class NeuralMemory(Module):
         self,
         token: Tensor,
         state = None,
+        prev_layer_updates: dict[str, Tensor] | None = None,
+        return_values = False,
+        value_residual = None,
     ):
 
         # unpack previous state
@@ -863,12 +866,18 @@ class NeuralMemory(Module):
         else:
             updates = updates.apply(lambda t: t[:, -1:])
 
+        if exists(prev_layer_updates):
+            prev_layer_updates = TensorDict(prev_layer_updates)
+            prev_layer_updates = prev_layer_updates.apply(lambda t: t[:, -1:])
+
         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
+                past_state = past_states,
+                prev_layer_updates = prev_layer_updates,
+                value_residual = value_residual
             )
 
             updates = next_updates
@@ -882,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,
@@ -894,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
     ):
@@ -927,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.33 → titans_pytorch-0.1.35}/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()