在大模型开发的领域中，注意力机制无疑是核心基石之一，而自注意力机制与多头注意力机制更是其中的关键组成部分。从 Transformer 架构开始，这两种机制就凭借其强大的序列建模能力，深刻地改变了自然语言处理、计算机视觉等多个领域的发展格局。本文将深入剖析自注意力机制和多头注意力机制的原理，并结合代码实现帮助大家更好地理解和掌握。

自注意力机制：序列内部的关联捕捉

自注意力机制（Self-Attention）的核心目标是让序列中的每个元素都能关注到序列中其他相关元素的信息，从而更好地理解序列的上下文语义。在传统的循环神经网络（RNN）中，处理序列数据时存在着长距离依赖难以捕捉、无法并行计算等问题。而自注意力机制的出现，很好地解决了这些痛点。

自注意力机制的核心原理

自注意力机制通过计算序列中每个位置与其他所有位置之间的关联程度（即注意力权重），来实现对序列内部依赖关系的建模。具体来说，对于输入序列中的每个元素，都会生成三个向量：查询向量（Query，Q）、键向量（Key，K）和值向量（Value，V）。

查询向量用于表示当前元素想要 “查询” 的信息，键向量用于表示其他元素能够提供的 “信息标识”，通过计算查询向量与键向量的相似度，就能得到当前元素对其他元素的注意力权重。然后，再用这些注意力权重对值向量进行加权求和，就得到了当前元素经过注意力机制处理后的输出。

自注意力机制的计算步骤

自注意力机制的计算过程可以分为以下几个关键步骤：

• 生成 Q、K、V 向量：对于输入序列的每个元素向量\(x_i\)，通过三个不同的线性变换矩阵\(W_Q\)、\(W_K\)、\(W_V\)分别生成对应的查询向量\(q_i = x_iW_Q\)、键向量\(k_i = x_iW_K\)和值向量\(v_i = x_iW_V\)。

• 计算注意力分数：计算每个查询向量\(q_i\)与所有键向量\(k_j\)的点积，得到注意力分数\(score_{i,j}=q_i \cdot k_j\)。为了避免点积结果过大导致 softmax 函数梯度消失，通常会对注意力分数进行缩放处理，即\(score_{i,j}=\frac{q_i \cdot k_j}{\sqrt{d_k}}\)，其中\(d_k\)是键向量的维度。

• 计算注意力权重：对注意力分数应用 softmax 函数进行归一化处理，得到注意力权重\(attention\_weight_{i,j}=\frac{exp(score_{i,j})}{\sum_{k}exp(score_{i,k})}\)，归一化后的权重之和为 1，代表了当前元素对其他元素的关注程度。

• 计算输出向量：将注意力权重与对应的值向量进行加权求和，得到自注意力机制的输出向量\(output_i=\sum_{j}attention\_weight_{i,j} \cdot v_j\)。

多头注意力机制：多角度的特征学习

多头注意力机制（Multi-Head Attention）是在自注意力机制基础上的扩展，它通过将注意力机制分成多个并行的 “头”（Head），让模型能够从不同的角度捕捉序列中的依赖关系，从而增强模型的表达能力。

多头注意力机制的核心原理

单头自注意力机制可能只能捕捉到序列中某一种特定类型的依赖关系，而多头注意力机制通过设置多个并行的注意力头，每个注意力头都独立地进行自注意力计算，然后将各个头的输出进行拼接并通过线性变换得到最终结果。这样一来，模型就能够同时关注到序列中不同维度、不同类型的关联信息，极大地提升了模型的性能。

多头注意力机制的计算步骤

多头注意力机制的计算过程在自注意力机制的基础上增加了多头拆分和结果拼接的步骤：

• 参数拆分：将生成的 Q、K、V 向量按照头的数量进行拆分，假设共有\(h\)个头，则每个头对应的 Q、K、V 向量维度为\(d_k/h\)、\(d_k/h\)、\(d_v/h\)（其中\(d_k\)为键向量原始维度，\(d_v\)为值向量原始维度）。

• 多头并行计算：每个头独立地执行自注意力计算过程，得到每个头的输出结果。

• 结果拼接：将所有头的输出结果按照顺序进行拼接，得到一个维度为\(d_v\)的向量。

• 线性变换：对拼接后的向量进行一次线性变换，得到多头注意力机制的最终输出。

代码实现：基于 PyTorch 的自注意力和多头注意力

下面我们将基于 PyTorch 框架分别实现自注意力机制和多头注意力机制。

自注意力机制的代码实现

import torch

import torch.nn as nn

import torch.nn.functional as F

class SelfAttention(nn.Module):

def __init__(self, embed_dim):

super(SelfAttention, self).__init__()

self.embed_dim = embed_dim

# 定义生成Q、K、V的线性层

self.q_linear = nn.Linear(embed_dim, embed_dim)

self.k_linear = nn.Linear(embed_dim, embed_dim)

self.v_linear = nn.Linear(embed_dim, embed_dim)

# 输出线性层

self.out_linear = nn.Linear(embed_dim, embed_dim)

def forward(self, x, mask=None):

# x的形状：[batch_size, seq_len, embed_dim]

batch_size, seq_len, embed_dim = x.size()

# 生成Q、K、V向量

q = self.q_linear(x) # [batch_size, seq_len, embed_dim]

k = self.k_linear(x) # [batch_size, seq_len, embed_dim]

v = self.v_linear(x) # [batch_size, seq_len, embed_dim]

# 计算注意力分数：Q与K的转置点积，并进行缩放

scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(embed_dim, dtype=torch.float32)) # [batch_size, seq_len, seq_len]

# 应用掩码（如果有的话），掩码位置的分数设为极小值

if mask is not None:

scores = scores.masked_fill(mask == 0, -1e9)

# 计算注意力权重

attn_weights = F.softmax(scores, dim=-1) # [batch_size, seq_len, seq_len]

# 计算输出向量

output = torch.matmul(attn_weights, v) # [batch_size, seq_len, embed_dim]

output = self.out_linear(output) # [batch_size, seq_len, embed_dim]

return output, attn_weights

多头注意力机制的代码实现

class MultiHeadAttention(nn.Module):

def __init__(self, embed_dim, num_heads):

super(MultiHeadAttention, self).__init__()

self.embed_dim = embed_dim

self.num_heads = num_heads

# 每个头的维度

self.head_dim = embed_dim // num_heads

# 确保嵌入维度能被头数整除

assert self.head_dim * num_heads == embed_dim, "Embedding dimension must be divisible by number of heads"

# 定义生成Q、K、V的线性层

self.q_linear = nn.Linear(embed_dim, embed_dim)

self.k_linear = nn.Linear(embed_dim, embed_dim)

self.v_linear = nn.Linear(embed_dim, embed_dim)

# 输出线性层

self.out_linear = nn.Linear(embed_dim, embed_dim)

def forward(self, x, mask=None):

# x的形状：[batch_size, seq_len, embed_dim]

batch_size, seq_len, embed_dim = x.size()

# 生成Q、K、V向量

q = self.q_linear(x) # [batch_size, seq_len, embed_dim]

k = self.k_linear(x) # [batch_size, seq_len, embed_dim]

v = self.v_linear(x) # [batch_size, seq_len, embed_dim]

# 将Q、K、V拆分成多个头

# 形状变为：[batch_size, num_heads, seq_len, head_dim]

q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

# 计算注意力分数：Q与K的转置点积，并进行缩放

scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)) # [batch_size, num_heads, seq_len, seq_len]

# 应用掩码

if mask is not None:

# 掩码形状需调整为：[batch_size, 1, seq_len, seq_len]，以适配多头结构

scores = scores.masked_fill(mask.unsqueeze(1) == 0, -1e9)

# 计算注意力权重

attn_weights = F.softmax(scores, dim=-1) # [batch_size, num_heads, seq_len, seq_len]

# 计算每个头的输出

output = torch.matmul(attn_weights, v) # [batch_size, num_heads, seq_len, head_dim]

# 将多个头的输出拼接起来

output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim) # [batch_size, seq_len, embed_dim]

# 应用输出线性层

output = self.out_linear(output) # [batch_size, seq_len, embed_dim]

return output, attn_weights

总结与拓展

自注意力机制通过计算序列内部元素的关联权重，实现了对上下文信息的有效捕捉，解决了传统序列模型的诸多缺陷。而多头注意力机制则进一步通过多视角学习，增强了模型对复杂依赖关系的建模能力，成为了现代大模型架构中不可或缺的核心组件。

在实际应用中，我们还可以对注意力机制进行各种优化和拓展，比如引入相对位置编码来增强位置信息的建模、采用稀疏注意力来提高计算效率等。深入理解和掌握自注意力机制与多头注意力机制，对于从事大模型开发和研究的人员来说至关重要，它们是打开大模型奥秘之门的一把关键钥匙。

希望通过本文的讲解和代码实现，大家能够对自注意力机制和多头注意力机制有更清晰、更深入的认识，为后续的大模型开发学习和实践打下坚实的基础。