titans-pytorch 0.3.15__py3-none-any.whl → 0.3.20__py3-none-any.whl

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/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
         )
 

{titans_pytorch-0.3.15.dist-info → titans_pytorch-0.3.20.dist-info}/METADATA

@@ -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.20.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
+titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
+titans_pytorch/mac_transformer.py,sha256=M45tLp_CzYVm9R33zkUHmWdzIuaNxNYQLCTAtyMechg,25294
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=hLMDbLosxgQZvvfuJM3X7w23npcoL9AWBQfbtQHOSiA,30637
+titans_pytorch-0.3.20.dist-info/METADATA,sha256=JILf0r0bT4KXAZD4V2nmYRIr5Y7bGGjjNlpfIuT0UNs,6817
+titans_pytorch-0.3.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.20.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.20.dist-info/RECORD,,

titans_pytorch-0.3.15.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
-titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
-titans_pytorch/mac_transformer.py,sha256=HIB3S3JBA8Fe1EBITvDZSHXtn-1_fF1rwlw-MzqagKY,25085
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=VmUAS1xOM0ZfearWIzQrX_P7HI69viuwrg9M7BQByeE,29349
-titans_pytorch-0.3.15.dist-info/METADATA,sha256=RPw9JXenAI7cGpVP3hQZlj0OA5-xsvXvXHvxyhWdgpg,6817
-titans_pytorch-0.3.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.15.dist-info/RECORD,,