UNet 架构

相关源文件

• ldm/models/diffusion/ddpm.py
• ldm/modules/diffusionmodules/openaimodel.py
• ldm/modules/encoders/modules.py
• ldm/modules/x_transformer.py
• ldm/util.py

本文档介绍了 Stable Diffusion 的扩散模型中使用的 UNet 架构。UNet 作为去噪过程的骨干，接收带噪的潜在表示，并根据时间步长和条件信息逐步去除噪声。有关文本如何编码和用于条件的信息，请参阅 文本条件。有关整体潜在扩散模型的详细信息，请参阅 潜在扩散模型。

概述

Stable Diffusion 中的 UNet 是一个时间条件化的 U 形神经网络，它在潜在空间中运行。它在每个扩散时间步长预测添加到图像的噪声或直接预测去噪后的输出。其架构结合了一系列下采样块、一个中间块和具有跳跃连接的上采样块，同时融入了残差块和注意力机制。

来源：ldm/modules/diffusionmodules/openaimodel.py413-742

关键组件

1. 时间嵌入

时间嵌入是使 UNet 能够根据扩散时间步长条件化其行为的关键组成部分。这使得模型能够根据噪声级别应用不同的去噪策略。

时间嵌入过程

1. 使用正弦函数将标量时间步长转换为高维嵌入
2. 将此嵌入通过一个 2 层 MLP 以获得最终的时间嵌入
3. 然后，此嵌入在整个 UNet 中用于条件化各种块

来源：ldm/modules/diffusionmodules/openaimodel.py506-511 ldm/modules/diffusionmodules/util.py18-19

2. 构建块

UNet 架构由几个关键的构建块组成

ResBlock（残差块）

ResBlock 是一个基本组件，有助于梯度流动和特征传播

来源：ldm/modules/diffusionmodules/openaimodel.py163-275

注意力块

注意力块允许不同的空间位置相互关注

来源：ldm/modules/diffusionmodules/openaimodel.py278-324

空间变换器（交叉注意力）

空间变换器通过交叉注意力实现基于文本嵌入的条件化

来源：ldm/modules/diffusionmodules/openaimodel.py557-559 ldm/modules/attention.py

3. UNet 架构组件

下采样（输入块）

下采样路径由一系列输入块组成，这些块逐渐减小空间尺寸同时增加通道深度

对于每个下采样级别

1. 多个 ResBlocks 在当前分辨率下处理特征
2. 在指定分辨率下添加注意力层
3. 通过跨步卷积或平均池化进行下采样

来源：ldm/modules/diffusionmodules/openaimodel.py516-587

中间块

中间块处理瓶颈表示

来源：ldm/modules/diffusionmodules/openaimodel.py597-623

上采样（输出块）

上采样路径通过跳跃连接结合了来自下采样路径的特征

上采样过程镜像了下采样路径

1. 来自下采样路径相应级别的特征被串联起来
2. 多个 ResBlocks 处理组合后的特征
3. 通过最近邻插值然后卷积进行上采样

来源：ldm/modules/diffusionmodules/openaimodel.py626-686

4. 前向传播

UNet 的前向传播将所有组件连接起来

来源：ldm/modules/diffusionmodules/openaimodel.py710-742

配置参数

UNet 架构是高度可配置的。以下是关键参数

参数	描述	默认值
image_size	输入图像的大小	64（潜在空间）
in_channels	输入通道数	4
model_channels	基本通道数	320
out_channels	输出通道数	4
num_res_blocks	每个下采样级别的 ResBlocks 数量	2
attention_resolutions	应用注意力的分辨率	(4, 2, 1)
dropout	Dropout 率	0.0
channel_mult	每个级别的通道乘数	(1, 2, 4, 8)
conv_resample	使用学习到的卷积进行采样	True
dims	信号的维度	2
num_heads	注意力头的数量	8
num_head_channels	每个注意力头的通道数	64
use_scale_shift_norm	使用 FiLM 式条件化	True
use_spatial_transformer	使用交叉注意力进行文本条件化	True
transformer_depth	交叉注意力 Transformer 的深度	1
context_dim	上下文（文本）嵌入的维度	768 或 1024

来源：ldm/modules/diffusionmodules/openaimodel.py443-468

与扩散过程的集成

UNet 模型通过 DiffusionWrapper 类集成到扩散过程中，该类在 DDPM 类中初始化。

集成工作如下

1. DDPM 类创建一个包装 UNet 模型的 DiffusionWrapper 实例
2. 在训练过程中，会生成带噪的样本并将其传递给 UNet
3. UNet 根据参数化预测噪声或干净样本
4. 在采样过程中，UNet 会迭代地去噪潜在表示

来源：ldm/models/diffusion/ddpm.py44-116 ldm/models/diffusion/ddpm.py231-265

条件化机制

UNet 支持不同的条件化机制，这些机制决定了辅助信息如何影响扩散过程

无条件

在无条件模式下，模型仅依赖时间步长，没有外部条件。

基于串联的条件化

条件数据直接与输入串联

基于交叉注意力的条件化

文本嵌入通过交叉注意力层对 UNet 进行条件化

来源：ldm/models/diffusion/ddpm.py29-31 ldm/modules/diffusionmodules/openaimodel.py557-559

性能考量

UNet 架构包含几项优化

1. 梯度检查点：通过重新计算中间激活而不是存储它们来减少反向传播期间的内存使用
2. 分辨率缩放：在潜在空间中以降低的分辨率运行（对于 512×512 图像通常为 64×64）
3. 参数效率：在特定分辨率下策略性地使用注意力机制

来源：ldm/modules/diffusionmodules/openaimodel.py315-316 ldm/modules/diffusionmodules/openaimodel.py250-252

相关组件

UNet 架构与 Stable Diffusion 系统中的其他几个组件紧密协作

1. CLIP 文本编码器：为交叉注意力条件化提供上下文嵌入
2. VAE 编码器/解码器：在像素空间和 UNet 运行的潜在空间之间进行转换
3. 扩散调度器：定义扩散过程的噪声调度和采样算法
4. 条件化机制：确定文本提示如何影响生成过程

来源：ldm/modules/encoders/modules.py137-162 ldm/models/diffusion/ddpm.py424-469