titans-pytorch 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

titans_pytorch/neural_memory.py

@@ -303,6 +303,7 @@ class NeuralMemory(Module):
         init_momentum_bias = None,
         init_decay_bias = None,
         accept_weight_residual = False,
+        gated_transition = False,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -464,6 +465,11 @@ class NeuralMemory(Module):
             Rearrange('b n h -> (b h) n 1')
         )
 
+        # learned transition, as seeing instability when decreasing neural mem batch size
+        # perhaps it can slowly learn to adjust from early residual to fully transitioning to new weights every batch size
+
+        self.transition_gate = nn.Parameter(tensor(-5.)) if gated_transition else None
+
         # inits
 
         if exists(init_adaptive_step_bias):
@@ -587,7 +593,7 @@ class NeuralMemory(Module):
 
             prev_weights = prev_weights.apply(lambda t: t[:, start_index:end_index])
 
-            if exists(self.to_learned_weight_residual_mix):
+            if exists(self.to_learned_weight_residual_mix) and num_chunks > 0:
                 mix = self.to_learned_weight_residual_mix(chunked_seq)
                 mix = rearrange(mix, 'b h n -> (b h) n')
                 prev_weights = prev_weights.apply(lambda t: einx.multiply('bh n, bh n ... -> bh n ...', mix, t))
@@ -832,6 +838,13 @@ class NeuralMemory(Module):
 
         store_seqs = store_seq.split(split_sizes, dim = -2)
 
+        # whether to allow network to slowly adjust from initial weight throughout (residual path) to fully updating weights every batch
+
+        gate = None
+
+        if exists(self.transition_gate):
+            gate = self.transition_gate.sigmoid()
+
         for ind, store_seq_chunk in enumerate(store_seqs):
             is_last = ind == (len(store_seqs) - 1)
 
@@ -857,7 +870,10 @@ class NeuralMemory(Module):
 
             last_update, _ = past_state
 
-            weights = last_update
+            if exists(gate):
+                weights = TensorDict({param_name: v1.lerp(v2, gate) for (param_name, v1), (_, v2) in zip(weights.items(), last_update.items())})
+            else:
+                weights = last_update
 
             next_neural_mem_state = tuple_index_set(next_neural_mem_state, 1, last_update)
 

{titans_pytorch-0.2.20.dist-info → titans_pytorch-0.2.22.dist-info}/METADATA

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

{titans_pytorch-0.2.20.dist-info → titans_pytorch-0.2.22.dist-info}/RECORD

@@ -2,8 +2,8 @@ titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,27
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
 titans_pytorch/mac_transformer.py,sha256=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
 titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
-titans_pytorch/neural_memory.py,sha256=e1TENsK0IptpapFRtQux-Uii2MFvisoiQLzTQboHc50,26709
-titans_pytorch-0.2.20.dist-info/METADATA,sha256=8v5puM3CwbPaYDt2aZaqWruPj8whCt_moiQUVE1vk5U,6816
-titans_pytorch-0.2.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.20.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.20.dist-info/RECORD,,
+titans_pytorch/neural_memory.py,sha256=1xKafkiJa5HxNGeD4Aw86Et_o-tePViBsiyKw-l7LR0,27511
+titans_pytorch-0.2.22.dist-info/METADATA,sha256=G-iinXLii0iXtbc8-sPDZqTNBukHX4fy66CJ9l-h7hU,6816
+titans_pytorch-0.2.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.2.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.2.22.dist-info/RECORD,,