Metadata-Version: 2.4
Name: titans-torch
Version: 0.1.1
Summary: PyTorch implementation of Titans: Learning to Memorize at Test Time
Home-page: https://github.com/buiquangdat1710/titans-torch
Author: Bui Quang Dat
Author-email: buiquangdat1458@gmail.com
Project-URL: Bug Tracker, https://github.com/buiquangdat1710/titans-torch/issues
Project-URL: Source Code, https://github.com/buiquangdat1710/titans-torch
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/plain
License-File: LICENSE.txt
Requires-Dist: torch>=2.0.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: icecream>=2.1.0; extra == "dev"
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license-file
Dynamic: project-url
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# titans-torch

> PyTorch implementation of **"Titans: Learning to Memorize at Test Time"**
> — Behrouz, Zhong & Daliri, arXiv:2501.00663, 2025.

Titans là kiến trúc kết hợp **short-term attention** và **Neural Long-Term Memory (LMM)** — một module bộ nhớ dạng MLP có thể được cập nhật tại test time thông qua inner-loop gradient descent, cho phép mô hình thích nghi với chuỗi mới mà không cần huấn luyện lại.

---

## Mục lục

- [Tính năng](#tính-năng)
- [Cài đặt](#cài-đặt)
- [Kiến trúc tổng quan](#kiến-trúc-tổng-quan)
- [Sử dụng nhanh](#sử-dụng-nhanh)
- [API Reference](#api-reference)
  - [NeuralLongTermMemory](#neurallongtermmemory)
  - [PersistentMemory](#persistentmemory)
  - [TitansMAC](#titansmac)
  - [TitansMAG](#titansmag)
  - [TitansMAL](#titansmal)
- [Ví dụ training](#ví-dụ-training)
- [Cấu trúc thư mục](#cấu-trúc-thư-mục)
- [Tham khảo](#tham-khảo)
- [License](#license)

---

## Tính năng

- ✅ **Neural Long-Term Memory (LMM)** với inner-loop differentiable — huấn luyện end-to-end
- ✅ **Test-time adaptation** — cập nhật memory theo từng token mà không cần giữ computation graph
- ✅ **3 kiến trúc Titans**: MAC · MAG · MAL
- ✅ **Persistent Memory** — learnable token bank prepend vào sequence
- ✅ **Data-dependent coefficients** (α, θ, η) điều khiển learning rate và momentum
- ✅ **1D depthwise conv** tùy chọn trên K/V/Q
- ✅ `num_heads` và `dropout` hoàn toàn cấu hình được

---

## Cài đặt

**Yêu cầu:** Python ≥ 3.9 · PyTorch ≥ 2.0.0

```bash
# Cài từ source
git clone https://github.com/buiquangdat1710/titans-torch.git
cd titans-pytorch
pip install -e .
```

```bash
# Hoặc khi đã publish lên PyPI
pip install titans-torch
```

---

## Kiến trúc tổng quan

```
                        ┌─────────────────────────────────────┐
                        │        Neural Long-Term Memory       │
                        │                                      │
                        │   W_K, W_V, W_Q  (+ optional conv)  │
                        │           ↓                          │
                        │       Memory MLP                     │
                        │     M(k) ≈ v  (associative)         │
                        │           ↓                          │
                        │   Inner-loop gradient descent        │
                        │   S_t = η·S_{t-1} − θ·∇L            │
                        │   W_t = (1−α)·W_{t-1} + S_t         │
                        └─────────────────────────────────────┘

  ┌──────────┐    ┌──────────┐    ┌──────────┐
  │  MAC     │    │  MAG     │    │  MAL     │
  │          │    │          │    │          │
  │ persist  │    │ persist  │    │ persist  │
  │    ‖     │    │    +     │    │    +     │
  │ h_t=M(q) │    │  SWAttn  │    │   LMM    │
  │    ‖     │    │  ──┬──   │    │    ↓     │
  │ segment  │    │  gate    │    │  SWAttn  │
  │    ↓     │    │  ──┴──   │    │          │
  │  Attn    │    │   LMM    │    │          │
  │    ↓     │    │          │    │          │
  │  y⊗M(y)  │    │  out     │    │   out    │
  └──────────┘    └──────────┘    └──────────┘
  Memory as       Memory as       Memory as
  Context         Gate            Layer
```

---

## Sử dụng nhanh

```python
import torch
from titans-torch import TitansMAC, TitansMAG, TitansMAL, NeuralLongTermMemory

batch, seq_len, dim = 2, 16, 64
x = torch.randn(batch, seq_len, dim)

# Memory as a Context
model = TitansMAC(input_dim=dim, num_heads=4, dropout=0.1)
out = model(x, return_all=True)    # (2, 16, 64)

# Memory as a Gate
model = TitansMAG(input_dim=dim, num_heads=4, dropout=0.0)
out = model(x, return_all=True)    # (2, 16, 64)

# Memory as a Layer
model = TitansMAL(input_dim=dim, num_heads=4, dropout=0.1)
out = model(x, return_all=True)    # (2, 16, 64)

# Standalone LMM — test-time update
lmm = NeuralLongTermMemory(input_dim=dim)
out = lmm.test_time_update(x, return_all_outputs=True)   # (2, 16, 64)
```

---

## API Reference

### `NeuralLongTermMemory`

```python
NeuralLongTermMemory(
    input_dim: int,
    mem_dim: int = None,       # mặc định = input_dim
    num_layers: int = 2,
    hidden_dim: int = 32,
    use_conv: bool = True,
    conv_kernel: int = 3,
)
```

| Method | Mô tả |
|--------|-------|
| `forward_trainable(x, chunk_size=4, return_all_outputs=False)` | Inner-loop differentiable, dùng cho outer-loop training |
| `forward(x, return_all_outputs=False)` | Inference thuần túy, không cập nhật memory |
| `test_time_update(x, return_all_outputs=False)` | Cập nhật tuần tự tại test time, không giữ computation graph |

**Inner-loop update rule:**

```
S_t  =  η_t · S_{t-1}  −  θ_t · ∇L(W; k_t, v_t)
W_t  =  (1 − α_t) · W_{t-1}  +  S_t
```

Các hệ số α, θ, η được tạo ra từ dữ liệu (data-dependent) qua các MLP nhỏ.

---

### `PersistentMemory`

```python
PersistentMemory(
    num_tokens: int,
    dim: int,
    init_scale: float = 0.02,
)
```

Prepend `num_tokens` learnable token vào đầu sequence trước khi đưa vào attention.

```python
pm = PersistentMemory(num_tokens=8, dim=64)
x_extended = pm(x)   # (batch, 8 + seq_len, 64)
```

---

### `TitansMAC`

**Memory as a Context** — truy xuất thông tin lịch sử từ LMM, đưa vào context của attention.

```python
TitansMAC(
    input_dim: int,
    num_persistent: int = 8,
    mem_dim: int = None,
    num_memory_layers: int = 2,
    hidden_memory_dim: int = 32,
    chunk_size: int = 4,
    use_conv: bool = True,
    num_heads: int = 4,       # số attention head
    dropout: float = 0.1,     # dropout rate
)
```

**Forward flow (mỗi segment):**

```
q = W_Q(S^(t))
h_t = M_{t-1}*(q)                          # retrieve từ LMM
input = [persistent ‖ h_t ‖ S^(t)]
y_t = CausalAttention(input)
M_t = update(M_{t-1}, y_t)                 # forward_trainable
o_t = y_t ⊗ M_t*(y_t)                      # Hadamard gating
```

---

### `TitansMAG`

**Memory as a Gate** — kết hợp sliding-window attention và LMM qua learned gate.

```python
TitansMAG(
    input_dim: int,
    num_persistent: int = 8,
    mem_dim: int = None,
    num_memory_layers: int = 2,
    hidden_memory_dim: int = 32,
    window_size: int = 4,
    use_conv: bool = True,
    num_heads: int = 4,       # số attention head
    dropout: float = 0.1,     # dropout rate
)
```

**Forward flow:**

```
attn_out   = SlidingWindowAttention([persistent ‖ x])
memory_out = LMM.forward_trainable(x)
gate       = sigmoid([attn_out ‖ memory_out] · W)
out        = gate ⊙ memory_out + (1 − gate) ⊙ attn_out
```

---

### `TitansMAL`

**Memory as a Layer** — LMM xử lý trước, kết quả đưa vào sliding-window attention.

```python
TitansMAL(
    input_dim: int,
    num_persistent: int = 4,
    mem_dim: int = None,
    num_memory_layers: int = 2,
    hidden_memory_dim: int = 32,
    window_size: int = 4,
    use_conv: bool = True,
    num_heads: int = 4,       # số attention head
    dropout: float = 0.1,     # dropout rate
)
```

**Forward flow:**

```
x_persist = [persistent ‖ x]
LMM.forward_trainable(x_persist)           # update memory
out = SlidingWindowAttention(x_persist)    # attention trên full context
```

---

## Ví dụ training

```python
import torch
import torch.nn.functional as F
from titans-torch import TitansMAL

model = TitansMAL(
    input_dim=64,
    num_persistent=4,
    num_heads=4,
    dropout=0.1,
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(200):
    x      = torch.randn(4, 32, 64)
    target = torch.randn(4, 32, 64)

    optimizer.zero_grad()
    out  = model(x, return_all=True)
    loss = F.mse_loss(out, target)
    loss.backward()
    optimizer.step()

    if (step + 1) % 20 == 0:
        print(f"Step {step+1:3d}  loss={loss.item():.4f}")
```

**Test-time adaptation với LMM:**

```python
from titans-torch import NeuralLongTermMemory

lmm = NeuralLongTermMemory(input_dim=64)
lmm.eval()

x_new = torch.randn(2, 128, 64)
# Cập nhật memory tại test time — không cần gradient
with torch.no_grad():
    out = lmm.test_time_update(x_new, return_all_outputs=True)
```

---

## Cấu trúc thư mục

```
titans-torch/
├── __init__.py        # Public exports
├── memory.py          # NeuralLongTermMemory · PersistentMemory
├── mac.py          # TitansMAC 
├── mag.py          # TitansMAG 
├── mal.py          # TitansMAL 

setup.py
README.md
MANIFEST.in
CHANGELOG.txt
LICENSE.txt
```

---

## Tham khảo

```bibtex
@article{behrouz2025titans,
  title   = {Titans: Learning to Memorize at Test Time},
  author  = {Behrouz, Ali and Zhong, Peilin and Daliri, Majid},
  journal = {arXiv preprint arXiv:2501.00663},
  year    = {2025},
  url     = {https://arxiv.org/abs/2501.00663}
}
```

---

## License

Phát hành theo giấy phép **MIT** — xem [`LICENSE.txt`](LICENSE.txt) để biết thêm chi tiết.
