titans-pytorch 0.2.4__tar.gz → 0.2.6__tar.gz

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.2.4
+Version: 0.2.6
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -56,7 +56,7 @@ Description-Content-Type: text/markdown
 
 <img src="./fig1.png" width="400px"></img>
 
-## Titans - Pytorch
+## Titans - Pytorch (wip)
 
 Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytorch. Will also contain some explorations into architectures beyond their simple 1-4 layer MLP for the neural memory module, if it works well to any degree.
 

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/README.md

@@ -2,7 +2,7 @@
 
 <img src="./fig1.png" width="400px"></img>
 
-## Titans - Pytorch
+## Titans - Pytorch (wip)
 
 Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytorch. Will also contain some explorations into architectures beyond their simple 1-4 layer MLP for the neural memory module, if it works well to any degree.
 

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.2.4"
+version = "0.2.6"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/titans_pytorch/mac_transformer.py

@@ -488,7 +488,7 @@ class MemoryAsContextTransformer(Module):
         neural_memory_model: Module | None = None,
         neural_memory_kwargs: dict = dict(),
         neural_memory_layers: tuple[int, ...] | None = None,
-        aux_kv_recon_loss_weight = 0.,
+        aux_kv_recon_loss_weight = 1.,
         use_flex_attn = False,
         sliding_window_attn = False,
         weight_tie_memory_model = False,

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/titans_pytorch/memory_models.py

@@ -1,8 +1,15 @@
 import torch
-from torch import nn
+from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 
+# functions
+
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
+# memory mlp proposed in TTT
+
 class MemoryMLP(Module):
     def __init__(
         self,
@@ -19,15 +26,17 @@ class MemoryMLP(Module):
         self,
         x
     ):
+        residual = x
+
         for ind, weight in enumerate(self.weights):
             is_first = ind == 0
 
             if not is_first:
-                x = F.silu(x)
+                x = F.gelu(x)
 
             x = x @ weight
 
-        return x
+        return x + residual
 
 # memory mlp, but with gated residual + final projection
 
@@ -36,7 +45,7 @@ class GatedResidualMemoryMLP(Module):
         self,
         dim,
         depth,
-        expansion_factor = 2.
+        expansion_factor = 4.
     ):
         super().__init__()
         dim_hidden = int(dim * expansion_factor)
@@ -58,11 +67,13 @@ class GatedResidualMemoryMLP(Module):
         self,
         x
     ):
+        residual = x
+
         for weight1, weight2, to_gates in self.weights:
             res = x
 
             hidden = x @ weight1
-            hidden = F.silu(hidden)
+            hidden = F.gelu(hidden)
             branch_out = hidden @ weight2
 
             # gated residual
@@ -70,7 +81,9 @@ class GatedResidualMemoryMLP(Module):
             gates = cat((branch_out, res), dim = -1) @ to_gates
             x = res.lerp(branch_out, gates.sigmoid())
 
-        return x @ self.final_proj
+        out = x @ self.final_proj
+
+        return out + residual
 
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
@@ -98,15 +111,17 @@ class FactorizedMemoryMLP(Module):
         self,
         x
     ):
+        residual = x
+
         for ind, (weight1, weight2) in enumerate(self.weights):
             is_first = ind == 0
 
             if not is_first:
-                x = F.silu(x)
+                x = F.gelu(x)
 
             x = x @ weight1 @ weight2
 
-        return x
+        return x + residual
 
 # improvised attention as memory module
 
@@ -133,10 +148,12 @@ class MemoryAttention(Module):
             nn.init.xavier_uniform_(weight)
 
     def forward(self, x):
+        residual = x
+
         wq, wk, wv, ffw1, ffw2 = self.weights
 
-        q = F.normalize(x @ wq, dim = -1)
-        k = F.normalize(x @ wk, dim = -1)
+        q = l2norm(x @ wq)
+        k = l2norm(x @ wk)
         v = x @ wv
 
         attn_out = F.scaled_dot_product_attention(
@@ -145,9 +162,10 @@ class MemoryAttention(Module):
             is_causal = True
         )
 
