titans-pytorch 0.0.65__py3-none-any.whl → 0.1.1__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -3,6 +3,8 @@ from typing import Callable
 from math import ceil
 from functools import partial
 
+import tqdm
+
 import torch
 from torch import nn, cat
 import torch.nn.functional as F
@@ -88,12 +90,6 @@ def pad_at_dim(t, pad, dim = -1, value = 0.):
 
 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
-
-    if not need_segment:
-        return seq, identity
-
     next_seq_len_mult = round_up_multiple(seq_len, segment_len)
 
     padding = next_seq_len_mult - seq_len
@@ -116,6 +112,29 @@ def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
 
     return seq, inverse
 
+# sampling related
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def gumbel_noise(t):
+    noise = torch.rand_like(t)
+    return -log(-log(noise))
+
+def gumbel_sample(t, temperature = 1.):
+    if temperature > 0.:
+        t = t / temperature + gumbel_noise(t)
+    return t.argmax(dim = -1, keepdim = True)
+
+# min_p
+# https://arxiv.org/abs/2407.01082
+
+def min_p_filter(logits, min_p = 0.1):
+    probs = logits.softmax(dim = -1)
+    max_probs = probs.amax(dim = -1, keepdim = True)
+    limit = min_p * max_probs
+    return torch.where(probs < limit, float('-inf'), logits)
+
 # feedforward and attention
 
 class GEGLU(Module):
@@ -500,6 +519,36 @@ class MemoryAsContextTransformer(Module):
         self.segment_len = segment_len
         self.num_persist_mem_tokens = num_persist_mem_tokens
 
+    @torch.no_grad()
+    def sample(
+        self,
+        prompt: Tensor,
+        seq_len: int,
+        temperature = 1.5,
+        filter_fn: Callable = min_p_filter,
+        filter_kwargs: dict = dict(
+            min_p = 0.1,
+        )
+    ):
+        was_training = self.training
+        self.eval()
+
+        prompt_seq_len, out = prompt.shape[-1], prompt.clone()
+        sample_num_times = max(0, seq_len - prompt_seq_len)
+
+        for _ in tqdm.tqdm(range(sample_num_times)):
+            logits = self.forward(out, disable_flex_attn = True)
+            logits = logits[:, -1]
+
+            logits = filter_fn(logits, **filter_kwargs)
+            sample = gumbel_sample(logits, temperature = temperature)
+
+            out = torch.cat((out, sample), dim = -1)
+
+        self.train(was_training)
+
+        return out[..., prompt_seq_len:]
+
     def forward(
         self,
         x,

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)
@@ -566,7 +592,12 @@ class NeuralMemory(Module):
         batch, seq_len = seq.shape[:2]
 
         if seq_len < self.chunk_size:
-            return self.init_empty_memory_embed(batch, seq_len)
+            out = self.init_empty_memory_embed(batch, seq_len)
+
+            if not return_aux_kv_loss:
+                return out
+
+            return out, self.zero
 
         if exists(past_state):
             past_state = tuple(TensorDict(d) for d in past_state)
@@ -580,7 +611,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.65.dist-info → titans_pytorch-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.65
+Version: 0.1.1
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -37,14 +37,15 @@ 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
 Requires-Dist: tensordict
 Requires-Dist: torch>=2.2
+Requires-Dist: tqdm
 Requires-Dist: x-transformers
 Provides-Extra: examples
-Requires-Dist: tqdm; extra == 'examples'
 Requires-Dist: wandb; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
@@ -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
@@ -99,7 +104,12 @@ transformer = MemoryAsContextTransformer(
 
 token_ids = torch.randint(0, 256, (1, 1023))
 
-logits = transformer(token_ids) # (1, 1023, 256)
+loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)
+loss.backward()
+
+# after much training
+
+sampled = transformer.sample(token_ids[:, :4], 512)
 ```
 
 ## Experiments

titans_pytorch-0.1.1.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=f_m559p9PLW1yxW6tQfrD43v0eMpeTUxpYdtb32UFgg,19031
+titans_pytorch/titans.py,sha256=cGWJHkOYmIeE6X383mZvyjusECBwbplVvK0cfgfhBxg,18634
+titans_pytorch-0.1.1.dist-info/METADATA,sha256=fxIRTFPGxq9RZMK4yRaAnoG3ym7dfOdaYiatBwmkS6Q,4684
+titans_pytorch-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.1.dist-info/RECORD,,

titans_pytorch-0.0.65.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=7PHBCbeB1LhHY5s3zAyYF0L3Mm7CNy4TOBbcpLX6bNE,17686
-titans_pytorch/titans.py,sha256=y6lJRErIoM6T2aTVFlf1GxSB0cpsmBZdSIj1DCHUCQ8,17486
-titans_pytorch-0.0.65.dist-info/METADATA,sha256=oDjEiufwOninsFDoCGbu691LXc1mey2OT7j6PNzkz0Q,4457
-titans_pytorch-0.0.65.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.65.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.65.dist-info/RECORD,,