Attention Mechanisms Deep Dive

1. Attention Complexity Overview

1.1 Standard Attention Bottleneck

For a sequence of length with model dimension :

This quadratic scaling limits context length. For and :

1.2 Approaches to Scaling

2. Sparse Attention Patterns

2.1 Fixed Sparse Patterns

Local (Sliding Window) Attention: Each token attends only to neighbours:

Complexity: .

Global Attention: A few "global" tokens attend to all positions:

where is the set of global token indices.

2.2 Longformer (Beltagy et al., 2020)

Combines local and global attention:

• Local: Sliding window of size
• Global: Selected tokens (e.g., [CLS])
• Dilated: Strided windows with dilation

Total complexity: .

2.3 BigBird (Zaheer et al., 2020)

Proven to be a universal approximator with:

1. Random attention: random connections per token
2. Window attention: local window of size
3. Global attention: global tokens

Theorem (Zaheer et al., 2020): BigBird with , , is Turing complete.

3. Linear Attention

3.1 Kernel Trick for Attention

Standard attention:

Linear attention replaces softmax with a kernel :

The key insight is to compute first, then multiply by :

3.2 Random Feature Maps

Performer (Choromanski et al., 2021):

The softmax kernel can be approximated using random features:

where

with .

Positive random features (FAVOR+):

This ensures .

3.3 Random Feature Attention (RFA)

This is simpler and works well in practice:

where and are maintained as running states.

4. Flash Attention

4.1 IO-Aware Algorithm

Flash Attention (Dao et al., 2022) is an IO-optimal algorithm that minimises HBM (High Bandwidth Memory) accesses:

Standard attention:

1. Load from HBM
2. Compute
3. Write to HBM
4. Load from HBM
5. Compute
6. Write to HBM
7. Load from HBM
8. Compute

HBM accesses:

4.2 Tiling and Recomputation

Flash Attention uses tiling to keep data in SRAM (on-chip memory):

1. Divide into blocks
2. For each :
   • Load to SRAM
   • For each :
     • Load to SRAM
     • Compute
     • Compute (online softmax)
     • Update with running statistics

Online softmax:

4.3 Complexity Comparison

where is SRAM size. For :

4.4 Flash Attention 2 and 3

Flash Attention 2: Improves parallelism across sequence length dimension.

Flash Attention 3 (Hopper GPU): Uses asynchronous operations and warp specialization.

5. Multi-Query and Grouped-Query Attention

5.1 Multi-Query Attention (MQA)

In MQA (Shazeer, 2019), all heads share the same key and value projections:

where

Parameter reduction:

For , : ~97% reduction in K/V parameters.

5.2 Grouped-Query Attention (GQA)

GQA (Ainslie et al., 2023) is a middle ground: query heads, key-value heads:

5.3 Performance Comparison

6. Ring Attention

6.1 Distributed Long-Context Attention

Ring Attention (Liu et al., 2023) distributes attention computation across multiple devices for very long sequences:

Setup:

• Sequence length , distributed across devices
• Each device holds tokens
• Devices arranged in a ring topology

Algorithm:

1. Each device starts with
2. Each device computes partial attention with local K, V
3. K, V blocks are passed around the ring
4. After steps, each device has computed full attention

Communication cost:

Computation:

6.2 Load Balancing

To ensure equal load across devices, the sequence is sharded with stride :

This ensures each device has a representative sample of the sequence.

7. Code Examples

7.1 Flash Attention Implementation

import torch
import torch.nn.functional as F

