titans-pytorch 0.3.10__tar.gz → 0.3.11__tar.gz

{titans_pytorch-0.3.10 → titans_pytorch-0.3.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.3.10
+Version: 0.3.11
 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.9
+Requires-Dist: hyper-connections>=0.1.10
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict

{titans_pytorch-0.3.10 → titans_pytorch-0.3.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.3.10"
+version = "0.3.11"
 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.9",
+    "hyper-connections>=0.1.10",
     "Ninja",
     "rotary-embedding-torch",
     "tensordict",

{titans_pytorch-0.3.10 → titans_pytorch-0.3.11}/titans_pytorch/mac_transformer.py

@@ -62,6 +62,7 @@ from rotary_embedding_torch import RotaryEmbedding
 # hyper connections / attend from x-transformers, which handles different queries and key lengths better
 
 from x_transformers.attend import Attend
+
 from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 # proposed neural memory
@@ -515,7 +516,7 @@ class MemoryAsContextTransformer(Module):
 
         # hyper conection
 
-        init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
+        init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, dim = dim, add_stream_embed = True, disable = num_residual_streams == 1)
 
         self.layers = ModuleList([])
 
@@ -553,7 +554,7 @@ class MemoryAsContextTransformer(Module):
             mem_hyper_conn = None
 
             if layer in neural_memory_layers:
-                mem_hyper_conn = init_hyper_conn(dim = dim, 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)
 
                 mem = NeuralMemory(
                     dim = dim,
@@ -571,8 +572,8 @@ class MemoryAsContextTransformer(Module):
             self.layers.append(ModuleList([
                 mem_hyper_conn,
                 mem,
-                init_hyper_conn(dim = dim, branch = attn),
-                init_hyper_conn(dim = dim, branch = ff)
+                init_hyper_conn(branch = attn),
+                init_hyper_conn(branch = ff)
             ]))
 
         self.norm = nn.RMSNorm(dim)

{titans_pytorch-0.3.10 → titans_pytorch-0.3.11}/titans_pytorch/neural_memory.py

@@ -39,7 +39,7 @@ o - momentum orders
 
 LinearNoBias = partial(Linear, bias = False)
 
-NeuralMemCache = namedtuple('NeuralMemCache', [
+NeuralMemState = namedtuple('NeuralMemState', [
     'seq_index',
     'weights',
     'cache_store_segment',
@@ -629,7 +629,7 @@ class NeuralMemory(Module):
 
         if num_chunks == 0:
             updates = rearrange_dict_values(weights, 'bh ... -> bh 1 ...')
-            next_store_state = NeuralMemCache(next_seq_len_index, weights, remainder, past_state, updates)
+            next_store_state = NeuralMemState(next_seq_len_index, weights, remainder, past_state, updates)
 
             output = (updates, next_store_state)
 
@@ -682,13 +682,11 @@ class NeuralMemory(Module):
 
         next_state = (next_last_update, next_last_momentum)
 
-        next_store_state = NeuralMemCache(next_seq_len_index, weights, remainder, next_state, updates)
+        next_store_state = NeuralMemState(next_seq_len_index, weights, remainder, next_state, updates)
 
-        # returns
+        # return updates to neural memory at all chunked timesteps + neural mem cache / state to be fed back
 
-        output = (updates, next_store_state)
-
-        return output
+        return updates, next_store_state
 
     def retrieve_memories(
         self,
@@ -785,7 +783,7 @@ class NeuralMemory(Module):
         self,
         seq,
         store_seq = None,
-        state: NeuralMemCache | None = None,
+        state: NeuralMemState | None = None,
         prev_weights = None
     ):
         if seq.ndim == 2:

{titans_pytorch-0.3.10 → titans_pytorch-0.3.11}/train_mac.py

@@ -37,6 +37,7 @@ NUM_LONGTERM_MEM = 4
 NEURAL_MEM_LAYERS = (2, 4, 6)                   # layers 2, 4, 6 have neural memory, can add more
 NEURAL_MEM_GATE_ATTN_OUTPUT = False
 NEURAL_MEM_MOMENTUM = True
+NEURAL_MEM_MOMENTUM_ORDER = 1
 NEURAL_MEM_QK_NORM = True
 NEURAL_MEM_MAX_LR = 1e-1
 USE_MEM_ATTENTION_MODEL = False
@@ -115,6 +116,7 @@ model = MemoryAsContextTransformer(
         attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
         qk_rmsnorm = NEURAL_MEM_QK_NORM,
         momentum = NEURAL_MEM_MOMENTUM,
+        momentum_order = NEURAL_MEM_MOMENTUM_ORDER,
         default_step_transform_max_lr = NEURAL_MEM_MAX_LR,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
         per_parameter_lr_modulation = MEMORY_MODEL_PER_LAYER_LEARNED_LR