当前位置：首页 > news >正文

详细到用手撕transformer上半部分

news 来源：原创 2025/6/2 4:52:20

1. 剖析 Transformer 中的自注意力机制

在深度学习的广袤领域中，Transformer 架构宛如一颗璀璨明星，自 Google 团队于 2017 年在《Attention Is All You Need》论文中提出后，便彻底革新了自然语言处理领域，并且其影响力不断外延，在计算机视觉、语音识别等众多领域开疆拓土。Transformer 的核心驱动力便是自注意力机制（Self - Attention），它为序列数据的处理带来了全新的视角与高效的解决方案。

自注意力机制的核心运算元素

在深入探讨自注意力机制的运算流程前，先了解三个关键概念：查询（Query，简称 Q）、键（Keys，简称 K）、值（Values，简称 V）。

查询（Query）

Q 可被视为一个特征向量，它明确了我们在序列中期望探寻的内容，是引导模型关注重点信息的 “指南针”。

键（Keys）

对于输入序列中的每个元素，都对应着一个键向量，同样作为特征向量，它刻画了该元素所蕴含的信息，以及其在整个序列中可能具有重要性的时刻。键向量的设计目的在于，能够与查询向量相互作用，从而让模型识别出需要重点关注的元素。

值（Values）

每个输入元素同样关联着一个值向量，模型的终极目标便是借助这些值向量的加权组合，生成能够精准反映复杂上下文关系的表示，进而为诸如翻译、分类等下游任务的预测提供有力支撑。

自注意力机制的运算流程

假设我们已经通过一系列步骤，将原始输入数据成功转换为所需的 Q、K、V 矩阵或向量，接下来便进入自注意力机制的核心运算环节。

计算注意力分数矩阵

首先，对 Q 和 K 的转置进行点积（在矩阵乘法语境下），得到的结果用于衡量查询（Query）与键（Key）之间的相似性或关联程度，这一结果矩阵被称为注意力分数矩阵（Score Matrix）。用公式表示为：计算 Q 和 K 转置的点积，即 $QK^T$ 。

随后，将点积后的结果除以 $\sqrt{d_k}$ ，这里的 $d_k$ 是 Q 和 K 的维度。之所以进行这一操作，背后有着严谨的数学原理：假设 Q 和 K 均采样于均值为 0，方差为 1 的正态分布。设两个向量 $\vec{a}$ 和 $\vec{b}$ ，它们的点积为 $\vec{a}\cdot\vec{b}=\sum_{i = 1}^{d_k}a_ib_i$ 。根据方差的性质，其方差为 $Var(\vec{a}\cdot\vec{b}) = d_k$ $Var(a_i)=Var(b_i) = 1$ 时）。可以发现，点积后的方差与 $d_k$ 呈线性相关，当 $d_k$ 较大时，方差会显著增大，这会导致数据分布不均衡，输入到 softmax 函数中的 “能量” 过高，使得 softmax 输出趋近于 one - hot 向量（仅一个最大值为 1，其余元素接近 0）。在这种情况下，softmax 函数处于饱和状态，其梯度几乎为零，这将阻碍模型通过反向传播进行参数更新。而除以 $\sqrt{d_k}$ ，能够使缩放后的方差回归到合理区间（若初始化时方差为 1，缩放后方差仍为 1 ），有效规避了数值爆炸和梯度消失问题。

归一化操作

得到注意力分数矩阵并完成缩放后，下一步是对其进行 softmax 操作。softmax 函数的作用是实现归一化，它针对矩阵中的每一行进行运算，确保输出的每个值都介于 (0,1) 区间内，并且每一行所有输出值之和为 1 。通过 softmax 操作，注意力分数矩阵转变为注意力权重矩阵（Attention Weights Matrix），它清晰地表明了在计算注意力时，各个位置的重要程度。

计算最终输出

最后，将注意力权重矩阵与值矩阵 V 进行点积。这一步的本质是，以每个位置的注意力权重作为加权系数，对 V 中对应的值向量进行加权求和，从而得到能够充分反映输入序列上下文信息的最终表示。

代码实现与验证

以下是按照上述运算流程实现的代码：

