拆解Transformer本源：350行源码吃透Attention底层原理

文章目录前言一、Scaled Dot-Product AttentionAI界的查户口大师二、Multi-Head Attention一个人同时开八个脑洞三、Position-wise FFN每个token的健身房私教课四、Positional Encoding给token发座位号五、Layer Normalization给数据穿统一制服六、Encoder Layer自注意力FFN的组合拳七、Decoder Layer戴眼罩的传话游戏八、完整Transformer组装高达的时刻九、写在最后350行代码八年AI霸权P.S. 无意间发现了一个巨牛的人工智能教程非常通俗易懂对AI感兴趣的朋友强烈推荐去看看传送门https://blog.csdn.net/HHX_01前言2017年Google那帮大佬甩出一篇论文标题叫《Attention Is All You Need》。翻译成人话就是“Attention就够了别的都是弟弟。“我当时一看好家伙这口气比我家楼下烧烤摊老板还大。老板至少还谦虚地说我家羊肉串全市第二”Google直接说我只需要注意力”。结果八年过去了GPT、BERT、LLaMA全是从这玩意儿肚子里爬出来的。现在大模型卷得跟春运抢票似的但你敢信吗这祖宗的源码纯PyTorch写出来就350行。350行我上次写个登录页面都不止350行。Google这帮人是真狠用个博客文章的长度把整个AI行业的地基给打好了。今天我就当一回源码拆弹专家把这350行代码一行一行掰开揉碎。放心不催眠不念经全程脱口秀节奏。你要是看完还犯困那我……那我下次换个更吵的BGM。一、Scaled Dot-Product AttentionAI界的查户口大师Transformer的核心就一句话Query问KeyKey回答Value。听起来像不像相亲Query就是男方问Key你有房吗有车吗存款几位数Key一一回答然后Value就是女方实际的嫁妆——哦不是实际的语义信息。公式长这样Attention(Q,K,V) softmax(QK^T / √d_k) V。别跑这玩意儿翻译成中文就是先把Query和Key的点积算出来除以一个√d_k再softmax一下最后跟Value乘一块儿。简单吧就像你问相亲对象三个问题打分归一化最后决定要不要继续聊。那为什么非要除以√d_k呢因为维度一高点积的数值容易膨胀softmax直接社死——梯度消失得比我的头发还快。除以√d_k就相当于给数值减肥保持身材匀称训练才不会崩盘。这操作跟我过年狂吃后上秤前先脱鞋脱外套一个逻辑。核心代码classScaledDotProductAttention(nn.Module):def__init__(self,dropout:float0.1):super().__init__()self.dropoutnn.Dropout(dropout)defforward(self,Q,K,V,maskNone):d_kQ.size(-1)scorestorch.matmul(Q,K.transpose(-2,-1))/math.sqrt(d_k)ifmaskisnotNone:scoresscores.masked_fill(mask0,float(-inf))attn_weightsF.softmax(scores,dim-1)attn_weightsself.dropout(attn_weights)outputtorch.matmul(attn_weights,V)returnoutput,attn_weights看到masked_fill没这就是拉黑操作。padding位置直接填-infsoftmax后权重归零相当于在相亲现场把不符合条件的直接请出去。dropout更是狠训练时随机闭麦一些注意力权重防止模型过拟合——就像你同时聊十个相亲对象突然随机断网几个逼你认真跟剩下的谈。复杂度是O(n²·d_k)n是序列长度。这也是Transformer被吐槽最多的地方序列一长计算量爆炸。GPT-4处理长文档时那算力消耗比我交房租时的心绞痛还真实。二、Multi-Head Attention一个人同时开八个脑洞单头注意力就像你只用一只眼看世界虽然能看但立体感差点意思。Multi-Head Attention呢相当于给你脑袋上装八个摄像头同时从八个角度观察同一个对象。语法关系、语义关联、指代消解、情感倾向……每个头负责一块最后把八份报告拼一起交一份综合情报。代码里有个神操作不是真的定义8组独立的Q/K/V投影层而是各用一个Linear层投影完再view拆成8份。数学上等价但参数从3h个降到4个。Google这帮人是真会过日子省下来的显存够我多跑两轮实验了。多头注意力代码classMultiHeadAttention(nn.Module):def__init__(self,d_model:int,n_heads:int,dropout:float0.1):super().__init__()assertd_model%n_heads0self.d_modeld_model self.n_headsn_heads self.d_kd_model//n_heads self.W_Qnn.Linear(d_model,d_model,biasFalse)self.W_Knn.Linear(d_model,d_model,biasFalse)self.W_Vnn.Linear(d_model,d_model,biasFalse)self.W_Onn.Linear(d_model,d_model,biasFalse)self.attentionScaledDotProductAttention(dropout)defforward(self,Q,K,V,maskNone):batch_sizeQ.size(0)Qself.W_Q(Q).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)Kself.W_K(K).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)Vself.W_V(V).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)attn_output,attn_weightsself.attention(Q,K,V,mask)attn_outputattn_output.transpose(1,2).contiguous().view(batch_size,-1,self.d_model)outputself.W_O(attn_output)returnoutput注意那个.contiguous()transpose只是换了个看数据的姿势内存里还是老样子。你要是不contiguous一下后面的view直接报错PyTorch的脾气比你女朋友还难猜。W_O是最后的总编辑把八个头的输出揉巴揉巴合成一份d_model维度的终稿。mask的广播机制也很有意思。mask形状是(batch, 1, 1, seq_len)scores是(batch, n_heads, seq_len, seq_len)。PyTorch自动帮你广播不用手动unsqueeze。这感觉就像你去餐厅吃饭服务员主动问你要不要加辣——细节到位体验丝滑。三、Position-wise FFN每个token的健身房私教课注意力层处理完每个token还得去FFN里撸个铁。FFN(x) ReLU(xW_1 b_1)W_2 b_2。两层全连接中间夹个ReLU跟夹心饼干似的。关键是position-wise——同一个参数矩阵给序列里每个token轮流用公平得很跟健身房私教同时带十个学员但训练计划一模一样。内部维度d_ff通常是d_model的四倍512变2048。这就好比把数据从单人间塞进四人间折腾一番再搬回单人间。折腾的过程就是非线性变换给模型增加表达能力。论文实验说一层不够三层多余两层刚刚好。Google这帮人是懂中庸之道的比我家楼下奶茶店半糖的拿捏还精准。FFN代码classPositionWiseFeedForward(nn.Module):def__init__(self,d_model:int,d_ff:int,dropout:float0.1):super().__init__()self.linear1nn.Linear(d_model,d_ff)self.linear2nn.Linear(d_ff,d_model)self.dropoutnn.Dropout(dropout)defforward(self,x):returnself.linear2(self.dropout(F.relu(self.linear1(x))))dropout放在ReLU之后、第二次线性之前这是行业惯例。原始论文用ReLU后来GPT系列换成了GELU。GELU更平滑像给ReLU做了个SPA从硬切换变成软着陆。不过ReLU胜在简单直接就像直男表白虽然生硬但好歹把意思传达到了。四、Positional Encoding给token发座位号Self-Attention有个致命bug它分不清张三打了李四和李四打了张三。你把句子里的词随便换位置它输出一模一样。这就像一个脸盲症患者看谁都像同一个人完全靠衣服颜色区分——但你要是给他换件衣服他就彻底懵了。所以必须给每个token发张座位号告诉它你在第几个位置。Google用的是正余弦编码公式长得跟高数期末考最后一道大题似的。但核心思想就一条不同位置的编码不同而且相对位置可以通过线性变换推导出来。sin(αΔ) sinα·cosΔ cosα·sinΔ三角函数恒等式高中数学的遗产现在被Google拿来给AI指路。位置编码代码classPositionalEncoding(nn.Module):def__init__(self,d_model:int,max_len:int5000,dropout:float0.1):super().__init__()self.dropoutnn.Dropout(dropout)petorch.zeros(max_len,d_model)positiontorch.arange(0,max_len,dtypetorch.float).unsqueeze(1)div_termtorch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))pe[:,0::2]torch.sin(position*div_term)pe[:,1::2]torch.cos(position*div_term)pepe.unsqueeze(0)self.register_buffer(pe,pe)defforward(self,x):xxself.pe[:,:x.size(1),:]returnself.dropout(x)div_term用指数形式计算是为了数值稳定性。你要是直接算10000^(2i/d_model)浮点数精度早崩了跟用计算器算1除以3然后乘3结果不是1一样让人抓狂。register_buffer让pe跟着模型到处跑CPU、GPU随便切但不会被优化器盯上——相当于公司里的保洁阿姨到处都有她但KPI考核里没她。为什么不用可学习的位置嵌入三个原因能外推更长序列训练时没见过5000长度但编码天然支持、省参数、相对位置有数学保证。Google这波叫用数学省算力跟我用优惠券点外卖一个思路但人家省出来的是几台A100的电费。五、Layer Normalization给数据穿统一制服Batch Norm在CV界混得风生水起但到了NLP这儿就水土不服。为啥序列长度不一样batch大小也不稳定Batch Norm统计的均值方差跟过山车似的。Layer Norm说“算了我不管batch了我每个样本自己跟自己比。”对每个样本的所有维度先减均值、除标准差再乘个γ加个β。γ和β是可学习的相当于制服虽然统一但允许你微调尺寸。这跟学校发校服一个道理大家都穿蓝白相间但胖瘦可以自己调。LayerNorm代码classLayerNorm(nn.Module):def__init__(self,d_model:int,eps:float1e-6):super().__init__()self.gammann.Parameter(torch.ones(d_model))self.betann.Parameter(torch.zeros(d_model))self.epsepsdefforward(self,x):meanx.mean(dim-1,keepdimTrue)stdx.std(dim-1,keepdimTrue,unbiasedFalse)returnself.gamma*(x-mean)/(stdself.eps)self.betaunbiasedFalse用的是样本标准差跟原始论文保持一致。eps1e-6是防止除零的保险丝虽然实际数据几乎不会遇到全零向量但代码里不防一手就跟开车不系安全带一样——大概率没事但出事就是大事。生产环境直接用nn.LayerNorm就行手写版纯粹是为了让你看清校服是怎么裁剪的。六、Encoder Layer自注意力FFN的组合拳Encoder层就是自注意力打完FFN补刀。每个子层后面都跟一个残差连接和Layer Norm。残差连接x sublayer(x)是深度网络的救命稻草——梯度可以沿着shortcut直接传回去不用一层一层慢慢爬。这感觉就像你住30楼电梯坏了但旁边有个滑梯直通一楼。虽然滑梯有点陡但好歹比爬楼梯快。Encoder层代码classEncoderLayer(nn.Module):def__init__(self,d_model,n_heads,d_ff,dropout0.1):super().__init__()self.self_attnMultiHeadAttention(d_model,n_heads,dropout)self.ffnPositionWiseFeedForward(d_model,d_ff,dropout)self.norm1LayerNorm(d_model)self.norm2LayerNorm(d_model)self.dropout1nn.Dropout(dropout)self.dropout2nn.Dropout(dropout)defforward(self,x,maskNone):attn_outputself.self_attn(x,x,x,mask)xxself.dropout1(attn_output)xself.norm1(x)ffn_outputself.ffn(x)xxself.dropout2(ffn_output)xself.norm2(x)returnx注意self_attn的三个参数都是x这叫自注意力——Query、Key、Value全来自同一个序列自己查自己自己关注自己。有点像你深夜翻自己三年前的朋友圈一边看一边自我剖析“我当时怎么会发这个”这是Post-LN模式先残差再归一化。后来有些变体改成Pre-LN先归一化再残差训练更稳定。但原始论文是Post-LN咱们尊重经典就像吃北京烤鸭必须配甜面酱虽然有人爱蘸白糖但传统不能丢。七、Decoder Layer戴眼罩的传话游戏Decoder比Encoder多一个子层叫Cross-Attention。Encoder把输入序列的信息压缩成一份参考手册Decoder一边看自己之前生成的token一边翻这份手册决定下一个词输出啥。这像极了我写代码时的状态一边回忆自己上一行写了啥一边查Stack Overflow。但Decoder有个特殊规矩自注意力层必须戴眼罩只能看当前位置及之前的token不能偷看未来。这叫causal mask下三角矩阵上三角全填-inf。为什么因为翻译时你还没生成后面的词要是让模型提前看答案跟考试作弊有什么区别GPT就是这么自律地长大的虽然它后来学会了不少作弊技巧比如背题库。Decoder层代码classDecoderLayer(nn.Module):def__init__(self,d_model,n_heads,d_ff,dropout0.1):super().__init__()self.self_attnMultiHeadAttention(d_model,n_heads,dropout)self.cross_attnMultiHeadAttention(d_model,n_heads,dropout)self.ffnPositionWiseFeedForward(d_model,d_ff,dropout)self.norm1LayerNorm(d_model)self.norm2LayerNorm(d_model)self.norm3LayerNorm(d_model)self.dropout1nn.Dropout(dropout)self.dropout2nn.Dropout(dropout)self.dropout3nn.Dropout(dropout)defforward(self,x,enc_output,src_maskNone,tgt_maskNone):attn_outputself.self_attn(x,x,x,tgt_mask)xxself.dropout1(attn_output)xself.norm1(x)attn_outputself.cross_attn(x,enc_output,enc_output,src_mask)xxself.dropout2(attn_output)xself.norm2(x)ffn_outputself.ffn(x)xxself.dropout3(ffn_output)xself.norm3(x)returnxCross-Attention的Q来自DecoderK和V来自Encoder。Decoder每生成一个词就拿着这个词去Encoder的手册里查前面输入的句子哪个部分跟我现在最相关这机制让翻译准确率直接起飞比传统RNN的传话游戏强太多了。RNN传话传到最后一个词第一个词的信息早就失真得跟谣言一样了。八、完整Transformer组装高达的时刻最后一步把N层Encoder和N层Decoder堆起来加上嵌入层、位置编码、输出投影一台完整的Transformer就组装完毕。论文里N6d_model512n_heads8d_ff2048。这些数字不是拍脑袋定的是Google烧了不少TPU试出来的黄金比例。完整Transformer代码classTransformer(nn.Module):def__init__(self,src_vocab,tgt_vocab,d_model512,n_heads8,d_ff2048,n_layers6,dropout0.1,max_len5000):super().__init__()self.encoder_embednn.Embedding(src_vocab,d_model)self.decoder_embednn.Embedding(tgt_vocab,d_model)self.pos_encodingPositionalEncoding(d_model,max_len,dropout)self.encoder_layersnn.ModuleList([EncoderLayer(d_model,n_heads,d_ff,dropout)for_inrange(n_layers)])self.decoder_layersnn.ModuleList([DecoderLayer(d_model,n_heads,d_ff,dropout)for_inrange(n_layers)])self.fc_outnn.Linear(d_model,tgt_vocab)defforward(self,src,tgt,src_maskNone,tgt_maskNone):src_embself.pos_encoding(self.encoder_embed(src))forlayerinself.encoder_layers:src_emblayer(src_emb,src_mask)tgt_embself.pos_encoding(self.decoder_embed(tgt))forlayerinself.decoder_layers:tgt_emblayer(tgt_emb,src_emb,src_mask,tgt_mask)returnself.fc_out(tgt_emb)nn.ModuleList确保每一层的参数都被PyTorch登记在册不会变成黑户。Encoder和Decoder各自有独立的嵌入层虽然理论上可以共享但分开更灵活。fc_out把d_model投影到词表大小输出就是下一个token的概率分布——相当于给词典里每个词打个分分最高的就是天选之子。九、写在最后350行代码八年AI霸权你看完这350行可能会觉得就这GPT-4、Claude、LLaMA这些动辄千亿参数的怪物祖宗居然这么简洁没错伟大的架构往往简单到让人怀疑人生。就像爱因斯坦的Emc²就五个字符但改变了整个物理学。理解了这些基础组件你再去看GPT系列只用Decoder、BERT系列只用Encoder、LLaMA把ReLU换成SwiGLU、把LayerNorm换成RMSNorm——这些变体就不再是黑魔法而是在祖宗的基础上装修房子。有人拆墙有人加隔断但地基永远是这350行。所以下次有人跟你吹大模型多神秘你可以淡定地喝口咖啡说“神秘啥我看过它祖宗的源码就350行还没我微信聊天记录长。”当然看完这篇你要是还写不出Transformer那很正常。我看完《舌尖上的中国》也没学会做佛跳墙。但起码你再打开GitHub上那些开源大模型的代码时不会一脸懵了。这就是拆穿底裤的意义。P.S. 无意间发现了一个巨牛的人工智能教程非常通俗易懂对AI感兴趣的朋友强烈推荐去看看传送门https://blog.csdn.net/HHX_01