-        x = x + attn_out
+        # parallel attention + feedforward block
+        # as in PaLM + Gpt-J
 
-        h = F.silu(x @ ffw1)
-        out = h @ ffw2
+        h = F.gelu(x @ ffw1)
+        ff_out = h @ ffw2
 
-        return out
+        return attn_out + ff_out + residual

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/titans_pytorch/neural_memory.py

@@ -304,13 +304,26 @@ class NeuralMemory(Module):
             nn.Sigmoid()
         ) if heads > 1 else None
 
-        # memory mlp
+        # memory model
 
         if not exists(model):
             model = MemoryMLP(dim_head, **default_model_kwargs)
 
+        # validate memory model
+
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
+        test_shape = (3, 2, dim_head)
+
+        with torch.no_grad():
+            try:
+                test_input = torch.randn(test_shape)
+                mem_model_output = model(test_input)
+            except:
+                raise RuntimeError(f'memory model unable to accept a tensor of shape {test_shape}')
+
+            assert mem_model_output.shape == test_shape, 'output of memory model needs to be same shape as input'
+
         # the memory is the weights of the model
 
         self.memory_model = model
@@ -772,10 +785,14 @@ class NeuralMemory(Module):
 
         # retrieve
 
-        if exists(prev_layer_updates):
-            prev_layer_updates = prev_layer_updates.apply(lambda t: pad_at_dim(t, (1, 0), dim = 1))
+        retrieve_chunk_size = default(chunk_size, self.retrieve_chunk_size)
+
+        if retrieve_chunk_size != 1:
+            if exists(prev_layer_updates):
+                prev_layer_updates = prev_layer_updates.apply(lambda t: pad_at_dim(t, (1, 0), dim = 1))
+
+            updates = updates.apply(lambda t: pad_at_dim(t, (1, 0), dim = 1))
 
-        updates = updates.apply(lambda t: pad_at_dim(t, (1, 0), dim = 1))
 
         retrieved = self.retrieve_memories(
             seq,

{titans_pytorch-0.2.4 → titans_pytorch-0.2.6}/train_mac.py

@@ -30,17 +30,18 @@ SEQ_LEN = 512
 NEURAL_MEMORY_DEPTH = 2
 NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
-NEURAL_MEM_LAYERS = (2, 4, 6)               # layers 2, 4, 6 have neural memory, can add more
+NEURAL_MEM_LAYERS = (2, 4, 6)                   # layers 2, 4, 6 have neural memory, can add more
 NEURAL_MEM_GATE_ATTN_OUTPUT = False
 NEURAL_MEM_MOMENTUM = True
-NEURAL_MEM_QK_NORM = False
+NEURAL_MEM_QK_NORM = True
 WINDOW_SIZE = 32
-NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2 # set smaller for more granularity for learning rate / momentum etc
+NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2       # set smaller for more granularity for learning rate / momentum etc
 SLIDING_WINDOWS = True
-WEIGHT_TIE_MEMORY_MODEL = True  # set to have memory MLP shared across layers
-STORE_ATTN_POOL_CHUNKS = True # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
+WEIGHT_TIE_MEMORY_MODEL = False                 # set to have memory MLP shared across layers
+PREV_MEM_UPDATE_FOR_WEIGHTS = True,
+STORE_ATTN_POOL_CHUNKS = True                   # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
 MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
-KV_RECON_LOSS_WEIGHT = 0.
+KV_RECON_LOSS_WEIGHT = 1.
 
 # experiment related
 
@@ -90,6 +91,7 @@ model = MemoryAsContextTransformer(
     use_flex_attn = USE_FLEX_ATTN,
     sliding_window_attn = SLIDING_WINDOWS,
     weight_tie_memory_model = WEIGHT_TIE_MEMORY_MODEL,
+    prev_neural_mem_update_for_weights = PREV_MEM_UPDATE_FOR_WEIGHTS,
     neural_memory_model = MemoryMLP(
         dim = 64,
         depth = NEURAL_MEMORY_DEPTH