ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

242、【AI】【模型部署】基座模型研究:注意力机制深入

242、【AI】【模型部署】基座模型研究:注意力机制深入 【声明】本博客所有内容均为个人业余时间创作所述技术案例均来自公开开源项目如GithubApache基金会不涉及任何企业机密或未公开技术如有侵权请联系删除标题242、【AI】【模型部署】基座模型研究注意力机制深入背景上篇 blog【AI】【模型部署】基座模型研究loss 面面观把损失这一族概念理清了信息量− log ⁡ p -\log p−logp按真实分布加权平均得到熵H ( P ) H(P)H(P)交叉熵H ( P , Q ) H ( P ) K L ( P ∥ Q ) H(P,Q)H(P)\mathrm{KL}(P\|Q)H(P,Q)H(P)KL(P∥Q)其中H ( P ) H(P)H(P)与模型无关所以最小化交叉熵 ⇔ 最小化 KL真实标签 one-hot 时交叉熵退化为负对数似然而从最大似然连乘、取负对数也走到同一个目标。到这里损失、logits、softmax、反向都讲完了但模型内部最核心的部件注意力一直被一笔带过——第 238 篇只说Block 归一化 注意力 前馈第 240 篇只提到softmax 也用在注意力里。本篇把它展开查询、键、值到底是什么为什么要除以d \sqrt{d}d​因果掩码在遮什么多头与 GQA 又省了什么。模型部署注意力Attention是 Transformer 的心脏。它的核心思想只有一句话让每个位置按相似度从其他位置收集信息。本篇顺这条主线把从q / k / v q/k/vq/k/v到加权求和的每一步讲清并解释那些看起来没道理、其实都有原因的设计——缩放、掩码、多头、位置编码。注意力要解决什么一句话每个位置都要从序列里的其他位置取信息而取多少由相似度决定。最朴素的做法是取平均但那意味着所有位置同等重要显然不合理。注意力做的改进是加权平均——与当前词越相关的位置权重越大越不相关的权重越小。至于相关怎么算交给模型自己学。举两个例子在小明把书递给小红因为她要复习里理解她指向谁就需要把她与前文每个词分别比对写代码时某个变量该取哪一行的定义也靠这种相似度检索。这类当前词该向谁取信息的问题正是注意力要回答的。查询、键、值一次带权检索把注意力想成一次带权检索手里拿一个**查询Queryq qq去和每个位置的键Keyk kk比对比得越像就从对应的值Valuev vv**里取得越多图 1 里查询与四个键的相似度决定了四个权重0.50 / 0.30 / 0.15 / 0.05 0.50/0.30/0.15/0.050.50/0.30/0.15/0.05和为 1最终输出就是四个值按这些权重的加权和。一个经典类比是图书馆检索查询是要找什么键是每本书的标签值是书的内容——先用查询匹配标签再按匹配度把内容汇总。要注意q qq、k kk、v vv都不是原向量而是原向量经三个线性层投影得到的投影矩阵是要训练的。为什么是三个投影为什么不直接拿输入向量当q / k / v q/k/vq/k/v而要各自过一次线性层因为用什么特征去检索“用什么特征被检索”提取什么内容本来就是三件不同的事同一句话在不同任务里需要的检索角度、提取角度都不一样。三个独立的投影矩阵让模型分别学习这三种角色灵活性远高于三者共用一套表示——这也是注意力能学到丰富关系的前提。一次注意力的完整流程把上面的直觉写成计算步骤图 2 是整条链线性投影得到q , k , v q,k,vq,k,v→ 用q k ⊤ q k^{\top}qk⊤算相似度分数 → 除以d \sqrt{d}d​缩放 → 加掩码遮住不该看的位置 → softmax 变成权重 → 用权重对v vv加权求和。对应的代码model.py的Attention.forward核心就几行att(q k.transpose(2,3))/math.sqrt(self.head_dim)# 相似度 / 缩放attatt.masked_fill(~mask,float(-inf))# 因果掩码attF.softmax(att,dim-1)# 归一化成权重out(att v)# 加权求和用一个极小的例子走一遍设某个q qq与三个k kk的相似度分数是[ 3 , 1 , 0 ] [3,\,1,\,0][3,1,0]softmax 后约得权重[ 0.84 , 0.11 , 0.05 ] [0.84,\,0.11,\,0.05][0.84,0.11,0.05]再用这三个权重对三个v vv加权求和就得到该位置的输出。整个注意力层对每个位置做的就是这么一次看别人、取信息的操作。顺带把涉及到的张量形状过一次H HH为头数、d h d / H d_hd/Hdh​d/H步骤形状作用投影( B , H , S , d h ) (B,H,S,d_h)(B,H,S,dh​)得到多头的q / k / v q/k/vq/k/v相似度q k ⊤ q k^{\top}qk⊤( B , H , S , S ) (B,H,S,S)(B,H,S,S)两两位置比对softmax( B , H , S , S ) (B,H,S,S)(B,H,S,S)变成权重加权( att ) v (\text{att})v(att)v( B , H , S , d h ) (B,H,S,d_h)(B,H,S,dh​)收集信息拼接( B , S , d ) (B,S,d)(B,S,d)融合多头注意力的临时账本正是那个( B , H , S , S ) (B,H,S,S)(B,H,S,S)矩阵——它随S SS平方增长是显存的大头后面讲 KV cache 也绕不开它。为什么要除以d \sqrt{d}d​这是最容易被当成魔法数字的一步。原因是量级q qq和k kk的点积是对d dd个数相乘再求和若各分量近似独立、方差为 1则点积的方差约为d dd标准差约为d \sqrt{d}d​——维度越高点积越大。图 3 画出这个增长d 1024 d1024d1024时点积的量级可达32 3232。分数一大softmax 就会被推到饱和——最大的那个占尽几乎全部权重、其余接近 0于是梯度消失、学不动。除以d \sqrt{d}d​把量级拉回约 1让 softmax 处在能平滑分配权重的工作区间。它不是可调超参而是按方差推导出的固定常数。这里可以和第 240 篇的温度对照着看温度是在 softmax 前对 logits 除以T TT同样在调送进 softmax 的分数尺度区别在于温度是推理时人为可调的旋钮而d \sqrt{d}d​是训练时固定、用于稳定的常数。因果掩码不许看未来语言模型是自回归的预测第t tt个 token 时只能看到它之前的内容不能偷看后面。做法是把未来位置的相似度分数在 softmax之前设成− ∞ -\infty−∞softmax 之后它们的权重就变成 0图 4 是掩码矩阵下三角可见、上三角全遮。注意必须在 softmax 之前遮若先 softmax 再置零权重和就不再是 1等于破坏了加权平均。对角线也保留每个位置能看自己所以是看自己和左边。多头与 GQA把注意力做多份多头注意力Multi-Head AttentionMHA的思路是一组q / k / v q/k/vq/k/v只能捕捉一种相似度那就并行做多组——每个头负责一种关系有的头可能看语法、有的看指代最后把各头输出拼接、再经一个输出投影融合。这让模型能同时关注多种模式表达力显著增强。代价是显存推理时需要缓存每个头的k kk和v vv第 243 篇要讲头越多缓存越大。分组查询注意力Grouped-Query AttentionGQA的对策是查询头保持多个键值头只留少数几组让多个查询头共享同一组k / v k/vk/v图 5 里 8 个查询头只对应 2 组键值头每组被 4 个查询头共享既省下k / v k/vk/v缓存又基本保留多头的表达力——这也是 Qwen2 等模型普遍采用的折中第 228 篇展开过。自注意力与交叉注意力上面讲的q / k / v q/k/vq/k/v都取自同一段序列这叫自注意力self-attention——“在自己内部互相看”。若查询来自一端、键值来自另一端如解码器看编码器、或模型看外部检索结果就是交叉注意力cross-attention。语言模型解码器架构的主体是自注意力这也是上下文里的 token 互相交换信息的实现方式。位置从哪来RoPE注意力本身只算相似度、不看顺序——如果打乱 token 顺序结果不变。所以必须把位置信息注入进去。旋转位置编码Rotary Position EmbeddingRoPE的做法是按位置给q qq、k kk的每一对分量做旋转让相对位置体现在点积里第 227 篇展开。这样注意力在算相似度时就天然带上了谁在前、谁在后。拼回一个 Block把这些部件装回第 238 篇的 Block前向顺序是xxself.attn(self.input_layernorm(x),cos,sin)# 注意力子层xxself.mlp(self.post_attention_layernorm(x))# 前馈子层每个子层都是归一化 → 子层 → 残差相加归一化RMSNorm稳定数值残差连接让梯度能顺畅回传。注意力负责在位置之间交换信息前馈网络负责在每个位置内部做变换——一横一纵构成了 Transformer 的基本单元。注意力的代价随长度平方增长最后提一句工程账相似度矩阵是S × S S\times SS×S的S SS是序列长度所以注意力的计算量与显存都随S SS平方增长——序列翻倍、开销变四倍。这正是长上下文昂贵、以及各种高效注意力方案被不断提出的根本原因。理解这一点第 243 篇讲的 KV cache 才显得顺理成章。三个常见误区注意力不是把整段记住它是一次加权平均输出维度与单个v vv相同并不会把上下文存进输出掩码必须在 softmax 之前加否则权重不再归一化加权平均就失真了d \sqrt{d}d​不是玄学超参它是把点积方差拉回 1 的固定常数目的是避免 softmax 饱和。一句话记忆注意力让每个位置按相似度从其他位置收集信息用线性投影得到查询q qq、键k kk、值v vv以q k ⊤ q k^{\top}qk⊤算相似度除以d \sqrt{d}d​把量级拉回 1否则 softmax 饱和、梯度消失加因果掩码遮住未来且必须在 softmax 之前softmax 得到权重后对v vv加权求和多头并行捕捉多种关系GQA 用共享的键值头省显存RoPE 以旋转注入位置它的开销随序列长度平方增长最后与归一化、残差、前馈一起拼成 Block。OK本篇先到这里如有疑问欢迎评论区留言讨论祝各位功力大涨技术更上一层楼更多内容见下篇 blog【AI】【模型部署】基座模型研究KV cache 与自回归生成
返回列表