说起扩散模型，绕不开的一个词就是自注意力机制。我在过去两年里接触过好几个图像生成类的项目，每次遇到生成结果“局部像样、整体散架”的毛病——多物体空间关系错乱、前景和背景之间缺乏呼应、重复结构满天飞——最后排查来排查去，问题几乎都出在注意力层的设计上。要么是注意力插错了位置，要么是头数和维度没配合好，要么是训练时没有给足计算资源。这篇就把扩散模型中的自注意力机制从原理、实现、布局、优化到调参经验一次讲清楚。刚入门的朋友可以照着配置跑通一条基线，已经在用的老手也能从里面找到几条值得试的优化思路。
1. 为什么扩散模型离不开自注意力机制
1.1 卷积的局部瓶颈与全局感受野的需求
主流扩散模型的主干大多还是类 U-Net 结构，卷积层负责提取局部特征。卷积天生的缺陷是感受野有限，虽然可以通过堆叠层数来扩大，但每一层看到的仍然是“周围一小块”。这在去噪过程中会带来一个很实际的问题：生成一张包含多个人物的图片时，左侧人物和右侧人物的比例关系、遮挡关系、远近关系，都是跨越大半个画幅的依赖，卷积层很难单独建立起这种长距离约束。
我做过一个直观对比实验：同一个去噪网络，只去掉所有注意力层，FID 指标其实下降得不算特别夸张，但生成图片里“对象数量不对”“肢体互相穿模”的现象明显增多。这说明注意力机制影响的不是低层纹理质量，而是全局结构的一致性。卷积负责“把每一块画好”，自注意力负责“让所有块之间互相商量好”，两者缺一不可。
1.2 自注意力在去噪过程中的具体职责
扩散模型的去噪过程并不是均匀地修复图像。早期时间步对应高噪声，网络必须先确定大致的构图和主体布局；后期时间步噪声小，网络主要在填补细节。在实际训练中你会发现，早期去噪阶段自注意力的注意力图信息量最大、变化也最剧烈，到了后期注意力图会逐渐稳定下来，几乎只在高语义区域之间保留较强连接。
这说明注意力机制的主要职责是“全局关系建模”。具体来说有三类关系非常重要：一是空间关系，比如某物体应该出现在另一物体的上方还是下方；二是尺度关系，比如远处的人脸应该小一些，近处的人脸应该大一些；三是语义一致性，比如两只眼睛应该保持对称、同一条裙子在不同位置应该保持统一纹理。这些都是卷积难以单独建模的，但自注意力可以通过“每个位置去看所有位置”的方式一步到位。
1.3 自注意力与交叉注意力如何分工
提到扩散模型里的注意力，很多人会把自注意力和交叉注意力混在一起。实际工程中两者处理的对象完全不同。自注意力的 Q、K、V 都来自同一个图像特征，它解决的是图像内部的关系；交叉注意力的 K、V 来自文本或类别条件，Q 来自图像特征，它解决的是“图像内容是否符合外部提示”的语义对齐。
在文生图模型里，两者的分工很清晰：交叉注意力负责“语义绑定”，比如提示词里有“红色的雨伞”，网络会把雨伞特征和红色特征在对应区域对齐；自注意力负责“空间一致”，比如两把雨伞之间不能重叠、远近大小要合理。只做交叉注意力不做自注意力，会出现物体位置错乱；只做自注意力不做交叉注意力，会出现图像自然但语义完全不符合提示词。所以你在看模型结构时，如果只有交叉注意力没有自注意力，基本可以断定这个框架在空间一致性上会出问题。
2. 自注意力模块核心实现拆解
2.1 QKV 与注意力分数的数学本质
自注意力的输入是一个特征序列，对图像来说就是把H x W的特征图展平成N = H*W个 token。每个 token 的向量经过三个线性投影得到 Q、K、V。你可以把 Q 理解成“我现在想找什么”，K 是“每个位置声称自己有什么”，V 是“每个位置实际能提供的特征”。通过计算 Q 和所有 K 的点积相似度，再做一次 softmax 归一化，就得到了每个位置应该从所有其他位置聚合多少信息，最后用这个权重去加权求和 V。
整条公式可以直接写成：
Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d)) @ V其中d是每个头的维度。这里为什么要除以sqrt(d)是一个高频考点：如果向量维度很大，点积结果的方差也会变大，softmax 会很快进入饱和区，导致绝大多数注意力权重接近 0，梯度很难传回去。除以sqrt(d)之后，点积结果的方差被拉回到 1 左右，softmax 曲线仍然处于有效梯度区间。
2.2 多头机制为什么能提升表达能力
多头注意力做的事情是把 Q、K、V 在特征维度上切分成多个组，每组独立计算注意力，最后再拼回去。直观理解是每个头可以关注不同类型的关系：一个头可能主要按空间距离分配权重，一个头可能专门抓颜色相似性，一个头可能对物体语义更敏感。相比单头注意力把所有关系混在一组权重里，多头机制相当于并行运行多套“关系探测器”。
但头数并不是越大越好。我在实际调参中发现，当每个头的维度低于 32 时，表达容量会明显不足，注意力图看起来“毛刺感”很强；当头数超过 16 时，边际收益开始下降，显存和计算开销却继续线性增长。一个比较稳妥的默认配置是dim_head=64、heads=8，也就是输出维度 512。如果你模型很大，可以提高到 16 个头，但每个头维度不建议低于 64。
2.3 归一化、残差与时间步条件怎么放
现在主流扩散模型里，注意力模块基本采用“Pre-LN”结构，也就是先做 LayerNorm，再进注意力，最后通过残差连接加回原始输入。这个顺序比 Post-LN 稳定得多，尤其当网络加深以后，Pre-LN 能避免梯度信号在归一化之前就被削弱。典型的模块形态是：
h = x + Attn(LayerNorm(x))残差连接不能省。扩散模型网络往往很深，没有残差时自注意力层的梯度很容易在多层传递中衰减或者爆炸。我见过有人把注意力模块里自带的残差和 U-Net block 的残差重复叠加，导致训练时特征尺度越来越大，最后损失曲线一路飙高。建议注意力模块内部只输出注意力结果，残差统一放在外层管理，避免重复加。
时间步条件一般不直接送进自注意力层。最常见的做法是在每个 ResBlock 里用时间步 embedding 生成一组 scale 和 shift，对 GroupNorm 后的特征做仿射变换，即 FiLM 调制。注意力层前面接收到的已经是携带了时间信息的特征，不需要再多加一路条件。如果你用的是偏 Transformer 的结构，也可以用 AdaLN 的方式，在每个 block 内对归一化后的特征做类似的调制。
2.4 可直接落地的 PyTorch 实现
下面这个实现可以作为你自己模型的起点。它支持普通 softmax 注意力和简化版线性注意力切换，输入输出都是(batch, sequence_length, channels)，图像特征需要在传入前展平。
import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, use_linear=False): super().__init__() self.heads = heads self.dim_head = dim_head inner_dim = heads * dim_head self.scale = dim_head ** -0.5 self.use_linear = use_linear self.norm = nn.LayerNorm(dim) self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) self.to_out = nn.Linear(inner_dim, dim, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x): b, n, c = x.shape inner_dim = self.heads * self.dim_head x = self.norm(x) q, k, v = self.to_qkv(x).chunk(3, dim=-1) q = q.reshape(b, n, self.heads, self.dim_head).transpose(1, 2) k = k.reshape(b, n, self.heads, self.dim_head).transpose(1, 2) v = v.reshape(b, n, self.heads, self.dim_head).transpose(1, 2) if self.use_linear: # 简化版线性注意力，用 ReLU 核函数做近似，复杂度 O(N) q = F.relu(q) k = F.relu(k) q = q / (q.sum(dim=-1, keepdim=True) + 1e-6) k = k / (k.sum(dim=-1, keepdim=True) + 1e-6) kv = torch.einsum("bhnd,bhme->bhde", k, v) out = torch.einsum("bhnd,bhde->bhne", q, kv) else: dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale attn = F.softmax(dots, dim=-1) attn = self.dropout(attn) out = torch.matmul(attn, v) out = out.transpose(1, 2).reshape(b, n, inner_dim) return self.to_out(out)调用时的注意点：x必须已经是三维的 token 序列；如果你的特征图是(B, C, H, W)，需要先permute再reshape；残差连接放在模块外层，模块内部不包含残差。训练时如果开了 fp16，标准的matmul + softmax路径在部分显卡上容易溢出，建议直接用F.scaled_dot_product_attention接管这一段的计算，或者至少在注意力内部使用 fp32。
3. 注意力层在 U-Net 里的布局与配置
3.1 该在哪几个分辨率插入注意力
很多同学拿到 U-Net 就直接往所有层级加注意力，结果显存爆炸，生成质量却没有提升。这里要先搞清楚复杂度：全局自注意力的计算量和显存占用与 token 数 N 的平方成正比，而 N 是特征图的H * W。分辨率翻一倍，注意力开销就变成原来的四倍。所以工程上的通用策略是只在较粗糙的特征图上使用全局注意力。
以 64x64 输入为例，U-Net 下采样过程中典型的分辨率是 64、32、16、8。我通常只在 32x32 和 16x16 这两个尺度加全局自注意力，8x8 的尺度由于 token 本身只有 64 个，加了也不贵，但信息收益有限，可加可不加。64x64 那一层基本上不建议加全局注意力，token 数 4096 会让计算量直接压垮显存，而且从实际效果看，这一层注意力对最终结构一致性的贡献很小，原因是高分辨率特征主要处理高频细节，全局关系在低分辨率层已经确立过了。
如果你训练的是 256x256 甚至更高分辨率，情况会更激进。一般只会在 32x32 和 16x16 层做全局注意力，更高分辨率分支如果要加注意力，应该用窗口注意力而不是全局注意力。把特征图切成8x8或16x16的窗口，每个窗口内部做局部自注意力，窗口之间不交互。这样可以保留注意力的自适应聚合能力，又把计算量控制在可接受范围。
3.2 头数、通道数与窗口的取舍
通道数的选择基本上跟着 U-Net 各层的宽度走。假设你有一个 channels 为[128, 256, 256, 512]的小模型，32x32 特征层是 256 通道，16x16 特征层是 512 通道。每个层插入注意力时，直接用这一层现有的通道数作为注意力投影的输入输出维度，不要额外扩宽通道，否则参数量翻倍，训练难度增加，收益却不明显。
头数和dim_head的搭配遵循我前面说的规则：每个头维度 64 左右，头数在 8 到 16 之间。如果某一层通道数只有 128，硬要拆 8 个头，每个头只有 16 维，基本学不到有效关系；这种情况下建议降低头数到 4，或者接受更小的dim_head。窗口注意力里窗口大小也要根据分辨率调整，32x32 特征层用8x8窗口就是 4x4 个窗口，16x16 特征层用16x16窗口就退化成全局注意力，所以要用8x8窗口保留一点窗口间交互的需求。
3.3 一个可复现的基线配置示例
下面是我常用的一份小型扩散模型注意力配置，可以作为你第一版基线的起点。
| 特征层分辨率 | 是否加注意力 | 通道数 | 头数 | 每个头维度 | 说明 |
|---|---|---|---|---|---|
| 64x64 | 否 | 128 | - | - | 只使用卷积，保留局部高频 |
| 32x32 | 是 | 256 | 8 | 32 | 全局注意力，token 数 1024 |
| 16x16 | 是 | 512 | 8 | 64 | 全局注意力，token 数 256 |
| 8x8 | 可选 | 512 | 8 | 64 | token 数少，开销低 |
在这个配置下，我试过把 32x32 层的 8 个头每个头改到 64 维，也就是把注意力输出维度扩到 512，结果显存增加比较明显，但结构一致性指标提升很小。反过来把 16x16 层头数从 8 减到 4，多物体生成时就开始出现物体关系不稳定的情况。所以这份配置在高性价比和生成质量之间算是比较平衡的点。
4. 显存与计算优化：从全局到可落地
4.1 复杂度瓶颈的量化分析
要理解注意力为什么“吃显存”，最好做一次手算。假设特征图是32x32，展平后 N=1024。一个多头注意力层中，最占显存的结构是注意力权重矩阵(batch, heads, N, N)。如果batch=8、heads=8、每个元素用 fp16 保存，这个矩阵就有8 * 8 * 1024 * 1024 = 67108864个元素，每个元素 2 字节，单层约 128MB。训练时反向传播往往还需要保留注意力权重用于梯度计算，如果开启梯度检查点，还可以通过重算来省显存，但默认情况下这块开销是实打实的。
计算量同样不能忽略。标准的注意力计算量大约在4N d^2 + 2N^2 d，其中前半部分是 QKV 投影和输出投影，后半部分是注意力矩阵计算。N=1024、d=512 时，2N^2 d大概是 10 亿次乘加操作，放在单层上问题不大，但 U-Net 里如果加三四层注意力，再乘上扩散模型常用的较大 batch，训练速度会肉眼可见地下降。这也是为什么高分辨率层插全局注意力非常不划算。
4.2 Flash Attention 与线性注意力的工程选型
工程上首选的是 Flash Attention 这类 IO 感知的注意力实现。它核心思路是不把完整注意力矩阵一次性写入显存，而是分块计算、分块存储，配合反向传播时重算，把显存开销从 O(N^2) 降到 O(N)，计算速度也有提升。实际使用中最简单的接入方式是调用F.scaled_dot_product_attention，PyTorch 的 SDPA 会在支持的时候自动选择 Flash Attention 后端，代码几乎不用改。
线性注意力是另一个方向，通过核函数把 softmax 近似成phi(Q) @ phi(K)^T @ V的形式，把计算复杂度从 O(N^2) 降到 O(N)。我在超分辨率模型里用过简化版线性注意力，速度确实快，但生成质量会有轻微下降，特别是细节纹理容易出现“模糊化”的倾向。如果你对质量要求很高，还是优先用 Flash Attention；只有当显存实在不够、或者需要处理超长 token 序列时，再考虑线性注意力作为降级方案。
4.3 混合精度与显存估算实操
训练扩散模型时混合精度几乎是标配，但注意力部分要特别小心。fp16 的表示范围有限，注意力矩阵里的点积结果很容易溢出到非正规数，导致梯度消失或者 NaN。我遇到过好多次训练跑到几千步突然出现 loss 尖刺，把注意力相关计算强制切到 fp32 之后问题就消失了。如果你的显卡支持 bf16，直接全程用 bf16 是最省心的选择，它比 fp16 稳很多，几乎不需要额外处理。
显存估算有一个经验公式。单个全局注意力层在训练时的显存开销大致是batch * heads * N * N * 2字节，原因是除了注意力矩阵本身，Q、K、V 和输出投影的中间变量也占不少空间，但注意力矩阵是最大头。你可以用这个公式倒推自己的 batch 能开多大。比如想控制在 2GB 以内，N=1024、heads=8、fp16 时，batch * 8 * 1024 * 1024 * 2 <= 2 * 1024^3，算出来 batch 最多 128，看起来很大，但这是单层单卡的理论值，实际加上卷积和优化器显存，能开的 batch 要小得多。因此最实用的做法是先把模型结构跑一个最小 batch，观察显存占用曲线，再按比例反推最大 batch。
5. 训练稳定性与常见故障排查
5.1 训练尖刺、损失发散怎么办
自注意力层的梯度常常是 U-Net 里最“暴躁”的部分。症状很典型：整体 loss 在平稳下降时，突然某一步跳到原来的好几倍，然后慢慢回落，严重时直接变成 NaN。排查顺序建议从三方面入手。第一，检查混合精度，把注意力相关的matmul和softmax强制用 fp32 计算，这是最简单也最高频的解决方案。第二，检查梯度裁剪，默认clip_grad_norm设为 1.0 起步，如果还出现尖刺就降到 0.5，代价是训练速度略慢。第三，检查学习率，注意力模块的更新幅度如果太大，可以给整个模型设置一个比较低的学习率，比如从 2e-4 降到 1e-4，再用 warmup 长一点来补偿。
我个人的习惯是：新项目第一次训练时，先跑一个 3 到 5 千步的 smoke test，专门盯 loss 曲线里有没有尖刺。如果有，第一时间切 fp32 attention，而不是反复调学习率。大多数不稳定问题都不是学习率不合适，而是数值精度和初始化尺度的问题。
5.2 生成重复结构与伪影的定位
生成结果出现周期性重复纹理、网格状伪影或者同一物体被复制多份，往往和注意力层的位置以及位置编码有关。全局注意力本身有能力建立长距离依赖，但如果训练数据里某一类结构重复出现，注意力图可能会收敛到一种“周期模式”，导致生成时不断重复相同结构。
遇到这类问题我一般做两个实验。第一个实验是减少高分辨率层的注意力，因为高分辨率层注意力更容易记住高频周期模式，去掉或换成窗口注意力能显著缓解重复现象。第二个实验是检查位置编码，如果模型完全没有位置信息，自注意力会把空间位置差异忽略掉，特征图上相近的结构就会被当成同一个对象处理。给 token 加上正弦位置编码或者可学习位置编码之后，重复伪影通常会有明显改善。
另外，推理阶段的无分类器引导强度过高也会放大重复伪影。引导强度从 7 调到 12 以上时，模型会更激进地追逐训练分布里的高频模式，重复现象会变重。如果你训练时已经确认模型本身没问题，可以先降一降推理引导强度试试。
5.3 用注意力图做可视化诊断
注意力图可视化是排查问题的利器。实现不复杂，在前向计算里把attn权重拿回来，形状是(batch, heads, N, N)，然后针对某个特征层做平均或取最大值，再 reshape 成空间维度。
# attn 形状: (batch, heads, N, N) # N = H * W，特征图为 H x W attn_map = attn[0].mean(dim=0) # 对 batch 内第一张图、所有头取平均 attn_map = attn_map.reshape(H, W, H, W) # 查看位置 (i, j) 的注意力来源 weight = attn_map[i, j] # 形状 (H, W)把weight画出来就能看到某个位置在关注全图的哪些区域。健康的注意力图应该有明显的语义边界：比如在生成人脸时，左眼的 query 会强烈关注右眼、鼻子和脸部轮廓。如果注意力图几乎均匀分布，说明这一层没有学到有效关系；如果注意力图集中在一个很小的局部区域，说明该层退化成了某种“局部卷积”，自注意力的全局建模能力没有发挥出来。我还会对比不同 head 的注意力图，如果所有 head 的分布几乎一样，说明头部多样性不足，可以考虑调整初始化或者增加dim_head。
6. 效果影响与架构演进
6.1 注意力层贡献的评估方法
很多同学只关心 FID，但 FID 对结构一致性并不敏感，有时候注意力层删掉一两层，FID 变化很小，实际看图却已经塌了。我建议做注意力层消融实验时，同时记录三组指标：FID、LPIPS、以及你任务里最关心的结构指标（比如目标检测准确率、分割 mIoU）。一组典型的对比是：无注意力、只有低分辨率注意力、多分辨率注意力。测试结果通常会发现 FID 差异不大，但结构类指标差异明显。
在图像编辑和可控生成场景里，注意力层的价值更突出。自注意力图本质上编码了对象之间的对应关系，利用这个特性可以做很多有意思的事：比如设计空间注意力损失来约束不同生成样本之间的结构一致性，或者通过替换注意力图来保持多个样本之间的空间布局。如果你后续要做图像编辑、视频生成或者多视角一致性的工作，提前把注意力图的可视化和导出接口写好，会节省大量调试时间。
6.2 从 U-Net 到 DiT：自注意力成为主菜
近两年架构演进的一个明显趋势是：自注意力从 U-Net 的辅助模块，变成了生成模型主干的核心算子。DiT 一类的方法直接把图像切成 patch，丢进纯 Transformer 的 block 序列里，每个 block 都包含自注意力、条件调制和前馈网络。相比 U-Net 里只在特定分辨率插入注意力层的做法，DiT 的自注意力在整个网络深度上更密集，全局关系建模的能力更强。
如果算力允许，我现在会优先建议新项目考虑 DiT 方向的网络，而不是在 U-Net 里继续堆注意力层。原因很简单：U-Net 里的自注意力只是“局部增强”，纯 Transformer 里的自注意力才是完全体，尤其在生成高分辨率图像、视频这类强结构需求任务上，DiT 的扩展性明显更好。具体配置上，patch size 越小 token 越多，模型表达能力越强但显存压力越大；较深的 DiT 一般会用 AdaLN 做条件注入，让模型 block 都能感知时间步和类别条件。自注意力机制本身是同一个东西，只是从配角变成了主角，理解好它，迁移到任何新架构里都能派上用场。
我个人现在的习惯是：接到一个生成任务，先画一张“分辨率-注意力分布表”，想清楚哪些层用全局注意力、哪些层用窗口注意力、哪些层完全不用，再动手写模型。如果算力允许，优先把主干往 DiT 方向推，让自注意力更密集地发挥作用。如果你暂时还留在 U-Net，那至少记住一个原则：注意力不是越多越好，而是要放在最关键的低分辨率结构层，并且尽早把头数、维度、归一化方式和混合精度验证好，否则后面调参的坑会一个接一个地冒出来。
