自注意力机制可以通过嵌入到负责去噪的神经网络架构中来整合到扩散过程中。在扩散模型中，通常使用 U-Net 在生成过程的每一步预测并去除噪声。自注意力层被添加到 U-Net 中，以帮助模型捕获卷积层单独可能无法捕捉的数据中的长程依赖关系。例如，在图像生成中，自注意力机制允许像素或特征图区域直接相互影响，无论它们在空间上有多接近，从而使模型能够在诸如人脸或物体之类的结构中保持全局连贯性。
这种整合涉及在 U-Net 的特定分辨率下插入自注意力层。例如，在将输入下采样到可管理的图像尺寸（例如，16x16 或 8x8）后，将自注意力应用于展平的特征图。特征图中的每个位置都计算相对于所有其他位置的注意力得分，并按其重要性加权。这个过程有助于模型理解遥远元素之间的关系，例如确保生成的人脸具有对称性，或对齐图像中的纹理。为了降低计算成本，自注意力通常应用于较低分辨率，并且可能会使用分组注意力或稀疏注意力模式等技术。此外，时间步长嵌入（表示当前的去噪步骤）会被注入到网络中，通常通过调节注意力权重或特征图来实现，确保模型在扩散步骤中调整其行为。
一个实际的例子是 Stable Diffusion，它在其 U-Net 中使用自注意力来处理空间关系。该模型在下采样后的中间层应用自注意力，平衡了局部细节处理（通过卷积）和全局上下文建模。例如，在生成狗的图像时，自注意力机制可能会将尾巴的位置与身体连接起来，即使它们在特征图中相距很远。这种整合通过结合卷积操作（局部模式）和注意力机制（全局结构）的优势，提高了模型生成连贯、高质量输出的能力。开发者可以通过将自注意力模块添加到现有的 U-Net 代码库中来实现这一点，确保其与扩散框架的噪声预测和训练循环兼容。