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

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

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

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

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

{titans_pytorch-0.2.16 → titans_pytorch-0.2.18}/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,
@@ -121,6 +149,7 @@ def test_neural_mem_chaining_with_batch_size():
 @pytest.mark.parametrize('num_longterm_mem_tokens', (0, 16))
 @pytest.mark.parametrize('neural_mem_gate_attn_output', (False, True))
 @pytest.mark.parametrize('neural_mem_segment_len', (8, 16))
+@pytest.mark.parametrize('neural_mem_weight_residual', (False, True))
 @pytest.mark.parametrize('neural_mem_batch_size', (None, 64))
 def test_mac(
     seq_len,
@@ -128,7 +157,8 @@ def test_mac(
     num_longterm_mem_tokens,
     neural_mem_gate_attn_output,
     neural_mem_segment_len,
-    neural_mem_batch_size
+    neural_mem_weight_residual,
+    neural_mem_batch_size,
 ):
     transformer = MemoryAsContextTransformer(
         num_tokens = 256,
@@ -140,6 +170,7 @@ def test_mac(
         neural_mem_gate_attn_output = neural_mem_gate_attn_output,
         neural_memory_segment_len = neural_mem_segment_len,
         neural_memory_batch_size = neural_mem_batch_size,
+        neural_mem_weight_residual = neural_mem_weight_residual
     )
 
     x = torch.randint(0, 256, (1, seq_len))

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

@@ -491,7 +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
+        neural_mem_weight_residual = False,
     ):
         super().__init__()
 
@@ -525,7 +525,10 @@ class MemoryAsContextTransformer(Module):
 
         neural_memory_layers = default(neural_memory_layers, layers)
 
+        # weight residual related
+
         self.neural_mem_weight_residual = neural_mem_weight_residual
+        is_first_neural_mem = True
 
         # mem, attn, and feedforward layers
 
@@ -557,9 +560,12 @@ 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,
                     **neural_memory_kwargs
                 )
 
+                is_first_neural_mem = False
+
             ff = FeedForward(dim = dim, mult = ff_mult)
 
             self.layers.append(ModuleList([

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

@@ -296,6 +296,7 @@ class NeuralMemory(Module):
         init_adaptive_step_bias = None,
         init_momentum_bias = None,
         init_decay_bias = None,
+        learned_weight_residual = False,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -438,6 +439,14 @@ class NeuralMemory(Module):
 
         self.max_mem_layer_modulation = max_mem_layer_modulation
 
+        # learned weight residual
+
+        self.to_learned_weight_residual_mix = Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b n h -> b h n'),
+            nn.Sigmoid()
+        ) if learned_weight_residual else None
+
         # allow for softclamp the gradient norms for storing memories
 
         self.max_grad_norm = max_grad_norm
@@ -563,8 +572,21 @@ class NeuralMemory(Module):
 
         # maybe add previous layer weight
 
+        assert xnor(exists(self.to_learned_weight_residual_mix), exists(prev_weights))
+
         if exists(prev_weights):
-            prev_weights = prev_weights.apply(lambda t: t[:, -1:])
+
+            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])
+
+            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
 
         # flatten batch and time if surprise depends on previous layer memory model

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

@@ -38,7 +38,7 @@ 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 = True
+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