多头注意力机制实战：如何高效定义4个全连接函数实现wq,wk,wv矩阵及多头拼接

1次阅读

没有评论

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

在实现 Transformer 的多头注意力机制时，我们常常会遇到两个主要问题：

代码冗余 ：手动为每个头定义独立的 Query(Wq)、Key(Wk)、Value(Wv) 权重矩阵，会导致大量重复代码
性能瓶颈：多头拼接时的显存占用高，特别是在处理长序列时，内存消耗会显著增加

传统实现方式往往需要为每个头单独定义三个全连接层，这不仅增加了代码量，还可能导致显存使用效率低下。

我们提出使用 4 个全连接函数来优化实现：

三个全连接层分别处理 Q、K、V 矩阵
第四个全连接层用于多头拼接后的最终投影

前三个全连接层分别对应 Wq、Wk、Wv 矩阵变换
第四个全连接层是多头注意力拼接后的投影层，将多个头的输出映射回原始维度
这样可以保持模块化设计，同时避免显存碎片化

假设输入维度为d_model，头数为h，每个头的维度为d_k = d_model/h

Wq/Wk/Wv 层的输入：(batch, seq_len, d_model)
Wq/Wk/Wv 层的输出：(batch, seq_len, h, d_k)
投影层的输入：(batch, seq_len, h*d_k)
投影层的输出：(batch, seq_len, d_model)

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h):
        super().__init__()
        assert d_model % h == 0, "d_model must be divisible by h"

        self.d_model = d_model
        self.h = h
        self.d_k = d_model // h

        # 定义 4 个全连接层
        self.wq = nn.Linear(d_model, d_model)  # Query 转换
        self.wk = nn.Linear(d_model, d_model)  # Key 转换
        self.wv = nn.Linear(d_model, d_model)  # Value 转换
        self.wo = nn.Linear(d_model, d_model)  # 多头拼接后投影

        # 初始化权重
        self._init_weights()

    def _init_weights(self):
        # 使用 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.wo.weight)

        # 偏置初始化为 0
        nn.init.zeros_(self.wq.bias)
        nn.init.zeros_(self.wk.bias)
        nn.init.zeros_(self.wv.bias)
        nn.init.zeros_(self.wo.bias)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 1. 线性投影得到 Q,K,V
        q = self.wq(q)  # (batch, seq_len, d_model)
        k = self.wk(k)
        v = self.wv(v)

        # 2. 分割多头
        q = q.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)  # (batch, h, seq_len, d_k)
        k = k.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        v = v.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)

        # 3. 计算 Scaled Dot-Product Attention
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)  # (batch, h, seq_len, d_k)

        # 4. 合并多头
        output = output.transpose(1, 2).contiguous()  # (batch, seq_len, h, d_k)
        output = output.view(batch_size, -1, self.d_model)  # (batch, seq_len, d_model)

        # 5. 最终投影
        output = self.wo(output)

        return output

我们对比了传统实现和优化后的 4 全连接层实现，结果如下：