open-titans 0.0.2

Memory-Augmented Sequence Models in Pytorch

🌌 OpenTitans

The Open-Source Framework for Memory-Augmented Sequence Models

---

Democratizing Test-Time Memorization and Neural Memory Architectures.

Introduction • Features • Quick Start • Usage • Citations

🌟 Introduction

OpenTitans is a modular, high-performance framework designed to implement and explore the next generation of sequence models. While Transformers revolutionized AI, their quadratic context limitations have met their match.

Inspired by groundbreaking research from Google and other top labs, OpenTitans focuses on Memory-Augmented Models that learn to memorize, optimize, and cache their internal states at test time. Our goal is to provide a "HuggingFace-like" experience for researchers and engineers building the future of infinite-context modeling.

---

🚀 Key Features

• 🧠 Neural Memory Modules: Sophisticated implementations of memory systems that evolve during inference.
• ⚡ Associative Scan: Optimized support for fast sequence processing, breaking the linear bottleneck.
• 🛠️ Test-Time Training (TTT): Seamless integration of models that optimize their parameters as they consume context.
• 🧩 Modular & Extensible: A clean, object-oriented API that allows you to swap memory models, poolers, and norms with ease.
• 📊 Performance Benchmarks: Built-in tools to measure latency, throughput, and memory efficiency on modern GPUs.

---

📦 Quick Start

Installation

[!TIP] We recommend using a virtual environment (venv or conda) for the best experience.

pip install open-titans

Development Installation

# Clone the repository
git clone https://github.com/Neeze/OpenTitans.git
cd OpenTitans

# Install in editable mode with dependencies
pip install -e .

---

🤝 Contributing

We are looking for "Titans" to help us build! 🚀

Whether you want to implement a new paper, optimize a CUDA kernel, or just fix a typo, your contributions are welcome. Check out our CONTRIBUTING.md to get started.

---

📚 Citations & Acknowledgements

OpenTitans stands on the shoulders of giants. We acknowledge the authors of the following papers for their foundational work:

@misc{behrouz2024titanslearningmemorizetest,
      title={Titans: Learning to Memorize at Test Time}, 
      author={Ali Behrouz and Peilin Zhong and Vahab Mirrokni},
      year={2024},
      url={https://arxiv.org/abs/2501.00663}
}

@misc{behrouz2025atlaslearningoptimallymemorize,
      title={ATLAS: Learning to Optimally Memorize the Context at Test Time}, 
      author={Ali Behrouz and Zeman Li and Praneeth Kacham and Majid Daliri and Yuan Deng and Peilin Zhong and Meisam Razaviyayn and Vahab Mirrokni},
      year={2025},
      url={https://arxiv.org/abs/2505.23735}
}

@misc{behrouz2025itsconnectedjourneytesttime,
      title={It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization}, 
      author={Ali Behrouz and Meisam Razaviyayn and Peilin Zhong and Vahab Mirrokni},
      year={2025},
      url={https://arxiv.org/abs/2504.13173}
}

---

📄 License

OpenTitans is released under the MIT License. See LICENSE for more details.