titans-pytorch 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

titans_pytorch/titans.py

@@ -323,6 +323,7 @@ class NeuralMemory(Module):
         per_parameter_lr_modulation = False, # allow outer network to control learning rate per weight matrix of memory network
         max_mem_layer_modulation = 1e1, # max of 10.
         attn_pool_chunks = False,
+        momentum = True,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         learned_mem_model_weights = True,
@@ -423,7 +424,7 @@ class NeuralMemory(Module):
         self.to_momentum = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
-        )
+        ) if momentum else None
 
         self.to_adaptive_step = Sequential(
             LinearNoBias(dim, heads),
@@ -462,12 +463,15 @@ class NeuralMemory(Module):
 
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
-    def init_weights_and_momentum(self):
+    def init_weights_and_momentum(self, zero_weights = False):
         params = TensorDict(dict(self.memory_model.named_parameters()))
 
-        init_weights = params.clone().zero_()
+        init_weights = params
         init_momentum = params.clone().zero_()
 
+        if zero_weights:
+            init_weights = params.clone().zero_()
+
         return init_weights, init_momentum
 
     def init_empty_memory_embed(self, batch, seq_len):
@@ -497,14 +501,10 @@ class NeuralMemory(Module):
 
         seq = seq[:, :round_down_seq_len]
 
-        # curr weights + past weights, in the case that the initial weights are learned
-
-        curr_weights = TensorDict(dict(self.memory_model.named_parameters()))
+        # get the weights of the memory network
 
         past_state = tuple(TensorDict(d) for d in past_state)
-        past_weights, past_momentum = past_state
-
-        curr_weights = curr_weights + past_weights
+        curr_weights, past_momentum = past_state
 
         # derive learned hparams for optimization of memory network
 
@@ -513,10 +513,13 @@ class NeuralMemory(Module):
 
         chunked_seq = self.reduce_to_chunk_rep(seq, chunk_size = chunk_size)
 
-        adaptive_momentum = self.to_momentum(chunked_seq).sigmoid()
         decay_factor = self.to_decay_factor(chunked_seq).sigmoid()
 
         need_layer_lr_mod = exists(self.to_layer_modulation)
+        has_momentum = exists(self.to_momentum)
+
+        if has_momentum:
+            adaptive_momentum = self.to_momentum(chunked_seq).sigmoid()
 
         if need_layer_lr_mod:
             layer_lr_mod = self.to_layer_modulation(chunked_seq) * self.max_mem_layer_modulation
@@ -595,23 +598,29 @@ class NeuralMemory(Module):
 
         # momentum + weight decay - momentum is the new contribution, as most linear RNNs have learned forgetting gates
 
-        next_momentum = TensorDict()
+        next_momentum = TensorDict() if has_momentum else None
         updates = TensorDict()
 
         for param_name, surprise in surprises.items():
 
             surprise, inverse_pack = pack_one_with_inverse(surprise, 'b n *')
 
+            update = surprise
+
             # derive momentum with associative scan - eq (10)
 
-            momentum = scan_fn(adaptive_momentum, surprise) # momentum is S / surprise in the paper
+            if has_momentum:
+                update = scan_fn(adaptive_momentum, surprise) # momentum is S / surprise in the paper
+                momentum = update
 
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
-            update = scan_fn(1. - decay_factor, momentum)
+            update = scan_fn(1. - decay_factor, update)
 
             updates[param_name] = inverse_pack(update)
-            next_momentum[param_name] = inverse_pack(momentum)
+
+            if has_momentum:
+                next_momentum[param_name] = inverse_pack(momentum)
 
         # compute the next weight per batch
 

{titans_pytorch-0.1.12.dist-info → titans_pytorch-0.1.14.dist-info}/METADATA

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

titans_pytorch-0.1.14.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=zxknstaI_Uz47Y8WvZ3S7geJ-TNdqKV5Rvj0Jlw8njs,19271
+titans_pytorch/titans.py,sha256=J7UbhhL0YKoQjrKPCgwUYwJd8-YfCZrHTSVjIFvvRRw,22077
+titans_pytorch-0.1.14.dist-info/METADATA,sha256=lM_x9obTibKKMXXdsW0U5GcQLTR4aXCtemEKPzYPqYo,6340
+titans_pytorch-0.1.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.14.dist-info/RECORD,,

titans_pytorch-0.1.12.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=zxknstaI_Uz47Y8WvZ3S7geJ-TNdqKV5Rvj0Jlw8njs,19271
-titans_pytorch/titans.py,sha256=eDTqAIDZjSLd34t8M-dCaqVf_s0wZ9jhVIOfXF7E9ts,21887
-titans_pytorch-0.1.12.dist-info/METADATA,sha256=dL8HpHt6V5gN8p8px7sc2IgJGqXthE7rULKIrRFCwF8,6340
-titans_pytorch-0.1.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.12.dist-info/RECORD,,