titans-pytorch 0.0.22__py3-none-any.whl → 0.0.24__py3-none-any.whl

titans_pytorch/titans.py

@@ -73,6 +73,17 @@ def softclamp_grad_norm(t, max_value):
     t = t * (clamped_norm / norm)
     return inverse(t)
 
+# multi head rmsnorm
+
+class MultiheadRMSNorm(Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.rmsnorm = nn.RMSNorm(dim, elementwise_affine = False)
+        self.gamma = nn.Parameter(torch.zeros(heads, 1, dim))
+
+    def forward(self, x):
+        return self.rmsnorm(x) * (self.gamma + 1.)
+
 # classes
 
 class MemoryMLP(Module):
@@ -100,11 +111,11 @@ class MemoryMLP(Module):
 
 # main neural memory
 
-def default_adaptive_step_transform(adaptive_step, max_lr = 1e-1):
+def default_adaptive_step_transform(adaptive_step, max_lr = 1e-2):
     return adaptive_step.sigmoid() * max_lr
 
 def default_loss_fn(pred, target):
-    return (pred - target).pow(2).mean(dim = -1).sum()
+    return (pred - target).pow(2).mean(dim = -1)
 
 class NeuralMemory(Module):
     def __init__(
@@ -131,17 +142,25 @@ class NeuralMemory(Module):
         self.retrieve_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
         self.store_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
 
-        self.post_rmsnorm = nn.RMSNorm(dim) if post_rmsnorm else nn.Identity()
+        self.multihead_rmsnorm = MultiheadRMSNorm(dim_head, heads) if post_rmsnorm else nn.Identity()
 
         # maybe multi-headed
 
         dim_head = default(dim_head, dim)
         dim_inner = dim_head * heads
 
+        self.heads = heads
+
         self.split_heads = Rearrange('b n (h d) -> (b h) n d', h = heads)
-        self.merge_heads = Rearrange('(b h) n d -> b n (h d)', h = heads)
+        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(
+            LinearNoBias(dim, heads),
+            Rearrange('b n h -> b h n 1'),
+            nn.Sigmoid()
+        ) if heads > 1 else None
+
         # memory mlp
 
         if not exists(model):
@@ -159,12 +178,13 @@ class NeuralMemory(Module):
 
         # prepare function for per sample gradients from model above, using torch.func
 
-        def forward_and_loss(params, inputs, target):
+        def forward_and_loss(params, inputs, loss_weights, target):
             pred = functional_call(self.memory_model, params, inputs)
             loss = self.store_memory_loss_fn(pred, target) # simple mse loss in paper - eq (12) - |M(k) - v|²
-            return loss
+            loss = loss * loss_weights
+            return loss.sum()
 
-        self.per_sample_grad_fn = vmap(grad(forward_and_loss), in_dims = (None, 0, 0))
+        self.per_sample_grad_fn = vmap(grad(forward_and_loss), in_dims = (None, 0, 0, 0))
 
         # queries for retrieving from the model
 
@@ -190,7 +210,6 @@ class NeuralMemory(Module):
         )
 
         self.to_adaptive_step = nn.Sequential(
-            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n')
         )
@@ -271,9 +290,11 @@ class NeuralMemory(Module):
 
         keys, values = tuple(rearrange(t, 'b (n c) d -> (b n) c d', c = self.chunk_size) for t in (keys, values))
 
+        adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = self.chunk_size)
+
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads = self.per_sample_grad_fn(dict(curr_weights), keys, values)
+        grads = self.per_sample_grad_fn(dict(curr_weights), keys, adaptive_lr, values)
 
         grads = TensorDict(grads)
 
@@ -286,9 +307,9 @@ class NeuralMemory(Module):
 
         grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))
 
-        # multiply gradients with learned adaptive step size
+        # negative gradients, adaptive lr already applied as loss weight
 
-        surprises = grads.apply(lambda t: einx.multiply('b n ..., b n -> b n ...', t, -adaptive_lr))
+        surprises = grads.apply(lambda t: -t)
 
         # determine scan function
 
@@ -356,7 +377,7 @@ class NeuralMemory(Module):
         past_weights: dict[str, Tensor] | None = None,
     ):
         chunk_size = self.chunk_size
-        seq_len = seq.shape[1]
+        batch, seq_len = seq.shape[:2]
 
         seq = self.retrieve_norm(seq)
 
@@ -390,8 +411,6 @@ class NeuralMemory(Module):
 
         queries = self.split_heads(queries)
 
-        batch = queries.shape[0]
-
         # fetch values from memory model
 
         curr_weights = curr_weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
@@ -403,7 +422,14 @@ class NeuralMemory(Module):
 
         # reconstitute batch dimension
 
-        values = rearrange(values, '(b n) c d -> b (n c) d', b = batch)
+        values = rearrange(values, '(b h n) c d -> b h (n c) d', b = batch, h = self.heads)
+
+        values = self.multihead_rmsnorm(values)
+
+        # maybe gate
+
+        if exists(self.retrieve_gate):
+            values = values * self.retrieve_gate(seq)
 
         # maybe merge heads and combine
 
@@ -411,10 +437,6 @@ class NeuralMemory(Module):
 
         values = self.combine_heads(values)
 
-        # post norm, somehow could not stabilize this without it, not in paper
-
-        values = self.post_rmsnorm(values)
-
         # restore, pad with empty memory embed
 
         empty_memory_embeds = self.init_empty_memory_embed(values.shape[0], chunk_size - 1)

{titans_pytorch-0.0.22.dist-info → titans_pytorch-0.0.24.dist-info}/METADATA

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

titans_pytorch-0.0.24.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/titans.py,sha256=3UxRZl_uwQBly11jQAWjfnNzHSoOUKiw-Ux2lXu2ilI,14304
+titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
+titans_pytorch-0.0.24.dist-info/METADATA,sha256=WGxo4oVx9HCq7LvSH8u_isp1tjxVXb3Ao_GrgjdFzSo,3811
+titans_pytorch-0.0.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.24.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.24.dist-info/RECORD,,

titans_pytorch-0.0.22.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/titans.py,sha256=s7KVqId7hhHw1Ck77FJUXfKC5rpwS4N7Nw2mKFtP8-s,13677
-titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.22.dist-info/METADATA,sha256=acTS0vWh84sGtTIEBzdTsxNi-S4V-Cr8NL4U4Vg0eY0,3811
-titans_pytorch-0.0.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.22.dist-info/RECORD,,