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

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

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

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

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

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

@@ -284,15 +284,18 @@ 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,
-        post_rmsnorm = True,
+        post_rmsnorm = False,
         qk_rmsnorm = False,
         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

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

@@ -34,6 +34,7 @@ NEURAL_MEM_LAYERS = (2, 4, 6)                   # layers 2, 4, 6 have neural mem
 NEURAL_MEM_GATE_ATTN_OUTPUT = False
 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_BATCH_SIZE = 128                     # set smaller to update the neural memory weights more often as it traverses the sequence
@@ -86,6 +87,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,
     use_flex_attn = USE_FLEX_ATTN,
     sliding_window_attn = SLIDING_WINDOWS,
     neural_memory_model = MemoryMLP(