def flash_attention_forward(Q, K, V, block_size=64):
    """
    Simplified Flash Attention implementation.
    
    Parameters
    ----------
    Q, K, V : Tensor (batch, num_heads, seq_len, d_k)
    block_size : int
        Tile size for computation
    
    Returns
    -------
    O : Tensor (batch, num_heads, seq_len, d_k)
    """
    batch, num_heads, seq_len, d_k = Q.shape
    
    # Output accumulator
    O = torch.zeros_like(Q)
    L = torch.zeros(batch, num_heads, seq_len, 1, device=Q.device)
    M = torch.full((batch, num_heads, seq_len, 1), float('-inf'), device=Q.device)
    
    # Iterate over K, V blocks
    for j in range(0, seq_len, block_size):
        # Load K, V block
        K_block = K[:, :, j:j+block_size, :]
        V_block = V[:, :, j:j+block_size, :]
        
        # Iterate over Q blocks
        for i in range(0, seq_len, block_size):
            # Load Q block
            Q_block = Q[:, :, i:i+block_size, :]
            
            # Compute S = QK^T / sqrt(d_k)
            S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_k ** 0.5)
            
            # Online softmax update
            M_block = S_block.max(dim=-1, keepdim=True).values
            M_new = torch.maximum(M[:, :, i:i+block_size], M_block)
            
            # Update L and O
            alpha = torch.exp(M[:, :, i:i+block_size] - M_new)
            beta = torch.exp(M_block - M_new)
            
            L_block = torch.exp(S_block - M_new).sum(dim=-1, keepdim=True)
            L_new = alpha * L[:, :, i:i+block_size] + beta * L_block
            
            O[:, :, i:i+block_size] = (
                alpha * O[:, :, i:i+block_size] + 
                beta * torch.matmul(torch.exp(S_block - M_new), V_block)
            )
            
            M[:, :, i:i+block_size] = M_new
            L[:, :, i:i+block_size] = L_new
    
    # Normalise
    O = O / L
    
    return O


# Example
batch, heads, seq_len, d_k = 2, 8, 1024, 64
Q = torch.randn(batch, heads, seq_len, d_k)
K = torch.randn(batch, heads, seq_len, d_k)
V = torch.randn(batch, heads, seq_len, d_k)

# Flash attention
O_flash = flash_attention_forward(Q, K, V, block_size=64)

# Standard attention for comparison
O_std = F.scaled_dot_product_attention(Q, K, V)

# Compare
max_diff = (O_flash - O_std).abs().max().item()
print(f"Max difference: {max_diff:.6f}")

7.2 Multi-Query Attention

import torch
import torch.nn as nn

class MultiQueryAttention(nn.Module):
    """Multi-Query Attention: shared K, V across heads."""
    
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, self.d_k)  # Single head
        self.W_v = nn.Linear(d_model, self.d_k)  # Single head
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # Q: multiple heads
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # K, V: single head, broadcast to all Q heads
        K = self.W_k(x).unsqueeze(1).expand(-1, self.num_heads, -1, -1)
        V = self.W_v(x).unsqueeze(1).expand(-1, self.num_heads, -1, -1)
        
        # Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
        
        # Reshape and project
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        output = self.W_o(output)
        
        return output


class GroupedQueryAttention(nn.Module):
    """Grouped-Query Attention: G KV heads for H Q heads."""
    
    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.d_k = d_model // num_heads
        self.num_groups = num_heads // num_kv_heads
        
        self.W_q = nn.Linear(d_model, num_heads * self.d_k)
        self.W_k = nn.Linear(d_model, num_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, num_kv_heads * self.d_k)
        self.W_o = nn.Linear(num_heads * self.d_k, d_model)
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # Project
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_kv_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_kv_heads, self.d_k).transpose(1, 2)
        
        # Expand K, V to match Q heads
        K = K.repeat_interleave(self.num_groups, dim=1)
        V = V.repeat_interleave(self.num_groups, dim=1)
        
        # Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
        
        # Reshape and project
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        output = self.W_o(output)
        
        return output


# Example: Compare MHA, MQA, GQA
d_model = 512
num_heads = 8

mha = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
mqa = MultiQueryAttention(d_model, num_heads)
gqa = GroupedQueryAttention(d_model, num_heads, num_kv_heads=2)

x = torch.randn(2, 100, d_model)

# Count parameters
print("Parameter counts:")
print(f"  MHA: {sum(p.numel() for p in mha.parameters()):,}")
print(f"  MQA: {sum(p.numel() for p in mqa.parameters()):,}")
print(f"  GQA: {sum(p.numel() for p in gqa.parameters()):,}")

