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

高焕堂

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

本期以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）。

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

2 Diffusion隐空间里的UNet角色

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

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

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

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

3 认识噪音（Noise）

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

图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）。

4.2 训练很多回合

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

上述SD 的UNet 训练方法比较特殊之处是：这UNet的功能是输出是噪音，而不是去除噪音。

5 结束语

本期说明了SD 隐空间里UNet 模型的角色及其训练步流程的起步。下一期，将继续介绍及欣赏SD 隐空间的架构之美，也领会AIGC 生成艺术的源点。