titans-pytorch 0.1.5__tar.gz → 0.1.6__tar.gz

{titans_pytorch-0.1.5 → titans_pytorch-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.5
+Version: 0.1.6
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -46,6 +46,7 @@ Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-transformers
 Provides-Extra: examples
+Requires-Dist: adam-atan2-pytorch>=0.1.18; extra == 'examples'
 Requires-Dist: wandb; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'

{titans_pytorch-0.1.5 → titans_pytorch-0.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.5"
+version = "0.1.6"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -45,6 +45,7 @@ Repository = "https://github.com/lucidrains/titans-pytorch"
 [project.optional-dependencies]
 
 examples = [
+    "adam-atan2-pytorch>=0.1.18",
     "wandb"
 ]
 

{titans_pytorch-0.1.5 → titans_pytorch-0.1.6}/titans_pytorch/__init__.py

@@ -2,7 +2,8 @@ from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
     MemoryAttention,
-    FactorizedMemoryMLP
+    FactorizedMemoryMLP,
+    GatedResidualMemoryMLP
 )
 
 from titans_pytorch.mac_transformer import (

{titans_pytorch-0.1.5 → titans_pytorch-0.1.6}/titans_pytorch/titans.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 from typing import Callable
+
 import math
 from functools import partial
 
 import torch
-from torch import nn, Tensor
+from torch import nn, cat, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList
 from torch.func import functional_call, vmap, grad
@@ -154,6 +155,46 @@ class MemoryMLP(Module):
 
         return x
 
+# memory mlp, but with gated residual + final projection
+
+class GatedResidualMemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth
+    ):
+        super().__init__()
+        self.depth = depth
+
+        self.weights = ParameterList([
+            ParameterList([
+                Parameter(torch.randn(dim, dim)),
+                Parameter(torch.randn(dim * 2, dim)),
+            ]) for _ in range(depth)
+        ])
+
+        self.final_proj = Parameter(torch.randn(dim, dim))
+
+        for param in self.parameters():
+            nn.init.xavier_uniform_(param)
+
+    def forward(
+        self,
+        x
+    ):
+        for weight, to_gates in self.weights:
+            res = x
+
+            x = x @ weight
+            x = F.silu(x)
+
+            # gated residual
+
+            gates = cat((x, res), dim = -1) @ to_gates
+            x = res.lerp(x, gates.sigmoid())
+
+        return x @ self.final_proj
+
 # memory mlp with factorized weights
 # so can tradeoff capacity for smaller chunk sizes
 

{titans_pytorch-0.1.5 → titans_pytorch-0.1.6}/train_mac.py

@@ -5,10 +5,10 @@ import numpy as np
 
 import torch
 from torch import nn, Tensor
-from torch.optim import Adam
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
+from adam_atan2_pytorch import AdoptAtan2
 from titans_pytorch import MemoryAsContextTransformer
 
 # constants
@@ -34,6 +34,7 @@ NEURAL_MEM_GATE_ATTN_OUTPUT = True
 WINDOW_SIZE = 32
 NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2 # set smaller for more granularity for learning rate / momentum etc
 SLIDING_WINDOWS = True
+STORE_ATTN_POOL_CHUNKS = True # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
 KV_RECON_LOSS_WEIGHT = 0.
 LEARNED_MEM_MODEL_WEIGHTS = True
 
@@ -86,6 +87,7 @@ model = MemoryAsContextTransformer(
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,
+        attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
         learned_mem_model_weights = LEARNED_MEM_MODEL_WEIGHTS,
         default_model_kwargs = dict(
@@ -122,11 +124,11 @@ val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE))
 
 # optimizer
 
-optim = Adam(model.parameters(), lr=LEARNING_RATE)
+optim = AdoptAtan2(model.parameters(), lr = LEARNING_RATE)
 
 # training
 
-for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10., desc = 'training'):
     model.train()
 
     for __ in range(GRADIENT_ACCUMULATE_EVERY):