trim-transformer 0.1.1

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, but with an attention function of the form Attn(Q,K,V) = QK^TV, which we call multi-linear attention. This implementation has time complexity O(nd^2), where n is the sequence length and d is the model dimension. Since the time complexity is linear in the sequence length, this implementation is well suited for high sequence length tasks. Attention in this form has shown success in operator learning tasks, see Choose a Transformer: Fourier or Galerkin.

This implementation is particularly relevent for training physics models where high sequence length can come from large grid sizes, long time periods, or both.

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

---

Benchmarks

Below are some benchmark plots demonstrating model performance and resource usage on the Navier-Stokes dataset from https://arxiv.org/abs/2010.08895:

The Trim Transformer achives more than 90% reduction in memory usage compared to a standard Pytorch transformer using softmax attention and 3.5x faster time per epoch while maintaining very similar validation loss. As grid size and sequence length increase these gains become even more drastic.

---

Quickstart

import torch
from trim_transformer.transformer_layers import TrimTransformerEncoderLayer, TrimTransformerEncoder

layer = TrimTransformerEncoderLayer(d_model=EMBED_DIM, nhead=NUM_HEADS, batch_first=True)
model = TrimTransformerEncoder(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 trim_vs_softmax.ipynb 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). To illustrate further, consider the one-dimensional mask [2, 0, 1]. This corresponds to the following two-dimensional mask, where, following the PyTorch convention, True indicates that an element is masked.

	Key 0	Key 1	Key 2
Query 0	False	False	False
Query 1	False	True	True
Query 2	False	False	True

Key-value caching

Inference with key-value caching 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.