titans-pytorch 0.0.14__py3-none-any.whl → 0.0.15__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,3 +1,5 @@
 from titans_pytorch.titans import (
-    NeuralMemory
+    NeuralMemory,
+    MemoryMLP,
+    MemoryAttention
 )

titans_pytorch/titans.py

@@ -6,7 +6,7 @@ import torch
 from torch import nn, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module
-from torch.func import functional_call, vmap, grad_and_value
+from torch.func import functional_call, vmap, grad
 
 from tensordict import TensorDict
 
@@ -57,7 +57,7 @@ def pack_one_with_inverse(t, pattern):
 
 # classes
 
-class MLP(Module):
+class MemoryMLP(Module):
     def __init__(
         self,
         dim,
@@ -80,6 +80,35 @@ class MLP(Module):
 
         return x
 
+# improvised attention as memory module
+# todo - expand if see signal in experiments
+
+class MemoryAttention(Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+        self.weights = nn.ParameterList([
+            nn.Parameter(torch.randn(dim, dim)), # queries
+            nn.Parameter(torch.randn(dim, dim)), # keys
+            nn.Parameter(torch.randn(dim, dim)), # values
+        ])
+
+    def forward(self, x):
+        wq, wk, wv = self.weights
+
+        q = x @ wq
+        k = x @ wk
+        v = x @ wv
+
+        sim = q @ k.transpose(-1, -2)
+
+        attn = sim.softmax(dim = -1)
+
+        out = attn @ v
+        return out
+
 # main neural memory
 
 def default_loss_fn(pred, target):
@@ -122,7 +151,7 @@ class NeuralMemory(Module):
         # memory mlp
 
         if not exists(model):
-            model = MLP(dim_head, **default_mlp_kwargs)
+            model = MemoryMLP(dim_head, **default_mlp_kwargs)
 
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
@@ -141,7 +170,7 @@ class NeuralMemory(Module):
             loss = self.store_memory_loss_fn(pred, target) # simple mse loss in paper - eq (12) - |M(k) - v|²
             return loss
 
-        self.per_sample_grad_and_value_fn = vmap(grad_and_value(forward_and_loss), in_dims = (None, 0, 0))
+        self.per_sample_grad_fn = vmap(grad(forward_and_loss), in_dims = (None, 0, 0))
 
         # queries for retrieving from the model
 
@@ -235,7 +264,7 @@ class NeuralMemory(Module):
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads, aux_store_loss = self.per_sample_grad_and_value_fn(dict(curr_weights), keys, values)
+        grads = self.per_sample_grad_fn(dict(curr_weights), keys, values)
 
         grads = TensorDict(grads)
 
@@ -305,7 +334,7 @@ class NeuralMemory(Module):
 
         next_state = (curr_weights + last_update, next_momentum)
 
-        return updates, next_state, aux_store_loss.mean() / chunk_size
+        return updates, next_state
 
     def retrieve_memories(
         self,
@@ -382,6 +411,7 @@ class NeuralMemory(Module):
     def forward(
         self,
         seq,
+        store_seq = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
         return_next_memories = False
     ):
@@ -396,7 +426,9 @@ class NeuralMemory(Module):
         if not exists(past_state):
             past_state = self.init_weights_and_momentum()
 
-        updates, next_memories, aux_kv_mse_loss = self.store_memories(seq, past_state)
+        store_seq = default(store_seq, seq)
+
+        updates, next_memories = self.store_memories(store_seq, past_state)
 
         past_weights, _ = past_state
 
@@ -405,4 +437,4 @@ class NeuralMemory(Module):
         if not return_next_memories:
             return retrieved
 
-        return retrieved, next_memories, aux_kv_mse_loss
+        return retrieved, next_memories

{titans_pytorch-0.0.14.dist-info → titans_pytorch-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.14
+Version: 0.0.15
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -109,3 +109,13 @@ $ python train.py
     url     = {https://api.semanticscholar.org/CorpusID:275212078}
 }
 ```
+
+```bibtex
+@software{Kyrylov_Accelerated_Scan_2024,
+    author  = {Kyrylov, Volodymyr},
+    doi     = {10.5281/zenodo.10600962},
+    title   = {Accelerated Scan},
+    version = {0.1.2},
+    year    = {2024}
+}
+```

titans_pytorch-0.0.15.dist-info/RECORD

@@ -0,0 +1,7 @@
+titans_pytorch/__init__.py,sha256=bu8p8kUA24EVrTz-ojixHTwV-6KTY9Y0cNJkaMW4Whw,91
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/titans.py,sha256=ArFYKgI0p7N3mmv8b4ncxkl3gkKAXWrFrnl2quh2RqE,12930
+titans_pytorch-0.0.15.dist-info/METADATA,sha256=3IC7BT7J3BYx23wUOYeuGgtz769dsIczBnWbm6oi0Tw,3811
+titans_pytorch-0.0.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.15.dist-info/RECORD,,

titans_pytorch-0.0.14.dist-info/RECORD

@@ -1,7 +0,0 @@
-titans_pytorch/__init__.py,sha256=QKuJPCOJCdgtaPeKoHEkYkiQe65_LV9_8-cIMbBPU30,55
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/titans.py,sha256=O4WO2I6GhxyZobqbgfFx01saFEKUxhA0BBwt70m2yeQ,12306
-titans_pytorch-0.0.14.dist-info/METADATA,sha256=HKDSJ3sWc54sN1_fOEYU7i5TjiQff49vcZG9G8EU6z4,3598
-titans_pytorch-0.0.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.14.dist-info/RECORD,,