Metadata-Version: 2.4
Name: titans-torch
Version: 0.1.0
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"
================================================================================

titans-torch is a Python library implementing the Titans architecture,
a family of neural network models that combine short-term attention with a
differentiable long-term memory module (Neural Long-Term Memory, LMM).
The memory can be updated at test time via inner-loop gradient descent,
allowing the model to adapt to new sequences without retraining.


--------------------------------------------------------------------------------
INSTALLATION
--------------------------------------------------------------------------------

Install from source:

    git clone https://github.com/buiquangdat1710/titans-torch.git
    cd titans-torch
    pip install -e .

Or install directly via pip (once published):

    pip install titans-torch

Requirements:
    - Python >= 3.9
    - PyTorch >= 2.0.0


--------------------------------------------------------------------------------
QUICK START
--------------------------------------------------------------------------------

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

    batch, seq_len, dim = 2, 16, 64

    # --- Memory as a Context (MAC) ---
    model = TitansMAC(
        input_dim=dim,
        num_persistent=8,
        mem_dim=dim,
        chunk_size=4,
        num_heads=4,
        dropout=0.1,
    )
    x = torch.randn(batch, seq_len, dim)
    out = model(x, return_all=True)   # (2, 16, 64)

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

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

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


--------------------------------------------------------------------------------
MODULE OVERVIEW
--------------------------------------------------------------------------------

titans-torch/
  __init__.py       — Public exports
  memory.py         — Core building blocks
                        · NeuralLongTermMemory  (LMM)
                        · PersistentMemory
  mac.py         — TitansMAC  (Memory as a Context)
  mag.py         — TitansMAG  (Memory as a Gate)
  mal.py         — TitansMAL  (Memory as a Layer)



--------------------------------------------------------------------------------
CLASS REFERENCE
--------------------------------------------------------------------------------

NeuralLongTermMemory(
    input_dim,
    mem_dim=None,
    num_layers=2,
    hidden_dim=32,
    use_conv=True,
    conv_kernel=3,
)
    Methods:
        forward_trainable(x, chunk_size=4, return_all_outputs=False)
            Differentiable inner-loop update. Use during outer-loop training.

        forward(x, return_all_outputs=False)
            Pure inference with no memory update.

        test_time_update(x, return_all_outputs=False)
            Sequential per-token update at test time. No gradient graph kept.

----

PersistentMemory(num_tokens, dim, init_scale=0.02)
    Prepends `num_tokens` learnable tokens to the input sequence.

----

TitansMAC(
    input_dim,
    num_persistent=8,
    mem_dim=None,
    num_memory_layers=2,
    hidden_memory_dim=32,
    chunk_size=4,
    use_conv=True,
    num_heads=4,       # <-- configurable
    dropout=0.1,       # <-- configurable
)

TitansMAG(
    input_dim,
    num_persistent=8,
    mem_dim=None,
    num_memory_layers=2,
    hidden_memory_dim=32,
    window_size=4,
    use_conv=True,
    num_heads=4,       # <-- configurable
    dropout=0.1,       # <-- configurable
)

TitansMAL(
    input_dim,
    num_persistent=4,
    mem_dim=None,
    num_memory_layers=2,
    hidden_memory_dim=32,
    window_size=4,
    use_conv=True,
    num_heads=4,       # <-- configurable
    dropout=0.1,       # <-- configurable
)

    All three models expose:
        forward(x, return_all=False)
            x            : (batch, seq_len, input_dim)
            return_all   : True  → (batch, seq_len, input_dim)
                           False → (batch, input_dim)  [last token only]


--------------------------------------------------------------------------------
TRAINING EXAMPLE
--------------------------------------------------------------------------------

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

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

    for step in range(100):
        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) % 10 == 0:
            print(f"Step {step+1:3d}  loss={loss.item():.4f}")


--------------------------------------------------------------------------------
REFERENCE
--------------------------------------------------------------------------------

    Ali Behrouz, Peilin Zhong, Majid Daliri.
    "Titans: Learning to Memorize at Test Time."
    arXiv:2501.00663, 2025.
    https://arxiv.org/abs/2501.00663


--------------------------------------------------------------------------------
LICENSE
--------------------------------------------------------------------------------

MIT License — see LICENSE.txt for details.
