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

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/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.dist-info → titans_pytorch-0.3.11.dist-info}/METADATA

@@ -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.11.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=EyqA53HBqvAr4UNZUs37LR6IltyEfA7FKEV54YzVYlg,24945
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=7YglrQaDpKS2hbpBBwx7PmqhJdjyvFEPZDt_QXmnUMM,28878
+titans_pytorch-0.3.11.dist-info/METADATA,sha256=xAEvavDiCj__5Bl_5UXaG__BycdUB2DzHOud-nwsn1c,6817
+titans_pytorch-0.3.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.11.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.11.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,,