local-attention

You can install the local-attention package using pip.

$ pip install local-attention

To use local attention in a shared query/key space (similar to the Reformer architecture), set shared_qk = True. The module handles key normalization and token masking automatically.

import torch
from local_attention import LocalAttention

qk = torch.randn(2, 8, 2048, 64)
v  = torch.randn(2, 8, 2048, 64)

attn = LocalAttention(
    dim = 64,
    window_size = 512,
    shared_qk = True,
    causal = True
)

mask = torch.ones(2, 2048).bool()
out = attn(qk, qk, v, mask = mask) # (2, 8, 2048, 64)

The LocalAttention module implements local windowed attention. It can be used with standard query, key, and value tensors, or in a shared query/key space (Reformer-style).

import torch
from local_attention import LocalAttention

q = torch.randn(2, 8, 2048, 64)
k = torch.randn(2, 8, 2048, 64)
v = torch.randn(2, 8, 2048, 64)

attn = LocalAttention(
    dim = 64,                # dimension of each head
    window_size = 512,       # window size
    causal = True,           # auto-regressive or not
    look_backward = 1,       # each window looks at the window before
    look_forward = 0,        # for non-auto-regressive case, defaults to 1
    dropout = 0.1,           # post-attention dropout
    exact_windowsize = False # if True, each query sees max window_size keys in causal setting
)

mask = torch.ones(2, 2048).bool()
out = attn(q, k, v, mask = mask) # (2, 8, 2048, 64)

The LocalTransformer class provides a full transformer implementation utilizing local attention.

import torch
from local_attention import LocalTransformer

model = LocalTransformer(
    num_tokens = 256,
    dim = 512,
    depth = 6,
    max_seq_len = 8192,
    causal = True,
    local_attn_window_size = 256
).cuda()

x = torch.randint(0, 256, (1, 8192)).cuda()

logits = model(x) # (1, 8192, 256)

If you set autopad = True in LocalAttention, the module will automatically pad the query, key, value, and mask tensors, and then truncate the output to the appropriate size.

import torch
from local_attention import LocalAttention

q = torch.randn(8, 2057, 64)
k = torch.randn(8, 2057, 64)
v = torch.randn(8, 2057, 64)

attn = LocalAttention(
    window_size = 512,
    causal = True,
    autopad = True      # auto pads both inputs and mask, then truncates output appropriately
)

mask = torch.ones(1, 2057).bool()
out = attn(q, k, v, mask = mask) # (8, 2057, 64)

When initializing LocalAttention, you can use the following parameters:

• dim: Dimension of each head (required for relative positional encoding).
• window_size: The size of the local attention window.
• causal: Boolean indicating if the attention should be auto-regressive.
• look_backward: Number of previous windows each window looks at.
• look_forward: Number of subsequent windows each window looks at (defaults to 1 for non-auto-regressive cases).
• dropout: Post-attention dropout rate.
• exact_windowsize: If True, in causal settings, each query sees a maximum of window_size keys.
• shared_qk: Boolean to enable shared query/key space (Reformer style).
• autopad: Boolean to enable automatic padding of inputs and masks.