titans-pytorch 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

titans_pytorch/titans.py

@@ -40,6 +40,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def identity(t):
+    return t
+
 def round_down_multiple(seq, mult):
     return seq // mult * mult
 
@@ -55,6 +58,21 @@ def pack_one_with_inverse(t, pattern):
 
     return packed, inverse
 
+# softclamping gradients
+
+def softclamp_max(t, max_value):
+    half_max_value = max_value / 2
+    return ((t / half_max_value).tanh() * half_max_value) + half_max_value
+
+def softclamp_grad_norm(t, max_value):
+    t, inverse = pack_one_with_inverse(t, 'bn *')
+
+    norm = t.norm(dim = -1, keepdim = True)
+    clamped_norm = softclamp_max(norm, max_value)
+
+    t = t * (clamped_norm / norm)
+    return inverse(t)
+
 # classes
 
 class MemoryMLP(Module):
@@ -82,6 +100,9 @@ class MemoryMLP(Module):
 
 # main neural memory
 
+def default_adaptive_step_transform(adaptive_step):
+    return torch.exp(adaptive_step.sigmoid() * -15) # from 1. - 1e-7
+
 def default_loss_fn(pred, target):
     return (pred - target).pow(2).mean(dim = -1).sum()
 
@@ -94,8 +115,10 @@ 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,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         default_mlp_kwargs: dict = dict(
             depth = 4
@@ -172,6 +195,12 @@ class NeuralMemory(Module):
             Rearrange('b n h -> (b h) n')
         )
 
+        self.adaptive_step_transform = adaptive_step_transform
+
+        # allow for softclamp the gradient norms for storing memories
+
+        self.max_grad_norm = max_grad_norm
+
         # weight decay factor
 
         self.to_decay_factor = nn.Sequential(
@@ -222,7 +251,8 @@ class NeuralMemory(Module):
 
         # pack batch and sequence dimension
 
-        adaptive_lr = (self.to_adaptive_step(seq).sigmoid() * -15).exp() # from 1. - 1e-7
+        adaptive_lr = self.to_adaptive_step(seq)
+        adaptive_lr = self.adaptive_step_transform(adaptive_lr)
 
         adaptive_momentum = self.to_momentum(seq).sigmoid()
         decay_factor = self.to_decay_factor(seq).sigmoid()
@@ -247,6 +277,11 @@ class NeuralMemory(Module):
 
         grads = TensorDict(grads)
 
+        # maybe softclamp grad norm
+
+        if exists(self.max_grad_norm):
+            grads = grads.apply(lambda t: softclamp_grad_norm(t, self.max_grad_norm))
+
         # restore batch and sequence dimension
 
         grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))

{titans_pytorch-0.0.19.dist-info → titans_pytorch-0.0.21.dist-info}/METADATA

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

titans_pytorch-0.0.21.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/titans.py,sha256=rsgpGs-nnqvsvF6Mu2GfmfNlRgYfHAV9MCEmd1ohUUI,13687
+titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
+titans_pytorch-0.0.21.dist-info/METADATA,sha256=3z2LGviHUicQq4_ThlRasOUCDXIe0xX2aIhxtKLWI0Q,3811
+titans_pytorch-0.0.21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.21.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.21.dist-info/RECORD,,

titans_pytorch-0.0.19.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/titans.py,sha256=CxbJgNdIS9NbbCDdgotFXAnrV16xmvufUErerKe7qJA,12636
-titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.19.dist-info/METADATA,sha256=5Wpk79HYI4z8LeNRV__UaamKppiGcJ2HdIlll1JSZr8,3811
-titans_pytorch-0.0.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.19.dist-info/RECORD,,