titans-pytorch 0.2.18__tar.gz → 0.2.20__tar.gz

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.2.18
+Version: 0.2.20
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -204,6 +204,6 @@ $ python train_mac.py
     eprint  = {2501.12352},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url = {https://arxiv.org/abs/2501.12352},
+    url     = {https://arxiv.org/abs/2501.12352},
 }
 ```

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/README.md

@@ -150,6 +150,6 @@ $ python train_mac.py
     eprint  = {2501.12352},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url = {https://arxiv.org/abs/2501.12352},
+    url     = {https://arxiv.org/abs/2501.12352},
 }
 ```

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/pyproject.toml

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

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/tests/test_titans.py

@@ -105,7 +105,8 @@ def test_neural_mem_chaining_with_weight_residual():
         dim = 384,
         dim_head = 64,
         heads = 2,
-        chunk_size = 64
+        chunk_size = 64,
+        accept_weight_residual = True
     )
 
     seq = torch.randn(2, 256, 384)

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/titans_pytorch/mac_transformer.py

@@ -560,7 +560,7 @@ class MemoryAsContextTransformer(Module):
                     chunk_size = self.neural_memory_segment_len,
                     batch_size = neural_memory_batch_size,
                     model = deepcopy(neural_memory_model),
-                    learned_weight_residual = neural_mem_weight_residual and not is_first_neural_mem,
+                    accept_weight_residual = neural_mem_weight_residual and not is_first_neural_mem,
                     **neural_memory_kwargs
                 )
 

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/titans_pytorch/neural_memory.py

@@ -66,6 +66,12 @@ def xnor(x, y):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def tuple_index_set(t: tuple, index, value):
+    klass = type(t)
+    t = list(t)
+    t[index] = value
+    return klass(*t)
+
 def safe_cat(inputs, dim = -2):
     inputs = tuple(filter(exists, inputs))
 
@@ -296,7 +302,7 @@ class NeuralMemory(Module):
         init_adaptive_step_bias = None,
         init_momentum_bias = None,
         init_decay_bias = None,
-        learned_weight_residual = False,
+        accept_weight_residual = False,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -445,7 +451,7 @@ class NeuralMemory(Module):
             nn.Linear(dim, heads),
             Rearrange('b n h -> b h n'),
             nn.Sigmoid()
-        ) if learned_weight_residual else None
+        ) if accept_weight_residual else None
 
         # allow for softclamp the gradient norms for storing memories
 
@@ -584,7 +590,6 @@ class NeuralMemory(Module):
             if exists(self.to_learned_weight_residual_mix):
                 mix = self.to_learned_weight_residual_mix(chunked_seq)
                 mix = rearrange(mix, 'b h n -> (b h) n')
-
                 prev_weights = prev_weights.apply(lambda t: einx.multiply('bh n, bh n ... -> bh n ...', mix, t))
 
             weights_for_surprise = weights_for_surprise + prev_weights
@@ -854,9 +859,9 @@ class NeuralMemory(Module):
 
             weights = last_update
 
-            next_neural_mem_state = list(next_neural_mem_state)
-            next_neural_mem_state[1] = last_update
-            next_neural_mem_state = NeuralMemCache(*next_neural_mem_state)
+            next_neural_mem_state = tuple_index_set(next_neural_mem_state, 1, last_update)
+
+        next_neural_mem_state = tuple_index_set(next_neural_mem_state, -1, updates)
 
         # retrieve
 

{titans_pytorch-0.2.18 → titans_pytorch-0.2.20}/train_mac.py

@@ -38,10 +38,10 @@ NEURAL_MEM_MAX_LR = 1e-1
 WINDOW_SIZE = 32
 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
+NEURAL_MEM_WEIGHT_RESIDUAL = True               # learning to accept contributions from the weights of the previous neural mem layer brings about significant improvements. this was improvised and not in the paper, but inspired by the value residual learning free lunch paper
 
 # experiment related