titans-pytorch 0.0.46__tar.gz → 0.0.49__tar.gz

{titans_pytorch-0.0.46 → titans_pytorch-0.0.49}/PKG-INFO

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

{titans_pytorch-0.0.46 → titans_pytorch-0.0.49}/pyproject.toml

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

{titans_pytorch-0.0.46 → titans_pytorch-0.0.49}/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.46 → titans_pytorch-0.0.49}/titans_pytorch/titans.py

@@ -217,10 +217,10 @@ class NeuralMemory(Module):
         def forward_and_loss(params, inputs, loss_weights, target):
             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|²
-            loss = loss * loss_weights
-            return loss.sum()
+            weighted_loss = loss * loss_weights
+            return weighted_loss.sum(), weighted_loss.mean()
 
-        self.per_sample_grad_fn = vmap(grad(forward_and_loss), in_dims = (None, 0, 0, 0))
+        self.per_sample_grad_fn = vmap(grad(forward_and_loss, has_aux = True), in_dims = (None, 0, 0, 0))
 
         # queries for retrieving from the model
 
@@ -282,7 +282,8 @@ class NeuralMemory(Module):
     def store_memories(
         self,
         seq,
-        past_state: tuple[dict[str, Tensor], dict[str, Tensor]]
+        past_state: tuple[dict[str, Tensor], dict[str, Tensor]],
+        return_aux_kv_loss = False
     ):
 
         seq = self.store_norm(seq)
@@ -330,7 +331,7 @@ class NeuralMemory(Module):
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads = self.per_sample_grad_fn(dict(curr_weights), keys, adaptive_lr, values)
+        grads, aux_kv_recon_loss = self.per_sample_grad_fn(dict(curr_weights), keys, adaptive_lr, values)
 
         grads = TensorDict(grads)
 
@@ -405,7 +406,10 @@ class NeuralMemory(Module):
 
         next_state = (curr_weights + last_update, next_momentum)
 
-        return updates, next_state
+        if not return_aux_kv_loss:
+            return updates, next_state
+
+        return updates, next_state, aux_kv_recon_loss.mean()
 
     def retrieve_memories(
         self,
@@ -484,7 +488,7 @@ class NeuralMemory(Module):
         seq,
         store_seq = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
-        return_next_memories = False
+        return_aux_kv_loss = False
     ):
         batch, seq_len = seq.shape[:2]
 
@@ -499,13 +503,13 @@ class NeuralMemory(Module):
 
         store_seq = default(store_seq, seq)
 
-        updates, next_memories = self.store_memories(store_seq, past_state)
+        updates, next_memories, aux_kv_recon_loss = self.store_memories(store_seq, past_state, return_aux_kv_loss = True)
 
         past_weights, _ = past_state
 
         retrieved = self.retrieve_memories(seq, past_weights + updates)
 
-        if not return_next_memories:
+        if not return_aux_kv_loss:
             return retrieved
 
-        return retrieved, next_memories
+        return retrieved, aux_kv_recon_loss

{titans_pytorch-0.0.46 → titans_pytorch-0.0.49}/train_mac.py

@@ -24,13 +24,14 @@ 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
-RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS}'
+KV_RECON_LOSS_WEIGHT = 0.1
+RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS} - kv recon loss 0.1'
 
 # 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()