titans-pytorch 0.1.22__tar.gz → 0.1.23__tar.gz

{titans_pytorch-0.1.22 → titans_pytorch-0.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.22
+Version: 0.1.23
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -78,7 +78,7 @@ from titans_pytorch import NeuralMemory
 
 mem = NeuralMemory(
     dim = 384,
-    chunk_size = 64
+    chunk_size = 64 # set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage)
 ).cuda()
 
 seq = torch.randn(2, 1024, 384).cuda()

{titans_pytorch-0.1.22 → titans_pytorch-0.1.23}/README.md

@@ -24,7 +24,7 @@ from titans_pytorch import NeuralMemory
 
 mem = NeuralMemory(
     dim = 384,
-    chunk_size = 64
+    chunk_size = 64 # set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage)
 ).cuda()
 
 seq = torch.randn(2, 1024, 384).cuda()

{titans_pytorch-0.1.22 → titans_pytorch-0.1.23}/pyproject.toml

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

{titans_pytorch-0.1.22 → titans_pytorch-0.1.23}/tests/test_titans.py

@@ -8,6 +8,9 @@ from titans_pytorch.mac_transformer import flex_attention, SegmentedAttention, M
 def exists(v):
     return v is not None
 
+def diff(x, y):
+    return (x - y).abs().amax()
+
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
 @pytest.mark.parametrize('learned_mem_model_weights', (False, True))
@@ -133,6 +136,39 @@ def test_mac_sampling(sliding):
 
     assert torch.allclose(sampled, sampled_with_cache)
 
+@pytest.mark.parametrize('seq_len', (2, 64))
+def test_neural_mem_inference(
+    seq_len
+):
+    mem = NeuralMemory(
+        dim = 384,
+        chunk_size = 64,
+    )
+
+    seq = torch.randn(2, seq_len, 384)
+    parallel_retrieved = mem(seq)
+
+    assert seq.shape == parallel_retrieved.shape
+
+    mem_model_state = None
+    cache_store_seq = None
+    sequential_retrieved = []
+
+    for ind, token in enumerate(seq.unbind(dim = 1)):
+
+        one_retrieved, cache_store_seq, mem_model_state = mem.forward_inference(
+            token,
+            seq_index = ind,
+            cache_store_seq = cache_store_seq,
+            mem_model_state = mem_model_state
+        )
+
+        sequential_retrieved.append(one_retrieved)
+
+    sequential_retrieved = torch.cat(sequential_retrieved, dim = -2)
+
+    assert torch.allclose(parallel_retrieved, sequential_retrieved, atol = 1e-5)
+
 @pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('sliding', (True, False))
 def test_flex(
@@ -157,3 +193,23 @@ def test_flex(
     out_non_flex, _ = attn(seq, disable_flex_attn = True)
 
     assert torch.allclose(out_flex, out_non_flex, atol = 1e-5)
+
+def test_assoc_scan():
+    from titans_pytorch.titans import AssocScan
+    import torch.nn.functional as F
+
+    scan = AssocScan()
+
+    gates = torch.randn(2, 1024, 512).sigmoid()
+    inputs = torch.randn(2, 1024, 512)
+
+    output = scan(gates, inputs)
+
+    gates1, gates2 = gates[:, :512], gates[:, 512:]
+    inputs1, inputs2 = inputs[:, :512], inputs[:, 512:]
+
+    first_half = scan(gates1, inputs1)
+
+    second_half = scan(gates2, inputs2, prev = inputs2[:, -1])
+
+    assert torch.allclose(output[:, -1], second_half[:, -1], atol = 1e-5)

{titans_pytorch-0.1.22 → titans_pytorch-0.1.23}/titans_pytorch/titans.py

@@ -44,6 +44,16 @@ def default(v, d):
 def xnor(x, y):
     return not (x ^ y)
 
+def safe_cat(inputs, dim = -2):
+    inputs = tuple(filter(exists, inputs))
+
+    if len(inputs) == 0:
+        return None
+    elif len(inputs) == 1:
+        return inputs[0]
+
+    return cat(inputs, dim = dim)
+
 def identity(t):
     return t
 
@@ -314,7 +324,11 @@ class AssocScan(Module):
         super().__init__()
         self.use_accelerated = use_accelerated
 
-    def forward(self, gates, inputs):
+    def forward(self, gates, inputs, prev = None):
+
+        if exists(prev):
+            inputs, _ = pack([prev, inputs], 'b * d')
+            gates = pad_at_dim(gates, (1, 0), value = 1., dim = -2)
 
         if not self.use_accelerated:
             _, outputs = associative_scan(binary_operator, (gates, inputs))
@@ -678,7 +692,7 @@ class NeuralMemory(Module):
     def retrieve_memories(
         self,
         seq,
-        past_weights: dict[str, Tensor] | None = None,
+        past_weights: dict[str, Tensor],
         chunk_size = None
     ):
         chunk_size = default(chunk_size, self.retrieve_chunk_size)
@@ -700,13 +714,7 @@ class NeuralMemory(Module):
         # the parameters of the memory model stores the memories of the key / values
         # when the MLP has only 1 weight matrix, it is equivalent to `kv` fast weight memories from linear attention literature (recall fetching of memories is q @ (kv)) / schmidhuber's paper
 
-        curr_weights = TensorDict(dict(self.memory_model.named_parameters()))
-
-        if exists(past_weights):
-            past_weights = TensorDict(past_weights)
-            assert past_weights.keys() == curr_weights.keys()
-
-            curr_weights = curr_weights + past_weights
+        curr_weights = TensorDict(past_weights)
 
         # sequence Float['b n d'] to queries
 
@@ -753,6 +761,56 @@ class NeuralMemory(Module):
 
         return values[:, :seq_len]
 
+    def forward_inference(
+        self,
+        token: Tensor,
+        seq_index = None, # the index of the token in the sequence, starts at 0
+        mem_model_state = None,
+        cache_store_seq = None
+    ):
+        seq_index = default(seq_index, 0)
+        curr_seq_len = seq_index + 1
+        batch = token.shape[0]
+
+        if token.ndim == 2:
+            token = rearrange(token, 'b d -> b 1 d')
+
+        # init memory model if needed
+
+        if not exists(mem_model_state):
+            mem_model_state = self.init_weights_and_momentum()
+
+        # increment the sequence cache which is at most the chunk size
+
+        cache_store_seq = safe_cat((cache_store_seq, token), dim = -2)
+
+        # early return empty memory, when no memories are stored for steps < first chunk size
+
+        if curr_seq_len < self.chunk_size:
+            empty_mem = self.init_empty_memory_embed(batch, 1)
+
+            return empty_mem, cache_store_seq, mem_model_state
+
+        # store if storage sequence cache hits the chunk size
+
+        store_seq_cache_len = cache_store_seq.shape[-2]
+
+        if store_seq_cache_len == self.chunk_size:
+            updates, _ = self.store_memories(cache_store_seq, mem_model_state)
+
+            past_weights, past_momentum = mem_model_state
+            mem_model_state = (past_weights + updates, past_momentum)
+
+            cache_store_seq = None
+
+        # retrieve
+
+        past_weights, _ = mem_model_state
+
+        retrieved = self.retrieve_memories(token, past_weights, chunk_size = 1)
+
+        return retrieved, cache_store_seq, mem_model_state
+
     def forward(
         self,
         seq,