Stable Diffsuion还能用来压缩图像？压缩率更高，清晰度超越JPEG等算法

基于文本的图像生成模型火了，出圈的不止有扩散模型，还有开源的Stable Diffusion模型。

最近一位瑞士的软件工程师Matthias Bühlmann无意间发现，Stable Diffusion不仅能用来生成图像，还可以用来压缩位图图像，甚至比JPEG和WebP的压缩率更高。

比如一张美洲骆驼的照片，原图为768KB，使用JPEG压缩到5.66KB，而Stable Diffusion可以进一步压缩到4.98KB，而且能够保留更多高分辨率的细节以及更少的压缩伪影，肉眼可见地优于其他压缩算法。

不过这种压缩方式也存在缺陷，即不适合压缩人脸和文本图像，在某些情况下，甚至会生成一些原图并不存在内容。

虽然重新训练一个自编码器也能做到类似于Stable Diffusion的压缩效果，但使用Stable Diffusion的一个主要优势在于，有人已经投入了上百万的资金帮你训练了一个，你又何必重新花钱训练一个压缩模型呢？

Stable Diffusion如何压缩图像

扩散模型正在挑战生成模型的霸主地位，对应的开源Stable Diffusion模型也在机器学习社区掀起一场艺术革命。

Stable Diffusion由三个训练后的神经网络串联得到，即一个变分自编码器（VAE），U-Net模型和一个文本编码器。

变分自编码器对图像空间中的图像进行编码和解码，从而获得该图像在潜空间的表征向量，以一个分辨率更低（64×64)具有更高精度(4x32bit)的向量来表示源图像(3×8或4x8bit的512×512)。

VAE在将图像编码到潜空间的训练过程主要依赖自监督学习，即输入和输出都是源图像，因此随着模型进一步训练，不同版本的模型的潜空间表征可能会看起来不同。

使用Stable Diffusion v1.4的潜空间表征通过重新映射和解释为4通道彩色图像后，看起来就是下图的中间图像，源图像中的主要特征仍然可见。

需要注意的是，VAE往返编码一次并不是无损的。

比如在解码之后，蓝色带子上的ANNA名字就没有源图像那么清晰了，可读性显著降低。

Stable Diffusion v1.4中的变分自编码器不太擅长表示小文本以及人脸图像，不知道在v1.5版本中是否会改善。

Stable Diffusion的主要压缩算法就是利用图像的这种潜空间表征，从短文本描述中生成新的图像。

从潜空间表征的随机噪声开始，使用充分训练的U-Net迭代去除潜空间图像的噪声，用一种更简单的表征输出模型认为它在这个噪声中「看到」的预测，有点像我们在看云的时候，从不规则的图形中还原出脑海里的形状或面孔。

当使用Stable Diffusion来生成图像时，这个迭代去噪步骤是由第三个组件，即文本编码器引导的，该编码器为U-Net提供关于它应该尝试在噪声中看到什么的信息。

不过对于压缩任务来说，并不需要文本编码器，所以实验过程只创建了一个空字符串的编码用于告诉U-Net在图像重建过程中进行非引导去噪。

为了使用Stable Diffusion作为图像压缩编解码器，算法需要有效地压缩由VAE产生的潜表征。

在实验中可以发现，对潜表征进行下采样或者直接使用现有的有损图像压缩方法，都会大大降低重建图像的质量。

但作者发现 VAE 的解码似乎对潜表征的量化（quantization）非常有效。

通过对从浮点到8位无符号整数的潜量化进行缩放、拖拽（clamping）和重新映射，只会产生很小的可见重构错误。

通过量化8位的潜表征，图像表示的数据大小现在是64*64*4*8bit=16kB ，远小于未压缩源图像的512*512*3*8bit=768kB

如果潜表征的位数小于8bit，无法产生比较好的效果。

如果对图像进一步执行调色板（palettizing）和抖动（dithering），则量化效果就会再次提升。

使用256*4*8位向量和Floyd-Steinberg抖动的潜表征创建了一个调色板表示，使数据大小进一步压缩到64*64*8+256*4*8bit=5kB

潜空间调色板的抖动会引入噪声，从而扭曲了解码结果。但由于Stable Diffusion是基于潜噪声的去除，所以可以使用U-Net去除抖动引起的噪声。

经过4次迭代，重建结果在视觉上非常接近未量化的版本。

虽然数据量大大减少了（源图像为压缩图像的155倍大），但效果是非常好的，不过也引入了一些伪影（比如原图的心形图案中不存在伪影）。

有趣的是，这种压缩方案引入的伪影对图像内容的影响比对图像质量的影响更大，而且以这种方式压缩的图像可能包含这些类型的压缩伪影。

作者还用zlib对调色板和索引进行了无损压缩，在测试样本中，大多数的压缩结果都小于5kb，但这种压缩方法仍然存在更多的优化空间。

为了评估该压缩编解码器，作者没有使用任何在网上找到的标准测试图像，因为网上的图像都有可能在Stable Diffusion的训练集中出现过，而压缩这类图像可能会导致不公平的对比优势。

为了尽可能公平地进行比较，作者使用了Python图像库中最高质量的编码器设置，以及使用mozjpeg库添加了压缩后的JPG数据的无损数据压缩。

值得注意的是，虽然Stable Diffusion的结果主观上看起来比JPG和WebP压缩的图像要好得多，但在标准测量指标(如PSNR或SSIM)方面，它们并没有明显更好，但也没有更差。

只是引入的伪影类型不那么明显，因为它们对图像内容的影响大于对图像质量的影响。

这种压缩方法也有一点危险，虽然重建特征的质量很高，但内容可能会受到压缩伪影的影响，即使它看起来非常清晰。

例如，在一张测试图像中，虽然Stable Diffusion作为编解码器在保持图像的质量方面要好得多，甚至连相机颗粒纹理（camera grain）都能保留下来(这是大多数传统压缩算法难以做到的) ，但其内容仍然受到压缩伪影的影响，像建筑物形状这样的精细特征可能会发生变化。

虽然在JPG压缩图像中当然不可能比在Stable Diffusion压缩图像中识别出更多的真实值，但是Stable Diffusion压缩结果的高视觉质量可能具有欺骗性，因为JPG和WebP中的压缩伪影更容易识别。

如果你也想动手复现一遍实验，作者在Colab上开源了代码。

代码链接：​https://colab.research.google.com/drive/1Ci1VYHuFJK5eOX9TB0Mq4NsqkeDrMaaH?usp=sharing​

最后，作者表示，文章中设计的实验仍然是相当浅显的，但效果仍然令人惊喜，未来仍然有很大的改进空间。