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

titans_pytorch/titans.py

@@ -41,6 +41,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def xnor(x, y):
+    return not (x ^ y)
+
 def identity(t):
     return t
 
@@ -365,6 +368,8 @@ class NeuralMemory(Module):
         momentum = True,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        qk_rmsnorm = False,
+        accept_value_residual = False,
         learned_mem_model_weights = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
@@ -389,13 +394,16 @@ 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
 
         self.heads = heads
 
-        self.split_heads = Rearrange('b n (h d) -> (b h) n d', h = 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)')
         self.combine_heads = LinearNoBias(dim_inner, dim) if heads > 1 else nn.Identity()
 
@@ -444,6 +452,14 @@ class NeuralMemory(Module):
         self.to_keys_values = Sequential(LinearNoBias(dim, dim_inner * 2), activation)
         self.store_memory_loss_fn = store_memory_loss_fn
 
+        # value residual learning
+
+        self.learned_value_residual = Sequential(
+            LinearNoBias(dim, heads),
+            Rearrange('b n h -> b h n 1'),
+            nn.Sigmoid()
+        ) if accept_value_residual else None
+
         # empty memory embed
 
         self.empty_memory_embed = nn.Parameter(torch.zeros(dim))
@@ -525,8 +541,11 @@ class NeuralMemory(Module):
         seq,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]],
         return_aux_kv_loss = False,
-        chunk_size = None
+        chunk_size = None,
+        value_residual = None
     ):
+        assert xnor(exists(value_residual), exists(self.learned_value_residual))
+
         seq_len, chunk_size = seq.shape[-2], default(chunk_size, self.store_chunk_size)
 
         # handle edge case
@@ -577,9 +596,21 @@ class NeuralMemory(Module):
 
         batch = keys.shape[0]
 
+        # maybe qk rmsnorm
+
+        keys = self.k_norm(keys)
+
+        # maybe value residual learning
+
+        orig_values = values
+
+        if exists(self.learned_value_residual):
+            mix = self.learned_value_residual(seq)
+            values = values.lerp(value_residual, mix)
+
         # 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))
+        keys, values = tuple(rearrange(t, 'b h (n c) d -> (b h n) c d', c = chunk_size) for t in (keys, values))
 
         adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = chunk_size)
 
@@ -637,10 +668,12 @@ class NeuralMemory(Module):
 
         last_update = updates.apply(lambda t: t[:, -1])
 
+        output = (updates, orig_values)
+
         if not return_aux_kv_loss:
-            return updates
+            return output
 
-        return updates, aux_kv_recon_loss.mean()
+        return output, aux_kv_recon_loss.mean()
 
     def retrieve_memories(
         self,
@@ -683,10 +716,14 @@ 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) ...'))
-        queries = rearrange(queries, 'b (n c) d -> (b n) c d', c = chunk_size)
+        queries = rearrange(queries, 'b h (n c) d -> (b h n) c d', c = chunk_size)
 
         # forward functional call
 
@@ -723,7 +760,8 @@ class NeuralMemory(Module):
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
         return_aux_kv_loss = False,
         chunk_size = None,
-        store_chunk_size = None
+        store_chunk_size = None,
+        return_values = False
     ):
         batch, seq_len = seq.shape[:2]
 
@@ -744,13 +782,18 @@ class NeuralMemory(Module):
         store_seq = default(store_seq, seq)
         store_chunk_size = default(store_chunk_size, chunk_size)
 
-        updates, aux_kv_recon_loss = self.store_memories(store_seq, past_state, chunk_size = store_chunk_size, return_aux_kv_loss = True)
+        (updates, values), aux_kv_recon_loss = self.store_memories(store_seq, past_state, chunk_size = store_chunk_size, return_aux_kv_loss = True)
 
         past_weights, _ = past_state
 
         retrieved = self.retrieve_memories(seq, past_weights + updates, chunk_size = chunk_size)
 
+        output = retrieved
+
+        if return_values:
+            output = (retrieved, values)
+
         if not return_aux_kv_loss:
-            return retrieved
+            return output
 
-        return retrieved, aux_kv_recon_loss
+        return output, aux_kv_recon_loss

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.20
+Version: 0.1.22
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -78,8 +78,7 @@ from titans_pytorch import NeuralMemory
 
 mem = NeuralMemory(
     dim = 384,
-    chunk_size = 64,
-    pre_rmsnorm = True
+    chunk_size = 64
 ).cuda()
 
 seq = torch.randn(2, 1024, 384).cuda()
@@ -196,3 +195,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.22.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=7PGnZxdKq6T6e51RL7-QV43wp-46YmrytTZLt0McMco,23407
+titans_pytorch-0.1.22.dist-info/METADATA,sha256=HCOAqLK605c8R2mvgQ80kwE9jayZ2CwJqHLsJtFx7Vs,6718
+titans_pytorch-0.1.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.22.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,,