titans-pytorch 0.3.10__py3-none-any.whl → 0.3.12__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

@@ -7,7 +7,7 @@ from itertools import zip_longest
 from collections import namedtuple
 
 import torch
-from torch import nn, cat, tensor, Tensor
+from torch import nn, stack, cat, tensor, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList, ParameterDict
 from torch.func import functional_call, vmap, grad
@@ -39,7 +39,7 @@ o - momentum orders
 
 LinearNoBias = partial(Linear, bias = False)
 
-NeuralMemCache = namedtuple('NeuralMemCache', [
+NeuralMemState = namedtuple('NeuralMemState', [
     'seq_index',
     'weights',
     'cache_store_segment',
@@ -230,6 +230,7 @@ class NeuralMemory(Module):
         momentum = True,
         momentum_order = 1,
         learned_momentum_combine = False,
+        learned_combine_include_zeroth = False,
         pre_rmsnorm = True,
         post_rmsnorm = False,
         qk_rmsnorm = False,
@@ -399,12 +400,17 @@ class NeuralMemory(Module):
             assert momentum
             assert momentum_order > 1, 'only second order momentum allowed for now, but may allow learned combination of zeroth'
 
+            if learned_combine_include_zeroth:
+                momentum_order += 1
+
             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)
             )
 
+            self.learned_combine_include_zeroth = learned_combine_include_zeroth
+
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
@@ -629,7 +635,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)
 
@@ -662,10 +668,14 @@ class NeuralMemory(Module):
 
                     momentums.append(momentum)
 
-                momentums = torch.stack(momentums)
+                momentums = stack(momentums)
 
                 next_last_momentum[param_name] = momentums[:, :, -1] # momentums shape is Float['o bh n 1']
 
+                if learned_combine and self.learned_combine_include_zeroth:
+                    # add the original surprise if learned combination of momentums
+                    momentums = cat((rearrange(surprise, '... -> 1 ...'), momentums), dim = 0)
+
                 if not learned_combine:
                     update = momentums[-1]
                 else:
@@ -682,13 +692,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)
-
-        # 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,
@@ -785,7 +793,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.dist-info → titans_pytorch-0.3.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.3.10
+Version: 0.3.12
 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.12.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=VmUAS1xOM0ZfearWIzQrX_P7HI69viuwrg9M7BQByeE,29349
+titans_pytorch-0.3.12.dist-info/METADATA,sha256=02OsMYNITFLjnKJgis8eUHxwcdH2aVbA_D-QK24TYbg,6817
+titans_pytorch-0.3.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.12.dist-info/RECORD,,

titans_pytorch-0.3.10.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=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,,