titans-pytorch 0.0.64__py3-none-any.whl → 0.1.0__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -20,13 +20,21 @@ try:
 except ImportError:
     pass
 
-def create_mac_block_mask(seq_len, window_size, persist_mem_len):
+def create_mac_block_mask(seq_len, window_size, persist_mem_len, sliding = False):
 
-    def create_mac_mask(b, h, q_idx, kv_idx):
+    def create_mac_mask(_, __, q_idx, kv_idx):
         is_persist_mem = kv_idx < persist_mem_len
-        causal_mask = q_idx >= (kv_idx - persist_mem_len)
-        block_diagonal = (q_idx // window_size) == ((kv_idx - persist_mem_len) // window_size)
-        return is_persist_mem | (~is_persist_mem & (causal_mask & block_diagonal))
+        kv_without_mem = kv_idx - persist_mem_len
+        causal_mask = q_idx >= kv_without_mem
+
+        if not sliding:
+            block_diagonal = (q_idx // window_size) == (kv_without_mem // window_size)
+            causal_mask = causal_mask & block_diagonal
+        else:
+            sliding_mask = (q_idx - kv_without_mem) <= window_size
+            causal_mask = causal_mask & sliding_mask
+
+        return is_persist_mem | (~is_persist_mem & causal_mask)
 
     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
@@ -73,7 +81,12 @@ def identity(t):
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
-def pad_and_segment_with_inverse(seq, segment_len):
+def pad_at_dim(t, pad, dim = -1, value = 0.):
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
+
+def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
     batch, seq_len = seq.shape[:2]
 
     need_segment = seq_len >= segment_len
@@ -89,13 +102,15 @@ def pad_and_segment_with_inverse(seq, segment_len):
     if needs_pad:
         seq = F.pad(seq, (0, 0, 0, padding))
 
-    seq = rearrange(seq, 'b (w n) d -> (b w) n d', n = segment_len)
+    if fold_into_batch:
+        seq = rearrange(seq, 'b (w n) d -> (b w) n d', n = segment_len)
 
     def inverse(out):
-        out = rearrange(out, '(b w) n d -> b (w n) d', b = batch)
+        if fold_into_batch:
+            out = rearrange(out, '(b w) n d -> b (w n) d', b = batch)
 
         if needs_pad:
-            out = out[:, :-padding]
+            out = out[..., :-padding, :]
 
         return out
 
@@ -127,6 +142,7 @@ class SegmentedAttention(Module):
         num_longterm_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
+        sliding = False,
         accept_value_residual = False,
         attend_kwargs: dict = dict(),
         use_flex_attn = False
@@ -155,6 +171,8 @@ class SegmentedAttention(Module):
         total_segment_len = segment_len + num_longterm_mem_tokens
         self.total_segment_len = total_segment_len
 
+        self.sliding = sliding # sliding window attn - doubt their non-sliding results being the best. local attention with overlapping windows is very strong
+
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
 
@@ -210,7 +228,7 @@ class SegmentedAttention(Module):
         # prep flex attention
 
         if not exists(flex_attn_fn):
-            block_mask = create_mac_block_mask(seq_len, self.total_segment_len, self.num_persist_mem_tokens)
+            block_mask = create_mac_block_mask(seq_len, self.total_segment_len, self.num_persist_mem_tokens, self.sliding)
 
             flex_attn_fn = partial(flex_attention, block_mask = block_mask)
 
@@ -243,7 +261,7 @@ class SegmentedAttention(Module):
 
         # auto pad to multiple
 
-        seq, inverse_segment = pad_and_segment_with_inverse(seq, total_segment_len)
+        seq, inverse_segment = pad_and_segment_with_inverse(seq, total_segment_len, fold_into_batch = False)
 
         # attention
 
@@ -260,14 +278,45 @@ class SegmentedAttention(Module):
             mix = self.to_learned_v_mix(seq)
             v = v.lerp(value_residual, mix)
 
-        # take care of persistent memory key / values
-
-        pmk, pmv = repeat(self.persistent_memory, 'kv ... -> kv b ...', b = seq.shape[0])
-
         # relative positions
 
         q, k = self.rotary_emb.rotate_queries_with_cached_keys(q, k)
 
+        # fold
+
+        q, k, v = tuple(rearrange(t, 'b h (w n) d -> (b w) h n d', n = total_segment_len) for t in (q, k, v))
+
+        # maybe sliding for cpu
+
+        attend_kwargs = dict()
+
+        if self.sliding:
+            k, v = tuple(rearrange(t, '(b w) ... -> b w ...', b = batch) for t in (k, v))
+            k, v = tuple(pad_at_dim(t, (1, 0), value = 0., dim = 1) for t in (k, v))
+            k = cat((k[:, :-1], k[:, 1:]), dim = -2)
+            v = cat((v[:, :-1], v[:, 1:]), dim = -2)
+            k, v = tuple(rearrange(t, 'b w ... -> (b w) ...') for t in (k, v))
+
+            # take care of masking
+
+            idx = torch.arange(seq.shape[-2], device = seq.device)
+            q_idx = rearrange(idx, '(w n) -> w n', n = total_segment_len)
+            k_idx = pad_at_dim(q_idx, (1, 0), dim = 0, value = -1e4)
+            k_idx = cat((k_idx[:-1], k_idx[1:]), dim = -1)
+
+            q_idx = rearrange(q_idx, 'w i -> w i 1')
+            k_idx = rearrange(k_idx, 'w j -> w 1 j')
+
+            sliding_mask = (q_idx - k_idx) <= total_segment_len
+            sliding_mask = F.pad(sliding_mask, (self.num_persist_mem_tokens, 0), value = True)
+
+            sliding_mask = repeat(sliding_mask, 'w i j -> (b w) 1 i j', b = batch)
+            attend_kwargs.update(mask = sliding_mask)
+
+        # take care of persistent memory key / values
+
+        pmk, pmv = repeat(self.persistent_memory, 'kv ... -> kv b ...', b = k.shape[0])
+
         # persistent memory
 
         k = cat((pmk, k), dim = -2)
@@ -275,12 +324,14 @@ class SegmentedAttention(Module):
 
         # attention
 
-        out, _ = self.attend(q, k, v)
+        out, _ = self.attend(q, k, v, **attend_kwargs)
 
         out = self.merge_heads(out)
 
         out = self.to_out(out)
 
+        out = rearrange(out, '(b w) n d -> b (w n) d', b = batch)
+
         out = inverse_segment(out)
 
         return out, orig_v
@@ -349,7 +400,8 @@ class MemoryAsContextTransformer(Module):
         neural_memory_kwargs: dict = dict(),
         neural_memory_layers: tuple[int, ...] | None = None,
         aux_kv_recon_loss_weight = 0.,
-        use_flex_attn = False
+        use_flex_attn = False,
+        sliding_window_attn = False
     ):
         super().__init__()
 
@@ -366,6 +418,10 @@ class MemoryAsContextTransformer(Module):
 
         self.longterm_mems = nn.Parameter(torch.randn(num_longterm_mem_tokens, dim) * 0.02)
 
+        # maybe sliding window attn
+
+        self.sliding_window_attn = sliding_window_attn
+
         # hyper conection
 
         init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
@@ -396,7 +452,8 @@ class MemoryAsContextTransformer(Module):
                 use_flex_attn = use_flex_attn,
                 accept_value_residual = not is_first,
                 num_longterm_mem_tokens = num_longterm_mem_tokens,
-                num_persist_mem_tokens = num_persist_mem_tokens
+                num_persist_mem_tokens = num_persist_mem_tokens,
+                sliding = sliding_window_attn
             )
 
             mem = None
