titans-pytorch 0.1.20__py3-none-any.whl → 0.1.21__py3-none-any.whl

titans_pytorch/titans.py

@@ -365,6 +365,7 @@ class NeuralMemory(Module):
         momentum = True,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        qk_rmsnorm = False,
         learned_mem_model_weights = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
@@ -389,6 +390,9 @@ class NeuralMemory(Module):
 
         self.multihead_rmsnorm = MultiheadRMSNorm(dim_head, heads) if post_rmsnorm else nn.Identity()
 
+        self.q_norm = MultiheadRMSNorm(dim_head, heads) if qk_rmsnorm else nn.Identity()
+        self.k_norm = MultiheadRMSNorm(dim_head, heads) if qk_rmsnorm else nn.Identity()
+
         # maybe multi-headed
 
         dim_inner = dim_head * heads
@@ -577,6 +581,10 @@ class NeuralMemory(Module):
 
         batch = keys.shape[0]
 
+        # maybe qk rmsnorm
+
+        keys = self.k_norm(keys)
+
         # take care of chunking
 
         keys, values = tuple(rearrange(t, 'b (n c) d -> (b n) c d', c = chunk_size) for t in (keys, values))
@@ -683,6 +691,10 @@ class NeuralMemory(Module):
 
         queries = self.split_heads(queries)
 
+        # maybe qk rmsnorm
+
+        queries = self.q_norm(queries)
+
         # fetch values from memory model
 
         curr_weights = curr_weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))

{titans_pytorch-0.1.20.dist-info → titans_pytorch-0.1.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.20
+Version: 0.1.21
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -196,3 +196,15 @@ $ python train_mac.py
     year    = {2024}
 }
 ```
+
+```bibtex
+@misc{wang2025testtimeregressionunifyingframework,
+    title   = {Test-time regression: a unifying framework for designing sequence models with associative memory},
+    author  = {Ke Alexander Wang and Jiaxin Shi and Emily B. Fox},
+    year    = {2025},
+    eprint  = {2501.12352},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url = {https://arxiv.org/abs/2501.12352},
+}
+```

titans_pytorch-0.1.21.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=Ejq1r3GQQnlT1Fo4McaOOie19t1HjwVlYbD90GLQCYI,22859
+titans_pytorch/titans.py,sha256=YYt6O5EiBVvyxWM4R1JuLLJH3bGm1V-74aB7VhbsWQ0,22577
+titans_pytorch-0.1.21.dist-info/METADATA,sha256=ixbJisycB0MgSIcOvDRM1PIMs3l1TM_AmQ88aWZYEsM,6742
+titans_pytorch-0.1.21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.21.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.21.dist-info/RECORD,,

titans_pytorch-0.1.20.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=Ejq1r3GQQnlT1Fo4McaOOie19t1HjwVlYbD90GLQCYI,22859
-titans_pytorch/titans.py,sha256=R0e25ly2uTHkHSZEb-9Eqb0DqtFq8wFBB8iH1T6bYVg,22240
-titans_pytorch-0.1.20.dist-info/METADATA,sha256=Y0TmkfpKQ4LAyhr6SmAGeLHs3H4ZiZ4lg-gevvUDmjI,6340
-titans_pytorch-0.1.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.20.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.20.dist-info/RECORD,,