import torch
import torch.nn.functional as F
import math
import pytorch_lightning as pldef scaled_dot_product(q, k, v, mask=None):d_k = q.size()[-1] #获取特征维度，q与k应为相同维度attn_logits = torch.matmul(q, k.transpose(-2, -1))#q与k的转置做点积attn_logits = attn_logits / math.sqrt(d_k)#缩放因子if mask is not None:attn_logits = attn_logits.masked_fill(mask == 0, -9e15)attention = F.softmax(attn_logits, dim=-1)#softmax归一化values = torch.matmul(attention, v)#与值矩阵点积return values, attention

为了验证代码的正确性，我们生成一组 Q、K、V 矩阵：

seq_len, d_k = 3, 2
pl.seed_everything(42)
q = torch.randn(seq_len, d_k).round(decimals=1)  
k = torch.randn(seq_len, d_k).round(decimals=1)  
v = torch.randn(seq_len, d_k).round(decimals=1)  
values, attention = scaled_dot_product(q, k, v)
print("Q\n", q)
print("K\n", k)
print("V\n", v)
print("Values\n", values)
print("Attention\n", attention)

输出如下：

手动验算过程

首先计算 Q 和 K 转置的点积，得到注意力分数矩阵。
然后将注意力分数矩阵除以 $\sqrt{d_k}$ （这里 $d_k = 2$ ，即除以 $\sqrt{2}$ ）。
接着进行 softmax 运算，softmax 的公式为: $Softmax(x_i)=\frac{e^{x_i}}{\sum_{j = 1}^{n}e^{x_j}}$ ，其中 $x_i$ 是输入向量中的第i 个元素，n 是向量长度。按照公式对每行元素分别计算，先计算每行元素的指数值，再对每行元素求和得到分母，最后将每行的每个元素除以分母，得到归一化后的注意力权重。
最后将注意力权重矩阵与值矩阵 V 进行点积，得到最终的输出值。通过手动计算结果与代码自动计算结果进行对比（由于手动计算时的精度保留以及设备运算精度差异，可能存在细微差别），可以验证代码实现的正确性。

至此，我们详细剖析了 Transformer 中自注意力机制的核心概念、运算流程，并通过代码实现与手动验算进行了验证。自注意力机制作为 Transformer 架构的基石，其高效的信息处理能力为众多深度学习任务带来了突破性的进展。我们将进一步深入探索 Transformer 的其他组成部分。

2. 从词嵌入到位置编码

在上一篇文章中，我们详细探讨了 Transformer 模型中自注意力机制的核心运算流程，了解了 Q、K、V 三个矩阵在注意力计算中的关键作用。但 Q、K、V 矩阵究竟是如何从原始输入数据中生成的呢？本文将深入探讨这一问题，并详细解析词嵌入（Word Embedding）和位置编码（Position Encoding）的原理与实现。

词嵌入：从单词到向量的映射

在自然语言处理中，计算机无法直接理解文本的语义信息，因此需要将文本转换为数字表示。词嵌入（Word Embedding）就是这样一种技术，它将单词映射到低维连续的向量空间中，使得语义相似的单词在向量空间中距离相近。

1. 词嵌入的基本概念

以猫狗分类任务为例，我们可以将猫和狗的特征（如大小、颜色）作为二维平面上的坐标轴，将每个样本投影到这个特征空间中。类似地，词嵌入将每个单词映射到一个 n 维空间中，使得语义相关的单词在这个空间中彼此靠近。

例如，"面条" 和 "米饭" 作为食物类词汇，在嵌入空间中可能指向相近的方向；而 "面条" 和 "汽车" 属于不同类别的词汇，它们的嵌入向量则可能指向不同的方向。这种 n 维空间不是人为定义的，而是模型在训练过程中自动学习得到的语义空间。

2. PyTorch 实现词嵌入

在 PyTorch 中，可以使用nn.Embedding层来实现词嵌入：