@@ -489,7 +546,7 @@ class MemoryAsContextTransformer(Module):
         flex_attn_fn = None
 
         if use_flex_attn:
-            block_mask = create_mac_block_mask(seq_len_with_mem, segment_len + num_longterm_mem_tokens, self.num_persist_mem_tokens)
+            block_mask = create_mac_block_mask(seq_len_with_mem, segment_len + num_longterm_mem_tokens, self.num_persist_mem_tokens, self.sliding_window_attn)
             flex_attn_fn = partial(flex_attention, block_mask = block_mask)
 
         # value residual

titans_pytorch/titans.py

@@ -17,6 +17,7 @@ from titans_pytorch.associative_scan import (
     pad_at_dim
 )
 
+import einx
 from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
@@ -26,6 +27,7 @@ b - batch
 n - sequence
 d - feature dimension
 c - intra-chunk
+w - num memory network weight parameters
 """
 
 LinearNoBias = partial(Linear, bias = False)
@@ -220,6 +222,8 @@ class NeuralMemory(Module):
         store_memory_loss_fn: Callable = default_loss_fn,
         adaptive_step_transform: Callable | None = None,
         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.
         pre_rmsnorm = True,
         post_rmsnorm = True,
         learned_mem_model_weights = True,
@@ -250,7 +254,7 @@ class NeuralMemory(Module):
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
         self.combine_heads = LinearNoBias(dim_inner, dim) if heads > 1 else nn.Identity()
 
-        self.retrieve_gate = nn.Sequential(
+        self.retrieve_gate = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> b h n 1'),
             nn.Sigmoid()
@@ -270,6 +274,8 @@ class NeuralMemory(Module):
 
         self.memory_model = model
 
+        self.num_memory_parameter_tensors = len(set(model.parameters()))
+
         # the chunk size within the paper where adaptive step, momentum, weight decay are shared
 
         self.chunk_size = chunk_size
@@ -299,15 +305,14 @@ class NeuralMemory(Module):
         nn.init.normal_(self.empty_memory_embed, std = 0.02)
 
         # learned adaptive learning rate and momentum
-        # todo - explore mlp layerwise learned lr / momentum
 
-        self.to_momentum = nn.Sequential(
+        self.to_momentum = Sequential(
             Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         )
 
-        self.to_adaptive_step = nn.Sequential(
+        self.to_adaptive_step = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n')
         )
@@ -317,13 +322,24 @@ class NeuralMemory(Module):
 
         self.adaptive_step_transform = adaptive_step_transform
 
+        # per layer learning rate modulation
+
+        self.to_layer_modulation = Sequential(
+            Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
+            LinearNoBias(dim, heads * self.num_memory_parameter_tensors),
+            Rearrange('b n (h w) -> w (b h) n', h = heads),
+            nn.Sigmoid()
+        ) if per_parameter_lr_modulation else None
+
+        self.max_mem_layer_modulation = max_mem_layer_modulation
+
         # allow for softclamp the gradient norms for storing memories
 
         self.max_grad_norm = max_grad_norm
 
         # weight decay factor
 
-        self.to_decay_factor = nn.Sequential(
+        self.to_decay_factor = Sequential(
             Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size),
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
@@ -387,6 +403,11 @@ class NeuralMemory(Module):
         adaptive_momentum = self.to_momentum(seq).sigmoid()
         decay_factor = self.to_decay_factor(seq).sigmoid()
 
+        need_layer_lr_mod = exists(self.to_layer_modulation)
+
+        if need_layer_lr_mod:
+            layer_lr_mod = self.to_layer_modulation(seq) * self.max_mem_layer_modulation
+
         # keys and values
 
         keys, values = self.to_keys_values(seq).chunk(2, dim = -1)
@@ -418,6 +439,11 @@ class NeuralMemory(Module):
 
         grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))
 
+        # maybe per layer modulation
+
+        if need_layer_lr_mod:
+            grads = TensorDict({name: einx.multiply('b h, b h ... -> b h ...', layer_lr_mod, t) for layer_lr_mod, (name, t) in zip(layer_lr_mod, grads.items())})
+
         # negative gradients, adaptive lr already applied as loss weight
 
         surprises = grads.apply(lambda t: -t)
@@ -469,7 +495,7 @@ class NeuralMemory(Module):
 
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
-            update = scan_fn(1. - decay_factor, momentum) # momentum is S / surprise in the paper
+            update = scan_fn(1. - decay_factor, momentum)
 
             updates[param_name] = inverse_pack(update)
             next_momentum[param_name] = inverse_pack(momentum)
@@ -580,7 +606,6 @@ class NeuralMemory(Module):
 
         past_weights, _ = past_state
 
-
         retrieved = self.retrieve_memories(seq, past_weights + updates)
 
         if not return_aux_kv_loss:

{titans_pytorch-0.0.64.dist-info → titans_pytorch-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.64
+Version: 0.1.0
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -37,6 +37,7 @@ Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
 Requires-Dist: axial-positional-embedding>=0.3.5
 Requires-Dist: einops>=0.8.0
+Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
@@ -58,6 +59,10 @@ Description-Content-Type: text/markdown
 
 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.
 
+## Appreciation
+
+- [@sentialx](https://github.com/sentialx) for sharing his early experimental results with me
+
 ## Install
 
 ```bash

titans_pytorch-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=7PHBCbeB1LhHY5s3zAyYF0L3Mm7CNy4TOBbcpLX6bNE,17686
+titans_pytorch/titans.py,sha256=L3Mu6pDnimD4MNn_832trFEJgXOPjxSdTrB9jiSUSTk,18533
+titans_pytorch-0.1.0.dist-info/METADATA,sha256=LuWDzv-NbGxYKOMThN_WKQWDueyIsOAMSwwiE_BDraI,4595
+titans_pytorch-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.0.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.0.dist-info/RECORD,,

titans_pytorch-0.0.64.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=7voYtbD_ErCD_JjvwhAiunUWtSIsIxGJAaf2aRB3c2s,15349
-titans_pytorch/titans.py,sha256=y6lJRErIoM6T2aTVFlf1GxSB0cpsmBZdSIj1DCHUCQ8,17486
-titans_pytorch-0.0.64.dist-info/METADATA,sha256=K63jobSfTdn-aFpEpZgolu4zSIvgUzF2rDuoCHGXkgE,4457
-titans_pytorch-0.0.64.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.64.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.64.dist-info/RECORD,,