为反向扩散步骤设计神经网络涉及创建一个模型,该模型可以通过学习目标分布的结构来迭代地消除数据中的噪声。核心思想是训练一个网络来预测去噪过程中每一步的噪声或干净数据。这通常使用 U-Net 架构完成,该架构非常适合通过其带有跳跃连接的编码器-解码器结构来捕获局部和全局上下文。该网络接收噪声输入和时间步长(指示扩散过程中的当前步骤),并输出噪声或干净数据的估计值。关键组件包括残差块、注意力机制和时间嵌入层,以根据当前扩散步骤调节模型。
U-Net 的编码器在减少空间分辨率的同时增加特征深度,从而捕获高级模式。然后,解码器通过上采样并将特征与来自编码器的跳跃连接相结合来重建数据,从而保留精细的细节。例如,在图像生成中,每个块可能由卷积层、组归一化和 SiLU 激活组成。时间步长条件通过将时间步长投影到嵌入向量中来实现,该嵌入向量被添加到每一层的特征图中。注意力层,例如自注意力或交叉注意力,通常插入到 U-Net 的中间以对长程依赖性进行建模。这种设置允许网络根据噪声水平(时间步长)调整其行为,并专注于数据中的相关结构。
实际考虑包括平衡模型容量与计算效率。例如,较小的 U-Net 可能足以处理低分辨率数据,而高分辨率任务需要更深层的架构或诸如检查点之类的技术来管理内存。训练涉及最小化预测噪声和实际噪声之间的损失函数(例如,均方误差)。为了提高稳定性,通常使用诸如梯度裁剪或模型权重的指数移动平均 (EMA) 之类的技术。例如,在 DDPM(去噪扩散概率模型)中,网络预测每一步的噪声分量,并且损失在时间步长上均匀加权。开发人员还应该试验用于添加噪声的方差时间表(线性、余弦),因为这会影响网络如何逐步学习去噪。在扩展到复杂数据之前,使用简化的数据集(例如,MNIST)进行测试有助于验证设计。