扩散模型火了两三年现在回头补“自注意力机制”这门课的人其实特别多。大家刚开始跑通一个扩散模型的训练流程时注意力模块往往只是当作“U型网络里的一层”原样抄下来直到哪天开始调大分辨率、加长文本条件、或者想让生成结果严格对齐布局才发现自己根本说不清这一层到底在干什么。这篇文章就从一个实际调模型的角度把扩散模型里的自注意力机制从头到尾捋一遍包括它的运算本质、为什么非它不可、训练和推理时怎么突破算力瓶颈、注意力图如何用做调试工具以及新架构对注意力模块的替代思路。1. 为什么扩散模型离不开自注意力从图像生成管线说起1.1 扩散模型到底在做什么扩散模型的生成过程可以拆成两个方向。前向过程是往一张干净图片上逐步叠加高斯噪声直到图像完全变成一片白噪反向过程则是训练一个去噪网络让它可以一步、一步地把噪声“刮掉”。采样的时候我们从纯噪声出发连续执行几十步到上千步的去噪预测最终还原出一张图像。这里有个容易被忽视的核心点去噪网络每一小步都在做一个全局判断。假设画面里有一辆停在草地上的红色汽车纯噪声状态下没有任何局部线索但当去噪进行到中段车身轮廓隐隐出现时网络必须同时知道三件事车身该是红色、车轮在车身下方、草坪应该围绕在四周。这些信息分散在整张特征图的不同区域如果网络只看局部邻域就容易出现“车身补得挺好但轮胎长到车顶”这种结构错乱。扩散模型生成的是整幅图像的联合分布每一次预测噪声都需要对全局结构做出隐式建模这是自注意力机制能发挥价值的最根本原因。早期的卷积主干网络把图像看成许多局部感受野的拼接每一层卷积所能看到的信息受限于卷积核大小。假如没有注意力机制要让网络“看”到遥远位置的信息只能靠不断堆叠卷积层来扩大感受野但这样既增加参数量又会损失精确的长距离依赖。在扩散模型里加入自注意力本质上是给局部卷积模型装上了一个“全局搜索”的插件让每个位置都可以直接读取整幅特征图的信息再决定怎样修正当前像素的噪声预测。1.2 卷积的天然短板与自注意力的补位卷积有几个天生的优势参数共享、平移等变、计算高效这让它非常适合处理像素级细节。但它也有一个痛点就是感受野受限。严格地说叠加足够深的卷积层后感受野可以覆盖全图但那种覆盖是通过多次非线性变换间接实现的远距离像素之间的关系已经被层层压缩很难精确保持。自注意力走的是另一条路它直接把特征图上任意两个位置的联系显式算出来。某个位置的更新值是所有其他位置根据相似度加权汇总的结果。这个操作没有“距离”的概念空间上隔得很远的像素只要特征相似就会被赋予较高的权重。扩散模型特别需要这种性质因为噪声预测本身就是一种全局一致性任务前景物体的结构、背景纹理的延续、物体之间的遮挡关系都需要跨区域协调。从生成模型对比的角度也能理解。像早期的对抗生成网络生成器可以通过上采样直接输出图像不需要在中间特征上做那么多精细的全局推理而扩散模型的每次去噪都相当于要做一次“图像的逆向推算”中间特征图承载了从粗粒度结构到细粒度纹理的全部信息。在关键分辨率层级上假如缺少全局注意力模型会倾向于过度关注局部统计量最终导致生成结果出现结构漂移——典型的症状是单看每个小区域都合理整体却完全不成形。1.3 从U型网络到注意力层的演化路线大多数扩散模型的骨架采用了一种编码器-解码器形态先逐步下采样提取高层语义再逐步上采样恢复细节同时通过跳连把同分辨率特征拼回来。早期版本完全由卷积构成后来从业者发现单纯卷积难以控制强文本条件下的空间布局才开始在中间低分辨率层插入注意力模块。为什么要放在低分辨率层原因很朴素注意力计算量随序列长度平方增长而低分辨率特征图的像素数少。以常见的特征图尺寸序列来看如果模型在64×64分辨率特征图上做全局注意力序列长度为4096注意力矩阵规模是4096×4096还能接受如果放在128×128上序列长度变成16384矩阵规模翻了16倍显存立刻爆掉。所以经典的做法是在最低分辨率几层之间使用注意力而高分辨率层继续用卷积去修细节。这条路线在后来被进一步推向极端有人干脆去掉了卷积和U型结构让所有层都用注意力堆叠这也就是后文要聊的纯注意力扩散架构。2. 自注意力在扩散模型里的三种打开方式2.1 底层运算Q、K、V与注意力分数把自注意力的计算拆开看没有想象中那么玄。假设输入是一张特征图形状可以简化为批量大小、通道数、高度、宽度。注意力机制的第一步是把特征图重新排列成一系列“词向量”每个空间位置对应一个词向量序列长度就等于高度乘宽度。随后用三组可学习的线性投影分别把每个词向量映射成查询向量、键向量和值向量。接下来计算注意力分数矩阵查询向量与所有键向量做点积再除以一个缩放因子然后经过Softmax归一化。这一步得到的矩阵每一行表示“当前这个位置应该关注其他位置的程度”。最后用这个权重矩阵去加权求和所有值向量得到当前位置的更新表示。这个流程可以用一句话记自注意力让每个位置“问一圈所有人再按照相关度汇总信息”。缩放因子为什么这么重要当维度较大时向量点积的结果方差也会变大Softmax之后容易出现梯度消失或注意力分布过于尖锐的问题。除以键向量维度的平方根可以把点积结果的方差拉回到一个稳定的范围这在扩散模型训练中尤其关键因为去噪网络本身是一个多步迭代的回归任务注意力分布如果过于极端训练早期就会不稳定。上述过程可以通过矩阵乘法一次性完成。不过很多实现会拆分成多次头注意力让不同的注意力头各自在不同的语义子空间里计算。有的头可能更关注前景物体本身的完整结构有的头可能关注背景与前景的边界有的头则负责建立跨区域的纹理一致性。把这些头的结果拼起来再过一层线性映射就是多头注意力的基本流程。2.2 空间注意力单个特征图内部的全局推理空间注意力是最常见的一种形态查询、键、值都来自同一个特征图。它的作用是在图像区域之间建立联系相当于让模型自己回答“哪些地方应该长得相似”。举个例子生成一头长颈鹿时脖子和腿虽然离得远但应该共享相似的花纹纹理生成一个人物时左手和右手应当有类似的皮肤质感。空间注意力可以轻松捕捉这种跨区域的相关性。空间注意力在扩散模型中的位置往往不止一层。通常编码器端和下采样之后的中层各有一组注意力块解码器端在恢复高分辨率之前还会再算一次。这样做的意义是让信息在压缩语义和高频细节之间反复交换避免全局结构信息只在单一尺度上被建模。不过空间注意力的问题也随之而来。在分辨率高的特征图上每个位置的词向量维度如果又比较大那么中间注意力矩阵的存储开销会非常可观。实际操作中大家会去数“有效序列长度”如果特征图是32×32序列长度只有1024如果特征图是16×16序列长度只有256。所以许多实现把注意力限制在16×16和32×32两层而64×64之上仍然由卷积负责。这里很容易产生一个误区以为注意力层放得越多越好结果序列长度失控训练速度断崖式下降。2.3 交叉注意力把文本、类别和时间条件“画”进图像扩散模型的条件生成一般离不开交叉注意力。它的计算方式与空间注意力基本相同唯一区别是查询向量来自图像特征分支键和值向量来自另一个条件序列。如果条件序列是文本嵌入那么注意力矩阵里的每个元素就代表了“某个图像区域”与“某个文本词元”之间的关联强度。这个过程直观上可以理解为模型每次画一个图像区域时都要去文本序列里“查阅”最相关的几个词再把那些词对应的信息融合进来。比如条件文本里有“红色”和“汽车”当模型处理车身的图像区域时注意力权重会更多地落在“红色”“汽车”这些词元上当处理背景区域时权重则会落在“草地”“天空”等词元上。整个生成过程实质上就是文本词元把语义“投射”到二维空间上的过程。除了文本交叉注意力也可以用在类别条件、深度图条件、边缘图条件等场景。还有一点值得注意时间步嵌入虽然是扩散模型必不可少的条件但它通常不参与交叉注意力而是通过加法或缩放方式注入到每个残差块中因为它表示全局噪声水平不需要精确对应某个空间区域。交叉注意力的参数量和计算量会随条件序列长度线性增长当文本序列变得很长时它也会成为显存优化需要考虑的对象。3. 训练与推理中的注意力瓶颈算力、显存和工程取舍3.1 序列长度决定一切注意力复杂度的数学直觉自注意力最经典的短板就是计算复杂度随着序列长度平方增长。假设输入图片分辨率扩大一倍特征图的长宽都翻倍序列长度就变成原来的四倍注意力矩阵的元素数量直接变成十六倍。这个增长速度是工程上最需要警惕的。以一个具体例子估算在一个批量大小为1的样本里64×64特征图生成4096个词元注意力矩阵需要存储约1600万个元素。如果使用单精度浮点数仅一个注意力矩阵就要占64MB左右再加上多头、多层、反向传播保存的中间变量显存压力会迅速膨胀。一旦分辨率提升到128×128矩阵元素数变成约6700万个显存占用轻松突破数百MB。这也是为什么很多人在训练过程中稍微把分辨率调大就直接遇到“CUDA out of memory”的根本原因——不是显存真的小而是注意力矩阵太能吃了。解决思路不外乎两条路要么压缩序列长度要么把注意力的复杂度从平方级别降下来。压缩序列长度的常见做法包括使用更大的下采样倍数、把相邻像素合并成一个局部块等改良注意力算法则是工程上更主流的做法。3.2 工程上常见的四种减负方案第一种是滑动窗口注意力。把每个位置的注意力范围限制在周围固定大小的窗口内复杂度与序列长度呈线性关系。代价是远距离信息需要用多层窗口堆叠或额外层来传播感受野不如全局注意力直接。第二种是轴向注意力。把二维注意力拆成水平方向和垂直方向两次一维注意力复杂度从序列长度的平方降为序列长度乘以单边长度。这种方案在保持近似全局视野的同时显著减少计算量比较适合分辨率较大的图像任务。第三种是分块局部注意力与全局稀疏采样的混合方案。部分头只在局部块内计算部分头每隔几个位置做一次全局采样既保留一定的全局沟通又避免矩阵爆炸。许多高效扩散模型采用类似设计。第四种是线性复杂度近似。用一个核函数把查询和键先映射成非负特征再改变计算顺序把注意力的结果近似为线性复杂度。这类方法在极长序列下优势明显但精度会有一定损失后面章节再展开讨论。3.3 推理阶段的重复计算与缓存策略扩散模型采样不是算一次就结束而是反复调用去噪网络几十次甚至上百次。注意力的计算结果存在严重的时间冗余相邻两步去噪时图像内容差别不大网络中间层的特征其实高度相似。可惜常见的推理框架并不会自动缓存这些中间注意力结果因为输入是逐步变化的缓存策略需要额外判断哪些层可以安全复用。实践中有人尝试每两步缓存一次中间特征能省不少计算但画面稳定性需要额外验证。更稳妥的节省手段是把注意力计算从训练状态切换到推理状态。训练时为了反向传播需要保存大量中间激活值推理时这些激活值算完就可以立刻释放配合半精度推理注意力层的显存占用可以降到训练时的四分之一以下。另一个经验是条件序列的键、值向量在整次采样中其实基本不变完全可以把条件编码器的输出提前算好并缓存避免每个采样步都重新计算一遍文本编码。这一步在工程实现中收益非常直接特别是在条件文本特别长的时候。实践中还有一个容易被忽略的调参选项减少采样步数。加速采样器让模型可以在二十步以内完成生成每一步的注意力计算量虽然没变但总计算次数大幅下降。注意力层数多、特征图尺寸大的模型配合加速采样器通常能获得比“换更省内存的注意力模块”更明显的提速效果。3.4 训练稳定性和注意力层参数初始化注意力机制虽然强大但也号称“训练不稳定重灾区”。如果初始参数设置不当注意力分数矩阵在训练早期就可能变成极端分布导致梯度非常稀疏U型网络的深层几乎没有有效更新。常见对策是确保注意力输出层在初始化时接近恒等映射线性投影的方差要按维度缩放处理同时网络整体采用残差连接与层归一化的组合。训练预热也对稳定注意力层有明显帮助。扩散模型一般先训练低分辨率再逐步切到高分辨率切分辨率时原本在低分辨率下已经稳定的注意力分布会因为序列长度拉长而重新出现方差变化。这个阶段特别容易损失突然飙升。我自己的经验是每次提升分辨率后把学习率临时降到原来的五分之一再逐步恢复可以有效避开注意力层重新适配带来的震荡。4. 注意力图不只是黑盒可视化、调试与条件控制4.1 把注意力图变成可诊断信号注意力矩阵除了用于前向计算本身就是一个现成的诊断工具。如果你想搞清楚“模型为什么没有按照提示词来画”可以直接取出交叉注意力层的权重矩阵把每个文本词元对应的注意力响应缩放成空间热力图覆盖到生成结果上观察。最常见的现象是文本遗漏提示词里有三件物体模型最终只画了两件。这时去看交叉注意力图大概率会发现被漏掉的那个词元对应的热力图非常弥散没有集中在任何区域说明模型在去噪过程中始终没有把那个词元映射到具体位置。另一种现象是注意力“粘连”两个不同物体共享了同一片高激活区域生成的图像就会出现错乱——比如“猫”和“狗”两个词元的热力图重叠严重结果画出一个四不像。这类问题排查方向一般有两个一是条件文本编码器本身没有把两个词元区分开二是模型中注意力头的数量不够没有足够容量分别处理多个物体的独立语义。空间注意力热力图同样可以诊断。空间注意力体现了模型对某个位置的“自相关性判断”如果注意力图杂乱无章、每个位置只关注自己周围的几个点说明模型网络结构还没有被完全训练好或者当前时间步的噪声水平过高全局结构信息还没有浮现出来。反过来如果注意力图过于集中在少数几个点又会让细节区域缺少必要的上下文交互生成纹理容易变得单调。4.2 把注意力图变成可编辑的画笔注意力图本身也是控制生成内容的有力工具。例如想要让生成过程保持某张参考图的布局可以在每步去噪时提取参考图像的交叉注意力图或空间注意力图把它注入到当前生成过程的对应层强制新图像的区域关注关系与参考图一致。这类做法在对结构一致性要求高的编辑任务里很常见核心优势是不需要额外训练。另一种用法是直接修改注意力权重。在生成中途如果你想微调某个物体的空间位置可以把对应词元的注意力热力图整体往某个方向平移再继续采样。不过这种操作要克制因为注意力权重的整体归一化会连带到其他词元粗暴修改容易导致画面崩毁一般只改某一层、某几个注意力头并且配合较低的注入强度。注意力图也可以用于图像修复。先让模型根据提示词生成完整图像同时记录下目标区域在空间注意力层里的交互关系随后在原图上替换目标区域内容时可以让新内容尽量复用之前的注意力分布维持和周围环境的一致性。这类思路的好处是无需重新训练整个模型适合快速验证效果。4.3 无分类器引导与注意力分布的相互作用无分类器引导是现在扩散模型里最普遍的条件增强手段。做法很简单同时预测有条件和无条件两种噪声再把两个结果拉开距离从而放大提示词的影响。这个增强在注意力层面带来一个副作用——它会同步放大注意力分布的差异。严格来说无分类器引导计算的是“有条件输出减去无条件输出的差值再乘以一个缩放系数后加回去”。注意力模块的输出包含在整体网络输出里所以注意力分布也被间接放大了。放大过头时本来温和的语义关联会变成过度锐化的强绑定生成结果在物体边缘容易出现伪影和色彩溢出。这不是注意力模块本身的Bug而是引导系数过高导致的条件信息过曝。解决手段通常是在调制时对注意力权重做软化处理或者在浅层和深层分别使用不同的引导强度。浅层主要控制整体布局深层主要控制局部语义细节分开设置可以让注意力在不同阶段被更合理地放大。如果你发现用无分类器引导后图像对比度异常、局部纹理发糊可以先试着降低引导系数再观察交叉注意力热力图是否回归到合理分布。5. 新结构冲击下的注意力纯Transformer扩散架构与线性替代5.1 从U型网络到“全注意力”扩散架构近年的一个明显趋势是把扩散模型从“卷积为主、注意力为辅”变成“注意力堆叠为主”的结构。在这种思路里输入图像先被切成固定大小的块每个块展开成一个词元整个图像变成一段词元序列。每一层都在序列上做全局注意力再接一个前馈网络完全不依赖卷积。这类设计的优势非常直观每一层都拥有全局视野结构信息的建模不再受卷积感受野限制同时图像的分辨率不再通过U型网络的逐步下采样来适应而是通过改变块大小来调整序列长度。它让扩散模型的骨干网络变成了一个更通用、更贴近大语言模型经验的架构跨任务迁移和缩放规律都更好套用。但它也有代价。全注意力结构往往需要更大的数据和更长的训练时间才能收敛因为局部先验被大大削弱了。卷积天然自带的平滑性和邻域相关性在纯注意力架构里需要模型自己从数据里学。假如你的数据量不大或者只是做小规模的风格化生成传统的卷积与注意力混合骨架往往更容易出效果而数据规模足够大、目标分辨率足够高时全注意力结构的上限通常更高。5.2 线性注意力与长序列的妥协全注意力结构的痛点回到老问题序列长度平方级增长。当处理视频或三维体素数据时序列轻松上万标准注意力矩阵几乎无法计算。线性注意力方案用核函数的思路把查询和键先映射到高维特征空间再调换矩阵乘法的结合顺序让复杂度从平方降低为线性。理论上线性注意力很好但它有两个明显的工程麻烦。第一是Softmax归一化被核函数近似替代之后注意力权重不再满足行和为1的性质需要额外的归一化修正。第二是线性注意力对“局部强相关”的建模能力偏弱容易把特征过度平滑。我的实测经验是对256×256左右的中低分辨率图片线性注意力的收益不明显还会引入精度损失但面对长序列的视频生成任务时线性方案是性价比极高的选择。另一个长序列方案是选择性状态空间模型它用一个隐状态来压缩序列信息复杂度也与序列长度成线性关系。它在某些任务上的表现接近注意力但推理速度更快。实用心态来看这两种方向并不冲突短序列、高精度场景用标准注意力长序列、实时性要求高的场景用线性替代还可以把二者混合在同一模型里浅层用局部注意力深层用全局注意力或线性注意力。5.3 我目前的工程选型经验如果让我给后入场的人一个可执行的参考建议会是这样数据量和算力一般目标分辨率不超过512×512优先采用卷积与注意力混合的U型骨架只在低分辨率层放少量全局注意力头训练稳定性和速度都能兼顾。数据量充足目标分辨率在1024×1024或更高考虑全注意力堆叠结构图片分块后直接序列建模配合逐步提高分辨率的训练策略最终细节和全局一致性都更好。做视频、长图、三维数据尽快把部分注意力层替换为局部窗口注意力或线性注意力模块否则序列长度会拖垮训练吞吐。调试阶段无论用哪种骨架我都会先做一次交叉注意力可视化。这一步比看损失曲线更直观结构化问题在注意力热力图上一眼就能看出来。模型里的注意力层数量并不是越多越好。大部分需求在最底层的两三处空间注意力加上交叉注意力就够了盲目增加注意力头的数量只会放大显存占用而不会成比例地提升生成质量。我自己这几年的体会是注意力机制的使用原则可以归结为六个字给在关键位置上。扩散模型的能力边界很大程度上由全局一致性决定而全局一致性并不需要每一层都做全局搜索。把注意力花在最能影响结构信息的尺度上同时在边缘细节上保留卷积的计算效率这样得到的模型既稳定又省显存也能让后续调试变得轻松不少。如果非要再提一个落地小技巧那就是每次改完注意力相关配置后一定记得对比同一张种子图在不同配置下的注意力热力图模型的每一次变化几乎都能在那里找到线索。方法不难关键是养成这个习惯。