titans-pytorch 0.0.47__tar.gz → 0.0.50__tar.gz

{titans_pytorch-0.0.47 → titans_pytorch-0.0.50}/PKG-INFO

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

{titans_pytorch-0.0.47 → titans_pytorch-0.0.50}/pyproject.toml

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

{titans_pytorch-0.0.47 → titans_pytorch-0.0.50}/titans_pytorch/mac_transformer.py

@@ -95,7 +95,7 @@ class SegmentedAttention(Module):
         dim_head = 64,
         heads = 8,
         accept_value_residual = False,
-        attend_kwargs: dict = dict()
+        attend_kwargs: dict = dict(),
     ):
         super().__init__()
         self.norm = nn.RMSNorm(dim)
@@ -201,6 +201,7 @@ class MemoryAsContextTransformer(Module):
         num_residual_streams = 4,
         neural_memory_kwargs: dict = dict(),
         neural_memory_layers: tuple[int, ...] | None = None,
+        aux_kv_recon_loss_weight = 0.
     ):
         super().__init__()
 
@@ -276,10 +277,18 @@ class MemoryAsContextTransformer(Module):
 
         self.to_logits = LinearNoBias(dim, num_tokens)
 
+        # auxiliary loss on kv recon
+
+        self.has_aux_kv_recon_loss = aux_kv_recon_loss_weight > 0.
+        self.aux_kv_recon_loss_weight = aux_kv_recon_loss_weight
+
+        self.register_buffer('zero', torch.tensor(0.), persistent = False)
+
     def forward(
         self,
         x,
-        return_loss = False
+        return_loss = False,
+        return_loss_breakdown = False
     ):
 
         if return_loss:
@@ -317,6 +326,10 @@ class MemoryAsContextTransformer(Module):
 
         value_residual = None
 
+        # aux losses
+
+        kv_recon_losses = self.zero
+
         # expand and reduce streams for hyper connections
 
         x = self.expand_streams(x)
@@ -324,7 +337,8 @@ class MemoryAsContextTransformer(Module):
         for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
 
             if exists(maybe_neural_mem):
-                x = maybe_neural_mem(x)
+                x, aux_kv_loss = maybe_neural_mem(x, return_aux_kv_loss = True)
+                kv_recon_losses = kv_recon_losses + aux_kv_loss
 
             x, values = attn(x, value_residual = value_residual)
 
@@ -351,4 +365,14 @@ class MemoryAsContextTransformer(Module):
         if not return_loss:
             return logits
 
-        return F.cross_entropy(rearrange(logits, 'b n l -> b l n'), labels)
+        ar_loss = F.cross_entropy(rearrange(logits, 'b n l -> b l n'), labels)
+
+        losses = ar_loss
+
+        if self.has_aux_kv_recon_loss:
+            losses = losses + kv_recon_losses * self.aux_kv_recon_loss_weight
+
+        if not return_loss_breakdown:
+            return losses
+
+        return losses, (ar_loss, kv_recon_losses)

{titans_pytorch-0.0.47 → titans_pytorch-0.0.50}/titans_pytorch/titans.py

@@ -162,7 +162,8 @@ class NeuralMemory(Module):
         heads = 1,
         model: Module | None = None,
         store_memory_loss_fn: Callable = default_loss_fn,
-        adaptive_step_transform: Callable = default_adaptive_step_transform,
+        adaptive_step_transform: Callable | None = None,
+        default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         max_grad_norm: float | None = None,
@@ -218,7 +219,7 @@ class NeuralMemory(Module):
             pred = functional_call(self.memory_model, params, inputs)
             loss = self.store_memory_loss_fn(pred, target) # simple mse loss in paper - eq (12) - |M(k) - v|²
             weighted_loss = loss * loss_weights
-            return weighted_loss.sum(), loss.mean()
+            return weighted_loss.sum(), weighted_loss.mean()
 
         self.per_sample_grad_fn = vmap(grad(forward_and_loss, has_aux = True), in_dims = (None, 0, 0, 0))
 
@@ -250,6 +251,9 @@ class NeuralMemory(Module):
             Rearrange('b n h -> (b h) n')
         )
 
+        if not exists(adaptive_step_transform):
+            adaptive_step_transform = partial(default_adaptive_step_transform, max_lr = default_step_transform_max_lr)
+
         self.adaptive_step_transform = adaptive_step_transform
 
         # allow for softclamp the gradient norms for storing memories
@@ -409,7 +413,7 @@ class NeuralMemory(Module):
         if not return_aux_kv_loss:
             return updates, next_state
 
-        return updates, next_state, aux_kv_recon_loss
+        return updates, next_state, aux_kv_recon_loss.mean()
 
     def retrieve_memories(
         self,

{titans_pytorch-0.0.47 → titans_pytorch-0.0.50}/train_mac.py

@@ -24,12 +24,13 @@ SHOULD_GENERATE = False
 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.
 RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS}'
 
 # wandb experiment tracker
@@ -63,6 +64,7 @@ model = MemoryAsContextTransformer(
     num_longterm_mem_tokens = NUM_LONGTERM_MEM,
     neural_memory_layers = NEURAL_MEM_LAYERS,
     neural_memory_segment_len = WINDOW_SIZE // 2,
+    aux_kv_recon_loss_weight = KV_RECON_LOSS_WEIGHT,
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,
@@ -108,20 +110,20 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
     model.train()
 
     for __ in range(GRADIENT_ACCUMULATE_EVERY):
-        loss = model(next(train_loader), return_loss = True)
+        loss, (ar_loss, kv_recon_losses) = model(next(train_loader), return_loss = True, return_loss_breakdown = True)
         loss.backward()
 
-    print(f'training loss: {loss.item()}')
+    print(f'training loss: {ar_loss.item()}')
     torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
     optim.step()
     optim.zero_grad()
-    wandb.log(dict(loss = loss.item()))
+    wandb.log(dict(loss = ar_loss.item()))
 
     if i % VALIDATE_EVERY == 0:
         model.eval()
         with torch.no_grad():
-            loss = model(next(val_loader), return_loss = True)
-            print(f'validation loss: {loss.item()}')
+            loss, (ar_loss, _) = model(next(val_loader), return_loss = True, return_loss_breakdown = True)
+            print(f'validation loss: {ar_loss.item()}')
 
     if SHOULD_GENERATE and i % GENERATE_EVERY == 0:
         model.eval()