titans-pytorch 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

titans_pytorch/memory_models.py

@@ -3,11 +3,33 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 
+from einops import rearrange
+
 # functions
 
 def l2norm(t):
     return F.normalize(t, dim = -1)
 
+# norms
+
+class LayerNorm(Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+
+        self.ln = nn.LayerNorm(dim, elementwise_affine = False)
+        self.gamma = Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        gamma = self.gamma
+
+        if gamma.ndim == 2:
+            gamma = rearrange(gamma, 'b d -> b 1 d')
+
+        return self.ln(x) * (gamma + 1.)
+
 # memory mlp proposed in TTT
 
 class MemoryMLP(Module):
@@ -19,6 +41,8 @@ class MemoryMLP(Module):
         super().__init__()
         self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
 
+        self.ln = LayerNorm(dim)
+
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 
@@ -36,7 +60,7 @@ class MemoryMLP(Module):
 
             x = x @ weight
 
-        return x + residual
+        return self.ln(x) + residual
 
 # memory mlp, but with gated residual + final projection
 
@@ -60,6 +84,8 @@ class GatedResidualMemoryMLP(Module):
 
         self.final_proj = Parameter(torch.randn(dim, dim))
 
+        self.ln = LayerNorm(dim)
+
         for param in self.parameters():
             nn.init.xavier_uniform_(param)
 
@@ -83,7 +109,7 @@ class GatedResidualMemoryMLP(Module):
 
         out = x @ self.final_proj
 
-        return out + residual
+        return self.ln(out) + residual
 
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
@@ -103,6 +129,8 @@ class FactorizedMemoryMLP(Module):
             ]) for _ in range(depth)
         ])
 
+        self.ln = LayerNorm(dim)
+
         for weight1, weight2 in self.weights:
             nn.init.xavier_uniform_(weight1)
             nn.init.xavier_uniform_(weight2)
@@ -121,7 +149,7 @@ class FactorizedMemoryMLP(Module):
 
             x = x @ weight1 @ weight2
 
-        return x + residual
+        return self.ln(x) + residual
 
 # improvised attention as memory module
 
@@ -144,6 +172,8 @@ class MemoryAttention(Module):
             nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
         ])
 
+        self.ln = LayerNorm(dim)
+
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 
@@ -168,4 +198,4 @@ class MemoryAttention(Module):
         h = F.gelu(x @ ffw1)
         ff_out = h @ ffw2
 
-        return attn_out + ff_out + residual
+        return self.ln(attn_out + ff_out) + residual

titans_pytorch/neural_memory.py

@@ -742,7 +742,7 @@ class NeuralMemory(Module):
 
         # retrieve
 
-        retrieved = self.retrieve_memories(token, weights, chunk_size = 1)
+        retrieved = self.retrieve_memories(token, updates, chunk_size = 1)
 
         # next state tuple
 
@@ -801,10 +801,9 @@ class NeuralMemory(Module):
 
             updates = updates.apply(lambda t: pad_at_dim(t, (1, 0), dim = 1))
 
-
         retrieved = self.retrieve_memories(
             seq,
-            mem_model_weights,
+            updates,
             chunk_size = chunk_size,
             prev_layer_updates = prev_layer_updates
         )

{titans_pytorch-0.2.7.dist-info → titans_pytorch-0.2.9.dist-info}/METADATA

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

titans_pytorch-0.2.9.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=UOJAMv7nTgkefBB7M7K3U0NnFkz75tFRG5WLXRdfnLw,26039
+titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
+titans_pytorch/neural_memory.py,sha256=YVbKl7DYKFWUgCawDTxXIEgJAcl7nq5OaZytmovIl8Q,24899
+titans_pytorch-0.2.9.dist-info/METADATA,sha256=fSFt54zXLKB5gRhLTJd9551O0pF2qcYNlR7039yJiD0,6811
+titans_pytorch-0.2.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.2.9.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.2.9.dist-info/RECORD,,

titans_pytorch-0.2.7.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=UOJAMv7nTgkefBB7M7K3U0NnFkz75tFRG5WLXRdfnLw,26039
-titans_pytorch/memory_models.py,sha256=CD8pQ-IUfTDvPmekuPTsZHE3Vy265QtbiUn_siJhA78,4064
-titans_pytorch/neural_memory.py,sha256=WAeR-nOpy1XbBP590By1-tCgirulqPbFGut4H1B77-g,24910
-titans_pytorch-0.2.7.dist-info/METADATA,sha256=ndFb28pAe8xWmNU6oncV8VJDDPImo3aCuBv0d0JylIs,6811
-titans_pytorch-0.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.7.dist-info/RECORD,,