Longformer › Implementing Longformer Attention
Language AI HandbookAll chapters Longformer › Implementing Longformer AttentionOne section · Sections in this chapter8 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
With the mathematical foundation in place, let's translate these ideas into working code. Building the implementation from scratch solidifies understanding and reveals the practical choices that make Longformer work. We'll start with the simpler sliding window mechanism, then layer on global attention support.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class LongformerSlidingWindowAttention(nn.Module):
"""
Sliding window attention for Longformer.
Each position attends only to tokens within a fixed window.
"""
def __init__(self, embed_dim, num_heads, window_size):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.window_size = window_size
assert embed_dim % num_heads == 0, (
"embed_dim must be divisible by num_heads"
)
# Projections for local attention
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x, attention_mask=None):
"""
Args:
x: Input tensor of shape (batch, seq_len, embed_dim)
attention_mask: Optional mask for padding tokens
Returns:
Output tensor of shape (batch, seq_len, embed_dim)
"""
batch_size, seq_len, _ = x.shape
# Project queries, keys, and values
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# Reshape for multi-head attention
q = q.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
k = k.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
v = v.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
# For simplicity, we'll use a loop-based implementation
# Production code would use more efficient sparse operations
half_window = self.window_size // 2
scale = 1.0 / math.sqrt(self.head_dim)
outputs = []
for i in range(seq_len):
start = max(0, i - half_window)
end = min(seq_len, i + half_window + 1)
# Get local keys and values
local_k = k[:, :, start:end, :] # (batch, heads, window, head_dim)
local_v = v[:, :, start:end, :]
# Compute attention scores
query = q[:, :, i : i + 1, :] # (batch, heads, 1, head_dim)
scores = torch.matmul(query, local_k.transpose(-2, -1)) * scale
# Apply softmax and compute weighted sum
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, local_v)
outputs.append(output)
# Concatenate outputs
output = torch.cat(outputs, dim=2) # (batch, heads, seq_len, head_dim)
output = (
output.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.embed_dim)
)
return self.out_proj(output)This implementation demonstrates the core sliding window mechanism. Notice how the loop over positions extracts only the local window of keys and values for each query. In production code, this loop would be replaced with efficient sparse matrix operations, but the conceptual structure remains the same: each position computes attention only over its local window, keeping memory usage bounded regardless of sequence length.
The key insight from the code is that the window size half_window determines how far each token can "see." Positions near the edges of the sequence have smaller effective windows (they can't attend to positions that don't exist), but this is handled gracefully by the max and min bounds.
Now let's extend this to the full Longformer attention with global token support:
class LongformerAttention(nn.Module):
"""
Full Longformer attention with both sliding window and global attention.
"""
def __init__(self, embed_dim, num_heads, window_size):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.window_size = window_size
# Local attention projections
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
# Separate projections for global attention
self.q_proj_global = nn.Linear(embed_dim, embed_dim)
self.k_proj_global = nn.Linear(embed_dim, embed_dim)
self.v_proj_global = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x, global_attention_mask=None):
"""
Args:
x: Input tensor of shape (batch, seq_len, embed_dim)
global_attention_mask: Boolean mask where True indicates global attention
Shape: (batch, seq_len)
Returns:
Output tensor of shape (batch, seq_len, embed_dim)
"""
batch_size, seq_len, _ = x.shape
# Compute local projections
q_local = self.q_proj(x)
k_local = self.k_proj(x)
v_local = self.v_proj(x)
# Compute global projections
q_global = self.q_proj_global(x)
k_global = self.k_proj_global(x)
v_global = self.v_proj_global(x)
# Reshape for multi-head attention
def reshape(t):
return t.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
q_local = reshape(q_local)
k_local = reshape(k_local)
v_local = reshape(v_local)
q_global = reshape(q_global)
k_global = reshape(k_global)
v_global = reshape(v_global)
scale = 1.0 / math.sqrt(self.head_dim)
half_window = self.window_size // 2
# Determine which positions have global attention
if global_attention_mask is None:
global_positions = []
else:
global_positions = (
global_attention_mask[0].nonzero(as_tuple=True)[0].tolist()
)
outputs = torch.zeros_like(q_local)
for i in range(seq_len):
is_global = i in global_positions
if is_global:
# Global token: attend to entire sequence using global projections
query = q_global[:, :, i : i + 1, :]
scores = torch.matmul(query, k_global.transpose(-2, -1)) * scale
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, v_global)
else:
# Local token: sliding window + attend to global tokens
start = max(0, i - half_window)
end = min(seq_len, i + half_window + 1)
# Gather local and global keys/values
local_indices = list(range(start, end))
attend_indices = list(set(local_indices + global_positions))
attend_indices.sort()
local_k = k_local[:, :, attend_indices, :]
local_v = v_local[:, :, attend_indices, :]
# Use global k,v for global positions within the attended set
for g in global_positions:
if g in attend_indices:
idx = attend_indices.index(g)
local_k[:, :, idx, :] = k_global[:, :, g, :]
local_v[:, :, idx, :] = v_global[:, :, g, :]
query = q_local[:, :, i : i + 1, :]
scores = torch.matmul(query, local_k.transpose(-2, -1)) * scale
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, local_v)
outputs[:, :, i : i + 1, :] = output
# Reshape and project output
outputs = (
outputs.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.embed_dim)
)
return self.out_proj(outputs)The implementation reveals several important design decisions:
-
Separate projection matrices: We maintain two complete sets of Q, K, V projections. The local projections (
q_proj,k_proj,v_proj) are used for sliding window attention, while the global projections (q_proj_global,k_proj_global,v_proj_global) are used when global tokens are involved. -
Asymmetric handling: Global tokens use global projections for everything. Local tokens use local projections for local attention, but when attending to global tokens, they use the global key and value vectors. This asymmetry ensures consistent representations regardless of which token is doing the attending.
-
Dynamic attention sets: For local tokens, we compute the union of the local window and global positions. This ensures that even if a global token is outside the local window, every token can still attend to it.
Let's verify our implementation works correctly:
# Create a small deterministic test using the chapter-level plot seed
batch_size = 2
seq_len = 32
embed_dim = 64
num_heads = 4
window_size = 8
model = LongformerAttention(embed_dim, num_heads, window_size)
x = torch.randn(batch_size, seq_len, embed_dim)
# Mark positions 0 and 16 as global
global_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool)
global_mask[:, 0] = True
global_mask[:, 16] = True
output = model(x, global_attention_mask=global_mask)Input shape: (2, 32, 64) Output shape: (2, 32, 64) Global positions: [0, 16] Window size: 8 Embedding dimension: 64 Number of attention heads: 4
The implementation correctly handles both local sliding window attention and global attention for designated tokens. The output maintains the same shape as the input (batch size 2, sequence length 32, embedding dimension 64), confirming that our attention mechanism preserves dimensionality while routing attention based on the global attention mask. The two global positions (0 and 16) receive full sequence attention, while all other positions use the sliding window of size 8.