titans-pytorch 0.0.53__tar.gz → 0.0.54__tar.gz

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/PKG-INFO

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

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/pyproject.toml

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

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/tests/test_titans.py

@@ -6,17 +6,20 @@ from titans_pytorch import NeuralMemory
 
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
+@pytest.mark.parametrize('learned_mem_model_weights', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 def test_titans(
     seq_len,
     silu,
+    learned_mem_model_weights,
     max_grad_norm,
 ):
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
         activation = nn.SiLU() if silu else None,
-        max_grad_norm = max_grad_norm
+        max_grad_norm = max_grad_norm,
+        learned_mem_model_weights = learned_mem_model_weights
     )
 
     seq = torch.randn(2, seq_len, 384)

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/titans_pytorch/__init__.py

@@ -1,6 +1,7 @@
 from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
+    MemoryAttention
 )
 
 from titans_pytorch.mac_transformer import (

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/titans_pytorch/titans.py

@@ -123,9 +123,12 @@ class MemoryMLP(Module):
 class MemoryAttention(Module):
     def __init__(
         self,
-        dim
+        dim,
+        scale = 8.
     ):
         super().__init__()
+        self.scale = scale
+
         self.weights = nn.ParameterList([
             nn.Parameter(torch.randn(dim, dim)), # queries
             nn.Parameter(torch.randn(dim, dim)), # keys
@@ -143,6 +146,7 @@ class MemoryAttention(Module):
 
         attn_out = F.scaled_dot_product_attention(
             q, k, v,
+            scale = self.scale,
             is_causal = True
         )
 
@@ -174,6 +178,7 @@ class NeuralMemory(Module):
         default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        learned_mem_model_weights = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         activation: Module | None = None,
@@ -212,6 +217,9 @@ class NeuralMemory(Module):
         if not exists(model):
             model = MemoryMLP(dim_head, **default_model_kwargs)
 
+        if not learned_mem_model_weights:
+            model.requires_grad_(False)
+
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
         # the memory is the weights of the model

{titans_pytorch-0.0.53 → titans_pytorch-0.0.54}/train_mac.py

@@ -9,7 +9,7 @@ from torch.optim import Adam
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
-from titans_pytorch.mac_transformer import MemoryAsContextTransformer
+from titans_pytorch import MemoryAsContextTransformer
 
 # constants
 
@@ -25,13 +25,14 @@ SHOULD_GENERATE = True
 SEQ_LEN = 512
 
 PROJECT_NAME = 'titans-mac-transformer'
-WANDB_ONLINE = False # turn this on to pipe experiment to cloud
+WANDB_ONLINE = True # turn this on to pipe experiment to cloud
 NEURAL_MEMORY_DEPTH = 2
 NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
 NEURAL_MEM_LAYERS = (2, 4)
 WINDOW_SIZE = 32
 KV_RECON_LOSS_WEIGHT = 0.
+LEARNED_MEM_MODEL_WEIGHTS = True
 RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS}'
 
 # wandb experiment tracker
@@ -115,6 +116,7 @@ model = MemoryAsContextTransformer(
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,
+        learned_mem_model_weights = LEARNED_MEM_MODEL_WEIGHTS,
         default_model_kwargs = dict(
             depth = NEURAL_MEMORY_DEPTH,
         )