当前位置：首页 > news >正文

扩散模型中的注意力机制实现详解：从Self-Attention到Cross-Attention的代码解析

news 2026/5/12 18:41:24

扩散模型中的注意力机制都是怎么实现的？

一、`Diffusers`中的注意力机制实现

diffusers中U-Net的注意力机制实现主要是UNet2DConditionModel这个类，通过在特定的下采样、中间和上采样阶段加入CrossAttnDownBlock2D类的模块，从而实现注意力层的加入。

这些模块内部的核心单元是Transformer2DModel，它负责把输入的2D特征图转换成序列形式。而在最底层的BasicTransformerBlock中，通常包含了两层注意力：

• 第一层是处理图像内部关系的self-attention（代码中通常叫attn1）；
• 第二层则是融合文本条件的cross-attention（代码中通常叫attn2），这两者在代码层面都由Attention类定义，并将具体的计算逻辑委托给AttnProcessor来执行。

二、Self-Attention的实现

关于self-attention的实现，其核心目的是建模图像不同部分之间的correlation，从而捕捉全局结构上的关联。

在diffusers的官方代码中，self-attention的实现如下：

query = attn.to_q(hidden_states)
encoder_hidden_states = hidden_states
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)

query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)

attention_scores = torch.baddbmm(
torch.empty(…), query, key.transpose(-1, -2), alpha=attn.scale
)
attention_probs = attention_scores.softmax(dim=-1)

hidden_states = torch.bmm(attention_probs, value)

由于是自注意力，query、key和value都是输入图像特征`hidden_states`的投影。大致的逻辑是：首先通过`to_q`、`to_k`、`to_v`线性层将`hidden_states`映射到特征空间，然后通过`head_to_batch_dim`操作将head维度折叠进batch维度，剩下的就是标准attention中的计算了——attention weights、softmax归一化、加权到value上。 从数学角度看，上面代码的过程如下： 1. 1. 形状原本都是，经过`head_to_batch_dim`的操作后变为的张量形式； 2. 2. 计算注意力分数为：，得到的是一个的矩阵； 3. 3. 通过将attention weights整合到value上； 4. 4. 经由`batch_to_head_dim`还原维度，并通过输出投影层，结束self-attention的计算过程。 三、Cross-Attention的实现 ==================== 在代码实现上，cross-attention复用了与self-attention完全相同的`AttnProcessor`逻辑，唯一的区别在于key和value的来源。 在`attn2`的前向传播中，`encoder_hidden_states`不再是None，而是传入了由CLIP编码好的text embedding（记形状为）。 ```plaintext 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 # From image feature query = attn.to_q(hidden_states) # From text encoder key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) query = attn.head_to_batch_dim(query) key = attn.head_to_batch_dim(key) value = attn.head_to_batch_dim(value) attention_scores = torch.baddbmm(..., query, key.transpose(-1, -2), alpha=attn.scale) attention_probs = attention_scores.softmax(dim=-1) hidden_states = torch.bmm(attention_probs, value)

这里的交互过程产生了一个的注意力矩阵（），它表示了图像特征中每一个位置与text prompt中每一个token的关联程度。