首页 / 资讯中心 / 文章详情

注意力机制从原理到代码:QKV、多头、掩码与KV Cache实战指南

注意力机制从原理到代码:QKV、多头、掩码与KV Cache实战指南 ★ FEATURED ARTICLE
简介这是一份面向深度学习初学者和自然语言处理研究者的注意力机制专题课件系统讲解序列到序列模型中的编码器—解码器结构为何需要注意力以及固定长度语义向量在长输入下的信息丢失问题。课件从循环神经网络只能处理等长序列的局限出发逐步引出注意力思想的来源并重点剖析Bahdanau提出的加性注意力、Luong提出的注意力变体以及硬注意力、局部注意力与自注意力等主流分支分别说明其计算逻辑、优缺点和适用任务。结合机器翻译、文本摘要、语音识别、图像描述等典型场景通过图解和权重配色示意直观呈现注意力分数如何在不同输入位置上分配帮助读者理解每类注意力的工作机制。资源为1个PPT文件大小约2.7MB结构紧凑、图文并茂既适合教师备课参考也适合学生课前预习、课堂展示或考研复试快速回顾。目前已有1183人学习下载能够显著降低理解门槛也为后续学习Transformer架构打下扎实基础。1. 注意力机制不是玄学你在Transformer、Seq2Seq和图像里早就见过它搜索“深度学习中attention注意力机制.ppt”的人多半是要准备技术分享或者被面试题逼到了。注意力机制看起来像一套优雅的数学公式但真正困扰开发者的不是读不懂论文而是“我怎么把它用起来”。这篇笔记会把你从“知道Attention”拉到“能写最小实现、能调参数、能排坑”的状态。我不打算复制PPT里的定义而是把QKV、多头、mask、缓存这些实操点拆开每个环节都给出能直接跑的代码和踩过的坑。你在Seq2Seq、Transformer、图像分类里见到的各种注意力底层都是同一套加权求和逻辑。新手可以照步骤做熟手可以对照参数和边界。2. 从“对齐”到“加权求和”Attention的核心数学直觉与QKV拆解2.1 早期Seq2Seq的Attention为什么要在译每个词时“回头看”注意力机制最早出圈是在机器翻译里。当时常见的做法是用一个编码器把整个输入句子压成一个固定长度的向量解码器每一步都从这个向量里取信息。句子一长最后那个向量扛不住信息量前面的内容基本被后面盖掉了。后来某篇工作提出解码时每一步不要把目光锁死在那个固定向量上而是回到编码器的每个隐藏状态上“按需取用”。具体做法是计算当前解码状态和每个编码器隐藏状态的相关程度得到一个权重分布然后用这个权重对编码器隐藏状态做加权求和得到这一步的上下文向量。这个“回头看”的操作本质就是给序列里的每个位置分配一个重要性分数。分数高的位置对当前输出的影响大分数低的位置几乎被忽略。这就是Attention最初的含义注意力权重告诉模型“该看哪个位置”加权和把“看到了什么”送到解码器。它不需要额外的监督信号权重完全由数据驱动学出来所以经常被形容为“网络学会了自动对齐”。现在很多文章一上来就讲Q、K、V但如果你没看过这个背景就会觉得QKV只是花哨的改名。其实Q就是“当前的查询”K是“序列里每个位置的钥匙”V是“序列里每个位置的信息”。相关程度就是拿查询和每一把钥匙做匹配匹配度越高对应信息被取出的比例就越大。这个视角贯穿所有变体包括Transformer里的自注意力。2.2 缩放点积Attention的完整公式为什么除以根号d_k现在主流公式长这样Attention(Q,K,V)softmax(QK^T/√d_k)V其中Q、K、V分别是查询、键、值矩阵d_k是键向量的维度。计算分四步先把Q和K做点积得到相似度分数再除以√d_k做缩放然后softmax归一化成权重最后用权重对V做加权求和。QK^T如果形状分别是[B, T_q, d_k]和[B, d_k, T_k]得到的就是[B, T_q, T_k]的分数矩阵表示每个查询对所有键的匹配程度。为什么要除以√d_k这是很多初学者最容易忽略的点。假设Q和K的元素是独立的随机变量均值0、方差1那么点积后的期望是0方差却会变成d_k。d_k越大分数的方差越大softmax输入值会落到梯度极小的区域训练时梯度很容易消失。除以√d_k后方差被拉回接近1softmax的输入范围更稳定。这个缩放是固定操作不需要学习也不要自作主张改成其他值它已经是被反复验证过的默认项。还有一种常见的变体叫加性注意力用一个小网络来算匹配度而不是直接用点积。它在早期Seq2Seq里常见但缩放点积因为计算更高效、更容易向量化后来成了主流。所以你现在看到“Attention”默认都是指缩放点积注意力除非特别说明。2.3 用PyTorch 20行实现单头Attention最小可用代码import torch import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, maskNone): d_k query.size(-1) # 点积分数形状 [B, T_q, T_k] scores torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) output torch.matmul(weights, value) return output, weights这段代码里query、key、value的最后一维分别是d_k最前面两维可以是[B, seq_len]。如果query长度是T_q、key长度是T_k那么输出形状是[B, T_q, d_v]。注意value的维度d_v可以和d_k不同但实际工程里通常相等方便做残差拼接。mask参数采用“0表示屏蔽”的约定。把对应分数填充成负无穷softmax之后权重就是0。为什么要用负无穷而不是直接置0因为如果把分数直接置0softmax后仍然会有小概率分配到该位置负无穷可以保证数学上彻底忽略。这个习惯在所有Transformer实现里都成立后面排查问题时你也会用到它。写这个函数时我一般会顺手打印出中间张量形状。你拿一个B2、T_q3、T_k4的输入跑一次就能看到scores变成[2,3,4]。先确认维度再谈效果这是避免后面对着头数调半天却没用的一步。3. Self-Attention与Multi-HeadTransformer的基石怎么搭起来3.1 Self-Attention如何让每个token看到全序列自注意力Self-Attention是Transformer的核心。一句话概括序列里的每个token都要看着整个序列重新计算自己的表示。它跟上一章的Seq2Seq注意力最大的区别是Q、K、V都来自同一个输入。每个token先经过三个线性投影变成自己的Q、K、V然后执行缩放点积注意力。这样一来每个位置都能直接和任意远的位置建立依赖卷积和循环网络那种“隔着好几层才能接触”的问题被绕开了。举个例子“他喜欢吃苹果”和“她喜欢吃苹果”只差一个字但语义完全变了。要理解“他”网络需要看“喜欢”和“吃”等远距离信息。用自注意力这句话里每个词两两之间都算一个分数模型可以自由决定哪个词对当前词更重要。训练时不带任何人工规则全靠任务loss反向传播来塑造这些分数。但自注意力有个代价分数矩阵是O(T^2)的。序列长度从512涨到1024显存占用直接翻4倍这在后面第5章会提到。所以工程上经常会限制序列长度或者用稀疏注意力、线形注意力等近似方案来处理超长文本。3.2 Multi-Head Attention的拼接与投影多子空间带来的增益单头注意力只能学一种“匹配模式”实际任务里需要同时关注语法关系、指代关系、局部连贯性等多种信息。多头注意力就是把d_model维拆成h个d_k维子空间每个头独立做注意力最后拼接起来再投影回d_model。谷歌最早提出的配置是d_model512、h8每个子空间d_k64。h个头并行计算彼此不共享参数所以不同的头会倾向于学到不同的注意力模式有的头关注前后词有的头关注主语谓语这就是“多子空间”的价值。实现时要注意每个头的维度d_k d_model / h所以d_model必须能整除h。这不只是一个数学要求更是一个工程约束。如果你选了不规则的头数比如h7、d_model512就麻烦了。所以选头数前先看d_model能不能被它整除。为了兼容预训练模型一般不要试图去改这个比例。多头注意力在计算结束后并不是简单把h个结果相加而是把所有头的输出在特征维度上拼接起来生成一个长度为d_model的向量再做一次线性投影W_o。这次投影的目的是让不同头的信息做一次交叉融合而不是各说各话。如果你在调试时发现效果不好检查一下是不是忘了最后这个投影或者维度拼接时contiguous搞丢了这类问题很常见。3.3 从零拼一个多头注意力模块代码与维度核对import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_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.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.size() # 投影并拆出头维度: [B, T, h, d_k] - [B, h, T, d_k] Q self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 缩放点积注意力scores形状 [B, h, T, T] scores Q K.transpose(-2, -1) / (self.d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) # 注意dropout在softmax之后 context attn V # [B, h, T, d_k] # 把头拼回去: [B, T, h, d_k] - [B, T, d_model] context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.W_o(context) return output这段代码里最容易被搞混的是view和transpose的顺序。先用view把[B, T, d_model]拆成[B, T, h, d_k]再transpose成[B, h, T, d_k]这样Q和K做矩阵乘时每个头是独立的。如果你先transpose再view会得到完全错乱的排列算出来的score没有任何意义。mask参数可以传两种形状。一是[B, 1, 1, T]的padding mask代表每个位置是否是有效token二是[B, 1, T, T]的causal mask下三角为1上三角为0用于自回归生成。因为维度广播机制[B, 1, T, T]会自动广播到[B, h, T, T]所以不需要复制多份。attention dropout放在softmax之后这是很多实现里的标准位置它随机让一部分注意力权重归零相当于正则化。我调试这段代码时习惯用一个小例子batch_size2、seq_len5、d_model8、num_heads2跑一次打印Q、K、V的shape确认最后输出shape又变回[2,5,8]。如果哪一步维度对不上问题一定出现在view或transpose上不要瞎调dropout。4. 训练与推理中的Attention参数怎么设d_model、头数、dropout与mask4.1 头数、维度和投影一个表格看清常见配置模型规模d_modelnum_headsd_k典型用途微型64416临时Demo、小数据集小号128432轻量分类、短文本基础512864常见Transformer基准大号10241664长文本、预训练模型头数不是越大越好。头越多每个头的维度越小单个头能表达的信息容量就越少。头数增大到一定程度后不同头会开始“冗余”。很多实验表明8头两头效果接近但不是翻倍就有翻倍效果。常见做法是保持每个头d_k在32到64之间。如果你自己搭模型建议先从d_model128、heads4开始跑通性能达标后再往大扩。直接照搬5128组合在小数据集上很容易过拟合。还有一个容易被忽略的参数输出投影W_o是否带bias。默认nn.Linear是带bias的但在某些实现里为了追求算力优化会把它设成False。这不会显著影响效果但如果你要加载预训练权重必须和原配置保持一致。另一个相关参数是注意力dropout和FFN的dropout它们通常是两个独立值。注意别把attention dropout设成0.5这会严重干扰学习一般0.1是个稳定起点。4.2 训练阶段必调的三个参数dropout、缩放因子和mask训练时最需要关注的三个点是dropout、缩放因子、mask它们各自都可能成为性能瓶颈。第一dropout。模型里通常有两个dropout注意力权重上的dropout和FeedForward后的dropout。注意力权重dropout作用于softmax之后的权重它会让某些位置的整条信息传递被阻断防止个别token过度主导。如果你发现训练loss下降很慢先看它是不是设太大。我的习惯是注意力dropout0.1FFN dropout0.1到0.2。若你的数据集非常小可以考虑降到0.05但不要完全去掉。第二缩放因子。公式里固定用√d_k这个不需要调。不过有一种常见误用把缩放因子当作超参数去搜或者直接去掉。去掉缩放后大d_k会导致softmax梯度消失这在长序列上尤其明显。所以缩放因子不是可选项而是标准操作。你看到的某些第三方代码里写了“scale”参数其实只是把除以√d_k换成了乘以scale本质一样。第三mask。训练时最常见的是padding mask和causal mask。padding mask负责把不足固定长度的批量样本填0的位置屏蔽掉causal mask负责保证自回归模型看不到未来。如果你两个都加需要把两个mask做位与操作一个位置要想参与注意力必须同时满足不是padding且在因果范围内。我见过很多翻车现场就是只加了padding mask忘了causal mask导致训练时模型提前看到了答案loss异常低但推理效果一塌糊涂。检查方法很简单打印scores矩阵里被mask后的分布确认右侧上三角全是负无穷。4.3 推理阶段的速度优化缓存KV与attention的代价到了推理阶段尤其是自回归生成Attention的代价会变得非常扎眼。假设当前生成了第n个token普通实现会把这n个token重新过一遍编码器重新计算所有Q、K、V再做一次完整n×n的注意力。这就导致生成第n步的计算开销是O(n^2)总开销O(n^3)。实际项目里n稍微上千延迟和显存就撑不住了。常用的优化手段是KV Cache。因为解码第n个token时前面所有token的K、V其实已经算过而且不会变。只有Q是当前步新算的。所以可以把每一步的K、V缓存下来下一步只需要算新token的K、V然后拼接到缓存矩阵上注意力分数矩阵虽然还是n×n但计算量从“重新算所有n个token”降为“只算一个新的token”。在很多推理框架里你看到的关键词“past_key_values”就是这个东西。实现KV Cache时要注意缓存的形状是[B, h, T_prev, d_k]每一步新token的K/V是[B, h, 1, d_k]拼接后变成[B, h, T_prev1, d_k]。拼接发生在时间维度不要搞错。还有缓存必须和输入长度对得上如果你在批量生成不同长度的序列需要给每个样本维护一个独立的cache否则短的序列会读到长的序列的K/V输出直接崩掉。显存方面即便有KV Cacheattention矩阵本身仍是O(T^2)。长序列生成时这一步就会成为显存瓶颈。常见的处理是把attention计算切成块或者用FlashAttention这类kernel融合把中间矩阵不落地到全局显存。作为一个独立开发者你最快能做的优化是推理时把batch_size调小或者限制max_new_tokens让attention矩阵不要爆掉。要是还不满足再考虑FlashAttention或改用线性复杂度注意力。5. Attention的5个常见“玄学”坑排查顺序与修正方案5.1 现象训练loss不降——原因mask把pad位置也参与注意力计算新手第一次搭Transformer最常遇到的问题是loss一直在一个高位震荡怎么调学习率都不动。我见过某开发者调试了一整晚最后发现是padding位置没屏蔽。原因是批量输入先pad到统一长度比如一批里有3个样本长度分别是3、5、4pad到5。没加mask的话第2个样本的attention会看到后面两个pad位置pad位置通常被初始化成0于是模型学会去关注这些废位置有效信息被稀释loss自然降不下去。解决方法是构造padding mask把pad位置置0然后传入attention函数在softmax之前把那些位置的分数改为负无穷。mask形状可以是[B, 1, T]广播到所有头。我在代码里习惯用mask 0作为填充条件这样mask里1表示有效、0表示pad。如果你用的库是0表示有效、1表示pad记得反转一下否则会把有效位置全屏蔽了loss直接飞起来。排查时先打印scores矩阵看pad位置的分数是不是负无穷。如果是-1e9或者某个有限值说明mask根本没生效。再检查mask是否在Q、K、V投影之后才传入或者在多头模块里用的是mask而不是~mask。这些细节不看代码一眼根本发现不了。5.2 现象多卡训练时显存爆炸——原因attention矩阵O(n^2)被复制到每卡有人用单卡训练512长度的序列很正常换成多卡DDP后明明batch_size没变显存却直线上涨。attention机制本身有一个O(T^2)的分数矩阵如果序列长度是1024就是100万个浮点数单卡还能扛。但多卡DDP不是专人负责某一块序列而是每张卡保留完整的batch副本和梯度副本所以attention矩阵也会在每个进程里各自存一份。显存占用是单卡的总和不会因为你加了卡就分掉。还有一个隐性膨胀点如果用PyTorch原生DDP且没有开启gradient checkpointing反向传播会保存所有中间激活包括注意力分数和softmax权重。序列越长激活越大。我对付这个问题的顺序是先开自动混合精度AMP把中间张量从FP32减到FP16显存能砍一半左右然后开gradient checkpointing牺牲一点重算时间换显存最后如果还爆就只能降低batch_size或序列长度。不要一开始就怀疑是多卡配置问题多数时候是attention矩阵大小乘以卡数超出物理显存。5.3 现象推理结果重复——原因softmax温度与attention dropout不一致生成任务里模型经常输出同一个词反复循环比如“我爱吃我爱吃我爱吃”。很多人以为是解码策略的惩罚参数没调好但背后还有一个更隐蔽的坑训练时attention dropout随机把一部分权重置零而推理时没有开dropout这会导致推理时的注意力分布和训练时不一样尤其对短序列影响更大。如果训练时dropout较高模型在推理时看到的信息过于确定会更容易陷入重复模式。解决思路有两类。一类是推理时也开启dropout但这会让生成结果随机性变大通常不建议。另一类是把注意力softmax的温度调到小于1让分布更尖锐减少随机重复。具体做法是在softmax之前把scores除以一个温度系数温度越低输出越确定。但温度太低又会坍缩到最高概率的token反而重复。我的经验是先确认训练和推理的dropout设置一致再调温度为0.8到0.9结合重复惩罚因子比单独调一个参数有效。如果你用的是预训练模型还得检查加载权重时是否强制设置了model.eval()。很多库在推理时会自动关闭dropout但如果你手工切换模型状态晚了整个推理结果都会带随机性看起来就像“模型变笨了”。这种问题耗时很长但排查起来很简单打印推理次数如果两次完全相同输入给不同结果说明dropout没有关干净。5.4 现象位置信息丢失——原因没有加位置编码或者加了但没参与attention自注意力是置换等变的你把序列顺序打乱得到的输出只是同样位置被打乱不会携带任何顺序信息。这一点在理解上很容易踩坑。某同学在搭建情感分类模型时直接拿Embedding后的向量进Transformer没加任何位置编码。模型训练出来准确率很低因为他觉得attention已经“记住了”每个词的位置。其实attention只关注词与词之间的匹配关系并不会告诉你“第一个词”是哪个绝对位置。解决方案是加位置编码。常见的有固定正弦编码和可学习位置编码。正弦编码根据公式生成不同频率的正弦余弦值叠加到embedding上可学习编码则是构造一个形状为[max_seq_len, d_model]的变量随模型一起训练。无论哪种都要确保位置编码在attention计算之前被加到词嵌入向量上而不是加在输出之后。否则attention看到的仍是原始embedding位置信息进不了Q、K、V。调试时可以做一个简单验证输入两个完全相同词但顺序不同的句子比如“猫追狗”和“狗追猫”检查模型最后的输出表示是否不同。如果两者完全一样问题几乎肯定出在位置编码没生效。还要注意位置编码的角色在多头注意力中会被每个头以不同方式调制所以加位置编码的幅度不能太大一般和embedding同尺度否则会淹没语义信息。5.5 现象可视化attention像个噪声图——原因多头平均后互相抵消或者没对齐tokenizer做模型可视化是注意力机制最有成就感的一步但经常会得到一团乱麻。有人把多头注意力权重直接平均后画出来发现完全看不出对角线或者局部模式。原因很简单每个头关注的信息可能完全不同一个头关注语法关系另一个头关注长距离依赖强行平均会把这些模式全部抵消掉。正确做法是分头可视化先确认每个头的权重分布再合并。另一个崩溃点是tokenizer对齐。如果原始句子是中文按字切分但训练时用的tokenizer是BPE可视化时你拿字粒度的句子去对应attention权重位置就会错位。需要先用同一个tokenizer把输入转换再拿解码后的tokenid去映射注意力矩阵的行列。不然你图上显示的“注意力”根本不在这个字上看起来当然像噪声。最后的常见问题是对角线权重过高。如果attention权重集中在对角线说明模型几乎没有跨位置交互这通常发生在训练不足或序列长度不够时。这时不要急于下结论说代码有问题先检查是不是位置编码没有参与注意力或者是不是层数太少。我的习惯是同时打印scores的标准差如果标准差接近0说明softmax已经接近均匀分布模型还在“偷懒”需要增加训练步数。6. 把这张“PPT”讲成落地技能三张图讲清原理一个脚本做原型验证如果你是要做分享与其堆公式不如准备三张图。第一张画Seq2Seq不带attention的瓶颈就是那个被固定向量压扁的过程第二张画QKV的映射关系用三个不同颜色的方块表示Q、K、V再把点积到softmax到加权求和的过程用箭头连起来第三张画多头在d_model维度上切出4条竖条每条算一次attention再拼回去。这三张图能让听的人五分钟内建立起直觉剩下的时间留给代码和现象。验证自己是否真的理解了注意力机制我建议写一个极小的脚本随机生成一组整数序列给每个位置接一个可学习的embedding然后只套一个单头attention不加位置编码用“预测序列中是否出现某个数字”这种玩具任务训练。如果attention能学会把权重放到那个数字出现的位置说明你的实现是通的。这比直接跑一个大模型更省时间也能快速暴露mask或维度问题。参考脚本的输入输出形状对照第2章的函数即可。我早期做这个方向时曾经为了一条维度对不上的错误把多头实现重写了三遍。后来养成一个习惯每个新模块先跑一个batch把中间张量的shape全部打印出来确认无误再往上搭。注意力机制本身并不复杂复杂的是它和周围模块的边界条件。你真正需要记住的是QKV是三种投影scores是匹配度mask是边界dropout是正则KV Cache是代价。这些点串联起来你再去读任何Transformer源码都会顺畅很多。希望这篇笔记能帮你在注意力机制这条路上少走几段弯路。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助?
咨询建站