titans-pytorch 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

titans_pytorch/__init__.py

@@ -6,6 +6,7 @@ from titans_pytorch.memory_models import (
     MemoryMLP,
     MemoryAttention,
     FactorizedMemoryMLP,
+    MemorySwiGluMLP,
     GatedResidualMemoryMLP
 )
 

titans_pytorch/memory_models.py

@@ -103,8 +103,6 @@ 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)
 
@@ -145,8 +143,6 @@ 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)
@@ -166,6 +162,44 @@ class FactorizedMemoryMLP(Module):
 
         return x
 
+# an MLP modelled after the popular swiglu ff in modern transformers
+
+class MemorySwiGluMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth = 1, # default to 2 layer MLP from TTT, depth of 2 would be 4 layer MLP, but done as 2 feedforwards with residual
+        expansion_factor = 4.
+    ):
+        super().__init__()
+
+        dim_inner = int(dim * expansion_factor * 2 / 3)
+
+        weights = []
+
+        for _ in range(depth):
+            weights.append(ParameterList([
+                Parameter(torch.randn(dim, dim_inner * 2)),
+                Parameter(torch.randn(dim_inner, dim)),
+            ]))
+
+        self.weights = ParameterList(weights)
+
+    def forward(self, x):
+
+        for w1, w2 in self.weights:
+            residual = x
+
+            x, gates = (x @ w1).chunk(2, dim = -1)
+
+            x = x * F.gelu(gates)
+
+            x = x @ w2
+
+            x = x + residual
+
+        return x
+
 # improvised attention as memory module
 
 class MemoryAttention(Module):
@@ -179,16 +213,14 @@ class MemoryAttention(Module):
         self.scale = scale
         dim_ff_hidden = int(dim * expansion_factor)
 
-        self.weights = nn.ParameterList([
-            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_ff_hidden)), # ff w1
-            nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
+        self.weights = ParameterList([
+            Parameter(torch.randn(dim, dim)), # queries
+            Parameter(torch.randn(dim, dim)), # keys
+            Parameter(torch.randn(dim, dim)), # values
+            Parameter(torch.randn(dim, dim_ff_hidden)), # ff w1
+            Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
         ])
 
-        self.ln = LayerNorm(dim)
-
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 

{titans_pytorch-0.3.5.dist-info → titans_pytorch-0.3.7.dist-info}/METADATA

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

titans_pytorch-0.3.7.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
+titans_pytorch/associative_scan.py,sha256=CEPXaZ2fEPWF8ZBe5wihCqPSGi8PNyL0uVSgvY7eV-s,5147
+titans_pytorch/mac_transformer.py,sha256=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
+titans_pytorch/memory_models.py,sha256=fC84MZNVfQxWqRx1rZKliqiTLsyT3eIH1iCxGRw8bpI,5789
+titans_pytorch/neural_memory.py,sha256=2Ffq5fob6_vJVUs1jIai3BJhfsfypKffJEb0QwRRdMk,27325
+titans_pytorch-0.3.7.dist-info/METADATA,sha256=hNbS9PgPGNFrNdZwG8nTgExvH1OQTHlOWGQNj4eb0ZU,6815
+titans_pytorch-0.3.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.7.dist-info/RECORD,,

titans_pytorch-0.3.5.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
-titans_pytorch/associative_scan.py,sha256=CEPXaZ2fEPWF8ZBe5wihCqPSGi8PNyL0uVSgvY7eV-s,5147
-titans_pytorch/mac_transformer.py,sha256=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
-titans_pytorch/memory_models.py,sha256=2fma9u0NQmDabgbpG6CLDGBRYzX99yIDQCSYIB0etkU,4989
-titans_pytorch/neural_memory.py,sha256=2Ffq5fob6_vJVUs1jIai3BJhfsfypKffJEb0QwRRdMk,27325
-titans_pytorch-0.3.5.dist-info/METADATA,sha256=R6EL4q-zgW7DV5OyLzqz5XP2IvLNpJkaBylwH8GsyII,6815
-titans_pytorch-0.3.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.5.dist-info/RECORD,,