DOCSEARCHbyAlgolia
Submit Docs

x-transformers

repository·main·Indexed 26 days ago

https://github.com/lucidrains/x-transformers

A concise, feature-rich library providing various transformer architectures, including encoder-only, decoder-only, encoder-decoder, and vision transformers. It supports advanced features such as Flash Attention, Multi-Query and Grouped-Query Attention, Transformer-XL recurrence, and various normalization methods like RMSNorm and ScaleNorm. The library includes specialized wrappers like ViTransformerWrapper for image classification and tools for implementing image-to-caption and PaLI-style language-vision models.

Tokens
29.3K
Snippets
53
Records
128
Agent score
91%

What's inside x-transformers

  1. Implement Transformer-XL Recurrence

    main

    To implement Transformer-XL recurrence:

    1. In TransformerWrapper, set max_mem_len (e.g., 2048).
    2. In the Decoder or Encoder, set rel_pos_bias = True or rotary_pos_emb = True.
    3. Use return_mems = True during the forward pass to get memories.
    4. Pass the retrieved memories back into the next iteration using the mems keyword.
    import torch
from x_transformers import TransformerWrapper, Decoder

model_xl = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 512,
    max_mem_len = 2048,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        rel_pos_bias = True
    )
)

seg1 = torch.randint(0, 20000, (1, 512))
seg2 = torch.randint(0, 20000, (1, 512))

logits1, mems1  = model_xl(seg1, return_mems = True)
logits2, mems2  = model_xl(seg2, mems = mems1, return_mems = True)

  3. Use Explicit Sparse Attention (Top-K)

    main

    To sparsify attention by only keeping the top $k$ values before the softmax, use attn_sparse_topk. You can also use attn_sparse_topk_straight_through = True to allow gradients to flow through the original values.

    Recommended value for $k$ is typically 8.

    import torch
from x_transformers import TransformerWrapper, Decoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        attn_sparse_topk = 8,
        attn_sparse_topk_straight_through = True
    )
)

  4. Use Memory Tokens (Register Tokens)

    main

    You can add learned tokens (similar to CLS tokens) that are passed through the attention layers alongside input tokens by setting num_memory_tokens in the TransformerWrapper. This is compatible with both encoder and decoder training and can help alleviate attention outliers.

    import torch
from x_transformers import TransformerWrapper, Decoder, Encoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    num_memory_tokens = 20, # 20 memory tokens
    attn_layers = Encoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

  5. Configure Dropouts in Transformer models

    main

    You can apply various dropout layers to control regularization:

    • emb_dropout: Dropout after the embedding layer (set in TransformerWrapper).
    • layer_dropout: Stochastic depth (dropout entire layer) in Encoder or Decoder layers.
    • attn_dropout: Dropout applied post-attention in Encoder or Decoder layers.
    • ff_dropout: Dropout in the feedforward network of Encoder or Decoder layers.
    import torch
from x_transformers import TransformerWrapper, Decoder, Encoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    emb_dropout = 0.1,         # dropout after embedding
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        layer_dropout = 0.1,   # stochastic depth - dropout entire layer
        attn_dropout = 0.1,    # dropout post-attention
        ff_dropout = 0.1       # feedforward dropout
    )
)

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

  6. Configure Feedforward (FF) variants

    main

    You can modify the feedforward layers using the following options in the attention layer:

    • GLU (Gated Linear Units): Set ff_glu = True to use gated feedforwards (e.g., GELU gating).
    • SwiGLU (PaLM style): Set both ff_swish = True and ff_glu = True to use the Swish GLU variant.
    • No Bias: Set ff_no_bias = True to remove biases from the feedforward layers, which can increase throughput.
    • ReLU²: Set ff_relu_squared = True to use the ReLU Squared activation (note: if using GLU, GELU is generally preferred).
    # Example: Using SwiGLU
model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        ff_swish = True,
        ff_glu = True
    )
)

  7. Augment self-attention with persistent memory (KV)

    main

    To add learned memory key/values prior to the attention mechanism, use the attn_num_mem_kv parameter in your attention layer. This allows the model to utilize a set number of learned memory slots.

    from x_transformers import Decoder, Encoder

