2.2 注意力机制的变体与改进

文档摘要

2.2 注意力机制的变体与改进读者读完这节，能够理解并应用Local Attention、Multi-Query Attention等改进方案，针对不同场景选择合适的注意力变体。 2.2.1 Local Attention (局部注意力) 传统自注意力机制最大的问题在于计算复杂度为O(n²)，当序列长度达到几千甚至上万时，内存和计算开销变得不可接受。Local Attention通过限制每个查询只能关注局部窗口内的键值对来解决这个问题。核心思想 Local Attention的基本原理是：对于序列中的每个位置，只关注周围一定范围内的其他位置，而非整个序列。

2.2 注意力机制的变体与改进

读者读完这节，能够理解并应用Local Attention、Multi-Query Attention等改进方案，针对不同场景选择合适的注意力变体。

2.2.1 Local Attention (局部注意力)

传统自注意力机制最大的问题在于计算复杂度为O(n²)，当序列长度达到几千甚至上万时，内存和计算开销变得不可接受。Local Attention通过限制每个查询只能关注局部窗口内的键值对来解决这个问题。

核心思想

Local Attention的基本原理是：对于序列中的每个位置，只关注周围一定范围内的其他位置，而非整个序列。


def local_attention(query, key, value, window_size=5):
    """
    Local Attention实现
    Args:
        query: [batch_size, seq_len, d_model]
        key: [batch_size, seq_len, d_model]
        value: [batch_size, seq_len, d_model]
        window_size: 局部窗口大小
    Returns:
        output: [batch_size, seq_len, d_model]
    """
    batch_size, seq_len, d_model = query.shape
    output = torch.zeros_like(query)
    
    for i in range(seq_len):
        # 定义局部窗口范围
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        
        # 只计算局部窗口内的注意力
        local_query = query[:, i:i+1, :]
        local_key = key[:, start:end, :]
        local_value = value[:, start:end, :]
        
        # 计算局部注意力
        local_output = scaled_dot_product_attention(
            local_query, local_key, local_value
        )
        output[:, i:i+1, :] = local_output
    
    return output

应用场景分析

适用场景：

长文本处理：处理长文档、书籍等，避免O(n²)复杂度
实时系统：对延迟敏感的实时应用
资源受限环境：移动端、嵌入式设备等计算资源有限的场景
保持局部结构：在语音识别、视频处理等需要保持局部时序信息的任务

性能对比分析

方法	复杂度	内存使用	特点
标准注意力	O(n²d)	O(n²)	全局关注，计算量大
局部注意力	O(wnd)	O(w²)	局部关注，计算量小
其中w为窗口大小

实际应用案例

在长文本摘要生成任务中，Local Attention表现出色：


