titans-pytorch 0.2.6__tar.gz → 0.2.8__tar.gz

{titans_pytorch-0.2.6 → titans_pytorch-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.2.6
+Version: 0.2.8
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -27,7 +27,7 @@ License: MIT License
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
 License-File: LICENSE
-Keywords: artificial intelligence,deep learning,linear attention,neural memory module,test time training
+Keywords: artificial intelligence,deep learning,linear attention,memory,test time training
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License

{titans_pytorch-0.2.6 → titans_pytorch-0.2.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.2.6"
+version = "0.2.8"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -11,9 +11,9 @@ license = { file = "LICENSE" }
 keywords = [
     'artificial intelligence',
     'deep learning',
-    'neural memory module',
     'test time training',
-    'linear attention'
+    'linear attention',
+    'memory',
 ]
 
 classifiers=[

{titans_pytorch-0.2.6 → titans_pytorch-0.2.8}/tests/test_titans.py

@@ -24,9 +24,9 @@ def torch_default_dtype(dtype):
 
 # main test
 
-@pytest.mark.parametrize('seq_len', (32, 1024, 77))
+@pytest.mark.parametrize('seq_len', (32, 512, 77))
 @pytest.mark.parametrize('silu', (False, True))
-@pytest.mark.parametrize('attn_pool_chunks', (False, True))
+@pytest.mark.parametrize('chunk_size, attn_pool_chunks', ((64, True), (64, False), (1, False)))
 @pytest.mark.parametrize('momentum', (False, True))
 @pytest.mark.parametrize('qk_rmsnorm', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
@@ -35,6 +35,7 @@ def test_titans(
     seq_len,
     silu,
     attn_pool_chunks,
+    chunk_size,
     momentum,
     qk_rmsnorm,
     max_grad_norm,
@@ -42,7 +43,7 @@ def test_titans(
 ):
     mem = NeuralMemory(
         dim = 384,
-        chunk_size = 64,
+        chunk_size = chunk_size,
         activation = nn.SiLU() if silu else None,
         attn_pool_chunks = attn_pool_chunks,
         max_grad_norm = max_grad_norm,

{titans_pytorch-0.2.6 → titans_pytorch-0.2.8}/titans_pytorch/memory_models.py

@@ -3,11 +3,33 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 
+from einops import rearrange
+
 # functions
 
 def l2norm(t):
     return F.normalize(t, dim = -1)
 
+# norms
+
+class LayerNorm(Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+
+        self.ln = nn.LayerNorm(dim, elementwise_affine = False)
+        self.gamma = Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        gamma = self.gamma
+
+        if gamma.ndim == 2:
+            gamma = rearrange(gamma, 'b d -> b 1 d')
+
+        return self.ln(x) * (gamma + 1.)
+
 # memory mlp proposed in TTT
 
 class MemoryMLP(Module):
@@ -19,6 +41,8 @@ class MemoryMLP(Module):
         super().__init__()
         self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
 
+        self.ln = LayerNorm(dim)
+
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 
@@ -36,7 +60,7 @@ class MemoryMLP(Module):
 
             x = x @ weight
 
-        return x + residual
+        return self.ln(x) + residual
 
 # memory mlp, but with gated residual + final projection
 
@@ -60,6 +84,8 @@ class GatedResidualMemoryMLP(Module):
 
         self.final_proj = Parameter(torch.randn(dim, dim))
 
+        self.ln = LayerNorm(dim)
+
         for param in self.parameters():
             nn.init.xavier_uniform_(param)
 
@@ -83,7 +109,7 @@ class GatedResidualMemoryMLP(Module):
 
         out = x @ self.final_proj
 
-        return out + residual
+        return self.ln(out) + residual
 
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
@@ -103,6 +129,8 @@ class FactorizedMemoryMLP(Module):
             ]) for _ in range(depth)
         ])
 
+        self.ln = LayerNorm(dim)
+
         for weight1, weight2 in self.weights:
             nn.init.xavier_uniform_(weight1)
             nn.init.xavier_uniform_(weight2)
@@ -121,7 +149,7 @@ class FactorizedMemoryMLP(Module):
 
             x = x @ weight1 @ weight2
 
-        return x + residual
+        return self.ln(x) + residual
 
 # improvised attention as memory module
 
@@ -144,6 +172,8 @@ class MemoryAttention(Module):
             nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
         ])
 
+        self.ln = LayerNorm(dim)
+
         for weight in self.weights:
             nn.init.xavier_uniform_(weight)
 
@@ -168,4 +198,4 @@ class MemoryAttention(Module):
         h = F.gelu(x @ ffw1)
         ff_out = h @ ffw2
 
-        return attn_out + ff_out + residual
+        return self.ln(attn_out + ff_out) + residual

{titans_pytorch-0.2.6 → titans_pytorch-0.2.8}/titans_pytorch/neural_memory.py

@@ -51,6 +51,9 @@ def default(*args):
             return arg
     return None
 
+def identity(t):
+    return t
+
 def xnor(x, y):
     return not (x ^ y)
 
@@ -64,9 +67,6 @@ def safe_cat(inputs, dim = -2):
 
     return cat(inputs, dim = dim)
 
-def identity(t):
-    return t
-
 def dict_get_shape(td):
     return {k: v.shape for k, v in td.items()}
 
@@ -454,14 +454,14 @@ class NeuralMemory(Module):
 
         weights = TensorDict(weights)
 
-        # allow for neural memory of a previous layer and the past to produce gradients that become the weights of the current one generating the surprise
-        # think this is necessary otherwise the memory model is static (unless if paper is misunderstood)
-        # improvise (or perhaps correcting to) a solution
+        # allow for neural memory of a previous layer to influence surprise of current layer
+
+        weights_for_surprise = weights
 
         if exists(prev_layer_updates):
             prev_layer_updates = TensorDict(prev_layer_updates)
 
-            weights = weights + prev_layer_updates
+            weights_for_surprise = weights_for_surprise + prev_layer_updates
 
             per_sample_grad_fn = self.per_sample_grad_fn_expanded_weights # the weights will now have a batch * chunk dimension
 
@@ -506,11 +506,11 @@ class NeuralMemory(Module):
         # flatten batch and time if surprise depends on previous layer memory model
 
         if exists(prev_layer_updates):
-            weights = weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
+            weights_for_surprise = weights_for_surprise.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads, aux_kv_recon_loss = per_sample_grad_fn(dict(weights), keys, adaptive_lr, values)
+        grads, aux_kv_recon_loss = per_sample_grad_fn(dict(weights_for_surprise), keys, adaptive_lr, values)
 
         grads = TensorDict(grads)
 
@@ -536,7 +536,15 @@ class NeuralMemory(Module):
 
         if not exists(past_state):
             empty_dict = {key: None for key in weights.keys()}
-            past_state = (empty_dict, empty_dict)
+
+            # minibatch_init_weight corresponds to W0 in figure 7 of TTT paper
+
+            minibatch_init_weight = weights
+
+            if dict_get_shape(weights) == self.init_weight_shape:
+                minibatch_init_weight = weights.apply(lambda t: repeat(t, '... -> b 1 (...)', b = batch * heads))
+
+            past_state = (minibatch_init_weight, empty_dict)
 
         past_last_update, past_last_momentum = past_state
 
@@ -734,7 +742,7 @@ class NeuralMemory(Module):
 
         # retrieve
 
-        retrieved = self.retrieve_memories(token, updates + weights, chunk_size = 1)
+        retrieved = self.retrieve_memories(token, weights, chunk_size = 1)
 
         # next state tuple
 
@@ -796,7 +804,7 @@ class NeuralMemory(Module):
 
         retrieved = self.retrieve_memories(
             seq,
-            mem_model_weights + updates,
+            mem_model_weights,
             chunk_size = chunk_size,
             prev_layer_updates = prev_layer_updates
         )