Transformer核心原理:从Token到Attention的原子级拆解

发布时间:2026/6/22 5:18:56
Transformer核心原理:从Token到Attention的原子级拆解 1. 这个标题不是在“降维打击”而是在拆解一个被过度神化的黑箱“Transformer 其实很简单”——看到这个标题你第一反应可能是怀疑甚至带点嘲讽一个让整个AI行业翻天覆地、催生了ChatGPT、Gemini、Claude、Sora的架构能“简单”它背后是上百页的论文、动辄千亿参数、需要千张A100训练的庞然大物怎么就简单了别急。这里的“简单”不是指“小学生都能手推反向传播”而是指它的核心思想、主干逻辑和关键模块完全可以用一套清晰、可具象、可图示、可手动演算的流程讲清楚。它不像早期RNN那样依赖难以捉摸的“隐藏状态演化”也不像CNN那样需要靠大量实验调参才能理解感受野如何叠加。Transformer 的“简单”是一种结构上的简洁性、计算上的并行性、以及原理上的可解释性。我做AI工程落地十年从2015年用LSTM跑金融时序预测到2018年第一批用BERT做风控文本分类再到2022年亲手把ViT部署进工业质检产线踩过所有坑也亲手把Transformer从论文里“抠”出来变成能跑在边缘设备上的代码。我敢说90%的工程师不是学不会Transformer而是被铺天盖地的术语、堆叠的公式、混乱的图示和“必须先懂矩阵论”的恐吓给劝退了。那些“Attention is All You Need”里的矩阵乘法本质上就是三步查表、打分、加权求和——和你在Excel里用VLOOKUPSUMPRODUCT完成一次动态报表逻辑内核一模一样。这个标题的价值就在于它直指要害我们不需要先成为数学家才能用好一个工具就像你不需要懂电磁波原理也能熟练操作手机。Transformer 的“难”主要在工程规模显存、通信、调度而不是概念本身。它的“简单”体现在四个刚性骨架上Token是原子、Position是坐标、Attention是关系网、FFN是放大器。这四块拼图每一块都足够直观组合起来却产生了涌现能力。本文不讲“为什么Transformer改变了世界”只讲“它到底在干什么、每一步在算什么、为什么这么算”。全文没有一行代码但你读完能自己画出前向传播的完整流程图能解释清楚为什么“狗咬人”和“人咬狗”在模型里是两个完全不同的向量序列能看懂《The Annotated Transformer》里那个著名的矩阵形状转换图到底在转什么。适合谁读如果你是刚接触NLP的算法新人别被“多头”“层归一化”“残差连接”吓住本文会带你从零搭起第一块积木如果你是转行过来的后端/前端工程师想快速理解大模型底层逻辑本文用你熟悉的“数据库JOIN”“缓存Key-Value”来类比如果你是已经调过BERT微调但总卡在loss不降的中级同学本文会指出你忽略的那个最关键的“位置编码注入时机”甚至如果你是硬件工程师正为KV Cache的内存布局发愁本文第三节的实操细节会直接告诉你为什么FlashAttention要重排QKV的内存顺序。它不承诺让你一夜之间写出LLaMA但它保证读完之后你再看到任何一篇Transformer相关论文或技术文档不会再有“每个字都认识连起来不知道在说什么”的窒息感。2. 内容整体设计与思路拆解为什么“简单”必须从“原子操作”开始讲起要真正讲清“Transformer其实很简单”绝不能一上来就甩出那张经典的Encoder-Decoder框图然后说“看这就是全部”。那不是讲解是供奉。真正的“简单”必须回归到最原始的计算单元像拆解一台机械钟表一样从游丝、齿轮、擒纵叉开始一层层还原它的计时逻辑。我的整体设计思路就是严格遵循信息流的物理路径以“一个token的生命周期”为叙事主线拒绝任何跳跃式概括。2.1 为什么必须放弃“架构图先行”的惯性思维几乎所有入门教程包括那篇划时代的《The Illustrated Transformer》都是从宏观架构图切入左边Encoder一堆方块右边Decoder一堆方块中间箭头飞来飞去。这种讲法对建立整体印象有帮助但对理解“为什么这样设计”是灾难性的。它掩盖了三个致命问题时间错位图中Encoder和Decoder看似并行但实际计算是严格串行的——Embedding层输出必须等Positional Encoding加完才能进AttentionAttention输出必须等LayerNorm和残差加完才能进FFN。图没体现这个“流水线节拍”。空间混淆图中一个“Multi-Head Attention”方块背后是12个或更多完全独立的子网络每个子网络内部又有Q/K/V三套权重矩阵。读者看到一个方块潜意识认为“这是一个操作”实际上它是“12个并行操作的集合”而每个操作又包含3次矩阵乘1次Softmax1次加权求和。这种“一个方块多个原子操作”的抽象是理解的第一道高墙。数据失真图中箭头标注“Sequence of Vectors”但没说明这个Sequence的维度是多少、每个Vector长什么样、它在GPU显存里是怎么排布的。而恰恰是这些“枯燥”的shape信息比如[batch, seq_len, d_model]决定了你能否写出正确的PyTorch代码决定了你的显存会不会爆。所以我的设计反其道而行之不画任何宏观框图只画一张“单token单步计算”的微观流程图。从一个原始字符“a”开始它如何被切分成subword token如何查embedding表得到768维向量如何与sin/cos位置向量相加如何被12个不同的W^Q矩阵分别乘出12个query向量……每一步都标注清楚输入shape、输出shape、运算类型矩阵乘/加法/Softmax、以及这个运算在GPU上实际耗时占比。这张图就是你调试模型时torch.profiler输出里每一行的真实映射。2.2 为什么“Token”是唯一可信的起点很多教程从“文本预处理”讲起罗列BPE、WordPiece、SentencePiece的区别。这很重要但不是“简单”的入口。真正的入口是承认“Token”是Transformer世界的唯一原生单位。它不关心你是中文、英文还是代码不关心你是一个字、一个词还是一段emoji。它只认一个整数ID。这个ID就是它世界的“原子序数”。Tokenization不是预处理而是世界观设定BPE算法生成的vocabulary.txt本质上就是Transformer的“元素周期表”。[CLS]是氢[SEP]是氦the是碳dog是铁……模型的所有知识都建立在这个离散符号系统之上。它没有“语义”只有“ID关联”。当你看到model(input_ids)这个APIinput_ids就是一个纯整数数组比如[101, 2023, 2003, 102]模型根本不“知道”这对应“[CLS] I am [SEP]”它只知道“查表取第101行、第2023行……的向量”。Embedding Layer是唯一的“翻译官”它的工作极其单纯把整数ID翻译成稠密向量。这个翻译表weight matrix的大小是[vocab_size, d_model]比如[30522, 768]。查表操作就是一次简单的索引indexing等价于embedding_weight[input_id]。没有任何魔法就是数组寻址。这也是为什么Embedding层通常是模型里参数最多的一层——它要为字典里每一个词都分配一个“身份向量”。Positional Encoding是“时空坐标系”Transformer没有RNN的时序记忆所以必须给每个token打上“我在第几个位置”的标签。Sinusoidal编码的公式看起来复杂但它的物理意义极其朴素为每个位置生成一个独一无二的、可区分的、且能表达相对距离的向量。pos0的向量和pos100的向量它们的点积应该很小pos5和pos6的向量点积应该很大。这个“相对距离可计算”的特性是后续Attention能捕捉“动词-宾语”远距离依赖的数学基础。它不是玄学就是一个精心设计的坐标函数。2.3 为什么“Attention”必须被还原为“数据库查询”这是全篇最关键的认知跃迁。我把Scaled Dot-Product Attention的公式softmax(QK^T / sqrt(d_k)) * V强行翻译成一个SQL查询-- 假设我们有一个“token关系表” CREATE TABLE token_relations ( query_token_id INT, key_token_id INT, attention_score FLOAT, value_vector BLOB ); -- 对于当前tokenquery_id 123我想知道它该关注哪些其他token SELECT value_vector FROM token_relations WHERE query_token_id 123 ORDER BY attention_score DESC LIMIT 5;Q就是你的查询条件WHERE query_token_id 123K就是表里所有可能的key_token_id即所有其他token的位置QK^T就是计算“当前token与每个其他token的匹配度”类似JOIN ON condition/ sqrt(d_k)是一个缩放因子防止点积过大导致Softmax梯度消失就像SQL里给score加个归一化softmax(...)就是把匹配度转成0~1之间的概率权重ORDER BY ... DESCV就是你要最终取出的value_vectorSELECT value_vector。这个类比不是为了简化而简化它精准抓住了Attention的本质它不是一个“学习注意力”的过程而是一个“根据当前内容动态检索最相关信息”的过程。模型并不“决定”要关注什么它只是用Query向量在Key向量构成的“索引库”里高效地找到最匹配的Value向量。这和你在ElasticSearch里用向量相似度搜索图片逻辑完全一致。所谓“自注意力”就是让每个token既当查询者Query又当被查询的索引项Key和Value。这种“全员皆可查、全员皆可被查”的平等结构正是它能并行计算、摆脱RNN时序枷锁的根本原因。2.4 为什么“FFN”是被严重低估的“特征放大器”很多人把Feed-Forward Network看作Attention的“附属品”一个简单的两层MLP。这是巨大的误解。FFN才是Transformer里真正负责“深度非线性变换”和“特征空间重塑”的核心引擎。它不是“加个非线性”那么简单标准FFN结构是Linear - GELU - Linear。第一个Linear层通常d_model - 4*d_model是特征升维把768维向量映射到3072维的高维空间GELU激活函数在高维空间里进行复杂的非线性切割第二个Linear层3072 - 768是特征降维与重组把高维空间里学到的复杂模式压缩回原始维度但此时的768维已经蕴含了远超原始输入的语义组合能力。它和Attention是“分工协作”的黄金搭档Attention负责“全局关系建模”——告诉我“这句话里‘银行’和‘抢劫’的关系最密切”FFN负责“局部特征深化”——基于这个关系生成一个全新的、融合了“金融”、“犯罪”、“紧急”等多重语义的强化向量。你可以把Attention看作一个“关系路由器”把FFN看作一个“特征加工厂”。没有FFNAttention输出的向量会非常“扁平”缺乏深度语义没有AttentionFFN就只能看到孤立的token无法建模长程依赖。它的参数量占比惊人在一个标准Transformer层中Attention的参数量约为3 * d_model * d_k d_model * d_vQ/K/V投影输出投影而FFN的参数量是d_model * 4*d_model 4*d_model * d_model 8 * d_model^2。以d_model768计算FFN参数量约4.7MAttention约1.8M。FFN贡献了层内超过70%的可学习参数。说Transformer的“智能”主要来自FFN并不为过。这个设计思路确保了“简单”不是肤浅的简化而是通过回归本质、剥离幻觉、建立精准类比让每一个模块都变得可触摸、可推理、可验证。它不回避复杂性而是把复杂性分解为一系列清晰、独立、可验证的原子步骤。3. 核心细节解析与实操要点从“纸上谈兵”到“动手拆解”光有思路还不够真正的“简单”必须落实到每一个可触摸、可验证、可调试的细节上。下面我将带着你像一个硬件工程师拆解电路板一样逐层剖析Transformer最核心的四个环节。所有参数、shape、计算过程都基于真实开源模型如BERT-base的配置绝不虚构。3.1 Tokenization与Embedding不是魔法是查表与拼接我们以句子“I love transformers”为例走一遍最底层的数据旅程。Step 1: Tokenization (BPE)BERT使用WordPiece但原理相通。假设我们的vocabulary很小只有[CLS], [SEP], I, love, ##trans, ##form, ##ers, transformers注意##表示子词subword。transformers被切分为##trans,##form,##ers。最终token IDs序列是[101, 1332, 2764, 2102, 2103, 2104, 102]101[CLS], 1332I, 2764love, 2102##trans, 2103##form, 2104##ers, 102[SEP]提示len(token_ids) 7这就是后续所有计算的seq_len。这个数字决定了Attention矩阵的大小7x7也决定了显存占用的基线。任何padding如pad到128都会无谓增加计算量。Step 2: Embedding LookupEmbedding weight matrix shape是[vocab_size, d_model] [30522, 768]。对于每个token ID我们执行# 伪代码实际是GPU上的向量化操作 token_vector embedding_weight[token_id] # shape: [768]对7个token我们得到7个768维向量堆叠成一个tensortoken_embeddings torch.stack([vec1, vec2, ..., vec7])shape变为[7, 768]。这就是Embedding层的输出。Step 3: Positional Encoding AdditionPositional encoding不是单独一层而是直接加到token embedding上。它的shape必须和token_embeddings完全一致[7, 768]。Sinusoidal编码的计算如下简化版# 对于位置pos (0 to 6) 和维度i (0 to 767) # PE(pos, 2i) sin(pos / 10000^(2i/d_model)) # PE(pos, 2i1) cos(pos / 10000^(2i/d_model)) pos_encoding torch.zeros(7, 768) for pos in range(7): for i in range(0, 768, 2): div_term 10000 ** (i / 768) pos_encoding[pos, i] math.sin(pos / div_term) if i 1 768: pos_encoding[pos, i1] math.cos(pos / div_term)然后final_input token_embeddings pos_encoding。注意这里是element-wise addition不是concatenate。这是关键很多初学者误以为Positional Encoding是另一个通道其实是“叠加”在原有向量上给它注入位置信息。注意Positional Encoding是固定值不可学习在原始Transformer中。BERT后来改用“learned positional embeddings”即一个可训练的[max_position, d_model]lookup table但逻辑相同加到embedding上。无论哪种它都必须和embedding同shape否则无法相加。3.2 Scaled Dot-Product Attention一场精确的“向量匹配”游戏现在我们有了[7, 768]的输入。Attention层的第一步是把它线性投影成Q、K、V三个矩阵。Step 1: Linear Projections每个projection都是一个Linear(d_model, d_k)层。在BERT-base中d_model768,d_kd_v64head数12。所以W^Qshape:[768, 64]W^Kshape:[768, 64]W^Vshape:[768, 64]计算Q input W^Q # [7, 768] [768, 64] [7, 64] K input W^K # [7, 768] [768, 64] [7, 64] V input W^V # [7, 768] [768, 64] [7, 64]注意这里input是[7, 768]但Q/K/V的shape是[7, 64]不是[7, 768]。这是因为每个head只关注64维的子空间。这是“多头”的物理基础把768维大向量拆成12个64维小向量分别处理。Step 2: Scaled Dot-Product Softmax计算注意力分数矩阵Attention Score Matrix# Q K.T 得到 [7, 7] 矩阵每个元素 score[i,j] Q_i · K_j scores Q K.T # [7, 64] [64, 7] [7, 7] # 缩放除以 sqrt(d_k) sqrt(64) 8 scores scores / 8.0 # 防止点积过大Softmax梯度消失 # Softmax对每一行即每个query做归一化 attention_weights torch.softmax(scores, dim-1) # shape [7, 7]attention_weights[i, j]就是“第i个token对第j个token的关注程度”所有j的权重和为1。例如attention_weights[2, 5]可能很高意味着“love”这个词特别关注“##ers”这个子词。Step 3: Weighted Sum of Values最后一步用权重去加权求和V# attention_weights [7, 7] V [7, 64] [7, 64] output attention_weights V # shape [7, 64]这个output就是单个head的输出。它是一个[7, 64]的矩阵每一行代表一个token经过“关注他人”后得到的新向量。注意Q K.T是[7, 7]而attention_weights V是[7, 64]。这个shape变化就是“从关系矩阵7x7到特征向量7x64”的转换。它把全局关系压缩回了每个token的个体表示。3.3 Multi-Head Attention12个“平行宇宙”的协同单个head的输出是[7, 64]但我们需要把它变回[7, 768]以匹配下一层的输入。这就是Multi-Head的精髓并行运行12个独立的Attention然后把结果拼接起来。# 假设我们有12个head每个head输出 [7, 64] head_outputs [] # list of 12 tensors, each [7, 64] for head_i in range(12): # ... 计算第i个head的 Q_i, K_i, V_i ... head_output attention_function(Q_i, K_i, V_i) # [7, 64] head_outputs.append(head_output) # 拼接在最后一个维度dim-1拼接 # [7, 64] x 12 - [7, 12*64] [7, 768] concatenated torch.cat(head_outputs, dim-1) # [7, 768] # 最后用一个线性层 W^O 投影回 [7, 768] W_O torch.nn.Linear(768, 768) # shape [768, 768] multihead_output W_O(concatenated) # [7, 768]这个W^O层至关重要。它不是可有可无的而是学习如何将12个不同视角heads的信息最优地融合成一个统一表示。没有它12个head的输出就是12个独立的向量无法形成合力。实操心得为什么是12个head这不是玄学。768 / 64 12这是一个完美的整除。它确保了计算效率最大化。如果你强行设为13个headd_k就得是768/13≈59.07必须向上取整到60那么13*60780768你就得padding或者丢弃维度徒增计算开销。所以head数是d_model的因数这是工程上的硬约束。3.4 Feed-Forward Network一场“升维-非线性-降维”的炼金术Multi-Head Attention的输出[7, 768]进入FFN。它的结构是Linear_1 (768 - 3072) - GELU - Linear_2 (3072 - 768)Step 1: First Linear Projection (升维)# W1 shape: [768, 3072], b1 shape: [3072] hidden input W1 b1 # [7, 768] [768, 3072] [7, 3072]这一步把每个token的768维向量“展开”成3072维的高维空间。为什么要升维因为高维空间提供了更丰富的“表达自由度”让GELU激活函数能刻画更复杂的非线性边界。想象一下你在二维平面上画一条曲线很难但在三维空间里你可以用一个曲面轻松拟合。Step 2: GELU Activation (非线性)GELUGaussian Error Linear Unit的公式是GELU(x) x * Φ(x)其中Φ(x)是标准正态分布的累积分布函数。它的效果比ReLU更平滑能保留一部分负值信息对训练稳定性有帮助。在PyTorch中就是torch.nn.GELU()。Step 3: Second Linear Projection (降维与重组)# W2 shape: [3072, 768], b2 shape: [768] output hidden W2 b2 # [7, 3072] [3072, 768] [7, 768]这一步把3072维的“丰富信息”重新压缩回768维。但此时的768维已经不再是原始输入而是经过高维空间“淬炼”后的、富含语义组合的新向量。它可能包含了“I love”和“transformers”之间的深层关联这种关联在原始embedding里是不存在的。注意FFN的两个Linear层权重矩阵W1和W2是完全独立、互不共享的。这和Embedding-Unembedding的权重绑定weight tying完全不同。FFN的参数是模型里最“奢侈”的部分也是最容易过拟合的地方所以Dropout通常加在这里。4. 实操过程与核心环节实现手把手复现一个“极简Transformer Block”理论讲完现在我们用最精简的PyTorch代码实现一个功能完整的Transformer Encoder Block。目标不依赖Hugging Face不调用任何高级API只用torch.nn.Linear、torch.nn.functional和基础张量操作跑通前向传播。这会让你彻底看清每一行代码在做什么。4.1 完整代码一个可运行的Blockimport torch import torch.nn as nn import torch.nn.functional as F import math class SimpleTransformerBlock(nn.Module): def __init__(self, d_model768, n_head12, dropout0.1): super().__init__() self.d_model d_model self.n_head n_head self.d_k d_model // n_head # 64 self.d_v self.d_k # Attention weights: Q, K, V for each head # Well use a single big matrix for all heads, then split self.W_q nn.Linear(d_model, d_model) # [768, 768] self.W_k nn.Linear(d_model, d_model) # [768, 768] self.W_v nn.Linear(d_model, d_model) # [768, 768] self.W_o nn.Linear(d_model, d_model) # [768, 768] # FFN layers self.ffn1 nn.Linear(d_model, 4 * d_model) # [768, 3072] self.ffn2 nn.Linear(4 * d_model, d_model) # [3072, 768] # Normalization and Dropout self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x): x: [batch, seq_len, d_model] e.g., [1, 7, 768] # --- Step 1: Self-Attention Sublayer --- # 1.1 Normalize input norm_x self.norm1(x) # [1, 7, 768] # 1.2 Project to Q, K, V Q self.W_q(norm_x) # [1, 7, 768] K self.W_k(norm_x) # [1, 7, 768] V self.W_v(norm_x) # [1, 7, 768] # 1.3 Reshape for multi-head: [batch, seq_len, n_head, d_k] - [batch, n_head, seq_len, d_k] batch_size, seq_len, _ Q.shape Q Q.view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) # [1, 12, 7, 64] K K.view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) # [1, 12, 7, 64] V V.view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) # [1, 12, 7, 64] # 1.4 Scaled Dot-Product Attention # scores Q K.T / sqrt(d_k) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # [1, 12, 7, 7] # Optional: add mask here for causal or padding # scores scores.masked_fill(mask 0, float(-inf)) # Softmax over last dimension (key positions) attn_weights F.softmax(scores, dim-1) # [1, 12, 7, 7] attn_weights self.dropout(attn_weights) # Weighted sum of V context torch.matmul(attn_weights, V) # [1, 12, 7, 64] # 1.5 Concatenate heads and project back # [1, 12, 7, 64] - [1, 7, 12*64] [1, 7, 768] context context.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) attn_output self.W_o(context) # [1, 7, 768] # 1.6 Residual connection and dropout x x self.dropout(attn_output) # [1, 7, 768] # --- Step 2: FFN Sublayer --- # 2.1 Normalize norm_x self.norm2(x) # [1, 7, 768] # 2.2 FFN: Linear - GELU - Linear - Dropout ffn_output self.ffn2(F.gelu(self.ffn1(norm_x))) # [1, 7, 768] ffn_output self.dropout(ffn_output) # 2.3 Residual connection output x ffn_output # [1, 7, 768] return output # --- 测试代码 --- if __name__ __main__: # 创建一个极简输入batch1, seq_len7, d_model768 # 用随机数模拟实际中是Embedding输出 x torch.randn(1, 7, 768) block SimpleTransformerBlock(d_model768, n_head12) output block(x) print(fInput shape: {x.shape}) print(fOutput shape: {output.shape}) print(Block forward pass successful!)4.2 关键步骤详解与参数验证这段代码每一行都对应着前面理论的精确实现。我们来逐行验证其物理意义self.W_q nn.Linear(d_model, d_model)这是“12个head共用一个大矩阵”的工程实现。它等价于12个[768, 64]的小矩阵拼在一起。Q.view(...).transpose(1,2)这行就是把[1, 7, 768]的输出按[12, 64]的块重新排列成[1, 12, 7, 64]为并行计算做准备。这是GPU高效计算的关键技巧。torch.matmul(Q, K.transpose(-2, -1))K.transpose(-2, -1)把[1, 12, 7, 64]变成[1, 12, 64, 7]这样Q K.T才能得到[1, 12, 7, 7]的注意力分数矩阵。-2和-1是PyTorch中指定最后两个维度的写法确保代码在任意batch size下都鲁棒。context.transpose(1, 2).contiguous().view(...)这是Multi-Head的“拼接”操作。transpose(1,2)把[1, 12, 7, 64]变成[1, 7, 12, 64]contiguous()确保内存连续否则view会报错view(batch_size, seq_len, -1)把最后两个维度[12, 64]压平成[768]。-1是PyTorch的自动推导非常实用。F.gelu(self.ffn1(norm_x))这里明确调用了F.gelu而不是nn.GELU()因为前者是函数式接口后者是模块。在forward里我们倾向于用函数式避免创建不必要的对象。x x self.dropout(attn_output)残差连接Residual Connection是Transformer稳定训练的基石。它让梯度可以直接绕过整个Attention子层避免深层网络的梯度消失。self.dropout加在残差之前是标准做法Post-Dropout。4.3 运行结果与Shape追踪运行上述代码你会看到Input shape: torch.Size([1, 7, 768]) Output shape: torch.Size([1, 7, 768]) Block forward pass successful!这个[1, 7, 768]的shape贯穿始终是Transformer的“脊柱”。它证明了输入和输出维度一致保证了Block可以无限堆叠所有中间计算Q/K/V、attention scores、FFN hidden都服务于这个最终shape的维护没有任何“维度爆炸”或“信息丢失”一切都在可控的线性代数框架内。实操心得如果你想亲眼看到Attention权重只需在attn_weights F.softmax(scores, dim-1)后面加一句print(attn_weights[0, 0])