enc = Encoder(
    dim = 512,
    depth = 6,
    heads = 8,
    attn_num_mem_kv = 16 # 16 memory key / values
)

  8. Configure alternative normalization methods

    main

    The library supports several normalization variants to improve convergence and stability:

    • ScaleNorm: Set use_scalenorm = True in the attention layer.
    • L2 Normalized Embeddings: Set l2norm_embed = True in TransformerWrapper. This also applies a small initialization. It is recommended to use this instead of post_emb_norm.
    • Post-Embedding LayerNorm: Set post_emb_norm = True in TransformerWrapper to apply layer normalization to the sum of token and positional embeddings.
    • RMSNorm: Set use_rmsnorm = True in the attention layer.
    • Simple RMSNorm: Set use_simple_rmsnorm = True in the attention layer for a version without the learned multiplicative gamma.
    # Example: Using ScaleNorm
model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        use_scalenorm = True
    )
)

# Example: Using L2 Normalized Embeddings
model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    l2norm_embed = True,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

# Example: Using RMSNorm
model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        use_rmsnorm = True
    )
)

  9. Enable Flash Attention for speed and memory efficiency

    main

    You can use Flash Attention to achieve significant memory savings and increased speed, provided you have PyTorch 2.0+ installed. This is done by setting attn_flash = True within your attention layer configuration (e.g., in Decoder or Encoder).

    Note: Avoid using Flash Attention if your model requires operating directly on the attention matrix (e.g., dynamic positional bias, talking heads, or residual attention).

    import torch
from x_transformers import TransformerWrapper, Decoder, Encoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        attn_flash = True # just set this to True if you have pytorch 2.0 installed
    )
)

  10. Implement Image-to-Caption models

    main

    Create an image-to-caption pipeline by using a ViTransformerWrapper as an encoder and a TransformerWrapper with cross_attend = True as a decoder. Use return_embeddings = True on the encoder to pass visual features to the decoder via the context argument.

    import torch
from x_transformers import ViTransformerWrapper, TransformerWrapper, Encoder, Decoder

encoder = ViTransformerWrapper(
    image_size = 256,
    patch_size = 32,
    attn_layers = Encoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

decoder = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8,
        cross_attend = True
    )
)

img = torch.randn(1, 3, 256, 256)
caption = torch.randint(0, 20000, (1, 1024))

encoded = encoder(img, return_embeddings = True)
decoder(caption, context = encoded) # (1, 1024, 20000)

  11. Implement PaLI-style language-vision models

    main

    To implement a PaLI-style model, compose a ViTransformerWrapper (vision encoder) and an XTransformer (encoder-decoder). During training, pass the image embeddings to the XTransformer using the src_prepend_embeds argument to prepend them to the encoder text embeddings before attention.

    import torch
from x_transformers import ViTransformerWrapper, XTransformer, Encoder

vit = ViTransformerWrapper(
    image_size = 256,
    patch_size = 32,
    attn_layers = Encoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

pali = XTransformer(
    dim = 512,
    enc_num_tokens = 256,
    enc_depth = 6,
    enc_heads = 8,
    enc_max_seq_len = 1024,
    dec_num_tokens = 256,
    dec_depth = 6,
    dec_heads = 8,
    dec_max_seq_len = 1024
)

img = torch.randn(1, 3, 256, 256)
prompt = torch.randint(0, 256, (1, 1024))
prompt_mask = torch.ones(1, 1024).bool()
output_text = torch.randint(0, 256, (1, 1024))

img_embeds = vit(
    img,
    return_embeddings = True
)

loss = pali(
    prompt,
    output_text,
    mask = prompt_mask,
    src_prepend_embeds = img_embeds             # will preprend image embeddings to encoder text embeddings before attention
)

loss.backward()

  12. Use TransformerWrapper with Encoder for BERT-like models

    main

    Implement an encoder-only (BERT-like) model by wrapping an Encoder layer inside a TransformerWrapper. You can pass a mask to the model call to control attention.

    import torch
from x_transformers import TransformerWrapper, Encoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Encoder(
        dim = 512,
        depth = 12,
        heads = 8
    )
).cuda()

x = torch.randint(0, 256, (1, 1024)).cuda()
mask = torch.ones_like(x).bool()

model(x, mask = mask) # (1, 1024, 20000)