class LocalAttentionEncoder(nn.Module):
    def __init__(self, d_model, window_size=5, num_heads=8):
        super().__init__()
        self.window_size = window_size
        self.num_heads = num_heads
        self.d_model = d_model
        self.d_k = d_model // num_heads
        
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 线性变换
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)
        
        # 多头注意力
        outputs = []
        for i in range(seq_len):
            # 局部窗口
            start = max(0, i - self.window_size // 2)
            end = min(seq_len, i + self.window_size // 2 + 1)
            
            # 分割多头
            q_heads = Q[:, i:i+1, :].view(1, 1, -1, self.d_k)
            k_heads = K[:, start:end, :].view(1, -1, self.num_heads, self.d_k)
            v_heads = V[:, start:end, :].view(1, -1, self.num_heads, self.d_k)
            
            # 计算注意力
            attn_output = self._scaled_dot_product_attention(q_heads, k_heads, v_heads)
            outputs.append(attn_output)
        
        # 合并结果
        output = torch.cat(outputs, dim=1)
        output = output.view(batch_size, seq_len, -1)
        return self.out_linear(output)

2.2.2 Multi-Query Attention (多查询注意力)

MQA是Transformer架构的重要改进，它通过使用多个查询向量来关注不同的信息维度，显著提升了模型的表达能力。

设计理念

MQA的核心思想是：每个位置可以使用多个查询来关注不同的语义维度，而不再是单一的查询向量。


def multi_query_attention(queries, key, value):
    """
    Multi-Query Attention实现
    Args:
        queries: [batch_size, seq_len, num_queries, d_k] 多个查询
        key: [batch_size, seq_len, d_k] 单个key
        value: [batch_size, seq_len, d_v] 单个value
    Returns:
        output: [batch_size, seq_len, num_queries, d_v]
    """
    batch_size, seq_len, num_queries, d_k = queries.shape
    
    # 展平查询维度
    queries_flat = queries.view(batch_size * seq_len * num_queries, d_k)
    
    # 计算注意力分数
    scores = torch.matmul(queries_flat, key.transpose(-2, -1)) / (d_k ** 0.5)
    attn_weights = F.softmax(scores, dim=-1)
    
    # 加权求和
    output = torch.matmul(attn_weights, value)
    
    # 恢复维度
    output = output.view(batch_size, seq_len, num_queries, -1)
    return output

实际应用场景

MQA特别适合以下场景：

多模态任务：需要同时关注视觉和文本信息
代码生成：同时关注语法结构和语义信息
多语言翻译：需要同时关注源语言和目标语言的信息
复杂推理：需要同时关注多个抽象层级的特征

性能优势

优势分析：

表达能力强：通过多个查询向量捕获更丰富的信息
计算效率高：共享key和value矩阵，减少参数量
灵活配置：可以根据任务需求调整查询数量


class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, num_queries=4):
        super().__init__()
        self.num_queries = num_queries
        self.q_linear = nn.Linear(d_model, d_model * num_queries)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model * num_queries, d_model)
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 生成多个查询
        Q = self.q_linear(x)  # [batch, seq, num_queries * d_k]
        Q = Q.view(batch_size, seq_len, self.num_queries, -1)
        
        # 单个key和value
        K = self.k_linear(x)
        V = self.v_linear(x)
        
        # 计算多查询注意力
        attn_output = multi_query_attention(Q, K, V)
        
        # 合并查询结果
        attn_output = attn_output.view(batch_size, seq_len, -1)
        return self.out_linear(attn_output)

2.2.3 Grouped Query Attention (分组查询注意力)

GQA是MQA的改进版本，它通过将查询分组进行并行计算，在保持表达能力的同时进一步优化了计算效率。

核心改进

GQA的主要改进在于：

分组并行：将查询分成多个组，每组独立计算
共享优化：同一组内的查询共享部分计算资源
可配置性：可以根据硬件特性调整分组策略


def grouped_query_attention(queries, key, value, group_size=2):
    """
    Grouped Query Attention实现
    Args:
        queries: [batch_size, seq_len, num_queries, d_k]
        key: [batch_size, seq_len, d_k]
        value: [batch_size, seq_len, d_v]
        group_size: 每组包含的查询数量
    Returns:
        output: [batch_size, seq_len, num_queries, d_v]
    """
    batch_size, seq_len, num_queries, d_k = queries.shape
    
    # 分组处理
    num_groups = (num_queries + group_size - 1) // group_size
    outputs = []
    
    for group_idx in range(num_groups):
        start_idx = group_idx * group_size
        end_idx = min((group_idx + 1) * group_size, num_queries)
        
        # 获取当前组的查询
        group_queries = queries[:, :, start_idx:end_idx, :]
        
        # 计算当前组的注意力
        group_output = multi_query_attention(group_queries, key, value)
        outputs.append(group_output)
    
    # 合并所有组的结果
    output = torch.cat(outputs, dim=2)
    return output

工程实现细节

在实际工程中，GQA的实现需要考虑以下优化点：


