title: "Attention Mechanisms in AI Agents: Focus, Context, and Selective Processing" description: "Discover how attention mechanisms enable AI agents to focus on relevant information, process context effectively, and make intelligent decisions by selectively attending to the most important elements in complex environments."
Attention Mechanisms in AI Agents: Focus, Context, and Selective Processing
Welcome to part 27 of our AI Agent Engineering series. In this comprehensive examination, we'll explore the transformative power of attention mechanisms that enable AI agents to selectively focus on relevant information, process contextual cues effectively, and make intelligent decisions in complex environments.
Introduction
In a world inundated with information, the ability to focus attention on what matters most separates intelligent beings from simple reactive systems. Just as humans can tune out background noise in a crowded room while focusing on a single conversation, intelligent AI agents must develop sophisticated attention mechanisms to navigate complex environments filled with competing stimuli.
Attention mechanisms have revolutionized artificial intelligence, enabling breakthrough performance in natural language processing, computer vision, and reinforcement learning applications. For AI agents operating in dynamic, multifaceted environments, attention serves as the cognitive lens that filters relevant information from noise, coordinates processing resources, and guides decision-making processes.
Consider a robotic assistant in a busy hospital ward. The robot must simultaneously monitor patient vital signs, track staff movements, process verbal instructions, avoid obstacles, and maintain awareness of emergency protocols. Without attention mechanisms, the robot would struggle to prioritize critical signals from overwhelming sensor data. With properly implemented attention, however, it can dynamically allocate processing resources to the most urgent tasks while maintaining peripheral awareness of background activities.
Core Concepts of Attention in AI
Defining Attention Mechanisms
Attention in artificial systems draws inspiration from cognitive science while incorporating computational efficiency considerations. At its core, attention is a mechanism that allows models to focus on specific parts of input data rather than processing everything uniformly.
The fundamental principle is selective processing: instead of treating all input elements equally, attention mechanisms assign varying degrees of importance (weights) to different components, amplifying relevant information while suppressing irrelevant details.
Mathematically, attention can be expressed as:
Attention(Q, K, V) = softmax(QK^T / √d_k) × V
Where:
- Q (Query): Represents what we're looking for
- K (Key): Represents what we can find
- V (Value): Represents what we actually get
- d_k: Dimension of key vectors (for scaling)
Historical Evolution
The development of attention mechanisms has followed a clear progression:
Early Approaches: Simple hard attention that selects one element at a time (e.g., visual attention in early computer vision systems).
Soft Attention: Weighted combinations allowing gradients to flow through all elements, enabling end-to-end training.
Multi-Head Attention: Parallel attention mechanisms capturing different aspects of relationships.
Transformers: Self-attention architectures that dispensed with recurrence while achieving superior performance.
Each evolution brought increased modeling capacity and flexibility in handling diverse types of relationships in data.
Types of Attention Mechanisms
Self-Attention
Self-attention mechanisms allow elements within a sequence to attend to each other, capturing dependencies regardless of distance:
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads=8):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
self.qkv_projection = nn.Linear(embed_dim, embed_dim * 3)
self.output_projection = nn.Linear(embed_dim, embed_dim)
def forward(self, x, mask=None):
batch_size, seq_length, embed_dim = x.shape
# Project to query, key, value
qkv = self.qkv_projection(x) # [batch, seq, embed_dim * 3]
qkv = qkv.reshape(batch_size, seq_length, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4) # [3, batch, heads, seq, head_dim]
queries, keys, values = qkv[0], qkv[1], qkv[2]
# Compute attention scores
attention_scores = torch.matmul(queries, keys.transpose(-2, -1))
attention_scores = attention_scores / (self.head_dim ** 0.5)
# Apply mask if provided
if mask is not None:
attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))
# Apply softmax to get attention weights
attention_weights = F.softmax(attention_scores, dim=-1)
# Apply attention to values
attended_values = torch.matmul(attention_weights, values)
# Reshape and project output
attended_values = attended_values.transpose(1, 2).reshape(
batch_size, seq_length, embed_dim)
output = self.output_projection(attended_values)
return output, attention_weights
# Example usage
attention_layer = SelfAttention(embed_dim=512, num_heads=8)
sequence = torch.randn(32, 100, 512) # Batch of 32 sequences, each 100 tokens, 512-dim embeddings
attended_sequence, weights = attention_layer(sequence)
Cross-Attention
Cross-attention mechanisms enable interaction between different sequences or modalities:
class CrossAttention(nn.Module):
def __init__(self, query_dim, context_dim, num_heads=8):
super().__init__()
self.query_dim = query_dim
self.context_dim = context_dim
self.num_heads = num_heads
self.head_dim = query_dim // num_heads
self.query_projection = nn.Linear(query_dim, query_dim)
self.key_projection = nn.Linear(context_dim, query_dim)
self.value_projection = nn.Linear(context_dim, query_dim)
self.output_projection = nn.Linear(query_dim, query_dim)
def forward(self, queries, context, mask=None):
batch_size, query_seq_len, _ = queries.shape
context_seq_len = context.shape[1]
# Project queries, keys, and values
proj_queries = self.query_projection(queries).view(
batch_size, query_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
proj_keys = self.key_projection(context).view(
batch_size, context_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
proj_values = self.value_projection(context).view(
batch_size, context_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# Compute attention
attention_scores = torch.matmul(proj_queries, proj_keys.transpose(-2, -1))
attention_scores = attention_scores / (self.head_dim ** 0.5)
if mask is not None:
attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))
attention_weights = F.softmax(attention_scores, dim=-1)
attended_values = torch.matmul(attention_weights, proj_values)
# Reshape and project
output = attended_values.transpose(1, 2).contiguous().view(
batch_size, query_seq_len, self.query_dim)
output = self.output_projection(output)
return output, attention_weights
# Example: Image captioning where queries are text tokens and context is image features
caption_tokens = torch.randn(16, 20, 512) # 16 captions, 20 tokens each
image_features = torch.randn(16, 49, 512) # 16 images, 49 patches each
cross_attention = CrossAttention(query_dim=512, context_dim=512)
attended_captions, attention_map = cross_attention(caption_tokens, image_features)
Causal Attention
Causal attention ensures that predictions only depend on past and present information:
class CausalAttention(nn.Module):
def __init__(self, embed_dim, num_heads=8):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.qkv_projection = nn.Linear(embed_dim, embed_dim * 3)
self.output_projection = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
batch_size, seq_length, embed_dim = x.shape
# Project to QKV
qkv = self.qkv_projection(x)
qkv = qkv.reshape(batch_size, seq_length, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
queries, keys, values = qkv[0], qkv[1], qkv[2]
# Causal masking
causal_mask = torch.tril(torch.ones(seq_length, seq_length)).bool()
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0) # Add batch and head dimensions
# Compute attention with causal mask
attention_scores = torch.matmul(queries, keys.transpose(-2, -1))
attention_scores = attention_scores / (self.head_dim ** 0.5)
attention_scores = attention_scores.masked_fill(~causal_mask, float('-inf'))
attention_weights = F.softmax(attention_scores, dim=-1)
attended_values = torch.matmul(attention_weights, values)
# Reshape output
output = attended_values.transpose(1, 2).reshape(batch_size, seq_length, embed_dim)
output = self.output_projection(output)
return output, attention_weights
# Example usage for autoregressive sequence modeling
autoregressive_model = CausalAttention(embed_dim=512)
input_sequence = torch.randn(8, 50, 512) # 8 sequences of length 50
output_sequence, attention_weights = autoregressive_model(input_sequence)
Multi-Head Attention Architecture
Multi-head attention allows models to focus on different aspects of the input simultaneously:
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim, num_heads, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.dropout = nn.Dropout(dropout)
assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
self.scaling = self.head_dim ** -0.5
def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
batch_size, tgt_len, embed_dim = query.shape
src_len = key.shape[1]
# Project to QKV
q = self.q_proj(query) * self.scaling
k = self.k_proj(key)
v = self.v_proj(value)
# Reshape for multi-head attention
q = q.view(batch_size, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(batch_size, src_len, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(batch_size, src_len, self.num_heads, self.head_dim).transpose(1, 2)
# Compute attention
attn_output_weights = torch.matmul(q, k.transpose(-2, -1))
# Apply masks
if attn_mask is not None:
attn_output_weights += attn_mask
if key_padding_mask is not None:
attn_output_weights = attn_output_weights.masked_fill(
key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf')
)
# Apply softmax and dropout
attn_output_weights = F.softmax(attn_output_weights, dim=-1)
attn_output_weights = self.dropout(attn_output_weights)
# Apply attention to values
attn_output = torch.matmul(attn_output_weights, v)
# Reshape and project output
attn_output = attn_output.transpose(1, 2).contiguous().view(
batch_size, tgt_len, embed_dim)
attn_output = self.out_proj(attn_output)
return attn_output, attn_output_weights
# Example with different query and key dimensions
encoder_hidden = torch.randn(16, 30, 512) # Encoder outputs
decoder_hidden = torch.randn(16, 20, 512) # Decoder inputs
attention_layer = MultiHeadAttention(embed_dim=512, num_heads=8)
attended_decoder, attention_weights = attention_layer(
decoder_hidden, encoder_hidden, encoder_hidden)
Attention in Different Modalities
Vision Transformers with Attention
Vision transformers apply attention mechanisms to image processing:
class VisionTransformer(nn.Module):
def __init__(self, img_size=224, patch_size=16, num_classes=1000,
embed_dim=768, depth=12, num_heads=12):
super().__init__()
self.patch_embed = PatchEmbedding(img_size, patch_size, embed_dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.randn(1, self.patch_embed.num_patches + 1, embed_dim))
self.dropout = nn.Dropout(0.1)
self.blocks = nn.ModuleList([
TransformerBlock(embed_dim, num_heads) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
def forward(self, x):
# Patch embedding
x = self.patch_embed(x) # [batch, num_patches, embed_dim]
# Add class token
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
# Add positional embedding
x = x + self.pos_embed
x = self.dropout(x)
# Apply transformer blocks
for block in self.blocks:
x = block(x)
# Classification
x = self.norm(x)
cls_token_final = x[:, 0] # Use class token for classification
return self.head(cls_token_final)
class PatchEmbedding(nn.Module):
def __init__(self, img_size, patch_size, embed_dim):
super().__init__()
self.img_size = img_size
self.patch_size = patch_size
self.num_patches = (img_size // patch_size) ** 2
self.proj = nn.Conv2d(3, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
x = self.proj(x) # [batch, embed_dim, h, w]
x = x.flatten(2) # [batch, embed_dim, num_patches]
x = x.transpose(1, 2) # [batch, num_patches, embed_dim]
return x
class TransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads, mlp_ratio=4, dropout=0.1):
super().__init__()
self.attn = MultiHeadAttention(embed_dim, num_heads, dropout)
self.norm1 = nn.LayerNorm(embed_dim)
self.mlp = MLP(embed_dim, int(embed_dim * mlp_ratio), dropout)
self.norm2 = nn.LayerNorm(embed_dim)
def forward(self, x):
# Self-attention with residual connection
attn_out, _ = self.attn(x, x, x)
x = x + attn_out
x = self.norm1(x)
# MLP with residual connection
mlp_out = self.mlp(x)
x = x + mlp_out
x = self.norm2(x)
return x
class MLP(nn.Module):
def __init__(self, in_features, hidden_features, dropout=0.1):
super().__init__()
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = nn.GELU()
self.fc2 = nn.Linear(hidden_features, in_features)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.dropout(x)
x = self.fc2(x)
x = self.dropout(x)
return x
# Example usage
vit = VisionTransformer(img_size=224, patch_size=16, num_classes=1000)
sample_image = torch.randn(8, 3, 224, 224) # Batch of 8 RGB images
output = vit(sample_image)
Attention in Natural Language Processing
Transformer-based language models extensively use attention mechanisms:
class BERTLikeModel(nn.Module):
def __init__(self, vocab_size, max_length=512, embed_dim=768,
num_layers=12, num_heads=12):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, embed_dim)
self.position_embedding = nn.Embedding(max_length, embed_dim)
self.segment_embedding = nn.Embedding(2, embed_dim)
self.dropout = nn.Dropout(0.1)
self.layers = nn.ModuleList([
TransformerEncoderLayer(embed_dim, num_heads) for _ in range(num_layers)
])
self.norm = nn.LayerNorm(embed_dim)
def forward(self, input_ids, token_type_ids=None, attention_mask=None):
batch_size, seq_length = input_ids.shape
# Token embeddings
token_embeds = self.token_embedding(input_ids)
# Position embeddings
position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
position_embeds = self.position_embedding(position_ids)
# Segment embeddings
if token_type_ids is None:
token_type_ids = torch.zeros_like(input_ids)
segment_embeds = self.segment_embedding(token_type_ids)
# Combine embeddings
embeddings = token_embeds + position_embeds + segment_embeds
embeddings = self.dropout(embeddings)
# Apply transformer layers
for layer in self.layers:
embeddings = layer(embeddings, attention_mask)
embeddings = self.norm(embeddings)
return embeddings
class TransformerEncoderLayer(nn.Module):
def __init__(self, embed_dim, num_heads, dropout=0.1):
super().__init__()
self.self_attn = MultiHeadAttention(embed_dim, num_heads, dropout)
self.norm1 = nn.LayerNorm(embed_dim)
self.mlp = MLP(embed_dim, embed_dim * 4, dropout)
self.norm2 = nn.LayerNorm(embed_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, src, src_mask=None):
# Self-attention
attn_out, _ = self.self_attn(src, src, src, attn_mask=src_mask)
src = src + self.dropout(attn_out)
src = self.norm1(src)
# Feed-forward
mlp_out = self.mlp(src)
src = src + self.dropout(mlp_out)
src = self.norm2(src)
return src
# Example usage for text classification
bert_model = BERTLikeModel(vocab_size=30522, max_length=512)
sample_input = torch.randint(0, 30522, (16, 128)) # 16 sentences, 128 tokens each
attention_mask = torch.ones(16, 128) # No padding in this example
output_embeddings = bert_model(sample_input, attention_mask=attention_mask)
Advanced Attention Mechanisms
Sparse Attention
Sparse attention mechanisms reduce computational complexity for long sequences:
class SparseAttention(nn.Module):
def __init__(self, embed_dim, num_heads, block_size=64):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.block_size = block_size
self.head_dim = embed_dim // num_heads
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
batch_size, seq_length, embed_dim = x.shape
# Project to QKV
q = self.q_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
# Divide sequence into blocks
num_blocks = (seq_length + self.block_size - 1) // self.block_size
outputs = []
attention_maps = []
for i in range(num_blocks):
start_idx = i * self.block_size
end_idx = min((i + 1) * self.block_size, seq_length)
# Local attention within block and adjacent blocks
local_start = max(0, start_idx - self.block_size)
local_end = min(seq_length, end_idx + self.block_size)
local_q = q[:, start_idx:end_idx]
local_k = k[:, local_start:local_end]
local_v = v[:, local_start:local_end]
# Compute sparse attention
attn_scores = torch.einsum('bihd,bjhd->bijh', local_q, local_k)
attn_scores = attn_scores / (self.head_dim ** 0.5)
attn_weights = F.softmax(attn_scores, dim=2)
attended_values = torch.einsum('bijh,bjhd->bihd', attn_weights, local_v)
outputs.append(attended_values.contiguous().view(
batch_size, end_idx - start_idx, embed_dim))
attention_maps.append(attn_weights)
# Combine outputs
output = torch.cat(outputs, dim=1)
output = self.out_proj(output)
return output, attention_maps
# Example usage for long sequences
sparse_attention = SparseAttention(embed_dim=512, num_heads=8, block_size=64)
long_sequence = torch.randn(4, 1024, 512) # Much longer sequence
sparse_output, _ = sparse_attention(long_sequence)
Learned Sparse Attention
Attention mechanisms that learn which positions to attend to:
class LearnedSparseAttention(nn.Module):
def __init__(self, embed_dim, num_heads, top_k=32):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.top_k = top_k
self.head_dim = embed_dim // num_heads
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Learnable routing network
self.routing_network = nn.Sequential(
nn.Linear(embed_dim, embed_dim // 2),
nn.ReLU(),
nn.Linear(embed_dim // 2, num_heads * top_k)
)
def forward(self, x):
batch_size, seq_length, embed_dim = x.shape
# Project to QKV
q = self.q_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
# Learn sparse attention patterns
routing_scores = self.routing_network(x) # [batch, seq, heads * top_k]
routing_scores = routing_scores.view(batch_size, seq_length, self.num_heads, self.top_k)
# For each head and position, select top-k keys to attend to
outputs = []
for head in range(self.num_heads):
head_q = q[:, :, head, :] # [batch, seq, head_dim]
head_k = k[:, :, head, :] # [batch, seq, head_dim]
head_v = v[:, :, head, :] # [batch, seq, head_dim]
head_routing = routing_scores[:, :, head, :] # [batch, seq, top_k]
# Select top-k keys for each query
_, top_k_indices = torch.topk(head_routing, self.top_k, dim=2)
# Gather selected keys and values
batch_indices = torch.arange(batch_size).unsqueeze(1).unsqueeze(2)
seq_indices = torch.arange(seq_length).unsqueeze(0).unsqueeze(2)
selected_k = head_k[batch_indices, top_k_indices, :] # [batch, seq, top_k, head_dim]
selected_v = head_v[batch_indices, top_k_indices, :] # [batch, seq, top_k, head_dim]
# Compute attention with selected keys
attn_scores = torch.einsum('bsd,bskd->bsk', head_q, selected_k)
attn_scores = attn_scores / (self.head_dim ** 0.5)
attn_weights = F.softmax(attn_scores, dim=2)
# Apply attention to selected values
attended_values = torch.einsum('bsk,bskd->bsd', attn_weights, selected_v)
outputs.append(attended_values.unsqueeze(2))
# Combine heads
output = torch.cat(outputs, dim=2).view(batch_size, seq_length, embed_dim)
output = self.out_proj(output)
return output, routing_scores
# Example usage
learned_sparse = LearnedSparseAttention(embed_dim=512, num_heads=8, top_k=32)
medium_sequence = torch.randn(8, 256, 512) # Medium-length sequence
sparse_output, learned_patterns = learned_sparse(medium_sequence)
Attention in Reinforcement Learning Agents
Reinforcement learning agents benefit significantly from attention mechanisms for processing complex environmental states:
class AttentionBasedRLAgent(nn.Module):
def __init__(self, observation_dim, action_dim, hidden_dim=512, num_heads=8):
super().__init__()
self.observation_encoder = nn.Linear(observation_dim, hidden_dim)
self.attention_processor = MultiHeadAttention(
embed_dim=hidden_dim, num_heads=num_heads)
self.policy_head = nn.Linear(hidden_dim, action_dim)
self.value_head = nn.Linear(hidden_dim, 1)
# For processing temporal sequences of observations
self.temporal_attention = MultiHeadAttention(
embed_dim=hidden_dim, num_heads=num_heads)
def forward(self, observations, action_masks=None):
"""
Process observations using attention mechanisms
Args:
observations: Tensor of shape [batch_size, seq_len, observation_dim]
or [batch_size, observation_dim] for single observations
action_masks: Optional binary masks for valid actions
"""
if observations.dim() == 2:
# Single observation
encoded_obs = self.observation_encoder(observations)
attended_obs, attention_weights = self.attention_processor(
encoded_obs.unsqueeze(1), encoded_obs.unsqueeze(1), encoded_obs.unsqueeze(1))
attended_obs = attended_obs.squeeze(1)
else:
# Sequence of observations
batch_size, seq_len, obs_dim = observations.shape
encoded_obs = self.observation_encoder(observations)
# Apply self-attention across sequence
attended_obs, attention_weights = self.temporal_attention(
encoded_obs, encoded_obs, encoded_obs)
# Use last timestep for decision making
attended_obs = attended_obs[:, -1, :]
# Policy and value computation
policy_logits = self.policy_head(attended_obs)
value = self.value_head(attended_obs).squeeze(-1)
# Apply action mask if provided
if action_masks is not None:
policy_logits = policy_logits.masked_fill(~action_masks, float('-inf'))
return F.softmax(policy_logits, dim=-1), value, attention_weights
def compute_attention_importance(self, attention_weights):
"""Analyze which parts of observations were most important"""
# For self-attention, compute average attention weights
if attention_weights.dim() == 4: # [batch, heads, seq, seq]
importance_scores = attention_weights.mean(dim=1).mean(dim=1) # [batch, seq]
else: # [batch, seq, seq]
importance_scores = attention_weights.mean(dim=1) # [batch, seq]
return importance_scores
# Example usage for a game-playing agent
rl_agent = AttentionBasedRLAgent(
observation_dim=128, # Features from game state
action_dim=10, # Number of possible actions
hidden_dim=512,
num_heads=8
)
# Single observation processing
single_obs = torch.randn(32, 128) # Batch of 32 observations
action_probs, values, attn_weights = rl_agent(single_obs)
# Sequential observation processing (e.g., for Atari games)
sequential_obs = torch.randn(16, 20, 128) # 16 sequences of 20 observations
seq_action_probs, seq_values, seq_attn_weights = rl_agent(sequential_obs)
Multi-Agent Attention Systems
Attention mechanisms that coordinate between multiple agents:
class MultiAgentAttentionSystem(nn.Module):
def __init__(self, num_agents, agent_dim, communication_dim=256, num_heads=4):
super().__init__()
self.num_agents = num_agents
self.agent_dim = agent_dim
self.communication_dim = communication_dim
# Individual agent processors
self.agent_encoders = nn.ModuleList([
nn.Linear(agent_dim, communication_dim) for _ in range(num_agents)
])
# Cross-attention for inter-agent communication
self.inter_agent_attention = nn.ModuleList([
MultiHeadAttention(communication_dim, num_heads) for _ in range(num_agents)
])
# Individual agent decoders
self.agent_decoders = nn.ModuleList([
nn.Linear(communication_dim, agent_dim) for _ in range(num_agents)
])
def forward(self, agent_states):
"""
Process states from multiple agents with attention-based communication
Args:
agent_states: Tensor of shape [batch_size, num_agents, agent_dim]
"""
batch_size = agent_states.shape[0]
# Encode agent states to common communication space
encoded_states = []
for i, encoder in enumerate(self.agent_encoders):
encoded_states.append(encoder(agent_states[:, i, :]))
encoded_states = torch.stack(encoded_states, dim=1) # [batch, num_agents, comm_dim]
# Inter-agent attention and communication
communicated_states = []
attention_maps = []
for i, attention_layer in enumerate(self.inter_agent_attention):
# Agent i attends to all other agents
query = encoded_states[:, i:i+1, :] # [batch, 1, comm_dim]
attended_state, attn_weights = attention_layer(
query, encoded_states, encoded_states)
communicated_states.append(attended_state.squeeze(1))
attention_maps.append(attn_weights)
communicated_states = torch.stack(communicated_states, dim=1)
# Decode back to agent-specific spaces
final_states = []
for i, decoder in enumerate(self.agent_decoders):
final_states.append(decoder(communicated_states[:, i, :]))
final_states = torch.stack(final_states, dim=1)
return final_states, attention_maps
def analyze_communication_patterns(self, attention_maps):
"""Analyze how agents communicate with each other"""
# Convert attention maps to communication matrices
communication_matrices = []
for attn_map in attention_maps:
if attn_map.dim() == 4: # [batch, heads, 1, num_agents]
comm_matrix = attn_map.mean(dim=1).mean(dim=1) # [batch, num_agents]
else: # [batch, 1, num_agents]
comm_matrix = attn_map.mean(dim=1) # [batch, num_agents]
communication_matrices.append(comm_matrix)
return torch.stack(communication_matrices, dim=1) # [batch, num_agents, num_agents]
# Example usage for cooperative multi-agent system
multi_agent_system = MultiAgentAttentionSystem(
num_agents=5, # 5 cooperating agents
agent_dim=64, # Each agent has 64-dimensional state
communication_dim=256, # 256-dimensional communication space
num_heads=4
)
# Agent states from environment
agent_states = torch.randn(16, 5, 64) # Batch of 16 scenarios with 5 agents each
coordinated_states, communication_patterns = multi_agent_system(agent_states)
# Analyze communication
comm_analysis = multi_agent_system.analyze_communication_patterns(communication_patterns)
agent_attention_to_others = comm_analysis.mean(dim=0) # Average attention patterns
Hierarchical Attention Networks
Attention mechanisms operating at multiple levels of abstraction:
class HierarchicalAttentionNetwork(nn.Module):
def __init__(self, vocab_size, embed_dim=256, hidden_dim=512):
super().__init__()
self.word_embedding = nn.Embedding(vocab_size, embed_dim)
# Word-level attention
self.word_attention = nn.MultiheadAttention(embed_dim, num_heads=8)
self.word_lstm = nn.LSTM(embed_dim, hidden_dim // 2, bidirectional=True, batch_first=True)
# Sentence-level attention
self.sentence_attention = nn.MultiheadAttention(hidden_dim, num_heads=8)
self.sentence_lstm = nn.LSTM(hidden_dim, hidden_dim // 2, bidirectional=True, batch_first=True)
# Document-level attention
self.document_attention = nn.MultiheadAttention(hidden_dim, num_heads=8)
self.classifier = nn.Linear(hidden_dim, 2) # Binary classification example
def forward(self, documents):
"""
Process hierarchical text data (documents containing sentences containing words)
Args:
documents: Tensor of shape [batch_size, max_sentences, max_words]
"""
batch_size, max_sentences, max_words = documents.shape
# Process words in each sentence
sentence_representations = []
word_attention_weights = []
for sent_idx in range(max_sentences):
# Extract words for current sentence
sentence_words = documents[:, sent_idx, :] # [batch, max_words]
# Skip empty sentences
if torch.all(sentence_words == 0):
sentence_representations.append(
torch.zeros(batch_size, self.word_lstm.hidden_size * 2,
device=documents.device))
word_attention_weights.append(None)
continue
# Embed words
word_embeddings = self.word_embedding(sentence_words) # [batch, words, embed]
# LSTM processing
lstm_output, _ = self.word_lstm(word_embeddings)
# Word-level attention (attend to important words in sentence)
# Use mean of LSTM output as query for attention
query = lstm_output.mean(dim=1, keepdim=True).transpose(0, 1) # [1, batch, hidden]
key = lstm_output.transpose(0, 1) # [words, batch, hidden]
value = lstm_output.transpose(0, 1) # [words, batch, hidden]
attended_words, attn_weights = self.word_attention(query, key, value)
sentence_representation = attended_words.transpose(0, 1).squeeze(1) # [batch, hidden]
sentence_representations.append(sentence_representation)
word_attention_weights.append(attn_weights)
# Combine sentence representations
sentence_tensor = torch.stack(sentence_representations, dim=1) # [batch, sentences, hidden]
# Sentence-level processing
sentence_lstm_output, _ = self.sentence_lstm(sentence_tensor)
# Sentence-level attention
sent_query = sentence_lstm_output.mean(dim=1, keepdim=True).transpose(0, 1)
sent_key = sentence_lstm_output.transpose(0, 1)
sent_value = sentence_lstm_output.transpose(0, 1)
attended_sentences, sentence_attn_weights = self.sentence_attention(
sent_query, sent_key, sent_value)
document_representation = attended_sentences.transpose(0, 1).squeeze(1)
# Final classification
logits = self.classifier(document_representation)
return {
'logits': logits,
'document_representation': document_representation,
'sentence_attentions': sentence_attn_weights,
'word_attentions': word_attention_weights
}
# Example usage for document classification
han = HierarchicalAttentionNetwork(vocab_size=10000, embed_dim=256, hidden_dim=512)
# Sample documents (batch_size=8, max_sentences=10, max_words=20)
sample_documents = torch.randint(0, 10000, (8, 10, 20))
output = han(sample_documents)
document_logits = output['logits']
document_embeddings = output['document_representation']
sentence_attention = output['sentence_attentions']
Attention Visualization and Interpretation
Tools for understanding what attention mechanisms focus on:
import matplotlib.pyplot as plt
import seaborn as sns
class AttentionVisualizer:
def __init__(self):
pass
def visualize_attention_weights(self, attention_weights, tokens=None, title="Attention Weights"):
"""
Visualize attention weights as a heatmap
Args:
attention_weights: Attention weight tensor [batch, heads, seq_len, seq_len]
tokens: List of token strings for labeling
title: Title for the visualization
"""
# Take mean across batch and heads for simplicity
if attention_weights.dim() == 4:
avg_weights = attention_weights.mean(dim=0).mean(dim=0).cpu().detach()
else:
avg_weights = attention_weights.cpu().detach()
fig, ax = plt.subplots(figsize=(10, 8))
if tokens:
sns.heatmap(avg_weights.numpy(),
xticklabels=tokens[:avg_weights.shape[1]],
yticklabels=tokens[:avg_weights.shape[0]],
cmap='viridis', ax=ax)
else:
sns.heatmap(avg_weights.numpy(), cmap='viridis', ax=ax)
ax.set_title(title)
ax.set_xlabel('Key Positions')
ax.set_ylabel('Query Positions')
plt.tight_layout()
return fig
def visualize_attention_flow(self, attention_weights, layer_names=None):
"""Visualize attention flow through transformer layers"""
if isinstance(attention_weights, list):
# Multiple layers
fig, axes = plt.subplots(len(attention_weights), 1, figsize=(12, 4*len(attention_weights)))
if len(attention_weights) == 1:
axes = [axes]
for i, layer_attn in enumerate(attention_weights):
avg_attn = layer_attn.mean(dim=0).mean(dim=0).cpu().detach()
sns.heatmap(avg_attn.numpy(), cmap='viridis', ax=axes[i])
layer_name = layer_names[i] if layer_names else f"Layer {i+1}"
axes[i].set_title(layer_name)
axes[i].set_xlabel('Key Tokens')
axes[i].set_ylabel('Query Tokens')
else:
# Single layer
fig, ax = plt.subplots(figsize=(12, 8))
avg_weights = attention_weights.mean(dim=0).mean(dim=0).cpu().detach()
sns.heatmap(avg_weights.numpy(), cmap='viridis', ax=ax)
ax.set_title("Attention Flow")
ax.set_xlabel('Key Tokens')
ax.set_ylabel('Query Tokens')
plt.tight_layout()
return fig
def find_attention_peaks(self, attention_weights, threshold=0.1):
"""Find positions with high attention weights"""
avg_weights = attention_weights.mean(dim=0).mean(dim=0).cpu().detach()
# Find peaks above threshold
peaks = (avg_weights > threshold).nonzero(as_tuple=False)
peak_info = []
for peak in peaks:
query_pos, key_pos = peak[0].item(), peak[1].item()
weight = avg_weights[query_pos, key_pos].item()
peak_info.append({
'query_position': query_pos,
'key_position': key_pos,
'attention_weight': weight
})
return sorted(peak_info, key=lambda x: x['attention_weight'], reverse=True)
# Example usage with attention visualization
visualizer = AttentionVisualizer()
# Assuming we have attention weights from a model
# attention_weights = model_output['attention_weights'] # [batch, heads, seq, seq]
# Example attention weights for visualization
sample_attn_weights = torch.rand(16, 8, 10, 10) # Simulated attention weights
tokens = [f"Token_{i}" for i in range(10)]
# Visualize attention
# fig = visualizer.visualize_attention_weights(sample_attn_weights, tokens, "Sample Attention")
# plt.show()
# Find high attention connections
peaks = visualizer.find_attention_peaks(sample_attn_weights, threshold=0.8)
# print(f"Found {len(peaks)} high-attention connections")
class AttentionAnalyzer:
def __init__(self, model):
self.model = model
self.hooks = []
self.attention_outputs = []
def register_attention_hooks(self):
"""Register hooks to capture attention weights during forward pass"""
def attention_hook(module, input, output):
# Assuming output is (attended_values, attention_weights)
if isinstance(output, tuple) and len(output) > 1:
self.attention_outputs.append(output[1])
# Register hooks on attention modules
for name, module in self.model.named_modules():
if 'attention' in name.lower() or 'attn' in name.lower():
hook = module.register_forward_hook(attention_hook)
self.hooks.append(hook)
def remove_hooks(self):
"""Remove registered hooks"""
for hook in self.hooks:
hook.remove()
self.hooks.clear()
self.attention_outputs.clear()
def analyze_input_importance(self, input_tensor):
"""Analyze which input positions are most important based on attention"""
self.register_attention_hooks()
try:
with torch.no_grad():
output = self.model(input_tensor)
# Aggregate attention importance
importance_scores = torch.zeros(input_tensor.shape[1]) # [seq_len]
for attn_weights in self.attention_outputs:
if attn_weights.dim() == 4: # [batch, heads, seq, seq]
# Sum attention going to each position
importance = attn_weights.mean(dim=0).mean(dim=0).sum(dim=0) # [seq]
else: # [batch, seq, seq]
importance = attn_weights.mean(dim=0).sum(dim=0) # [seq]
importance_scores += importance.cpu()
# Normalize
importance_scores = importance_scores / len(self.attention_outputs)
return importance_scores
finally:
self.remove_hooks()
# Example usage of attention analyzer
# analyzer = AttentionAnalyzer(model)
# importance = analyzer.analyze_input_importance(input_sequence)
# important_positions = torch.topk(importance, k=5).indices
Applications in Specific Agent Types
Conversational Agents with Attention
Attention mechanisms enhance conversational agents by enabling context-aware responses:
class ConversationalAttentionAgent(nn.Module):
def __init__(self, vocab_size, embed_dim=512, max_context_length=512):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, embed_dim)
self.position_embedding = nn.Embedding(max_context_length, embed_dim)
# Multi-layer attention for context processing
self.context_attention = nn.ModuleList([
nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=8,
dim_feedforward=2048,
batch_first=True
) for _ in range(6)
])
# Response generation with attention to context
self.response_generator = nn.TransformerDecoderLayer(
d_model=embed_dim,
nhead=8,
dim_feedforward=2048,
batch_first=True
)
self.output_projection = nn.Linear(embed_dim, vocab_size)
self.dropout = nn.Dropout(0.1)
def forward(self, context_tokens, response_prefix=None, generate_length=50):
"""
Generate conversational responses using attention to context
Args:
context_tokens: Previous conversation context [batch, context_len]
response_prefix: Beginning of response (if any) [batch, prefix_len]
generate_length: Maximum length of generated response
"""
batch_size, context_len = context_tokens.shape
# Process context with attention
context_embeds = self.token_embedding(context_tokens)
context_positions = torch.arange(context_len, device=context_tokens.device)
context_positions = context_positions.unsqueeze(0).expand(batch_size, -1)
context_pos_embeds = self.position_embedding(context_positions)
context_repr = context_embeds + context_pos_embeds
context_repr = self.dropout(context_repr)
# Apply context attention layers
for layer in self.context_attention:
context_repr = layer(context_repr)
if response_prefix is not None:
# Generate response conditioned on prefix
return self._generate_conditioned_response(
context_repr, context_tokens, response_prefix)
else:
# Generate response from scratch
return self._generate_unconditioned_response(
context_repr, context_tokens, generate_length)
def _generate_conditioned_response(self, context_repr, context_tokens, response_prefix):
"""Generate response given a prefix"""
batch_size, prefix_len = response_prefix.shape
# Embed response prefix
prefix_embeds = self.token_embedding(response_prefix)
prefix_positions = torch.arange(prefix_len, device=response_prefix.device)
prefix_positions = prefix_positions.unsqueeze(0).expand(batch_size, -1)
prefix_pos_embeds = self.position_embedding(prefix_positions)
prefix_repr = prefix_embeds + prefix_pos_embeds
prefix_repr = self.dropout(prefix_repr)
# Use context as key/value, prefix as query
generated_repr = self.response_generator(prefix_repr, context_repr)
logits = self.output_projection(generated_repr)
return F.log_softmax(logits, dim=-1)
def _generate_unconditioned_response(self, context_repr, context_tokens, max_length):
"""Generate response from scratch using autoregressive decoding"""
batch_size = context_tokens.shape[0]
generated_tokens = []
current_input = torch.full((batch_size, 1),
self.token_embedding.padding_idx + 1, # Start token
device=context_tokens.device,
dtype=torch.long)
for _ in range(max_length):
# Embed current input
input_embeds = self.token_embedding(current_input)
input_positions = torch.arange(current_input.shape[1],
device=current_input.device)
input_positions = input_positions.unsqueeze(0).expand(batch_size, -1)
input_pos_embeds = self.position_embedding(input_positions)
input_repr = input_embeds + input_pos_embeds
input_repr = self.dropout(input_repr)
# Generate next token
next_repr = self.response_generator(input_repr, context_repr)
logits = self.output_projection(next_repr[:, -1:, :]) # Only last token
# Sample next token
next_token_logits = logits.squeeze(1)
next_token_probs = F.softmax(next_token_logits, dim=-1)
next_tokens = torch.multinomial(next_token_probs, 1)
generated_tokens.append(next_tokens)
current_input = torch.cat([current_input, next_tokens], dim=1)
# Stop if end token generated
if torch.any(next_tokens == self.token_embedding.padding_idx + 2): # End token
break
return torch.cat(generated_tokens, dim=1) if generated_tokens else None
# Example usage for chatbot
chatbot = ConversationalAttentionAgent(vocab_size=30000, embed_dim=512)
# Context conversation history
context = torch.randint(0, 30000, (8, 100)) # 8 conversations, 100 tokens each
# Generate response from scratch
generated_response = chatbot(context, generate_length=30)
# Or generate conditioned on prefix
response_prefix = torch.randint(0, 30000, (8, 5)) # 5-token prefixes
conditioned_response = chatbot(context, response_prefix=response_prefix)
Game-Playing Agents with Strategic Attention
Attention mechanisms help game agents focus on strategically relevant board positions:
class StrategicGameAttentionAgent(nn.Module):
def __init__(self, board_size, action_space, num_channels=128):
super().__init__()
self.board_size = board_size
self.action_space = action_space
self.num_channels = num_channels
# Convolutional layers for board feature extraction
self.conv_layers = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(128, num_channels, kernel_size=3, padding=1),
nn.ReLU()
)
# Spatial attention for board positions
self.spatial_attention = MultiHeadAttention(
embed_dim=num_channels,
num_heads=8
)
# Channel attention for feature importance
self.channel_attention = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(num_channels, num_channels // 8),
nn.ReLU(),
nn.Linear(num_channels // 8, num_channels),
nn.Sigmoid()
)
# Move prediction with attention to relevant positions
self.move_predictor = nn.Sequential(
nn.Linear(num_channels * board_size * board_size, 1024),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(1024, 512),
nn.ReLU(),
nn.Linear(512, action_space)
)
# Value estimation
self.value_estimator = nn.Sequential(
nn.Linear(num_channels * board_size * board_size, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, 256),
nn.ReLU(),
nn.Linear(256, 1),
nn.Tanh()
)
def forward(self, board_state):
"""
Process game board state and predict moves
Args:
board_state: Tensor of shape [batch, channels, height, width]
"""
batch_size = board_state.shape[0]
# Extract features with convolutional layers
conv_features = self.conv_layers(board_state) # [batch, channels, height, width]
# Apply spatial attention
# Reshape for attention: [batch, height*width, channels]
features_flat = conv_features.view(batch_size, self.num_channels, -1).transpose(1, 2)
# Self-attention on spatial positions
attended_features, spatial_weights = self.spatial_attention(
features_flat, features_flat, features_flat)
attended_features = attended_features.transpose(1, 2) # [batch, channels, height*width]
# Reshape back to spatial format
attended_spatial = attended_features.view(batch_size, self.num_channels,
self.board_size, self.board_size)
# Apply channel attention
channel_weights = self.channel_attention(conv_features) # [batch, channels, 1, 1]
channel_attended = conv_features * channel_weights
# Combine spatial and channel attention
final_features = attended_spatial + channel_attended
# Flatten for fully connected layers
flattened_features = final_features.view(batch_size, -1)
# Predict moves and estimate value
move_logits = self.move_predictor(flattened_features)
value_estimate = self.value_estimator(flattened_features).squeeze(-1)
return F.softmax(move_logits, dim=-1), value_estimate, spatial_weights
def analyze_strategic_focus(self, spatial_attention_weights):
"""Analyze which board positions received most attention"""
# Convert attention weights to spatial importance map
if spatial_attention_weights.dim() == 4: # [batch, heads, positions, positions]
# Average attention going TO each position
importance_map = spatial_attention_weights.mean(dim=1).mean(dim=1) # [batch, positions]
else: # [batch, positions, positions]
importance_map = spatial_attention_weights.mean(dim=1) # [batch, positions]
# Reshape to board dimensions
batch_size, num_positions = importance_map.shape
height = width = int(num_positions ** 0.5)
spatial_importance = importance_map.view(batch_size, height, width)
return spatial_importance
# Example usage for chess-like game
game_agent = StrategicGameAttentionAgent(
board_size=8, # 8x8 board
action_space=4672, # Chess action space size
num_channels=128
)
# Sample board states (batch_size=4, channels=3, height=8, width=8)
board_states = torch.randn(4, 3, 8, 8)
# Predict moves and values
move_probabilities, value_estimates, attention_weights = game_agent(board_states)
# Analyze strategic focus
strategic_focus = game_agent.analyze_strategic_focus(attention_weights)
Optimization and Efficiency Considerations
Efficient Attention Implementations
Memory-efficient and compute-optimized attention mechanisms:
class EfficientAttention(nn.Module):
def __init__(self, embed_dim, num_heads, use_flash_attention=False):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.use_flash_attention = use_flash_attention
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
# Try to import flash attention if available
self.flash_attn_fn = None
if use_flash_attention:
try:
from flash_attn import flash_attn_func
self.flash_attn_fn = flash_attn_func
except ImportError:
print("Flash attention not available, using standard attention")
self.use_flash_attention = False
def forward(self, query, key, value, mask=None):
batch_size = query.shape[0]
# Project to QKV
q = self.q_proj(query).view(batch_size, -1, self.num_heads, self.head_dim)
k = self.k_proj(key).view(batch_size, -1, self.num_heads, self.head_dim)
v = self.v_proj(value).view(batch_size, -1, self.num_heads, self.head_dim)
if self.use_flash_attention and self.flash_attn_fn is not None:
# Use flash attention for better efficiency
q = q.transpose(1, 2) # [batch, heads, seq, head_dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# Flash attention expects [batch, seq, heads, head_dim]
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
attn_output = self.flash_attn_fn(q, k, v, causal=False)
# Convert back to standard format
attn_output = attn_output.transpose(1, 2).contiguous()
else:
# Standard attention implementation
q = q.transpose(1, 2) # [batch, heads, seq, head_dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# Scaled dot-product attention
attn_scores = torch.matmul(q, k.transpose(-2, -1))
attn_scores = attn_scores / (self.head_dim ** 0.5)
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(attn_scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
# Reshape and project output
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, -1, self.embed_dim)
output = self.out_proj(attn_output)
return output
class LinearAttention(nn.Module):
"""Linear attention approximation for long sequences"""
def __init__(self, embed_dim, num_heads):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
self.elu = nn.ELU()
def forward(self, x):
batch_size, seq_length, embed_dim = x.shape
# Project to QKV
q = self.q_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
# Apply ELU activation and add 1 for positivity
q = self.elu(q) + 1
k = self.elu(k) + 1
# Compute linear attention
# This approximates softmax attention with linear complexity
kv = torch.einsum('bshd,bshv->bhvd', k, v) # Key-value interactions
z = 1 / (torch.einsum('bshd,bhd->bsh', q, k.sum(dim=1)) + 1e-6) # Normalization
attn_output = torch.einsum('bshd,bhvd,bsh->bshv', q, kv, z)
# Reshape and project
attn_output = attn_output.contiguous().view(batch_size, seq_length, embed_dim)
output = self.out_proj(attn_output)
return output
# Memory-efficient attention for very long sequences
class MemoryEfficientAttention(nn.Module):
def __init__(self, embed_dim, num_heads, chunk_size=1024):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.chunk_size = chunk_size
self.head_dim = embed_dim // num_heads
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
batch_size, seq_length, embed_dim = x.shape
if seq_length <= self.chunk_size:
# Use standard attention for shorter sequences
return self._standard_attention(x)
else:
# Chunk long sequences for memory efficiency
return self._chunked_attention(x)
def _standard_attention(self, x):
batch_size, seq_length, embed_dim = x.shape
# Standard multi-head attention
q = self.q_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
k = self.k_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
v = self.v_proj(x).view(batch_size, seq_length, self.num_heads, self.head_dim)
q = q.transpose(1, 2) # [batch, heads, seq, head_dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)
attn_scores = torch.matmul(q, k.transpose(-2, -1))
attn_scores = attn_scores / (self.head_dim ** 0.5)
attn_weights = F.softmax(attn_scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_length, embed_dim)
return self.out_proj(attn_output)
def _chunked_attention(self, x):
batch_size, seq_length, embed_dim = x.shape
# Process in chunks to save memory
outputs = []
for i in range(0, seq_length, self.chunk_size):
chunk_end = min(i + self.chunk_size, seq_length)
chunk = x[:, i:chunk_end, :]
# Process chunk with standard attention
chunk_output = self._standard_attention(chunk)
outputs.append(chunk_output)
# Concatenate chunk outputs
return torch.cat(outputs, dim=1)
# Example usage comparing efficiency
efficient_attn = EfficientAttention(embed_dim=512, num_heads=8)
linear_attn = LinearAttention(embed_dim=512, num_heads=8)
memory_efficient_attn = MemoryEfficientAttention(embed_dim=512, num_heads=8, chunk_size=512)
# Test with different sequence lengths
short_seq = torch.randn(8, 128, 512) # Short sequence
long_seq = torch.randn(4, 2048, 512) # Long sequence
# Efficient attention handles both sizes well
short_output = memory_efficient_attn(short_seq)
long_output = memory_efficient_attn(long_seq)
Training and Fine-Tuning Strategies
Attention-Aware Training Techniques
Methods for effectively training attention mechanisms:
class AttentionRegularization(nn.Module):
"""Regularization techniques specifically for attention mechanisms"""
def __init__(self, attention_heads):
super().__init__()
self.attention_heads = attention_heads
def entropy_regularization(self, attention_weights, lambda_entropy=0.01):
"""Encourage diverse attention patterns"""
# Compute entropy of attention distributions
if attention_weights.dim() == 4: # [batch, heads, seq, seq]
entropy = -torch.sum(attention_weights * torch.log(attention_weights + 1e-8), dim=-1)
mean_entropy = entropy.mean()
else: # [batch, seq, seq]
entropy = -torch.sum(attention_weights * torch.log(attention_weights + 1e-8), dim=-1)
mean_entropy = entropy.mean()
# Maximize entropy to encourage exploration
return -lambda_entropy * mean_entropy
def sparsity_regularization(self, attention_weights, lambda_sparse=0.01):
"""Encourage sparse attention patterns"""
# L1 regularization on attention weights
l1_penalty = torch.abs(attention_weights).mean()
return lambda_sparse * l1_penalty
def smoothness_regularization(self, attention_weights, lambda_smooth=0.01):
"""Encourage smooth attention transitions"""
if attention_weights.dim() == 4:
# Smooth across sequence dimension
diff = attention_weights[:, :, 1:, :] - attention_weights[:, :, :-1, :]
else:
diff = attention_weights[:, 1:, :] - attention_weights[:, :-1, :]
smoothness_penalty = torch.abs(diff).mean()
return lambda_smooth * smoothness_penalty
class CurriculumLearningScheduler:
"""Curriculum learning for attention training"""
def __init__(self, initial_seq_len=32, max_seq_len=512, growth_steps=10000):
self.initial_seq_len = initial_seq_len
self.max_seq_len = max_seq_len
self.growth_steps = growth_steps
self.current_step = 0
def get_current_seq_len(self):
"""Get current maximum sequence length based on training progress"""
progress = min(self.current_step / self.growth_steps, 1.0)
seq_len = self.initial_seq_len + int(
progress * (self.max_seq_len - self.initial_seq_len))
return seq_len
def step(self):
"""Update training step"""
self.current_step += 1
class AttentionTrainingWrapper:
"""Wrapper for training models with attention mechanisms"""
def __init__(self, model, criterion, optimizer, attention_reg=None):
self.model = model
self.criterion = criterion
self.optimizer = optimizer
self.attention_reg = attention_reg
self.scheduler = CurriculumLearningScheduler()
def training_step(self, batch, lambda_attn=0.01):
"""Perform one training step with attention-aware regularization"""
self.optimizer.zero_grad()
# Forward pass
outputs = self.model(batch['input'])
if isinstance(outputs, tuple) and len(outputs) > 1:
predictions = outputs[0]
attention_weights = outputs[1]
else:
predictions = outputs
attention_weights = None
# Main loss
main_loss = self.criterion(predictions, batch['target'])
# Attention regularization
attn_loss = 0
if attention_weights is not None and self.attention_reg is not None:
attn_loss += self.attention_reg.entropy_regularization(attention_weights)
attn_loss += self.attention_reg.sparsity_regularization(attention_weights)
attn_loss += self.attention_reg.smoothness_regularization(attention_weights)
# Total loss
total_loss = main_loss + lambda_attn * attn_loss
# Backward pass
total_loss.backward()
self.optimizer.step()
return {
'loss': total_loss.item(),
'main_loss': main_loss.item(),
'attn_loss': attn_loss.item() if isinstance(attn_loss, torch.Tensor) else attn_loss
}
def validation_step(self, batch):
"""Perform validation step"""
with torch.no_grad():
outputs = self.model(batch['input'])
if isinstance(outputs, tuple):
predictions = outputs[0]
else:
predictions = outputs
loss = self.criterion(predictions, batch['target'])
return {'loss': loss.item()}
# Example training setup
model = SomeAttentionModel(embed_dim=512, num_heads=8)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
attention_regularizer = AttentionRegularization(attention_heads=8)
trainer = AttentionTrainingWrapper(model, criterion, optimizer, attention_regularizer)
# Training loop example
# for epoch in range(num_epochs):
# for batch in train_loader:
# train_metrics = trainer.training_step(batch, lambda_attn=0.01)
#
# # Validation
# for batch in val_loader:
# val_metrics = trainer.validation_step(batch)
Case Study: Real-World Attention Implementation
Large-Scale Language Agent with Attention
Implementation of a practical large language agent with attention mechanisms:
class ProductionLanguageAgent(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.tokenizer = self._initialize_tokenizer(config.vocab_size)
self.embedding_layer = nn.Embedding(config.vocab_size, config.embed_dim)
self.position_encoding = self._create_positional_encoding(
config.max_seq_len, config.embed_dim)
# Transformer layers with attention
self.transformer_layers = nn.ModuleList([
self._create_transformer_block(config)
for _ in range(config.num_layers)
])
self.layer_norm = nn.LayerNorm(config.embed_dim)
self.output_projection = nn.Linear(config.embed_dim, config.vocab_size)
# Attention analysis components
self.attention_analyzer = AttentionVisualizer()
self.performance_monitor = PerformanceMonitor()
def _initialize_tokenizer(self, vocab_size):
"""Initialize tokenizer for the agent"""
# This would integrate with actual tokenizer implementation
return {
'vocab_size': vocab_size,
'pad_token_id': 0,
'bos_token_id': 1,
'eos_token_id': 2
}
def _create_positional_encoding(self, max_len, embed_dim):
"""Create positional encodings"""
pe = torch.zeros(max_len, embed_dim)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, embed_dim, 2).float() *
(-math.log(10000.0) / embed_dim))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
return pe.unsqueeze(0)
def _create_transformer_block(self, config):
"""Create a single transformer block"""
return nn.TransformerEncoderLayer(
d_model=config.embed_dim,
nhead=config.num_heads,
dim_feedforward=config.ff_dim,
dropout=config.dropout,
batch_first=True
)
def forward(self, input_ids, attention_mask=None):
"""
Process input and generate outputs with attention
Args:
input_ids: Token IDs [batch_size, seq_len]
attention_mask: Mask for padded tokens [batch_size, seq_len]
"""
batch_size, seq_len = input_ids.shape
# Embedding and positional encoding
embeddings = self.embedding_layer(input_ids)
position_encodings = self.position_encoding[:, :seq_len, :].to(embeddings.device)
hidden_states = embeddings + position_encodings
# Apply transformer layers with attention
attention_weights = []
for layer in self.transformer_layers:
# Store attention weights for analysis
if hasattr(layer.self_attn, 'register_forward_hook'):
def hook_fn(module, input, output):
if isinstance(output, tuple) and len(output) > 1:
attention_weights.append(output[1])
layer.self_attn.register_forward_hook(hook_fn)
hidden_states = layer(hidden_states, src_key_padding_mask=~attention_mask.bool() if attention_mask is not None else None)
# Final layer normalization and output projection
hidden_states = self.layer_norm(hidden_states)
logits = self.output_projection(hidden_states)
return {
'logits': logits,
'hidden_states': hidden_states,
'attention_weights': attention_weights
}
def generate_text(self, prompt, max_new_tokens=100, temperature=1.0, top_k=50):
"""Generate text using attention-aware decoding"""
# Tokenize prompt
input_ids = self._tokenize(prompt)
# Generate tokens one by one
generated_ids = input_ids.tolist()[0] if input_ids.dim() == 2 else input_ids.tolist()
for _ in range(max_new_tokens):
# Prepare input for current generation step
current_input = torch.tensor([generated_ids], device=input_ids.device)
# Forward pass
with torch.no_grad():
outputs = self.forward(current_input)
logits = outputs['logits'][0, -1, :] # Last token predictions
# Apply temperature scaling
logits = logits / temperature
# Top-k filtering
if top_k > 0:
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = float('-inf')
# Sample next token
probabilities = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probabilities, 1).item()
# Stop if end token generated
if next_token == self.tokenizer['eos_token_id']:
break
generated_ids.append(next_token)
# Decode generated tokens
return self._detokenize(generated_ids)
def _tokenize(self, text):
"""Convert text to token IDs"""
# Placeholder implementation
# In practice, this would use a real tokenizer like SentencePiece or BPE
import random
tokens = text.split()
token_ids = [random.randint(3, self.tokenizer['vocab_size']-1) for _ in tokens]
# Add BOS and EOS tokens
token_ids = [self.tokenizer['bos_token_id']] + token_ids + [self.tokenizer['eos_token_id']]
return torch.tensor([token_ids])
def _detokenize(self, token_ids):
"""Convert token IDs back to text"""
# Placeholder implementation
# In practice, this would use the tokenizer's decode method
return " ".join([f"token_{tid}" for tid in token_ids])
class PerformanceMonitor:
"""Monitor performance and attention quality during inference"""
def __init__(self):
self.metrics = {
'inference_time': [],
'attention_entropy': [],
'memory_usage': [],
'response_quality': []
}
def record_inference(self, start_time, end_time, attention_weights=None):
"""Record inference metrics"""
inference_time = end_time - start_time
self.metrics['inference_time'].append(inference_time)
if attention_weights is not None:
# Compute attention entropy
if attention_weights.dim() == 4:
entropy = -torch.sum(attention_weights * torch.log(attention_weights + 1e-8), dim=-1).mean()
else:
entropy = -torch.sum(attention_weights * torch.log(attention_weights + 1e-8), dim=-1).mean()
self.metrics['attention_entropy'].append(entropy.item())
def get_performance_report(self):
"""Generate performance report"""
report = {}
for metric_name, values in self.metrics.items():
if values:
report[metric_name] = {
'mean': np.mean(values),
'std': np.std(values),
'min': np.min(values),
'max': np.max(values)
}
return report
# Configuration for the production agent
class AgentConfig:
def __init__(self):
self.vocab_size = 50257 # GPT-2 vocabulary size
self.embed_dim = 768
self.num_heads = 12
self.num_layers = 12
self.ff_dim = 3072
self.max_seq_len = 1024
self.dropout = 0.1
# Initialize and use the production agent
config = AgentConfig()
language_agent = ProductionLanguageAgent(config)
# Example usage
sample_input = torch.randint(0, config.vocab_size, (4, 512)) # Batch of 4 sequences
attention_mask = torch.ones(4, 512) # No padding
# Forward pass with attention analysis
outputs = language_agent(sample_input, attention_mask)
logits = outputs['logits']
attention_weights = outputs['attention_weights']
# Text generation
# prompt = "Once upon a time"
# generated_text = language_agent.generate_text(prompt, max_new_tokens=50)
# print(f"Generated: {generated_text}")
# Performance monitoring
monitor = language_agent.performance_monitor
# start_time = time.time()
# outputs = language_agent(sample_input)
# end_time = time.time()
# monitor.record_inference(start_time, end_time, attention_weights[0] if attention_weights else None)
# performance_report = monitor.get_performance_report()
Future Directions and Emerging Trends
Multimodal Attention Systems
Next-generation attention mechanisms that integrate multiple modalities:
class MultimodalAttentionFusion(nn.Module):
"""Advanced multimodal attention systems"""
def __init__(self, modalities_config):
super().__init__()
self.modalities = list(modalities_config.keys())
self.modality_dims = {k: v['dim'] for k, v in modalities_config.items()}
# Modality-specific encoders
self.encoders = nn.ModuleDict({
modality: nn.Linear(dim, 512)
for modality, dim in self.modality_dims.items()
})
# Cross-modal attention layers
self.cross_attention = nn.ModuleDict({
f"{m1}_to_{m2}": MultiHeadAttention(512, 8)
for m1 in self.modalities
for m2 in self.modalities
if m1 != m2
})
# Self-attention within each modality
self.self_attention = nn.ModuleDict({
modality: MultiHeadAttention(512, 8)
for modality in self.modalities
})
# Fusion layer
self.fusion_attention = MultiHeadAttention(512, 8)
self.output_projection = nn.Linear(512, 1024) # Unified representation
def forward(self, inputs):
"""
Process multimodal inputs with cross-attention fusion
Args:
inputs: Dictionary with keys as modalities and values as tensors
"""
batch_size = next(iter(inputs.values())).shape[0]
# Encode each modality
encoded_modalities = {}
for modality, tensor in inputs.items():
encoded_modalities[modality] = self.encoders[modality](tensor)
# Apply self-attention within each modality
attended_modalities = {}
modality_attention_weights = {}
for modality, encoded in encoded_modalities.items():
attended, weights = self.self_attention[modality](encoded, encoded, encoded)
attended_modalities[modality] = attended
modality_attention_weights[modality] = weights
# Cross-modal attention
cross_attention_weights = {}
cross_attended = {}
for source_modality in self.modalities:
for target_modality in self.modalities:
if source_modality != target_modality:
key = f"{source_modality}_to_{target_modality}"
source_features = attended_modalities[source_modality]
target_features = attended_modalities[target_modality]
attended, weights = self.cross_attention[key](
target_features, source_features, source_features)
if target_modality not in cross_attended:
cross_attended[target_modality] = []
cross_attention_weights[target_modality] = []
cross_attended[target_modality].append(attended)
cross_attention_weights[target_modality].append(weights)
# Fuse all modalities with attention
# Stack all modality features
all_features = torch.cat(list(attended_modalities.values()), dim=1)
# Apply fusion attention
fused_features, fusion_weights = self.fusion_attention(
all_features, all_features, all_features)
# Global pooling and projection
global_features = fused_features.mean(dim=1) # Pool across sequence dimension
output = self.output_projection(global_features)
return {
'fused_representation': output,
'modality_attention': modality_attention_weights,
'cross_attention': cross_attention_weights,
'fusion_attention': fusion_weights
}
# Example configuration for multimodal system
multimodal_config = {
'text': {'dim': 768}, # Text embeddings
'image': {'dim': 2048}, # Image features (e.g., from CNN)
'audio': {'dim': 1024}, # Audio features
'video': {'dim': 4096} # Video features
}
multimodal_fusion = MultimodalAttentionFusion(multimodal_config)
# Example inputs
sample_inputs = {
'text': torch.randn(8, 50, 768), # 8 samples, 50 tokens, 768-dim embeddings
'image': torch.randn(8, 1, 2048), # 8 samples, 1 image feature vector
'audio': torch.randn(8, 100, 1024), # 8 samples, 100 audio frames
'video': torch.randn(8, 1, 4096) # 8 samples, 1 video feature vector
}
# Process multimodal inputs
multimodal_output = multimodal_fusion(sample_inputs)
fused_representation = multimodal_output['fused_representation']
Adaptive Attention Scaling
Attention mechanisms that dynamically adjust their behavior based on input complexity:
class AdaptiveAttention(nn.Module):
"""Attention mechanism that adapts its behavior based on input"""
def __init__(self, embed_dim, max_heads=16, min_heads=4):
super().__init__()
self.embed_dim = embed_dim
self.max_heads = max_heads
self.min_heads = min_heads
# Multiple attention heads with different configurations
self.attention_heads = nn.ModuleList([
MultiHeadAttention(embed_dim, num_heads=heads)
for heads in range(min_heads, max_heads + 1, 2)
])
# Complexity predictor to choose appropriate attention mechanism
self.complexity_predictor = nn.Sequential(
nn.Linear(embed_dim, embed_dim // 2),
nn.ReLU(),
nn.Linear(embed_dim // 2, len(self.attention_heads)),
nn.Softmax(dim=-1)
)
# Adaptive fusion mechanism
self.fusion_weights = nn.Parameter(torch.ones(len(self.attention_heads)))
def forward(self, query, key, value, complexity_hint=None):
"""
Forward pass with adaptive attention selection
Args:
query, key, value: Input tensors for attention
complexity_hint: Optional hint about input complexity
"""
batch_size, seq_len, embed_dim = query.shape
# Predict complexity and select appropriate attention
if complexity_hint is not None:
# Use provided complexity hint
complexity_scores = complexity_hint
else:
# Predict complexity from input statistics
input_stats = torch.cat([
query.mean(dim=1), # Mean across sequence
query.std(dim=1), # Std across sequence
torch.abs(query).max(dim=1)[0] # Max absolute values
], dim=-1)
complexity_scores = self.complexity_predictor(input_stats) # [batch, num_configs]
# Apply multiple attention mechanisms
attention_outputs = []
attention_weights_collections = []
for i, attention_head in enumerate(self.attention_heads):
output, weights = attention_head(query, key, value)
attention_outputs.append(output)
attention_weights_collections.append(weights)
# Adaptive fusion based on complexity scores
# Expand complexity scores to match sequence dimensions
complexity_expanded = complexity_scores.unsqueeze(1).unsqueeze(-1) # [batch, 1, configs, 1]
complexity_expanded = complexity_expanded.expand(-1, seq_len, -1, embed_dim)
# Weighted combination of attention outputs
stacked_outputs = torch.stack(attention_outputs, dim=2) # [batch, seq, configs, embed]
fused_output = (stacked_outputs * complexity_expanded).sum(dim=2)
return fused_output, attention_weights_collections
# Example usage
adaptive_attn = AdaptiveAttention(embed_dim=512, max_heads=12, min_heads=4)
# Inputs of varying complexity
simple_input = torch.randn(16, 50, 512) # Simple input
complex_input = torch.randn(16, 200, 512) # Complex input
# Process with adaptive attention
simple_output, simple_weights = adaptive_attn(simple_input, simple_input, simple_input)
complex_output, complex_weights = adaptive_attn(complex_input, complex_input, complex_input)
Conclusion
Attention mechanisms represent one of the most transformative innovations in artificial intelligence, enabling AI agents to focus on relevant information, process context effectively, and make intelligent decisions in complex environments. From the foundational scaled dot-product attention to advanced multimodal fusion systems, attention mechanisms have evolved to meet the growing demands of sophisticated AI applications.
The journey through attention mechanisms—from basic self-attention to hierarchical, sparse, and adaptive variants—reveals how crucial selective processing is for artificial cognition. Whether in natural language understanding, computer vision, game playing, or multi-agent coordination, attention provides the mechanism for intelligent prioritization and contextual reasoning.
Key takeaways for implementing attention in AI agents include:
Contextual Focus: Attention allows agents to dynamically allocate processing resources to the most relevant input elements while filtering noise.
Architectural Flexibility: Different attention variants (self-attention, cross-attention, causal attention) serve specialized purposes while maintaining mathematical coherence.
Computational Considerations: Efficient implementations like sparse attention and linear attention enable scaling to long sequences and large models.
Interpretability Benefits: Attention weights provide insights into model decision-making processes, enhancing transparency and trustworthiness.
Multimodal Integration: Cross-modal attention mechanisms enable seamless integration of multiple sensory inputs and knowledge sources.
Looking forward, the evolution of attention mechanisms promises even more sophisticated capabilities. Emerging trends in quantum attention, neuromorphic implementations, and biologically-inspired architectures suggest that attention will continue to be central to advancing artificial intelligence toward more human-like reasoning abilities.
The practical implementation examples provided—from conversation agents to game players to multimodal fusion systems—demonstrate how attention mechanisms can be tailored to specific application requirements while maintaining their essential capability for selective focus and contextual processing.
For practitioners developing AI agents, mastering attention mechanisms represents an essential step toward building truly intelligent, adaptive systems. By carefully selecting appropriate attention variants, optimizing for computational constraints, and leveraging attention for interpretable and robust agent behavior, engineers can create AI systems that approach the flexible, context-sensitive intelligence characteristic of biological cognition.
As attention mechanisms continue advancing—becoming more efficient, expressive, and integrated with other AI capabilities—they will play an ever-more-crucial role in realizing the vision of artificial general intelligence. The attention revolution in AI is far from complete, and its continuing evolution promises to reshape how we think about machine cognition and artificial intelligence systems.
References
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473.
Xu, K., Ba, J., Kiros, R., Cho, K., Courville, A., Salakhudinov, R., ... & Bengio, Y. (2015). Show, attend and tell: Neural image caption generation with visual attention. In International conference on machine learning (pp. 2048-2057).
Mnih, V., Heess, N., Graves, A., & others. (2014). Recurrent models of visual attention. Advances in neural information processing systems, 27.
Luong, M. T., Pham, H., & Manning, C. D. (2015). Effective approaches to attention-based neural machine translation. arXiv preprint arXiv:1508.04025.
Parikh, A. P., Täckström, O., Das, D., & Uszkoreit, J. (2016). A decomposable attention model for natural language inference. arXiv preprint arXiv:1606.01933.
Kaiser, Ł., & Bengio, S. (2016). Can active memory replace attention?. Advances in neural information processing systems, 29.
Child, R., Gray, S., Radford, A., & Sutskever, I. (2019). Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509.
Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150.
Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.
Published as part of the AI Agent Engineering series