细观Diffusion隐空间里UNet的训练流程(一)-厦门雄霸

细观Diffusion隐空间里UNet的训练流程(一)

行业动态 | 2023-08-18 09:45:43 | 作者

在上一期里，您已经领会Encoder 和Decoder 两者的涵意和功能了。请您看看图1 里，在两者之间有个空间，通称为：隐空间(latent space)。

图1

（来源：https://arxiv.org/pdf/1909.11573.pdf）

本期以SD(Stable Diffusion) 模型为例，来欣赏其隐空间里的架构之美。由于Diffusion 隐空间里的主角是：UNet模型。于是，先来细观Diffusion 隐空间里UNet模型的训练流程。

1 简介UNet模型

在本专栏里，曾经详细介绍过AE(autoencoder)，相信您对AE 已经很熟悉了，不再重复介绍AE模型。

UNet是AE模型的延伸型式。由于典型AE 模型的特性是：其前段的Encoder( 编码器) 是特征提取的过程，会过滤掉一些信息。这些被过滤掉的信息，就没有办法传递到后段的Decoder( 解码器) 了。

于是，UNet 就来延伸AE 模型，将原有的Encoder与Decoder 之间增加了一些连结，例如Encoder 的每一层输出都与对映层级的译码器连接。因此使得编码器每一层的信息，额外输入到译码器的对映层，让Decoder在重建的过程中，比较不会遗失掉重要的信息了。而Unet 模型的架构如其名呈现一个U 字形( 图2)。

图2

（来源：https://www.researchgate.net/)

如此，UNet 使用对映层级的连接来保留小且精细尺度的特征。

2 Diffusion隐空间里的UNet角色

在AIGC 潮流中，SD(Stable Diffusion) 产品的推出是AIGC 图像生成发展历程中的一个里程碑，提供了高性能模型，能快速生成创意十足的图像( 图3)。

图3

（来源：https://jalammar.github.io/illustrated-stable-diffusion/)

SD 提供两项主要功能：输入文本(Text) 提示来生成图像( 即text2img)；输入图像，然后根据文本描述来修改图像( 即输入text + img)。在SD 隐空间里，UNet 扮演关键性角色( 图4)。

图4

（来源：https://jalammar.github.io/illustrated-stable-diffusion/)

在SD 隐空间里，使用了1 个UNet 模型，并搭配1 个时间调度(scheduling) 器，来担任图像生成的核心任务。扩散(diffusion) 一词描述了SD 隐空间里进行的图像生成情形，整个过程都是在隐空间里逐步推进(stepby step) 的，每一步都会增加更多的噪音(noise)。例如，我们准备了一张图像，接着产生随机噪音加入到图像中，然后将这含有噪音的图像输入给UNet，让它来预测噪音。

为了加速图像的生成效率，SD 使用VAE 的Encoder 先将原图像压缩为4×64×64型式的隐空间数组(array)，然后才展开上述的扩散、生成过程。换句话说，整个扩散过程都是在隐空间里进行的。

3 认识噪音(Noise)

刚才提到了，我们将随机噪音添加到原图像里，成为含噪音图像(noisyimage)。然后把它输入给UNet 模型，让它预测出此图像所含的噪音。在SD里，使用数学运算，根据时间步数(time step) 来决定将多少噪音量添加到原图像里( 图5)。

图5

来源(https://jalammar.github.io/illustrated-stable-diffusion/)

图4 里的step-1 是：准备1 张原图像，并使用VAE的Encoder将原图压缩为4x64x64 型式的数组。接着，step-2 是：产生1 张也是4x64x64 型式的随机噪音数组；step-3 是：选取1 个随机整数( 如图里的0~3)，即是时间步数(time step)，并由此整数( 即时间步数)来决定欲添加的噪音量；step-4 是：将所决定的噪音量实际添加到原图像的压缩数组里，就成为含有噪音的图像了。

4 训练UNet模型

4.1 输入1张训练图像

刚才提到：我们会根据时间步数(time step) 来计算出多少噪音量将添加到原图像里。也就是，我们掌握了两项数据：含噪音的图像(noisy image) 数组；时间步数(time step)。我们就把这两项数据输入给UNet 模型，来训练它，让它学习预测所添加的噪音，所以它输出是：预测噪音(predicted noise)。在训练过程中，希望它输出的预测噪音愈接近所添加的噪音就愈好，也就是预测误差值(loss) 愈小愈好( 图6)。

图6

(来源：https://medium.com/@steinsfu/diffusion-model-clearlyexplained-cd331bd41166)

4.2 训练很多回合

在实际训练UNet模型时，把全部图像的压缩数组都输入给UNet，让它学习一遍，通称为：训练1 回合(epoch)。在每一回合里，都会为每一张图像搭配1 个随机整数( 即时间步数) 值。并根据时间步数值来计算出噪音量，于是得到两项数据：含噪音的图像数组和时间步数。然后，把这两项数据输入给UNet 模型，来训练它，让它学习预测噪音( 图7)。