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

titans_pytorch/mac_transformer.py

@@ -483,7 +483,7 @@ class MemoryAsContextTransformer(Module):
         num_longterm_mem_tokens = 0,
         num_persist_mem_tokens = 0,
         neural_memory_batch_size = None,
-        neural_memory_kv_receives_diff_views = False,
+        neural_memory_qkv_receives_diff_views = False,
         dim_head = 64,
         heads = 8,
         ff_mult = 4,
@@ -561,14 +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, num_input_views = 2 if neural_memory_kv_receives_diff_views else 1)
+                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),
-                    kv_receives_diff_views = neural_memory_kv_receives_diff_views,
+                    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,7 +231,7 @@ class NeuralMemory(Module):
         momentum_order = 1,
         learned_momentum_combine = False,
         learned_combine_include_zeroth = False,
-        kv_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
+        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,
@@ -268,7 +268,7 @@ class NeuralMemory(Module):
 
         # key values receiving different views
 
-        self.kv_receives_diff_views = kv_receives_diff_views
+        self.qkv_receives_diff_views = qkv_receives_diff_views
 
         # norms
 
@@ -511,7 +511,7 @@ class NeuralMemory(Module):
         seq_index = 0,
         prev_weights = None
     ):
-        if self.kv_receives_diff_views:
+        if self.qkv_receives_diff_views:
             _, batch, seq_len = seq.shape[:3]
         else:
             batch, seq_len = seq.shape[:2]
@@ -550,7 +550,7 @@ class NeuralMemory(Module):
 
         values_seq = seq
 
-        if self.kv_receives_diff_views:
+        if self.qkv_receives_diff_views:
             seq, values_seq = seq
 
         # derive learned hparams for optimization of memory network
@@ -820,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
 
-        is_single_token = seq.shape[1] == 1
+        # 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
+
+        # 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)
@@ -839,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
 
@@ -939,11 +950,8 @@ class NeuralMemory(Module):
             last_update, _ = next_neural_mem_state.states
             updates = rearrange_dict_values(last_update, 'b ... -> b 1 ...')
 
-        if self.kv_receives_diff_views:
-            seq = seq[0]
-
         retrieved = self.retrieve_memories(
-            seq,
+            retrieve_seq,
             updates
         )
 

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

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

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.19.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=4N8WcoqPYNhMGGQAZjpm-djVsLnU7VADH_l06qFPuOk,25290
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=KKj8r3oRoDYf-vFQFnI4Rd4DfMH1f5QFs4vdcS35og8,30374
-titans_pytorch-0.3.19.dist-info/METADATA,sha256=5h4f5gsO1emX5LEwe8cgpH35Rtjydz2UTlpK4DKSntI,6817
-titans_pytorch-0.3.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.19.dist-info/RECORD,,