近来,Stable Diffusion 成为一个新兴的研究方向。一位名为 Matthias Bühlmann 的博主尝试实验探究这种模型的强大功能,结果发现 Stable Diffusion 是一个非常强大的有损图像压缩编解码器。他撰写了一篇博客描述了这个实验分析过程,以下是博客原文。
首先 Matthias Bühlmann 给出在高压缩因子条件下,Stable Diffusion 方法与 JPG、WebP 的压缩结果,所有结果都是 512×512 像素的分辨率:
旧金山风景图,从左至右:JPG (6.16kB), WebP (6.80kB), Stable Diffusion: (4.96kB)。
糖果店,从左至右:JPG (5.68kB), WebP (5.71kB), Stable Diffusion (4.98kB)。
动物照片,从左至右:JPG (5.66kB), WebP (6.74kB), Stable Diffusion (4.97kB)。
这些例子明显表明,与 JPG 和 WebP 相比,使用 Stable Diffusion 压缩图像可以在更小的文件大小下保留更出色的图像质量。
探究实验
Matthias Bühlmann 分析了一下其中的工作原理,Stable Diffusion 使用三个串联的训练好的人工神经网络:
- 变分自编码器(Variational Auto Encoder,VAE)
- U-Net
- 文本编码器(Text Encoder)
VAE 将图像空间中的图像编码和解码为某种潜在的空间表征。源图像(512 x 512,3×8 或 4×8 bit)的潜在空间表征会分辨率更低(64 x 64)、精度更高(4×32 bit)。
VAE 在训练过程中自行学习,随着模型的逐步训练,不同版本模型的潜在空间表征看起来可能会有所不同,例如 Stable Diffusion v1.4 的潜在空间表征如下(重映射为 4-channel 彩色图像):
当重新扩展和将潜在特征解释为颜色值(使用 alpha channel)时,图像的主要特征仍然可见,并且 VAE 还将更高分辨率的特征编码到像素值中。
例如,通过一次 VAE 编码 / 解码 roundtrip 得到如下结果:
值得注意的是,这种 roundtrip 不是无损的。例如,图中蓝色带子上白色的字在解码后可读性稍差了一些。Stable Diffusion v1.4 模型的 VAE 一般不太擅长表征小型文本和人脸。
我们知道,Stable Diffusion 的主要用途是根据文本描述生成图像,这就要求该模型要对图像的潜在空间表征进行操作。该模型使用经过训练的 U-Net 迭代地对潜在空间图像进行去噪,输出它在噪声中「看到」(预测)的内容,类似于我们有时把云看成某种形状或面孔。在迭代去噪步骤中,第三个 ML 模型(文本编码器)指导 U-Net 来尝试看到不同的信息。
Matthias Bühlmann 分析了 VAE 生成的潜在表征(latent representation)是如何进行有效压缩的。他发现对 VAE 中的潜在表征进行采样或对潜在表征应用已有的有损图像压缩方法,都会极大地降低重构图像的质量,而 VAE 解码过程似乎对潜在表征的质量鲁棒性较高。
Matthias Bühlmann 将潜在表征从浮点数量化为 8-bit 无符号整数,结果发现只有非常小的重构误差。如下图所示,左:32-bit 浮点潜在表征;中:ground truth;右:8-bit 整数潜在表征。
他还发现通过 palette 和抖动算法进一步量化,得到的结果会出乎意料的好。然而,当直接使用 VAE 解码时,palettized 表征会导致一些可见的伪影:
左:32-bit 潜在表征;中:8-bit 量化潜在表征;右:带有 Floyd-Steinberg 抖动的 palettized 8-bit 潜在表征
带有 Floyd-Steinberg 抖动的 palettized 表征引入了噪声,使解码结果失真。于是 Matthias Bühlmann 使用 U-Net 来去除抖动带来的噪声。经过 4 次迭代,重构结果在视觉上非常接近未量化的版本:
重构结果(左:带有 Floyd-Steinberg 抖动的 palettized 表征;中:经过四次迭代去噪;右:Ground Truth)。
虽然结果非常好,但还是会引入一些伪影,例如上图中心形符号上的光泽阴影。
虽然从主观上看,Stable Diffusion 压缩图像的结果比 JPG 和 WebP 好很多,但从 PSNR、SSIM 等指标看,Stable Diffusion 并没有明显的优势。
如下图所示,虽然作为编解码器的 Stable Diffusion 在保留图像粒度方面比其他方法要好得多,但受压缩伪影的影响,图像中物体形状等特征可能会发生变化。
左:JPG 压缩;中:Ground Truth;右:Stable Diffusion 压缩。
值得注意的是,当前的 Stable Diffusion v1.4 模型在压缩过程中无法很好地保留字体很小的文本信息和人脸特征,但 Stable Diffusion v1.5 模型在人脸生成方面有所改进。
左:Ground Truth;中:经过 VAE roundtrip (32-bit 潜在特征) ;右:从 palettized 去噪 8-bit 潜在特征解码的结果。
博客发布后,Matthias Bühlmann 的实验分析引起了大家的讨论。
Matthias Bühlmann 自己认为 Stable Diffusion 的图像压缩效果比预期好,U-Net 似乎能够有效消除抖动引入的噪声。不过,Stable Diffusion 模型未来的版本可能不会再有这种图像压缩特性。
然而有网友质疑道:「VAE 本身就被用于图像压缩」,例如基于 Transformer 的图像压缩方法 TIC 就用到了 VAE 架构,所以 Matthias Bühlmann 的实验似乎是大材小用了。
对此,你有什么看法?