深入理解多头注意力机制：原理与实现-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00174/article/details/148323877

深入理解多头注意力机制：原理与实现

d2l-zh 项目地址: https://gitcode.com/gh_mirrors/d2l/d2l-zh

多头注意力(Multi-head Attention)是现代深度学习模型中最重要的组件之一，尤其在Transformer架构中扮演着核心角色。本文将深入解析多头注意力的工作原理，并通过代码实现帮助读者掌握这一关键技术。

多头注意力的基本概念

传统的注意力机制在处理查询(Query)、键(Key)和值(Value)时，通常只学习单一的注意力模式。而多头注意力则通过并行计算多个注意力"头"(Head)，让模型能够同时关注输入序列的不同方面。

这种设计有三大优势：

可以捕捉序列中不同范围的依赖关系（如短距离和长距离依赖）
每个头可以学习不同的注意力模式
通过并行计算提高效率

多头注意力的数学原理

给定查询q ∈ ℝᵈᵩ、键k ∈ ℝᵈᵏ和值v ∈ ℝᵈᵥ，每个注意力头hᵢ的计算过程为：

hᵢ = f(Wᵢ⁽ᵠ⁾q, Wᵢ⁽ᵏ⁾k, Wᵢ⁽ᵛ⁾v) ∈ ℝᵖᵛ

其中：

Wᵢ⁽ᵠ⁾、Wᵢ⁽ᵏ⁾、Wᵢ⁽ᵛ⁾是可学习的线性变换矩阵
f是注意力汇聚函数（如缩放点积注意力）

最终输出是所有头的拼接结果经过另一个线性变换Wₒ：

输出 = Wₒ [h₁;...;hₕ] ∈ ℝᵖᵒ

代码实现解析

我们以PyTorch框架为例，展示多头注意力的实现细节：

class MultiHeadAttention(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super().__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = DotProductAttention(dropout)
        # 定义查询、键、值的线性变换
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        # 定义输出线性变换
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

关键实现技巧包括：

使用转置函数实现并行计算
将输入分割到多个头
计算注意力后重新拼接结果

多头注意力的应用示例

让我们通过一个简单例子测试多头注意力：

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                             num_hiddens, num_heads, 0.5)

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
output = attention(X, Y, Y, valid_lens)
print(output.shape)  # 输出形状应为(2,4,100)