从零实现Transformer多头注意力：详解4个全连接函数的定义与应用

1次阅读

共计 2317 个字符，预计需要花费 6 分钟才能阅读完成。

多头注意力是 Transformer 的核心组件，但实现时经常遇到维度混乱和计算效率问题。具体来说：

需要同时维护 wq、wk、wv 三个权重矩阵和最后的输出全连接层
矩阵乘法需要考虑 batch、head、sequence 三个维度的并行计算
维度变换容易引发 hidden_size 和 head_num 不匹配的错误

我们需要定义 4 个全连接层：

query 的变换矩阵 wq
key 的变换矩阵 wk
value 的变换矩阵 wv
多头 concat 后的输出变换矩阵

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim=512, num_heads=8):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 定义 4 个全连接层
        self.wq = nn.Linear(embed_dim, embed_dim)  # query 变换
        self.wk = nn.Linear(embed_dim, embed_dim)  # key 变换
        self.wv = nn.Linear(embed_dim, embed_dim)  # value 变换
        self.fc_out = nn.Linear(embed_dim, embed_dim)  # 输出变换

        # Xavier 初始化
        nn.init.xavier_uniform_(self.wq.weight)
        nn.init.xavier_uniform_(self.wk.weight)
        nn.init.xavier_uniform_(self.wv.weight)
        nn.init.xavier_uniform_(self.fc_out.weight)

多头注意力的核心是并行计算多个头的注意力：

def forward(self, query, key, value, mask=None):
    batch_size = query.shape[0]

    # 1. 通过全连接层线性变换
    Q = self.wq(query)  # [batch, seq_len, embed_dim]
    K = self.wk(key)
    V = self.wv(value)

    # 2. 重组为多头形式
    Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
    K = K.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
    V = V.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

    # 3. 计算缩放点积注意力
    energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / (self.head_dim ** 0.5)
    if mask is not None:
        energy = energy.masked_fill(mask == 0, -1e10)
    attention = torch.softmax(energy, dim=-1)

    # 4. 应用注意力权重
    x = torch.matmul(attention, V)
    x = x.permute(0, 2, 1, 3).contiguous()
    x = x.view(batch_size, -1, self.embed_dim)

    # 5. 通过输出全连接层
    x = self.fc_out(x)
    return x

当 embed_dim 不能被 num_heads 整除时，会报错：

# 错误示例
model = MultiHeadAttention(embed_dim=512, num_heads=7)  # 512 不能被 7 整除

未正确初始化权重会导致梯度问题：

# 错误示例 - 未初始化权重
self.wq = nn.Linear(embed_dim, embed_dim)
nn.init.zeros_(self.wq.weight)  # 全零初始化会导致梯度消失

view+permute 方法（如上文实现）
einops 方法（需要安装 einops 库）

from einops import rearrange

# einops 实现
Q = rearrange(Q, 'b s (h d) -> b h s d', h=self.num_heads)

实测表明：

小矩阵：view+permute 更快（节省 einops 解析时间）
大矩阵：einops 更优（优化了内存布局）

可以通过以下方式估算计算量：

def count_flops(batch_size, seq_len, embed_dim, num_heads):
    # 矩阵乘法 FLOPs: 2 * batch * seq_len * embed_dim^2
    # 注意力计算 FLOPs: 2 * batch * num_heads * seq_len^2 * head_dim
    total = 2 * batch_size * seq_len * (embed_dim**2) * 3  # Q,K,V
    total += 2 * batch_size * num_heads * (seq_len**2) * (embed_dim//num_heads)
    total += 2 * batch_size * seq_len * (embed_dim**2)  # 输出全连接
    return total