titans-pytorch 0.3.23__py3-none-any.whl → 0.3.25__py3-none-any.whl

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

{titans_pytorch-0.3.23.dist-info → titans_pytorch-0.3.25.dist-info}/METADATA

@@ -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.dist-info → titans_pytorch-0.3.25.dist-info}/RECORD

@@ -2,8 +2,8 @@ titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,29
 titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
 titans_pytorch/mac_transformer.py,sha256=grD327B3OCIy7d23jNUWIoUo1bIgXUqD26dXWCjdi28,25565
 titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=idjjW_K5qZS29isy7pI_pzGDq95vONwaPyXNLnDaj1w,31161
-titans_pytorch-0.3.23.dist-info/METADATA,sha256=KfG0t2AGKAWoifI81y2Bugte3EihJs3w4sjxEYompsw,6817
-titans_pytorch-0.3.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.23.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.23.dist-info/RECORD,,
+titans_pytorch/neural_memory.py,sha256=uh5NbtAAzfPeZPFe7uhgnpUF6qyP0zjP0eXPIgY5pfc,31929
+titans_pytorch-0.3.25.dist-info/METADATA,sha256=SZwazbaNFe1GstoF45zI_aNMpzgXAqv4mOh78gMN5-U,6817
+titans_pytorch-0.3.25.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.25.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.25.dist-info/RECORD,,