DOCSEARCHbyAlgolia
Submit Docs

Linear Attention Transformer

repository·master·Indexed 21 days ago

https://github.com/lucidrains/linear-attention-transformer

A library providing efficient Transformer architectures that combine local (QKᵀ)V attention with global Q(KᵀV) attention for linear scaling with sequence length. It includes LinearAttentionTransformerLM for language modeling, LinearAttentionTransformer for general tasks, and ImageLinearAttention for image processing. Supports features such as reversible nets, GLU variants, Linformer settings, and encoder-decoder architectures.

Tokens
2.2K
Snippets
7
Records
7
Agent score
24%

What's inside linear-attention-transformer

  1. Use Linformer settings for linear complexity

    master

    Linformer is a variant of attention with linear complexity that works with non-autoregressive models of a fixed sequence length. You can enable it by passing LinformerSettings to the linformer_settings argument of LinearAttentionTransformerLM.

    To use Linformer for the contextual attention layer (when the context has a fixed sequence length), use LinformerContextSettings and pass it to context_linformer_settings in a decoder with receives_context = True.

    from linear_attention_transformer import LinearAttentionTransformerLM, LinformerSettings

# For standard Linformer attention
settings = LinformerSettings(k = 256)
enc = LinearAttentionTransformerLM(
    num_tokens = 20000,
    dim = 512,
    heads = 8,
    depth = 6,
    max_seq_len = 4096,
    linformer_settings = settings
).cuda()

# For Linformer contextual attention
from linear_attention_transformer import LinformerContextSettings

context_settings = LinformerContextSettings(
  seq_len = 2048,
  k = 256
)

dec = LinearAttentionTransformerLM(
    num_tokens = 20000,
    dim = 512,
    heads = 8,
    depth = 6,
    max_seq_len = 4096,
    causal = True,
    context_linformer_settings = context_settings,
    receives_context = True
).cuda()

  3. Train Linear Attention Transformer with Deepspeed on Enwik8

    master

    To train the Linear Attention Transformer using Microsoft's Deepspeed framework on the Enwik8 dataset, you must first install Deepspeed following the official instructions from the DeepSpeed repository. Once installed, you can initiate training by running the train.py script with the --deepspeed flag and providing a configuration file via --deepspeed_config.

    deepspeed train.py --deepspeed --deepspeed_config ds_config.json

  4. Implement Encoder-Decoder architecture

    master

    You can build an encoder-decoder setup using LinearAttentionTransformerLM.

    • The Encoder should typically have receives_context = False (default) and can use return_embeddings = True to pass context to the decoder.
    • The Decoder should have causal = True, receives_context = True, and accepts a context tensor produced by the encoder.
    • Both components support input_mask and context_mask for handling padding.
    import torch
from linear_attention_transformer import LinearAttentionTransformerLM

enc = LinearAttentionTransformerLM(
    num_tokens = 20000,
    dim = 512,
    heads = 8,
    depth = 6,
    max_seq_len = 4096,
    reversible = True,
    n_local_attn_heads = 4,
    return_embeddings = True
).cuda()

dec = LinearAttentionTransformerLM(
    num_tokens = 20000,
    dim = 512,
    heads = 8,
    depth = 6,
    causal = True,
    max_seq_len = 4096,
    reversible = True,
    receives_context = True,
    n_local_attn_heads = 4
).cuda()

src = torch.randint(0, 20000, (1, 4096)).cuda()
src_mask = torch.ones_like(src).bool().cuda()

tgt = torch.randint(0, 20000, (1, 4096)).cuda()
tgt_mask = torch.ones_like(tgt).bool().cuda()

context = enc(src, input_mask = src_mask)
logits = dec(tgt, context = context, input_mask = tgt_mask, context_mask = src_mask)

  5. Use ImageLinearAttention for efficient image processing

    master

    The ImageLinearAttention module provides an efficient implementation of linear attention specifically for image data. It expects input in the shape (batch, channels, height, width).

    import torch
from linear_attention_transformer.images import ImageLinearAttention

attn = ImageLinearAttention(
  chan = 32,
  heads = 8,
  key_dim = 64       # can be decreased to 32 for more memory savings
)

img = torch.randn(1, 32, 256, 256)
output = attn(img) # (1, 32, 256, 256)

  6. Use LinearAttentionTransformerLM for language modeling

    master

    The LinearAttentionTransformerLM class is designed for language modeling. It supports causal (auto-regressive) modeling and includes several advanced features like reversible nets, GLU variants, and blindspot settings for memory efficiency.

    Key parameters:

    • num_tokens: Vocabulary size.
    • dim: Embedding dimension.
    • heads: Number of attention heads.
    • max_seq_len: Maximum sequence length.
    • causal: Set to True for auto-regressive modeling.
    • blindspot_size: Provides a blindspot for the $Q(K^TV)$ attention in causal mode to save memory. Should be paired with a local_attn_window_size at least as large as this value.
    • n_local_attn_heads: Number of local attention heads (can be a tuple for depth-specific counts).
    • local_attn_window_size: Receptive field of the local attention.
    • reversible: Enables reversible networks (from Reformer).
    • ff_glu: Enables GLU variant for feedforward layers.
    import torch
from linear_attention_transformer import LinearAttentionTransformerLM

model = LinearAttentionTransformerLM(
    num_tokens = 20000,
    dim = 512,
    heads = 8,
    depth = 1,
    max_seq_len = 8192,
    causal = True,
    ff_dropout = 0.1,
    attn_layer_dropout = 0.1,
    attn_dropout = 0.1,
    emb_dim = 128,
    dim_head = 128,
    blindspot_size = 64,
    n_local_attn_heads = 4,
    local_attn_window_size = 128,
    reversible = True,
    ff_chunks = 2,
    ff_glu = True,
    attend_axially = False,
    shift_tokens = True
).cuda()

x = torch.randint(0, 20000, (1, 8192)).cuda()
output = model(x) # (1, 8192, 512)

  7. Use LinearAttentionTransformer for general transformer tasks

    master

    For non-causal transformer tasks, use the LinearAttentionTransformer class. This is suitable for tasks where the entire sequence is available at once.

    import torch
from linear_attention_transformer import LinearAttentionTransformer

model = LinearAttentionTransformer(
    dim = 512,
    heads = 8,
    depth = 1,
    max_seq_len = 8192,
    n_local_attn_heads = 4
).cuda()

x = torch.randn(1, 8192, 512).cuda()
output = model(x) # (1, 8192, 512)