titans-pytorch 0.2.14__tar.gz → 0.2.16__tar.gz

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/PKG-INFO

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

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/pyproject.toml

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

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/tests/test_titans.py

@@ -148,7 +148,7 @@ def test_mac(
     assert logits.shape == (1, seq_len, 256)
 
 @pytest.mark.parametrize('sliding', (False, True))
-@pytest.mark.parametrize('mem_layers', (()))
+@pytest.mark.parametrize('mem_layers', ((), None))
 @pytest.mark.parametrize('longterm_mems', (0, 4, 16))
 @pytest.mark.parametrize('prompt_len', (4, 16))
 @torch_default_dtype(torch.float64)
@@ -190,7 +190,6 @@ def test_neural_mem_inference(
     prompt_len,
     mem_chunk_size
 ):
-    pytest.skip()
 
     mem = NeuralMemory(
         dim = 384,
@@ -218,7 +217,7 @@ def test_neural_mem_inference(
 
     for token in seq.unbind(dim = 1):
 
-        one_retrieved, state = mem.forward_inference(
+        one_retrieved, state = mem.forward(
             token,
             state = state,
         )

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/titans_pytorch/mac_transformer.py

@@ -491,6 +491,7 @@ class MemoryAsContextTransformer(Module):
         neural_memory_layers: tuple[int, ...] | None = None,
         use_flex_attn = False,
         sliding_window_attn = False,
+        neural_mem_weight_residual = False
     ):
         super().__init__()
 
@@ -524,6 +525,8 @@ class MemoryAsContextTransformer(Module):
 
         neural_memory_layers = default(neural_memory_layers, layers)
 
+        self.neural_mem_weight_residual = neural_mem_weight_residual
+
         # mem, attn, and feedforward layers
 
         for layer in layers:
@@ -739,6 +742,10 @@ class MemoryAsContextTransformer(Module):
 
         value_residual = None
 
+        # neural mem weight residual
+
+        mem_weight_residual = None
+
         # when inferencing, only do one token at a time
 
         if is_inferencing:
@@ -761,16 +768,14 @@ class MemoryAsContextTransformer(Module):
 
                 mem_input, add_residual = mem_hyper_conn(x)
 
-                if not is_inferencing:
-                    retrieved, next_neural_mem_cache = mem(
-                        mem_input
-                    )
+                retrieved, next_neural_mem_cache = mem.forward(
+                    mem_input,
+                    state = next(neural_mem_caches, None),
+                    prev_weights = mem_weight_residual
+                )
 
-                else:
-                    (retrieved, next_neural_mem_cache) = mem.forward_inference(
-                        mem_input,
-                        state = next(neural_mem_caches, None),
-                    )
+                if self.neural_mem_weight_residual:
+                    mem_weight_residual = next_neural_mem_cache.updates
 
                 if self.gate_attn_output:
                     attn_out_gates = retrieved.sigmoid()

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/titans_pytorch/neural_memory.py

@@ -494,7 +494,8 @@ class NeuralMemory(Module):
         seq,
         weights: dict[str, Tensor] | None = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
-        seq_index = 0
+        seq_index = 0,
+        prev_weights = None
     ):
         batch, seq_len, heads, chunk_size = *seq.shape[:2], self.heads, self.store_chunk_size
 
@@ -560,6 +561,12 @@ class NeuralMemory(Module):
 
         adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = chunk_size)
 
+        # maybe add previous layer weight
+
+        if exists(prev_weights):
+            prev_weights = prev_weights.apply(lambda t: t[:, -1:])
+            weights_for_surprise = weights_for_surprise + prev_weights
+
         # flatten batch and time if surprise depends on previous layer memory model
 
         weights_for_surprise = rearrange_dict_values(weights_for_surprise, 'b n ... -> (b n) ...')
@@ -655,15 +662,19 @@ class NeuralMemory(Module):
         self,
         seq,
         past_weights: dict[str, Tensor],
+        chunk_size = None,
+        need_pad = True
     ):
-        chunk_size = self.retrieve_chunk_size
+        chunk_size = default(chunk_size, self.retrieve_chunk_size)
         batch, seq_len = seq.shape[:2]
 
         seq = self.retrieve_norm(seq)
 
-        needs_pad = chunk_size > 1
+        need_pad = need_pad or chunk_size > 1
+
+        if need_pad:
+            seq = pad_at_dim(seq, (1, 0), dim = 1)
 
-        seq = pad_at_dim(seq, (1, 0), dim = 1)
         seq_len_plus_one = seq.shape[-2]
 
         next_seq_len = round_up_multiple(seq_len_plus_one, chunk_size)
@@ -718,94 +729,37 @@ class NeuralMemory(Module):
 
         # restore, pad with empty memory embed
 