7.3 Linear Attention

class LinearAttention(nn.Module):
    """
    Linear attention using kernel feature maps.
    
    Complexity: O(n·d²) instead of O(n²·d)
    """
    
    def __init__(self, d_model, num_heads, feature_map='elu'):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        if feature_map == 'elu':
            self.feature_map = lambda x: F.elu(x) + 1
        elif feature_map == 'relu':
            self.feature_map = lambda x: F.relu(x) + 1e-6
        elif feature_map == 'softmax':
            self.feature_map = lambda x: torch.softmax(x, dim=-1)
        else:
            raise ValueError(f"Unknown feature map: {feature_map}")
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # Project
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # Apply feature maps
        Q = self.feature_map(Q)
        K = self.feature_map(K)
        
        if mask is not None:
            # For linear attention, we need to handle masking differently
            mask = mask.unsqueeze(1)  # (batch, 1, 1, seq_len)
            K = K * mask
        
        # Compute KV state: O(d²) not O(n·d)
        KV = torch.einsum('bhnd,bhnv->bhdv', K, V)  # (batch, heads, d_k, d_k)
        Z = torch.einsum('bhnd,bhn->bhd', K, torch.ones_like(K[..., 0]))  # (batch, heads, d_k)
        
        # Compute output
        numerator = torch.einsum('bhnd,bhdv->bhnv', Q, KV)  # (batch, heads, seq_len, d_k)
        denominator = torch.einsum('bhnd,bhd->bhn', Q, Z)  # (batch, heads, seq_len)
        denominator = denominator.unsqueeze(-1) + 1e-6
        
        output = numerator / denominator
        
        # Reshape and project
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        output = self.W_o(output)
        
        return output
    
    def compute_complexity(self, seq_len):
        """Compute FLOPs and memory usage."""
        n = seq_len
        d = self.d_k
        h = self.num_heads
        
        # FLOPs
        flops_qkv = 3 * 2 * n * h * d * d  # Q, K, V projections
        flops_kv = h * n * d * d  # K^T V
        flops_qkv_out = h * n * d * d  # Q (K^T V)
        flops_total = flops_qkv + flops_kv + flops_qkv_out
        
        # Memory
        memory_kv = h * d * d  # KV state
        memory_q = h * n * d  # Q
        memory_total = memory_kv + memory_q
        
        return {
            'flops': flops_total,
            'memory': memory_total,
            'flops_ratio_vs_quadratic': flops_total / (5 * n * n * h * d)
        }


# Example: Compare linear vs quadratic attention
linear_attn = LinearAttention(d_model=256, num_heads=4)

for seq_len in [256, 1024, 4096]:
    complexity = linear_attn.compute_complexity(seq_len)
    print(f"\nseq_len={seq_len}:")
    print(f"  FLOPs: {complexity['flops']:,}")
    print(f"  Memory: {complexity['memory']:,}")
    print(f"  FLOPs ratio vs quadratic: {complexity['flops_ratio_vs_quadratic']:.3f}")

8. Summary

1. Sparse attention reduces complexity to through structured patterns (local, global, dilated).

2. Linear attention achieves via kernel approximations, trading accuracy for efficiency.

3. Flash attention is IO-optimal, reducing HBM accesses from to .

4. Multi-query/grouped-query attention reduce K/V parameters and memory bandwidth by sharing across heads.

5. Ring attention enables distributed computation for sequences exceeding single-device memory.

References

• Dao, T., et al. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. NeurIPS.
• Choromanski, K., et al. (2021). Rethinking attention with performers. ICLR.
• Beltagy, I., et al. (2020). Longformer: The long-document transformer. arXiv.
• Zaheer, M., et al. (2020). Big Bird: Transformers for longer sequences. NeurIPS.
• Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need. arXiv.
• Ainslie, J., et al. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. EMNLP.
• Liu, H., et al. (2023). Ring attention with blockwise transformers for near-infinite context. ICLR.