多头注意力机制(CV)

码上就来-0

2761人浏览 · 2024-12-06 09:58:14

码上就来-0 · 2024-12-06 09:58:14 发布

简单记录学习~

多头注意力代码通过 reshape 和 permute 来实现多个头的计算，是因为这种实现方式使用了 向量化计算 技术，在单个大张量上同时处理所有注意力头。然而这种方法并没有真正独立分开每个头的权重，而是将 多头权重共享计算 合并到一个更高效的矩阵运算中实现。

1. 多头注意力机制的数学原理

多头注意力的核心公式是：

${Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

在多头机制中，每个注意力头会有一套独立的权重 $W_Q, W_K, W_V$ ，分别用于计算查询 Q、键 K 和值 V。传统理解中，这意味着每个头需要单独计算这些矩阵的投影。

在实际实现中，使用了一个共享的大权重矩阵，通过向量化操作来同时计算所有头。

2. 如何通过 `reshape` 和 `permute` 实现多头注意力

a) 共享大矩阵计算

实际实现中，会将所有头的权重合并为一个大的投影矩阵 $W_Q$ （或 $W_K, W_V$ ）。假设输入的特征维度是 $d_{\text{model}}$ （例如 512），注意力头数是 h（例如 8），每头的维度是 dh= $d_{\text{model}} / h$ 。则：

原始投影矩阵的大小是 $d_{\text{model}} \times d_{\text{model}}$ 。
将其拆分为 h 个头时，权重矩阵整体大小保持不变，仍是 $d_{\text{model}} \times d_{\text{model}}$ ，但可以理解为每个头共享计算。

具体实现：

输入张量形状： $B, L, d_{\text{model}}$ （批量大小 B，序列长度 L，特征维度 $d_{\text{model}}$ ）。
通过矩阵乘法 $XW_Q$ ，得到一个新的张量，形状为 $B, L, d_{\text{model}}$ 。
然后通过 reshape 和 permute：
- reshape：将张量形状变成 $B, L, h, d_h$ ，表示将 $d_{\text{model}}$ 划分到 h 个头中。
- permute：调整维度顺序为 $B, h, L, d_h$ ，以方便并行处理每个头。

b) 并行计算多个头的注意力

在调整维度后，所有头的数据可以并行计算。由于所有头的权重是共享计算出来的矩阵的一部分，使用张量运算（例如矩阵乘法、softmax）可以同时处理所有头，而不需要显式区分每个头。

计算注意力分数： ${Attention Scores} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_h}}\right)V$ 这里的 Q 和 K 和 V 都是形状为 $B, h, L, d_h$ ，点积和 softmax 会在所有头上并行计算。

c) 输出的合并

在完成每个头的注意力计算后，结果张量形状为 $B, h, L, L$ 。接下来再与 $V$ 进行矩阵乘法最后得到输出 $B, h, L, d_h$ 。使用 permute 将张量维度调整回 $B, L, h, d_h$ 。最后再合并多头： $B, L, d_{model}$ 。最后，对合并后的结果再通过一个全连接层（权重矩阵 $W_O$ ）生成最终的输出 $B, L, d_{model}$ 。

3. 为什么没有单独的权重矩阵？

实际上，多头的权重矩阵是隐含存在的，只不过通过一次矩阵乘法计算 $Q, K, V$ 的大矩阵包含了所有头的权重。这种方式：

不是真的省略了独立权重，而是将权重的计算通过张量维度划分隐藏在共享的大权重矩阵中。
通过向量化操作，一次性计算所有头的结果，避免显式的逐头计算，提升了效率。

4. 代码实现与直观公式的区别

在直观公式中，每个头需要单独的 $W_Q, W_K, W_V$ 。而在代码实现中：

权重矩阵合并：代码将 $W_Q, W_K, W_V$ 合并为一个大矩阵，通过 reshape 将结果划分给不同的头。
向量化运算：并行处理所有头，利用了矩阵运算的高效性。
结果还原：通过 reshape 和 permute 实现了多头结果的拼接和恢复。

这种实现不仅减少了计算量，还充分利用了硬件的并行能力。

5. 总结

虽然代码中通过 reshape 和 permute 计算多个头，但这背后仍然有隐含的权重矩阵，只是将其向量化处理后显得统一：

没有独立权重矩阵 是因为权重已合并在一个大的权重矩阵中。
代码高效性：通过向量化操作，避免显式的逐头计算和合并。
这种实现方式既符合多头注意力机制的数学原理，也优化了计算性能。

实现代码如下~：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0):
        """
        多头注意力机制实现
        Args:
            embed_dim (int): 输入特征的维度（模型维度）。
            num_heads (int): 注意力头的数量。
            dropout (float): 注意力权重的dropout概率。
        """
        super(MultiHeadAttention, self).__init__()
        assert embed_dim % num_heads == 0, "embed_dim 必须能被 num_heads 整除"
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads  # 每个头的维度
        self.scale = self.head_dim ** -0.5  # 缩放因子，用于稳定点积注意力的梯度
        
        # 定义 Q, K, V 的线性变换
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        
        # 输出线性变换
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        """
        前向传播
        Args:
            query, key, value: [batch_size, seq_len, embed_dim]
            mask: [batch_size, 1, seq_len, seq_len] or None
        Returns:
            output: [batch_size, seq_len, embed_dim]
        """
        batch_size, seq_len, _ = query.size()
        
        # 线性变换 + 分头
        Q = self.q_proj(query).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)  # [B, num_heads, L, head_dim]
        K = self.k_proj(key).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)    # [B, num_heads, L, head_dim]
        V = self.v_proj(value).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)  # [B, num_heads, L, head_dim]

        # 点积注意力
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale  # [B, num_heads, L, L]
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))  # 应用注意力mask

        attn_probs = F.softmax(attn_scores, dim=-1)  # [B, num_heads, L, L]
        attn_probs = self.dropout(attn_probs)  # Dropout
        
        # 加权求和
        attn_output = torch.matmul(attn_probs, V)  # [B, num_heads, L, head_dim]
        
        # 合并多头
        attn_output = attn_output.permute(0, 2, 1, 3).contiguous()  # [B, L, num_heads, head_dim]
        attn_output = attn_output.view(batch_size, seq_len, self.embed_dim)  # [B, L, embed_dim]
        
        # 输出线性变换
        output = self.out_proj(attn_output)  # [B, L, embed_dim]
        
        return output