titans-torch 0.1.0

PyTorch implementation of Titans: Learning to Memorize at Test Time

================================================================================

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.