多头注意力机制的例子

### 多头注意力机制实例

在深度学习中，多头注意力机制是一种用于处理序列数据的强大工具。通过并行运行多个不同的注意力层，模型可以从不同角度捕捉输入特征之间的关系。

#### 单头注意力计算过程
单头注意力机制可以表示为：

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

其中 \( Q \)，\( K \)，和 \( V \) 分别代表查询矩阵、键矩阵以及值矩阵；而 \( d_k \) 是键向量的维度大小[^1]。

#### 多头注意力实现方式
为了增强表达能力，引入了多头设计，在此架构下会创建若干组独立工作的 (query-key-value) 三元组，并最终将它们的结果拼接起来形成输出:

import torch
import torch.nn as nn

class MultiHeadedAttention(nn.Module):
    def __init__(self, num_heads, model_dim):
        super(MultiHeadedAttention, self).__init__()
        
        assert model_dim % num_heads == 0
        
        self.head_dim = model_dim // num_heads
        self.num_heads = num_heads
        
        self.qkv_linear = nn.Linear(model_dim, 3 * model_dim)
        self.out_proj = nn.Linear(model_dim, model_dim)

    def forward(self, x):
        batch_size, seq_len, _ = x.size()
        
        qkv = self.qkv_linear(x).reshape(batch_size, seq_len, self.num_heads, 3*self.head_dim)
        q,k,v = torch.chunk(qkv, chunks=3,dim=-1)
        
        scores = torch.einsum('bqhd,bkhd->bhqk', [q, k]) / math.sqrt(self.head_dim)
        attn_weights = F.softmax(scores, dim=-1)
        
        context = torch.einsum('bhqv,bvhd->bqhd', [attn_weights, v])
        output = self.out_proj(context.flatten(start_dim=2))
        
        return output

该代码片段定义了一个简单的 `MultiHeadedAttention` 类来展示如何构建一个多头注意力模块。这里使用 PyTorch 实现了线性变换、缩放点积操作及最后投影回原始空间的过程[^2]。