titans-pytorch 0.2.12__py3-none-any.whl → 0.2.15__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -761,16 +761,10 @@ class MemoryAsContextTransformer(Module):
 
                 mem_input, add_residual = mem_hyper_conn(x)
 
-                if not is_inferencing:
-                    retrieved, next_neural_mem_cache = mem(
-                        mem_input
-                    )
-
-                else:
-                    (retrieved, next_neural_mem_cache) = mem.forward_inference(
-                        mem_input,
-                        state = next(neural_mem_caches, None),
-                    )
+                retrieved, next_neural_mem_cache = mem.forward(
+                    mem_input,
+                    state = next(neural_mem_caches, None),
+                )
 
                 if self.gate_attn_output:
                     attn_out_gates = retrieved.sigmoid()

titans_pytorch/neural_memory.py

@@ -284,7 +284,7 @@ class NeuralMemory(Module):
         adaptive_step_transform: Callable | None = None,
         default_step_transform_max_lr = 1.,
         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.
+        max_mem_layer_modulation = 1., # max of 10.
         attn_pool_chunks = False,
         momentum = True,
         pre_rmsnorm = True,
@@ -293,6 +293,9 @@ class NeuralMemory(Module):
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         activation: Module | None = None,
+        init_adaptive_step_bias = None,
+        init_momentum_bias = None,
+        init_decay_bias = None,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -411,12 +414,12 @@ class NeuralMemory(Module):
         # learned adaptive learning rate and momentum
 
         self.to_momentum = Sequential(
-            LinearNoBias(dim, heads),
+            nn.Linear(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         ) if momentum else None
 
         self.to_adaptive_step = Sequential(
-            LinearNoBias(dim, heads),
+            nn.Linear(dim, heads),
             Rearrange('b n h -> (b h) n')
         )
 
@@ -428,7 +431,7 @@ class NeuralMemory(Module):
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
-            LinearNoBias(dim, heads * self.num_memory_parameter_tensors),
+            nn.Linear(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
@@ -442,10 +445,27 @@ class NeuralMemory(Module):
         # weight decay factor
 
         self.to_decay_factor = Sequential(
-            LinearNoBias(dim, heads),
+            nn.Linear(dim, heads),
             Rearrange('b n h -> (b h) n 1')
         )
 
+        # inits
+
+        if exists(init_adaptive_step_bias):
+            linear = self.to_adaptive_step[0]
+            nn.init.zeros_(linear.weight)
+            nn.init.constant_(linear.bias, init_adaptive_step_bias)
+
+        if exists(init_momentum_bias):
+            linear = self.to_momentum[0]
+            nn.init.zeros_(linear.weight)
+            nn.init.constant_(linear.bias, init_momentum_bias)
+
+        if exists(init_decay_bias):
+            linear = self.to_decay_factor[0]
+            nn.init.zeros_(linear.weight)
+            nn.init.constant_(linear.bias, init_decay_bias)
+
         # maybe use accelerated scan
 
         self.use_accelerated_scan = use_accelerated_scan
@@ -635,15 +655,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)
@@ -698,76 +722,10 @@ 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)
+        if need_pad:
+            values = values[:, 1:]
 
-            return output
-
-        # store if storage sequence cache hits the chunk size
-
-        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,
@@ -775,17 +733,25 @@ class NeuralMemory(Module):
         store_seq = None,
         state: NeuralMemCache | None = 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
@@ -863,9 +829,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.12.dist-info → titans_pytorch-0.2.15.dist-info}/METADATA

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

titans_pytorch-0.2.15.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=Udu-9mtPy9sDeDyXKo95YMel3ELv5quJXINW-JG-hdk,24357
+titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
+titans_pytorch/neural_memory.py,sha256=iu9lnRrqWmtFw3QYyJlS7mOP2zI2HJFuhs3TyfkKV3o,25482
+titans_pytorch-0.2.15.dist-info/METADATA,sha256=vOb0Tt6-egnqtNXMfrJVibHwm8VuWQMlPw3C7Y_L4Wg,6812
+titans_pytorch-0.2.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.2.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.2.15.dist-info/RECORD,,

titans_pytorch-0.2.12.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=RfJ1SvQH5_4PmlB7g-13wPAqYtCCUJxfmtaL0oBrRCU,24563
-titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
-titans_pytorch/neural_memory.py,sha256=PbPjpm7BZIkl1z4Zyt3OcWW6MxWKAnTlmDvJthcmlp4,26151
-titans_pytorch-0.2.12.dist-info/METADATA,sha256=xZvhE0thLfasfF7DHO2FcwNo-bD0T5_R01pvL4sWd4A,6812
-titans_pytorch-0.2.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.12.dist-info/RECORD,,