-        values = values[:, 1:(seq_len + 1)]
-
-        return values
-
-    @torch.no_grad()
-    def forward_inference(
-        self,
-        token: Tensor,
-        state: NeuralMemCache | None = None,
-    ):
-        # unpack previous state
-
-        if not exists(state):
-            state = (0, None, None, None, None)
-
-        seq_index, weights, cache_store_seq, past_states, updates = state
-
-        curr_seq_len = seq_index + 1
-        batch = token.shape[0]
-
-        if token.ndim == 2:
-            token = rearrange(token, 'b d -> b 1 d')
-
-        assert token.shape[1] == 1
-
-        # 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:
-            retrieve = self.retrieve_memories(token, weights, chunk_size = 1)
-
-            output = retrieve, NeuralMemCache(curr_seq_len, weights, cache_store_seq, past_states, updates)
-
-            return output
-
-        # store if storage sequence cache hits the chunk size
+        if need_pad:
+            values = values[:, 1:]
 
-        next_states = past_states
-        store_seq_cache_len = cache_store_seq.shape[-2]
-
-        if not exists(updates):
-            updates = weights.clone().zero_()
-            updates = repeat_dict_values(updates, '... -> b 1 ...', b = batch)
-        else:
-            updates = updates.apply(lambda t: t[:, -1:])
-
-        if store_seq_cache_len == self.chunk_size:
-
-            next_updates, store_state = self.store_memories(
-                cache_store_seq,
-                weights,
-                past_state = past_states,
-            )
-
-            updates = next_updates
-            cache_store_seq = None
-            next_states = store_state.states
-
-        # retrieve
-
-        retrieved = self.retrieve_memories(token, updates, chunk_size = 1)
-
-        # next state tuple
-
-        next_store_state = NeuralMemCache(curr_seq_len, weights, cache_store_seq, next_states, updates)
-
-        return retrieved, next_store_state
+        return values[:, :seq_len]
 
     def forward(
         self,
         seq,
         store_seq = None,
         state: NeuralMemCache | None = None,
+        prev_weights = None
     ):
+        if seq.ndim == 2:
+            seq = rearrange(seq, 'b d -> b 1 d')
+
+        is_single_token = seq.shape[1] == 1
+
         if not exists(state):
             state = (0, None, None, None, None)
 
         seq_index, weights, cache_store_seq, past_state, updates = state
 
-        assert not exists(cache_store_seq) or is_empty_tensor(cache_store_seq)
-
         # store
 
         store_seq = default(store_seq, seq)
 
+        # take care of cache
+
+        if exists(cache_store_seq):
+            store_seq = safe_cat((cache_store_seq, store_seq))
+
         # functions
 
         # compute split sizes of sequence
@@ -861,6 +815,7 @@ class NeuralMemory(Module):
                 weights,
                 seq_index = seq_index,
                 past_state = past_state,
+                prev_weights = prev_weights
             )
 
             seq_index = next_neural_mem_state.seq_index
@@ -883,9 +838,21 @@ class NeuralMemory(Module):
 
         # retrieve
 
+        need_pad = True
+        retrieve_chunk_size = None
+
+        if is_single_token:
+            retrieve_chunk_size = 1
+            need_pad = False
+
+            last_update, _ = past_state
+            updates = rearrange_dict_values(last_update, 'b ... -> b 1 ...')
+
         retrieved = self.retrieve_memories(
             seq,
-            updates
+            updates,
+            chunk_size = retrieve_chunk_size,
+            need_pad = need_pad,
         )
 
         return retrieved, next_neural_mem_state

{titans_pytorch-0.2.14 → titans_pytorch-0.2.16}/train_mac.py

@@ -36,8 +36,9 @@ NEURAL_MEM_MOMENTUM = True
 NEURAL_MEM_QK_NORM = True
 NEURAL_MEM_MAX_LR = 1e-1
 WINDOW_SIZE = 32
-NEURAL_MEM_SEGMENT_LEN = 2                      # set smaller for more granularity for learning rate / momentum etc
+NEURAL_MEM_SEGMENT_LEN = 4                      # set smaller for more granularity for learning rate / momentum etc
 NEURAL_MEM_BATCH_SIZE = 128                     # set smaller to update the neural memory weights more often as it traverses the sequence
+NEURAL_MEM_WEIGHT_RESIDUAL = True
 SLIDING_WINDOWS = True
 STORE_ATTN_POOL_CHUNKS = True                   # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
 MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
@@ -87,7 +88,7 @@ model = MemoryAsContextTransformer(
     neural_memory_segment_len = NEURAL_MEM_SEGMENT_LEN,
     neural_memory_batch_size = NEURAL_MEM_BATCH_SIZE,
     neural_mem_gate_attn_output = NEURAL_MEM_GATE_ATTN_OUTPUT,
-    default_step_transform_max_lr = NEURAL_MEM_MAX_LR,
+    neural_mem_weight_residual = NEURAL_MEM_WEIGHT_RESIDUAL,
     use_flex_attn = USE_FLEX_ATTN,
     sliding_window_attn = SLIDING_WINDOWS,
     neural_memory_model = MemoryMLP(
@@ -100,6 +101,7 @@ model = MemoryAsContextTransformer(
         attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
         qk_rmsnorm = NEURAL_MEM_QK_NORM,
         momentum = NEURAL_MEM_MOMENTUM,
+        default_step_transform_max_lr = NEURAL_MEM_MAX_LR,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
         per_parameter_lr_modulation = MEMORY_MODEL_PER_LAYER_LEARNED_LR
     )