titans-pytorch 0.3.15__tar.gz → 0.3.20__tar.gz

{titans_pytorch-0.3.15 → titans_pytorch-0.3.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.3.15
+Version: 0.3.20
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -38,7 +38,7 @@ Requires-Dist: accelerated-scan>=0.2.0
 Requires-Dist: axial-positional-embedding>=0.3.10
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
-Requires-Dist: hyper-connections>=0.1.10
+Requires-Dist: hyper-connections>=0.1.11
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict

{titans_pytorch-0.3.15 → titans_pytorch-0.3.20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.3.15"
+version = "0.3.20"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -29,7 +29,7 @@ dependencies = [
     "axial_positional_embedding>=0.3.10",
     "einops>=0.8.0",
     "einx>=0.3.0",
-    "hyper-connections>=0.1.10",
+    "hyper-connections>=0.1.11",
     "Ninja",
     "rotary-embedding-torch",
     "tensordict",

{titans_pytorch-0.3.15 → titans_pytorch-0.3.20}/tests/test_titans.py

@@ -200,6 +200,7 @@ def test_neural_mem_chaining_with_batch_size():
 @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))
+@pytest.mark.parametrize('neural_mem_qkv_receives_diff_views', (False, True))
 @pytest.mark.parametrize('neural_mem_momentum', (False, True))
 def test_mac(
     seq_len,
@@ -209,6 +210,7 @@ def test_mac(
     neural_mem_segment_len,
     neural_mem_weight_residual,
     neural_mem_batch_size,
+    neural_mem_qkv_receives_diff_views,
     neural_mem_momentum
 ):
     transformer = MemoryAsContextTransformer(
@@ -221,6 +223,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_memory_qkv_receives_diff_views = neural_mem_qkv_receives_diff_views,
         neural_mem_weight_residual = neural_mem_weight_residual,
         neural_memory_kwargs = dict(
             momentum = neural_mem_momentum

{titans_pytorch-0.3.15 → titans_pytorch-0.3.20}/titans_pytorch/mac_transformer.py

@@ -483,6 +483,7 @@ class MemoryAsContextTransformer(Module):
         num_longterm_mem_tokens = 0,
         num_persist_mem_tokens = 0,
         neural_memory_batch_size = None,
+        neural_memory_qkv_receives_diff_views = False,
         dim_head = 64,
         heads = 8,
         ff_mult = 4,
@@ -560,13 +561,14 @@ class MemoryAsContextTransformer(Module):
             mem_hyper_conn = None
 
             if layer in neural_memory_layers:
-                mem_hyper_conn = init_hyper_conn(add_branch_out_to_residual = not neural_mem_gate_attn_output)
+                mem_hyper_conn = init_hyper_conn(add_branch_out_to_residual = not neural_mem_gate_attn_output, num_input_views = 3 if neural_memory_qkv_receives_diff_views else 1)
 
                 mem = NeuralMemory(
                     dim = dim,
                     chunk_size = self.neural_memory_segment_len,
                     batch_size = neural_memory_batch_size,
                     model = deepcopy(neural_memory_model),
+                    qkv_receives_diff_views = neural_memory_qkv_receives_diff_views,
                     accept_weight_residual = neural_mem_weight_residual and not is_first_neural_mem,
                     **neural_memory_kwargs
                 )

{titans_pytorch-0.3.15 → titans_pytorch-0.3.20}/titans_pytorch/neural_memory.py

@@ -231,6 +231,7 @@ class NeuralMemory(Module):
         momentum_order = 1,
         learned_momentum_combine = False,
         learned_combine_include_zeroth = False,
+        qkv_receives_diff_views = False, # to address an issue raised by a phd student (who will be credited if experiments are green). basically the issue raised is that the memory MLP is only learning Wk @ Wv linear mapping and that may not be expressive enough. we will use hyper connections to allow the network to choose different previous layer inputs as keys / values and see if that does anything
         pre_rmsnorm = True,
         post_rmsnorm = False,
         qk_rmsnorm = False,
@@ -265,6 +266,10 @@ class NeuralMemory(Module):
 
         self.assoc_scan = AssocScan(use_accelerated = use_accelerated_scan)
 
+        # key values receiving different views
+
+        self.qkv_receives_diff_views = qkv_receives_diff_views
+
         # norms
 
         self.retrieve_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
@@ -358,7 +363,9 @@ class NeuralMemory(Module):
 
         # keys and values for storing to the model
 
-        self.to_keys_values = Sequential(LinearNoBias(dim, dim_inner * 2), activation)
+        self.to_keys = Sequential(LinearNoBias(dim, dim_inner), activation)
+        self.to_values = Sequential(LinearNoBias(dim, dim_inner), activation)
+
         self.store_memory_loss_fn = store_memory_loss_fn
 
         # `chunk_size` refers to chunk size used for storing to memory model weights
@@ -504,7 +511,14 @@ class NeuralMemory(Module):
         seq_index = 0,
         prev_weights = None
     ):
-        batch, seq_len, heads, chunk_size = *seq.shape[:2], self.heads, self.store_chunk_size
+        if self.qkv_receives_diff_views:
+            _, batch, seq_len = seq.shape[:3]
+        else:
+            batch, seq_len = seq.shape[:2]
+
+        # shapes and variables
+
+        heads, chunk_size = self.heads, self.store_chunk_size
 
         # curtail sequence by multiple of the chunk size
         # only a complete chunk of the sequence provides the memory for the next chunk
@@ -512,7 +526,7 @@ class NeuralMemory(Module):
         round_down_seq_len = round_down_multiple(seq_len, chunk_size)
         num_chunks = round_down_seq_len // chunk_size
 
-        seq, remainder = seq[:, :round_down_seq_len], seq[:, round_down_seq_len:]
+        seq, remainder = seq[..., :round_down_seq_len, :], seq[..., round_down_seq_len:, :]
 
         next_seq_len_index = seq_index + round_down_seq_len
 
@@ -528,10 +542,19 @@ class NeuralMemory(Module):
 
         weights_for_surprise = repeat_dict_values(weights, 'b ... -> b n ...', n = num_chunks)
 
-        # derive learned hparams for optimization of memory network
+        # initial norm
 
         seq = self.store_norm(seq)
 
+        # handle keys and values coming from different sequences from hyper connection
+
+        values_seq = seq
+
+        if self.qkv_receives_diff_views:
+            seq, values_seq = seq
+
+        # derive learned hparams for optimization of memory network
+
         adaptive_lr = self.to_adaptive_step(seq)
         adaptive_lr = self.adaptive_step_transform(adaptive_lr)
 
@@ -555,7 +578,8 @@ class NeuralMemory(Module):
 
         # keys and values
 
-        keys, values = self.to_keys_values(seq).chunk(2, dim = -1)
+        keys = self.to_keys(seq)
+        values = self.to_values(values_seq)
 
         # maybe multi head
 
@@ -796,10 +820,23 @@ class NeuralMemory(Module):
         state: NeuralMemState | None = None,
         prev_weights = None
     ):
-        if seq.ndim == 2:
-            seq = rearrange(seq, 'b d -> b 1 d')
+        is_multi_input = self.qkv_receives_diff_views
+
+        # handle single token
+
+        if seq.ndim == 2 or (is_multi_input and seq.ndim == 3):
+            seq = rearrange(seq, '... b d -> ... b 1 d')
+
+        is_single_token = seq.shape[-2] == 1
 
-        is_single_token = seq.shape[1] == 1
+        # if different views for qkv, then
+
+        if is_multi_input:
+            retrieve_seq, seq = seq[0], seq[1:]
+        else:
+            retrieve_seq = seq
+
+        # handle previous state init
 
         if not exists(state):
             state = (0, None, None, None, None)
@@ -815,8 +852,6 @@ class NeuralMemory(Module):
         if exists(cache_store_seq):
             store_seq = safe_cat((cache_store_seq, store_seq))
 
-        # functions
-
         # compute split sizes of sequence
         # for now manually update weights to last update at the correct boundaries
 
@@ -916,7 +951,7 @@ class NeuralMemory(Module):
             updates = rearrange_dict_values(last_update, 'b ... -> b 1 ...')
 
         retrieved = self.retrieve_memories(
-            seq,
+            retrieve_seq,
             updates
         )