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

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/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.dist-info → titans_pytorch-0.0.50.dist-info}/METADATA

@@ -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.50.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=EMhxPt86Vr6LFvPm0OLMFYLaIY19khU9yIHkIhl2EMA,10316
+titans_pytorch/titans.py,sha256=TklMAxNDxgFBpJZFJa8hEhqA_DITmT6EM0p0ueE1jo8,15712
+titans_pytorch-0.0.50.dist-info/METADATA,sha256=KU7TTrH89eNVPP10NKKTDKnW-ik344_kVQkAXW7NRL8,4210
+titans_pytorch-0.0.50.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.50.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.50.dist-info/RECORD,,

titans_pytorch-0.0.47.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=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,,