titans-pytorch 0.1.12__tar.gz → 0.1.15__tar.gz

{titans_pytorch-0.1.12 → titans_pytorch-0.1.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.12
+Version: 0.1.15
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -35,7 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
-Requires-Dist: axial-positional-embedding>=0.3.7
+Requires-Dist: axial-positional-embedding>=0.3.9
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8

{titans_pytorch-0.1.12 → titans_pytorch-0.1.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.12"
+version = "0.1.15"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -26,7 +26,7 @@ classifiers=[
 
 dependencies = [
     "accelerated-scan>=0.2.0",
-    "axial_positional_embedding>=0.3.7",
+    "axial_positional_embedding>=0.3.9",
     "einops>=0.8.0",
     "einx>=0.3.0",
     "hyper-connections>=0.1.8",

{titans_pytorch-0.1.12 → titans_pytorch-0.1.15}/tests/test_titans.py

@@ -12,6 +12,7 @@ def exists(v):
 @pytest.mark.parametrize('silu', (False, True))
 @pytest.mark.parametrize('learned_mem_model_weights', (False, True))
 @pytest.mark.parametrize('attn_pool_chunks', (False, True))
+@pytest.mark.parametrize('momentum', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 @pytest.mark.parametrize('per_parameter_lr_modulation', (False, True))
 def test_titans(
@@ -19,6 +20,7 @@ def test_titans(
     silu,
     learned_mem_model_weights,
     attn_pool_chunks,
+    momentum,
     max_grad_norm,
     per_parameter_lr_modulation
 ):
@@ -28,6 +30,7 @@ def test_titans(
         activation = nn.SiLU() if silu else None,
         attn_pool_chunks = attn_pool_chunks,
         max_grad_norm = max_grad_norm,
+        momentum = momentum,
         per_parameter_lr_modulation = per_parameter_lr_modulation,
         learned_mem_model_weights = learned_mem_model_weights
     )

{titans_pytorch-0.1.12 → titans_pytorch-0.1.15}/titans_pytorch/titans.py

@@ -301,6 +301,45 @@ class MemoryAttention(Module):
 
         return out
 
+# associative scan wrapper
+
+class AssocScan(Module):
+    def __init__(
+        self,
+        use_accelerated = False
+    ):
+        super().__init__()
+        self.use_accelerated = use_accelerated
+
+    def forward(self, gates, inputs):
+
+        if not self.use_accelerated:
+            _, outputs = associative_scan(binary_operator, (gates, inputs))
+            return outputs
+
+        from accelerated_scan.triton import scan as triton_scan
+        from accelerated_scan.warp import scan as warp_scan
+
+        scan = triton_scan if gates.is_cuda else warp_scan
+
+        def accelerate_scan_fn(gates, inputs):
+            gates = gates.expand_as(inputs)
+            gates, inputs = tuple(rearrange(t, 'b n d -> b d n') for t in (gates, inputs))
+
+            seq_len = gates.shape[-1]
+            next_power_two_seq_len = 2 ** max(5, int(math.ceil(math.log2(seq_len))))
+
+            gates = F.pad(gates, (0, next_power_two_seq_len - seq_len))
+            inputs = F.pad(inputs, (0, next_power_two_seq_len - seq_len))
+
+            outputs = scan(gates.contiguous(), inputs.contiguous())
+
+            outputs = outputs[..., :seq_len]
+            outputs = rearrange(outputs, 'b d n -> b n d')
+            return outputs
+
+        return accelerate_scan_fn(gates, inputs)
+
 # main neural memory
 
 def default_adaptive_step_transform(adaptive_step, max_lr = 1e-2):
@@ -323,6 +362,7 @@ class NeuralMemory(Module):
         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.
         attn_pool_chunks = False,
+        momentum = True,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         learned_mem_model_weights = True,
@@ -338,6 +378,10 @@ class NeuralMemory(Module):
 
         self.retrieve_chunk_size, self.store_chunk_size = pair(chunk_size)
 
+        # associative scan
+
+        self.assoc_scan = AssocScan(use_accelerated = use_accelerated_scan)
+
         # norms
 
         self.retrieve_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
@@ -423,7 +467,7 @@ class NeuralMemory(Module):
         self.to_momentum = Sequential(
             LinearNoBias(dim, heads),
             Rearrange('b n h -> (b h) n 1')
-        )
+        ) if momentum else None
 
         self.to_adaptive_step = Sequential(
             LinearNoBias(dim, heads),
@@ -462,12 +506,15 @@ class NeuralMemory(Module):
 
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
-    def init_weights_and_momentum(self):
+    def init_weights_and_momentum(self, zero_weights = False):
         params = TensorDict(dict(self.memory_model.named_parameters()))
 
-        init_weights = params.clone().zero_()
+        init_weights = params
         init_momentum = params.clone().zero_()
 
+        if zero_weights:
+            init_weights = params.clone().zero_()
+
         return init_weights, init_momentum
 
     def init_empty_memory_embed(self, batch, seq_len):
@@ -497,14 +544,10 @@ class NeuralMemory(Module):
 
         seq = seq[:, :round_down_seq_len]
 
-        # curr weights + past weights, in the case that the initial weights are learned
-
-        curr_weights = TensorDict(dict(self.memory_model.named_parameters()))
+        # get the weights of the memory network
 
         past_state = tuple(TensorDict(d) for d in past_state)
-        past_weights, past_momentum = past_state
-
-        curr_weights = curr_weights + past_weights
+        curr_weights, past_momentum = past_state
 
         # derive learned hparams for optimization of memory network
 
@@ -513,10 +556,13 @@ class NeuralMemory(Module):
 
         chunked_seq = self.reduce_to_chunk_rep(seq, chunk_size = chunk_size)
 
-        adaptive_momentum = self.to_momentum(chunked_seq).sigmoid()
         decay_factor = self.to_decay_factor(chunked_seq).sigmoid()
 
         need_layer_lr_mod = exists(self.to_layer_modulation)
+        has_momentum = exists(self.to_momentum)
+
+        if has_momentum:
+            adaptive_momentum = self.to_momentum(chunked_seq).sigmoid()
 
         if need_layer_lr_mod:
             layer_lr_mod = self.to_layer_modulation(chunked_seq) * self.max_mem_layer_modulation
@@ -561,57 +607,31 @@ class NeuralMemory(Module):
 
         surprises = grads.apply(lambda t: -t)
 
-        # determine scan function
-
-        def default_associative_scan(gates, inputs):
-            _, outputs = associative_scan(binary_operator, (gates, inputs))
-            return outputs
-
-        if self.use_accelerated_scan:
-            from accelerated_scan.triton import scan as triton_scan
-            from accelerated_scan.warp import scan as warp_scan
-
-            scan = triton_scan if seq.is_cuda else warp_scan
-
-            def accelerate_scan_fn(gates, inputs):
-                gates = gates.expand_as(inputs)
-                gates, inputs = tuple(rearrange(t, 'b n d -> b d n') for t in (gates, inputs))
-
-                seq_len = gates.shape[-1]
-                next_power_two_seq_len = 2 ** max(5, int(math.ceil(math.log2(seq_len))))
-
-                gates = F.pad(gates, (0, next_power_two_seq_len - seq_len))
-                inputs = F.pad(inputs, (0, next_power_two_seq_len - seq_len))
-
-                outputs = scan(gates.contiguous(), inputs.contiguous())
-
-                outputs = outputs[..., :seq_len]
-                outputs = rearrange(outputs, 'b d n -> b n d')
-                return outputs
-
-            scan_fn = accelerate_scan_fn
-        else:
-            scan_fn = default_associative_scan
-
         # momentum + weight decay - momentum is the new contribution, as most linear RNNs have learned forgetting gates
 
-        next_momentum = TensorDict()
+        next_momentum = TensorDict() if has_momentum else None
         updates = TensorDict()
 
         for param_name, surprise in surprises.items():
 
             surprise, inverse_pack = pack_one_with_inverse(surprise, 'b n *')
 
+            update = surprise
+
             # derive momentum with associative scan - eq (10)
 
-            momentum = scan_fn(adaptive_momentum, surprise) # momentum is S / surprise in the paper
+            if has_momentum:
+                update = self.assoc_scan(adaptive_momentum, surprise) # momentum is S / surprise in the paper
+                momentum = update
 
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
-            update = scan_fn(1. - decay_factor, momentum)
+            update = self.assoc_scan(1. - decay_factor, update)
 
             updates[param_name] = inverse_pack(update)
-            next_momentum[param_name] = inverse_pack(momentum)
+
+            if has_momentum:
+                next_momentum[param_name] = inverse_pack(momentum)
 
         # compute the next weight per batch
 

{titans_pytorch-0.1.12 → titans_pytorch-0.1.15}/train_mac.py

@@ -31,6 +31,7 @@ NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
 NEURAL_MEM_LAYERS = (2, 4)
 NEURAL_MEM_GATE_ATTN_OUTPUT = True
+NEURAL_MEM_MOMENTUM = True
 WINDOW_SIZE = 32
 NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2 # set smaller for more granularity for learning rate / momentum etc
 SLIDING_WINDOWS = True
@@ -88,6 +89,7 @@ model = MemoryAsContextTransformer(
         dim_head = 64,
         heads = 4,
         attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
+        momentum = NEURAL_MEM_MOMENTUM,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
         learned_mem_model_weights = LEARNED_MEM_MODEL_WEIGHTS,
         default_model_kwargs = dict(