titans-pytorch 0.3.0__tar.gz → 0.3.2__tar.gz

{titans_pytorch-0.3.0 → titans_pytorch-0.3.2}/PKG-INFO

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

{titans_pytorch-0.3.0 → titans_pytorch-0.3.2}/pyproject.toml

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

{titans_pytorch-0.3.0 → titans_pytorch-0.3.2}/tests/test_titans.py

@@ -156,6 +156,7 @@ def test_neural_mem_chaining_with_batch_size():
 @pytest.mark.parametrize('neural_mem_segment_len', (8, 16))
 @pytest.mark.parametrize('neural_mem_weight_residual', (False, True))
 @pytest.mark.parametrize('neural_mem_batch_size', (None, 64))
+@pytest.mark.parametrize('neural_mem_momentum', (False, True))
 def test_mac(
     seq_len,
     num_persist_mem_tokens,
@@ -164,6 +165,7 @@ def test_mac(
     neural_mem_segment_len,
     neural_mem_weight_residual,
     neural_mem_batch_size,
+    neural_mem_momentum
 ):
     transformer = MemoryAsContextTransformer(
         num_tokens = 256,
@@ -175,7 +177,10 @@ def test_mac(
         neural_mem_gate_attn_output = neural_mem_gate_attn_output,
         neural_memory_segment_len = neural_mem_segment_len,
         neural_memory_batch_size = neural_mem_batch_size,
-        neural_mem_weight_residual = neural_mem_weight_residual
+        neural_mem_weight_residual = neural_mem_weight_residual,
+        neural_memory_kwargs = dict(
+            momentum = neural_mem_momentum
+        )
     )
 
     x = torch.randint(0, 256, (1, seq_len))

{titans_pytorch-0.3.0 → titans_pytorch-0.3.2}/titans_pytorch/memory_models.py

@@ -36,10 +36,14 @@ class MemoryMLP(Module):
     def __init__(
         self,
         dim,
-        depth
+        depth,
+        expansion_factor = 4.
     ):
         super().__init__()
-        self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
+        dim_hidden = int(dim * expansion_factor)
+        dims = (dim, *((dim_hidden,) * (depth - 1)), dim)
+
+        self.weights = ParameterList([Parameter(torch.randn(dim_in, dim_out)) for dim_in, dim_out in zip(dims[:-1], dims[1:])])
 
         self.ln = LayerNorm(dim)
 

{titans_pytorch-0.3.0 → titans_pytorch-0.3.2}/titans_pytorch/neural_memory.py

@@ -299,7 +299,8 @@ class NeuralMemory(Module):
         accept_weight_residual = False,
         gated_transition = False,
         default_model_kwargs: dict = dict(
-            depth = 2
+            depth = 2,
+            expansion_factor = 4.
         )
     ):
         super().__init__()
@@ -652,13 +653,14 @@ class NeuralMemory(Module):
         next_last_update = TensorDict()
         next_last_momentum = TensorDict()
 
-        for (param_name, surprise), (_, last_update), (_, last_momentum) in zip(surprises.items(), past_last_update.items(), past_last_momentum.items()):
+        for (param_name, surprise), (_, last_update) in zip(surprises.items(), past_last_update.items()):
 
             update = surprise
 
             # derive momentum with associative scan - eq (10)
 
             if has_momentum:
+                last_momentum = past_last_momentum[param_name]
                 update = self.assoc_scan(adaptive_momentum, surprise, prev = last_momentum) # momentum is S / surprise in the paper
                 momentum = update
                 next_last_momentum[param_name] = momentum[:, -1]