深入理解注意力机制：MHA、MQA与GQA的演进与对比

没有评论

在大语言模型 (LLM) 快速发展的今天，注意力机制 (Attention Mechanism) 始终是核心技术之一。本文将深入探讨三种重要的注意力机制变体：多头注意力 (MHA)、多查询注意力(MQA) 和分组查询注意力(GQA)，分析它们的演进历程、技术特点及应用场景。

1. 多头注意力（Multi-Head Attention，MHA）

1.1 起源与发展

多头注意力机制最早由 Google 在 2017 年的论文《Attention Is All You Need》中提出，作为 Transformer 架构的核心组件。这个创新性的设计为后来的 BERT、GPT 等模型奠定了基础。

1.2 技术原理

MHA 的核心思想是将注意力机制的查询 (Q)、键(K) 和值 (V) 分成多个 ” 头 ”，每个头独立计算注意力，最后将结果合并。其数学表达式为：

def multi_head_attention(Q, K, V, num_heads):
    # 将输入分割成 num_heads 个头
    Q_heads = split_into_heads(Q, num_heads)  # [batch, num_heads, seq_len, d_k]
    K_heads = split_into_heads(K, num_heads)
    V_heads = split_into_heads(V, num_heads)

    # 对每个头计算注意力
    attention_heads = []
    for i in range(num_heads):
        score = dot_product(Q_heads[i], K_heads[i]) / sqrt(d_k)
        attention = softmax(score)
        head_output = dot_product(attention, V_heads[i])
        attention_heads.append(head_output)

    # 合并所有头的输出
    return concat(attention_heads)

1.3 优势与特点

并行性：多个头可以同时计算，提高了计算效率
特征多样性：不同的头可以关注输入数据的不同方面
强大的表达能力：能够捕获复杂的上下文关系

2. 多查询注意力（Multi-Query Attention，MQA）

2.1 背景介绍

MQA 是在 PaLM 模型中提出的一种优化方案，旨在提高大模型推理阶段的效率。这种设计在保持模型性能的同时，显著减少了计算开销。

2.2 技术创新

MQA 的核心创新在于让多个查询头共享同一组键值矩阵：

def multi_query_attention(Q, K, V, num_heads):
    # Q 仍然分成多个头
    Q_heads = split_into_heads(Q, num_heads)  # [batch, num_heads, seq_len, d_k]
    # K 和 V 不再分头，直接共享
    K_shared = K  # [batch, seq_len, d_k]
    V_shared = V

    attention_heads = []
    for i in range(num_heads):
        score = dot_product(Q_heads[i], K_shared) / sqrt(d_k)
        attention = softmax(score)
        head_output = dot_product(attention, V_shared)
        attention_heads.append(head_output)

    return concat(attention_heads)

2.3 性能提升

内存使用量降低约 30%
推理速度提升 40-50%
模型质量仅有轻微下降

3. 分组查询注意力（Grouped-Query Attention，GQA）

3.1 设计理念

GQA 是在 LLaMA 2 中引入的一种折衷方案，试图在 MHA 和 MQA 之间找到平衡点。它将查询头分组，每组内共享键值矩阵。

3.2 实现细节

def grouped_query_attention(Q, K, V, num_heads, num_groups):
    # 将头分组，每组共享 K /V
    heads_per_group = num_heads // num_groups
    Q_heads = split_into_heads(Q, num_heads)
    K_groups = split_into_heads(K, num_groups)
    V_groups = split_into_heads(V, num_groups)

    attention_heads = []
    for i in range(num_heads):
        group_idx = i // heads_per_group
        score = dot_product(Q_heads[i], K_groups[group_idx]) / sqrt(d_k)
        attention = softmax(score)
        head_output = dot_product(attention, V_groups[group_idx])
        attention_heads.append(head_output)

    return concat(attention_heads)

3.3 实际效果

根据 LLaMA 2 的实验数据：

比 MHA 减少了约 50% 的内存使用
比 MQA 提供更好的模型性能
训练稳定性优于 MQA

4. 对比分析

下面是三种机制在不同维度的详细对比：

特性	MHA	MQA	GQA
内存使用	高	低	中等
计算复杂度	O(H×N²×D)	O(N²×D)	O(G×N²×D)
推理速度	慢	快	较快
模型性能	最佳	轻微下降	接近 MHA
训练稳定性	高	中等	高