titans-unofficial 1.1.0
Unofficial PyTorch implementation of Titans: Learning to Memorize at Test Time
Requires Python: >=3.10

# Unofficial Implementation of Titans: Learning to Memorize at Test Time

[![PyPI version](https://img.shields.io/pypi/v/titans-unofficial.svg)](https://pypi.org/project/titans-unofficial/)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/release/python-3100/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

This is an unofficial PyTorch implementation of the paper ["Titans: Learning to Memorize at Test Time"](https://arxiv.org/abs/2501.00663) by Ali Behrouz, Peilin Zhong, and Vahab Mirrokni.

## Overview

Titans is a novel neural architecture that combines attention-based short-term memory with a neural long-term memory module. The architecture addresses the limitations of both recurrent models (which compress data into fixed-size memory) and attention mechanisms (which have quadratic complexity).

### Key Features

- **Neural Long-term Memory**: A module that learns to memorize historical context
- **Persistent Memory**: Learnable tokens that encode task-specific knowledge
- **Three Architectural Variants**:
  - MAC (Memory as Context): Uses memory as context for attention
  - MAG (Memory as Gate): Combines memory with core branch using gating
  - MAL (Memory as Layer): Integrates memory as a separate layer

## Installation

### From PyPI (Recommended)

```bash
# Install core package (minimal dependencies)
pip install titans-unofficial

# Install with examples dependencies (for running example scripts)
pip install "titans-unofficial[examples]"

# Install with development dependencies (for contributing)
pip install "titans-unofficial[dev]"
```

### From Source (Development)

```bash
# Clone the repository
git clone https://github.com/Shehryar718/titans-unofficial.git
cd titans-unofficial

# Install in development mode with all extras
pip install -e ".[dev,examples]"
```

## Project Structure

```
titans-unofficial/
├── titans/
│   ├── __init__.py
│   ├── models/
│   │   ├── titans_base.py    # Base class for all variants
│   │   ├── titans_mac.py     # Memory as Context implementation
│   │   ├── titans_mag.py     # Memory as Gate implementation
│   │   └── titans_mal.py     # Memory as Layer implementation
│   └── utils/
│       ├── memory.py         # Neural Memory Module
│       ├── attention.py      # Attention mechanisms
│       └── persistent_memory.py  # Persistent Memory implementation
├── examples/
│   ├── text_classification.py  # Text classification example
│   ├── language_modeling.py    # Language modeling example
│   └── fine_tuning.py         # Fine-tuning example
├── pytests/
│   └── test_memory.py         # Tests for memory module
├── requirements.txt
├── LICENSE
└── README.md
```

## Usage

### Text Classification

```python
from titans import TitansMAC, TitansMAG, TitansMAL
from examples.text_classification import TitansForClassification

# Initialize model
model = TitansForClassification(
    vocab_size=30000,
    d_model=128,
    n_layers=2,
    n_heads=4,
    num_classes=2,
    memory_depth=2,
    persistent_tokens=8,
    window_size=16,
    model_type="mal"  # Choose from: "mac", "mag", "mal"
)
```

```bash
# Train and evaluate
python examples/text_classification.py
```

### Language Modeling

```python
from titans import TitansMAC, TitansMAG, TitansMAL
from examples.language_modeling import TitansForLanguageModeling

# Initialize model
model = TitansForLanguageModeling(
    vocab_size=30000,
    d_model=128,
    n_layers=2,
    n_heads=4,
    memory_depth=2,
    persistent_tokens=16,
    window_size=128,
    model_type="mac"
)
```

```bash
# Train and generate text
python examples/fine_tuning.py
```

## Architecture Details

### Neural Memory Module

The neural memory module consists of:

- Key/Value/Query projections for memory access
- Multi-layer perceptron for memory processing
- Momentum-based update mechanism with configurable parameters
- Weight decay for forgetting mechanism
- Gradient scaling for numerical stability

### Variants

1. **MAC (Memory as Context)**

   - Memory output serves as additional context
   - Efficient for tasks requiring long-range dependencies
   - Parallel processing with chunked computation
   - Configurable chunk size and parallel processing

2. **MAG (Memory as Gate)**

   - Gating mechanism to combine memory with core processing
   - Adaptive balance between short and long-term memory
   - Enhanced numerical stability
   - Improved gradient flow through gating

3. **MAL (Memory as Layer)**
   - Memory integrated as a separate layer
   - Direct memory access at each layer
   - Sliding window attention for efficiency
   - Layer-wise memory updates

## Example Tasks

The repository includes implementations for:

- Text Classification (Binary and multi-class)
- Language Modeling with test-time adaptation
- Fine-tuning with early stopping

Each example demonstrates different aspects of the Titans architecture:

- Memory reset between epochs for fresh adaptation
- Efficient batch processing with dynamic batching
- Gradient scaling for numerical stability
- Early stopping and model checkpointing
- Proper memory state management

## Testing

```bash
# Run all tests
pytest pytests/

# Run specific test file
pytest pytests/test_memory.py
```

## Citation

This repository provides an **unofficial implementation** of the Titans architecture.  
If you reference this work, please **cite the original paper**:

```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}
}
```

## License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.