titans-pytorch 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

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
 

{titans_pytorch-0.3.9.dist-info → titans_pytorch-0.3.10.dist-info}/METADATA

@@ -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.10.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=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=BeOnq41gjZeq-XJFjkHE44F9dLzsg9mm36EBYZ4wHMA,28814
+titans_pytorch-0.3.10.dist-info/METADATA,sha256=sA_Dx_x5RMcpz5-vUPDHuz__tHYfKzs4W_BgY4CHPdk,6816
+titans_pytorch-0.3.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.10.dist-info/RECORD,,

titans_pytorch-0.3.9.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=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=2Ffq5fob6_vJVUs1jIai3BJhfsfypKffJEb0QwRRdMk,27325
-titans_pytorch-0.3.9.dist-info/METADATA,sha256=EUuuqHl8jPUPIG8m7xV5esN_2yDNaPdD-H8qFeDxWGo,6815
-titans_pytorch-0.3.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.9.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.9.dist-info/RECORD,,