titans-pytorch 0.0.65__tar.gz → 0.1.0__tar.gz

{titans_pytorch-0.0.65 → titans_pytorch-0.1.0}/PKG-INFO

@@ -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.0.65 → titans_pytorch-0.1.0}/README.md

@@ -6,6 +6,10 @@
 
 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.0.65 → titans_pytorch-0.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.65"
+version = "0.1.0"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -28,6 +28,7 @@ dependencies = [
     "accelerated-scan>=0.2.0",
     "axial_positional_embedding>=0.3.5",
     "einops>=0.8.0",
+    "einx>=0.3.0",
     "hyper-connections>=0.1.8",
     "Ninja",
     "rotary-embedding-torch",

{titans_pytorch-0.0.65 → titans_pytorch-0.1.0}/tests/test_titans.py

@@ -8,17 +8,20 @@ from titans_pytorch import NeuralMemory
 @pytest.mark.parametrize('silu', (False, True))
 @pytest.mark.parametrize('learned_mem_model_weights', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
+@pytest.mark.parametrize('per_parameter_lr_modulation', (False, True))
 def test_titans(
     seq_len,
     silu,
     learned_mem_model_weights,
     max_grad_norm,
+    per_parameter_lr_modulation
 ):
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
         activation = nn.SiLU() if silu else None,
         max_grad_norm = max_grad_norm,
+        per_parameter_lr_modulation = per_parameter_lr_modulation,
         learned_mem_model_weights = learned_mem_model_weights
     )
 

{titans_pytorch-0.0.65 → titans_pytorch-0.1.0}/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: