titans-pytorch 0.0.58__py3-none-any.whl → 0.0.62__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,7 +1,8 @@
 from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
-    MemoryAttention
+    MemoryAttention,
+    FactorizedMemoryMLP
 )
 
 from titans_pytorch.mac_transformer import (

titans_pytorch/mac_transformer.py

@@ -24,8 +24,8 @@ def create_mac_block_mask(seq_len, window_size, persist_mem_len):
 
     def create_mac_mask(b, h, q_idx, kv_idx):
         is_persist_mem = kv_idx < persist_mem_len
-        causal_mask = q_idx >= (kv_idx - is_persist_mem)
-        block_diagonal = (q_idx // window_size) == ((kv_idx - is_persist_mem) // window_size)
+        causal_mask = q_idx >= (kv_idx - persist_mem_len)
+        block_diagonal = (q_idx // window_size) == ((kv_idx - persist_mem_len) // window_size)
         return is_persist_mem | (~is_persist_mem & (causal_mask & block_diagonal))
 
     block_mask = create_block_mask(create_mac_mask, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len + persist_mem_len, _compile = True)
@@ -489,7 +489,7 @@ class MemoryAsContextTransformer(Module):
         flex_attn_fn = None
 
         if use_flex_attn:
-            block_mask = create_mac_block_mask(seq_len_with_mem, self.segment_len, self.num_persist_mem_tokens)
+            block_mask = create_mac_block_mask(seq_len_with_mem, segment_len + num_longterm_mem_tokens, self.num_persist_mem_tokens)
             flex_attn_fn = partial(flex_attention, block_mask = block_mask)
 
         # value residual

titans_pytorch/titans.py

@@ -6,7 +6,7 @@ from functools import partial
 import torch
 from torch import nn, Tensor
 import torch.nn.functional as F
-from torch.nn import Linear, Module
+from torch.nn import Linear, Module, Parameter, ParameterList
 from torch.func import functional_call, vmap, grad
 
 from tensordict import TensorDict
@@ -88,7 +88,7 @@ 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))
+        self.gamma = Parameter(torch.zeros(heads, 1, dim))
 
     def forward(self, x):
         return self.rmsnorm(x) * (self.gamma + 1.)
@@ -102,7 +102,10 @@ class MemoryMLP(Module):
         depth
     ):
         super().__init__()
-        self.weights = nn.ParameterList([nn.Parameter(torch.randn(dim, dim)) for _ in range(depth)])
+        self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
+
+        for weight in self.weights:
+            nn.init.xavier_uniform_(weight)
 
     def forward(
         self,
@@ -118,13 +121,50 @@ class MemoryMLP(Module):
 
         return x
 
+# memory mlp with factorized weights
+# so can tradeoff capacity for smaller chunk sizes
+
+class FactorizedMemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        k = 32
+    ):
+        super().__init__()
+        self.weights = ParameterList([
+            ParameterList([
+                Parameter(torch.randn(dim, k)),
+                Parameter(torch.randn(k, dim)),
+            ]) for _ in range(depth)
+        ])
+
+        for weight1, weight2 in self.weights:
+            nn.init.xavier_uniform_(weight1)
+            nn.init.xavier_uniform_(weight2)
+
+    def forward(
+        self,
+        x
+    ):
+        for ind, (weight1, weight2) in enumerate(self.weights):
+            is_first = ind == 0
+
+            if not is_first:
+                x = F.silu(x)
+
+            x = x @ weight1 @ weight2
+
+        return x
+
 # improvised attention as memory module
 
 class MemoryAttention(Module):
     def __init__(
         self,
         dim,
-        scale = 8.
+        scale = 8.,
+        expansion_factor = 2.
     ):
         super().__init__()
         self.scale = scale
@@ -133,10 +173,13 @@ class MemoryAttention(Module):
             nn.Parameter(torch.randn(dim, dim)), # queries
             nn.Parameter(torch.randn(dim, dim)), # keys
             nn.Parameter(torch.randn(dim, dim)), # values
-            nn.Parameter(torch.randn(dim, dim * 2)), # ff w1
-            nn.Parameter(torch.randn(dim * 2, dim)), # ff w2
+            nn.Parameter(torch.randn(dim, dim * expansion_factor)), # ff w1
+            nn.Parameter(torch.randn(dim * expansion_factor, dim)), # ff w2
         ])
 
+        for weight in self.weights:
+            nn.init.xavier_uniform_(weight)
+
     def forward(self, x):
         wq, wk, wv, ffw1, ffw2 = self.weights
 
@@ -289,6 +332,8 @@ class NeuralMemory(Module):
 
         self.use_accelerated_scan = use_accelerated_scan
 
+        self.register_buffer('zero', torch.tensor(0.), persistent = False)
+
     def init_weights_and_momentum(self):
         params = TensorDict(dict(self.memory_model.named_parameters()))
 
@@ -306,6 +351,13 @@ class NeuralMemory(Module):
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]],
         return_aux_kv_loss = False
     ):
+        seq_len = seq.shape[-2]
+
+        # handle edge case
+
+        if seq_len < self.chunk_size:
+            past_weight, _ = past_state
+            return TensorDict(past_weight).clone().zero_(), self.zero
 
         seq = self.store_norm(seq)
 
@@ -425,12 +477,10 @@ class NeuralMemory(Module):
 
         last_update = updates.apply(lambda t: t[:, -1])
 
-        next_state = (curr_weights + last_update, next_momentum)
-
         if not return_aux_kv_loss:
-            return updates, next_state
+            return updates
 
-        return updates, next_state, aux_kv_recon_loss.mean()
+        return updates, aux_kv_recon_loss.mean()
 
     def retrieve_memories(
         self,
@@ -442,7 +492,8 @@ class NeuralMemory(Module):
 
         seq = self.retrieve_norm(seq)
 
-        assert seq_len >= chunk_size
+        if seq_len < self.chunk_size:
+            return self.init_empty_memory_embed(batch, seq_len)
 
         seq = seq[:, (chunk_size - 1):]
         curtailed_seq_len = seq.shape[-2]
@@ -524,10 +575,11 @@ class NeuralMemory(Module):
 
         store_seq = default(store_seq, seq)
 
-        updates, next_memories, aux_kv_recon_loss = self.store_memories(store_seq, past_state, return_aux_kv_loss = True)
+        updates, aux_kv_recon_loss = self.store_memories(store_seq, past_state, return_aux_kv_loss = True)
 
         past_weights, _ = past_state
 
+
         retrieved = self.retrieve_memories(seq, past_weights + updates)
 
         if not return_aux_kv_loss:

{titans_pytorch-0.0.58.dist-info → titans_pytorch-0.0.62.dist-info}/METADATA

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

titans_pytorch-0.0.62.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=R9Xu-BjqAe9ZY60IGk4aNXBx_L8THsjJ4QrkbTnRNHo,15346
+titans_pytorch/titans.py,sha256=95J6UL44lOrdZSXdm7p36xC9tDeSmRBwdjig9T82PzI,17452
+titans_pytorch-0.0.62.dist-info/METADATA,sha256=08Blaa9Ehyv09rSA9uWguxbhKpbrd7C53Ya13E1VbpU,4457
+titans_pytorch-0.0.62.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.62.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.62.dist-info/RECORD,,

titans_pytorch-0.0.58.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=Q0MQA3RS8vqzs-KzSGZkvLR7afQ6ZW9uMOq1MeNuFoY,170
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=kk8s8Q2WmbJxCVi8PcqSUyJBc8-CDAHrVjt6M0d_kFs,15323
-titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
-titans_pytorch-0.0.58.dist-info/METADATA,sha256=a-Y6MV_89D44HlB7eKpurh-sw5DDiS-pIVei3Uw_uGE,4457
-titans_pytorch-0.0.58.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.58.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.58.dist-info/RECORD,,