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

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.3.9"
+version = "0.3.10"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.3.9 → titans_pytorch-0.3.10}/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.10}/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,16 +22,19 @@ 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)
@@ -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
 
@@ -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,20 +649,34 @@ 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