import torch
import torch.nn as nn# 定义词嵌入层
vocab_size = 10000  # 假设词表包含10,000个单词
d_model = 512       # 词嵌入维度
embedding = nn.Embedding(vocab_size, d_model)# 输入：一个批次的单词索引（batch_size=2, 序列长度=5）
input_ids = torch.LongTensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]])# 输出：嵌入后的向量（形状：[2, 5, 512]）
output = embedding(input_ids)
print(output.shape)  # torch.Size([2, 5, 512])

这里，vocab_size表示词汇表的大小，d_model表示每个单词的嵌入维度。输入是一个批次的单词索引，输出是对应的嵌入向量。例如，输出形状[2, 5, 512]表示批次大小为 2，每个序列包含 5 个单词，每个单词用 512 维的向量表示。

位置编码：为序列添加位置信息

1. 为什么需要位置编码？

虽然词嵌入能够捕捉单词的语义信息，但它无法表达单词在序列中的位置信息。例如，"小明教你机器学习" 和 "你教小明机器学习" 这两句话的语义完全不同，但如果仅使用词嵌入而不考虑顺序，它们的向量表示将是相同的。

为了解决这个问题，Transformer 引入了位置编码（Position Encoding），将序列中单词的位置信息附加到词嵌入中，使模型能够感知单词的顺序。

2. 位置编码的数学原理

Transformer 使用正弦和余弦函数来计算位置编码，其公式如下：

对于偶数维度： $PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

对于奇数维度： $PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

其中：

pos 是单词在序列中的位置（从 0 开始）
$d_{\text{model}}$ 是词向量的维度
i 是维度索引 $0 \leq 2i < d_{\text{model}}$

3. 位置编码的直观理解

让我们通过一个简单的例子来理解位置编码的工作原理。假设我们使用 4 维向量来表示每个字，对于句子 "你们好呀"，其中每个字的位置索引分别为 pos=0,1,2,3。根据位置编码公式，我们可以计算每个位置的编码向量：

对于位置 pos=0（"你"）：

类似地，可以计算其他位置的编码向量。最终，我们得到一个 4×4 的位置编码矩阵：

$\text{PE} = \begin{bmatrix} 0 & 1 & 0 & 1 \\ \sin(1/100) & \cos(1/100) & \sin(1/10000) & \cos(1/10000) \\ \sin(2/100) & \cos(2/100) & \sin(2/10000) & \cos(2/10000) \\ \sin(3/100) & \cos(3/100) & \sin(3/10000) & \cos(3/10000) \\ \end{bmatrix}$

4. 位置编码的代码实现

下面是位置编码的 Python 实现：

import numpy as np
import torch
import mathdef position_encoding(pos, d_model=4):"""计算单个位置的位置编码向量参数:pos: 单词在序列中的位置索引d_model: 词向量的维度返回:位置编码向量 (形状: [d_model])"""position_enc = np.zeros(d_model)for dim_idx in range(d_model):  # dim_idx表示维度索引if dim_idx % 2 == 0:# 偶数维度使用正弦函数position_enc[dim_idx] = math.sin(pos / (10000 ** (dim_idx / d_model)))else:# 奇数维度使用余弦函数position_enc[dim_idx] = math.cos(pos / (10000 ** ((dim_idx-1) / d_model)))return torch.tensor(position_enc, dtype=torch.float32)# 生成4x4位置编码矩阵
PE = torch.zeros(4, 4)
for pos in range(4):pe_vector = position_encoding(pos)PE[pos] = pe_vectorprint(f"Position {pos} encoding: {pe_vector}")

5. 位置编码的特点与优势

位置编码通过正弦和余弦函数的组合，为模型提供了以下优势：

绝对位置信息：每个位置都有唯一的编码向量
相对位置信息：不同位置之间的编码向量存在可计算的关系
泛化能力：可以处理任意长度的序列（超出训练长度的序列也能处理）
计算效率：无需额外训练参数，计算速度快

值得注意的是，Transformer 原始论文使用固定的位置编码公式，而后续研究（如 BERT）也探索了可学习的位置嵌入方法。

从嵌入向量到 Q、K、V 矩阵

通过词嵌入和位置编码，我们得到了包含语义和位置信息的嵌入向量X：

$X = \text{Embedding} + \text{PositionEncoding}$

