titans-memory 0.1.0__tar.gz

titans_memory-0.1.0/PKG-INFO

@@ -0,0 +1,264 @@
+Metadata-Version: 2.4
+Name: titans-memory
+Version: 0.1.0
+Summary: PyTorch implementation of Titans: Learning to Memorize at Test Time (Behrouz, Zhong & Mirrokni, 2024)
+Author: Neuranox, Implementation of arXiv:2501.00663
+License: MIT
+Project-URL: Homepage, https://github.com/Neuranox/titans-memory
+Project-URL: Bug Tracker, https://github.com/Neuranox/titans-memory/issues
+Project-URL: Paper, https://arxiv.org/abs/2501.00663
+Keywords: deep-learning,transformers,long-context,memory,titans,neural-memory
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.1.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Provides-Extra: train
+Requires-Dist: datasets>=2.16; extra == "train"
+Requires-Dist: transformers>=4.38; extra == "train"
+Requires-Dist: tqdm; extra == "train"
+Requires-Dist: tensorboard; extra == "train"
+Provides-Extra: all
+Requires-Dist: datasets>=2.16; extra == "all"
+Requires-Dist: transformers>=4.38; extra == "all"
+Requires-Dist: tqdm; extra == "all"
+Requires-Dist: tensorboard; extra == "all"
+Requires-Dist: pytest>=7.0; extra == "all"
+Requires-Dist: pytest-cov; extra == "all"
+
+# Titans: Learning to Memorize at Test Time
+
+[![Python](https://img.shields.io/badge/python-3.9%2B-blue)](https://python.org)
+[![PyTorch](https://img.shields.io/badge/pytorch-2.1%2B-orange)](https://pytorch.org)
+[![GitHub release](https://img.shields.io/github/v/release/Neuranox/titans-memory)](https://github.com/Neuranox/titans-memory)
+[![License: MIT](https://img.shields.io/badge/License-MIT-green)](LICENSE)
+[![arXiv](https://img.shields.io/badge/arXiv-2501.00663-red)](https://arxiv.org/abs/2501.00663)
+
+A clean, highly-optimized PyTorch implementation of the **Titans** architecture from:
+
+> **Titans: Learning to Memorize at Test Time**  
+> Ali Behrouz, Peilin Zhong, Vahab Mirrokni — Google Research, 2024  
+> [arXiv:2501.00663](https://arxiv.org/abs/2501.00663)
+
+<p align="center">
+  <img src="assets/image.png" alt="Titans Architecture Overview" width="80%">
+</p>
+
+---
+
+## What's Inside
+
+| Module | Description |
+|---|---|
+| `NeuralMemory` | Deep MLP that learns to memorize via gradient descent with **momentum** + **weight-decay forgetting** (§3) |
+| `PersistentMemory` | Learnable task-knowledge tokens prepended to every sequence (§3.3) |
+| `TitansMAC` | **Memory as a Context** — retrieves long-term memory as prefix to attention window (§4.1) |
+| `TitansMAG` | **Memory as a Gate** — SWA ⊗ NeuralMemory gated branch (§4.2) |
+| `TitansMAL` | **Memory as a Layer** — sequential LMM → SWA stack (§4.3) |
+| `TitansLMM` | **Standalone LMM** — neural memory without attention (§4.3) |
+
+---
+
+## Installation
+
+```bash
+# Install directly from GitHub
+pip install git+https://github.com/Neuranox/titans-memory.git
+
+# Or clone and install locally (editable — recommended for development)
+git clone https://github.com/Neuranox/titans-memory.git
+cd titans-memory
+pip install -e .
+```
+
+---
+
+## Quick Start
+
+```python
+import torch
+from titans import TitansMAC, TitansMAG, TitansMAL, TitansLMM
+from titans.utils import TitansConfig, build_model, count_parameters
+
+# ── Build from config ──────────────────────────────────────────────────
+cfg   = TitansConfig.small(variant="MAC")   # ~170 M params
+cfg.vocab_size = 32_000
+model = build_model(cfg)
+print(f"Parameters: {count_parameters(model):,}")
+
+# ── Forward pass ───────────────────────────────────────────────────────
+input_ids = torch.randint(0, 32_000, (2, 512))
+labels    = input_ids.clone()
+
+out = model(input_ids, labels=labels)
+print(out["logits"].shape)   # (2, 512, 32000)
+print(out["loss"].item())
+
+# ── Generation ─────────────────────────────────────────────────────────
+prompt    = torch.randint(0, 32_000, (1, 8))
+generated = model.generate(prompt, max_new_tokens=50, top_k=50)
+```
+
+---
+
+## All Four Variants
+
+```python
+VOCAB = 32_000
+D     = 512
+
+models = {
+    "LMM": TitansLMM(VOCAB, d_model=D, n_layers=12, mem_layers=2),
+    "MAC": TitansMAC(VOCAB, d_model=D, n_layers=12, mem_layers=2, chunk_size=128),
+    "MAG": TitansMAG(VOCAB, d_model=D, n_layers=12, mem_layers=2, window=512),
+    "MAL": TitansMAL(VOCAB, d_model=D, n_layers=12, mem_layers=2, window=512),
+}
+```
+
+---
+
+## TitansConfig — Paper-Scale Presets
+
+```python
+from titans.utils import TitansConfig, build_model
+
+cfg = TitansConfig.tiny(variant="MAC")    # ~30 M  — quick experiments
+cfg = TitansConfig.small(variant="MAC")   # ~170 M — paper Table 1
+cfg = TitansConfig.medium(variant="MAC")  # ~340 M — paper Table 1
+cfg = TitansConfig.large(variant="MAC")   # ~760 M — paper Table 1
+
+# JSON save / load
+cfg.to_json("config.json")
+cfg = TitansConfig.from_json("config.json")
+```
+
+---
+
+## Training
+
+```python
+from titans.utils.training import build_optimizer, get_cosine_schedule_with_warmup
+
+optim = build_optimizer(model, lr=4e-4, weight_decay=0.1)          # AdamW, no wd on bias/norm
+sched = get_cosine_schedule_with_warmup(optim,
+            warmup_steps=2000, total_steps=100_000, min_lr_ratio=0.1)
+
+for batch in dataloader:
+    out  = model(batch["input_ids"], labels=batch["labels"])
+    out["loss"].backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    optim.step(); sched.step(); optim.zero_grad()
+```
+
+See `examples/02_training_loop.py` for a complete runnable example.
+
+---
+
+## Architecture Overview
+
+<p align="center">
+  <img src="assets/image1.png" alt="Titans Detail" width="80%">
+</p>
+
+```
+Titans (MAC) — Memory as a Context
+───────────────────────────────────────────────────────────
+ For each segment S^(t):
+   h_t     = M*_{t-1}(q_t)           ← retrieve long-term memory
+   S̃^(t)  = [P || h_t || S^(t)]     ← augment with persistent + history
+   y_t     = Attention(S̃^(t))        ← full causal attention over window
+   M_t     = M_{t-1}.update(y_t)    ← write: gradient descent w/ momentum
+   o_t     = y_t ⊗ M*_t(y_t)        ← gated output
+
+Titans (MAG) — Memory as a Gate
+───────────────────────────────────────────────────────────
+   x̃  = [P || x]
+   y   = SW-Attn(x̃)  ← precise short-term memory (sliding window)
+   o   = y ⊗ M(x̃)   ← gated with neural long-term memory
+
+Titans (MAL) — Memory as a Layer
+───────────────────────────────────────────────────────────
+   x̃  = [P || x]
+   y   = M(x̃)         ← memory compresses context
+   o   = SW-Attn(y)   ← attention refines compressed representation
+```
+
+---
+
+## Neural Memory — Key Equations
+
+| Component | Equation | Description |
+|---|---|---|
+| Momentary surprise | `∇ℓ(M_{t-1}; x_t)` | How unexpected is `x_t`? |
+| Surprise with momentum | `S_t = η_t S_{t-1} − θ_t ∇ℓ` | Eq. 10 — carries information flow |
+| Forgetting gate | `M_t = (1−α_t) M_{t-1} + S_t` | Eq. 13 — weight-decay style |
+| Retrieval | `y_t = M*(q_t)` | Eq. 15 — inference, no update |
+
+---
+
+## Running Tests
+
+```bash
+cd "F:\Titan Model"
+pip install -e .[dev]
+pytest
+```
+
+---
+
+## Project Structure
+
+```
+Titan Model/
+├── titans/
+│   ├── __init__.py           ← public API
+│   ├── memory/
+│   │   ├── neural_memory.py  ← NeuralMemory (LMM core)
+│   │   └── persistent_memory.py
+│   ├── models/
+│   │   ├── lmm.py            ← TitansLMM
+│   │   ├── mac.py            ← TitansMAC
+│   │   ├── mag.py            ← TitansMAG
+│   │   └── mal.py            ← TitansMAL
+│   ├── ops/
+│   │   ├── scan.py           ← parallel associative scan
+│   │   └── attention.py      ← causal + sliding-window attention
+│   └── utils/
+│       ├── config.py         ← TitansConfig dataclass
+│       ├── factory.py        ← build_model()
+│       └── training.py       ← optimizer + LR schedule helpers
+├── tests/
+│   ├── test_scan.py
+│   ├── test_memory.py
+│   └── test_models.py
+├── examples/
+│   ├── 01_quickstart.py
+│   ├── 02_training_loop.py
+│   └── 03_memory_standalone.py
+├── pyproject.toml
+├── setup.py
+└── README.md
+```
+
+---
+
+## Citation
+
+```bibtex
+@article{behrouz2024titans,
+  title   = {Titans: Learning to Memorize at Test Time},
+  author  = {Behrouz, Ali and Zhong, Peilin and Mirrokni, Vahab},
+  journal = {arXiv preprint arXiv:2501.00663},
+  year    = {2024}
+}
+```

titans_memory-0.1.0/README.md

@@ -0,0 +1,226 @@
+# Titans: Learning to Memorize at Test Time
+
+[![Python](https://img.shields.io/badge/python-3.9%2B-blue)](https://python.org)
+[![PyTorch](https://img.shields.io/badge/pytorch-2.1%2B-orange)](https://pytorch.org)
+[![GitHub release](https://img.shields.io/github/v/release/Neuranox/titans-memory)](https://github.com/Neuranox/titans-memory)
+[![License: MIT](https://img.shields.io/badge/License-MIT-green)](LICENSE)
+[![arXiv](https://img.shields.io/badge/arXiv-2501.00663-red)](https://arxiv.org/abs/2501.00663)
+
+A clean, highly-optimized PyTorch implementation of the **Titans** architecture from:
+
+> **Titans: Learning to Memorize at Test Time**  
+> Ali Behrouz, Peilin Zhong, Vahab Mirrokni — Google Research, 2024  
+> [arXiv:2501.00663](https://arxiv.org/abs/2501.00663)
+
+<p align="center">
+  <img src="assets/image.png" alt="Titans Architecture Overview" width="80%">
+</p>
+
+---
+
+## What's Inside
+
+| Module | Description |
+|---|---|
+| `NeuralMemory` | Deep MLP that learns to memorize via gradient descent with **momentum** + **weight-decay forgetting** (§3) |
+| `PersistentMemory` | Learnable task-knowledge tokens prepended to every sequence (§3.3) |
+| `TitansMAC` | **Memory as a Context** — retrieves long-term memory as prefix to attention window (§4.1) |
+| `TitansMAG` | **Memory as a Gate** — SWA ⊗ NeuralMemory gated branch (§4.2) |
+| `TitansMAL` | **Memory as a Layer** — sequential LMM → SWA stack (§4.3) |
+| `TitansLMM` | **Standalone LMM** — neural memory without attention (§4.3) |
+
+---
+
+## Installation
+
+```bash
+# Install directly from GitHub
+pip install git+https://github.com/Neuranox/titans-memory.git
+
+# Or clone and install locally (editable — recommended for development)
+git clone https://github.com/Neuranox/titans-memory.git
+cd titans-memory
+pip install -e .
+```
+
+---
+
+## Quick Start
+
+```python
+import torch
+from titans import TitansMAC, TitansMAG, TitansMAL, TitansLMM
+from titans.utils import TitansConfig, build_model, count_parameters
+
+# ── Build from config ──────────────────────────────────────────────────
+cfg   = TitansConfig.small(variant="MAC")   # ~170 M params
+cfg.vocab_size = 32_000
+model = build_model(cfg)
+print(f"Parameters: {count_parameters(model):,}")
+
+# ── Forward pass ───────────────────────────────────────────────────────
+input_ids = torch.randint(0, 32_000, (2, 512))
+labels    = input_ids.clone()
+
+out = model(input_ids, labels=labels)
+print(out["logits"].shape)   # (2, 512, 32000)
+print(out["loss"].item())
+
+# ── Generation ─────────────────────────────────────────────────────────
+prompt    = torch.randint(0, 32_000, (1, 8))
+generated = model.generate(prompt, max_new_tokens=50, top_k=50)
+```
+
+---
+
+## All Four Variants
+
+```python
+VOCAB = 32_000
+D     = 512
+
+models = {
+    "LMM": TitansLMM(VOCAB, d_model=D, n_layers=12, mem_layers=2),
+    "MAC": TitansMAC(VOCAB, d_model=D, n_layers=12, mem_layers=2, chunk_size=128),
+    "MAG": TitansMAG(VOCAB, d_model=D, n_layers=12, mem_layers=2, window=512),
+    "MAL": TitansMAL(VOCAB, d_model=D, n_layers=12, mem_layers=2, window=512),
+}
+```
+
+---
+
+## TitansConfig — Paper-Scale Presets
+
+```python
+from titans.utils import TitansConfig, build_model
+
+cfg = TitansConfig.tiny(variant="MAC")    # ~30 M  — quick experiments
+cfg = TitansConfig.small(variant="MAC")   # ~170 M — paper Table 1
+cfg = TitansConfig.medium(variant="MAC")  # ~340 M — paper Table 1
+cfg = TitansConfig.large(variant="MAC")   # ~760 M — paper Table 1
+
+# JSON save / load
+cfg.to_json("config.json")
+cfg = TitansConfig.from_json("config.json")
+```
+
+---
+
+## Training
+
+```python
+from titans.utils.training import build_optimizer, get_cosine_schedule_with_warmup
+
+optim = build_optimizer(model, lr=4e-4, weight_decay=0.1)          # AdamW, no wd on bias/norm
+sched = get_cosine_schedule_with_warmup(optim,
+            warmup_steps=2000, total_steps=100_000, min_lr_ratio=0.1)
+
+for batch in dataloader:
+    out  = model(batch["input_ids"], labels=batch["labels"])
+    out["loss"].backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    optim.step(); sched.step(); optim.zero_grad()
+```
+
+See `examples/02_training_loop.py` for a complete runnable example.
+
+---
+
+## Architecture Overview
+
+<p align="center">
+  <img src="assets/image1.png" alt="Titans Detail" width="80%">
+</p>
+
+```
+Titans (MAC) — Memory as a Context
+───────────────────────────────────────────────────────────
+ For each segment S^(t):
+   h_t     = M*_{t-1}(q_t)           ← retrieve long-term memory
+   S̃^(t)  = [P || h_t || S^(t)]     ← augment with persistent + history
+   y_t     = Attention(S̃^(t))        ← full causal attention over window
+   M_t     = M_{t-1}.update(y_t)    ← write: gradient descent w/ momentum
+   o_t     = y_t ⊗ M*_t(y_t)        ← gated output
+
+Titans (MAG) — Memory as a Gate
+───────────────────────────────────────────────────────────
+   x̃  = [P || x]
+   y   = SW-Attn(x̃)  ← precise short-term memory (sliding window)
+   o   = y ⊗ M(x̃)   ← gated with neural long-term memory
+
+Titans (MAL) — Memory as a Layer
+───────────────────────────────────────────────────────────
+   x̃  = [P || x]
+   y   = M(x̃)         ← memory compresses context
+   o   = SW-Attn(y)   ← attention refines compressed representation
+```
+
+---
+
+## Neural Memory — Key Equations
+
+| Component | Equation | Description |
+|---|---|---|
+| Momentary surprise | `∇ℓ(M_{t-1}; x_t)` | How unexpected is `x_t`? |
+| Surprise with momentum | `S_t = η_t S_{t-1} − θ_t ∇ℓ` | Eq. 10 — carries information flow |
+| Forgetting gate | `M_t = (1−α_t) M_{t-1} + S_t` | Eq. 13 — weight-decay style |
+| Retrieval | `y_t = M*(q_t)` | Eq. 15 — inference, no update |
+
+---
+
+## Running Tests
+
+```bash
+cd "F:\Titan Model"
+pip install -e .[dev]
+pytest
+```
+
+---
+
+## Project Structure
+
+```
+Titan Model/
+├── titans/
+│   ├── __init__.py           ← public API
+│   ├── memory/
+│   │   ├── neural_memory.py  ← NeuralMemory (LMM core)
+│   │   └── persistent_memory.py
+│   ├── models/
+│   │   ├── lmm.py            ← TitansLMM
+│   │   ├── mac.py            ← TitansMAC
+│   │   ├── mag.py            ← TitansMAG
+│   │   └── mal.py            ← TitansMAL
+│   ├── ops/
+│   │   ├── scan.py           ← parallel associative scan
+│   │   └── attention.py      ← causal + sliding-window attention
+│   └── utils/
+│       ├── config.py         ← TitansConfig dataclass
+│       ├── factory.py        ← build_model()
+│       └── training.py       ← optimizer + LR schedule helpers
+├── tests/
+│   ├── test_scan.py
+│   ├── test_memory.py
+│   └── test_models.py
+├── examples/
+│   ├── 01_quickstart.py
+│   ├── 02_training_loop.py
+│   └── 03_memory_standalone.py
+├── pyproject.toml
+├── setup.py
+└── README.md
+```
+
+---
+
+## Citation
+
+```bibtex
+@article{behrouz2024titans,
+  title   = {Titans: Learning to Memorize at Test Time},
+  author  = {Behrouz, Ali and Zhong, Peilin and Mirrokni, Vahab},
+  journal = {arXiv preprint arXiv:2501.00663},
+  year    = {2024}
+}
+```

titans_memory-0.1.0/pyproject.toml

@@ -0,0 +1,66 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name            = "titans-memory"
+version         = "0.1.0"
+description     = "PyTorch implementation of Titans: Learning to Memorize at Test Time (Behrouz, Zhong & Mirrokni, 2024)"
+readme          = "README.md"
+license         = { text = "MIT" }
+requires-python = ">=3.9"
+keywords        = ["deep-learning", "transformers", "long-context", "memory", "titans", "neural-memory"]
+
+authors = [
+  { name = "Neuranox" },
+  { name = "Implementation of arXiv:2501.00663" },
+]
+
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+  "torch>=2.1.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=7.0",
+  "pytest-cov",
+]
+train = [
+  "datasets>=2.16",
+  "transformers>=4.38",
+  "tqdm",
+  "tensorboard",
+]
+all = [
+  "datasets>=2.16",
+  "transformers>=4.38",
+  "tqdm",
+  "tensorboard",
+  "pytest>=7.0",
+  "pytest-cov",
+]
+
+[project.urls]
+"Homepage"    = "https://github.com/Neuranox/titans-memory"
+"Bug Tracker" = "https://github.com/Neuranox/titans-memory/issues"
+"Paper"       = "https://arxiv.org/abs/2501.00663"
+
+[tool.setuptools.packages.find]
+where   = ["."]
+include = ["titans*"]
+
+[tool.pytest.ini_options]
+testpaths    = ["tests"]
+addopts      = "-v --tb=short"

titans_memory-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

titans_memory-0.1.0/setup.py

@@ -0,0 +1,8 @@
+"""
+Fallback setup.py for environments that don't support PEP 517 build.
+All configuration is in pyproject.toml.
+"""
+from setuptools import setup
+
+if __name__ == "__main__":
+    setup()

titans_memory-0.1.0/tests/test_memory.py

@@ -0,0 +1,85 @@
+"""
+Tests for NeuralMemory and PersistentMemory modules.
+"""
+import pytest
+import torch
+from titans.memory import NeuralMemory, PersistentMemory
+
+
+class TestNeuralMemory:
+
+    @pytest.fixture
+    def mem(self):
+        return NeuralMemory(d_model=32, n_layers=2, chunk_size=8)
+
+    def test_forward_shape(self, mem):
+        x   = torch.randn(2, 16, 32)
+        out = mem(x)
+        assert out.shape == x.shape, f"Expected {x.shape}, got {out.shape}"
+
+    def test_retrieve_shape(self, mem):
+        q   = torch.randn(2, 8, 32)
+        out = mem.retrieve(q)
+        assert out.shape == q.shape
+
+    @pytest.mark.parametrize("lm", [1, 2, 3])
+    def test_memory_depth(self, lm):
+        m   = NeuralMemory(d_model=16, n_layers=lm, chunk_size=4)
+        x   = torch.randn(1, 8, 16)
+        out = m(x)
+        assert out.shape == x.shape
+
+    def test_no_momentum(self):
+        m   = NeuralMemory(d_model=16, n_layers=1, chunk_size=4, use_momentum=False)
+        x   = torch.randn(1, 8, 16)
+        out = m(x)
+        assert out.shape == x.shape
+
+    def test_no_decay(self):
+        m   = NeuralMemory(d_model=16, n_layers=1, chunk_size=4, use_decay=False)
+        x   = torch.randn(1, 8, 16)
+        out = m(x)
+        assert out.shape == x.shape
+
+    def test_gradients_flow(self, mem):
+        x   = torch.randn(1, 8, 32, requires_grad=True)
+        out = mem(x)
+        out.sum().backward()
+        assert x.grad is not None
+        assert not torch.isnan(x.grad).any()
+
+    def test_chunk_size_larger_than_seq(self):
+        m   = NeuralMemory(d_model=16, n_layers=1, chunk_size=128)
+        x   = torch.randn(1, 10, 16)          # T < chunk_size
+        out = m(x)
+        assert out.shape == x.shape
+
+
+class TestPersistentMemory:
+
+    def test_prepend_shape(self):
+        pm  = PersistentMemory(n_tokens=8, d_model=32)
+        x   = torch.randn(2, 16, 32)
+        out = pm(x)
+        assert out.shape == (2, 24, 32)
+
+    def test_strip_shape(self):
+        pm  = PersistentMemory(n_tokens=8, d_model=32)
+        x   = torch.randn(2, 16, 32)
+        aug = pm(x)
+        stripped = pm.strip(aug)
+        assert stripped.shape == x.shape
+
+    def test_freeze_unfreeze(self):
+        pm = PersistentMemory(n_tokens=4, d_model=8)
+        pm.freeze()
+        assert not pm.P.requires_grad
+        pm.unfreeze()
+        assert pm.P.requires_grad
+
+    def test_gradient_through_persistent(self):
+        pm  = PersistentMemory(n_tokens=4, d_model=8)
+        x   = torch.randn(1, 4, 8)
+        aug = pm(x)
+        aug.sum().backward()
+        assert pm.P.grad is not None