titans-pytorch 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

titans_pytorch/neural_memory.py

@@ -15,7 +15,7 @@ from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 
 from tensordict import TensorDict
 
-from titans_pytorch.associative_scan import AssocScan
+from assoc_scan import AssocScan
 
 from titans_pytorch.memory_models import(
     MemoryMLP,
@@ -152,6 +152,30 @@ def softclamp_grad_norm(t, max_value):
     t = t * (clamped_norm / norm)
     return inverse(t)
 
+# spectral norming the surprise update w/ newton schulz matrix iter
+# Keller Jordan et al. from OSS w/ nanogpt, now being used for two works, Atlas and 'TTT done right'
+
+def newtonschulz5(
+    t,
+    steps = 5,
+    eps = 1e-7,
+    coefs = (3.4445, -4.7750, 2.0315)
+):
+    if t.ndim <= 3:
+        return t
+
+    t, inv_pack = pack_one_with_inverse(t, '* i j')
+    t = t / t.norm(dim = (-1, -2), keepdim = True).clamp(min = eps)
+
+    a, b, c = coefs
+
+    for _ in range(steps):
+        A = t @ t.transpose(-1, -2)
+        B = b * A + c * A @ A
+        t = a * t + B @ t
+
+    return inv_pack(t)
+
 # multi head rmsnorm
 
 class MultiheadRMSNorm(Module):
@@ -254,6 +278,7 @@ class NeuralMemory(Module):
         init_momentum_bias = None,
         init_decay_bias = None,
         accept_weight_residual = False,
+        spectral_norm_surprises = False,
         gated_transition = False,
         mem_model_norm_add_residual = True, # by default, layernorm output and add residual as proposed in TTT paper, but could be removed
         default_model_kwargs: dict = dict(
@@ -465,6 +490,10 @@ class NeuralMemory(Module):
 
         self.max_grad_norm = max_grad_norm
 
+        # spectral norming the surprises before update, a la Muon from Jordan et al.
+
+        self.spectral_norm_surprises = spectral_norm_surprises
+
         # weight decay factor
 
         self.to_decay_factor = Sequential(
@@ -748,6 +777,11 @@ class NeuralMemory(Module):
                 else:
                     update = einsum(combine_momentums, momentums, 'o b n, o b n ... -> b n ...')
 
+            # maybe spectral norm surprises
+
+            if self.spectral_norm_surprises:
+                update = newtonschulz5(update)
+
             # use associative scan again for learned forgetting (weight decay) - eq (13)
 
             update = self.assoc_scan(1. - decay_factor, update, prev = last_update, remove_prev = False)

{titans_pytorch-0.4.6.dist-info → titans_pytorch-0.4.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.4.6
+Version: 0.4.8
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -34,7 +34,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
-Requires-Dist: accelerated-scan>=0.2.0
+Requires-Dist: assoc-scan
 Requires-Dist: axial-positional-embedding>=0.3.10
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
@@ -207,3 +207,31 @@ $ python train_mac.py
     url     = {https://arxiv.org/abs/2501.12352},
 }
 ```
+
+```bibtex
+@misc{jordan2024muon,
+    author  = {Keller Jordan and Yuchen Jin and Vlado Boza and Jiacheng You and
+                    Franz Cesista and Laker Newhouse and Jeremy Bernstein},
+    title   = {Muon: An optimizer for hidden layers in neural networks},
+    year    = {2024},
+    url     = {https://kellerjordan.github.io/posts/muon/}
+}
+```
+
+```bibtex
+@inproceedings{Zhang2025TestTimeTD,
+    title   = {Test-Time Training Done Right},
+    author  = {Tianyuan Zhang and Sai Bi and Yicong Hong and Kai Zhang and Fujun Luan and Songlin Yang and Kalyan Sunkavalli and William T. Freeman and Hao Tan},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:279071244}
+}
+```
+
+```bibtex
+@inproceedings{Behrouz2025ATLASLT,
+    title  = {ATLAS: Learning to Optimally Memorize the Context at Test Time},
+    author = {Ali Behrouz and Ze-Minghui Li and Praneeth Kacham and Majid Daliri and Yuan Deng and Peilin Zhong and Meisam Razaviyayn and Vahab S. Mirrokni},
+    year   = {2025},
+    url    = {https://api.semanticscholar.org/CorpusID:278996373}
+}
+```

titans_pytorch-0.4.8.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=sVTOuRUkaIYabFExdLY6s1qXm1UwHHz_J19H8ZV-X74,338
+titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=ypWXN8koY8pXt7IvlcMR1QM7cYJnWK_iYLEHy2pjx88,34277
+titans_pytorch-0.4.8.dist-info/METADATA,sha256=BbhF0oiPGdcgxrBGJziZvbvXUmS5lVlGVxpUGPmP0O8,7873
+titans_pytorch-0.4.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.4.8.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.4.8.dist-info/RECORD,,

titans_pytorch/associative_scan.py

@@ -1,178 +0,0 @@
-from __future__ import annotations
-import math
-from typing import Callable
-
-import torch
-from torch import Tensor
-from torch.nn import Module
-import torch.nn.functional as F
-
-from einops import rearrange, repeat, reduce, pack, unpack
-
-# taken from S5-pytorch repository
-# https://github.com/i404788/s5-pytorch/blob/74e2fdae00b915a62c914bf3615c0b8a4279eb84/s5/jax_compat.py#L51-L134
-
-# helper functions
-
-def exists(v):
-    return v is not None
-
-def default(*args):
-    for arg in args:
-        if exists(arg):
-            return arg
-    return None
-
-def pad_at_dim(t, pad, dim = -1, value = 0.):
-    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
-    zeros = ((0, 0) * dims_from_right)
-    return F.pad(t, (*zeros, *pad), value = value)
-
-def pack_one_with_inverse(t, pattern):
-    packed, packed_shape = pack([t], pattern)
-
-    def inverse(out, inv_pattern = None):
-        inv_pattern = default(inv_pattern, pattern)
-        return unpack(out, packed_shape, inv_pattern)[0]
-
-    return packed, inverse
-
-# the operator that is needed
-
-@torch.jit.script
-def binary_operator(
-    a: tuple[Tensor, Tensor],
-    b: tuple[Tensor, Tensor]
-):
-    a_i, kv_i = a
-    a_j, kv_j = b
-    return a_j * a_i, torch.addcmul(kv_j, a_j, kv_i)
-
-# Pytorch impl. of jax.lax.associative_scan
-# made specifically for axis of 1 (sequence of tokens for autoregressive modeling)
-
-def associative_scan(
-    operator: Callable,
-    elems: tuple[Tensor, Tensor]
-):
-    num_elems = int(elems[0].shape[1])
-
-    if not all(int(elem.shape[1]) == num_elems for elem in elems[1:]):
-        raise ValueError('Array inputs to associative_scan must have the same '
-                         'first dimension. (saw: {})'
-                         .format([elem.shape for elem in elems]))
-
-    def _scan(elems):
-        """Perform scan on `elems`."""
-        num_elems = elems[0].shape[1]
-
-        if num_elems < 2:
-            return elems
-
-        # Combine adjacent pairs of elements.
-
-        reduced_elems = operator(
-          [elem[:, :-1:2] for elem in elems],
-          [elem[:, 1::2] for elem in elems])
-
-        # Recursively compute scan for partially reduced tensors.
-
-        odd_elems = _scan(reduced_elems)
-
-        if num_elems % 2 == 0:
-            even_elems = operator(
-                [e[:, :-1] for e in odd_elems],
-                [e[:, 2::2] for e in elems])
-        else:
-            even_elems = operator(
-                odd_elems,
-                [e[:, 2::2] for e in elems])
-
-        # The first element of a scan is the same as the first element
-        # of the original `elems`.
-
-        even_elems = [
-          torch.cat([elem[:, :1], result], dim=1)
-          for (elem, result) in zip(elems, even_elems)]
-
-        return list(map(_interleave, even_elems, odd_elems))
-
-    return _scan(elems)
-
-def _interleave(a, b):
-    a_axis_len, b_axis_len = a.shape[1], b.shape[1]
-    output_axis_len = a_axis_len + b_axis_len
-
-    if (a_axis_len == (b_axis_len + 1)):
-        b = pad_at_dim(b, (0, 1), dim = 1)
-
-    stacked = torch.stack([a, b], dim=2)
-    interleaved = torch.flatten(stacked, start_dim=1, end_dim=2)
-
-    return interleaved[:, :output_axis_len]
-
-# associative scan wrapper around naive and accelerated version
-
-class AssocScan(Module):
-    def __init__(
-        self,
-        use_accelerated = False
-    ):
-        super().__init__()
-        self.use_accelerated = use_accelerated
-
-    def forward(
-        self,
-        gates,
-        inputs,
-        prev = None,
-        remove_prev = None
-    ):
-        remove_prev = default(remove_prev, exists(prev))
-
-        inputs, inverse_pack_weight_shape = pack_one_with_inverse(inputs, 'b n *')
-        gates, _ = pack_one_with_inverse(gates, 'b n *')
-
-        if exists(prev):
-            prev, _ = pack_one_with_inverse(prev, 'b *')
-
-        if exists(prev):
-            inputs, _ = pack([prev, inputs], 'b * d')
-            gates = pad_at_dim(gates, (1, 0), value = 1., dim = -2)
-
-        if not self.use_accelerated:
-            _, out = associative_scan(binary_operator, (gates, inputs))
-
-            if remove_prev:
-                out = out[:, 1:]
-
-            return inverse_pack_weight_shape(out)
-
-        from accelerated_scan.triton import scan as triton_scan
-        from accelerated_scan.warp import scan as warp_scan
-
-        scan = triton_scan if gates.is_cuda else warp_scan
-
-        def accelerate_scan_fn(gates, inputs):
-            gates = gates.expand_as(inputs)
-            gates, inputs = tuple(rearrange(t, 'b n d -> b d n') for t in (gates, inputs))
-
-            seq_len = gates.shape[-1]
-            next_power_two_seq_len = 2 ** max(5, int(math.ceil(math.log2(seq_len))))
-
-            gates = F.pad(gates, (0, next_power_two_seq_len - seq_len))
-            inputs = F.pad(inputs, (0, next_power_two_seq_len - seq_len))
-
-            outputs = scan(gates.contiguous(), inputs.contiguous())
-
-            outputs = outputs[..., :seq_len]
-            outputs = rearrange(outputs, 'b d n -> b n d')
-
-            return outputs
-
-        out = accelerate_scan_fn(gates, inputs)
-
-        if remove_prev:
-            out = out[:, 1:]
-
-        return inverse_pack_weight_shape(out)

titans_pytorch-0.4.6.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=sVTOuRUkaIYabFExdLY6s1qXm1UwHHz_J19H8ZV-X74,338
-titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
-titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=CsB-Wd3T_N_HluRLjjwB2gPdlLymJmmhuxiRJhQuXkA,33377
-titans_pytorch-0.4.6.dist-info/METADATA,sha256=aWhQMQrjBLzUPmvtH-LY47r4ayz_ts70gGyDMMyJ6Sc,6810
-titans_pytorch-0.4.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.4.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.4.6.dist-info/RECORD,,