接下来，我们需要通过三个独立的可学习权重矩阵，将嵌入向量X转换为 Q、K、V 矩阵：

其中 $W^Q$ 、 $W^K$ 和 $W^V$ 是模型在训练过程中学习的权重矩阵。这一步骤将每个位置的嵌入向量映射到查询（Query）、键（Key）和值（Value）空间，为后续的自注意力计算奠定基础。

总结

这章节详细解析了 Transformer 模型中词嵌入和位置编码的原理与实现，以及如何从嵌入向量生成 Q、K、V 矩阵。通过词嵌入，我们将单词映射到语义空间；通过位置编码，我们为模型提供了序列的位置信息。这两者的结合使得 Transformer 能够有效处理序列数据。

接着，我们将探讨多头注意力（Multi-Head Attention）的概念，这是 Transformer 模型中的另一个核心创新点。多头注意力机制允许模型从不同的子空间中获取信息，从而捕获更丰富的特征表示。

3. 理解 Transformer 之多头注意力机制

在之前对注意力机制的探讨中，我们了解到它通过计算查询（Query）、键（Key）和值（Value）之间的关系，实现对输入序列的加权平均，从而捕捉序列中的重要信息。然而，在实际应用中，一个序列元素往往需要从多个不同角度去关注，例如语法结构、语义角色、指代关系等。单一的注意力机制难以满足这种复杂需求，因此，多头注意力机制应运而生。

多头注意力机制的原理

多头注意力机制的核心思想是，对同一组特征使用多个不同的查询 - 键 - 值（query - key - value）三元组。具体而言，给定查询矩阵（Q）、键矩阵（K）和值矩阵（V），我们将它们分别转换为 h 个子查询、子键和子值矩阵，然后分别对这些子矩阵进行缩放点积注意力计算。每个注意力头可以学习到不同的注意力模式，使得模型能够同时关注输入序列的不同方面，比如语法、语义、位置等信息。最后，将各个注意力头的输出在特征维度上进行拼接（concatenation），并通过一个可学习的权重矩阵进行融合，从而增强模型的表达能力。

其数学表达式为：

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h)W^O$

其中，

$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

这里的 $W_i^Q$ 、 $W_i^K$ 、 $W_i^V$ 分别是第 i 个注意力头对应的可学习权重矩阵，用于将原始的 Q、K、V 矩阵投影到不同的子空间中， $W^O$ 是用于融合各个注意力头输出的可学习权重矩阵。

多头注意力机制的实现

当我们没有预先给定的查询向量（query）、键向量（key）和值向量（value）时，一种常见且有效的做法是将神经网络当前层的特征图直接设定为 Q、K 和 V。具体来说，假设输入特征图的形状为 $[B, T, d_{\text{model}}]$ ，其中 B 表示批量大小（batch size），T 表示序列长度， $d_{\text{model}}$ 表示特征图的隐藏维度。我们通过参数化投影来实现 Q、K、V 的分离。

以下是使用 PyTorch 实现多头注意力机制的代码：

class MultiheadAttention(nn.Module):def __init__(self, input_dim, embed_dim, num_heads):super().__init__()assert embed_dim % num_heads == 0, "Embedding dimension must be 0 modulo number of heads."self.embed_dim = embed_dimself.num_heads = num_headsself.head_dim = embed_dim // num_heads# Stack all weight matrices 1...h together for efficiency# Note that in many implementations you see "bias=False" which is optionalself.qkv_proj = nn.Linear(input_dim, 3 * embed_dim)self.o_proj = nn.Linear(embed_dim, embed_dim)self._reset_parameters()def _reset_parameters(self):# Original Transformer initialization, see PyTorch documentationnn.init.xavier_uniform_(self.qkv_proj.weight)self.qkv_proj.bias.data.fill_(0)nn.init.xavier_uniform_(self.o_proj.weight)self.o_proj.bias.data.fill_(0)def forward(self, x, mask=None, return_attention=False):batch_size, seq_length, embed_dim = x.size()qkv = self.qkv_proj(x)# Separate Q, K, V from linear outputqkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3 * self.head_dim)qkv = qkv.permute(0, 2, 1, 3)  # [Batch, Head, SeqLen, Dims]q, k, v = qkv.chunk(3, dim=-1)# Determine value outputsvalues, attention = scaled_dot_product(q, k, v, mask=mask)values = values.permute(0, 2, 1, 3)  # [Batch, SeqLen, Head, Dims]values = values.reshape(batch_size, seq_length, embed_dim)o = self.o_proj(values)if return_attention:return o, attentionelse:return o