class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_queries=8, group_size=2):
        super().__init__()
        self.num_queries = num_queries
        self.group_size = group_size
        self.num_groups = (num_queries + group_size - 1) // group_size
        
        self.q_linear = nn.Linear(d_model, d_model * num_queries)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        
        # 每组独立的输出线性层
        self.group_linears = nn.ModuleList([
            nn.Linear(d_model * group_size, d_model)
            for _ in range(self.num_groups)
        ])
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 生成多个查询
        Q = self.q_linear(x)  # [batch, seq, num_queries * d_k]
        Q = Q.view(batch_size, seq_len, self.num_queries, -1)
        
        # 单个key和value
        K = self.k_linear(x)
        V = self.v_linear(x)
        
        # 分组计算
        group_outputs = []
        for group_idx in range(self.num_groups):
            start_idx = group_idx * self.group_size
            end_idx = min((group_idx + 1) * self.group_size, self.num_queries)
            
            group_queries = Q[:, :, start_idx:end_idx, :]
            group_output = self._group_attention(group_queries, K, V)
            group_outputs.append(group_output)
        
        # 合并组间结果
        output = torch.cat(group_outputs, dim=2)
        return self._merge_groups(output)
    
    def _group_attention(self, group_queries, key, value):
        """计算单个组的注意力"""
        batch_size, seq_len, num_queries, d_k = group_queries.shape
        
        # 计算注意力分数
        scores = torch.matmul(
            group_queries.view(-1, d_k),
            key.transpose(-2, -1)
        ) / (d_k ** 0.5)
        
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, value)
        
        return output.view(batch_size, seq_len, num_queries, -1)
    
    def _merge_groups(self, output):
        """合并不同组的结果"""
        batch_size, seq_len, num_queries, d_v = output.shape
        
        # 按组分离
        group_outputs = []
        for group_idx in range(self.num_groups):
            start_idx = group_idx * self.group_size
            end_idx = min((group_idx + 1) * self.group_size, num_queries)
            
            group_output = output[:, :, start_idx:end_idx, :]
            group_outputs.append(group_output)
        
        # 合并所有组
        merged = torch.cat(group_outputs, dim=2)
        return merged.view(batch_size, seq_len, -1)

2.2.4 Sliding Window Attention (滑动窗口注意力)

SWA结合了全局注意力和局部注意力的优势，通过滑动窗口机制实现了既关注全局又保持局部结构的平衡。

核心思想

SWA的工作机制：

全局关注：对序列中的重要位置进行全局关注
局部关注：对当前位置附近的区域进行局部关注
动态调整：根据任务需求动态调整全局和局部关注的比重


