尧图建网站 尧图建网站 YAOTU WEB BUILD 免费咨询
ARTICLE DETAIL

资讯详情

深耕网站建设与建站编程的一线实战洞察。

Transformer架构图解:从自注意力机制到PyTorch最小实现

Transformer架构图解:从自注意力机制到PyTorch最小实现 直接给结论这几年你看到的大模型产品——GPT 系列、LLaMA、Gemini、文心一言、DeepSeek不管产品层有多大差异主干网络几乎都建立在同一个架构上Transformer。Transformer 最早由 Google 在 2017 年发表的论文《Attention Is All You Need》中提出它的核心改动不是把神经网络做得更深而是把“注意力机制”拉到绝对主位放弃了循环神经网络RNN的时序递归和卷积神经网络CNN的局部感受野。这个决策带来三个非常直接的好处能并行训练、能建模长距离依赖、能靠堆参数量持续提升效果。今天的 GPT、BERT、ViT、Swin Transformer本质上都是同一个底层思路在不同数据模态上的延伸。这篇博客会从图解角度把 Transformer 拆开整体架构长什么样自注意力机制里的 Q、K、V 到底在算什么位置编码为什么不能省多头注意力又是怎么工作的以及从 NLP 到视觉领域它如何迁移。最后我会给出一份可以复制运行的 PyTorch 最小实现帮你完成一次“手撕 Transformer”的练习。适合这样的读者学过深度学习和 RNN/CNN但对 Transformer 一直停留在“看过概念图”的阶段准备面试大模型岗位需要系统梳理原理或者要开始读 GPT、LLaMA、ViT 源码却卡在注意力实现细节上。1. Transformer 核心要点速览先把关键信息放在前面方便快速判断这篇文章是否与你的需求对上。维度说明全称Transformer出自论文《Attention Is All You Need》提出时间2017 年Google核心创新自注意力机制Self-Attention取代循环和卷积基础组成输入嵌入、位置编码、多头注意力、残差连接、LayerNorm、前馈网络主流用法Encoder-Decoder、Encoder-only、Decoder-only典型模型BERT、GPT、T5、LLaMA、ViT、Swin Transformer处理模态文本、图像、语音、视频、多模态最大优势并行训练、长程依赖建模、规模化扩展主要短板计算复杂度随序列长度平方增长显存消耗高学习门槛需要理解矩阵运算、Softmax、残差连接和 LayerNorm需要提醒一点网上很多文章把 Transformer 等同于“注意力机制”这个理解过于简化。注意力是一个组件Transformer 是一整套可训练的深度网络架构。你可以用注意力做机器翻译、做推荐、做图像分类但只有把注意力、归一化、残差、前馈网络组织成特定结构它才叫 Transformer。2. 为什么最终是 Transformer它解决了什么问题理解 Transformer 最好的方式是先看它替代的模型到底卡在哪里。2.1 RNN 的问题顺序计算与长程遗忘RNN 处理序列时必须按时间步逐个输入。比如句子“小明昨天在上海今天已经飞到北京”模型要读完最后一个词才知道“飞到北京”的主语是“小明”。这个信息如果离当前预测位置很远梯度在反向传播中就会指数级衰减模型学不到长距离依赖。LSTM 和 GRU 引入门控机制缓解了一部分遗忘问题但本质上仍是顺序计算无法充分利用 GPU 并行能力。训练效率是另一个痛点。假设一个句子有 100 个 tokenRNN 需要 100 步才能计算完每一步都依赖上一步的隐状态GPU 即使有几千个计算核心也只能串行推进。当语料规模到几十亿 token 时这种串行限制会让训练成本变得不可接受。2.2 CNN 的问题感受野有限CNN 靠卷积核捕捉局部特征相邻位置的特征容易融合但距离很远的位置之间需要堆叠很多层才能建立联系。对图像来说这是一个合理的归纳偏置因为像素近邻相关性最强对长文本来说就很别扭一句话里主语和谓语可能隔了很长一段修饰成分。CNN 处理这类问题需要靠深度来补偿感受野网络结构因此越堆越深。2.3 Transformer 的解题思路Transformer 不按顺序读句子。它一次性接收整段序列通过自注意力机制让任意两个位置直接计算相关性。第一个字和最后一个字之间的关联不需要经过中间所有位置逐层传递一步就能建起来。这套设计让训练可以完全并行批量输入里多个 token 的矩阵运算同时执行模型规模越大并行优势越明显。这也是为什么大模型基本都沿着 Transformer 路线走而不是回到 LSTM——成本结构决定了技术路线。3. Transformer 整体架构图解原始 Transformer 分为 Encoder 和 Decoder 两部分结构上可以画成下面这样输入序列 │ ▼ Token Embedding ───┐ 位置编码 ──────────┼──► 相加 ──► 多头自注意力 ──► 残差连接 ──► LayerNorm │ │ │ ▼ │ 前馈网络FFN │ │ └────────────────── 残差连接 ──────────────────► LayerNorm │ ▼ Encoder 输出实际实现里每个大模块内部还会嵌套子层。一个标准 Encoder Block 通常包含四个关键操作多头自注意力Multi-Head Self-Attention残差连接Residual Connection层归一化LayerNorm前馈网络Feed-Forward Network通常由两个线性层加激活函数组成Decoder 与 Encoder 的区别在于它内部多了一个“掩码多头注意力”用来保证生成当前 token 时看不到未来的 token。这个掩码就是把未来位置的注意力分数设为一个非常大的负数Softmax 之后接近 0模型只能依赖当前位置及之前的信息。3.1 为什么需要残差连接深层网络在训练时容易出现梯度消失问题。残差连接的意思是输出 子层结果 输入。即使子层内部学不到有效变换梯度也能通过“输入”这一条捷径回传。这个设计在 ResNet 里被验证过Transformer 直接沿用让网络可以堆到几十层甚至上百层。3.2 为什么需要 LayerNormBatchNorm 在 NLP 里不好用因为不同句子的长度不一样批次统计不稳定。LayerNorm 对单个样本的所有特征维度做归一化不依赖 batch 大小训练和推理行为一致。GPT、LLaMA 等现代模型基本都用了 RMSNorm它是 LayerNorm 的简化版省去了均值中心化只做缩放计算开销更低。4. 自注意力机制与 Q、K、V 详解自注意力是整个 Transformer 最核心的公式。理解它其他部分都是围绕它服务的工程化设计。4.1 从数据库查询讲起Q、K、V 经常被类比成搜索引擎Query 是你输入的关键词Key 是文档的标题标签Value 是文档正文内容。Transformer 中的每个 token 都会生成三个向量QQuery表示“我想找什么”KKey表示“我是什么内容”VValue表示“我实际携带的信息”当前 token 的 Q 与所有 token 的 K 做点积得到注意力分数分数越高说明这个 token 与当前 token 的相关性越强分数经过 Softmax 变成权重再对所有 V 做加权求和得到当前 token 的新表示。公式可以写成Attention(Q, K, V) softmax( QK^T / sqrt(d_k) ) V其中d_k是 Key 向量的维度。除以sqrt(d_k)是为了防止点积结果过大导致 Softmax 进入饱和区梯度变小。4.2 为什么用点积而不是其他相似度计算点积运算在 GPU 上有高度优化的矩阵乘法支持计算效率比拼接向量再经过 MLP 高得多。而且点积在维度较高时数值波动会增大所以论文作者特意做了缩放处理。实践中也有人尝试过加性注意力效果相近但速度更慢所以缩放点积注意力成了事实标准。4.3 单头注意力最小实现下面用 PyTorch 写一个单头自注意力模块重点看张量形状变化import torch import torch.nn as nn import math class SelfAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x, maskNone): batch, seq_len, _ x.shape # 将 d_model 拆成 n_heads 个 d_k q self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) k self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) v self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) # 注意力分数 scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) out torch.matmul(attn, v) # 合并多头 out out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model) return self.out_proj(out)输入 x 的形状是(batch, seq_len, d_model)经过 Q、K、V 三个线性层后仍然保持这个形状。因为模型是把d_model拆成n_heads个d_k所以多个头的计算可以在同一个矩阵乘法里完成不需要显式循环。用这种方式计算整段序列任意两个位置之间的注意力分数都会同时被算出来这也是 Transformer 并行能力的来源。5. 多头注意力让模型看到不同子空间单头注意力相当于只从一个角度判断“哪些 token 相关”表达力有限。多头注意力把 Q、K、V 分别映射到多个子空间每个头可以关注不同关系有的头关注语法依赖有的头关注指代关系有的头关注位置距离。多头注意力的流程是将输入经过线性层映射成 Q、K、V。拆成n_heads组每组维度为d_k。每组独立做注意力计算。把多组结果拼接回d_model维度。经过一个输出线性层。常见配置是n_heads 8或n_heads 16。头数不是越大越好头数增加会带来计算开销而且如果d_model不变每个头的d_k会变小过于碎片化的信息反而可能损害效果。实际调参时可以从 8 头起步根据显存和数据规模调整。6. 位置编码给 Transformer 补上顺序信息自注意力机制本身是“位置无关”的。把句子里的词顺序打乱如果不加位置信息Transformer 计算出的注意力结果会完全一样。因为 Q、K、V 都是对同一个 token 做线性变换不含先后顺序信息。原始论文使用正弦函数位置编码import torch import math def sinusoidal_positional_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float32).unsqueeze(1) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float32) * (-math.log(10000.0) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe这个编码的特点是每个位置都有一个唯一的向量而且位置pos k的编码可以表示为位置pos编码的线性函数。模型有机会学到相对位置关系而不是只记住绝对位置。6.1 现代模型的位置编码演变BERT 使用可学习位置编码让模型从数据里自己学位置向量。GPT 系列用的是可学习绝对位置编码。但绝对位置编码对长文本外推能力差一旦输入长度超过训练时的最大长度模型会不适应。所以现在很多模型改用旋转位置编码RoPE它把位置信息直接融进 Q、K 向量里相对位置关系计算更自然外推能力也更强。LLaMA、ChatGLM 等模型都在用 RoPE 的变体。如果你不理解位置编码看源码时会觉得“每个 token 不是直接输入的吗”然后困惑为什么同样一批 token 放在不同位置会产生不同结果。答案就在位置编码这一步。7. token 视角Transformer 如何读文本大模型处理文本的第一步不是直接把句子塞进网络而是把句子切分成 token。7.1 什么是 tokentoken 是模型能处理的最小单位。中文可能一个字就是一个 token也可能一个词被拆成多个 token。英文常见的是子词切分比如 “unhappy” 可能被拆成 “un” 和 “happy”。BERT 的 WordPiece、GPT 的 BPE 都做这类事情。切分完成后每个 token 在词表中对应一个整数 ID。模型拿到的是一个整数数组比如[101, 234, 567, 102]。这个数组接着通过 Embedding 矩阵映射成稠密向量再进入 Transformer 主网络。7.2 上下文长度和注意力成本上下文长度context length就是模型一次能读的最大 token 数。为什么大模型都宣传上下文窗口因为 Transformer 自注意力的计算复杂度是O(n^2)序列长度翻倍计算量翻四倍。所以上下文从 2048 涨到 4096显存消耗远不止翻倍。用户看到的是“能读更多内容”研发看到的是“训练和推理成本急剧上升”。7.3 双向注意力和因果注意力BERT 这类 Encoder 模型做语言理解时可以看完整句话属于双向注意力。GPT 这类 Decoder 模型做生成时只能看左边的 token不能看到右边未来的 token属于因果注意力。这正是 GPT 可以做生成的本质原因——它在每个位置只预测下一个 token。8. Transformer 的训练与推理流程理解训练流程能帮你搞清楚为什么大模型训练费钱费 GPU为什么需要那么大的算力。8.1 预训练GPT 的预训练任务很简单给你一段文本遮住或者让你预测下一个 token模型通过大量文本学习语言规律。这个阶段消耗的算力巨大因为要在几十亿 token 的语料上反复迭代。大规模训练为什么需要大量 GPU就是因为 Transformer 的并行和参数量数据一多单卡根本不可能在合理时间内完成训练。8.2 损失函数预训练通常使用交叉熵损失只计算被预测 token 位置的损失。模型输出的是一个词表大小的概率分布目标是让正确答案对应的概率尽量高。优化器一般是 AdamW配合学习率预热warmup和余弦衰减。8.3 推理时自回归生成GPT 推理时一次生成一个 token生成完拼到输入后面再重新喂给模型预测下一个 token。这就是“自回归生成”。简化示例prompt 人工智能的核心是 input_ids tokenizer.encode(prompt) for _ in range(30): outputs model(input_ids) next_token_id sample_from_logits(outputs[-1]) input_ids.append(next_token_id)这个循环看起来简单实际大模型推理会做 KV Cache缓存历史 token 的 K、V 向量避免每生成一个 token 就重算整段上下文。所以现代推理引擎的关键优化都在 KV Cache 上显存占用、访存带宽、缓存淘汰策略。9. Transformer 从 NLP 走向视觉ViT 与 Swin TransformerTransformer 最早是为文本设计的但图像也可以变成“序列”。9.1 ViT 的基本思路ViTVision Transformer把一张图切成一堆固定大小的 patch比如 224x224 的图像分成 14x14 个 16x16 的 patch。每个 patch 展平后经过线性映射得到一个向量再加上位置编码然后送入标准 Transformer Encoder。这个思路打破了 CNN 拿卷积核扫图像的惯例让图像分类、目标检测任务也能享受 Transformer 的全局建模能力。但 ViT 也有代价patch 之间的全局自注意力计算量大小数据集上不如 CNN 好训练需要大规模预训练才能发挥优势。9.2 Swin Transformer 做了什么改进Swin Transformer 把注意力限制在局部窗口内例如 7x7 的窗口窗口内部做自注意力窗口之间通过移位窗口建立跨窗口连接。这样既保留了 Transformer 的表达能力又把计算复杂度从全局平方级降为局部可控。这也是视觉 Transformer 落地时非常实际的选择全局注意力建模能力强但显存开销大局部注意力虽然视野受限但可以通过多层堆叠和窗口移位弥补。10. Transformer 与 CNN、RNN 全面对比维度RNN / LSTMCNNTransformer并行性差必须顺序计算好很好长程依赖弱容易遗忘需要堆深度强任意位置直接建模位置信息天然由时间步体现卷积核隐含局部位置需要显式位置编码计算复杂度序列长度线性卷积核大小关联序列长度平方图像任务不常用强归纳偏置需要大规模数据文本建模曾是主流早期探索过当前事实标准可扩展性差中等极强这解释了为什么最后胜出的是 Transformer它不是某一项指标碾压而是在并行性、长程建模、规模化三方面同时胜出。11. 手撕一个最小 Transformer Block前面拆解了所有组件现在把它们拼成一个可运行的最小 Transformer Block。这个实现省略了完整 Encoder 的堆叠和复杂训练代码专注展示核心结构import torch import torch.nn as nn class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.attention SelfAttention(d_model, n_heads) self.norm1 nn.LayerNorm(d_model) self.ff nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # Pre-LN 结构先归一化再做注意力 x x self.dropout(self.attention(self.norm1(x), mask)) x x self.dropout(self.ff(self.norm2(x))) return x # 使用示例 d_model 128 seq_len 32 batch 2 x torch.randn(batch, seq_len, d_model) block TransformerBlock(d_modeld_model, n_heads8, d_ff512) out block(x) print(out.shape) # torch.Size([2, 32, 128])注意这里用的是 Pre-LN 结构先 LayerNorm再进入子层最后残差连接。原始 Transformer 论文是 Post-LN但现代大模型更偏爱 Pre-LN因为它训练更稳定不用对学习率做太多精细调整。想验证注意力是否真的学到了东西可以打印 attention 权重# 在 SelfAttention 的 forward 里临时返回 attn # attn.shape: (batch, n_heads, seq_len, seq_len)你会看到模型确实能学会“当前 token 关注句子中哪些其他 token”。比如在代词位置注意力可能集中在它所指代的名词上。12. 常见理解误区与排查思路误区实际情况Transformer 注意力机制注意力只是组件完整架构还包括 FFN、LayerNorm、残差连接注意力矩阵越大一定越好计算量平方增长显存极易打满头数越多效果越好头数过多会碎片化需要实验验证位置编码可以随便去掉去掉后模型完全失去顺序信息Transformer 不需要 CNN视觉任务中两者可结合很多模型仍用卷积做降采样上下文越长越好长上下文的显存和推理开销很高要平衡如果自己写代码训练 Transformer 时遇到问题按这个顺序排查先确认输入张量形状是否为(batch, seq_len, d_model)。确认d_model % n_heads 0否则无法拆分多头。确认注意力分数是否做过缩放不缩放可能导致训练初期不稳定。确认是否加了 LayerNorm 和残差连接去掉后深层网络容易不收敛。确认因果掩码是否实现正确生成任务中模型预测当前位置时不能看到未来 token。显存不足时优先减小 batch size、seq_len 和 n_heads不要一上来就调大模型。13. 给初学者的下一步学习建议如果你刚接触 Transformer建议按这个路径走不要一开始就硬啃大模型源码。第一步把本文的自注意力代码和 TransformerBlock 跑通打印中间张量的形状。确认每个组件输入输出维度一致。第二步找一个小数据集比如简单的文本分类任务用 PyTorch 搭一个单层 Transformer Encoder跑一轮训练。看到 loss 下降说明你对整个流程的组装是正确的。第三步用 Hugging Face 的 transformers 库跑现成模型接口。下面是通用调用示例实际使用时把模型路径换成你本地可访问的模型from transformers import AutoTokenizer, AutoModelForCausalLM # 实际使用时需要指定本地或远程可访问的模型名称 model_path your-model-path tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForCausalLM.from_pretrained(model_path) prompt Transformer 的核心是 inputs tokenizer(prompt, return_tensorspt) outputs model.generate(**inputs, max_new_tokens50) print(tokenizer.decode(outputs[0], skip_special_tokensTrue))第四步读一遍 PyTorch 官方实现的 Transformer 源码对照本文的结构图逐行看。这时候你会发现自己已经能看懂大部分代码了。第五步针对视觉方向读 ViT 的 patch embedding 实现针对大模型推理读 KV Cache 相关的量化文章和推理引擎代码。14. 总结Transformer 能在过去几年里成为 AI 领域的事实标准不是因为某个单一技巧惊艳而是它把并行训练、长程建模、规模化扩展三件事同时做对了。理解它不需要背论文里的每个公式关键是要抓住自注意力、多头、位置编码、残差和 LayerNorm 这五条主线。把这五条主线串起来你再看 GPT、LLaMA、ViT 甚至多模态模型会发现它们都是同一个骨架在不同数据形态上的变体。这篇文章最值得收藏的部分是第 4 节的单头注意力实现和第 11 节的最小 Transformer Block。建议你先把它们跑通再逐步加训练逻辑。最容易踩的坑集中在张量形状、缩放因子和掩码三处遇到不收敛或显存溢出优先检查这三个点。下一步可以继续沿着两条线深入一是研究 KV Cache 和量化推理理解大模型为什么能跑得快二是读 ViT 和 Swin Transformer 的光源码看视觉任务是怎么复用注意力机制的。Transformer 的图解到这里结束希望你能从“见过概念图”变成“能写出最小实现”。
返回列表