稳定扩散v2模型训练数据和过程详解

训练

训练数据：模型开发者使用以下数据集进行模型训练：

LAION-5B及其子集（详见下文）。使用LAION的NSFW检测器进一步过滤训练数据，使用'p_unsafe'分数为0.1（保守）。有关详细信息，请参阅LAION-5B的NeurIPS 2022论文以及有关此主题的审稿人讨论。

训练过程：稳定扩散v2是一种潜在扩散模型，它将自动编码器与在自动编码器的潜在空间中训练的扩散模型相结合。在训练期间，

图像通过编码器进行编码，将图像转换为潜在表示。自动编码器使用相对下采样因子8，并将形状为H x W x 3的图像映射到形状为H/f x W/f x 4的潜在状态中。

文本提示通过OpenCLIP-ViT/H文本编码器进行编码。

文本编码器的输出通过交叉注意力馈送到潜在扩散模型的UNet骨干网络中。

损失是潜在噪声添加到重建目标之间的重建目标和UNet预测之间的重构目标。我们还使用所谓的v目标，参见https://arxiv.org/abs/2202.00512。

我们目前提供以下检查点：

512-base-ema.ckpt：在分辨率256x256上的LAION-5B子集上进行了550k步训练，过滤了明确的色情材料，使用LAION-NSFW分类器，punafe = 0.1，美学得分>= 4.5。在相同的数据集上，分辨率>= 512x512的情况下进行了850k步训练。

768-v-ema.ckpt：从512-base-ema.ckpt中恢复并在相同数据集上使用v目标进行了150k步训练。在我们的数据集的768x768子集上恢复了另外140k步。

512-depth-ema.ckpt：从512-base-ema.ckpt中恢复并进行了200k步微调。添加了一个额外的输入通道，以处理MiDaS（dpt_hybrid）产生的（相对）深度预测，该预测用作附加条件。处理此额外信息的U-Net的附加输入通道被零初始化。

512-inpainting-ema.ckpt：从512-base-ema.ckpt中恢复并进行了另外200k步训练。遵循LAMA中提出的掩膜生成策略，该策略与掩膜图像的潜在VAE表示相结合，用作附加条件。处理此额外信息的U-Net的附加输入通道被零初始化。相同的策略用于训练1.5-inpainting检查点。

x4-upscaling-ema.ckpt：在包含图像> 2048x2048的LAION的10M子集上进行了125万步的训练。该模型在512x512的裁剪上进行训练，是一个文本引导的潜在上采样扩散模型。除了文本输入外，它还接收一个噪声级别作为输入参数，该参数可用于根据预定义的扩散计划向低分辨率输入添加噪声。

硬件：32 x 8 x A100 GPU

优化器：AdamW

梯度累积：1

批次：32 x 8 x 2 x 4 = 2048

学习率：热身到0.0001，进行10000步，然后保持不变

评估结果

使用不同的无分类器引导比例（1.5、2.0、3.0、4.0、5.0、6.0、7.0、8.0）和50个步骤的DDIM采样步骤进行评估，显示了检查点的相对改进：