titans-pytorch 0.1.2__py3-none-any.whl → 0.1.6__py3-none-any.whl

titans_pytorch/__init__.py

@@ -2,7 +2,8 @@ from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
     MemoryAttention,
-    FactorizedMemoryMLP
+    FactorizedMemoryMLP,
+    GatedResidualMemoryMLP
 )
 
 from titans_pytorch.mac_transformer import (

titans_pytorch/titans.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 from typing import Callable
+
 import math
 from functools import partial
 
 import torch
-from torch import nn, Tensor
+from torch import nn, cat, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList
 from torch.func import functional_call, vmap, grad
@@ -18,7 +19,7 @@ from titans_pytorch.associative_scan import (
 )
 
 import einx
-from einops import rearrange, repeat, pack, unpack
+from einops import rearrange, repeat, reduce, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
 """
@@ -95,6 +96,37 @@ class MultiheadRMSNorm(Module):
     def forward(self, x):
         return self.rmsnorm(x) * (self.gamma + 1.)
 
+# attention pool
+
+class AttentionPool(Module):
+    def __init__(
+        self,
+        dim,
+        chunk_size
+    ):
+        """
+        taken from Enformer https://www.nature.com/articles/s41592-021-01252-x , in turn taken from somewhere else
+        """
+        super().__init__()
+        self.split_chunks = Rearrange('b (n c) d -> b n c d', c = chunk_size)
+        self.to_attn_logits = nn.Linear(dim, dim)
+
+        # default to average pool
+
+        nn.init.zeros_(self.to_attn_logits.weight)
+        nn.init.zeros_(self.to_attn_logits.bias)
+
+    def forward(
+        self,
+        x
+    ):
+        x = self.split_chunks(x)
+        attn_logits = self.to_attn_logits(x)
+
+        attn = attn_logits.softmax(dim = -2)
+
+        return reduce(x * attn, 'b n c d -> b n d', 'sum')
+
 # classes
 
 class MemoryMLP(Module):
@@ -123,6 +155,46 @@ class MemoryMLP(Module):
 
         return x
 
+# memory mlp, but with gated residual + final projection
+
+class GatedResidualMemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth
+    ):
+        super().__init__()
+        self.depth = depth
+
+        self.weights = ParameterList([
+            ParameterList([
+                Parameter(torch.randn(dim, dim)),
+                Parameter(torch.randn(dim * 2, dim)),
+            ]) for _ in range(depth)
+        ])
+
+        self.final_proj = Parameter(torch.randn(dim, dim))
+
+        for param in self.parameters():
+            nn.init.xavier_uniform_(param)
+
+    def forward(
+        self,
+        x
+    ):
+        for weight, to_gates in self.weights:
+            res = x
+
+            x = x @ weight
+            x = F.silu(x)
+
+            # gated residual
+
+            gates = cat((x, res), dim = -1) @ to_gates
+            x = res.lerp(x, gates.sigmoid())
+
+        return x @ self.final_proj
+
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
 
@@ -224,6 +296,7 @@ class NeuralMemory(Module):
         default_step_transform_max_lr = 1e-2,
         per_parameter_lr_modulation = False, # allow outer network to control learning rate per weight matrix of memory network
         max_mem_layer_modulation = 1e1, # max of 10.
+        attn_pool_chunks = False,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         learned_mem_model_weights = True,
@@ -304,10 +377,17 @@ class NeuralMemory(Module):
         self.empty_memory_embed = nn.Parameter(torch.zeros(dim))
         nn.init.normal_(self.empty_memory_embed, std = 0.02)
 
+        # whether to use averaging of chunks, or attention pooling
+
+        if not attn_pool_chunks:
+            chunk_reduce_module = Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size)
+        else:
+            chunk_reduce_module = AttentionPool(dim, chunk_size = chunk_size)
+
         # learned adaptive learning rate and momentum
 
         self.to_momentum = Sequential(
-            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
+            chunk_reduce_module,
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         )
@@ -325,7 +405,7 @@ class NeuralMemory(Module):
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
-            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
+            chunk_reduce_module,
             LinearNoBias(dim, heads * self.num_memory_parameter_tensors),
             Rearrange('b n (h w) -> w (b h) n', h = heads),
             nn.Sigmoid()
@@ -340,7 +420,7 @@ class NeuralMemory(Module):
         # weight decay factor
 
         self.to_decay_factor = Sequential(
-            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
+            chunk_reduce_module,
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         )

{titans_pytorch-0.1.2.dist-info → titans_pytorch-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.2
+Version: 0.1.6
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -46,6 +46,7 @@ Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-transformers
 Provides-Extra: examples
+Requires-Dist: adam-atan2-pytorch>=0.1.18; extra == 'examples'
 Requires-Dist: wandb; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'

titans_pytorch-0.1.6.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
+titans_pytorch/titans.py,sha256=VMcPcKsoR3G13Um62Aa1HbdwrrV60ljPhP-yF40x90I,20555
+titans_pytorch-0.1.6.dist-info/METADATA,sha256=LJW26WfT9WB-0NfokLLHhcRpWnt76jwkXMt_FSTI3SM,4747
+titans_pytorch-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.6.dist-info/RECORD,,

titans_pytorch-0.1.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
-titans_pytorch/titans.py,sha256=cGWJHkOYmIeE6X383mZvyjusECBwbplVvK0cfgfhBxg,18634
-titans_pytorch-0.1.2.dist-info/METADATA,sha256=FWq5JIp1WY9dYpzatfGzfkcAGQFk-mEPwxF0wCrbM5w,4684
-titans_pytorch-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.2.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.2.dist-info/RECORD,,