Stable Diffusion 交叉注意力机制全解5 步看懂文本如何注入扩散模型【免费下载链接】stable-diffusionA latent text-to-image diffusion model项目地址: https://gitcode.com/GitHub_Trending/st/stable-diffusion在 Stable Diffusion潜空间文本生成图像扩散模型中交叉注意力是文本条件进入 U-Net 的唯一通道。读完本文你能看懂 attention.py 里 CrossAttention 的完整数据流理解它如何把一句话翻译成对每个图像位置的特征引导并知道二次开发时该动哪些参数。痛点文本和图像对不上怎么办早期扩散模型的条件注入方式很直接把类别编号或标签向量拼进特征图或做 AdaIN 式的调制。对一张 32 个位置的类别图来说够用但文本是变长序列a fire 和 a watercolor painting of a fire on canvas 长度完全不同直接拼接根本没法实现。更本质的问题是图像特征要逐位置参考文本。画fire该出现在画面哪一格画watercolor则影响整幅画的风格。这就要求一种细粒度的连接方式每个图像位置都要能问一句——这句话里哪些词跟我有关交叉注意力就是为这个问题设计的。先建直觉面试官与候选人打个比方。把图像的每个特征位置想象成面试官文本的每个词想象成候选人。面试官Query不改变自己的身份而是给每个候选人Key按相关性打分再按分数加权吸收他们带来的信息Value。一句话定义交叉注意力 Q 来自一方图像特征K/V 来自另一方文本特征的注意力。最容易混淆的是自注意力Self-AttentionQ、K、V 三者同源序列只关注自己内部交叉注意力则 Q 与 K/V 来自两个不同序列。Stable Diffusion 的 U-Net 里两者并存——自注意力负责画面内部的结构协调交叉注意力负责听文本的话。按数据流逐步拆解先看数据是怎么流动的文本经冻结的 CLIP 编码器变成 token 向量作为context传入 U-Net卷积主干产出的图像特征作为x两者在CrossAttention相遇。下面 5 步跟着走一遍。1. 逐段读懂 Q/K/V 投影三个独立的线性层把两边特征投进同一隐空间。注意context_dim与query_dim可以不同——这正是文本维度768和图像通道数512能共存的秘密。inner_dim dim_head * heads # 8 头 × 64 维 512 self.to_q nn.Linear(query_dim, inner_dim, biasFalse) # Q图像特征 self.to_k nn.Linear(context_dim, inner_dim, biasFalse) # K文本特征 self.to_v nn.Linear(context_dim, inner_dim, biasFalse) # V文本特征Q 描述我要找什么K 描述我能提供什么V 是我的实际内容。三者同维后点积才有意义。2. 多头拆分一次算完 8 组把最后一维拆成(heads, dim_head)并把 batch 与 head 合并成一个伪 batch后面所有运算对 8 个头完全并行零 for 循环。h self.heads # b批, n序列长度, h头数, d单头维度 q, k, v map(lambda t: rearrange(t, b n (h d) - (b h) n d, hh), (q, k, v))多头不是更深的注意力而是让不同头在不同子空间里各自学一种对齐方式——有的头盯主体名词有的头盯风格词。3. 相似度矩阵与 mask 拦截Q 与 K 做点积得到每个图像位置 × 每个词的打分矩阵。CLIP 固定输出 77 个 token短 prompt 的尾部是 pad必须在 softmax 之前屏蔽掉。sim einsum(b i d, b j d - b i j, q, k) * self.scale # 缩放后打分 if exists(mask): max_neg_value -torch.finfo(sim.dtype).max mask repeat(mask, b j - (b h) () j, hh) sim.masked_fill_(~mask, max_neg_value) # pad 位填成极不可能masked_fill_把 pad 位分数打成浮点数能表示的最小值softmax 之后权重趋近 0——pad token 从此对图像零贡献。4. softmax 加权取值文本注入的瞬间这一步就是文本信息真正进入图像特征的路口每个图像位置按权重从所有词里抽取内容。attn sim.softmax(dim-1) # 每个位置对 77 个词的权重和为 1 out einsum(b i j, b j d - b i d, attn, v) # 加权求和 V out rearrange(out, (b h) n d - b n (h d), hh) # 8 头拼回 return self.to_out(out) # 线性层 Dropout投回 512 维结论输出形状与输入x一致但每个位置都混入了它认为最相关的词的内容。5. 放回 U-Net一个 TransformerBlock 的完整走位CrossAttention从不单独工作它和自注意力、前馈层串成一个块全部走残差连接。def _forward(self, x, contextNone): x self.attn1(self.norm1(x)) x # 自注意力图像内部协调 x self.attn2(self.norm2(x), contextcontext) x # 交叉注意力吸收文本 x self.ff(self.norm3(x)) x # 前馈逐位置非线性加工 return x每层先看图、再听文本、后消化 x保证浅层信息不被深层冲掉。实现里的巧思与坑先缩放再点积顺序不能换。点积的方差随维度 d 线性增长d 越大分数绝对值越大softmax 会退化成一个词通吃的近似 one-hot其余词梯度几乎为零。dim_head ** -0.5把方差拉回 1让 softmax 保持灵敏。这就是每个实现里必有那个 scale 的原因。context缺省时的伪装。forward 里有context default(context, x)不传文本时K/V 也用 x 自己算交叉注意力自动退化成自注意力。这一个默认值让同一个类身兼两职也意味着测试时忘了传 context 不会报错——只是文本引导悄悄消失了。出口投影用零初始化。SpatialTransformer的最后一层卷积套了zero_module参数全置零。self.proj_out zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size1, stride1, padding0))训练开始时 Transformer 分支输出恒为 0网络等价于普通 U-Net等卷积主干稳定后文本引导再逐渐接管。这是扩散模型里对抗不稳定的经典手法删掉它训练很容易崩。它在架构中的位置上游是 FrozenCLIPEmbedder把 prompt 编码成 token 向量作为context一路传入 U-Net 的 forward。U-Net 的每个下采样/上采样块都收到context但真正启用交叉注意力的位置在最低分辨率的特征层——openaimodel.py 里的SpatialTransformer把特征图b c h w拉平成长度h*w的序列送进 TransformerBlock再恢复回特征图。下游U-Net 输出噪声预测交给采样器DDIM / PLMS / DPM-Solver见 scripts/txt2img.py采样完成后由 VAE 解码器把潜变量还原成像素。文本自始至终没进过卷积层它只通过context在注意力层被读了 N 次。使用与二次开发建议调参先动采样别动网络。文图不符通常该调 txt2img.py 里的--scaleguidance scale和--ddim_stepsscale过大出伪影过小贴不紧提示词。交叉注意力本身的结构参数V1 配置8 头 × 64 维、transformer_depth: 1、context_dim: 768见 v1-inference.yaml是随权重一起固定的改了也对不上预训练 ckpt。不要对齐 768 和 512。Q 维度由图像通道数决定K/V 维度由文本编码器输出决定二者天生不同靠to_k/to_v解耦。二次开发换文本编码器时只需在 yaml 里改cond_stage_config并把context_dim改成新编码器的输出维U-Net 一行不用动。想加深度先想显存。transformer_depth每加一层交叉注意力对每个 token 都做一遍全序列 softmax低分辨率层 token 数只有 64~256 还好但显存和时延是线性涨的。收束交叉注意力让文本以逐位置查询的方式进入潜空间 U-Net替代了无法处理变长序列的拼接式条件注入Q/K/V 角色分离Q 管图像、K/V 管文本维度靠投影层解耦768 与 512 各走各的两个小零件保住训练稳定点积前的dim_head ** -0.5缩放和出口卷积的零初始化mask 屏蔽 pad token否则 77 个固定槽位里的占位符会污染整张图二次开发看两个数context_dim和transformer_depth其余别碰相关源码交叉注意力与 Transformer 块ldm/modules/attention.pyU-Net 与 SpatialTransformer 注入点ldm/modules/diffusionmodules/openaimodel.py条件编码与潜空间模型ldm/models/diffusion/ddpm.py文本编码器CLIPldm/modules/encoders/modules.pyV1 推理配置configs/stable-diffusion/v1-inference.yaml【免费下载链接】stable-diffusionA latent text-to-image diffusion model项目地址: https://gitcode.com/GitHub_Trending/st/stable-diffusion创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考