titans-pytorch 0.2.17__py3-none-any.whl → 0.2.18__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -491,7 +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
+        neural_mem_weight_residual = False,
     ):
         super().__init__()
 
@@ -525,7 +525,10 @@ class MemoryAsContextTransformer(Module):
 
         neural_memory_layers = default(neural_memory_layers, layers)
 
+        # weight residual related
+
         self.neural_mem_weight_residual = neural_mem_weight_residual
+        is_first_neural_mem = True
 
         # mem, attn, and feedforward layers
 
@@ -557,9 +560,12 @@ class MemoryAsContextTransformer(Module):
                     chunk_size = self.neural_memory_segment_len,
                     batch_size = neural_memory_batch_size,
                     model = deepcopy(neural_memory_model),
+                    learned_weight_residual = neural_mem_weight_residual and not is_first_neural_mem,
                     **neural_memory_kwargs
                 )
 
+                is_first_neural_mem = False
+
             ff = FeedForward(dim = dim, mult = ff_mult)
 
             self.layers.append(ModuleList([

titans_pytorch/neural_memory.py

@@ -296,6 +296,7 @@ class NeuralMemory(Module):
         init_adaptive_step_bias = None,
         init_momentum_bias = None,
         init_decay_bias = None,
+        learned_weight_residual = False,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -438,6 +439,14 @@ class NeuralMemory(Module):
 
         self.max_mem_layer_modulation = max_mem_layer_modulation
 
+        # learned weight residual
+
+        self.to_learned_weight_residual_mix = Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b n h -> b h n'),
+            nn.Sigmoid()
+        ) if learned_weight_residual else None
+
         # allow for softclamp the gradient norms for storing memories
 
         self.max_grad_norm = max_grad_norm
@@ -563,6 +572,8 @@ class NeuralMemory(Module):
 
         # maybe add previous layer weight
 
+        assert xnor(exists(self.to_learned_weight_residual_mix), exists(prev_weights))
+
         if exists(prev_weights):
 
             start_index = math.ceil(seq_index / chunk_size)
@@ -570,6 +581,12 @@ class NeuralMemory(Module):
 
             prev_weights = prev_weights.apply(lambda t: t[:, start_index:end_index])
 
+            if exists(self.to_learned_weight_residual_mix):
+                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))
+
             weights_for_surprise = weights_for_surprise + prev_weights
 
         # flatten batch and time if surprise depends on previous layer memory model

{titans_pytorch-0.2.17.dist-info → titans_pytorch-0.2.18.dist-info}/METADATA

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

titans_pytorch-0.2.18.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=J9MHFViWWehBOFOKx7ry_X2k8nAXaAZFFeCwGtudZyk,24942
+titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
+titans_pytorch/neural_memory.py,sha256=A74v2g4OUXJIqXLu1Rk7FDsTlhveJIHsiamsWN9e7cI,26602
+titans_pytorch-0.2.18.dist-info/METADATA,sha256=QFJI4OoZZSdgZd6FJInAG5-z2ubpvO5I0aCRNKwhDnQ,6812
+titans_pytorch-0.2.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.2.18.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.2.18.dist-info/RECORD,,

titans_pytorch-0.2.17.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=QhdSPgWntfWILMJ1t0xLKgvZfPZWu9vhzZWaesftaPg,24724
-titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
-titans_pytorch/neural_memory.py,sha256=M7vB6AgrphY3_cdELyWXF8qS9cjYUvv1NVCUqQOlA8M,25928
-titans_pytorch-0.2.17.dist-info/METADATA,sha256=8oBe_7SkPmHvHVGXpoRvGWMTldfl2pzXIztOaM_qGrI,6812
-titans_pytorch-0.2.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.17.dist-info/RECORD,,