titans-pytorch 0.3.11__tar.gz → 0.3.12__tar.gz

{titans_pytorch-0.3.11 → titans_pytorch-0.3.12}/PKG-INFO

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

{titans_pytorch-0.3.11 → titans_pytorch-0.3.12}/pyproject.toml

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

{titans_pytorch-0.3.11 → titans_pytorch-0.3.12}/tests/test_titans.py

@@ -61,8 +61,10 @@ def test_titans(
     assert seq.shape == retrieved.shape
 
 @pytest.mark.parametrize('learned_momentum_combine', (False, True))
+@pytest.mark.parametrize('learned_combine_include_zeroth', (False, True))
 def test_titans_second_order_momentum(
-    learned_momentum_combine
+    learned_momentum_combine,
+    learned_combine_include_zeroth
 ):
 
     mem  = NeuralMemory(
@@ -72,7 +74,8 @@ def test_titans_second_order_momentum(
         chunk_size = 1,
         batch_size = 2,
         momentum_order = 2,
-        learned_momentum_combine = learned_momentum_combine
+        learned_momentum_combine = learned_momentum_combine,
+        learned_combine_include_zeroth = learned_combine_include_zeroth
     )
 
     seq = torch.randn(2, 5, 384)

{titans_pytorch-0.3.11 → titans_pytorch-0.3.12}/titans_pytorch/neural_memory.py

@@ -7,7 +7,7 @@ from itertools import zip_longest
 from collections import namedtuple
 
 import torch
-from torch import nn, cat, tensor, Tensor
+from torch import nn, stack, cat, tensor, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList, ParameterDict
 from torch.func import functional_call, vmap, grad
@@ -230,6 +230,7 @@ class NeuralMemory(Module):
         momentum = True,
         momentum_order = 1,
         learned_momentum_combine = False,
+        learned_combine_include_zeroth = False,
         pre_rmsnorm = True,
         post_rmsnorm = False,
         qk_rmsnorm = False,
@@ -399,12 +400,17 @@ class NeuralMemory(Module):
             assert momentum
             assert momentum_order > 1, 'only second order momentum allowed for now, but may allow learned combination of zeroth'
 
+            if learned_combine_include_zeroth:
+                momentum_order += 1
+
             self.to_learned_momentum_combine = Sequential(
                 nn.Linear(dim, heads * momentum_order),
                 nn.Softmax(dim = -1),
                 Rearrange('b n (h o) -> o (b h) n', h = heads)
             )
 
+            self.learned_combine_include_zeroth = learned_combine_include_zeroth
+
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
@@ -662,10 +668,14 @@ class NeuralMemory(Module):
 
                     momentums.append(momentum)
 
-                momentums = torch.stack(momentums)
+                momentums = stack(momentums)
 
                 next_last_momentum[param_name] = momentums[:, :, -1] # momentums shape is Float['o bh n 1']
 
+                if learned_combine and self.learned_combine_include_zeroth:
+                    # add the original surprise if learned combination of momentums
+                    momentums = cat((rearrange(surprise, '... -> 1 ...'), momentums), dim = 0)
+
                 if not learned_combine:
                     update = momentums[-1]
                 else: