trim-transformer 0.1.0

A linear-attention transformer implementation with KV caching.

Trim Transformer

trim-transformer is a lightweight PyPI package that replicates the familiar interface of torch.nn.TransformerEncoder. The attention kernel has the form Attn(Q,K,V) = QK^TV. This implementation has complexity O(nd^2), where n is the sequence length and d is the model dimension.

Additionally, this implementation supports key-value caching for inference that is also linear in the number of tokens generated. Finally, this implementation supports custom weight initialization functions for the query, key, and value projection matrices, and custom normalization layers for the query, key, and value activations.

---

Installation

The package is published on PyPI and only depends on PyTorch:

pip install trim-transformer

Alternatively, install the latest commit from GitHub:

pip install git+https://github.com/emanuel-nuclearsoftware/trim-transformer.git

---

Quickstart

import torch
from trim_transformer.transformer_layers import CumulativeTransformerKV

layer = CumulativeTransformerEncoderLayerKV(d_model=EMBED_DIM, nhead=NUM_HEADS, batch_first=True)
model = CumulativeTransformerEncoderKV(layer, num_layers=NUM_LAYERS)

x = torch.randn(8, 2048, 512)  # (batch, seq_len, dim)

# Standard forward pass (causal mask optional)
out = model(x, is_causal=True)  # (batch, seq_len, dim)

See tutorial_notebook.ipynb for demonstrations of each feature. And see demo_vs_baseline for an example of a complete pipeline and a comparison to a PyTorch baseline.

Masking

A significant departure from PyTorch syntax is the structure of the mask. Multi-linear attention with arbitrary boolean masks cannot be computed in time linear in the sequence length. Instead, this package supports masks such that the i-th query attends to all keys up to index m_i. Such masks can be specified by an integer array of length n, with values in [0, n-1], where n is the sequence length.

For example, a causal mask of length n is given by torch.arange(n). And the one dimensional mask [2, 0, 1] corresponds to the two dimensional mask [[False, False, False], [False, True, True], [False, False, True]], following the PyTorch convention that True means to set that element of the attention matrix to 0.

Key-value caching

Inference with key value can be performed with a simple loop.

def generate(model, initial_tokens, num_new_tokens=5):
    """Autoregressive generation with caching"""
    model.eval()  # optional
    generated_sequence = []
    new_token = initial_tokens.clone()

    for _ in range(num_new_tokens):
        # Must set use_kv_cache=True and update_kv_cache=True to use key-value caching.
        # use_kv_cache=True means that the model will key-value cache that is already stored.
        # update_kv_cache=True means that the model will update the key-value cache with the
        # keys and values from the new token.
        output = model(new_token, is_causal=True, use_kv_cache=True, update_kv_cache=True)
        generated_sequence.append(output)

    # Clear the key-value cache after generation. If you don't clear the cache, then if
    # use_kv_cache=True in the future, the model will continue to use the key-value cache
    # on future forward passes.
    model.clear_kv_cache()
    return generated_sequence

---

License

trim-transformer is released under the MIT License. See LICENSE for the full text.