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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.3.9
+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.9 → titans_pytorch-0.3.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.3.9"
+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.9 → titans_pytorch-0.3.11}/tests/test_titans.py

@@ -60,6 +60,26 @@ def test_titans(
 
     assert seq.shape == retrieved.shape
 
+@pytest.mark.parametrize('learned_momentum_combine', (False, True))
+def test_titans_second_order_momentum(
+    learned_momentum_combine
+):
+
+    mem  = NeuralMemory(
+        dim = 384,
+        dim_head = 64,
+        heads = 2,
+        chunk_size = 1,
+        batch_size = 2,
+        momentum_order = 2,
+        learned_momentum_combine = learned_momentum_combine
+    )
+
+    seq = torch.randn(2, 5, 384)
+
+    parallel_retrieved, state = mem(seq)
+    assert seq.shape == parallel_retrieved.shape
+
 def test_titans_attn_memory():
     from titans_pytorch.memory_models import MemoryAttention
 
@@ -318,12 +338,16 @@ def test_flex(
 
     assert torch.allclose(out_flex, out_non_flex, atol = 1e-5)
 
-@torch_default_dtype(torch.float64)
-def test_assoc_scan():
+@pytest.mark.parametrize('use_accelerated', (True, False))
+def test_assoc_scan(
+    use_accelerated
+):
     from titans_pytorch.neural_memory import AssocScan
-    torch.set_default_dtype(torch.float64)
 
-    scan = AssocScan()
+    if use_accelerated and not torch.cuda.is_available():
+        pytest.skip()
+
+    scan = AssocScan(use_accelerated = use_accelerated)
 
     seq_len = 128
     mid_point = seq_len // 2
@@ -331,6 +355,10 @@ def test_assoc_scan():
     gates = torch.randn(2, seq_len, 16).sigmoid()
     inputs = torch.randn(2, seq_len, 16)
 
+    if use_accelerated:
+        gates = gates.cuda()
+        inputs = inputs.cuda()
+
     output = scan(gates, inputs)
 
     gates1, gates2 = gates[:, :mid_point], gates[:, mid_point:]
@@ -341,4 +369,4 @@ def test_assoc_scan():
     second_half = scan(gates2, inputs2, prev = first_half[:, -1])
     assert second_half.shape == inputs2.shape
 
-    assert torch.allclose(output[:, -1], second_half[:, -1], atol = 1e-6)
+    assert torch.allclose(output[:, -1], second_half[:, -1], atol = 1e-5)

{titans_pytorch-0.3.9 → 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.9 → titans_pytorch-0.3.11}/titans_pytorch/neural_memory.py

@@ -3,6 +3,7 @@ from typing import Callable
 
 import math
 from functools import partial
+from itertools import zip_longest
 from collections import namedtuple
 
 import torch
@@ -21,21 +22,24 @@ from titans_pytorch.memory_models import(
 )
 
 import einx
-from einops import rearrange, repeat, reduce, pack, unpack
+from einops import einsum, rearrange, repeat, reduce, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
 """
 ein notation:
 b - batch
+h - heads
+bh - batch and heads
 n - sequence
 d - feature dimension
 c - intra-chunk
 w - num memory network weight parameters
+o - momentum orders
 """
 
 LinearNoBias = partial(Linear, bias = False)
 
-NeuralMemCache = namedtuple('NeuralMemCache', [
+NeuralMemState = namedtuple('NeuralMemState', [
     'seq_index',
     'weights',
     'cache_store_segment',
@@ -224,6 +228,8 @@ class NeuralMemory(Module):
         per_head_learned_parameters = True,
         attn_pool_chunks = False,
         momentum = True,
+        momentum_order = 1,
+        learned_momentum_combine = False,
         pre_rmsnorm = True,
         post_rmsnorm = False,
         qk_rmsnorm = False,
@@ -367,12 +373,7 @@ class NeuralMemory(Module):
         else:
             self.reduce_to_chunk_rep = AttentionPool(dim, chunk_size = chunk_size)
 
-        # learned adaptive learning rate and momentum
-
-        self.to_momentum = Sequential(
-            nn.Linear(dim, heads),
-            Rearrange('b n h -> (b h) n 1')
-        ) if momentum else None
+        # learned adaptive learning rate
 
         self.to_adaptive_step = Sequential(
             nn.Linear(dim, heads),
@@ -384,6 +385,26 @@ class NeuralMemory(Module):
 
         self.adaptive_step_transform = adaptive_step_transform
 
+        # momentum related
+
+        self.to_momentum = Sequential(
+            nn.Linear(dim, heads * momentum_order),
+            Rearrange('b n (h o) -> o (b h) n 1', o = momentum_order)
+        ) if momentum else None
+
+        self.momentum_order = momentum_order
+        self.to_learned_momentum_combine = None
+
+        if learned_momentum_combine:
+            assert momentum
+            assert momentum_order > 1, 'only second order momentum allowed for now, but may allow learned combination of zeroth'
+
+            self.to_learned_momentum_combine = Sequential(
+                nn.Linear(dim, heads * momentum_order),
+                nn.Softmax(dim = -1),
+                Rearrange('b n (h o) -> o (b h) n', h = heads)
+            )
+
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
@@ -463,9 +484,9 @@ class NeuralMemory(Module):
         zeros = self.memory_model_parameter_dict.clone().zero_()
 
         if self.per_head_learned_parameters:
-            zeros = repeat_dict_values(zeros, 'h ... -> (b h) ...', b = batch)
+            zeros = repeat_dict_values(zeros, 'h ... -> o (b h) ...', b = batch, o = self.momentum_order)
         else:
-            zeros = repeat_dict_values(zeros, '... -> bh ...', bh = batch * self.heads)
+            zeros = repeat_dict_values(zeros, '... -> o bh ...', bh = batch * self.heads, o = self.momentum_order)
 
         return zeros
 
@@ -518,6 +539,11 @@ class NeuralMemory(Module):
         if has_momentum:
             adaptive_momentum = self.to_momentum(chunked_seq).sigmoid()
 
+            learned_combine = exists(self.to_learned_momentum_combine)
+
+            if learned_combine:
+                combine_momentums = self.to_learned_momentum_combine(chunked_seq)
+
         if need_layer_lr_mod:
             layer_lr_mod = self.to_layer_modulation(chunked_seq) * self.max_mem_layer_modulation
 
@@ -585,7 +611,7 @@ class NeuralMemory(Module):
 
         # negative gradients, adaptive lr already applied as loss weight
 
-        surprises = grads.apply(lambda t: -t)
+        surprises = grads.mul(-1)
 
         # past states
 
@@ -603,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)
 
@@ -611,7 +637,6 @@ class NeuralMemory(Module):
 
         # momentum + weight decay - momentum is the new contribution, as most linear RNNs have learned forgetting gates
 
-        next_momentum = TensorDict() if has_momentum else None
         updates = TensorDict()
 
         next_last_update = TensorDict()
@@ -624,32 +649,44 @@ class NeuralMemory(Module):
             # derive momentum with associative scan - eq (10)
 
             if has_momentum:
+                momentum = surprise
+
+                momentums = [] # stores all momentum orders starting with first, to generalize to Nth order momentum
+
                 last_momentum = past_last_momentum[param_name]
-                update = self.assoc_scan(adaptive_momentum, surprise, prev = last_momentum) # momentum is S / surprise in the paper
-                momentum = update
-                next_last_momentum[param_name] = momentum[:, -1]
+
+                # go from first order momentum all the way to the Nth
+
+                for one_adaptive_momentum, one_last_momentum in zip_longest(adaptive_momentum, last_momentum):
+                    momentum = self.assoc_scan(one_adaptive_momentum, momentum, prev = one_last_momentum) # momentum is S / surprise in the paper
+
+                    momentums.append(momentum)
+
+                momentums = torch.stack(momentums)
+
+                next_last_momentum[param_name] = momentums[:, :, -1] # momentums shape is Float['o bh n 1']
+
+                if not learned_combine:
+                    update = momentums[-1]
+                else:
+                    update = einsum(combine_momentums, momentums, 'o b n, o b n ... -> b n ...')
 
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
             update = self.assoc_scan(1. - decay_factor, update, prev = last_update, remove_prev = False)
-            next_last_update[param_name] = update[:, -1]
 
             updates[param_name] = update
-
-            if has_momentum:
-                next_momentum[param_name] = momentum
+            next_last_update[param_name] = update[:, -1]
 
         # determine next state for the storing of memories
 
         next_state = (next_last_update, next_last_momentum)
 
-        next_store_state = NeuralMemCache(next_seq_len_index, weights, remainder, next_state, updates)
-
-        # returns
+        next_store_state = NeuralMemState(next_seq_len_index, weights, remainder, next_state, updates)
 
-        output = (updates, next_store_state)
+        # return updates to neural memory at all chunked timesteps + neural mem cache / state to be fed back
 
-        return output
+        return updates, next_store_state
 
     def retrieve_memories(
         self,
@@ -746,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.9 → 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