Chapter 3: Anatomy of an LLM
3.3 Advanced Architectures: SwiGLU, GQA, Attention Sparsity
As transformers evolved from small research models to trillion-parameter giants, engineers discovered that even small changes to internal components can yield big improvements in efficiency and performance. These architectural innovations became increasingly critical as models grew in size and complexity, addressing challenges in computation, memory usage, and training stability. In this section, we'll look at three important innovations that have fundamentally changed how modern LLMs are designed:
- SwiGLU activation functions (improving feedforward networks inside transformer blocks). These specialized activation functions replace traditional ReLU or GELU activations with a more sophisticated gating mechanism that allows for smoother gradient flow during training. By incorporating both multiplicative interactions and non-linear transformations, SwiGLU enables the model to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Grouped Query Attention (GQA) (making attention faster without losing much accuracy). This clever optimization reduces the memory and computational requirements of the attention mechanism by allowing multiple query heads to share the same key and value projections. This significantly decreases both the parameter count and memory bandwidth needed during inference, addressing one of the major bottlenecks in large language model deployment.
- Attention sparsity techniques (reducing compute by ignoring unnecessary connections). These approaches recognize that not all tokens need to attend to every other token in a sequence, especially for long documents. By strategically limiting which tokens can attend to which others, sparse attention patterns can reduce the quadratic complexity of standard attention to near-linear, enabling models to process much longer contexts efficiently.
These optimizations aren't just academic curiosities — they power modern models like LLaMA, Mistral, and GPT-5. Without these architectural advances, today's state-of-the-art models would be prohibitively expensive to train and deploy. Each innovation represents a careful balance between model capability and computational efficiency, addressing specific bottlenecks that emerge at different scales of model size and context length.
3.3.1 SwiGLU (Switched Gated Linear Units)
Every transformer block contains a feedforward network (FFN) after attention. Traditionally, this FFN uses a ReLU or GELU activation. But research found that SwiGLU (a variant of GLU) yields smoother optimization and better performance at scale.
This performance improvement is largely due to the gating mechanism that allows information to flow more selectively through the network, effectively letting the model adaptively control which features are emphasized in each forward pass. Unlike traditional activation functions that apply the same transformation to all inputs, SwiGLU introduces a dynamic, input-dependent filtering mechanism that can emphasize or suppress different aspects of the representation based on the content itself.
In technical terms, SwiGLU combines the benefits of multiplicative interactions (from gates) with non-linear transformations, creating a more expressive computation unit. The gating component (using the swish/SiLU function) produces values between 0 and 1 that act as "soft switches," controlling how much information from each dimension passes through. This adaptive behavior allows the model to create more complex functional mappings with fewer parameters, resulting in improved gradient flow during training and more efficient use of model capacity.
How SwiGLU works:
- Split the hidden dimension into two parts - this creates parallel pathways through the network, allowing the model to learn different aspects of the input simultaneously. This splitting mechanism is crucial because it enables the network to process information along two distinct channels that can later be recombined in a meaningful way. Each pathway can specialize in capturing different features or patterns in the data, similar to how different neurons in the brain might respond to different aspects of visual stimuli.
- Apply a linear transformation to both parts - each pathway gets its own weight matrix (W1 and W2), enabling the network to learn different feature mappings. These linear transformations are fully learned during training and adapt to the specific task at hand. W1 typically projects the input into a space suitable for gating decisions, while W2 creates representations that will be selectively passed through based on those gates. The separation of these transformations allows the model to develop specialized feature detectors that work in tandem.
- Pass one through a sigmoid (or swish) gate - the swish activation (x * sigmoid(x)) provides a smoother gradient than ReLU and allows for some small negative values to pass through, which helps prevent "dying neurons" that can occur with ReLU. The swish function combines the benefits of sigmoid (bounded output) with the non-saturating behavior of ReLU, creating an activation function with better mathematical properties for optimization. This smoother gradient flow helps address the vanishing gradient problem that can plague deep neural networks during training.
- Multiply the two parts together - this multiplicative interaction creates a gating mechanism where the output of the swish function controls how much of the second linear transformation's output passes through. This dynamic gating allows the network to selectively amplify or suppress different features based on the input, leading to more expressive representations. The multiplication operation enables complex, content-dependent filtering of information - effectively allowing some dimensions to be emphasized while others are attenuated based on the specific input pattern, creating a form of adaptive computation.
- The mathematical formula is: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication. In practice, this can be implemented efficiently in modern deep learning frameworks by computing both transformations in parallel and then combining them with a simple Hadamard product. This formulation creates a powerful non-linear transformation that combines the benefits of gating mechanisms (like those in LSTMs and GRUs) with the parallelizability and computational efficiency of feed-forward networks.
PyTorch Example: SwiGLU vs ReLU
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class SwiGLU(nn.Module):
"""
SwiGLU activation function as used in modern LLMs like LLaMA and PaLM.
Uses the Swish/SiLU gating mechanism for better gradient properties.
"""
def __init__(self, dim_in, dim_out):
super().__init__()
self.W1 = nn.Linear(dim_in, dim_out) # Transformation for the gate
self.W2 = nn.Linear(dim_in, dim_out) # Transformation for the content
def forward(self, x):
# SiLU (Swish) activation for gating: x * sigmoid(x)
return F.silu(self.W1(x)) * self.W2(x)
class StandardFFN(nn.Module):
"""
Standard Feedforward Network with ReLU activation
as used in original Transformer architecture.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.ReLU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
class GELUBasedFFN(nn.Module):
"""
Feedforward Network with GELU activation
as used in models like BERT and early GPT versions.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.GELU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
# Hyperparameters
batch_size = 8
seq_len = 32
embed_dim = 256
hidden_dim = 1024
output_dim = 256
# Create example input
x = torch.randn(batch_size, seq_len, embed_dim)
# Initialize models
swiglu = SwiGLU(embed_dim, hidden_dim)
relu_ffn = StandardFFN(embed_dim, hidden_dim, output_dim)
gelu_ffn = GELUBasedFFN(embed_dim, hidden_dim, output_dim)
# Model outputs for comparison
with torch.no_grad():
# Timing comparisons
start = time.time()
swiglu_out = swiglu(x)
swiglu_time = time.time() - start
start = time.time()
relu_out = relu_ffn(x)
relu_time = time.time() - start
start = time.time()
gelu_out = gelu_ffn(x)
gelu_time = time.time() - start
print(f"SwiGLU output shape: {swiglu_out.shape}")
print(f"ReLU FFN output shape: {relu_out.shape}")
print(f"GELU FFN output shape: {gelu_out.shape}")
# Print timing results
print(f"\nForward pass timing:")
print(f"SwiGLU: {swiglu_time*1000:.2f}ms")
print(f"ReLU: {relu_time*1000:.2f}ms")
print(f"GELU: {gelu_time*1000:.2f}ms")
# Print sample outputs
print("\nSample outputs (first 5 values):")
print(f"SwiGLU: {swiglu_out[0, 0, :5].numpy()}")
print(f"ReLU: {relu_out[0, 0, :5].numpy()}")
print(f"GELU: {gelu_out[0, 0, :5].numpy()}")
# Visualize activation functions for comparison
def swish(x):
return x * torch.sigmoid(x)
def relu(x):
return torch.maximum(torch.zeros_like(x), x)
def gelu(x):
return 0.5 * x * (1 + torch.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))
x_range = torch.linspace(-5, 5, 1000)
plt.figure(figsize=(10, 6))
plt.plot(x_range, swish(x_range), label='Swish/SiLU (used in SwiGLU)')
plt.plot(x_range, relu(x_range), label='ReLU')
plt.plot(x_range, gelu(x_range), label='GELU')
plt.grid(True)
plt.legend()
plt.title('Comparison of Activation Functions')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.savefig('activation_comparison.png')
plt.close()
# Simple gradient flow demonstration
def compare_gradient_flow():
"""Compare gradient flow through different activation functions"""
# Create synthetic data
x = torch.randn(100, 32, requires_grad=True)
y = torch.randn(100, 32)
models = {
'SwiGLU': nn.Sequential(
nn.Linear(32, 128),
SwiGLU(128, 128),
nn.Linear(128, 32)
),
'ReLU': nn.Sequential(
nn.Linear(32, 128),
nn.ReLU(),
nn.Linear(128, 32)
),
'GELU': nn.Sequential(
nn.Linear(32, 128),
nn.GELU(),
nn.Linear(128, 32)
)
}
results = {}
for name, model in models.items():
# Forward pass
pred = model(x)
loss = torch.nn.functional.mse_loss(pred, y)
# Backward pass
loss.backward()
# Record gradient statistics
grad_norms = []
for param in model.parameters():
if param.grad is not None:
grad_norms.append(param.grad.norm().item())
results[name] = {
'mean': np.mean(grad_norms),
'std': np.std(grad_norms),
'min': np.min(grad_norms),
'max': np.max(grad_norms)
}
print("\nGradient statistics after one backward pass:")
for name, stats in results.items():
print(f"{name}: mean={stats['mean']:.6f}, std={stats['std']:.6f}, min={stats['min']:.6f}, max={stats['max']:.6f}")
compare_gradient_flow()
Breakdown of SwiGLU
How SwiGLU Works
The implementation breaks down into these key steps:
- Two Parallel Pathways: SwiGLU splits the computation into two parallel linear transformations (W1 and W2).
- Gating Mechanism: One pathway (W1) is passed through a swish/SiLU activation function (x * sigmoid(x)), which provides smoother gradients than ReLU.
- Multiplicative Interaction: The outputs are multiplied together element-wise, allowing the swish-activated path to act as a gate that controls how much of the other path's output passes through.
- Mathematical Formula: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication.
Key Advantages of SwiGLU
- Smoother Gradients: The swish function provides better gradient flow during backpropagation, addressing the vanishing gradient problem that can affect deep networks.
- Dynamic Feature Selection: The gating mechanism allows the network to selectively emphasize or suppress different features based on input content.
- Better Performance Per Parameter: SwiGLU enables models to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Improved Training Dynamics: The smoother activation function and gating mechanism result in more stable and effective training, especially in deep networks.
Code Implementation Details
- The example code demonstrates SwiGLU alongside traditional ReLU and GELU-based feedforward networks for comparison.
- It includes timing comparisons to show computational efficiency differences.
- The visualization of activation functions illustrates how swish/SiLU differs from ReLU and GELU in shape and smoothness.
- The gradient flow demonstration highlights how SwiGLU affects gradient statistics during backpropagation.
This implementation showcases why SwiGLU has become a critical component in modern LLM architectures, offering a better balance of expressivity, computational efficiency, and training stability compared to earlier alternatives.
3.3.2 Grouped Query Attention (GQA)
Standard multi-head attention has a computational cost that grows with the number of heads. To optimize this, Grouped Query Attention (GQA) was introduced. This optimization technique addresses both memory usage and computational efficiency while maintaining model quality.
Key idea: Instead of each query head having its own set of key/value heads, multiple query heads can share the same key/value projections. This sharing mechanism substantially reduces the number of parameters and computation required during inference while preserving the model's ability to learn diverse representations.
The fundamental insight behind GQA is that we can achieve a better balance between computational efficiency and model expressivity by decoupling the number of query heads from the number of key-value heads. This allows models to maintain the benefits of multi-perspective querying while reducing the overhead associated with generating and storing separate key-value pairs for each head.
In traditional multi-head attention, if you have 8 attention heads, you would have 8 separate key projections and 8 separate value projections. With GQA, you might have 8 query heads but only 2 or 4 key-value head pairs that are shared among the query heads. This means that instead of maintaining 8Q+8K+8V projection matrices (24 total), you might only need 8Q+2K+2V (12 total), cutting parameter count and computation significantly.
This reduction becomes increasingly significant in larger models. For instance, in a model with 32 attention heads and an embedding dimension of 4096, traditional multi-head attention would require approximately 402 million parameters for the attention mechanism alone, whereas GQA with 8 KV heads could reduce this to roughly 268 million parameters - a 33% reduction in memory footprint.
To understand the mechanism better, consider how attention works: each query computes similarity scores with all keys, then uses these scores to create a weighted sum of values. In GQA, multiple different queries compute attention scores against the same set of keys, and use these scores to attend to the same set of values. This maintains the expressivity of having multiple query perspectives while economizing on the key-value computations.
The efficiency gains from GQA become particularly evident during inference, where the KV cache (storing pre-computed key-value pairs for autoregressive generation) is often the primary bottleneck. By reducing the size of this cache through key-value sharing, GQA enables models to handle much longer context windows and generate text more efficiently, which is crucial for practical applications in production environments.
- Reduces memory footprint by decreasing the number of parameters needed for key and value projections, which is particularly important for larger models with billions of parameters. For example, in a model with an embedding dimension of 4096 and 32 heads, GQA with 8 KV groups can save approximately 33 million parameters per transformer layer. This reduction is achieved because traditional multi-head attention requires separate projection matrices for each attention head (Q, K, V), whereas GQA allows multiple query heads to share the same key and value projections, dramatically reducing the total parameter count across the model.
- Speeds up inference, especially in long-context models, by reducing the computational burden of generating and storing separate key-value pairs for each attention head. This is critical for server-side deployment where latency directly impacts user experience. During autoregressive generation, the KV cache (which stores previously computed key-value pairs) can become a memory bottleneck. By sharing KV projections across multiple query heads, GQA significantly reduces this cache size, allowing for faster token generation and enabling longer context handling without proportional memory increases.
- Offers nearly the same representational power as full multi-head attention but with significantly improved efficiency — a crucial trade-off for production deployment. Empirical studies show that models with GQA can achieve 95-99% of the performance of models with full multi-head attention while using considerably fewer resources. This minimal performance drop occurs because while the number of key-value pairs is reduced, the model still maintains its full capacity to generate diverse queries, preserving much of its ability to attend from different representational perspectives. The slight performance trade-off is well worth the substantial efficiency gains in most real-world applications.
- Used in LLaMA-2 to balance efficiency with performance, contributing to its ability to handle longer contexts while maintaining reasonable inference speeds. Other models like PaLM 2 and Claude have also adopted variants of this technique to scale efficiently. The implementation in LLaMA-2 specifically helped it achieve significant improvements in context window handling (up to 4K tokens) compared to its predecessor, while keeping inference costs manageable. In PaLM 2, a similar approach enabled efficient scaling to much longer contexts without the quadratic computational explosion that would occur with standard attention mechanisms.
Example: GQA Principle
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class GroupedQueryAttention(nn.Module):
def __init__(self, embed_dim=512, num_query_heads=8, num_kv_heads=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_query_heads = num_query_heads
self.num_kv_heads = num_kv_heads
self.head_dim = embed_dim // num_query_heads
# Ensure dimensions are compatible
assert self.head_dim * num_query_heads == embed_dim, "embed_dim must be divisible by num_query_heads"
assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
# Query projections (many heads)
self.q_proj = nn.Linear(embed_dim, embed_dim)
# Key/Value projections (fewer heads - shared)
self.k_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
self.v_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
# Output projection
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Dropout for attention weights
self.dropout = nn.Dropout(dropout)
# Group size: how many query heads share one kv head
self.group_size = num_query_heads // num_kv_heads
def forward(self, x, mask=None):
"""
Args:
x: Input tensor of shape (batch_size, sequence_length, embed_dim)
mask: Optional attention mask
Returns:
output: Tensor after self-attention of shape (batch_size, sequence_length, embed_dim)
"""
batch_size, seq_len, _ = x.size()
# Project inputs to queries, keys, and values
q = self.q_proj(x).view(batch_size, seq_len, self.num_query_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Transpose for attention computation
q = q.transpose(1, 2) # (batch_size, num_query_heads, seq_len, head_dim)
k = k.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
v = v.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
# Expand k and v to match the number of query heads through repetition
# Each group of query heads shares the same k and v
k_expanded = torch.repeat_interleave(k, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
v_expanded = torch.repeat_interleave(v, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
# Compute scaled dot-product attention
# (batch_size, num_query_heads, seq_len, seq_len)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
# Apply mask if provided (useful for preventing attention to padding tokens)
if mask is not None:
attn_weights = attn_weights.masked_fill(mask == 0, float("-inf"))
# Apply softmax and dropout
attn_weights = F.softmax(attn_weights, dim=-1)
attn_weights = self.dropout(attn_weights)
# Apply attention weights to values
# (batch_size, num_query_heads, seq_len, head_dim)
attn_output = torch.matmul(attn_weights, v_expanded)
# Reshape and apply output projection
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
output = self.out_proj(attn_output)
return output
def compare_attention_mechanisms(seq_len=1024, embed_dim=512):
"""Compare memory usage and speed between standard MHA and GQA"""
batch_size = 1
# Create inputs
x = torch.randn(batch_size, seq_len, embed_dim)
# Standard Multi-Head Attention (8 heads)
class StandardMHA(nn.Module):
def __init__(self):
super().__init__()
self.mha = nn.MultiheadAttention(embed_dim, num_heads=8, batch_first=True)
def forward(self, x):
return self.mha(x, x, x)[0]
standard_mha = StandardMHA()
# GQA with 8 query heads, 2 KV heads
gqa = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=2)
# GQA with 8 query heads, 4 KV heads
gqa2 = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=4)
# Measure memory and speed
results = {}
for name, model in [("Standard MHA (8 heads)", standard_mha),
("GQA (8Q, 2KV heads)", gqa),
("GQA (8Q, 4KV heads)", gqa2)]:
# Warm up
for _ in range(5):
_ = model(x)
# Measure time
torch.cuda.synchronize() if torch.cuda.is_available() else None
start_time = time.time()
for _ in range(10):
_ = model(x)
torch.cuda.synchronize() if torch.cuda.is_available() else None
end_time = time.time()
# Count parameters
param_count = sum(p.numel() for p in model.parameters())
results[name] = {
"time_per_run_ms": (end_time - start_time) * 100, # ms per 10 runs
"parameters": param_count
}
# Print results
print("Performance Comparison (sequence length = {})".format(seq_len))
print("=" * 50)
for name, metrics in results.items():
print(f"{name}:")
print(f" Time per 10 runs: {metrics['time_per_run_ms']:.2f} ms")
print(f" Parameters: {metrics['parameters']:,}")
print("-" * 50)
# Visualize KV cache size comparison
kv_cache_sizes = {
"Standard MHA": seq_len * 2 * embed_dim, # Full KV cache (8 heads)
"GQA (2 KV heads)": seq_len * 2 * (embed_dim // 4), # 1/4 the size (2 heads)
"GQA (4 KV heads)": seq_len * 2 * (embed_dim // 2), # 1/2 the size (4 heads)
}
plt.figure(figsize=(10, 5))
plt.bar(kv_cache_sizes.keys(), [size/1e6 for size in kv_cache_sizes.values()])
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Size Comparison')
for i, v in enumerate(kv_cache_sizes.values()):
plt.text(i, v/1e6 + 0.1, f"{v/1e6:.2f} MB", ha='center')
# Show how KV cache grows with sequence length
seq_lengths = [1024, 2048, 4096, 8192, 16384]
std_cache_sizes = [seq_len * 2 * embed_dim / 1e6 for seq_len in seq_lengths]
gqa_cache_sizes = [seq_len * 2 * (embed_dim // 4) / 1e6 for seq_len in seq_lengths]
plt.figure(figsize=(10, 5))
plt.plot(seq_lengths, std_cache_sizes, 'bo-', label='Standard MHA')
plt.plot(seq_lengths, gqa_cache_sizes, 'ro-', label='GQA (2 KV heads)')
plt.xlabel('Sequence Length')
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Growth with Sequence Length')
plt.legend()
plt.grid(True)
# Demonstrate usage with a simple example
seq_len, batch_size, embed_dim = 5, 1, 32
x = torch.randn(batch_size, seq_len, embed_dim)
gqa = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=8, num_kv_heads=2)
output = gqa(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
# Compare with standard attention
compare_attention_mechanisms(seq_len=2048, embed_dim=512)
# Basic visualization of attention patterns
def visualize_attention_pattern():
seq_len = 10
embed_dim = 64
x = torch.randn(1, seq_len, embed_dim)
model = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=4, num_kv_heads=2)
# Get attention weights by modifying forward pass temporarily
with torch.no_grad():
q = model.q_proj(x).view(1, seq_len, model.num_query_heads, model.head_dim)
k = model.k_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
v = model.v_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
k_expanded = torch.repeat_interleave(k, model.group_size, dim=1)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(model.head_dim, dtype=torch.float32))
attn_weights = F.softmax(attn_weights, dim=-1)
# Plot attention patterns for each head
fig, axes = plt.subplots(1, model.num_query_heads, figsize=(15, 3))
for i in range(model.num_query_heads):
im = axes[i].imshow(attn_weights[0, i].cpu().numpy(), cmap='viridis')
axes[i].set_title(f'Head {i+1}')
axes[i].set_xlabel('Key position')
axes[i].set_ylabel('Query position')
fig.colorbar(im, ax=axes.ravel().tolist())
plt.tight_layout()
plt.suptitle('Attention Patterns with GQA (notice shared patterns within groups)')
visualize_attention_pattern()Here's a comprehensive breakdown:
Core GQA Implementation
The code example Grouped Query Attention, an optimization technique that reduces memory usage and computational cost compared to standard multi-head attention.
Class Structure
- The
GroupedQueryAttentionclass inherits fromnn.Moduleand takes parameters for embedding dimension, number of query heads, number of key-value heads, and dropout rate. - The key innovation is that multiple query heads share the same key-value heads, reducing parameter count and memory footprint.
- Two compatibility assertions ensure:
- embedding dimension is divisible by the number of query heads
- query heads are divisible by key-value heads
Projection Layers
- Query projection: Full dimension projection (
self.q_proj) - Key/Value projections: Reduced dimension projections (
self.k_proj,self.v_proj) - Output projection: Maps attention output back to original dimensions
Forward Pass
- Projects input into queries, keys and values with appropriate dimensions
- Transposes tensors for attention computation
- The critical step: expands key and value tensors to match query heads through repetition using
torch.repeat_interleave - Computes scaled dot-product attention with softmax normalization
- Applies attention weights to values and reshapes the output
Performance Comparison Functions
The code includes utilities to demonstrate GQA's advantages:
compare_attention_mechanisms(): Benchmarks standard MHA against GQA variants with different head configurations measuring:- Execution time
- Parameter count
- KV cache size - critical for inference efficiency
- Visualization functions for KV cache size comparisons and growth with sequence length
- The
visualize_attention_pattern()function demonstrates how attention patterns appear in GQA, showing how multiple query heads share the same key-value pairs
Key Benefits Demonstrated
- Memory efficiency: Reduces parameters by sharing key-value projections
- Inference speed: Smaller KV cache allows for faster token generation
- Context length: Enables handling longer sequences with minimal memory growth
- Used in modern models: The implementation resembles approaches used in LLaMA-2, PaLM 2, and Claude
This implementation provides both a practical demonstration of GQA and tools to visualize its benefits over traditional attention mechanisms, particularly in terms of memory usage and computational efficiency while maintaining most of the representational power of full multi-head attention.
3.3.3 Attention Sparsity
In full self-attention, each token attends to every other token in the sequence. This creates a computational complexity that scales quadratically as O(n²) with sequence length, which becomes prohibitively expensive for long sequences (think 100k+ tokens). For context, processing a sequence of 100,000 tokens would require 10 billion attention computations per layer!
To understand why this is problematic, consider what happens as we scale: if we double our context length from 4K to 8K tokens, the computational work quadruples from 16 million to 64 million connections per layer. This quadratic scaling quickly becomes a bottleneck for both training and inference.
Additionally, the memory requirements for storing the attention matrix also scale quadratically. For a sequence of length n, we need to store an n×n attention matrix, which for long sequences can exceed available GPU memory. For example, a 32K token sequence would require approximately 4GB of memory just to store a single attention matrix in 32-bit precision.
Sparse attention techniques reduce this computational burden by attending only to the most relevant positions, effectively pruning unnecessary connections. This transforms the scaling from quadratic to nearly linear in many implementations. By strategically limiting which tokens can attend to which other tokens, these techniques dramatically reduce both computation and memory requirements.
The key insight behind sparse attention is that not all token-to-token interactions are equally important. Many language phenomena are local in nature, while certain special tokens may need global context. By exploiting this pattern, sparse attention can preserve most of the model's capabilities while eliminating many unnecessary computations.
Local attention
Each token attends only to its neighbors within a fixed window size (e.g., ±128 tokens). This creates a sliding window of attention that moves with each token position. For example, with a window size of 128, token at position 500 would attend to tokens from positions 372 to 628.
This approach works particularly well for tasks where nearby context is most relevant, such as speech recognition where phonemes relate strongly to adjacent sounds, or DNA analysis where nearby nucleotides often form functional units together. Local attention is also effective for text processing tasks where most semantic relationships occur between words that are relatively close to each other in the sequence.
The efficiency gains are substantial - the computational complexity becomes O(n×w), where w is the fixed window size. Since w is a constant (like 128 or 256), this effectively makes the attention mechanism scale linearly with sequence length rather than quadratically. For a sequence of 100,000 tokens with a window size of 256, this reduces computations from 10 billion to just 25.6 million - a 390x improvement.
However, local attention does have limitations - it struggles with tasks requiring long-range dependencies, such as document-level reasoning where important information may be separated by thousands of tokens. This is why more sophisticated sparse attention patterns often combine local attention with other mechanisms to capture both local and global relationships.
Block-sparse attention
Tokens attend within defined chunks or blocks, with occasional global tokens that can see across the entire sequence. This creates a sparse attention pattern where most tokens have limited vision but a few sentinel tokens maintain global context. These blocks can be arranged in various patterns - diagonal blocks for local attention, or more complex structures that allow for hierarchical information flow.
For example, in a block-sparse approach, a document might be divided into chunks of 512 tokens, with each chunk having internal full attention, plus dedicated "summary tokens" that can see across all chunks. This creates an information highway where local details are processed efficiently within blocks, while global information flows through the designated global tokens.
Additionally, some implementations use strided patterns where tokens can attend to blocks at regular intervals throughout the sequence, capturing periodic patterns or relationships. Others employ random sparse patterns that theoretically allow information to flow between any two positions through a small number of hops.
This hybrid approach preserves most of the modeling power of full attention while dramatically reducing computation. By carefully designing which blocks can attend to which others, these models achieve an attention complexity closer to O(n√n) or even O(n log n) rather than O(n²), enabling processing of much longer sequences with the same computational resources.
BigBird and Longformer
BigBird and Longformer implement sophisticated sparse attention patterns combining local windows, global tokens, and random connections. These architectures can efficiently scale to sequences of 4,000–8,000+ tokens with minimal loss in performance compared to full attention models.
BigBird, for example, combines three distinct attention patterns:
- Window attention: Each token attends to its local neighborhood (similar to the sliding window approach). This allows the model to capture local context effectively by focusing on nearby tokens. For instance, in a document about climate change, this helps the model understand phrases and nearby semantic connections by creating a focused attention window around each token, typically spanning 256-512 tokens in each direction.
- Global attention: Special tokens like [CLS] attend to all tokens and are attended to by all tokens, creating information highways. These global tokens serve as aggregation points that collect information from the entire sequence and distribute it back, enabling document-level understanding. For example, in a long scientific paper, the [CLS] token might gather key conclusions from various sections and make this information available to all other tokens, facilitating cross-referencing across distant parts of the document.
- Random attention: Each token attends to a small set of randomly selected tokens, which theoretically allows information to flow between any two positions in logarithmic steps. This random connectivity creates shortcuts across the document, ensuring information can propagate efficiently between distant sections. Mathematical proofs show that with just O(log n) random connections, information can flow between any two tokens in the sequence. In practice, this means even tokens separated by thousands of positions can exchange information through just a few intermediate connections.
This tri-directional attention mechanism achieves near-linear scaling while maintaining strong performance on long-document tasks like summarization and question answering. Importantly, BigBird maintains the theoretical property of "universal approximation" - it can represent any sequence-to-sequence function that full attention can, but with dramatically reduced computational requirements.
Longformer employs a similar approach but with a slightly different pattern, using a combination of sliding window attention with global attention for special tokens. It has demonstrated particular effectiveness in tasks requiring both local precision and document-level understanding, such as long-document question answering and multi-document summarization, where it can process inputs of 16,000+ tokens.
Code Example: Local Attention (Sliding Window)
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import time
class LocalAttention(nn.Module):
def __init__(self, dim, window_size=128):
super().__init__()
self.dim = dim
self.window_size = window_size
self.query_proj = nn.Linear(dim, dim)
self.key_proj = nn.Linear(dim, dim)
self.value_proj = nn.Linear(dim, dim)
self.output_proj = nn.Linear(dim, dim)
self.scaling = dim ** -0.5
def forward(self, x):
B, T, D = x.size()
# Project inputs to queries, keys, values
queries = self.query_proj(x) * self.scaling # [B, T, D]
keys = self.key_proj(x) # [B, T, D]
values = self.value_proj(x) # [B, T, D]
# Initialize output tensor
output = torch.zeros_like(x)
# Compute local attention for each position
for i in range(T):
# Define local window boundaries
start = max(0, i - self.window_size)
end = min(T, i + self.window_size + 1)
# Extract local context
local_keys = keys[:, start:end, :] # [B, window_size*2, D]
local_values = values[:, start:end, :] # [B, window_size*2, D]
# Current query
query = queries[:, i:i+1, :] # [B, 1, D]
# Compute attention scores
scores = torch.bmm(query, local_keys.transpose(1, 2)) # [B, 1, window_size*2]
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, 1, window_size*2]
# Weight values by attention
context = torch.bmm(attn_weights, local_values) # [B, 1, D]
# Store in output
output[:, i:i+1, :] = context
return self.output_proj(output)
def naive_local_attention(x, window=2):
"""A simple implementation of local attention for educational purposes"""
B, T, D = x.size()
outputs = []
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
context = x[:, start:end, :]
weights = F.softmax(torch.bmm(x[:, i:i+1, :], context.transpose(1,2)), dim=-1)
out = torch.bmm(weights, context)
outputs.append(out)
return torch.cat(outputs, dim=1)
def vectorized_local_attention(x, window=2):
"""A more efficient implementation using vectorized operations"""
B, T, D = x.size()
# Create attention mask to implement sliding window
mask = torch.zeros(T, T, device=x.device)
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
mask[i, start:end] = 1
# Compute attention scores
scores = torch.bmm(x, x.transpose(1, 2)) # [B, T, T]
# Apply mask (setting padded values to -inf before softmax)
scores = scores.masked_fill(mask.unsqueeze(0) == 0, -1e9)
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, T, T]
# Weight values by attention
output = torch.bmm(attn_weights, x) # [B, T, D]
return output
def compare_performance(seq_lengths=[10, 50, 100, 200], window=2):
"""Compare performance of different local attention implementations"""
results = {'naive': [], 'vectorized': [], 'optimized': []}
for seq_len in seq_lengths:
# Generate random input tensor
x = torch.randn(1, seq_len, 64)
# Naive implementation
start_time = time.time()
naive_local_attention(x, window)
naive_time = time.time() - start_time
results['naive'].append(naive_time)
# Vectorized implementation
start_time = time.time()
vectorized_local_attention(x, window)
vectorized_time = time.time() - start_time
results['vectorized'].append(vectorized_time)
# Optimized implementation
model = LocalAttention(64, window)
start_time = time.time()
model(x)
optimized_time = time.time() - start_time
results['optimized'].append(optimized_time)
print(f"Sequence length {seq_len}:")
print(f" Naive: {naive_time:.5f}s")
print(f" Vectorized: {vectorized_time:.5f}s")
print(f" Optimized: {optimized_time:.5f}s")
# Plot results
plt.figure(figsize=(10, 6))
plt.plot(seq_lengths, results['naive'], 'o-', label='Naive')
plt.plot(seq_lengths, results['vectorized'], 's-', label='Vectorized')
plt.plot(seq_lengths, results['optimized'], '^-', label='Optimized')
plt.xlabel('Sequence Length')
plt.ylabel('Time (s)')
plt.title('Performance Comparison of Local Attention Implementations')
plt.legend()
plt.grid(True)
plt.show()
def visualize_attention_pattern(window=2, seq_len=10):
"""Visualize the sparse attention pattern created by local attention"""
attention_mask = torch.zeros(seq_len, seq_len)
for i in range(seq_len):
start = max(0, i - window)
end = min(seq_len, i + window + 1)
attention_mask[i, start:end] = 1
plt.figure(figsize=(8, 8))
plt.imshow(attention_mask, cmap='Blues')
plt.title(f'Local Attention Pattern (Window Size = {window})')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar(label='Attention Connection')
for i in range(seq_len):
for j in range(seq_len):
color = 'white' if attention_mask[i, j] > 0 else 'none'
plt.text(j, i, '1' if attention_mask[i, j] > 0 else '0',
ha='center', va='center', color=color)
plt.tight_layout()
plt.show()
# Example
if __name__ == "__main__":
# Basic functionality test
x = torch.randn(1, 6, 16)
model = LocalAttention(16, window_size=2)
out = model(x)
print(f"Input shape: {x.shape}, Output shape: {out.shape}")
# Compare implementations
compare_performance([10, 50, 100, 200], window=2)
# Visualize the attention pattern
visualize_attention_pattern(window=2, seq_len=10)Comprehensive Breakdown: Local Attention Implementation
This code example provides a complete toolkit for understanding, implementing and analyzing local attention mechanisms. Here's a detailed breakdown:
1. Core Implementations
- LocalAttention Class: A proper PyTorch module implementation with:
- Dedicated projection layers for queries, keys, and valuesDedicated projection layers for queries, keys, and values
- Window-based sliding attention with configurable window sizeWindow-based sliding attention with configurable window size
- Proper scaling factor (1/√d) for stable gradientsProper scaling factor (1/√d) for stable gradients
- Final output projection as in standard attentionFinal output projection as in standard attention
- Naive Implementation: The original function that:
- Processes each position sequentiallyProcesses each position sequentially
- Demonstrates the core sliding window concept clearlyDemonstrates the core sliding window concept clearly
- Uses simple tensor operations for educational purposesUses simple tensor operations for educational purposes
- Vectorized Implementation: A more efficient approach that:
- Uses a mask tensor to implement the sliding window patternUses a mask tensor to implement the sliding window pattern
- Computes all attention scores at onceComputes all attention scores at once
- Avoids explicit loops over sequence positionsAvoids explicit loops over sequence positions
2. Analysis Tools
- Performance Comparison Function: Benchmarks all three implementations:
- Measures execution time across different sequence lengthsMeasures execution time across different sequence lengths
- Generates performance plots to visualize scaling behaviorGenerates performance plots to visualize scaling behavior
- Demonstrates how vectorized operations improve efficiencyDemonstrates how vectorized operations improve efficiency
- Visualization Function: Illustrates the sparse attention pattern:
- Creates a visual representation of which tokens attend to which othersCreates a visual representation of which tokens attend to which others
- Shows the diagonal band pattern characteristic of local attentionShows the diagonal band pattern characteristic of local attention
- Helps intuitive understanding of how information flows in the modelHelps intuitive understanding of how information flows in the model
3. Key Technical Insights
- Masking Technique: The code demonstrates how to create and apply attention masks to restrict which tokens can attend to which others
- Computational Efficiency: Shows how the computational complexity becomes O(n·w) instead of O(n²), where w is the window size
- Implementation Trade-offs: Illustrates the balance between code clarity (naive implementation) and computational efficiency (vectorized implementation)
This implementation provides both theoretical understanding and practical tools for working with local attention, a key technique for making transformers more efficient with long sequences. The visualization and comparison functions make it especially valuable for educational purposes.
3.3.4 Why These Matter
SwiGLU (Swish-Gated Linear Unit) significantly improves learning dynamics, giving models richer representations with little extra computational cost. This sophisticated activation function combines the benefits of gating mechanisms with a simple weighted identity connection, allowing for more effective gradient flow during training. By replacing traditional ReLU or GELU activations, SwiGLU enables models to learn more complex patterns while maintaining computational efficiency.
The mathematical formulation of SwiGLU involves multiplying a linear projection of the input with a sigmoid-weighted version of another projection, creating a smooth, differentiable pathway for gradients that helps prevent vanishing gradient problems. Models using SwiGLU typically converge faster and achieve better performance across various natural language processing tasks, making it a preferred choice in modern LLM architectures like PaLM and Gemini.
GQA (Grouped Query Attention) makes attention mechanisms substantially more efficient, reducing memory use without significant accuracy loss. This innovative technique groups queries together to share the same keys and values, dramatically reducing the memory footprint during inference. Unlike standard multi-head attention that requires separate key-value pairs for each attention head (creating a parameter explosion), GQA significantly cuts down on parameters while preserving most of the model's reasoning capabilities.
This approach creates a middle ground between multi-head attention (MHA) and multi-query attention (MQA), finding an optimal balance between parameter efficiency and model capacity. In practice, GQA can reduce the key-value cache memory requirements by 2-4x compared to standard attention while maintaining 95-99% of the model's performance, making it possible to deploy larger models on the same hardware or increase batch sizes during inference. Models like PaLM 2 and Claude have successfully implemented GQA as a core architectural improvement.
Sparse attention fundamentally transforms how LLMs can handle very long contexts without suffering from quadratic computational blow-ups. Instead of having each token attend to every other token (which scales as O(n²) with sequence length), sparse attention patterns like local, dilated, or longformer attention enable selective focus on only the most relevant tokens. This reduces computational complexity to O(n) or O(n log n), making it feasible to process documents with thousands or even tens of thousands of tokens.
Local attention, as shown in the code example above, restricts each token to attend only to a window of neighboring tokens. Dilated attention extends this by allowing tokens to attend to positions at various distances, creating a wider receptive field without increasing computation proportionally. More advanced sparse attention patterns like Reformer's LSH attention or Longformer's global+local attention combine different strategies to balance computational efficiency with model capacity. These approaches have enabled breakthroughs in long-context models that can process entire books, codebases, or lengthy conversations while maintaining coherent understanding throughout the document.
Together, these architectural refinements are why today's LLMs can be faster, leaner, and more scalable than early transformers. They represent critical engineering breakthroughs that have transformed theoretical research models into practical, deployable systems capable of handling real-world tasks with unprecedented efficiency.
3.3 Advanced Architectures: SwiGLU, GQA, Attention Sparsity
As transformers evolved from small research models to trillion-parameter giants, engineers discovered that even small changes to internal components can yield big improvements in efficiency and performance. These architectural innovations became increasingly critical as models grew in size and complexity, addressing challenges in computation, memory usage, and training stability. In this section, we'll look at three important innovations that have fundamentally changed how modern LLMs are designed:
- SwiGLU activation functions (improving feedforward networks inside transformer blocks). These specialized activation functions replace traditional ReLU or GELU activations with a more sophisticated gating mechanism that allows for smoother gradient flow during training. By incorporating both multiplicative interactions and non-linear transformations, SwiGLU enables the model to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Grouped Query Attention (GQA) (making attention faster without losing much accuracy). This clever optimization reduces the memory and computational requirements of the attention mechanism by allowing multiple query heads to share the same key and value projections. This significantly decreases both the parameter count and memory bandwidth needed during inference, addressing one of the major bottlenecks in large language model deployment.
- Attention sparsity techniques (reducing compute by ignoring unnecessary connections). These approaches recognize that not all tokens need to attend to every other token in a sequence, especially for long documents. By strategically limiting which tokens can attend to which others, sparse attention patterns can reduce the quadratic complexity of standard attention to near-linear, enabling models to process much longer contexts efficiently.
These optimizations aren't just academic curiosities — they power modern models like LLaMA, Mistral, and GPT-5. Without these architectural advances, today's state-of-the-art models would be prohibitively expensive to train and deploy. Each innovation represents a careful balance between model capability and computational efficiency, addressing specific bottlenecks that emerge at different scales of model size and context length.
3.3.1 SwiGLU (Switched Gated Linear Units)
Every transformer block contains a feedforward network (FFN) after attention. Traditionally, this FFN uses a ReLU or GELU activation. But research found that SwiGLU (a variant of GLU) yields smoother optimization and better performance at scale.
This performance improvement is largely due to the gating mechanism that allows information to flow more selectively through the network, effectively letting the model adaptively control which features are emphasized in each forward pass. Unlike traditional activation functions that apply the same transformation to all inputs, SwiGLU introduces a dynamic, input-dependent filtering mechanism that can emphasize or suppress different aspects of the representation based on the content itself.
In technical terms, SwiGLU combines the benefits of multiplicative interactions (from gates) with non-linear transformations, creating a more expressive computation unit. The gating component (using the swish/SiLU function) produces values between 0 and 1 that act as "soft switches," controlling how much information from each dimension passes through. This adaptive behavior allows the model to create more complex functional mappings with fewer parameters, resulting in improved gradient flow during training and more efficient use of model capacity.
How SwiGLU works:
- Split the hidden dimension into two parts - this creates parallel pathways through the network, allowing the model to learn different aspects of the input simultaneously. This splitting mechanism is crucial because it enables the network to process information along two distinct channels that can later be recombined in a meaningful way. Each pathway can specialize in capturing different features or patterns in the data, similar to how different neurons in the brain might respond to different aspects of visual stimuli.
- Apply a linear transformation to both parts - each pathway gets its own weight matrix (W1 and W2), enabling the network to learn different feature mappings. These linear transformations are fully learned during training and adapt to the specific task at hand. W1 typically projects the input into a space suitable for gating decisions, while W2 creates representations that will be selectively passed through based on those gates. The separation of these transformations allows the model to develop specialized feature detectors that work in tandem.
- Pass one through a sigmoid (or swish) gate - the swish activation (x * sigmoid(x)) provides a smoother gradient than ReLU and allows for some small negative values to pass through, which helps prevent "dying neurons" that can occur with ReLU. The swish function combines the benefits of sigmoid (bounded output) with the non-saturating behavior of ReLU, creating an activation function with better mathematical properties for optimization. This smoother gradient flow helps address the vanishing gradient problem that can plague deep neural networks during training.
- Multiply the two parts together - this multiplicative interaction creates a gating mechanism where the output of the swish function controls how much of the second linear transformation's output passes through. This dynamic gating allows the network to selectively amplify or suppress different features based on the input, leading to more expressive representations. The multiplication operation enables complex, content-dependent filtering of information - effectively allowing some dimensions to be emphasized while others are attenuated based on the specific input pattern, creating a form of adaptive computation.
- The mathematical formula is: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication. In practice, this can be implemented efficiently in modern deep learning frameworks by computing both transformations in parallel and then combining them with a simple Hadamard product. This formulation creates a powerful non-linear transformation that combines the benefits of gating mechanisms (like those in LSTMs and GRUs) with the parallelizability and computational efficiency of feed-forward networks.
PyTorch Example: SwiGLU vs ReLU
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class SwiGLU(nn.Module):
"""
SwiGLU activation function as used in modern LLMs like LLaMA and PaLM.
Uses the Swish/SiLU gating mechanism for better gradient properties.
"""
def __init__(self, dim_in, dim_out):
super().__init__()
self.W1 = nn.Linear(dim_in, dim_out) # Transformation for the gate
self.W2 = nn.Linear(dim_in, dim_out) # Transformation for the content
def forward(self, x):
# SiLU (Swish) activation for gating: x * sigmoid(x)
return F.silu(self.W1(x)) * self.W2(x)
class StandardFFN(nn.Module):
"""
Standard Feedforward Network with ReLU activation
as used in original Transformer architecture.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.ReLU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
class GELUBasedFFN(nn.Module):
"""
Feedforward Network with GELU activation
as used in models like BERT and early GPT versions.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.GELU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
# Hyperparameters
batch_size = 8
seq_len = 32
embed_dim = 256
hidden_dim = 1024
output_dim = 256
# Create example input
x = torch.randn(batch_size, seq_len, embed_dim)
# Initialize models
swiglu = SwiGLU(embed_dim, hidden_dim)
relu_ffn = StandardFFN(embed_dim, hidden_dim, output_dim)
gelu_ffn = GELUBasedFFN(embed_dim, hidden_dim, output_dim)
# Model outputs for comparison
with torch.no_grad():
# Timing comparisons
start = time.time()
swiglu_out = swiglu(x)
swiglu_time = time.time() - start
start = time.time()
relu_out = relu_ffn(x)
relu_time = time.time() - start
start = time.time()
gelu_out = gelu_ffn(x)
gelu_time = time.time() - start
print(f"SwiGLU output shape: {swiglu_out.shape}")
print(f"ReLU FFN output shape: {relu_out.shape}")
print(f"GELU FFN output shape: {gelu_out.shape}")
# Print timing results
print(f"\nForward pass timing:")
print(f"SwiGLU: {swiglu_time*1000:.2f}ms")
print(f"ReLU: {relu_time*1000:.2f}ms")
print(f"GELU: {gelu_time*1000:.2f}ms")
# Print sample outputs
print("\nSample outputs (first 5 values):")
print(f"SwiGLU: {swiglu_out[0, 0, :5].numpy()}")
print(f"ReLU: {relu_out[0, 0, :5].numpy()}")
print(f"GELU: {gelu_out[0, 0, :5].numpy()}")
# Visualize activation functions for comparison
def swish(x):
return x * torch.sigmoid(x)
def relu(x):
return torch.maximum(torch.zeros_like(x), x)
def gelu(x):
return 0.5 * x * (1 + torch.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))
x_range = torch.linspace(-5, 5, 1000)
plt.figure(figsize=(10, 6))
plt.plot(x_range, swish(x_range), label='Swish/SiLU (used in SwiGLU)')
plt.plot(x_range, relu(x_range), label='ReLU')
plt.plot(x_range, gelu(x_range), label='GELU')
plt.grid(True)
plt.legend()
plt.title('Comparison of Activation Functions')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.savefig('activation_comparison.png')
plt.close()
# Simple gradient flow demonstration
def compare_gradient_flow():
"""Compare gradient flow through different activation functions"""
# Create synthetic data
x = torch.randn(100, 32, requires_grad=True)
y = torch.randn(100, 32)
models = {
'SwiGLU': nn.Sequential(
nn.Linear(32, 128),
SwiGLU(128, 128),
nn.Linear(128, 32)
),
'ReLU': nn.Sequential(
nn.Linear(32, 128),
nn.ReLU(),
nn.Linear(128, 32)
),
'GELU': nn.Sequential(
nn.Linear(32, 128),
nn.GELU(),
nn.Linear(128, 32)
)
}
results = {}
for name, model in models.items():
# Forward pass
pred = model(x)
loss = torch.nn.functional.mse_loss(pred, y)
# Backward pass
loss.backward()
# Record gradient statistics
grad_norms = []
for param in model.parameters():
if param.grad is not None:
grad_norms.append(param.grad.norm().item())
results[name] = {
'mean': np.mean(grad_norms),
'std': np.std(grad_norms),
'min': np.min(grad_norms),
'max': np.max(grad_norms)
}
print("\nGradient statistics after one backward pass:")
for name, stats in results.items():
print(f"{name}: mean={stats['mean']:.6f}, std={stats['std']:.6f}, min={stats['min']:.6f}, max={stats['max']:.6f}")
compare_gradient_flow()
Breakdown of SwiGLU
How SwiGLU Works
The implementation breaks down into these key steps:
- Two Parallel Pathways: SwiGLU splits the computation into two parallel linear transformations (W1 and W2).
- Gating Mechanism: One pathway (W1) is passed through a swish/SiLU activation function (x * sigmoid(x)), which provides smoother gradients than ReLU.
- Multiplicative Interaction: The outputs are multiplied together element-wise, allowing the swish-activated path to act as a gate that controls how much of the other path's output passes through.
- Mathematical Formula: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication.
Key Advantages of SwiGLU
- Smoother Gradients: The swish function provides better gradient flow during backpropagation, addressing the vanishing gradient problem that can affect deep networks.
- Dynamic Feature Selection: The gating mechanism allows the network to selectively emphasize or suppress different features based on input content.
- Better Performance Per Parameter: SwiGLU enables models to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Improved Training Dynamics: The smoother activation function and gating mechanism result in more stable and effective training, especially in deep networks.
Code Implementation Details
- The example code demonstrates SwiGLU alongside traditional ReLU and GELU-based feedforward networks for comparison.
- It includes timing comparisons to show computational efficiency differences.
- The visualization of activation functions illustrates how swish/SiLU differs from ReLU and GELU in shape and smoothness.
- The gradient flow demonstration highlights how SwiGLU affects gradient statistics during backpropagation.
This implementation showcases why SwiGLU has become a critical component in modern LLM architectures, offering a better balance of expressivity, computational efficiency, and training stability compared to earlier alternatives.
3.3.2 Grouped Query Attention (GQA)
Standard multi-head attention has a computational cost that grows with the number of heads. To optimize this, Grouped Query Attention (GQA) was introduced. This optimization technique addresses both memory usage and computational efficiency while maintaining model quality.
Key idea: Instead of each query head having its own set of key/value heads, multiple query heads can share the same key/value projections. This sharing mechanism substantially reduces the number of parameters and computation required during inference while preserving the model's ability to learn diverse representations.
The fundamental insight behind GQA is that we can achieve a better balance between computational efficiency and model expressivity by decoupling the number of query heads from the number of key-value heads. This allows models to maintain the benefits of multi-perspective querying while reducing the overhead associated with generating and storing separate key-value pairs for each head.
In traditional multi-head attention, if you have 8 attention heads, you would have 8 separate key projections and 8 separate value projections. With GQA, you might have 8 query heads but only 2 or 4 key-value head pairs that are shared among the query heads. This means that instead of maintaining 8Q+8K+8V projection matrices (24 total), you might only need 8Q+2K+2V (12 total), cutting parameter count and computation significantly.
This reduction becomes increasingly significant in larger models. For instance, in a model with 32 attention heads and an embedding dimension of 4096, traditional multi-head attention would require approximately 402 million parameters for the attention mechanism alone, whereas GQA with 8 KV heads could reduce this to roughly 268 million parameters - a 33% reduction in memory footprint.
To understand the mechanism better, consider how attention works: each query computes similarity scores with all keys, then uses these scores to create a weighted sum of values. In GQA, multiple different queries compute attention scores against the same set of keys, and use these scores to attend to the same set of values. This maintains the expressivity of having multiple query perspectives while economizing on the key-value computations.
The efficiency gains from GQA become particularly evident during inference, where the KV cache (storing pre-computed key-value pairs for autoregressive generation) is often the primary bottleneck. By reducing the size of this cache through key-value sharing, GQA enables models to handle much longer context windows and generate text more efficiently, which is crucial for practical applications in production environments.
- Reduces memory footprint by decreasing the number of parameters needed for key and value projections, which is particularly important for larger models with billions of parameters. For example, in a model with an embedding dimension of 4096 and 32 heads, GQA with 8 KV groups can save approximately 33 million parameters per transformer layer. This reduction is achieved because traditional multi-head attention requires separate projection matrices for each attention head (Q, K, V), whereas GQA allows multiple query heads to share the same key and value projections, dramatically reducing the total parameter count across the model.
- Speeds up inference, especially in long-context models, by reducing the computational burden of generating and storing separate key-value pairs for each attention head. This is critical for server-side deployment where latency directly impacts user experience. During autoregressive generation, the KV cache (which stores previously computed key-value pairs) can become a memory bottleneck. By sharing KV projections across multiple query heads, GQA significantly reduces this cache size, allowing for faster token generation and enabling longer context handling without proportional memory increases.
- Offers nearly the same representational power as full multi-head attention but with significantly improved efficiency — a crucial trade-off for production deployment. Empirical studies show that models with GQA can achieve 95-99% of the performance of models with full multi-head attention while using considerably fewer resources. This minimal performance drop occurs because while the number of key-value pairs is reduced, the model still maintains its full capacity to generate diverse queries, preserving much of its ability to attend from different representational perspectives. The slight performance trade-off is well worth the substantial efficiency gains in most real-world applications.
- Used in LLaMA-2 to balance efficiency with performance, contributing to its ability to handle longer contexts while maintaining reasonable inference speeds. Other models like PaLM 2 and Claude have also adopted variants of this technique to scale efficiently. The implementation in LLaMA-2 specifically helped it achieve significant improvements in context window handling (up to 4K tokens) compared to its predecessor, while keeping inference costs manageable. In PaLM 2, a similar approach enabled efficient scaling to much longer contexts without the quadratic computational explosion that would occur with standard attention mechanisms.
Example: GQA Principle
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class GroupedQueryAttention(nn.Module):
def __init__(self, embed_dim=512, num_query_heads=8, num_kv_heads=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_query_heads = num_query_heads
self.num_kv_heads = num_kv_heads
self.head_dim = embed_dim // num_query_heads
# Ensure dimensions are compatible
assert self.head_dim * num_query_heads == embed_dim, "embed_dim must be divisible by num_query_heads"
assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
# Query projections (many heads)
self.q_proj = nn.Linear(embed_dim, embed_dim)
# Key/Value projections (fewer heads - shared)
self.k_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
self.v_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
# Output projection
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Dropout for attention weights
self.dropout = nn.Dropout(dropout)
# Group size: how many query heads share one kv head
self.group_size = num_query_heads // num_kv_heads
def forward(self, x, mask=None):
"""
Args:
x: Input tensor of shape (batch_size, sequence_length, embed_dim)
mask: Optional attention mask
Returns:
output: Tensor after self-attention of shape (batch_size, sequence_length, embed_dim)
"""
batch_size, seq_len, _ = x.size()
# Project inputs to queries, keys, and values
q = self.q_proj(x).view(batch_size, seq_len, self.num_query_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Transpose for attention computation
q = q.transpose(1, 2) # (batch_size, num_query_heads, seq_len, head_dim)
k = k.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
v = v.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
# Expand k and v to match the number of query heads through repetition
# Each group of query heads shares the same k and v
k_expanded = torch.repeat_interleave(k, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
v_expanded = torch.repeat_interleave(v, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
# Compute scaled dot-product attention
# (batch_size, num_query_heads, seq_len, seq_len)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
# Apply mask if provided (useful for preventing attention to padding tokens)
if mask is not None:
attn_weights = attn_weights.masked_fill(mask == 0, float("-inf"))
# Apply softmax and dropout
attn_weights = F.softmax(attn_weights, dim=-1)
attn_weights = self.dropout(attn_weights)
# Apply attention weights to values
# (batch_size, num_query_heads, seq_len, head_dim)
attn_output = torch.matmul(attn_weights, v_expanded)
# Reshape and apply output projection
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
output = self.out_proj(attn_output)
return output
def compare_attention_mechanisms(seq_len=1024, embed_dim=512):
"""Compare memory usage and speed between standard MHA and GQA"""
batch_size = 1
# Create inputs
x = torch.randn(batch_size, seq_len, embed_dim)
# Standard Multi-Head Attention (8 heads)
class StandardMHA(nn.Module):
def __init__(self):
super().__init__()
self.mha = nn.MultiheadAttention(embed_dim, num_heads=8, batch_first=True)
def forward(self, x):
return self.mha(x, x, x)[0]
standard_mha = StandardMHA()
# GQA with 8 query heads, 2 KV heads
gqa = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=2)
# GQA with 8 query heads, 4 KV heads
gqa2 = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=4)
# Measure memory and speed
results = {}
for name, model in [("Standard MHA (8 heads)", standard_mha),
("GQA (8Q, 2KV heads)", gqa),
("GQA (8Q, 4KV heads)", gqa2)]:
# Warm up
for _ in range(5):
_ = model(x)
# Measure time
torch.cuda.synchronize() if torch.cuda.is_available() else None
start_time = time.time()
for _ in range(10):
_ = model(x)
torch.cuda.synchronize() if torch.cuda.is_available() else None
end_time = time.time()
# Count parameters
param_count = sum(p.numel() for p in model.parameters())
results[name] = {
"time_per_run_ms": (end_time - start_time) * 100, # ms per 10 runs
"parameters": param_count
}
# Print results
print("Performance Comparison (sequence length = {})".format(seq_len))
print("=" * 50)
for name, metrics in results.items():
print(f"{name}:")
print(f" Time per 10 runs: {metrics['time_per_run_ms']:.2f} ms")
print(f" Parameters: {metrics['parameters']:,}")
print("-" * 50)
# Visualize KV cache size comparison
kv_cache_sizes = {
"Standard MHA": seq_len * 2 * embed_dim, # Full KV cache (8 heads)
"GQA (2 KV heads)": seq_len * 2 * (embed_dim // 4), # 1/4 the size (2 heads)
"GQA (4 KV heads)": seq_len * 2 * (embed_dim // 2), # 1/2 the size (4 heads)
}
plt.figure(figsize=(10, 5))
plt.bar(kv_cache_sizes.keys(), [size/1e6 for size in kv_cache_sizes.values()])
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Size Comparison')
for i, v in enumerate(kv_cache_sizes.values()):
plt.text(i, v/1e6 + 0.1, f"{v/1e6:.2f} MB", ha='center')
# Show how KV cache grows with sequence length
seq_lengths = [1024, 2048, 4096, 8192, 16384]
std_cache_sizes = [seq_len * 2 * embed_dim / 1e6 for seq_len in seq_lengths]
gqa_cache_sizes = [seq_len * 2 * (embed_dim // 4) / 1e6 for seq_len in seq_lengths]
plt.figure(figsize=(10, 5))
plt.plot(seq_lengths, std_cache_sizes, 'bo-', label='Standard MHA')
plt.plot(seq_lengths, gqa_cache_sizes, 'ro-', label='GQA (2 KV heads)')
plt.xlabel('Sequence Length')
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Growth with Sequence Length')
plt.legend()
plt.grid(True)
# Demonstrate usage with a simple example
seq_len, batch_size, embed_dim = 5, 1, 32
x = torch.randn(batch_size, seq_len, embed_dim)
gqa = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=8, num_kv_heads=2)
output = gqa(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
# Compare with standard attention
compare_attention_mechanisms(seq_len=2048, embed_dim=512)
# Basic visualization of attention patterns
def visualize_attention_pattern():
seq_len = 10
embed_dim = 64
x = torch.randn(1, seq_len, embed_dim)
model = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=4, num_kv_heads=2)
# Get attention weights by modifying forward pass temporarily
with torch.no_grad():
q = model.q_proj(x).view(1, seq_len, model.num_query_heads, model.head_dim)
k = model.k_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
v = model.v_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
k_expanded = torch.repeat_interleave(k, model.group_size, dim=1)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(model.head_dim, dtype=torch.float32))
attn_weights = F.softmax(attn_weights, dim=-1)
# Plot attention patterns for each head
fig, axes = plt.subplots(1, model.num_query_heads, figsize=(15, 3))
for i in range(model.num_query_heads):
im = axes[i].imshow(attn_weights[0, i].cpu().numpy(), cmap='viridis')
axes[i].set_title(f'Head {i+1}')
axes[i].set_xlabel('Key position')
axes[i].set_ylabel('Query position')
fig.colorbar(im, ax=axes.ravel().tolist())
plt.tight_layout()
plt.suptitle('Attention Patterns with GQA (notice shared patterns within groups)')
visualize_attention_pattern()Here's a comprehensive breakdown:
Core GQA Implementation
The code example Grouped Query Attention, an optimization technique that reduces memory usage and computational cost compared to standard multi-head attention.
Class Structure
- The
GroupedQueryAttentionclass inherits fromnn.Moduleand takes parameters for embedding dimension, number of query heads, number of key-value heads, and dropout rate. - The key innovation is that multiple query heads share the same key-value heads, reducing parameter count and memory footprint.
- Two compatibility assertions ensure:
- embedding dimension is divisible by the number of query heads
- query heads are divisible by key-value heads
Projection Layers
- Query projection: Full dimension projection (
self.q_proj) - Key/Value projections: Reduced dimension projections (
self.k_proj,self.v_proj) - Output projection: Maps attention output back to original dimensions
Forward Pass
- Projects input into queries, keys and values with appropriate dimensions
- Transposes tensors for attention computation
- The critical step: expands key and value tensors to match query heads through repetition using
torch.repeat_interleave - Computes scaled dot-product attention with softmax normalization
- Applies attention weights to values and reshapes the output
Performance Comparison Functions
The code includes utilities to demonstrate GQA's advantages:
compare_attention_mechanisms(): Benchmarks standard MHA against GQA variants with different head configurations measuring:- Execution time
- Parameter count
- KV cache size - critical for inference efficiency
- Visualization functions for KV cache size comparisons and growth with sequence length
- The
visualize_attention_pattern()function demonstrates how attention patterns appear in GQA, showing how multiple query heads share the same key-value pairs
Key Benefits Demonstrated
- Memory efficiency: Reduces parameters by sharing key-value projections
- Inference speed: Smaller KV cache allows for faster token generation
- Context length: Enables handling longer sequences with minimal memory growth
- Used in modern models: The implementation resembles approaches used in LLaMA-2, PaLM 2, and Claude
This implementation provides both a practical demonstration of GQA and tools to visualize its benefits over traditional attention mechanisms, particularly in terms of memory usage and computational efficiency while maintaining most of the representational power of full multi-head attention.
3.3.3 Attention Sparsity
In full self-attention, each token attends to every other token in the sequence. This creates a computational complexity that scales quadratically as O(n²) with sequence length, which becomes prohibitively expensive for long sequences (think 100k+ tokens). For context, processing a sequence of 100,000 tokens would require 10 billion attention computations per layer!
To understand why this is problematic, consider what happens as we scale: if we double our context length from 4K to 8K tokens, the computational work quadruples from 16 million to 64 million connections per layer. This quadratic scaling quickly becomes a bottleneck for both training and inference.
Additionally, the memory requirements for storing the attention matrix also scale quadratically. For a sequence of length n, we need to store an n×n attention matrix, which for long sequences can exceed available GPU memory. For example, a 32K token sequence would require approximately 4GB of memory just to store a single attention matrix in 32-bit precision.
Sparse attention techniques reduce this computational burden by attending only to the most relevant positions, effectively pruning unnecessary connections. This transforms the scaling from quadratic to nearly linear in many implementations. By strategically limiting which tokens can attend to which other tokens, these techniques dramatically reduce both computation and memory requirements.
The key insight behind sparse attention is that not all token-to-token interactions are equally important. Many language phenomena are local in nature, while certain special tokens may need global context. By exploiting this pattern, sparse attention can preserve most of the model's capabilities while eliminating many unnecessary computations.
Local attention
Each token attends only to its neighbors within a fixed window size (e.g., ±128 tokens). This creates a sliding window of attention that moves with each token position. For example, with a window size of 128, token at position 500 would attend to tokens from positions 372 to 628.
This approach works particularly well for tasks where nearby context is most relevant, such as speech recognition where phonemes relate strongly to adjacent sounds, or DNA analysis where nearby nucleotides often form functional units together. Local attention is also effective for text processing tasks where most semantic relationships occur between words that are relatively close to each other in the sequence.
The efficiency gains are substantial - the computational complexity becomes O(n×w), where w is the fixed window size. Since w is a constant (like 128 or 256), this effectively makes the attention mechanism scale linearly with sequence length rather than quadratically. For a sequence of 100,000 tokens with a window size of 256, this reduces computations from 10 billion to just 25.6 million - a 390x improvement.
However, local attention does have limitations - it struggles with tasks requiring long-range dependencies, such as document-level reasoning where important information may be separated by thousands of tokens. This is why more sophisticated sparse attention patterns often combine local attention with other mechanisms to capture both local and global relationships.
Block-sparse attention
Tokens attend within defined chunks or blocks, with occasional global tokens that can see across the entire sequence. This creates a sparse attention pattern where most tokens have limited vision but a few sentinel tokens maintain global context. These blocks can be arranged in various patterns - diagonal blocks for local attention, or more complex structures that allow for hierarchical information flow.
For example, in a block-sparse approach, a document might be divided into chunks of 512 tokens, with each chunk having internal full attention, plus dedicated "summary tokens" that can see across all chunks. This creates an information highway where local details are processed efficiently within blocks, while global information flows through the designated global tokens.
Additionally, some implementations use strided patterns where tokens can attend to blocks at regular intervals throughout the sequence, capturing periodic patterns or relationships. Others employ random sparse patterns that theoretically allow information to flow between any two positions through a small number of hops.
This hybrid approach preserves most of the modeling power of full attention while dramatically reducing computation. By carefully designing which blocks can attend to which others, these models achieve an attention complexity closer to O(n√n) or even O(n log n) rather than O(n²), enabling processing of much longer sequences with the same computational resources.
BigBird and Longformer
BigBird and Longformer implement sophisticated sparse attention patterns combining local windows, global tokens, and random connections. These architectures can efficiently scale to sequences of 4,000–8,000+ tokens with minimal loss in performance compared to full attention models.
BigBird, for example, combines three distinct attention patterns:
- Window attention: Each token attends to its local neighborhood (similar to the sliding window approach). This allows the model to capture local context effectively by focusing on nearby tokens. For instance, in a document about climate change, this helps the model understand phrases and nearby semantic connections by creating a focused attention window around each token, typically spanning 256-512 tokens in each direction.
- Global attention: Special tokens like [CLS] attend to all tokens and are attended to by all tokens, creating information highways. These global tokens serve as aggregation points that collect information from the entire sequence and distribute it back, enabling document-level understanding. For example, in a long scientific paper, the [CLS] token might gather key conclusions from various sections and make this information available to all other tokens, facilitating cross-referencing across distant parts of the document.
- Random attention: Each token attends to a small set of randomly selected tokens, which theoretically allows information to flow between any two positions in logarithmic steps. This random connectivity creates shortcuts across the document, ensuring information can propagate efficiently between distant sections. Mathematical proofs show that with just O(log n) random connections, information can flow between any two tokens in the sequence. In practice, this means even tokens separated by thousands of positions can exchange information through just a few intermediate connections.
This tri-directional attention mechanism achieves near-linear scaling while maintaining strong performance on long-document tasks like summarization and question answering. Importantly, BigBird maintains the theoretical property of "universal approximation" - it can represent any sequence-to-sequence function that full attention can, but with dramatically reduced computational requirements.
Longformer employs a similar approach but with a slightly different pattern, using a combination of sliding window attention with global attention for special tokens. It has demonstrated particular effectiveness in tasks requiring both local precision and document-level understanding, such as long-document question answering and multi-document summarization, where it can process inputs of 16,000+ tokens.
Code Example: Local Attention (Sliding Window)
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import time
class LocalAttention(nn.Module):
def __init__(self, dim, window_size=128):
super().__init__()
self.dim = dim
self.window_size = window_size
self.query_proj = nn.Linear(dim, dim)
self.key_proj = nn.Linear(dim, dim)
self.value_proj = nn.Linear(dim, dim)
self.output_proj = nn.Linear(dim, dim)
self.scaling = dim ** -0.5
def forward(self, x):
B, T, D = x.size()
# Project inputs to queries, keys, values
queries = self.query_proj(x) * self.scaling # [B, T, D]
keys = self.key_proj(x) # [B, T, D]
values = self.value_proj(x) # [B, T, D]
# Initialize output tensor
output = torch.zeros_like(x)
# Compute local attention for each position
for i in range(T):
# Define local window boundaries
start = max(0, i - self.window_size)
end = min(T, i + self.window_size + 1)
# Extract local context
local_keys = keys[:, start:end, :] # [B, window_size*2, D]
local_values = values[:, start:end, :] # [B, window_size*2, D]
# Current query
query = queries[:, i:i+1, :] # [B, 1, D]
# Compute attention scores
scores = torch.bmm(query, local_keys.transpose(1, 2)) # [B, 1, window_size*2]
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, 1, window_size*2]
# Weight values by attention
context = torch.bmm(attn_weights, local_values) # [B, 1, D]
# Store in output
output[:, i:i+1, :] = context
return self.output_proj(output)
def naive_local_attention(x, window=2):
"""A simple implementation of local attention for educational purposes"""
B, T, D = x.size()
outputs = []
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
context = x[:, start:end, :]
weights = F.softmax(torch.bmm(x[:, i:i+1, :], context.transpose(1,2)), dim=-1)
out = torch.bmm(weights, context)
outputs.append(out)
return torch.cat(outputs, dim=1)
def vectorized_local_attention(x, window=2):
"""A more efficient implementation using vectorized operations"""
B, T, D = x.size()
# Create attention mask to implement sliding window
mask = torch.zeros(T, T, device=x.device)
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
mask[i, start:end] = 1
# Compute attention scores
scores = torch.bmm(x, x.transpose(1, 2)) # [B, T, T]
# Apply mask (setting padded values to -inf before softmax)
scores = scores.masked_fill(mask.unsqueeze(0) == 0, -1e9)
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, T, T]
# Weight values by attention
output = torch.bmm(attn_weights, x) # [B, T, D]
return output
def compare_performance(seq_lengths=[10, 50, 100, 200], window=2):
"""Compare performance of different local attention implementations"""
results = {'naive': [], 'vectorized': [], 'optimized': []}
for seq_len in seq_lengths:
# Generate random input tensor
x = torch.randn(1, seq_len, 64)
# Naive implementation
start_time = time.time()
naive_local_attention(x, window)
naive_time = time.time() - start_time
results['naive'].append(naive_time)
# Vectorized implementation
start_time = time.time()
vectorized_local_attention(x, window)
vectorized_time = time.time() - start_time
results['vectorized'].append(vectorized_time)
# Optimized implementation
model = LocalAttention(64, window)
start_time = time.time()
model(x)
optimized_time = time.time() - start_time
results['optimized'].append(optimized_time)
print(f"Sequence length {seq_len}:")
print(f" Naive: {naive_time:.5f}s")
print(f" Vectorized: {vectorized_time:.5f}s")
print(f" Optimized: {optimized_time:.5f}s")
# Plot results
plt.figure(figsize=(10, 6))
plt.plot(seq_lengths, results['naive'], 'o-', label='Naive')
plt.plot(seq_lengths, results['vectorized'], 's-', label='Vectorized')
plt.plot(seq_lengths, results['optimized'], '^-', label='Optimized')
plt.xlabel('Sequence Length')
plt.ylabel('Time (s)')
plt.title('Performance Comparison of Local Attention Implementations')
plt.legend()
plt.grid(True)
plt.show()
def visualize_attention_pattern(window=2, seq_len=10):
"""Visualize the sparse attention pattern created by local attention"""
attention_mask = torch.zeros(seq_len, seq_len)
for i in range(seq_len):
start = max(0, i - window)
end = min(seq_len, i + window + 1)
attention_mask[i, start:end] = 1
plt.figure(figsize=(8, 8))
plt.imshow(attention_mask, cmap='Blues')
plt.title(f'Local Attention Pattern (Window Size = {window})')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar(label='Attention Connection')
for i in range(seq_len):
for j in range(seq_len):
color = 'white' if attention_mask[i, j] > 0 else 'none'
plt.text(j, i, '1' if attention_mask[i, j] > 0 else '0',
ha='center', va='center', color=color)
plt.tight_layout()
plt.show()
# Example
if __name__ == "__main__":
# Basic functionality test
x = torch.randn(1, 6, 16)
model = LocalAttention(16, window_size=2)
out = model(x)
print(f"Input shape: {x.shape}, Output shape: {out.shape}")
# Compare implementations
compare_performance([10, 50, 100, 200], window=2)
# Visualize the attention pattern
visualize_attention_pattern(window=2, seq_len=10)Comprehensive Breakdown: Local Attention Implementation
This code example provides a complete toolkit for understanding, implementing and analyzing local attention mechanisms. Here's a detailed breakdown:
1. Core Implementations
- LocalAttention Class: A proper PyTorch module implementation with:
- Dedicated projection layers for queries, keys, and valuesDedicated projection layers for queries, keys, and values
- Window-based sliding attention with configurable window sizeWindow-based sliding attention with configurable window size
- Proper scaling factor (1/√d) for stable gradientsProper scaling factor (1/√d) for stable gradients
- Final output projection as in standard attentionFinal output projection as in standard attention
- Naive Implementation: The original function that:
- Processes each position sequentiallyProcesses each position sequentially
- Demonstrates the core sliding window concept clearlyDemonstrates the core sliding window concept clearly
- Uses simple tensor operations for educational purposesUses simple tensor operations for educational purposes
- Vectorized Implementation: A more efficient approach that:
- Uses a mask tensor to implement the sliding window patternUses a mask tensor to implement the sliding window pattern
- Computes all attention scores at onceComputes all attention scores at once
- Avoids explicit loops over sequence positionsAvoids explicit loops over sequence positions
2. Analysis Tools
- Performance Comparison Function: Benchmarks all three implementations:
- Measures execution time across different sequence lengthsMeasures execution time across different sequence lengths
- Generates performance plots to visualize scaling behaviorGenerates performance plots to visualize scaling behavior
- Demonstrates how vectorized operations improve efficiencyDemonstrates how vectorized operations improve efficiency
- Visualization Function: Illustrates the sparse attention pattern:
- Creates a visual representation of which tokens attend to which othersCreates a visual representation of which tokens attend to which others
- Shows the diagonal band pattern characteristic of local attentionShows the diagonal band pattern characteristic of local attention
- Helps intuitive understanding of how information flows in the modelHelps intuitive understanding of how information flows in the model
3. Key Technical Insights
- Masking Technique: The code demonstrates how to create and apply attention masks to restrict which tokens can attend to which others
- Computational Efficiency: Shows how the computational complexity becomes O(n·w) instead of O(n²), where w is the window size
- Implementation Trade-offs: Illustrates the balance between code clarity (naive implementation) and computational efficiency (vectorized implementation)
This implementation provides both theoretical understanding and practical tools for working with local attention, a key technique for making transformers more efficient with long sequences. The visualization and comparison functions make it especially valuable for educational purposes.
3.3.4 Why These Matter
SwiGLU (Swish-Gated Linear Unit) significantly improves learning dynamics, giving models richer representations with little extra computational cost. This sophisticated activation function combines the benefits of gating mechanisms with a simple weighted identity connection, allowing for more effective gradient flow during training. By replacing traditional ReLU or GELU activations, SwiGLU enables models to learn more complex patterns while maintaining computational efficiency.
The mathematical formulation of SwiGLU involves multiplying a linear projection of the input with a sigmoid-weighted version of another projection, creating a smooth, differentiable pathway for gradients that helps prevent vanishing gradient problems. Models using SwiGLU typically converge faster and achieve better performance across various natural language processing tasks, making it a preferred choice in modern LLM architectures like PaLM and Gemini.
GQA (Grouped Query Attention) makes attention mechanisms substantially more efficient, reducing memory use without significant accuracy loss. This innovative technique groups queries together to share the same keys and values, dramatically reducing the memory footprint during inference. Unlike standard multi-head attention that requires separate key-value pairs for each attention head (creating a parameter explosion), GQA significantly cuts down on parameters while preserving most of the model's reasoning capabilities.
This approach creates a middle ground between multi-head attention (MHA) and multi-query attention (MQA), finding an optimal balance between parameter efficiency and model capacity. In practice, GQA can reduce the key-value cache memory requirements by 2-4x compared to standard attention while maintaining 95-99% of the model's performance, making it possible to deploy larger models on the same hardware or increase batch sizes during inference. Models like PaLM 2 and Claude have successfully implemented GQA as a core architectural improvement.
Sparse attention fundamentally transforms how LLMs can handle very long contexts without suffering from quadratic computational blow-ups. Instead of having each token attend to every other token (which scales as O(n²) with sequence length), sparse attention patterns like local, dilated, or longformer attention enable selective focus on only the most relevant tokens. This reduces computational complexity to O(n) or O(n log n), making it feasible to process documents with thousands or even tens of thousands of tokens.
Local attention, as shown in the code example above, restricts each token to attend only to a window of neighboring tokens. Dilated attention extends this by allowing tokens to attend to positions at various distances, creating a wider receptive field without increasing computation proportionally. More advanced sparse attention patterns like Reformer's LSH attention or Longformer's global+local attention combine different strategies to balance computational efficiency with model capacity. These approaches have enabled breakthroughs in long-context models that can process entire books, codebases, or lengthy conversations while maintaining coherent understanding throughout the document.
Together, these architectural refinements are why today's LLMs can be faster, leaner, and more scalable than early transformers. They represent critical engineering breakthroughs that have transformed theoretical research models into practical, deployable systems capable of handling real-world tasks with unprecedented efficiency.
3.3 Advanced Architectures: SwiGLU, GQA, Attention Sparsity
As transformers evolved from small research models to trillion-parameter giants, engineers discovered that even small changes to internal components can yield big improvements in efficiency and performance. These architectural innovations became increasingly critical as models grew in size and complexity, addressing challenges in computation, memory usage, and training stability. In this section, we'll look at three important innovations that have fundamentally changed how modern LLMs are designed:
- SwiGLU activation functions (improving feedforward networks inside transformer blocks). These specialized activation functions replace traditional ReLU or GELU activations with a more sophisticated gating mechanism that allows for smoother gradient flow during training. By incorporating both multiplicative interactions and non-linear transformations, SwiGLU enables the model to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Grouped Query Attention (GQA) (making attention faster without losing much accuracy). This clever optimization reduces the memory and computational requirements of the attention mechanism by allowing multiple query heads to share the same key and value projections. This significantly decreases both the parameter count and memory bandwidth needed during inference, addressing one of the major bottlenecks in large language model deployment.
- Attention sparsity techniques (reducing compute by ignoring unnecessary connections). These approaches recognize that not all tokens need to attend to every other token in a sequence, especially for long documents. By strategically limiting which tokens can attend to which others, sparse attention patterns can reduce the quadratic complexity of standard attention to near-linear, enabling models to process much longer contexts efficiently.
These optimizations aren't just academic curiosities — they power modern models like LLaMA, Mistral, and GPT-5. Without these architectural advances, today's state-of-the-art models would be prohibitively expensive to train and deploy. Each innovation represents a careful balance between model capability and computational efficiency, addressing specific bottlenecks that emerge at different scales of model size and context length.
3.3.1 SwiGLU (Switched Gated Linear Units)
Every transformer block contains a feedforward network (FFN) after attention. Traditionally, this FFN uses a ReLU or GELU activation. But research found that SwiGLU (a variant of GLU) yields smoother optimization and better performance at scale.
This performance improvement is largely due to the gating mechanism that allows information to flow more selectively through the network, effectively letting the model adaptively control which features are emphasized in each forward pass. Unlike traditional activation functions that apply the same transformation to all inputs, SwiGLU introduces a dynamic, input-dependent filtering mechanism that can emphasize or suppress different aspects of the representation based on the content itself.
In technical terms, SwiGLU combines the benefits of multiplicative interactions (from gates) with non-linear transformations, creating a more expressive computation unit. The gating component (using the swish/SiLU function) produces values between 0 and 1 that act as "soft switches," controlling how much information from each dimension passes through. This adaptive behavior allows the model to create more complex functional mappings with fewer parameters, resulting in improved gradient flow during training and more efficient use of model capacity.
How SwiGLU works:
- Split the hidden dimension into two parts - this creates parallel pathways through the network, allowing the model to learn different aspects of the input simultaneously. This splitting mechanism is crucial because it enables the network to process information along two distinct channels that can later be recombined in a meaningful way. Each pathway can specialize in capturing different features or patterns in the data, similar to how different neurons in the brain might respond to different aspects of visual stimuli.
- Apply a linear transformation to both parts - each pathway gets its own weight matrix (W1 and W2), enabling the network to learn different feature mappings. These linear transformations are fully learned during training and adapt to the specific task at hand. W1 typically projects the input into a space suitable for gating decisions, while W2 creates representations that will be selectively passed through based on those gates. The separation of these transformations allows the model to develop specialized feature detectors that work in tandem.
- Pass one through a sigmoid (or swish) gate - the swish activation (x * sigmoid(x)) provides a smoother gradient than ReLU and allows for some small negative values to pass through, which helps prevent "dying neurons" that can occur with ReLU. The swish function combines the benefits of sigmoid (bounded output) with the non-saturating behavior of ReLU, creating an activation function with better mathematical properties for optimization. This smoother gradient flow helps address the vanishing gradient problem that can plague deep neural networks during training.
- Multiply the two parts together - this multiplicative interaction creates a gating mechanism where the output of the swish function controls how much of the second linear transformation's output passes through. This dynamic gating allows the network to selectively amplify or suppress different features based on the input, leading to more expressive representations. The multiplication operation enables complex, content-dependent filtering of information - effectively allowing some dimensions to be emphasized while others are attenuated based on the specific input pattern, creating a form of adaptive computation.
- The mathematical formula is: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication. In practice, this can be implemented efficiently in modern deep learning frameworks by computing both transformations in parallel and then combining them with a simple Hadamard product. This formulation creates a powerful non-linear transformation that combines the benefits of gating mechanisms (like those in LSTMs and GRUs) with the parallelizability and computational efficiency of feed-forward networks.
PyTorch Example: SwiGLU vs ReLU
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class SwiGLU(nn.Module):
"""
SwiGLU activation function as used in modern LLMs like LLaMA and PaLM.
Uses the Swish/SiLU gating mechanism for better gradient properties.
"""
def __init__(self, dim_in, dim_out):
super().__init__()
self.W1 = nn.Linear(dim_in, dim_out) # Transformation for the gate
self.W2 = nn.Linear(dim_in, dim_out) # Transformation for the content
def forward(self, x):
# SiLU (Swish) activation for gating: x * sigmoid(x)
return F.silu(self.W1(x)) * self.W2(x)
class StandardFFN(nn.Module):
"""
Standard Feedforward Network with ReLU activation
as used in original Transformer architecture.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.ReLU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
class GELUBasedFFN(nn.Module):
"""
Feedforward Network with GELU activation
as used in models like BERT and early GPT versions.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.GELU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
# Hyperparameters
batch_size = 8
seq_len = 32
embed_dim = 256
hidden_dim = 1024
output_dim = 256
# Create example input
x = torch.randn(batch_size, seq_len, embed_dim)
# Initialize models
swiglu = SwiGLU(embed_dim, hidden_dim)
relu_ffn = StandardFFN(embed_dim, hidden_dim, output_dim)
gelu_ffn = GELUBasedFFN(embed_dim, hidden_dim, output_dim)
# Model outputs for comparison
with torch.no_grad():
# Timing comparisons
start = time.time()
swiglu_out = swiglu(x)
swiglu_time = time.time() - start
start = time.time()
relu_out = relu_ffn(x)
relu_time = time.time() - start
start = time.time()
gelu_out = gelu_ffn(x)
gelu_time = time.time() - start
print(f"SwiGLU output shape: {swiglu_out.shape}")
print(f"ReLU FFN output shape: {relu_out.shape}")
print(f"GELU FFN output shape: {gelu_out.shape}")
# Print timing results
print(f"\nForward pass timing:")
print(f"SwiGLU: {swiglu_time*1000:.2f}ms")
print(f"ReLU: {relu_time*1000:.2f}ms")
print(f"GELU: {gelu_time*1000:.2f}ms")
# Print sample outputs
print("\nSample outputs (first 5 values):")
print(f"SwiGLU: {swiglu_out[0, 0, :5].numpy()}")
print(f"ReLU: {relu_out[0, 0, :5].numpy()}")
print(f"GELU: {gelu_out[0, 0, :5].numpy()}")
# Visualize activation functions for comparison
def swish(x):
return x * torch.sigmoid(x)
def relu(x):
return torch.maximum(torch.zeros_like(x), x)
def gelu(x):
return 0.5 * x * (1 + torch.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))
x_range = torch.linspace(-5, 5, 1000)
plt.figure(figsize=(10, 6))
plt.plot(x_range, swish(x_range), label='Swish/SiLU (used in SwiGLU)')
plt.plot(x_range, relu(x_range), label='ReLU')
plt.plot(x_range, gelu(x_range), label='GELU')
plt.grid(True)
plt.legend()
plt.title('Comparison of Activation Functions')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.savefig('activation_comparison.png')
plt.close()
# Simple gradient flow demonstration
def compare_gradient_flow():
"""Compare gradient flow through different activation functions"""
# Create synthetic data
x = torch.randn(100, 32, requires_grad=True)
y = torch.randn(100, 32)
models = {
'SwiGLU': nn.Sequential(
nn.Linear(32, 128),
SwiGLU(128, 128),
nn.Linear(128, 32)
),
'ReLU': nn.Sequential(
nn.Linear(32, 128),
nn.ReLU(),
nn.Linear(128, 32)
),
'GELU': nn.Sequential(
nn.Linear(32, 128),
nn.GELU(),
nn.Linear(128, 32)
)
}
results = {}
for name, model in models.items():
# Forward pass
pred = model(x)
loss = torch.nn.functional.mse_loss(pred, y)
# Backward pass
loss.backward()
# Record gradient statistics
grad_norms = []
for param in model.parameters():
if param.grad is not None:
grad_norms.append(param.grad.norm().item())
results[name] = {
'mean': np.mean(grad_norms),
'std': np.std(grad_norms),
'min': np.min(grad_norms),
'max': np.max(grad_norms)
}
print("\nGradient statistics after one backward pass:")
for name, stats in results.items():
print(f"{name}: mean={stats['mean']:.6f}, std={stats['std']:.6f}, min={stats['min']:.6f}, max={stats['max']:.6f}")
compare_gradient_flow()
Breakdown of SwiGLU
How SwiGLU Works
The implementation breaks down into these key steps:
- Two Parallel Pathways: SwiGLU splits the computation into two parallel linear transformations (W1 and W2).
- Gating Mechanism: One pathway (W1) is passed through a swish/SiLU activation function (x * sigmoid(x)), which provides smoother gradients than ReLU.
- Multiplicative Interaction: The outputs are multiplied together element-wise, allowing the swish-activated path to act as a gate that controls how much of the other path's output passes through.
- Mathematical Formula: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication.
Key Advantages of SwiGLU
- Smoother Gradients: The swish function provides better gradient flow during backpropagation, addressing the vanishing gradient problem that can affect deep networks.
- Dynamic Feature Selection: The gating mechanism allows the network to selectively emphasize or suppress different features based on input content.
- Better Performance Per Parameter: SwiGLU enables models to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Improved Training Dynamics: The smoother activation function and gating mechanism result in more stable and effective training, especially in deep networks.
Code Implementation Details
- The example code demonstrates SwiGLU alongside traditional ReLU and GELU-based feedforward networks for comparison.
- It includes timing comparisons to show computational efficiency differences.
- The visualization of activation functions illustrates how swish/SiLU differs from ReLU and GELU in shape and smoothness.
- The gradient flow demonstration highlights how SwiGLU affects gradient statistics during backpropagation.
This implementation showcases why SwiGLU has become a critical component in modern LLM architectures, offering a better balance of expressivity, computational efficiency, and training stability compared to earlier alternatives.
3.3.2 Grouped Query Attention (GQA)
Standard multi-head attention has a computational cost that grows with the number of heads. To optimize this, Grouped Query Attention (GQA) was introduced. This optimization technique addresses both memory usage and computational efficiency while maintaining model quality.
Key idea: Instead of each query head having its own set of key/value heads, multiple query heads can share the same key/value projections. This sharing mechanism substantially reduces the number of parameters and computation required during inference while preserving the model's ability to learn diverse representations.
The fundamental insight behind GQA is that we can achieve a better balance between computational efficiency and model expressivity by decoupling the number of query heads from the number of key-value heads. This allows models to maintain the benefits of multi-perspective querying while reducing the overhead associated with generating and storing separate key-value pairs for each head.
In traditional multi-head attention, if you have 8 attention heads, you would have 8 separate key projections and 8 separate value projections. With GQA, you might have 8 query heads but only 2 or 4 key-value head pairs that are shared among the query heads. This means that instead of maintaining 8Q+8K+8V projection matrices (24 total), you might only need 8Q+2K+2V (12 total), cutting parameter count and computation significantly.
This reduction becomes increasingly significant in larger models. For instance, in a model with 32 attention heads and an embedding dimension of 4096, traditional multi-head attention would require approximately 402 million parameters for the attention mechanism alone, whereas GQA with 8 KV heads could reduce this to roughly 268 million parameters - a 33% reduction in memory footprint.
To understand the mechanism better, consider how attention works: each query computes similarity scores with all keys, then uses these scores to create a weighted sum of values. In GQA, multiple different queries compute attention scores against the same set of keys, and use these scores to attend to the same set of values. This maintains the expressivity of having multiple query perspectives while economizing on the key-value computations.
The efficiency gains from GQA become particularly evident during inference, where the KV cache (storing pre-computed key-value pairs for autoregressive generation) is often the primary bottleneck. By reducing the size of this cache through key-value sharing, GQA enables models to handle much longer context windows and generate text more efficiently, which is crucial for practical applications in production environments.
- Reduces memory footprint by decreasing the number of parameters needed for key and value projections, which is particularly important for larger models with billions of parameters. For example, in a model with an embedding dimension of 4096 and 32 heads, GQA with 8 KV groups can save approximately 33 million parameters per transformer layer. This reduction is achieved because traditional multi-head attention requires separate projection matrices for each attention head (Q, K, V), whereas GQA allows multiple query heads to share the same key and value projections, dramatically reducing the total parameter count across the model.
- Speeds up inference, especially in long-context models, by reducing the computational burden of generating and storing separate key-value pairs for each attention head. This is critical for server-side deployment where latency directly impacts user experience. During autoregressive generation, the KV cache (which stores previously computed key-value pairs) can become a memory bottleneck. By sharing KV projections across multiple query heads, GQA significantly reduces this cache size, allowing for faster token generation and enabling longer context handling without proportional memory increases.
- Offers nearly the same representational power as full multi-head attention but with significantly improved efficiency — a crucial trade-off for production deployment. Empirical studies show that models with GQA can achieve 95-99% of the performance of models with full multi-head attention while using considerably fewer resources. This minimal performance drop occurs because while the number of key-value pairs is reduced, the model still maintains its full capacity to generate diverse queries, preserving much of its ability to attend from different representational perspectives. The slight performance trade-off is well worth the substantial efficiency gains in most real-world applications.
- Used in LLaMA-2 to balance efficiency with performance, contributing to its ability to handle longer contexts while maintaining reasonable inference speeds. Other models like PaLM 2 and Claude have also adopted variants of this technique to scale efficiently. The implementation in LLaMA-2 specifically helped it achieve significant improvements in context window handling (up to 4K tokens) compared to its predecessor, while keeping inference costs manageable. In PaLM 2, a similar approach enabled efficient scaling to much longer contexts without the quadratic computational explosion that would occur with standard attention mechanisms.
Example: GQA Principle
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class GroupedQueryAttention(nn.Module):
def __init__(self, embed_dim=512, num_query_heads=8, num_kv_heads=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_query_heads = num_query_heads
self.num_kv_heads = num_kv_heads
self.head_dim = embed_dim // num_query_heads
# Ensure dimensions are compatible
assert self.head_dim * num_query_heads == embed_dim, "embed_dim must be divisible by num_query_heads"
assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
# Query projections (many heads)
self.q_proj = nn.Linear(embed_dim, embed_dim)
# Key/Value projections (fewer heads - shared)
self.k_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
self.v_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
# Output projection
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Dropout for attention weights
self.dropout = nn.Dropout(dropout)
# Group size: how many query heads share one kv head
self.group_size = num_query_heads // num_kv_heads
def forward(self, x, mask=None):
"""
Args:
x: Input tensor of shape (batch_size, sequence_length, embed_dim)
mask: Optional attention mask
Returns:
output: Tensor after self-attention of shape (batch_size, sequence_length, embed_dim)
"""
batch_size, seq_len, _ = x.size()
# Project inputs to queries, keys, and values
q = self.q_proj(x).view(batch_size, seq_len, self.num_query_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Transpose for attention computation
q = q.transpose(1, 2) # (batch_size, num_query_heads, seq_len, head_dim)
k = k.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
v = v.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
# Expand k and v to match the number of query heads through repetition
# Each group of query heads shares the same k and v
k_expanded = torch.repeat_interleave(k, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
v_expanded = torch.repeat_interleave(v, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
# Compute scaled dot-product attention
# (batch_size, num_query_heads, seq_len, seq_len)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
# Apply mask if provided (useful for preventing attention to padding tokens)
if mask is not None:
attn_weights = attn_weights.masked_fill(mask == 0, float("-inf"))
# Apply softmax and dropout
attn_weights = F.softmax(attn_weights, dim=-1)
attn_weights = self.dropout(attn_weights)
# Apply attention weights to values
# (batch_size, num_query_heads, seq_len, head_dim)
attn_output = torch.matmul(attn_weights, v_expanded)
# Reshape and apply output projection
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
output = self.out_proj(attn_output)
return output
def compare_attention_mechanisms(seq_len=1024, embed_dim=512):
"""Compare memory usage and speed between standard MHA and GQA"""
batch_size = 1
# Create inputs
x = torch.randn(batch_size, seq_len, embed_dim)
# Standard Multi-Head Attention (8 heads)
class StandardMHA(nn.Module):
def __init__(self):
super().__init__()
self.mha = nn.MultiheadAttention(embed_dim, num_heads=8, batch_first=True)
def forward(self, x):
return self.mha(x, x, x)[0]
standard_mha = StandardMHA()
# GQA with 8 query heads, 2 KV heads
gqa = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=2)
# GQA with 8 query heads, 4 KV heads
gqa2 = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=4)
# Measure memory and speed
results = {}
for name, model in [("Standard MHA (8 heads)", standard_mha),
("GQA (8Q, 2KV heads)", gqa),
("GQA (8Q, 4KV heads)", gqa2)]:
# Warm up
for _ in range(5):
_ = model(x)
# Measure time
torch.cuda.synchronize() if torch.cuda.is_available() else None
start_time = time.time()
for _ in range(10):
_ = model(x)
torch.cuda.synchronize() if torch.cuda.is_available() else None
end_time = time.time()
# Count parameters
param_count = sum(p.numel() for p in model.parameters())
results[name] = {
"time_per_run_ms": (end_time - start_time) * 100, # ms per 10 runs
"parameters": param_count
}
# Print results
print("Performance Comparison (sequence length = {})".format(seq_len))
print("=" * 50)
for name, metrics in results.items():
print(f"{name}:")
print(f" Time per 10 runs: {metrics['time_per_run_ms']:.2f} ms")
print(f" Parameters: {metrics['parameters']:,}")
print("-" * 50)
# Visualize KV cache size comparison
kv_cache_sizes = {
"Standard MHA": seq_len * 2 * embed_dim, # Full KV cache (8 heads)
"GQA (2 KV heads)": seq_len * 2 * (embed_dim // 4), # 1/4 the size (2 heads)
"GQA (4 KV heads)": seq_len * 2 * (embed_dim // 2), # 1/2 the size (4 heads)
}
plt.figure(figsize=(10, 5))
plt.bar(kv_cache_sizes.keys(), [size/1e6 for size in kv_cache_sizes.values()])
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Size Comparison')
for i, v in enumerate(kv_cache_sizes.values()):
plt.text(i, v/1e6 + 0.1, f"{v/1e6:.2f} MB", ha='center')
# Show how KV cache grows with sequence length
seq_lengths = [1024, 2048, 4096, 8192, 16384]
std_cache_sizes = [seq_len * 2 * embed_dim / 1e6 for seq_len in seq_lengths]
gqa_cache_sizes = [seq_len * 2 * (embed_dim // 4) / 1e6 for seq_len in seq_lengths]
plt.figure(figsize=(10, 5))
plt.plot(seq_lengths, std_cache_sizes, 'bo-', label='Standard MHA')
plt.plot(seq_lengths, gqa_cache_sizes, 'ro-', label='GQA (2 KV heads)')
plt.xlabel('Sequence Length')
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Growth with Sequence Length')
plt.legend()
plt.grid(True)
# Demonstrate usage with a simple example
seq_len, batch_size, embed_dim = 5, 1, 32
x = torch.randn(batch_size, seq_len, embed_dim)
gqa = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=8, num_kv_heads=2)
output = gqa(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
# Compare with standard attention
compare_attention_mechanisms(seq_len=2048, embed_dim=512)
# Basic visualization of attention patterns
def visualize_attention_pattern():
seq_len = 10
embed_dim = 64
x = torch.randn(1, seq_len, embed_dim)
model = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=4, num_kv_heads=2)
# Get attention weights by modifying forward pass temporarily
with torch.no_grad():
q = model.q_proj(x).view(1, seq_len, model.num_query_heads, model.head_dim)
k = model.k_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
v = model.v_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
k_expanded = torch.repeat_interleave(k, model.group_size, dim=1)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(model.head_dim, dtype=torch.float32))
attn_weights = F.softmax(attn_weights, dim=-1)
# Plot attention patterns for each head
fig, axes = plt.subplots(1, model.num_query_heads, figsize=(15, 3))
for i in range(model.num_query_heads):
im = axes[i].imshow(attn_weights[0, i].cpu().numpy(), cmap='viridis')
axes[i].set_title(f'Head {i+1}')
axes[i].set_xlabel('Key position')
axes[i].set_ylabel('Query position')
fig.colorbar(im, ax=axes.ravel().tolist())
plt.tight_layout()
plt.suptitle('Attention Patterns with GQA (notice shared patterns within groups)')
visualize_attention_pattern()Here's a comprehensive breakdown:
Core GQA Implementation
The code example Grouped Query Attention, an optimization technique that reduces memory usage and computational cost compared to standard multi-head attention.
Class Structure
- The
GroupedQueryAttentionclass inherits fromnn.Moduleand takes parameters for embedding dimension, number of query heads, number of key-value heads, and dropout rate. - The key innovation is that multiple query heads share the same key-value heads, reducing parameter count and memory footprint.
- Two compatibility assertions ensure:
- embedding dimension is divisible by the number of query heads
- query heads are divisible by key-value heads
Projection Layers
- Query projection: Full dimension projection (
self.q_proj) - Key/Value projections: Reduced dimension projections (
self.k_proj,self.v_proj) - Output projection: Maps attention output back to original dimensions
Forward Pass
- Projects input into queries, keys and values with appropriate dimensions
- Transposes tensors for attention computation
- The critical step: expands key and value tensors to match query heads through repetition using
torch.repeat_interleave - Computes scaled dot-product attention with softmax normalization
- Applies attention weights to values and reshapes the output
Performance Comparison Functions
The code includes utilities to demonstrate GQA's advantages:
compare_attention_mechanisms(): Benchmarks standard MHA against GQA variants with different head configurations measuring:- Execution time
- Parameter count
- KV cache size - critical for inference efficiency
- Visualization functions for KV cache size comparisons and growth with sequence length
- The
visualize_attention_pattern()function demonstrates how attention patterns appear in GQA, showing how multiple query heads share the same key-value pairs
Key Benefits Demonstrated
- Memory efficiency: Reduces parameters by sharing key-value projections
- Inference speed: Smaller KV cache allows for faster token generation
- Context length: Enables handling longer sequences with minimal memory growth
- Used in modern models: The implementation resembles approaches used in LLaMA-2, PaLM 2, and Claude
This implementation provides both a practical demonstration of GQA and tools to visualize its benefits over traditional attention mechanisms, particularly in terms of memory usage and computational efficiency while maintaining most of the representational power of full multi-head attention.
3.3.3 Attention Sparsity
In full self-attention, each token attends to every other token in the sequence. This creates a computational complexity that scales quadratically as O(n²) with sequence length, which becomes prohibitively expensive for long sequences (think 100k+ tokens). For context, processing a sequence of 100,000 tokens would require 10 billion attention computations per layer!
To understand why this is problematic, consider what happens as we scale: if we double our context length from 4K to 8K tokens, the computational work quadruples from 16 million to 64 million connections per layer. This quadratic scaling quickly becomes a bottleneck for both training and inference.
Additionally, the memory requirements for storing the attention matrix also scale quadratically. For a sequence of length n, we need to store an n×n attention matrix, which for long sequences can exceed available GPU memory. For example, a 32K token sequence would require approximately 4GB of memory just to store a single attention matrix in 32-bit precision.
Sparse attention techniques reduce this computational burden by attending only to the most relevant positions, effectively pruning unnecessary connections. This transforms the scaling from quadratic to nearly linear in many implementations. By strategically limiting which tokens can attend to which other tokens, these techniques dramatically reduce both computation and memory requirements.
The key insight behind sparse attention is that not all token-to-token interactions are equally important. Many language phenomena are local in nature, while certain special tokens may need global context. By exploiting this pattern, sparse attention can preserve most of the model's capabilities while eliminating many unnecessary computations.
Local attention
Each token attends only to its neighbors within a fixed window size (e.g., ±128 tokens). This creates a sliding window of attention that moves with each token position. For example, with a window size of 128, token at position 500 would attend to tokens from positions 372 to 628.
This approach works particularly well for tasks where nearby context is most relevant, such as speech recognition where phonemes relate strongly to adjacent sounds, or DNA analysis where nearby nucleotides often form functional units together. Local attention is also effective for text processing tasks where most semantic relationships occur between words that are relatively close to each other in the sequence.
The efficiency gains are substantial - the computational complexity becomes O(n×w), where w is the fixed window size. Since w is a constant (like 128 or 256), this effectively makes the attention mechanism scale linearly with sequence length rather than quadratically. For a sequence of 100,000 tokens with a window size of 256, this reduces computations from 10 billion to just 25.6 million - a 390x improvement.
However, local attention does have limitations - it struggles with tasks requiring long-range dependencies, such as document-level reasoning where important information may be separated by thousands of tokens. This is why more sophisticated sparse attention patterns often combine local attention with other mechanisms to capture both local and global relationships.
Block-sparse attention
Tokens attend within defined chunks or blocks, with occasional global tokens that can see across the entire sequence. This creates a sparse attention pattern where most tokens have limited vision but a few sentinel tokens maintain global context. These blocks can be arranged in various patterns - diagonal blocks for local attention, or more complex structures that allow for hierarchical information flow.
For example, in a block-sparse approach, a document might be divided into chunks of 512 tokens, with each chunk having internal full attention, plus dedicated "summary tokens" that can see across all chunks. This creates an information highway where local details are processed efficiently within blocks, while global information flows through the designated global tokens.
Additionally, some implementations use strided patterns where tokens can attend to blocks at regular intervals throughout the sequence, capturing periodic patterns or relationships. Others employ random sparse patterns that theoretically allow information to flow between any two positions through a small number of hops.
This hybrid approach preserves most of the modeling power of full attention while dramatically reducing computation. By carefully designing which blocks can attend to which others, these models achieve an attention complexity closer to O(n√n) or even O(n log n) rather than O(n²), enabling processing of much longer sequences with the same computational resources.
BigBird and Longformer
BigBird and Longformer implement sophisticated sparse attention patterns combining local windows, global tokens, and random connections. These architectures can efficiently scale to sequences of 4,000–8,000+ tokens with minimal loss in performance compared to full attention models.
BigBird, for example, combines three distinct attention patterns:
- Window attention: Each token attends to its local neighborhood (similar to the sliding window approach). This allows the model to capture local context effectively by focusing on nearby tokens. For instance, in a document about climate change, this helps the model understand phrases and nearby semantic connections by creating a focused attention window around each token, typically spanning 256-512 tokens in each direction.
- Global attention: Special tokens like [CLS] attend to all tokens and are attended to by all tokens, creating information highways. These global tokens serve as aggregation points that collect information from the entire sequence and distribute it back, enabling document-level understanding. For example, in a long scientific paper, the [CLS] token might gather key conclusions from various sections and make this information available to all other tokens, facilitating cross-referencing across distant parts of the document.
- Random attention: Each token attends to a small set of randomly selected tokens, which theoretically allows information to flow between any two positions in logarithmic steps. This random connectivity creates shortcuts across the document, ensuring information can propagate efficiently between distant sections. Mathematical proofs show that with just O(log n) random connections, information can flow between any two tokens in the sequence. In practice, this means even tokens separated by thousands of positions can exchange information through just a few intermediate connections.
This tri-directional attention mechanism achieves near-linear scaling while maintaining strong performance on long-document tasks like summarization and question answering. Importantly, BigBird maintains the theoretical property of "universal approximation" - it can represent any sequence-to-sequence function that full attention can, but with dramatically reduced computational requirements.
Longformer employs a similar approach but with a slightly different pattern, using a combination of sliding window attention with global attention for special tokens. It has demonstrated particular effectiveness in tasks requiring both local precision and document-level understanding, such as long-document question answering and multi-document summarization, where it can process inputs of 16,000+ tokens.
Code Example: Local Attention (Sliding Window)
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import time
class LocalAttention(nn.Module):
def __init__(self, dim, window_size=128):
super().__init__()
self.dim = dim
self.window_size = window_size
self.query_proj = nn.Linear(dim, dim)
self.key_proj = nn.Linear(dim, dim)
self.value_proj = nn.Linear(dim, dim)
self.output_proj = nn.Linear(dim, dim)
self.scaling = dim ** -0.5
def forward(self, x):
B, T, D = x.size()
# Project inputs to queries, keys, values
queries = self.query_proj(x) * self.scaling # [B, T, D]
keys = self.key_proj(x) # [B, T, D]
values = self.value_proj(x) # [B, T, D]
# Initialize output tensor
output = torch.zeros_like(x)
# Compute local attention for each position
for i in range(T):
# Define local window boundaries
start = max(0, i - self.window_size)
end = min(T, i + self.window_size + 1)
# Extract local context
local_keys = keys[:, start:end, :] # [B, window_size*2, D]
local_values = values[:, start:end, :] # [B, window_size*2, D]
# Current query
query = queries[:, i:i+1, :] # [B, 1, D]
# Compute attention scores
scores = torch.bmm(query, local_keys.transpose(1, 2)) # [B, 1, window_size*2]
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, 1, window_size*2]
# Weight values by attention
context = torch.bmm(attn_weights, local_values) # [B, 1, D]
# Store in output
output[:, i:i+1, :] = context
return self.output_proj(output)
def naive_local_attention(x, window=2):
"""A simple implementation of local attention for educational purposes"""
B, T, D = x.size()
outputs = []
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
context = x[:, start:end, :]
weights = F.softmax(torch.bmm(x[:, i:i+1, :], context.transpose(1,2)), dim=-1)
out = torch.bmm(weights, context)
outputs.append(out)
return torch.cat(outputs, dim=1)
def vectorized_local_attention(x, window=2):
"""A more efficient implementation using vectorized operations"""
B, T, D = x.size()
# Create attention mask to implement sliding window
mask = torch.zeros(T, T, device=x.device)
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
mask[i, start:end] = 1
# Compute attention scores
scores = torch.bmm(x, x.transpose(1, 2)) # [B, T, T]
# Apply mask (setting padded values to -inf before softmax)
scores = scores.masked_fill(mask.unsqueeze(0) == 0, -1e9)
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, T, T]
# Weight values by attention
output = torch.bmm(attn_weights, x) # [B, T, D]
return output
def compare_performance(seq_lengths=[10, 50, 100, 200], window=2):
"""Compare performance of different local attention implementations"""
results = {'naive': [], 'vectorized': [], 'optimized': []}
for seq_len in seq_lengths:
# Generate random input tensor
x = torch.randn(1, seq_len, 64)
# Naive implementation
start_time = time.time()
naive_local_attention(x, window)
naive_time = time.time() - start_time
results['naive'].append(naive_time)
# Vectorized implementation
start_time = time.time()
vectorized_local_attention(x, window)
vectorized_time = time.time() - start_time
results['vectorized'].append(vectorized_time)
# Optimized implementation
model = LocalAttention(64, window)
start_time = time.time()
model(x)
optimized_time = time.time() - start_time
results['optimized'].append(optimized_time)
print(f"Sequence length {seq_len}:")
print(f" Naive: {naive_time:.5f}s")
print(f" Vectorized: {vectorized_time:.5f}s")
print(f" Optimized: {optimized_time:.5f}s")
# Plot results
plt.figure(figsize=(10, 6))
plt.plot(seq_lengths, results['naive'], 'o-', label='Naive')
plt.plot(seq_lengths, results['vectorized'], 's-', label='Vectorized')
plt.plot(seq_lengths, results['optimized'], '^-', label='Optimized')
plt.xlabel('Sequence Length')
plt.ylabel('Time (s)')
plt.title('Performance Comparison of Local Attention Implementations')
plt.legend()
plt.grid(True)
plt.show()
def visualize_attention_pattern(window=2, seq_len=10):
"""Visualize the sparse attention pattern created by local attention"""
attention_mask = torch.zeros(seq_len, seq_len)
for i in range(seq_len):
start = max(0, i - window)
end = min(seq_len, i + window + 1)
attention_mask[i, start:end] = 1
plt.figure(figsize=(8, 8))
plt.imshow(attention_mask, cmap='Blues')
plt.title(f'Local Attention Pattern (Window Size = {window})')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar(label='Attention Connection')
for i in range(seq_len):
for j in range(seq_len):
color = 'white' if attention_mask[i, j] > 0 else 'none'
plt.text(j, i, '1' if attention_mask[i, j] > 0 else '0',
ha='center', va='center', color=color)
plt.tight_layout()
plt.show()
# Example
if __name__ == "__main__":
# Basic functionality test
x = torch.randn(1, 6, 16)
model = LocalAttention(16, window_size=2)
out = model(x)
print(f"Input shape: {x.shape}, Output shape: {out.shape}")
# Compare implementations
compare_performance([10, 50, 100, 200], window=2)
# Visualize the attention pattern
visualize_attention_pattern(window=2, seq_len=10)Comprehensive Breakdown: Local Attention Implementation
This code example provides a complete toolkit for understanding, implementing and analyzing local attention mechanisms. Here's a detailed breakdown:
1. Core Implementations
- LocalAttention Class: A proper PyTorch module implementation with:
- Dedicated projection layers for queries, keys, and valuesDedicated projection layers for queries, keys, and values
- Window-based sliding attention with configurable window sizeWindow-based sliding attention with configurable window size
- Proper scaling factor (1/√d) for stable gradientsProper scaling factor (1/√d) for stable gradients
- Final output projection as in standard attentionFinal output projection as in standard attention
- Naive Implementation: The original function that:
- Processes each position sequentiallyProcesses each position sequentially
- Demonstrates the core sliding window concept clearlyDemonstrates the core sliding window concept clearly
- Uses simple tensor operations for educational purposesUses simple tensor operations for educational purposes
- Vectorized Implementation: A more efficient approach that:
- Uses a mask tensor to implement the sliding window patternUses a mask tensor to implement the sliding window pattern
- Computes all attention scores at onceComputes all attention scores at once
- Avoids explicit loops over sequence positionsAvoids explicit loops over sequence positions
2. Analysis Tools
- Performance Comparison Function: Benchmarks all three implementations:
- Measures execution time across different sequence lengthsMeasures execution time across different sequence lengths
- Generates performance plots to visualize scaling behaviorGenerates performance plots to visualize scaling behavior
- Demonstrates how vectorized operations improve efficiencyDemonstrates how vectorized operations improve efficiency
- Visualization Function: Illustrates the sparse attention pattern:
- Creates a visual representation of which tokens attend to which othersCreates a visual representation of which tokens attend to which others
- Shows the diagonal band pattern characteristic of local attentionShows the diagonal band pattern characteristic of local attention
- Helps intuitive understanding of how information flows in the modelHelps intuitive understanding of how information flows in the model
3. Key Technical Insights
- Masking Technique: The code demonstrates how to create and apply attention masks to restrict which tokens can attend to which others
- Computational Efficiency: Shows how the computational complexity becomes O(n·w) instead of O(n²), where w is the window size
- Implementation Trade-offs: Illustrates the balance between code clarity (naive implementation) and computational efficiency (vectorized implementation)
This implementation provides both theoretical understanding and practical tools for working with local attention, a key technique for making transformers more efficient with long sequences. The visualization and comparison functions make it especially valuable for educational purposes.
3.3.4 Why These Matter
SwiGLU (Swish-Gated Linear Unit) significantly improves learning dynamics, giving models richer representations with little extra computational cost. This sophisticated activation function combines the benefits of gating mechanisms with a simple weighted identity connection, allowing for more effective gradient flow during training. By replacing traditional ReLU or GELU activations, SwiGLU enables models to learn more complex patterns while maintaining computational efficiency.
The mathematical formulation of SwiGLU involves multiplying a linear projection of the input with a sigmoid-weighted version of another projection, creating a smooth, differentiable pathway for gradients that helps prevent vanishing gradient problems. Models using SwiGLU typically converge faster and achieve better performance across various natural language processing tasks, making it a preferred choice in modern LLM architectures like PaLM and Gemini.
GQA (Grouped Query Attention) makes attention mechanisms substantially more efficient, reducing memory use without significant accuracy loss. This innovative technique groups queries together to share the same keys and values, dramatically reducing the memory footprint during inference. Unlike standard multi-head attention that requires separate key-value pairs for each attention head (creating a parameter explosion), GQA significantly cuts down on parameters while preserving most of the model's reasoning capabilities.
This approach creates a middle ground between multi-head attention (MHA) and multi-query attention (MQA), finding an optimal balance between parameter efficiency and model capacity. In practice, GQA can reduce the key-value cache memory requirements by 2-4x compared to standard attention while maintaining 95-99% of the model's performance, making it possible to deploy larger models on the same hardware or increase batch sizes during inference. Models like PaLM 2 and Claude have successfully implemented GQA as a core architectural improvement.
Sparse attention fundamentally transforms how LLMs can handle very long contexts without suffering from quadratic computational blow-ups. Instead of having each token attend to every other token (which scales as O(n²) with sequence length), sparse attention patterns like local, dilated, or longformer attention enable selective focus on only the most relevant tokens. This reduces computational complexity to O(n) or O(n log n), making it feasible to process documents with thousands or even tens of thousands of tokens.
Local attention, as shown in the code example above, restricts each token to attend only to a window of neighboring tokens. Dilated attention extends this by allowing tokens to attend to positions at various distances, creating a wider receptive field without increasing computation proportionally. More advanced sparse attention patterns like Reformer's LSH attention or Longformer's global+local attention combine different strategies to balance computational efficiency with model capacity. These approaches have enabled breakthroughs in long-context models that can process entire books, codebases, or lengthy conversations while maintaining coherent understanding throughout the document.
Together, these architectural refinements are why today's LLMs can be faster, leaner, and more scalable than early transformers. They represent critical engineering breakthroughs that have transformed theoretical research models into practical, deployable systems capable of handling real-world tasks with unprecedented efficiency.
3.3 Advanced Architectures: SwiGLU, GQA, Attention Sparsity
As transformers evolved from small research models to trillion-parameter giants, engineers discovered that even small changes to internal components can yield big improvements in efficiency and performance. These architectural innovations became increasingly critical as models grew in size and complexity, addressing challenges in computation, memory usage, and training stability. In this section, we'll look at three important innovations that have fundamentally changed how modern LLMs are designed:
- SwiGLU activation functions (improving feedforward networks inside transformer blocks). These specialized activation functions replace traditional ReLU or GELU activations with a more sophisticated gating mechanism that allows for smoother gradient flow during training. By incorporating both multiplicative interactions and non-linear transformations, SwiGLU enables the model to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Grouped Query Attention (GQA) (making attention faster without losing much accuracy). This clever optimization reduces the memory and computational requirements of the attention mechanism by allowing multiple query heads to share the same key and value projections. This significantly decreases both the parameter count and memory bandwidth needed during inference, addressing one of the major bottlenecks in large language model deployment.
- Attention sparsity techniques (reducing compute by ignoring unnecessary connections). These approaches recognize that not all tokens need to attend to every other token in a sequence, especially for long documents. By strategically limiting which tokens can attend to which others, sparse attention patterns can reduce the quadratic complexity of standard attention to near-linear, enabling models to process much longer contexts efficiently.
These optimizations aren't just academic curiosities — they power modern models like LLaMA, Mistral, and GPT-5. Without these architectural advances, today's state-of-the-art models would be prohibitively expensive to train and deploy. Each innovation represents a careful balance between model capability and computational efficiency, addressing specific bottlenecks that emerge at different scales of model size and context length.
3.3.1 SwiGLU (Switched Gated Linear Units)
Every transformer block contains a feedforward network (FFN) after attention. Traditionally, this FFN uses a ReLU or GELU activation. But research found that SwiGLU (a variant of GLU) yields smoother optimization and better performance at scale.
This performance improvement is largely due to the gating mechanism that allows information to flow more selectively through the network, effectively letting the model adaptively control which features are emphasized in each forward pass. Unlike traditional activation functions that apply the same transformation to all inputs, SwiGLU introduces a dynamic, input-dependent filtering mechanism that can emphasize or suppress different aspects of the representation based on the content itself.
In technical terms, SwiGLU combines the benefits of multiplicative interactions (from gates) with non-linear transformations, creating a more expressive computation unit. The gating component (using the swish/SiLU function) produces values between 0 and 1 that act as "soft switches," controlling how much information from each dimension passes through. This adaptive behavior allows the model to create more complex functional mappings with fewer parameters, resulting in improved gradient flow during training and more efficient use of model capacity.
How SwiGLU works:
- Split the hidden dimension into two parts - this creates parallel pathways through the network, allowing the model to learn different aspects of the input simultaneously. This splitting mechanism is crucial because it enables the network to process information along two distinct channels that can later be recombined in a meaningful way. Each pathway can specialize in capturing different features or patterns in the data, similar to how different neurons in the brain might respond to different aspects of visual stimuli.
- Apply a linear transformation to both parts - each pathway gets its own weight matrix (W1 and W2), enabling the network to learn different feature mappings. These linear transformations are fully learned during training and adapt to the specific task at hand. W1 typically projects the input into a space suitable for gating decisions, while W2 creates representations that will be selectively passed through based on those gates. The separation of these transformations allows the model to develop specialized feature detectors that work in tandem.
- Pass one through a sigmoid (or swish) gate - the swish activation (x * sigmoid(x)) provides a smoother gradient than ReLU and allows for some small negative values to pass through, which helps prevent "dying neurons" that can occur with ReLU. The swish function combines the benefits of sigmoid (bounded output) with the non-saturating behavior of ReLU, creating an activation function with better mathematical properties for optimization. This smoother gradient flow helps address the vanishing gradient problem that can plague deep neural networks during training.
- Multiply the two parts together - this multiplicative interaction creates a gating mechanism where the output of the swish function controls how much of the second linear transformation's output passes through. This dynamic gating allows the network to selectively amplify or suppress different features based on the input, leading to more expressive representations. The multiplication operation enables complex, content-dependent filtering of information - effectively allowing some dimensions to be emphasized while others are attenuated based on the specific input pattern, creating a form of adaptive computation.
- The mathematical formula is: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication. In practice, this can be implemented efficiently in modern deep learning frameworks by computing both transformations in parallel and then combining them with a simple Hadamard product. This formulation creates a powerful non-linear transformation that combines the benefits of gating mechanisms (like those in LSTMs and GRUs) with the parallelizability and computational efficiency of feed-forward networks.
PyTorch Example: SwiGLU vs ReLU
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class SwiGLU(nn.Module):
"""
SwiGLU activation function as used in modern LLMs like LLaMA and PaLM.
Uses the Swish/SiLU gating mechanism for better gradient properties.
"""
def __init__(self, dim_in, dim_out):
super().__init__()
self.W1 = nn.Linear(dim_in, dim_out) # Transformation for the gate
self.W2 = nn.Linear(dim_in, dim_out) # Transformation for the content
def forward(self, x):
# SiLU (Swish) activation for gating: x * sigmoid(x)
return F.silu(self.W1(x)) * self.W2(x)
class StandardFFN(nn.Module):
"""
Standard Feedforward Network with ReLU activation
as used in original Transformer architecture.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.ReLU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
class GELUBasedFFN(nn.Module):
"""
Feedforward Network with GELU activation
as used in models like BERT and early GPT versions.
"""
def __init__(self, dim_in, dim_hidden, dim_out):
super().__init__()
self.fc1 = nn.Linear(dim_in, dim_hidden)
self.act = nn.GELU()
self.fc2 = nn.Linear(dim_hidden, dim_out)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
# Hyperparameters
batch_size = 8
seq_len = 32
embed_dim = 256
hidden_dim = 1024
output_dim = 256
# Create example input
x = torch.randn(batch_size, seq_len, embed_dim)
# Initialize models
swiglu = SwiGLU(embed_dim, hidden_dim)
relu_ffn = StandardFFN(embed_dim, hidden_dim, output_dim)
gelu_ffn = GELUBasedFFN(embed_dim, hidden_dim, output_dim)
# Model outputs for comparison
with torch.no_grad():
# Timing comparisons
start = time.time()
swiglu_out = swiglu(x)
swiglu_time = time.time() - start
start = time.time()
relu_out = relu_ffn(x)
relu_time = time.time() - start
start = time.time()
gelu_out = gelu_ffn(x)
gelu_time = time.time() - start
print(f"SwiGLU output shape: {swiglu_out.shape}")
print(f"ReLU FFN output shape: {relu_out.shape}")
print(f"GELU FFN output shape: {gelu_out.shape}")
# Print timing results
print(f"\nForward pass timing:")
print(f"SwiGLU: {swiglu_time*1000:.2f}ms")
print(f"ReLU: {relu_time*1000:.2f}ms")
print(f"GELU: {gelu_time*1000:.2f}ms")
# Print sample outputs
print("\nSample outputs (first 5 values):")
print(f"SwiGLU: {swiglu_out[0, 0, :5].numpy()}")
print(f"ReLU: {relu_out[0, 0, :5].numpy()}")
print(f"GELU: {gelu_out[0, 0, :5].numpy()}")
# Visualize activation functions for comparison
def swish(x):
return x * torch.sigmoid(x)
def relu(x):
return torch.maximum(torch.zeros_like(x), x)
def gelu(x):
return 0.5 * x * (1 + torch.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))
x_range = torch.linspace(-5, 5, 1000)
plt.figure(figsize=(10, 6))
plt.plot(x_range, swish(x_range), label='Swish/SiLU (used in SwiGLU)')
plt.plot(x_range, relu(x_range), label='ReLU')
plt.plot(x_range, gelu(x_range), label='GELU')
plt.grid(True)
plt.legend()
plt.title('Comparison of Activation Functions')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.savefig('activation_comparison.png')
plt.close()
# Simple gradient flow demonstration
def compare_gradient_flow():
"""Compare gradient flow through different activation functions"""
# Create synthetic data
x = torch.randn(100, 32, requires_grad=True)
y = torch.randn(100, 32)
models = {
'SwiGLU': nn.Sequential(
nn.Linear(32, 128),
SwiGLU(128, 128),
nn.Linear(128, 32)
),
'ReLU': nn.Sequential(
nn.Linear(32, 128),
nn.ReLU(),
nn.Linear(128, 32)
),
'GELU': nn.Sequential(
nn.Linear(32, 128),
nn.GELU(),
nn.Linear(128, 32)
)
}
results = {}
for name, model in models.items():
# Forward pass
pred = model(x)
loss = torch.nn.functional.mse_loss(pred, y)
# Backward pass
loss.backward()
# Record gradient statistics
grad_norms = []
for param in model.parameters():
if param.grad is not None:
grad_norms.append(param.grad.norm().item())
results[name] = {
'mean': np.mean(grad_norms),
'std': np.std(grad_norms),
'min': np.min(grad_norms),
'max': np.max(grad_norms)
}
print("\nGradient statistics after one backward pass:")
for name, stats in results.items():
print(f"{name}: mean={stats['mean']:.6f}, std={stats['std']:.6f}, min={stats['min']:.6f}, max={stats['max']:.6f}")
compare_gradient_flow()
Breakdown of SwiGLU
How SwiGLU Works
The implementation breaks down into these key steps:
- Two Parallel Pathways: SwiGLU splits the computation into two parallel linear transformations (W1 and W2).
- Gating Mechanism: One pathway (W1) is passed through a swish/SiLU activation function (x * sigmoid(x)), which provides smoother gradients than ReLU.
- Multiplicative Interaction: The outputs are multiplied together element-wise, allowing the swish-activated path to act as a gate that controls how much of the other path's output passes through.
- Mathematical Formula: SwiGLU(x) = swish(W1·x) ⊙ W2·x, where ⊙ represents element-wise multiplication.
Key Advantages of SwiGLU
- Smoother Gradients: The swish function provides better gradient flow during backpropagation, addressing the vanishing gradient problem that can affect deep networks.
- Dynamic Feature Selection: The gating mechanism allows the network to selectively emphasize or suppress different features based on input content.
- Better Performance Per Parameter: SwiGLU enables models to capture more complex patterns with fewer parameters, leading to better performance per compute dollar.
- Improved Training Dynamics: The smoother activation function and gating mechanism result in more stable and effective training, especially in deep networks.
Code Implementation Details
- The example code demonstrates SwiGLU alongside traditional ReLU and GELU-based feedforward networks for comparison.
- It includes timing comparisons to show computational efficiency differences.
- The visualization of activation functions illustrates how swish/SiLU differs from ReLU and GELU in shape and smoothness.
- The gradient flow demonstration highlights how SwiGLU affects gradient statistics during backpropagation.
This implementation showcases why SwiGLU has become a critical component in modern LLM architectures, offering a better balance of expressivity, computational efficiency, and training stability compared to earlier alternatives.
3.3.2 Grouped Query Attention (GQA)
Standard multi-head attention has a computational cost that grows with the number of heads. To optimize this, Grouped Query Attention (GQA) was introduced. This optimization technique addresses both memory usage and computational efficiency while maintaining model quality.
Key idea: Instead of each query head having its own set of key/value heads, multiple query heads can share the same key/value projections. This sharing mechanism substantially reduces the number of parameters and computation required during inference while preserving the model's ability to learn diverse representations.
The fundamental insight behind GQA is that we can achieve a better balance between computational efficiency and model expressivity by decoupling the number of query heads from the number of key-value heads. This allows models to maintain the benefits of multi-perspective querying while reducing the overhead associated with generating and storing separate key-value pairs for each head.
In traditional multi-head attention, if you have 8 attention heads, you would have 8 separate key projections and 8 separate value projections. With GQA, you might have 8 query heads but only 2 or 4 key-value head pairs that are shared among the query heads. This means that instead of maintaining 8Q+8K+8V projection matrices (24 total), you might only need 8Q+2K+2V (12 total), cutting parameter count and computation significantly.
This reduction becomes increasingly significant in larger models. For instance, in a model with 32 attention heads and an embedding dimension of 4096, traditional multi-head attention would require approximately 402 million parameters for the attention mechanism alone, whereas GQA with 8 KV heads could reduce this to roughly 268 million parameters - a 33% reduction in memory footprint.
To understand the mechanism better, consider how attention works: each query computes similarity scores with all keys, then uses these scores to create a weighted sum of values. In GQA, multiple different queries compute attention scores against the same set of keys, and use these scores to attend to the same set of values. This maintains the expressivity of having multiple query perspectives while economizing on the key-value computations.
The efficiency gains from GQA become particularly evident during inference, where the KV cache (storing pre-computed key-value pairs for autoregressive generation) is often the primary bottleneck. By reducing the size of this cache through key-value sharing, GQA enables models to handle much longer context windows and generate text more efficiently, which is crucial for practical applications in production environments.
- Reduces memory footprint by decreasing the number of parameters needed for key and value projections, which is particularly important for larger models with billions of parameters. For example, in a model with an embedding dimension of 4096 and 32 heads, GQA with 8 KV groups can save approximately 33 million parameters per transformer layer. This reduction is achieved because traditional multi-head attention requires separate projection matrices for each attention head (Q, K, V), whereas GQA allows multiple query heads to share the same key and value projections, dramatically reducing the total parameter count across the model.
- Speeds up inference, especially in long-context models, by reducing the computational burden of generating and storing separate key-value pairs for each attention head. This is critical for server-side deployment where latency directly impacts user experience. During autoregressive generation, the KV cache (which stores previously computed key-value pairs) can become a memory bottleneck. By sharing KV projections across multiple query heads, GQA significantly reduces this cache size, allowing for faster token generation and enabling longer context handling without proportional memory increases.
- Offers nearly the same representational power as full multi-head attention but with significantly improved efficiency — a crucial trade-off for production deployment. Empirical studies show that models with GQA can achieve 95-99% of the performance of models with full multi-head attention while using considerably fewer resources. This minimal performance drop occurs because while the number of key-value pairs is reduced, the model still maintains its full capacity to generate diverse queries, preserving much of its ability to attend from different representational perspectives. The slight performance trade-off is well worth the substantial efficiency gains in most real-world applications.
- Used in LLaMA-2 to balance efficiency with performance, contributing to its ability to handle longer contexts while maintaining reasonable inference speeds. Other models like PaLM 2 and Claude have also adopted variants of this technique to scale efficiently. The implementation in LLaMA-2 specifically helped it achieve significant improvements in context window handling (up to 4K tokens) compared to its predecessor, while keeping inference costs manageable. In PaLM 2, a similar approach enabled efficient scaling to much longer contexts without the quadratic computational explosion that would occur with standard attention mechanisms.
Example: GQA Principle
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import time
class GroupedQueryAttention(nn.Module):
def __init__(self, embed_dim=512, num_query_heads=8, num_kv_heads=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_query_heads = num_query_heads
self.num_kv_heads = num_kv_heads
self.head_dim = embed_dim // num_query_heads
# Ensure dimensions are compatible
assert self.head_dim * num_query_heads == embed_dim, "embed_dim must be divisible by num_query_heads"
assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
# Query projections (many heads)
self.q_proj = nn.Linear(embed_dim, embed_dim)
# Key/Value projections (fewer heads - shared)
self.k_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
self.v_proj = nn.Linear(embed_dim, self.head_dim * num_kv_heads)
# Output projection
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Dropout for attention weights
self.dropout = nn.Dropout(dropout)
# Group size: how many query heads share one kv head
self.group_size = num_query_heads // num_kv_heads
def forward(self, x, mask=None):
"""
Args:
x: Input tensor of shape (batch_size, sequence_length, embed_dim)
mask: Optional attention mask
Returns:
output: Tensor after self-attention of shape (batch_size, sequence_length, embed_dim)
"""
batch_size, seq_len, _ = x.size()
# Project inputs to queries, keys, and values
q = self.q_proj(x).view(batch_size, seq_len, self.num_query_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Transpose for attention computation
q = q.transpose(1, 2) # (batch_size, num_query_heads, seq_len, head_dim)
k = k.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
v = v.transpose(1, 2) # (batch_size, num_kv_heads, seq_len, head_dim)
# Expand k and v to match the number of query heads through repetition
# Each group of query heads shares the same k and v
k_expanded = torch.repeat_interleave(k, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
v_expanded = torch.repeat_interleave(v, self.group_size, dim=1) # (batch_size, num_query_heads, seq_len, head_dim)
# Compute scaled dot-product attention
# (batch_size, num_query_heads, seq_len, seq_len)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
# Apply mask if provided (useful for preventing attention to padding tokens)
if mask is not None:
attn_weights = attn_weights.masked_fill(mask == 0, float("-inf"))
# Apply softmax and dropout
attn_weights = F.softmax(attn_weights, dim=-1)
attn_weights = self.dropout(attn_weights)
# Apply attention weights to values
# (batch_size, num_query_heads, seq_len, head_dim)
attn_output = torch.matmul(attn_weights, v_expanded)
# Reshape and apply output projection
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
output = self.out_proj(attn_output)
return output
def compare_attention_mechanisms(seq_len=1024, embed_dim=512):
"""Compare memory usage and speed between standard MHA and GQA"""
batch_size = 1
# Create inputs
x = torch.randn(batch_size, seq_len, embed_dim)
# Standard Multi-Head Attention (8 heads)
class StandardMHA(nn.Module):
def __init__(self):
super().__init__()
self.mha = nn.MultiheadAttention(embed_dim, num_heads=8, batch_first=True)
def forward(self, x):
return self.mha(x, x, x)[0]
standard_mha = StandardMHA()
# GQA with 8 query heads, 2 KV heads
gqa = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=2)
# GQA with 8 query heads, 4 KV heads
gqa2 = GroupedQueryAttention(embed_dim, num_query_heads=8, num_kv_heads=4)
# Measure memory and speed
results = {}
for name, model in [("Standard MHA (8 heads)", standard_mha),
("GQA (8Q, 2KV heads)", gqa),
("GQA (8Q, 4KV heads)", gqa2)]:
# Warm up
for _ in range(5):
_ = model(x)
# Measure time
torch.cuda.synchronize() if torch.cuda.is_available() else None
start_time = time.time()
for _ in range(10):
_ = model(x)
torch.cuda.synchronize() if torch.cuda.is_available() else None
end_time = time.time()
# Count parameters
param_count = sum(p.numel() for p in model.parameters())
results[name] = {
"time_per_run_ms": (end_time - start_time) * 100, # ms per 10 runs
"parameters": param_count
}
# Print results
print("Performance Comparison (sequence length = {})".format(seq_len))
print("=" * 50)
for name, metrics in results.items():
print(f"{name}:")
print(f" Time per 10 runs: {metrics['time_per_run_ms']:.2f} ms")
print(f" Parameters: {metrics['parameters']:,}")
print("-" * 50)
# Visualize KV cache size comparison
kv_cache_sizes = {
"Standard MHA": seq_len * 2 * embed_dim, # Full KV cache (8 heads)
"GQA (2 KV heads)": seq_len * 2 * (embed_dim // 4), # 1/4 the size (2 heads)
"GQA (4 KV heads)": seq_len * 2 * (embed_dim // 2), # 1/2 the size (4 heads)
}
plt.figure(figsize=(10, 5))
plt.bar(kv_cache_sizes.keys(), [size/1e6 for size in kv_cache_sizes.values()])
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Size Comparison')
for i, v in enumerate(kv_cache_sizes.values()):
plt.text(i, v/1e6 + 0.1, f"{v/1e6:.2f} MB", ha='center')
# Show how KV cache grows with sequence length
seq_lengths = [1024, 2048, 4096, 8192, 16384]
std_cache_sizes = [seq_len * 2 * embed_dim / 1e6 for seq_len in seq_lengths]
gqa_cache_sizes = [seq_len * 2 * (embed_dim // 4) / 1e6 for seq_len in seq_lengths]
plt.figure(figsize=(10, 5))
plt.plot(seq_lengths, std_cache_sizes, 'bo-', label='Standard MHA')
plt.plot(seq_lengths, gqa_cache_sizes, 'ro-', label='GQA (2 KV heads)')
plt.xlabel('Sequence Length')
plt.ylabel('KV Cache Size (MB)')
plt.title('KV Cache Growth with Sequence Length')
plt.legend()
plt.grid(True)
# Demonstrate usage with a simple example
seq_len, batch_size, embed_dim = 5, 1, 32
x = torch.randn(batch_size, seq_len, embed_dim)
gqa = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=8, num_kv_heads=2)
output = gqa(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
# Compare with standard attention
compare_attention_mechanisms(seq_len=2048, embed_dim=512)
# Basic visualization of attention patterns
def visualize_attention_pattern():
seq_len = 10
embed_dim = 64
x = torch.randn(1, seq_len, embed_dim)
model = GroupedQueryAttention(embed_dim=embed_dim, num_query_heads=4, num_kv_heads=2)
# Get attention weights by modifying forward pass temporarily
with torch.no_grad():
q = model.q_proj(x).view(1, seq_len, model.num_query_heads, model.head_dim)
k = model.k_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
v = model.v_proj(x).view(1, seq_len, model.num_kv_heads, model.head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
k_expanded = torch.repeat_interleave(k, model.group_size, dim=1)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) / torch.sqrt(torch.tensor(model.head_dim, dtype=torch.float32))
attn_weights = F.softmax(attn_weights, dim=-1)
# Plot attention patterns for each head
fig, axes = plt.subplots(1, model.num_query_heads, figsize=(15, 3))
for i in range(model.num_query_heads):
im = axes[i].imshow(attn_weights[0, i].cpu().numpy(), cmap='viridis')
axes[i].set_title(f'Head {i+1}')
axes[i].set_xlabel('Key position')
axes[i].set_ylabel('Query position')
fig.colorbar(im, ax=axes.ravel().tolist())
plt.tight_layout()
plt.suptitle('Attention Patterns with GQA (notice shared patterns within groups)')
visualize_attention_pattern()Here's a comprehensive breakdown:
Core GQA Implementation
The code example Grouped Query Attention, an optimization technique that reduces memory usage and computational cost compared to standard multi-head attention.
Class Structure
- The
GroupedQueryAttentionclass inherits fromnn.Moduleand takes parameters for embedding dimension, number of query heads, number of key-value heads, and dropout rate. - The key innovation is that multiple query heads share the same key-value heads, reducing parameter count and memory footprint.
- Two compatibility assertions ensure:
- embedding dimension is divisible by the number of query heads
- query heads are divisible by key-value heads
Projection Layers
- Query projection: Full dimension projection (
self.q_proj) - Key/Value projections: Reduced dimension projections (
self.k_proj,self.v_proj) - Output projection: Maps attention output back to original dimensions
Forward Pass
- Projects input into queries, keys and values with appropriate dimensions
- Transposes tensors for attention computation
- The critical step: expands key and value tensors to match query heads through repetition using
torch.repeat_interleave - Computes scaled dot-product attention with softmax normalization
- Applies attention weights to values and reshapes the output
Performance Comparison Functions
The code includes utilities to demonstrate GQA's advantages:
compare_attention_mechanisms(): Benchmarks standard MHA against GQA variants with different head configurations measuring:- Execution time
- Parameter count
- KV cache size - critical for inference efficiency
- Visualization functions for KV cache size comparisons and growth with sequence length
- The
visualize_attention_pattern()function demonstrates how attention patterns appear in GQA, showing how multiple query heads share the same key-value pairs
Key Benefits Demonstrated
- Memory efficiency: Reduces parameters by sharing key-value projections
- Inference speed: Smaller KV cache allows for faster token generation
- Context length: Enables handling longer sequences with minimal memory growth
- Used in modern models: The implementation resembles approaches used in LLaMA-2, PaLM 2, and Claude
This implementation provides both a practical demonstration of GQA and tools to visualize its benefits over traditional attention mechanisms, particularly in terms of memory usage and computational efficiency while maintaining most of the representational power of full multi-head attention.
3.3.3 Attention Sparsity
In full self-attention, each token attends to every other token in the sequence. This creates a computational complexity that scales quadratically as O(n²) with sequence length, which becomes prohibitively expensive for long sequences (think 100k+ tokens). For context, processing a sequence of 100,000 tokens would require 10 billion attention computations per layer!
To understand why this is problematic, consider what happens as we scale: if we double our context length from 4K to 8K tokens, the computational work quadruples from 16 million to 64 million connections per layer. This quadratic scaling quickly becomes a bottleneck for both training and inference.
Additionally, the memory requirements for storing the attention matrix also scale quadratically. For a sequence of length n, we need to store an n×n attention matrix, which for long sequences can exceed available GPU memory. For example, a 32K token sequence would require approximately 4GB of memory just to store a single attention matrix in 32-bit precision.
Sparse attention techniques reduce this computational burden by attending only to the most relevant positions, effectively pruning unnecessary connections. This transforms the scaling from quadratic to nearly linear in many implementations. By strategically limiting which tokens can attend to which other tokens, these techniques dramatically reduce both computation and memory requirements.
The key insight behind sparse attention is that not all token-to-token interactions are equally important. Many language phenomena are local in nature, while certain special tokens may need global context. By exploiting this pattern, sparse attention can preserve most of the model's capabilities while eliminating many unnecessary computations.
Local attention
Each token attends only to its neighbors within a fixed window size (e.g., ±128 tokens). This creates a sliding window of attention that moves with each token position. For example, with a window size of 128, token at position 500 would attend to tokens from positions 372 to 628.
This approach works particularly well for tasks where nearby context is most relevant, such as speech recognition where phonemes relate strongly to adjacent sounds, or DNA analysis where nearby nucleotides often form functional units together. Local attention is also effective for text processing tasks where most semantic relationships occur between words that are relatively close to each other in the sequence.
The efficiency gains are substantial - the computational complexity becomes O(n×w), where w is the fixed window size. Since w is a constant (like 128 or 256), this effectively makes the attention mechanism scale linearly with sequence length rather than quadratically. For a sequence of 100,000 tokens with a window size of 256, this reduces computations from 10 billion to just 25.6 million - a 390x improvement.
However, local attention does have limitations - it struggles with tasks requiring long-range dependencies, such as document-level reasoning where important information may be separated by thousands of tokens. This is why more sophisticated sparse attention patterns often combine local attention with other mechanisms to capture both local and global relationships.
Block-sparse attention
Tokens attend within defined chunks or blocks, with occasional global tokens that can see across the entire sequence. This creates a sparse attention pattern where most tokens have limited vision but a few sentinel tokens maintain global context. These blocks can be arranged in various patterns - diagonal blocks for local attention, or more complex structures that allow for hierarchical information flow.
For example, in a block-sparse approach, a document might be divided into chunks of 512 tokens, with each chunk having internal full attention, plus dedicated "summary tokens" that can see across all chunks. This creates an information highway where local details are processed efficiently within blocks, while global information flows through the designated global tokens.
Additionally, some implementations use strided patterns where tokens can attend to blocks at regular intervals throughout the sequence, capturing periodic patterns or relationships. Others employ random sparse patterns that theoretically allow information to flow between any two positions through a small number of hops.
This hybrid approach preserves most of the modeling power of full attention while dramatically reducing computation. By carefully designing which blocks can attend to which others, these models achieve an attention complexity closer to O(n√n) or even O(n log n) rather than O(n²), enabling processing of much longer sequences with the same computational resources.
BigBird and Longformer
BigBird and Longformer implement sophisticated sparse attention patterns combining local windows, global tokens, and random connections. These architectures can efficiently scale to sequences of 4,000–8,000+ tokens with minimal loss in performance compared to full attention models.
BigBird, for example, combines three distinct attention patterns:
- Window attention: Each token attends to its local neighborhood (similar to the sliding window approach). This allows the model to capture local context effectively by focusing on nearby tokens. For instance, in a document about climate change, this helps the model understand phrases and nearby semantic connections by creating a focused attention window around each token, typically spanning 256-512 tokens in each direction.
- Global attention: Special tokens like [CLS] attend to all tokens and are attended to by all tokens, creating information highways. These global tokens serve as aggregation points that collect information from the entire sequence and distribute it back, enabling document-level understanding. For example, in a long scientific paper, the [CLS] token might gather key conclusions from various sections and make this information available to all other tokens, facilitating cross-referencing across distant parts of the document.
- Random attention: Each token attends to a small set of randomly selected tokens, which theoretically allows information to flow between any two positions in logarithmic steps. This random connectivity creates shortcuts across the document, ensuring information can propagate efficiently between distant sections. Mathematical proofs show that with just O(log n) random connections, information can flow between any two tokens in the sequence. In practice, this means even tokens separated by thousands of positions can exchange information through just a few intermediate connections.
This tri-directional attention mechanism achieves near-linear scaling while maintaining strong performance on long-document tasks like summarization and question answering. Importantly, BigBird maintains the theoretical property of "universal approximation" - it can represent any sequence-to-sequence function that full attention can, but with dramatically reduced computational requirements.
Longformer employs a similar approach but with a slightly different pattern, using a combination of sliding window attention with global attention for special tokens. It has demonstrated particular effectiveness in tasks requiring both local precision and document-level understanding, such as long-document question answering and multi-document summarization, where it can process inputs of 16,000+ tokens.
Code Example: Local Attention (Sliding Window)
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import time
class LocalAttention(nn.Module):
def __init__(self, dim, window_size=128):
super().__init__()
self.dim = dim
self.window_size = window_size
self.query_proj = nn.Linear(dim, dim)
self.key_proj = nn.Linear(dim, dim)
self.value_proj = nn.Linear(dim, dim)
self.output_proj = nn.Linear(dim, dim)
self.scaling = dim ** -0.5
def forward(self, x):
B, T, D = x.size()
# Project inputs to queries, keys, values
queries = self.query_proj(x) * self.scaling # [B, T, D]
keys = self.key_proj(x) # [B, T, D]
values = self.value_proj(x) # [B, T, D]
# Initialize output tensor
output = torch.zeros_like(x)
# Compute local attention for each position
for i in range(T):
# Define local window boundaries
start = max(0, i - self.window_size)
end = min(T, i + self.window_size + 1)
# Extract local context
local_keys = keys[:, start:end, :] # [B, window_size*2, D]
local_values = values[:, start:end, :] # [B, window_size*2, D]
# Current query
query = queries[:, i:i+1, :] # [B, 1, D]
# Compute attention scores
scores = torch.bmm(query, local_keys.transpose(1, 2)) # [B, 1, window_size*2]
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, 1, window_size*2]
# Weight values by attention
context = torch.bmm(attn_weights, local_values) # [B, 1, D]
# Store in output
output[:, i:i+1, :] = context
return self.output_proj(output)
def naive_local_attention(x, window=2):
"""A simple implementation of local attention for educational purposes"""
B, T, D = x.size()
outputs = []
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
context = x[:, start:end, :]
weights = F.softmax(torch.bmm(x[:, i:i+1, :], context.transpose(1,2)), dim=-1)
out = torch.bmm(weights, context)
outputs.append(out)
return torch.cat(outputs, dim=1)
def vectorized_local_attention(x, window=2):
"""A more efficient implementation using vectorized operations"""
B, T, D = x.size()
# Create attention mask to implement sliding window
mask = torch.zeros(T, T, device=x.device)
for i in range(T):
start = max(0, i - window)
end = min(T, i + window + 1)
mask[i, start:end] = 1
# Compute attention scores
scores = torch.bmm(x, x.transpose(1, 2)) # [B, T, T]
# Apply mask (setting padded values to -inf before softmax)
scores = scores.masked_fill(mask.unsqueeze(0) == 0, -1e9)
# Apply softmax to get attention weights
attn_weights = F.softmax(scores, dim=-1) # [B, T, T]
# Weight values by attention
output = torch.bmm(attn_weights, x) # [B, T, D]
return output
def compare_performance(seq_lengths=[10, 50, 100, 200], window=2):
"""Compare performance of different local attention implementations"""
results = {'naive': [], 'vectorized': [], 'optimized': []}
for seq_len in seq_lengths:
# Generate random input tensor
x = torch.randn(1, seq_len, 64)
# Naive implementation
start_time = time.time()
naive_local_attention(x, window)
naive_time = time.time() - start_time
results['naive'].append(naive_time)
# Vectorized implementation
start_time = time.time()
vectorized_local_attention(x, window)
vectorized_time = time.time() - start_time
results['vectorized'].append(vectorized_time)
# Optimized implementation
model = LocalAttention(64, window)
start_time = time.time()
model(x)
optimized_time = time.time() - start_time
results['optimized'].append(optimized_time)
print(f"Sequence length {seq_len}:")
print(f" Naive: {naive_time:.5f}s")
print(f" Vectorized: {vectorized_time:.5f}s")
print(f" Optimized: {optimized_time:.5f}s")
# Plot results
plt.figure(figsize=(10, 6))
plt.plot(seq_lengths, results['naive'], 'o-', label='Naive')
plt.plot(seq_lengths, results['vectorized'], 's-', label='Vectorized')
plt.plot(seq_lengths, results['optimized'], '^-', label='Optimized')
plt.xlabel('Sequence Length')
plt.ylabel('Time (s)')
plt.title('Performance Comparison of Local Attention Implementations')
plt.legend()
plt.grid(True)
plt.show()
def visualize_attention_pattern(window=2, seq_len=10):
"""Visualize the sparse attention pattern created by local attention"""
attention_mask = torch.zeros(seq_len, seq_len)
for i in range(seq_len):
start = max(0, i - window)
end = min(seq_len, i + window + 1)
attention_mask[i, start:end] = 1
plt.figure(figsize=(8, 8))
plt.imshow(attention_mask, cmap='Blues')
plt.title(f'Local Attention Pattern (Window Size = {window})')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar(label='Attention Connection')
for i in range(seq_len):
for j in range(seq_len):
color = 'white' if attention_mask[i, j] > 0 else 'none'
plt.text(j, i, '1' if attention_mask[i, j] > 0 else '0',
ha='center', va='center', color=color)
plt.tight_layout()
plt.show()
# Example
if __name__ == "__main__":
# Basic functionality test
x = torch.randn(1, 6, 16)
model = LocalAttention(16, window_size=2)
out = model(x)
print(f"Input shape: {x.shape}, Output shape: {out.shape}")
# Compare implementations
compare_performance([10, 50, 100, 200], window=2)
# Visualize the attention pattern
visualize_attention_pattern(window=2, seq_len=10)Comprehensive Breakdown: Local Attention Implementation
This code example provides a complete toolkit for understanding, implementing and analyzing local attention mechanisms. Here's a detailed breakdown:
1. Core Implementations
- LocalAttention Class: A proper PyTorch module implementation with:
- Dedicated projection layers for queries, keys, and valuesDedicated projection layers for queries, keys, and values
- Window-based sliding attention with configurable window sizeWindow-based sliding attention with configurable window size
- Proper scaling factor (1/√d) for stable gradientsProper scaling factor (1/√d) for stable gradients
- Final output projection as in standard attentionFinal output projection as in standard attention
- Naive Implementation: The original function that:
- Processes each position sequentiallyProcesses each position sequentially
- Demonstrates the core sliding window concept clearlyDemonstrates the core sliding window concept clearly
- Uses simple tensor operations for educational purposesUses simple tensor operations for educational purposes
- Vectorized Implementation: A more efficient approach that:
- Uses a mask tensor to implement the sliding window patternUses a mask tensor to implement the sliding window pattern
- Computes all attention scores at onceComputes all attention scores at once
- Avoids explicit loops over sequence positionsAvoids explicit loops over sequence positions
2. Analysis Tools
- Performance Comparison Function: Benchmarks all three implementations:
- Measures execution time across different sequence lengthsMeasures execution time across different sequence lengths
- Generates performance plots to visualize scaling behaviorGenerates performance plots to visualize scaling behavior
- Demonstrates how vectorized operations improve efficiencyDemonstrates how vectorized operations improve efficiency
- Visualization Function: Illustrates the sparse attention pattern:
- Creates a visual representation of which tokens attend to which othersCreates a visual representation of which tokens attend to which others
- Shows the diagonal band pattern characteristic of local attentionShows the diagonal band pattern characteristic of local attention
- Helps intuitive understanding of how information flows in the modelHelps intuitive understanding of how information flows in the model
3. Key Technical Insights
- Masking Technique: The code demonstrates how to create and apply attention masks to restrict which tokens can attend to which others
- Computational Efficiency: Shows how the computational complexity becomes O(n·w) instead of O(n²), where w is the window size
- Implementation Trade-offs: Illustrates the balance between code clarity (naive implementation) and computational efficiency (vectorized implementation)
This implementation provides both theoretical understanding and practical tools for working with local attention, a key technique for making transformers more efficient with long sequences. The visualization and comparison functions make it especially valuable for educational purposes.
3.3.4 Why These Matter
SwiGLU (Swish-Gated Linear Unit) significantly improves learning dynamics, giving models richer representations with little extra computational cost. This sophisticated activation function combines the benefits of gating mechanisms with a simple weighted identity connection, allowing for more effective gradient flow during training. By replacing traditional ReLU or GELU activations, SwiGLU enables models to learn more complex patterns while maintaining computational efficiency.
The mathematical formulation of SwiGLU involves multiplying a linear projection of the input with a sigmoid-weighted version of another projection, creating a smooth, differentiable pathway for gradients that helps prevent vanishing gradient problems. Models using SwiGLU typically converge faster and achieve better performance across various natural language processing tasks, making it a preferred choice in modern LLM architectures like PaLM and Gemini.
GQA (Grouped Query Attention) makes attention mechanisms substantially more efficient, reducing memory use without significant accuracy loss. This innovative technique groups queries together to share the same keys and values, dramatically reducing the memory footprint during inference. Unlike standard multi-head attention that requires separate key-value pairs for each attention head (creating a parameter explosion), GQA significantly cuts down on parameters while preserving most of the model's reasoning capabilities.
This approach creates a middle ground between multi-head attention (MHA) and multi-query attention (MQA), finding an optimal balance between parameter efficiency and model capacity. In practice, GQA can reduce the key-value cache memory requirements by 2-4x compared to standard attention while maintaining 95-99% of the model's performance, making it possible to deploy larger models on the same hardware or increase batch sizes during inference. Models like PaLM 2 and Claude have successfully implemented GQA as a core architectural improvement.
Sparse attention fundamentally transforms how LLMs can handle very long contexts without suffering from quadratic computational blow-ups. Instead of having each token attend to every other token (which scales as O(n²) with sequence length), sparse attention patterns like local, dilated, or longformer attention enable selective focus on only the most relevant tokens. This reduces computational complexity to O(n) or O(n log n), making it feasible to process documents with thousands or even tens of thousands of tokens.
Local attention, as shown in the code example above, restricts each token to attend only to a window of neighboring tokens. Dilated attention extends this by allowing tokens to attend to positions at various distances, creating a wider receptive field without increasing computation proportionally. More advanced sparse attention patterns like Reformer's LSH attention or Longformer's global+local attention combine different strategies to balance computational efficiency with model capacity. These approaches have enabled breakthroughs in long-context models that can process entire books, codebases, or lengthy conversations while maintaining coherent understanding throughout the document.
Together, these architectural refinements are why today's LLMs can be faster, leaner, and more scalable than early transformers. They represent critical engineering breakthroughs that have transformed theoretical research models into practical, deployable systems capable of handling real-world tasks with unprecedented efficiency.