titans-pytorch 0.2.15__tar.gz → 0.2.17__tar.gz

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/PKG-INFO

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

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/pyproject.toml

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

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/tests/test_titans.py

@@ -93,6 +93,34 @@ def test_neural_mem_chaining_chunks():
 
     assert torch.allclose(parallel_retrieved, torch.cat((first_retrieved, second_retrieved, third_retrieved), dim = 1), atol = 1e-5)
 
+def test_neural_mem_chaining_with_weight_residual():
+    mem  = NeuralMemory(
+        dim = 384,
+        dim_head = 64,
+        heads = 2,
+        chunk_size = 64
+    )
+
+    mem2 = NeuralMemory(
+        dim = 384,
+        dim_head = 64,
+        heads = 2,
+        chunk_size = 64
+    )
+
+    seq = torch.randn(2, 256, 384)
+
+    seq, state = mem(seq)
+
+    parallel_retrieved, _ = mem2(seq, prev_weights = state.updates)
+
+    seq_first, seq_second = seq[:, :128], seq[:, 128:]
+
+    first_retrieved, state1 = mem2(seq_first, prev_weights = state.updates)
+    second_retrieved, state2 = mem2(seq_second, state = state1, prev_weights = state.updates)
+
+    assert torch.allclose(parallel_retrieved, torch.cat((first_retrieved, second_retrieved), dim = 1), atol = 1e-6)
+
 def test_neural_mem_chaining_with_batch_size():
     mem  = NeuralMemory(
         dim = 384,

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/titans_pytorch/mac_transformer.py

@@ -491,6 +491,7 @@ class MemoryAsContextTransformer(Module):
         neural_memory_layers: tuple[int, ...] | None = None,
         use_flex_attn = False,
         sliding_window_attn = False,
+        neural_mem_weight_residual = False
     ):
         super().__init__()
 
@@ -524,6 +525,8 @@ class MemoryAsContextTransformer(Module):
 
         neural_memory_layers = default(neural_memory_layers, layers)
 
+        self.neural_mem_weight_residual = neural_mem_weight_residual
+
         # mem, attn, and feedforward layers
 
         for layer in layers:
@@ -739,6 +742,10 @@ class MemoryAsContextTransformer(Module):
 
         value_residual = None
 
+        # neural mem weight residual
+
+        mem_weight_residual = None
+
         # when inferencing, only do one token at a time
 
         if is_inferencing:
@@ -764,8 +771,12 @@ class MemoryAsContextTransformer(Module):
                 retrieved, next_neural_mem_cache = mem.forward(
                     mem_input,
                     state = next(neural_mem_caches, None),
+                    prev_weights = mem_weight_residual
                 )
 
+                if self.neural_mem_weight_residual:
+                    mem_weight_residual = next_neural_mem_cache.updates
+
                 if self.gate_attn_output:
                     attn_out_gates = retrieved.sigmoid()
                 else:

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/titans_pytorch/neural_memory.py

@@ -494,7 +494,8 @@ class NeuralMemory(Module):
         seq,
         weights: dict[str, Tensor] | None = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
-        seq_index = 0
+        seq_index = 0,
+        prev_weights = None
     ):
         batch, seq_len, heads, chunk_size = *seq.shape[:2], self.heads, self.store_chunk_size
 
@@ -560,6 +561,17 @@ class NeuralMemory(Module):
 
         adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = chunk_size)
 
+        # maybe add previous layer weight
+
+        if exists(prev_weights):
+
+            start_index = math.ceil(seq_index / chunk_size)
+            end_index = start_index + num_chunks
+
+            prev_weights = prev_weights.apply(lambda t: t[:, start_index:end_index])
+
+            weights_for_surprise = weights_for_surprise + prev_weights
+
         # flatten batch and time if surprise depends on previous layer memory model
 
         weights_for_surprise = rearrange_dict_values(weights_for_surprise, 'b n ... -> (b n) ...')
@@ -732,6 +744,7 @@ class NeuralMemory(Module):
         seq,
         store_seq = None,
         state: NeuralMemCache | None = None,
+        prev_weights = None
     ):
         if seq.ndim == 2:
             seq = rearrange(seq, 'b d -> b 1 d')
@@ -807,6 +820,7 @@ class NeuralMemory(Module):
                 weights,
                 seq_index = seq_index,
                 past_state = past_state,
+                prev_weights = prev_weights
             )
 
             seq_index = next_neural_mem_state.seq_index

{titans_pytorch-0.2.15 → titans_pytorch-0.2.17}/train_mac.py

@@ -36,8 +36,9 @@ NEURAL_MEM_MOMENTUM = True
 NEURAL_MEM_QK_NORM = True
 NEURAL_MEM_MAX_LR = 1e-1
 WINDOW_SIZE = 32
-NEURAL_MEM_SEGMENT_LEN = 2                      # set smaller for more granularity for learning rate / momentum etc
+NEURAL_MEM_SEGMENT_LEN = 4                      # set smaller for more granularity for learning rate / momentum etc
 NEURAL_MEM_BATCH_SIZE = 128                     # set smaller to update the neural memory weights more often as it traverses the sequence
+NEURAL_MEM_WEIGHT_RESIDUAL = False
 SLIDING_WINDOWS = True
 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
@@ -87,6 +88,7 @@ model = MemoryAsContextTransformer(
     neural_memory_segment_len = NEURAL_MEM_SEGMENT_LEN,
     neural_memory_batch_size = NEURAL_MEM_BATCH_SIZE,
     neural_mem_gate_attn_output = NEURAL_MEM_GATE_ATTN_OUTPUT,
+    neural_mem_weight_residual = NEURAL_MEM_WEIGHT_RESIDUAL,
     use_flex_attn = USE_FLEX_ATTN,
     sliding_window_attn = SLIDING_WINDOWS,
     neural_memory_model = MemoryMLP(