DOCSEARCHbyAlgolia
Submit Docs

flash-pytorch Documentation

repository·main·Indexed 18 days ago

https://github.com/lucidrains/flash-pytorch

A PyTorch implementation of the Transformer variant from the paper 'Transformer Quality in Linear Time'. It features the Gated Attention Unit (GAU), grouped linear attention, and the full FLASHTransformer model, which includes scaled sinusoidal absolute positional embeddings, T5 relative positional bias, and Rotary embeddings (RoPE).

Tokens
1.1K
Snippets
4
Records
4
Agent score
13%

What's inside flash-pytorch

  2. Use the Gated Attention Unit (GAU)

    main

    The GAU class implements the Gated Attention Unit, which replaces multi-headed attention with a single head using a relu squared activation (or Laplace attention for better stability).

    import torch
from flash_pytorch import GAU

gau = GAU(
    dim = 512,
    query_key_dim = 128,     # query / key dimension
    causal = True,           # autoregressive or not
    expansion_factor = 2,    # hidden dimension = dim * expansion_factor
    laplace_attn_fn = True   # use Laplace attention for better stability
)

x = torch.randn(1, 1024, 512)
out = gau(x) # (1, 1024, 512)

  3. Use the FLASH module

    main

    The FLASH module combines the quadratic Gated Attention Unit (GAU) with grouped linear attention to overcome issues with autoregressive linear attention. Sequences are automatically padded to the nearest group_size.

    import torch
from flash_pytorch import FLASH

flash = FLASH(
    dim = 512,
    group_size = 256,             # group size
    causal = True,                # autoregressive or not
    query_key_dim = 128,          # query / key dimension
    expansion_factor = 2.,        # hidden dimension = dim * expansion_factor
    laplace_attn_fn = True        # use Laplace attention for better stability
)

x = torch.randn(1, 1111, 512)     # sequence will be auto-padded to nearest group size
out = flash(x) # (1, 1111, 512)

  4. Use the full FLASHTransformer

    main

    The FLASHTransformer is the complete model implementation described in the paper. It includes:

    • Scaled sinusoidal absolute positional embeddings.
    • T5 relative positional bias for GAU quadratic attention.
    • Rotary embeddings (RoPE) for both GAU and linear attention.
    • Optional shift_tokens for improved convergence.
    • Support for both scalenorm and layernorm.
    import torch
from flash_pytorch import FLASHTransformer

model = FLASHTransformer(
    num_tokens = 20000,          # number of tokens
    dim = 512,                   # model dimension
    depth = 12,                  # depth
    causal = True,               # autoregressive or not
    group_size = 256,            # size of the groups
    query_key_dim = 128,         # dimension of queries / keys
    expansion_factor = 2.,       # hidden dimension = dim * expansion_factor
    norm_type = 'scalenorm',     # 'scalenorm' (faster training) or 'layernorm'
    shift_tokens = True          # shifts half of the feature space forward to improve convergence
)

x = torch.randint(0, 20000, (1, 1024))
logits = model(x) # (1, 1024, 20000)