代码解析

初始化部分
- 首先，我们创建一个名为 MultiheadAttention 的类，它继承自 nn.Module。在初始化函数 __init__ 中，我们进行一些必要的检查和参数设置。
- 确保 embed_dim 能被 num_heads 整除，这是因为在多头注意力机制中，我们希望每个注意力头处理相同维度的特征，若不能整除则无法平均分配。
- 存储总的嵌入维度 embed_dim、注意力头的数量 num_heads，并计算出每个头应该处理的维度数 head_dim（即 embed_dim // num_heads）。
- 初始化一个 QKV 投影矩阵 qkv_proj 和一个输出矩阵 o_proj。这里的 qkv_proj 实际上是将三个独立的权重矩阵合并为一个大的线性投影层，通过单次矩阵乘法同时生成 Q、K、V 的投影结果，再通过张量切分（split）得到三者，这样可以显著提升计算效率。o_proj 则用于融合各个注意力头的输出。
- 对投影矩阵的权重和偏置进行初始化，采用的是原始 Transformer 中的初始化方法，例如使用 nn.init.xavier_uniform_ 对权重进行初始化，将偏置数据填充为 0。
前向传播部分
- 在 forward 函数中，首先获取输入 x 的形状信息，即批量大小 batch_size、序列长度 seq_length 和嵌入维度 embed_dim。
- 使用 qkv_proj 对输入 x 进行投影，得到 qkv。此时 qkv 的形状为 [batch_size, seq_length, 3 * embed_dim]。
- 然后，通过 reshape 操作将 qkv 重新组织为多头形式，使其形状变为 [batch_size, seq_length, num_heads, 3 * head_dim]，保持前两个维度与输入 x 一致，将最后一个隐藏维度拆分为注意力头数量和每个头处理的维度（乘以 3 是因为将 Q、K、V 的维度合并在一起）。
- 接着，使用 permute 操作对维度进行重排，得到形状为 [batch_size, num_heads, seq_length, 3 * head_dim] 的张量，交换了 num_heads 和 seq_length 的顺序。
- 再使用 chunk 沿着张量的最后一个维度将 Q、K、V 平均分出来，得到三个矩阵，其形状均为 [batch_size, num_heads, seq_length, head_dim]。
- 得到 Q、K、V 后，就可以按照之前介绍的缩放点积注意力机制的计算方法进行计算，调用 scaled_dot_product 函数得到注意力输出 values 和注意力权重 attention。
- 由于之前进行了维度重排，现在需要将 values 的维度调整回来，先通过 permute 操作将其形状变为 [batch_size, seq_length, num_heads, head_dim]，再通过 reshape 操作将其投影到 [batch_size, seq_length, embed_dim]，得到与输入 x 相同维度顺序的张量。
- 最后，通过线性变换 o_proj 整合多头信息，得到最终输出 o，这一步相当于做逆变换，使输出便于后续的残差连接和层归一化等操作。

多头注意力机制的特性

多头注意力机制具有一个重要特性，即输入的置换等变性。也就是说，如果我们交换序列中的两个输入元素（暂不考虑批次维度），那么输出除了对应元素交换位置外，其余部分完全相同。这表明多头注意力机制实际上并非将输入严格视为一个序列，而是将其看作一组元素。正是由于这个特性，使得多头注意力模块和 Transformer 架构在处理各种序列数据时表现得非常强大且应用广泛。但也正因为如此，我们需要使用位置编码的方式将被忽略的顺序信息直接附加在各个元素中，这也是我们之前讨论位置编码的原因。

至此，我们详细梳理了多头注意力机制的关键部分。后续，我们将继续探讨 Transformer 中其他相关的深度学习流程。