ChipYTY/titans_NPC

1<img src="./fig2.png" width="400px"></img>2 3<img src="./fig1.png" width="400px"></img>4 5## Titans - Pytorch6 7Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytorch. Will also contain some explorations into architectures beyond their simple 1-4 layer MLP for the neural memory module, if it works well to any degree.8 9[Paper review by Yannic](https://www.youtube.com/watch?v=v67plFw1nMw)10 11[Quick Colab Run](https://colab.research.google.com/drive/11cGgSABykte3qbK-hjzPgLif3-9UUejm?usp=sharing)12 13## Appreciation14 15- [Eryk](https://github.com/sentialx) for sharing his early experimental results with me, positive for 2 layer MLP16 17## Install18 19```bash20$ pip install titans-pytorch21```22 23## Usage24 25```python26import torch27from titans_pytorch import NeuralMemory28 29mem = NeuralMemory(30    dim = 384,31    chunk_size = 64 # set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage)32).cuda()33 34seq = torch.randn(2, 1024, 384).cuda()35retrieved, mem_state = mem(seq)36 37assert seq.shape == retrieved.shape38```39 40A transformer with the `MAC` configuration can be used as41 42```python43import torch44from titans_pytorch import MemoryAsContextTransformer45 46transformer = MemoryAsContextTransformer(47    num_tokens = 256,48    dim = 256,49    depth = 2,50    segment_len = 128,              # local attention window size51    num_persist_mem_tokens = 4,52    num_longterm_mem_tokens = 16,53)54 55token_ids = torch.randint(0, 256, (1, 1023))56 57loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)58loss.backward()59 60# after much training61 62sampled = transformer.sample(token_ids[:, :4], 512)63```64 65## Experiments66 67```bash68$ pip install uv69```70 71Then modify `train_mac.py` and run it to query nature72 73```bash74$ uv run train_mac.py75```76 77## Citations78 79```bibtex80@inproceedings{Behrouz2024TitansLT,81    title   = {Titans: Learning to Memorize at Test Time},82    author  = {Ali Behrouz and Peilin Zhong and Vahab S. Mirrokni},83    year    = {2024},84    url     = {https://api.semanticscholar.org/CorpusID:275212078}85}86```87 88```bibtex89@article{Sun2024LearningT,90    title   = {Learning to (Learn at Test Time): RNNs with Expressive Hidden States},91    author  = {Yu Sun and Xinhao Li and Karan Dalal and Jiarui Xu and Arjun Vikram and Genghan Zhang and Yann Dubois and Xinlei Chen and Xiaolong Wang and Oluwasanmi Koyejo and Tatsunori Hashimoto and Carlos Guestrin},92    journal = {ArXiv},93    year    = {2024},94    volume  = {abs/2407.04620},95    url     = {https://api.semanticscholar.org/CorpusID:271039606}96}97```98 99```bibtex100@inproceedings{Yang2024GatedDN,101    title   = {Gated Delta Networks: Improving Mamba2 with Delta Rule},102    author  = {Songlin Yang and Jan Kautz and Ali Hatamizadeh},103    year    = {2024},104    url     = {https://api.semanticscholar.org/CorpusID:274598177}105}106```107 108```bibtex109@inproceedings{Nguyen2024TurningUT,110    title   = {Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs},111    author  = {Minh Nguyen and Andrew Baker and Clement Neo and Allen Roush and Andreas Kirsch and Ravid Shwartz-Ziv},112    year    = {2024},113    url     = {https://api.semanticscholar.org/CorpusID:270870613}114}115```116 117```bibtex118@article{Zhu2024HyperConnections,119    title   = {Hyper-Connections},120    author  = {Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou},121    journal = {ArXiv},122    year    = {2024},123    volume  = {abs/2409.19606},124    url     = {https://api.semanticscholar.org/CorpusID:272987528}125}126```127 128```bibtex129@article{Zhou2024ValueRL,130    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},131    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},132    journal = {ArXiv},133    year    = {2024},134    volume  = {abs/2410.17897},135    url     = {https://api.semanticscholar.org/CorpusID:273532030}136}137```138 139```bibtex140@software{Kyrylov_Accelerated_Scan_2024,141    author  = {Kyrylov, Volodymyr},142    doi     = {10.5281/zenodo.10600962},143    title   = {Accelerated Scan},144    version = {0.1.2},145    year    = {2024}146}147```148 149```bibtex150@misc{wang2025testtimeregressionunifyingframework,151    title   = {Test-time regression: a unifying framework for designing sequence models with associative memory},152    author  = {Ke Alexander Wang and Jiaxin Shi and Emily B. Fox},153    year    = {2025},154    eprint  = {2501.12352},155    archivePrefix = {arXiv},156    primaryClass = {cs.LG},157    url     = {https://arxiv.org/abs/2501.12352},158}159```160 161```bibtex162@misc{jordan2024muon,163    author  = {Keller Jordan and Yuchen Jin and Vlado Boza and Jiacheng You and164                    Franz Cesista and Laker Newhouse and Jeremy Bernstein},165    title   = {Muon: An optimizer for hidden layers in neural networks},166    year    = {2024},167    url     = {https://kellerjordan.github.io/posts/muon/}168}169```170 171```bibtex172@inproceedings{Zhang2025TestTimeTD,173    title   = {Test-Time Training Done Right},174    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},175    year    = {2025},176    url     = {https://api.semanticscholar.org/CorpusID:279071244}177}178```179 180```bibtex181@inproceedings{Behrouz2025ATLASLT,182    title  = {ATLAS: Learning to Optimally Memorize the Context at Test Time},183    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},184    year   = {2025},185    url    = {https://api.semanticscholar.org/CorpusID:278996373}186}187```188