titans-pytorch 0.1.1__tar.gz → 0.1.5__tar.gz

{titans_pytorch-0.1.1 → titans_pytorch-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.1
+Version: 0.1.5
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch

{titans_pytorch-0.1.1 → titans_pytorch-0.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.1"
+version = "0.1.5"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.1.1 → titans_pytorch-0.1.5}/tests/test_titans.py

@@ -11,12 +11,14 @@ def exists(v):
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
 @pytest.mark.parametrize('learned_mem_model_weights', (False, True))
+@pytest.mark.parametrize('attn_pool_chunks', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 @pytest.mark.parametrize('per_parameter_lr_modulation', (False, True))
 def test_titans(
     seq_len,
     silu,
     learned_mem_model_weights,
+    attn_pool_chunks,
     max_grad_norm,
     per_parameter_lr_modulation
 ):
@@ -24,6 +26,7 @@ def test_titans(
         dim = 384,
         chunk_size = 64,
         activation = nn.SiLU() if silu else None,
+        attn_pool_chunks = attn_pool_chunks,
         max_grad_norm = max_grad_norm,
         per_parameter_lr_modulation = per_parameter_lr_modulation,
         learned_mem_model_weights = learned_mem_model_weights
@@ -50,10 +53,12 @@ def test_titans_attn_memory():
 
     assert seq.shape == retrieved.shape
 
+@pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('num_persist_mem_tokens', (0, 16))
 @pytest.mark.parametrize('num_longterm_mem_tokens', (0, 16))
 @pytest.mark.parametrize('neural_mem_gate_attn_output', (False, True))
 def test_mac(
+    seq_len,
     num_persist_mem_tokens,
     num_longterm_mem_tokens,
     neural_mem_gate_attn_output
@@ -70,13 +75,15 @@ def test_mac(
         neural_mem_gate_attn_output = neural_mem_gate_attn_output
     )
 
-    x = torch.randint(0, 256, (1, 1023))
+    x = torch.randint(0, 256, (1, seq_len))
 
     logits = transformer(x)
-    assert logits.shape == (1, 1023, 256)
+    assert logits.shape == (1, seq_len, 256)
 
+@pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('sliding', (True, False))
 def test_flex(
+    seq_len,
     sliding
 ):
     if not (torch.cuda.is_available() and exists(flex_attention)):
@@ -91,7 +98,7 @@ def test_flex(
         sliding = sliding
     ).cuda()
 
-    seq = torch.randn(1, 1019, 512).cuda()
+    seq = torch.randn(1, seq_len, 512).cuda()
 
     out_flex, _ = attn(seq)
     out_non_flex, _ = attn(seq, disable_flex_attn = True)

{titans_pytorch-0.1.1 → titans_pytorch-0.1.5}/titans_pytorch/mac_transformer.py

@@ -528,7 +528,8 @@ class MemoryAsContextTransformer(Module):
         filter_fn: Callable = min_p_filter,
         filter_kwargs: dict = dict(
             min_p = 0.1,
-        )
+        ),
+        show_progress = True
     ):
         was_training = self.training
         self.eval()
@@ -536,7 +537,9 @@ class MemoryAsContextTransformer(Module):
         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)):
+        iter_wrap = tqdm.tqdm if show_progress else identity
+
+        for _ in iter_wrap(range(sample_num_times)):
             logits = self.forward(out, disable_flex_attn = True)
             logits = logits[:, -1]
 

{titans_pytorch-0.1.1 → titans_pytorch-0.1.5}/titans_pytorch/titans.py

@@ -18,7 +18,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 +95,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):
@@ -224,6 +255,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 +336,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 +364,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 +379,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')
         )