DOCSEARCHbyAlgolia
Submit Docs

Nyström Attention

repository·main·Indexed 40 days ago

https://github.com/lucidrains/nystrom-attention

An efficient implementation of Nyström Self-attention based on the Nyströmformer paper. It approximates self-attention using landmarks to reduce computational complexity for long sequences, providing the NystromAttention class and a Nystromformer stack for multiple layers of attention.

Tokens
819
Snippets
4
Records
4
Agent score
0%

What's inside nystrom-attention

  3. Use NystromAttention for self-attention

    main

    The NystromAttention class implements Nyström Self-attention as described in the Nyströmformer paper. It allows for efficient attention approximation using landmarks.

    Parameters:

    • dim: The dimension of the input.
    • dim_head: The dimension of each attention head.
    • heads: The number of attention heads.
    • num_landmarks: The number of landmarks used for approximation.
    • pinv_iterations: The number of Moore-Penrose iterations for approximating the pseudo-inverse. The paper recommends 6.
    • residual: Whether to perform an extra residual with the value. Turning this on can lead to faster convergence.
    import torch
from nystrom_attention import NystromAttention

attn = NystromAttention(
    dim = 512,
    dim_head = 64,
    heads = 8,
    num_landmarks = 256,    # number of landmarks
    pinv_iterations = 6,    # number of moore-penrose iterations for approximating pinverse. 6 was recommended by the paper
    residual = True         # whether to do an extra residual with the value or not. supposedly faster convergence if turned on
)

x = torch.randn(1, 16384, 512)
mask = torch.ones(1, 16384).bool()

attn(x, mask = mask) # (1, 16384, 512)

  4. Use Nystromformer for multiple layers of attention

    main

    The Nystromformer class provides a stack of Nyström attention layers, useful for building full transformer-like architectures.

    Parameters:

    • dim: The dimension of the input.
    • dim_head: The dimension of each attention head.
    • heads: The number of attention heads.
    • depth: The number of layers (depth) of the Nyströmformer.
    • num_landmarks: The number of landmarks used for approximation.
    • pinv_iterations: The number of Moore-Penrose iterations for approximating the pseudo-inverse.
    import torch
from nystrom_attention import Nystromformer

model = Nystromformer(
    dim = 512,
    dim_head = 64,
    heads = 8,
    depth = 6,
    num_landmarks = 256,
    pinv_iterations = 6
)

x = torch.randn(1, 16384, 512)
mask = torch.ones(1, 16384).bool()

model(x, mask = mask) # (1, 16384, 512)