【diffusers极速入门（五）】扩散模型中的 Scheduler（noise_scheduler）的作用是什么？

系列文章目录

• 【diffusers 极速入门（一）】pipeline 实际调用的是什么？ call 方法!
• 【diffusers 极速入门（二）】如何得到扩散去噪的中间结果？Pipeline callbacks 管道回调函数
• 【diffusers极速入门（三）】生成的图像尺寸与 UNet 和 VAE 之间的关系
• 【diffusers极速入门（四）】EMA 操作是什么？

引言

扩散模型作为一种强大的生成模型，在图像生成、文本生成等领域取得了显著的成果。

但需要注意，扩散模型不是一个神经网络，而是一套定义了加噪、去噪公式的模型。扩散模型中需要一个去噪模型来去噪，去噪模型一般是一个神经网络¹。

Scheduler 作为去噪和加噪过程中的核心组件，在整个生成过程中扮演着至关重要的角色。本文将深入探讨 Scheduler 的工作原理、作用以及不同类型的 Scheduler。

Scheduler 的概念与作用

• Scheduler，中文译为“调度器”，在扩散模型中负责控制噪声的添加和去除过程。
• 它定义了在每个扩散步骤中，向数据添加多少噪声，以及在去噪过程中如何逐步恢复原始数据。

Scheduler 的主要作用有：

1. 控制噪声强度: 决定在每个扩散步骤中添加多少噪声。
2. 影响模型训练: 不同的 Scheduler 会影响模型的训练速度和生成效果。
3. 提升采样质量: 通过调整 Scheduler 的参数，可以提高生成样本的质量。

Scheduler 的工作原理

前向扩散（推理阶段）

• 从原始数据开始，逐步向数据添加高斯噪声。
• 随着扩散步数的增加，噪声的强度逐渐增大，数据逐渐变得模糊。
• 这个过程是一个确定性的过程，不需要训练。

后向扩噪（训练阶段）

• 训练一个神经网络，这个神经网络的任务就是从带噪声的数据中恢复出原始数据。
• Scheduler 控制着噪声强度，使得神经网络能够从不同程度的噪声数据中学习。
• 在训练过程中，我们计算预测噪声和真实噪声之间的损失，并通过优化算法来更新神经网络的参数。

常用的 Scheduler

• DDPM Scheduler:
  • 最基础的 Scheduler。
  • 以线性方式增加噪声。
  • 在去噪过程中采用线性插值。
• DDIM Scheduler:
  • 在 DDPM 的基础上进行改进。
  • 通过调整噪声添加和去除的步数，可以实现更快的采样速度。
• PNDM Scheduler:
  • 引入预测噪声的方差。
  • 提高了模型的稳定性。

Scheduler 的参数和对应 diffuers 代码

下方代码位于/path/to/diffusers/examples/unconditional_image_generation/train_unconditional.py

    # Initialize the scheduleraccepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())if accepts_prediction_type:noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps,beta_schedule=args.ddpm_beta_schedule,prediction_type=args.prediction_type,)else:noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)...
parser.add_argument("--ddpm_num_steps", type=int, default=1000)
parser.add_argument("--ddpm_beta_schedule", type=str, default="linear")parser.add_argument("--prediction_type",type=str,default="epsilon",choices=["epsilon", "sample"],help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.",)

• num_train_timesteps (ddpm_num_steps): 训练过程中的扩散步数。参数决定了训练过程中所使用的扩散步数（timesteps）的数量。每个时间步（timestep）代表一个噪声注入或去噪的阶段。更多的时间步通常意味着更细粒度的噪声注入过程，但也会增加计算成本。
• ddpm_beta_schedule: 默认值为 linear，表示使用线性调度方式。线性（linear）： 噪声量随时间步线性增加。余弦（cosine）： 噪声量遵循余弦曲线变化，通常在训练初期变化较大，后期变化较小。
• prediction_type: 参数决定了模型的预测目标，是预测噪声（epsilon）还是直接预测重建图像（x0）。选择不同的预测目标会影响模型的训练策略和生成效果。
  • epsilon： 模型预测噪声误差。这种方法通常需要反向计算来得到重建图像。
  • sample： 模型直接预测重建图像 x0。这种方法在计算上可能更直接，但在某些情况下预测效果可能不如预测 epsilon。
  • 目前来说一般是预测噪声效果更好。

Scheduler 的设计原则

• 平滑过渡: 确保从纯数据到纯噪声的过渡是平滑的。
• 控制生成过程: 通过调整 Scheduler 的参数，可以控制生成样本的多样性和质量。
• 提升模型稳定性: 防止模型出现发散等问题。

损失函数

在后向扩噪过程中，通常采用均方误差（MSE）作为损失函数：

loss = mean((predicted_noise - noise) ** 2)

其中：

• predicted_noise：模型预测的噪声。
• noise：Scheduler 计算出的真实噪声。

总结

Scheduler 在扩散模型中扮演着重要的角色，它通过控制噪声的添加和去除，影响着模型的训练和生成效果。不同的 Scheduler 有不同的特点和适用场景，研究者可以根据具体任务选择合适的 Scheduler。

参考文献

• Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840-6851.