titans-pytorch 0.0.49__py3-none-any.whl → 0.0.51__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -7,16 +7,48 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear
 
+# flex attention
+# https://pytorch.org/blog/flexattention/
+
+flex_attention = None
+
+try:
+    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+    if torch.cuda.is_available():
+        flex_attention = torch.compile(flex_attention)
+except ImportError:
+    pass
+
+def create_mac_block_mask(seq_len, window_size, persist_mem_len):
+
+    def create_mac_mask(b, h, q_idx, kv_idx):
+        is_persist_mem = kv_idx < persist_mem_len
+        causal_mask = q_idx >= (kv_idx - is_persist_mem)
+        block_diagonal = (q_idx // window_size) == ((kv_idx - is_persist_mem) // window_size)
+        return is_persist_mem | (~is_persist_mem & (causal_mask & block_diagonal))
+
+    block_mask = create_block_mask(create_mac_mask, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len + persist_mem_len, _compile = True)
+    return block_mask
+
+# einstein notation related
+
 from einops import einsum, repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
 
-from hyper_connections import get_init_and_expand_reduce_stream_functions
+# b - batch
+# n - sequence
+# h - heads
+# d - feature dimension
 
 # absolute and relative positions
 
 from axial_positional_embedding import ContinuousAxialPositionalEmbedding
 from rotary_embedding_torch import RotaryEmbedding
+
+# hyper connections / attend from x-transformers, which handles different queries and key lengths better
+
 from x_transformers.attend import Attend
+from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 # proposed neural memory
 

titans_pytorch/titans.py

@@ -56,6 +56,17 @@ def pack_one_with_inverse(t, pattern):
 
     return packed, inverse
 
+def Sequential(*modules):
+    modules = [*filter(exists, modules)]
+
+    if len(modules) == 0:
+        return nn.Identity()
+
+    if len(modules) == 1:
+        return modules[0]
+
+    return nn.Sequential(*modules)
+
 # softclamping gradients
 
 def softclamp_max(t, max_value):
@@ -124,9 +135,6 @@ class MemoryAttention(Module):
         ])
 
     def forward(self, x):
-
-        assert x.shape[-2] > 1, 'chunk size needs to be greater than 1 for using attention as memory'
-
         wq, wk, wv, ffw1, ffw2 = self.weights
 
         q = F.normalize(x @ wq, dim = -1)
@@ -162,11 +170,13 @@ class NeuralMemory(Module):
         heads = 1,
         model: Module | None = None,
         store_memory_loss_fn: Callable = default_loss_fn,
-        adaptive_step_transform: Callable = default_adaptive_step_transform,
+        adaptive_step_transform: Callable | None = None,
+        default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
+        activation: Module | None = None,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -224,11 +234,11 @@ class NeuralMemory(Module):
 
         # queries for retrieving from the model
 
-        self.to_queries = LinearNoBias(dim, dim_inner)
+        self.to_queries = Sequential(LinearNoBias(dim, dim_inner), activation)
 
         # keys and values for storing to the model
 
-        self.to_keys_values = LinearNoBias(dim, dim_inner * 2)
+        self.to_keys_values = Sequential(LinearNoBias(dim, dim_inner * 2), activation)
         self.store_memory_loss_fn = store_memory_loss_fn
 
         # empty memory embed
@@ -250,6 +260,9 @@ class NeuralMemory(Module):
             Rearrange('b n h -> (b h) n')
         )
 
+        if not exists(adaptive_step_transform):
+            adaptive_step_transform = partial(default_adaptive_step_transform, max_lr = default_step_transform_max_lr)
+
         self.adaptive_step_transform = adaptive_step_transform
 
         # allow for softclamp the gradient norms for storing memories

{titans_pytorch-0.0.49.dist-info → titans_pytorch-0.0.51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.49
+Version: 0.0.51
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -135,3 +135,12 @@ $ python train_mac.py
     year    = {2024}
 }
 ```
+
+```bibtex
+@inproceedings{Yang2024GatedDN,
+    title   = {Gated Delta Networks: Improving Mamba2 with Delta Rule},
+    author  = {Songlin Yang and Jan Kautz and Ali Hatamizadeh},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:274598177}
+}
+```

titans_pytorch-0.0.51.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=R_lWRe_iBkiWZ1iZAt1tNyjaTUyB5mb80mcYZHUKww0,11369
+titans_pytorch/titans.py,sha256=T04onF0xhcrosS-Qkx7fcx-Cqgh0TdU5JLdq9l8ayGg,15911
+titans_pytorch-0.0.51.dist-info/METADATA,sha256=OKqRYWucpjsLKKOksBrdHnvCAqgHgv2FUtL6l8hN2-Y,4484
+titans_pytorch-0.0.51.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.51.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.51.dist-info/RECORD,,

titans_pytorch-0.0.49.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=EMhxPt86Vr6LFvPm0OLMFYLaIY19khU9yIHkIhl2EMA,10316
-titans_pytorch/titans.py,sha256=tV2ej2PGUhMjSmDFV_wowX5q9hyp4SM4Jv3eJNu7cy8,15518
-titans_pytorch-0.0.49.dist-info/METADATA,sha256=hEpYHDqm_gffXybcotEmsK6o-siKrE7HwT_UgbOd-4o,4210
-titans_pytorch-0.0.49.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.49.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.49.dist-info/RECORD,,