def sliding_window_attention(query, key, value, global_positions, window_size=5):
    """
    Sliding Window Attention实现
    Args:
        query: [batch_size, seq_len, d_model]
        key: [batch_size, seq_len, d_model]
        value: [batch_size, seq_len, d_model]
        global_positions: 需要全局关注的位置索引
        window_size: 局部窗口大小
    Returns:
        output: [batch_size, seq_len, d_model]
    """
    batch_size, seq_len, d_model = query.shape
    output = torch.zeros_like(query)
    
    for i in range(seq_len):
        if i in global_positions:
            # 全局关注
            local_output = scaled_dot_product_attention(
                query[:, i:i+1, :], key, value
            )
        else:
            # 局部关注
            start = max(0, i - window_size // 2)
            end = min(seq_len, i + window_size // 2 + 1)
            
            local_query = query[:, i:i+1, :]
            local_key = key[:, start:end, :]
            local_value = value[:, start:end, :]
            
            local_output = scaled_dot_product_attention(
                local_query, local_key, local_value
            )
        
        output[:, i:i+1, :] = local_output
    
    return output

应用场景

SWA特别适合以下场景：

长文本处理：在保持全局结构的同时处理长文本
多跳推理：需要同时关注直接相关和间接相关的信息
图像处理：需要同时关注局部纹理和全局结构
语音识别：需要同时关注音素级别的局部特征和语调级别的全局特征

性能优化策略


class EfficientSlidingWindowAttention(nn.Module):
    def __init__(self, d_model, window_size=5, global_ratio=0.1):
        super().__init__()
        self.window_size = window_size
        self.global_ratio = global_ratio
        self.d_model = d_model
        
        # 识别全局位置（可以根据启发式规则或学习）
        self.global_detector = nn.Linear(d_model, 1)
        
        # 局部注意力
        self.local_attention = LocalAttention(d_model, window_size)
        
        # 全局注意力
        self.global_attention = nn.MultiheadAttention(d_model, num_heads=8)
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 识别全局位置
        global_scores = self.global_detector(x).squeeze(-1)
        global_positions = torch.topk(
            global_scores, 
            int(seq_len * self.global_ratio)
        ).indices
        
        # 构建注意力掩码
        attention_mask = torch.ones(seq_len, seq_len)
        for i in range(seq_len):
            if i not in global_positions:
                start = max(0, i - self.window_size // 2)
                end = min(seq_len, i + self.window_size // 2 + 1)
                attention_mask[i, start:end] = 1
        
        # 执行注意力计算
        attn_output, _ = self.global_attention(
            x, x, x, 
            attn_mask=attention_mask
        )
        
        return attn_output

2.2.5 注意力变体的比较分析

变体	计算复杂度	内存使用	适用场景	表达能力
标准注意力	O(n²d)	O(n²)	短序列、高精度需求	最强
局部注意力	O(wnd)	O(w²)	长序列、实时系统	中等
多查询注意力	O(nqd)	O(n²)	多模态、复杂推理	较强
分组查询注意力	O(nqd/g)	O(n²/g)	大规模并行计算	较强
滑动窗口注意力	O(nwd + ngd)	O(n² + ng)	平衡全局与局部	中等

其中：w=窗口大小，q=查询数量，g=分组数量，n=序列长度，d=维度

实践建议

选择合适的注意力变体

1. 短序列任务（<512 tokens）：

标准注意力：最优精度
多查询注意力：增强表达能力

2. 中等序列任务（512-2048 tokens）：

局部注意力+全局门控：平衡效率与精度
滑动窗口注意力：保持结构信息

3. 长序列任务（>2048 tokens）：

纯局部注意力：最优先
分组查询注意力：并行计算优化

工程优化建议

批处理优化：利用GPU的并行计算能力
内存管理：对长序列采用流式处理
缓存机制：重复计算的中间结果缓存
混合精度：使用FP16/FP8加速计算

性能测试建议


def benchmark_attention_variants(sequence_lengths, variants):
    """测试不同注意力变体的性能"""
    results = {}
    
    for seq_len in sequence_lengths:
        for variant_name, variant_func in variants.items():
            # 生成测试数据
            x = torch.randn(32, seq_len, 512)
            
            # 测试时间
            start_time = time.time()
            with torch.no_grad():
                output = variant_func(x)
            end_time = time.time()
            
            # 记录结果
            if seq_len not in results:
                results[seq_len] = {}
            results[seq_len][variant_name] = end_time - start_time
    
    return results

# 使用示例
sequence_lengths = [512, 1024, 2048, 4096]
variants = {
    'standard': standard_attention,
    'local': local_attention,
    'multi_query': multi_query_attention,
    'grouped': grouped_query_attention,
    'sliding_window': sliding_window_attention
}

results = benchmark_attention_variants(sequence_lengths, variants)

2.2.6 注意力变体的组合使用

在实际应用中，不同注意力变体可以组合使用，形成更强大的注意力机制：


class HybridAttention(nn.Module):
    def __init__(self, d_model, window_size=5, num_queries=4):
        super().__init__()
        self.d_model = d_model
        
        # 多层次注意力
        self.local_attention = LocalAttention(d_model, window_size)
        self.multi_query_attention = MultiQueryAttention(d_model, num_queries)
        self.global_attention = nn.MultiheadAttention(d_model, num_heads=8)
        
        # 融合权重
        self.fusion_weights = nn.Linear(d_model, 3)
        self.gate = nn.Sigmoid()
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 计算三种注意力
        local_output = self.local_attention(x)
        multi_query_output = self.multi_query_attention(x)
        global_output, _ = self.global_attention(x, x, x)
        
        # 计算融合权重
        fusion_weights = self.fusion_weights(x)
        weights = self.gate(fusion_weights)
        
        # 加权融合
        output = (
            weights[:, :, 0:1] * local_output +
            weights[:, :, 1:2] * multi_query_output +
            weights[:, :, 2:3] * global_output
        )
        
        return output

本章总结

本节深入探讨了注意力机制的各种变体与改进：

Local Attention: 通过局部窗口解决了长序列的计算瓶颈
Multi-Query Attention: 使用多个查询向量增强了表达能力
Grouped Query Attention: 在保持表达能力的同时优化了计算效率
Sliding Window Attention: 实现了全局与局部关注的平衡

这些变体为不同场景下的注意力机制应用提供了丰富的选择工具。下一节将深入探讨注意力计算的具体优化策略。