titans-pytorch 0.0.46__py3-none-any.whl → 0.0.47__py3-none-any.whl

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(), 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
 
     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.dist-info → titans_pytorch-0.0.47.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.46
+Version: 0.0.47
 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.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=mF8PYAjeAjLas1gkYybgzZX1AVK82A_ps_LY00ofYYs,9565
+titans_pytorch/titans.py,sha256=ZcWxx6n-f8ttojRnK9fExavmT1bS-QSCRHQn7ldv7J0,15502
+titans_pytorch-0.0.47.dist-info/METADATA,sha256=HjZxbJlnqsSgbioQz6KHWJb--8n18WDdL2T-jz-CFKc,4210
+titans_pytorch-0.0.47.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.47.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.47.dist-info/RECORD,,

titans_pytorch-0.0.46.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=mF8PYAjeAjLas1gkYybgzZX1AVK82A_ps_LY00ofYYs,9565
-titans_pytorch/titans.py,sha256=qxQ8pZCz8GEDhKeJMEaeAEzH66GAGVBNaRdNam_-czg,15260
-titans_pytorch-0.0.46.dist-info/METADATA,sha256=Gg1-_Mmp9u_sJYEvaRt5GzKhhJoTNHjBL3efjSSDLL0,4210
-titans_pytorch-0.0.46.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.46.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.46.dist-info/RECORD,,