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

titans_pytorch/memory_models.py

@@ -3,6 +3,13 @@ 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/neural_memory.py

@@ -51,6 +51,9 @@ def default(*args):
             return arg
     return None
 
+def identity(t):
+    return t
+
 def xnor(x, y):
     return not (x ^ y)
 
@@ -64,9 +67,6 @@ def safe_cat(inputs, dim = -2):
 
     return cat(inputs, dim = dim)
 
-def identity(t):
-    return t
-
 def dict_get_shape(td):
     return {k: v.shape for k, v in td.items()}
 
@@ -454,14 +454,14 @@ class NeuralMemory(Module):
 
         weights = TensorDict(weights)
 
-        # allow for neural memory of a previous layer and the past to produce gradients that become the weights of the current one generating the surprise
-        # think this is necessary otherwise the memory model is static (unless if paper is misunderstood)
-        # improvise (or perhaps correcting to) a solution
+        # allow for neural memory of a previous layer to influence surprise of current layer
+
+        weights_for_surprise = weights
 
         if exists(prev_layer_updates):
             prev_layer_updates = TensorDict(prev_layer_updates)
 
-            weights = weights + prev_layer_updates
+            weights_for_surprise = weights_for_surprise + prev_layer_updates
 
             per_sample_grad_fn = self.per_sample_grad_fn_expanded_weights # the weights will now have a batch * chunk dimension
 
@@ -506,11 +506,11 @@ class NeuralMemory(Module):
         # flatten batch and time if surprise depends on previous layer memory model
 
         if exists(prev_layer_updates):
-            weights = weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
+            weights_for_surprise = weights_for_surprise.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads, aux_kv_recon_loss = per_sample_grad_fn(dict(weights), keys, adaptive_lr, values)
+        grads, aux_kv_recon_loss = per_sample_grad_fn(dict(weights_for_surprise), keys, adaptive_lr, values)
 
         grads = TensorDict(grads)
 
@@ -536,7 +536,15 @@ class NeuralMemory(Module):
 
         if not exists(past_state):
             empty_dict = {key: None for key in weights.keys()}
-            past_state = (empty_dict, empty_dict)
+
+            # minibatch_init_weight corresponds to W0 in figure 7 of TTT paper
+
+            minibatch_init_weight = weights
+
+            if dict_get_shape(weights) == self.init_weight_shape:
+                minibatch_init_weight = weights.apply(lambda t: repeat(t, '... -> b 1 (...)', b = batch * heads))
+
+            past_state = (minibatch_init_weight, empty_dict)
 
         past_last_update, past_last_momentum = past_state
 
@@ -734,7 +742,7 @@ class NeuralMemory(Module):
 
         # retrieve
 
-        retrieved = self.retrieve_memories(token, updates + weights, chunk_size = 1)
+        retrieved = self.retrieve_memories(token, weights, chunk_size = 1)
 
         # next state tuple
 
@@ -785,14 +793,18 @@ 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,
-            mem_model_weights + updates,
+            mem_model_weights,
             chunk_size = chunk_size,
             prev_layer_updates = prev_layer_updates
         )

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.2.5
+Version: 0.2.7
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -27,7 +27,7 @@ License: MIT License
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
 License-File: LICENSE
-Keywords: artificial intelligence,deep learning,linear attention,neural memory module,test time training
+Keywords: artificial intelligence,deep learning,linear attention,memory,test time training
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -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.7.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=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,,

titans_pytorch-0.2.5.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=Ew28waD9gf1wn-5Nkdc676u1I92IqzaOAw-tv0JXMwc,3777
-titans_pytorch/neural_memory.py,sha256=YiBsMiqYn-Hva4yhxfaqkGV857vZIASxi5Z0TT0FC10,24606
-titans_pytorch-0.2.5.dist-info/METADATA,sha256=x3RePuTDf3rUT3vtvge1X3Ry18Y3tV_swCgycbtSCjQ,6819
-titans_pytorch-0.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.5.dist-info/RECORD,,