titans-pytorch 0.3.23__tar.gz → 0.3.25__tar.gz

{titans_pytorch-0.3.23 → titans_pytorch-0.3.25}/PKG-INFO

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

{titans_pytorch-0.3.23 → titans_pytorch-0.3.25}/pyproject.toml

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

{titans_pytorch-0.3.23 → titans_pytorch-0.3.25}/tests/test_titans.py

@@ -29,10 +29,12 @@ def torch_default_dtype(dtype):
 @pytest.mark.parametrize('chunk_size, attn_pool_chunks', ((64, True), (64, False), (1, False)))
 @pytest.mark.parametrize('momentum', (False, True))
 @pytest.mark.parametrize('qk_rmsnorm', (False, True))
+@pytest.mark.parametrize('heads', (1, 4))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 @pytest.mark.parametrize('num_kv_per_token', (1, 2))
 @pytest.mark.parametrize('per_parameter_lr_modulation', (False, True))
 @pytest.mark.parametrize('per_head_learned_parameters', (False, True))
+@pytest.mark.parametrize('test_store_mask', (False, True))
 def test_titans(
     seq_len,
     silu,
@@ -40,10 +42,12 @@ def test_titans(
     chunk_size,
     momentum,
     qk_rmsnorm,
+    heads,
     max_grad_norm,
     num_kv_per_token,
     per_parameter_lr_modulation,
-    per_head_learned_parameters
+    per_head_learned_parameters,
+    test_store_mask
 ):
     mem = NeuralMemory(
         dim = 16,
@@ -54,12 +58,19 @@ def test_titans(
         num_kv_per_token = num_kv_per_token,
         momentum = momentum,
         qk_rmsnorm = qk_rmsnorm,
+        heads = heads,
         per_parameter_lr_modulation = per_parameter_lr_modulation,
         per_head_learned_parameters = per_head_learned_parameters
     )
 
     seq = torch.randn(2, seq_len, 16)
-    retrieved, _ = mem(seq)
+
+    store_mask = None
+
+    if test_store_mask:
+        store_mask = torch.randint(0, 2, (2, seq_len)).bool()
+
+    retrieved, _ = mem(seq, store_mask = store_mask)
 
     assert seq.shape == retrieved.shape
 

{titans_pytorch-0.3.23 → titans_pytorch-0.3.25}/titans_pytorch/neural_memory.py

@@ -289,6 +289,8 @@ class NeuralMemory(Module):
         self.heads = heads
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.split_kv_heads = Rearrange('b n (h u d) -> b h (n u) d', h = heads, u = num_kv_per_token)
+
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
         self.combine_heads = LinearNoBias(dim_inner, dim) if heads > 1 else nn.Identity()
 
@@ -522,7 +524,8 @@ class NeuralMemory(Module):
         weights: dict[str, Tensor] | None = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
         seq_index = 0,
-        prev_weights = None
+        prev_weights = None,
+        mask: Tensor | None = None,
     ):
         if self.qkv_receives_diff_views:
             _, batch, seq_len = seq.shape[:3]
@@ -596,22 +599,28 @@ class NeuralMemory(Module):
 
         # maybe multi head
 
-        keys, values = map(self.split_heads, (keys, values))
-
-        batch = keys.shape[0]
+        keys, values = map(self.split_kv_heads, (keys, values))
 
-        # take care of chunking
+        # maybe keys rmsnorm
 
-        keys, values = tuple(rearrange(t, 'b h (n c) (u d) -> (b h n) (c u) d', c = chunk_size, u = num_updates) for t in (keys, values))
+        keys = self.k_norm(keys)
 
-        # maybe qk rmsnorm
+        # take care of chunking
 
-        keys = self.k_norm(keys)
+        keys, values = tuple(rearrange(t, 'b h (n c u) d -> (b h n) (c u) d', c = chunk_size, u = num_updates) for t in (keys, values))
 
         # adaptive lr
 
         adaptive_lr = rearrange(adaptive_lr, 'b (n c u) -> (b n) (c u)', c = chunk_size, u = num_updates)
 
+        # optionally a storing memories mask can be passed in. if False, will set the learning rate to 0. for those positions
+
+        if exists(mask):
+            mask = mask[..., :round_down_seq_len]
+            mask = repeat(mask, 'b (n c) -> (b h n) (c u)', h = heads, u = num_updates, c = chunk_size)
+
+            adaptive_lr = torch.where(mask, adaptive_lr, 0.)
+
         # maybe add previous layer weight
 
         assert xnor(exists(self.to_learned_weight_residual_mix), exists(prev_weights))
@@ -833,7 +842,8 @@ class NeuralMemory(Module):
         seq,
         store_seq = None,
         state: NeuralMemState | None = None,
-        prev_weights = None
+        prev_weights = None,
+        store_mask: Tensor | None = None
     ):
         is_multi_input = self.qkv_receives_diff_views
 
@@ -910,6 +920,11 @@ class NeuralMemory(Module):
 
         store_seqs = store_seq.split(split_sizes, dim = -2)
 
+        if exists(store_mask):
+            store_masks = store_mask.split(split_sizes, dim = -1)
+        else:
+            store_masks = (None,) * len(split_sizes)
+
         # whether to allow network to slowly adjust from initial weight throughout (residual path) to fully updating weights every batch
 
         gate = None
@@ -917,7 +932,7 @@ class NeuralMemory(Module):
         if exists(self.transition_gate):
             gate = self.transition_gate.sigmoid()
 
-        for ind, store_seq_chunk in enumerate(store_seqs):
+        for ind, (store_seq_chunk, maybe_store_mask) in enumerate(zip(store_seqs, store_masks)):
             is_last = ind == (len(store_seqs) - 1)
 
             # store
@@ -927,7 +942,8 @@ class NeuralMemory(Module):
                 weights,
                 seq_index = seq_index,
                 past_state = past_state,
-                prev_weights = prev_weights
+                prev_weights = prev_weights,
+                mask = maybe_store_mask
             )
 
             weights = next_neural_mem_state.weights