titans-pytorch 0.0.65__py3-none-any.whl → 0.1.0__py3-none-any.whl

titans_pytorch/titans.py

@@ -17,6 +17,7 @@ from titans_pytorch.associative_scan import (
     pad_at_dim
 )
 
+import einx
 from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
@@ -26,6 +27,7 @@ b - batch
 n - sequence
 d - feature dimension
 c - intra-chunk
+w - num memory network weight parameters
 """
 
 LinearNoBias = partial(Linear, bias = False)
@@ -220,6 +222,8 @@ class NeuralMemory(Module):
         store_memory_loss_fn: Callable = default_loss_fn,
         adaptive_step_transform: Callable | None = None,
         default_step_transform_max_lr = 1e-2,
+        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.
         pre_rmsnorm = True,
         post_rmsnorm = True,
         learned_mem_model_weights = True,
@@ -250,7 +254,7 @@ class NeuralMemory(Module):
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
         self.combine_heads = LinearNoBias(dim_inner, dim) if heads > 1 else nn.Identity()
 
-        self.retrieve_gate = nn.Sequential(
+        self.retrieve_gate = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> b h n 1'),
             nn.Sigmoid()
@@ -270,6 +274,8 @@ class NeuralMemory(Module):
 
         self.memory_model = model
 
+        self.num_memory_parameter_tensors = len(set(model.parameters()))
+
         # the chunk size within the paper where adaptive step, momentum, weight decay are shared
 
         self.chunk_size = chunk_size
@@ -299,15 +305,14 @@ class NeuralMemory(Module):
         nn.init.normal_(self.empty_memory_embed, std = 0.02)
 
         # learned adaptive learning rate and momentum
-        # todo - explore mlp layerwise learned lr / momentum
 
-        self.to_momentum = nn.Sequential(
+        self.to_momentum = Sequential(
             Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         )
 
-        self.to_adaptive_step = nn.Sequential(
+        self.to_adaptive_step = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n')
         )
@@ -317,13 +322,24 @@ class NeuralMemory(Module):
 
         self.adaptive_step_transform = adaptive_step_transform
 
+        # per layer learning rate modulation
+
+        self.to_layer_modulation = Sequential(
+            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
+            LinearNoBias(dim, heads * self.num_memory_parameter_tensors),
+            Rearrange('b n (h w) -> w (b h) n', h = heads),
+            nn.Sigmoid()
+        ) if per_parameter_lr_modulation else None
+
+        self.max_mem_layer_modulation = max_mem_layer_modulation
+
         # allow for softclamp the gradient norms for storing memories
 
         self.max_grad_norm = max_grad_norm
 
         # weight decay factor
 
-        self.to_decay_factor = nn.Sequential(
+        self.to_decay_factor = Sequential(
             Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
@@ -387,6 +403,11 @@ class NeuralMemory(Module):
         adaptive_momentum = self.to_momentum(seq).sigmoid()
         decay_factor = self.to_decay_factor(seq).sigmoid()
 
+        need_layer_lr_mod = exists(self.to_layer_modulation)
+
+        if need_layer_lr_mod:
+            layer_lr_mod = self.to_layer_modulation(seq) * self.max_mem_layer_modulation
+
         # keys and values
 
         keys, values = self.to_keys_values(seq).chunk(2, dim = -1)
@@ -418,6 +439,11 @@ class NeuralMemory(Module):
 
         grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))
 
+        # maybe per layer modulation
+
+        if need_layer_lr_mod:
+            grads = TensorDict({name: einx.multiply('b h, b h ... -> b h ...', layer_lr_mod, t) for layer_lr_mod, (name, t) in zip(layer_lr_mod, grads.items())})
+
         # negative gradients, adaptive lr already applied as loss weight
 
         surprises = grads.apply(lambda t: -t)
@@ -469,7 +495,7 @@ class NeuralMemory(Module):
 
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
-            update = scan_fn(1. - decay_factor, momentum) # momentum is S / surprise in the paper
+            update = scan_fn(1. - decay_factor, momentum)
 
             updates[param_name] = inverse_pack(update)
             next_momentum[param_name] = inverse_pack(momentum)
@@ -580,7 +606,6 @@ class NeuralMemory(Module):
 
         past_weights, _ = past_state
 
-
         retrieved = self.retrieve_memories(seq, past_weights + updates)
 
         if not return_aux_kv_loss:

{titans_pytorch-0.0.65.dist-info → titans_pytorch-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.65
+Version: 0.1.0
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -37,6 +37,7 @@ Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
 Requires-Dist: axial-positional-embedding>=0.3.5
 Requires-Dist: einops>=0.8.0
+Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
@@ -58,6 +59,10 @@ Description-Content-Type: text/markdown
 
 Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytorch. Will also contain some explorations into architectures beyond their simple 1-4 layer MLP for the neural memory module, if it works well to any degree.
 
+## Appreciation
+
+- [@sentialx](https://github.com/sentialx) for sharing his early experimental results with me
+
 ## Install
 
 ```bash

titans_pytorch-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=7PHBCbeB1LhHY5s3zAyYF0L3Mm7CNy4TOBbcpLX6bNE,17686
+titans_pytorch/titans.py,sha256=L3Mu6pDnimD4MNn_832trFEJgXOPjxSdTrB9jiSUSTk,18533
+titans_pytorch-0.1.0.dist-info/METADATA,sha256=LuWDzv-NbGxYKOMThN_WKQWDueyIsOAMSwwiE_BDraI,4595
+titans_pytorch-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.0.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.0.dist-info/RECORD,,

titans_pytorch-0.0.65.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=7PHBCbeB1LhHY5s3zAyYF0L3Mm7CNy4TOBbcpLX6bNE,17686
-titans_pytorch/titans.py,sha256=y6lJRErIoM6T2aTVFlf1GxSB0cpsmBZdSIj1DCHUCQ8,17486
-titans_pytorch-0.0.65.dist-info/METADATA,sha256=oDjEiufwOninsFDoCGbu691LXc1mey2OT7j6PNzkz0Q,4457
-titans_pytorch-0.0.65.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.65.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.65.dist-info/RECORD,,