titans-pytorch 0.4.7__tar.gz → 0.4.9__tar.gz

{titans_pytorch-0.4.7 → titans_pytorch-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.4.7
+Version: 0.4.9
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -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.7 → titans_pytorch-0.4.9}/README.md

@@ -153,3 +153,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.7 → titans_pytorch-0.4.9}/pyproject.toml

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

{titans_pytorch-0.4.7 → titans_pytorch-0.4.9}/titans_pytorch/neural_memory.py

@@ -152,6 +152,39 @@ 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
+
+    shape = t.shape
+    should_transpose = shape[2] > shape[-1]
+
+    if should_transpose:
+        t = t.transpose(-1, -2)
+
+    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
+
+    if should_transpose:
+        t = t.transpose(-1, -2)
+
+    return inv_pack(t)
+
 # multi head rmsnorm
 
 class MultiheadRMSNorm(Module):
@@ -254,6 +287,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 +499,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 +786,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.7 → titans_pytorch-0.4.9}/train_mac.py

@@ -49,6 +49,7 @@ STORE_ATTN_POOL_CHUNKS = True                   # whether to use attention pooli
 MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
 NEURAL_MEM_WEIGHT_RESIDUAL = True               # learning to accept contributions from the weights of the previous neural mem layer brings about significant improvements. this was improvised and not in the paper, but inspired by the value residual learning free lunch paper
 NEURAL_MEM_QKV_RECEIVES_DIFF_VIEW = True        # will allow the neural memory to select what layers from which to derive queries / keys / values, effectively allowing it to graft itself to the transformer in any way to be beneficial. this is to address an issue from a phd student who noted that the mem network is learning nothing more than wk @ wv. this also generalizes all possible ways to connect the neural memory to a transformer, a sort of NAS
+NEURAL_MEM_SPEC_NORM_SURPRISES = True           # applying lessons from Muon optimizer to surprise updates, by spectral norming the surprises
 
 # experiment related
 
@@ -121,7 +122,8 @@ model = MemoryAsContextTransformer(
         momentum_order = NEURAL_MEM_MOMENTUM_ORDER,
         default_step_transform_max_lr = NEURAL_MEM_MAX_LR,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
-        per_parameter_lr_modulation = MEMORY_MODEL_PER_LAYER_LEARNED_LR
+        per_parameter_lr_modulation = MEMORY_MODEL_PER_LAYER_LEARNED_LR,
+        spectral_norm_surprises = NEURAL_MEM_SPEC_NORM_SURPRISES
     )
 ).cuda()