titans-pytorch 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

titans_pytorch/memory_models.py

@@ -30,6 +30,25 @@ class LayerNorm(Module):
 
         return self.ln(x) * (gamma + 1.)
 
+# norm + residual wrapper, as used in original TTT paper
+# but could be removed
+
+class ResidualNorm(Module):
+    def __init__(
+        self,
+        dim,
+        model: Module
+    ):
+        super().__init__()
+        self.norm = LayerNorm(dim)
+        self.model = model
+
+    def forward(self, x):
+
+        out = self.model(x)
+
+        return self.norm(out) + x
+
 # memory mlp proposed in TTT
 
 class MemoryMLP(Module):
@@ -45,8 +64,6 @@ class MemoryMLP(Module):
 
         self.weights = ParameterList([Parameter(torch.randn(dim_in, dim_out)) for dim_in, dim_out in zip(dims[:-1], dims[1:])])
 
-        self.ln = LayerNorm(dim)
-
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 
@@ -54,8 +71,6 @@ class MemoryMLP(Module):
         self,
         x
     ):
-        residual = x
-
         for ind, weight in enumerate(self.weights):
             is_first = ind == 0
 
@@ -64,7 +79,7 @@ class MemoryMLP(Module):
 
             x = x @ weight
 
-        return self.ln(x) + residual
+        return x
 
 # memory mlp, but with gated residual + final projection
 
@@ -88,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)
 
@@ -97,7 +110,6 @@ class GatedResidualMemoryMLP(Module):
         self,
         x
     ):
-        residual = x
 
         for weight1, weight2, to_gates in self.weights:
             res = x
@@ -111,9 +123,7 @@ class GatedResidualMemoryMLP(Module):
             gates = cat((branch_out, res), dim = -1) @ to_gates
             x = res.lerp(branch_out, gates.sigmoid())
 
-        out = x @ self.final_proj
-
-        return self.ln(out) + residual
+        return x @ self.final_proj
 
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
@@ -133,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)
@@ -143,7 +151,6 @@ class FactorizedMemoryMLP(Module):
         self,
         x
     ):
-        residual = x
 
         for ind, (weight1, weight2) in enumerate(self.weights):
             is_first = ind == 0
@@ -153,7 +160,7 @@ class FactorizedMemoryMLP(Module):
 
             x = x @ weight1 @ weight2
 
-        return self.ln(x) + residual
+        return x
 
 # improvised attention as memory module
 
@@ -176,13 +183,10 @@ 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)
 
     def forward(self, x):
-        residual = x
 
         wq, wk, wv, ffw1, ffw2 = self.weights
 
@@ -202,4 +206,4 @@ class MemoryAttention(Module):
         h = F.gelu(x @ ffw1)
         ff_out = h @ ffw2
 
-        return self.ln(attn_out + ff_out) + residual
+        return attn_out + ff_out

titans_pytorch/neural_memory.py

@@ -16,7 +16,8 @@ from tensordict import TensorDict
 from titans_pytorch.associative_scan import AssocScan
 
 from titans_pytorch.memory_models import(
-    MemoryMLP
+    MemoryMLP,
+    ResidualNorm
 )
 
 import einx
@@ -234,6 +235,7 @@ class NeuralMemory(Module):
         init_decay_bias = None,
         accept_weight_residual = False,
         gated_transition = False,
+        mem_model_norm_add_residual = True, # by default, layernorm output and add residual as proposed in TTT paper, but could be removed
         default_model_kwargs: dict = dict(
             depth = 2,
             expansion_factor = 4.
@@ -304,6 +306,9 @@ class NeuralMemory(Module):
 
         # the memory is the weights of the model
 
+        if mem_model_norm_add_residual:
+            model = ResidualNorm(dim = dim_head, model = model)
+
         self.memory_model = model
 
         mem_model_params = dict(model.named_parameters())

{titans_pytorch-0.3.4.dist-info → titans_pytorch-0.3.6.dist-info}/METADATA

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

titans_pytorch-0.3.6.dist-info/RECORD

@@ -0,0 +1,9 @@
+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=u4DjLU9gnUu5TeT6YEnqeMIuPDfdCyMDocPOHNQBo5U,4887
+titans_pytorch/neural_memory.py,sha256=2Ffq5fob6_vJVUs1jIai3BJhfsfypKffJEb0QwRRdMk,27325
+titans_pytorch-0.3.6.dist-info/METADATA,sha256=CFbbMESkeeScixULPFLoWrbWuIBk_0vHyB95uPKsGdU,6815
+titans_pytorch-0.3.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.6.dist-info/RECORD,,

titans_pytorch-0.3.4.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=0KLHZN-y_7lwrhWSnFRaYJ3GiUV3tzVjxS9CxIx_eI8,4843
-titans_pytorch/neural_memory.py,sha256=9eyeEvYsP5OFlwLDRyVut99uVYGvXAElFPabVoZnGJw,27063
-titans_pytorch-0.3.4.dist-info/METADATA,sha256=2ZD_DovSYkVejsTWHq7_IOTN-Je0of1f-HOiojaQBhQ,6815
-titans_pytorch-0.3.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.4.dist-info/RECORD,,