Transformer学习笔记

📅 2026/7/6 1:38:10
Transformer学习笔记
这段时间在找实习想着把之前学的东西都系统地复习一遍先从Transformer开始Transformer概述2017 年Google 在论文《 Attention is All you need 》中提出了 Transformer 模型其使用 Self-Attention 结构取代了在 NLP 任务中常用的 RNN 网络结构。相比 RNN 网络结构其最大的优点是可以并行计算和长距离信息捕捉。Transformer 的整体模型架构如下图所示目前这种 Encode-Decode 的 Transformer架构已经不常见了主流的模型如 GPT 使用的是 Decode-only 架构Bert 使用的是 Encode-only 架构。我们可以把 Transformer 简单理解为一个函数输入一个序列预测下一个token是什么。先来了解一下它由哪些部分组成输入Embedding 位置编码编码器/解码器Attention机制 LayerNorm 残差连接 → FFN 残差连接输出Liner层 Softmax层输入Embedding 位置编码Transformer 推理第一步文本embedding原始输入文本 → 分词处理Tokenization→ 得到 token 序列 → 词汇表映射Vocabulary Mapping → 得到 token ID 序列 → 嵌入层Embedding Layer因为 Transformer 自注意力是并行、无序的而不是像 RNN 一样串行处理数据序列所以必须引入外部位置编码来补全顺序。Transformer 推理第二步添加位置编码词嵌入矩阵 → 位置编码生成Positional Encoding→ 词嵌入矩阵位置编码 → 带位置信息的输入表示矩阵RoPE 的本质是利用复数乘法的几何意义模长相乘、角度相加不改变向量长度只通过旋转向量的角度来注入位置信息。 当带角度的Q和K进行内积时结果刚好只与它们的相对位置差有关。编码器/解码器Attention机制 LayerNorm 残差连接Attention机制概括来说就是考虑别的token对当前token在语义空间中的影响让它能够非常准确的用高纬度的向量表达语义特征。Transformer 推理第三步多头自注意力机制线性变换生成Query、Key、Value矩阵 → 计算注意力得分Attention Scores → 得到注意力得分矩阵 → Softmax归一化加权求和 → 多头并行处理 → 输出上下文感知的表示矩阵多头机制的本质是:与其用一个高维的注意力头去捕捉所有语义不如把维度切分成h个低维的”子头”(Head)让不同的头去关注不同的特征子空间(比如有的头关注语法有的头关注指代关系)最后再拼起来。多头注意力机制计算公式如下MultiHead(Q,K,V)Concat(head1,…,headh)WO \mathrm{MultiHead}(Q, K, V) \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h) W^OMultiHead(Q,K,V)Concat(head1​,…,headh​)WO其中每一个head的计算方式为headiAttention(Qi,Ki,Vi)softmax(QiKi⊤dk)Vi \mathrm{head}_i \mathrm{Attention}(Q_i, K_i, V_i) \mathrm{softmax}\left( \frac{Q_i K_i^\top}{\sqrt{d_k}} \right) V_iheadi​Attention(Qi​,Ki​,Vi​)softmax(dk​​Qi​Ki⊤​​)Vi​Concat把所有头计算出的结果在特征维度上拼接起来恢复成原来的维度。WO最后的线性投影矩阵。因为多个头是独立计算的拼接后需要经过一次线性变换混合各个头的信息输出最终结果。Transformer 推理第四步LayerNorm 残差连接LayerNorm即Layer Normalization是 Transformer中的层归一化主要是为了稳定训练过程防止梯度爆炸或消失。在《 Attention is All you need 》中位于残差连接之后前馈网络之前。但在现如今大多位于残差连接之前。先说一下残差连接首次在《Deep Residual Learning for Image Recognition》中被提出这篇论文是目前ai领域引用最高的论文。残差连接的核心思想是将输入直接“跳过子层加到输出上。也就是说遇到训练后效果下降的子层会直接赋予一个很低的权重将这个子层带来的训练效果的影响降到很低从而选择性地保留训练效果好的子层。残差连接解决了深层网络训练困难网络更深性能反而下降的问题。然后是LayerNorm输入矩阵经过注意力或者经过注意力和残差连接→ 计算均值 → 计算标准差 → 标准化 → 缩放和平移 → 得到归一化后的矩阵。就是正常的归一化过程没什么好多说的。Transformer 推理第五步FFN 残差连接引入前馈神经网络是为了通过非线性变换增加模型的表达能力。FFN的处理过程输入矩阵 → 线性层1升维→ 激活函数 → 线性层2降维→ 残差连接 → 最终输出输出Liner层 Softmax层Transformer 推理第六步Liner层 Softmax层解码器栈的输出是一个 float 向量。怎么把这个向量转换为一个词呢通过一个线性层再加上一个 Softmax 层实现。线性层是一个简单的全连接神经网络其将解码器栈的输出向量映射到一个更长的向量这个向量被称为 logits 向量。现在假设我们的模型有 10000 个英文单词模型的输出词汇表。因此 logits 向量有 10000 个数字每个数表示一个单词的分数。然后Softmax 层会把这些分数转换为概率把所有的分数转换为正数并且加起来等于 1。最后选择最高概率所对应的单词作为这个时间步的输出。手撕Transformer前面我们理清了整个 Transformer 的推理流现在来看看整个 Transformer 的代码是什么样的我们将依次实现Token Embedding、RoPE (旋转位置编码)、Multi-Head Attention、Feed Forward Network (FFN)最后把它们组装成一个完整的Transformer Block。1. Token Embeddingimporttorchimporttorch.nnasnnimportmathimporttorch.nn.functionalasFclassTokenEmbedding(nn.Module):def__init__(self,vocab_size,d_model):super().__init__()self.embeddingnn.Embedding(vocab_size,d_model)self.d_modeld_modeldefforward(self,x):# Embedding 结果需要乘以 sqrt(d_model) 进行放大# 目的避免后续加上位置编码时词向量自身的特征方差过小被位置编码淹没returnself.embedding(x)*math.sqrt(self.d_model)2. 旋转位置编码 (RoPE)现代大模型的标配利用复数乘法在 Attention 计算时注入相对位置信息。defprecompute_freqs_cis(dim:int,end:int,theta:float10000.0):预计算复数频率freqs1.0/(theta**(torch.arange(0,dim,2)[:(dim//2)].float()/dim))ttorch.arange(end,devicefreqs.device,dtypetorch.float32)freqstorch.outer(t,freqs).float()# 转换为复数张量 e^(ix) cos(x) i*sin(x)freqs_cistorch.polar(torch.ones_ones_like(freqs),freqs)returnfreqs_cisdefapply_rotary_emb(xq:torch.Tensor,xk:torch.Tensor,freqs_cis:torch.Tensor):应用旋转位置编码到 Q 和 K# 把最后两维转为复数xq_torch.view_as_complex(xq.float().reshape(*xq.shape[:-1],-1,2))xk_torch.view_as_complex(xk.float().reshape(*xk.shape[:-1],-1,2))# 调整 freqs_cis 形状以支持广播freqs_cisfreqs_cis.view(1,xq_.shape[1],1,xq_.shape[-1])# 复数乘法完成旋转xq_outxq_*freqs_cis xk_outxk_*freqs_cis# 转回实数并展平xq_outtorch.view_as_real(xq_out).flatten(3)xk_outtorch.view_as_real(xk_out).flatten(3)returnxq_out.type_as(xq),xk_out.type_as(xk)3. 多头自注意力 (Multi-Head Attention)多头注意力的核心在于利用view和transpose对矩阵进行拆分和重排使得各个“头”可以并行计算 Attention 分数。classMultiHeadAttention(nn.Module):def__init__(self,d_model,num_heads):super().__init__()assertd_model%num_heads0,d_model 必须能被 num_heads 整除self.num_headsnum_heads self.d_modeld_model self.head_dimd_model//num_heads# 线性映射层生成 Q, K, Vself.wqnn.Linear(d_model,d_model,biasFalse)self.wknn.Linear(d_model,d_model,biasFalse)self.wvnn.Linear(d_model,d_model,biasFalse)# 输出映射层self.wonn.Linear(d_model,d_model,biasFalse)defforward(self,x,freqs_cisNone,maskNone):batch_size,seq_len,_x.shape# 1. 线性映射并拆分为多头# 形状变化: (bs, seq_len, d_model) - (bs, seq_len, num_heads, head_dim)Qself.wq(x).view(batch_size,seq_len,self.num_heads,self.head_dim)Kself.wk(x).view(batch_size,seq_len,self.num_heads,self.head_dim)Vself.wv(x).view(batch_size,seq_len,self.num_heads,self.head_dim)# 2. 如果提供了 RoPE 的复数频率就在这里旋转 Q 和 Kiffreqs_cisisnotNone:Q,Kapply_rotary_emb(Q,K,freqs_cis)# 3. 维度重排准备计算 Attention# 形状变化: (bs, num_heads, seq_len, head_dim)QQ.transpose(1,2)KK.transpose(1,2)VV.transpose(1,2)# 4. 计算 Attention Score: Q * K^T / sqrt(d_k)# K.transpose(-2, -1) 把最后两维转置scorestorch.matmul(Q,K.transpose(-2,-1))/math.sqrt(self.head_dim)# 5. (可选) 加入 Mask比如在 Decoder 中防止看到未来信息ifmaskisnotNone:scoresscores.masked_fill(mask0,-1e9)# 6. Softmax 归一化并乘以 Vattn_weightsF.softmax(scores,dim-1)# out 形状: (bs, num_heads, seq_len, head_dim)outtorch.matmul(attn_weights,V)# 7. 拼接所有头 (Concat)# 形状恢复: (bs, seq_len, d_model)outout.transpose(1,2).contiguous().view(batch_size,seq_len,self.d_model)# 8. 最后经过一次线性投影returnself.wo(out)4. 前馈神经网络 (FFN)通常是先将维度放大 4 倍提取非线性特征再压缩回原来的维度。classFeedForward(nn.Module):def__init__(self,d_model,hidden_dimNone):super().__init__()ifhidden_dimisNone:hidden_dimd_model*4self.w1nn.Linear(d_model,hidden_dim)self.w2nn.Linear(hidden_dim,d_model)# 激活函数这里用经典的 ReLU大模型常换成 GeLU 或 SwiGLUself.actnn.ReLU()defforward(self,x):returnself.w2(self.act(self.w1(x)))5. 组装完整的 Transformer Block (Pre-LN 架构)我们采用现代大模型主流的Pre-LN先归一化再进入网络结构把上述组件拼装起来。classTransformerBlock(nn.Module):def__init__(self,d_model,num_heads):super().__init__()self.attentionMultiHeadAttention(d_model,num_heads)self.ffnFeedForward(d_model)# 定义两个 LayerNorm 层self.norm1nn.LayerNorm(d_model)self.norm2nn.LayerNorm(d_model)defforward(self,x,freqs_cisNone,maskNone):# Pre-LN 架构: x - LN - Attention - Add(x)hxself.attention(self.norm1(x),freqs_cis,mask)# x - LN - FFN - Add(x)outhself.ffn(self.norm2(h))returnout补一下分词器Tokenization当前的 LLM 几乎清一色采用子词级的分词算法以在完整单词和单字符之间取得平衡。以下介绍几种常用算法及其原理。子词分词的共同思想是:将词汇拆解为频繁出现的子单元从而压缩序列长度并缓解未登录词问题同时控制词表规模。字节对编码(BPE)BPE 从初始字符集开始迭代将文本中最频繁的相邻符号对合并为新符号重复该过程直到达到预定词汇大小。例如“hello”最初分成字母h、e、l、l、 oBPE 可能先合并频繁的he和1o成新符号得到he 1 1o进一步合并得到hel 1o乃至最终整体作为hello一个词元。现代LLM常用一种变体称为字节级BPE。它以256个字节作为基本单元确保任何Unicode文本都能表示。具体做法是将输入文本的每个字符拆解为UTF-8字节序列(例如汉字“中”可能表示为3个字节\xE4\xB8\xAD而英文字母‘A’则是单个字节\x41)然后对字节序列执行BPE合并。这使得初始词表固定为256个符号(对应单字节的所有可能取值)无需预先收集全部字符。WordPiece 词片段算法WordPiece 是另一种广泛使用的子词分词算法被BERT等模型采用。它与BPE的过程相似也是从字符集出发迭代合并符号但选择何种符号对进行合并的策略略有不同。具体而言WordPiece 并不总选择全局频次最高的符号对合并而是评估合并后对语言模型整体概率的提升选择**能最大化似然(或互信息)**的合并对。直观来说WordPiece在执行合并时会考虑若两个符号经常一起出现且单独出现时反而少见,则更有价值将它们合并。例如,假设合并 play 和 ##ing (##表示非词首)能显著提高语料的整体概率那么这个合并就会被优先选择。这种基于统计互信息的准则可以避免某些高频但意义松散的拼接理论上产生更有信息量的子词单元。不过总体而言WordPiece与BPE的效果和产生的词表相近二者都是通过固定数量的合并操作来生成子词词汇。