用基础扩散模型实现有损图像压缩

将扩散模型纳入图像压缩领域，有可能产生逼真而详细的重建效果，尤其是在比特率极低的情况下。以往的方法主要是将扩散模型用作对条件信号中的量化误差具有鲁棒性的表达式解码器。然而，要以这种方式获得有好的结果，需要对扩散模型进行大量的训练，而且由于迭代式的生成过程，推理时间很长。本文将消除量化误差作为一项去噪任务，利用扩散来恢复传输图像潜像中丢失的信息。本文的方法只需执行不到 10% 的完整扩散生成过程，并且不需要对扩散模型进行结构性更改，从而可以将基础模型作为强先验模型使用，而无需对主干模型进行额外的微调。本文提出的编解码器在定量真实度指标上优于以前的方法，而且本文验证了该方法重建在质量上更受终端用户的青睐，即使其他方法使用的比特率是本文的两倍。

来源：ECCV 2024
论文题目：Lossy Image Compression with Foundation Diffusion Models
论文链接：https://arxiv.org/pdf/2404.08580
论文作者：Lucas Relic, Roberto Azevedo, Markus Gross, Christopher Schroers
内容整理：阳浩宁

简介

在当今的数字时代，多媒体内容主导着全球互联网流量，因此开发高效的压缩算法变得越来越重要。数据驱动的神经图像压缩（NIC）方法对速率和失真都进行了优化，现已超越了使用手工制作变换的传统编解码器，但目前的大多数方法仍然会在极低的比特率设置下产生模糊和不真实的图像。生成压缩方法试图通过引入 GAN 架构和对抗性或感知损失来重建逼真的图像。扩散模型在性能上优于 GAN，但其适用性和采用受到了模型规模大和训练时间昂贵的阻碍。开源基础模型的引入有可能使这些功能强大的模型平民化，并提供强大的先验，可用于各种领域的特征提取或迁移学习，例如图像生成、深度估计，甚至音乐生成。这种方法目前还很少被探究过。

为了提升图像压缩的最优性能，本文提出了一种新型图像压缩编解码器，利用基础潜在扩散模型来合成丢失的细节，尤其是在低比特率情况下。利用量化误差和噪声之间的相似性，本文传输量化图像潜像，并在接收器上执行与潜像噪声水平（即量化误差）相对应的去噪步骤子集。本文提出的关键组成部分是：i) 基础潜在扩散模型中的自编码器，将输入图像转换到低维潜在空间；ii) 学习型自适应量化和熵编码器，在单一模型内实现对比特率的推理时间控制；iii) 学习型方法，预测理想的去噪时间步长，在传输成本和重建质量之间实现平衡；以及 iv) 扩散解码过程，合成量化过程中丢失的信息。与之前的工作不同，本文的方法只需要一小部分迭代扩散步骤，并且可以在少于 100k 幅图像的数据集上进行训练。本文还直接优化了输入图像和重建图像之间的失真目标，在确保输入图像一致性的同时，还能通过扩散骨干保持高度逼真的重建。

方法

图1展示了本文方法的高级结构。它由一个 VAE，一个量化和扩散模型时间步骤参数预测网络，一个熵模型和一个潜在扩散模型组成。

其中的基础模型选择的是 Stable Diffusion v2.1。本方法的一个关键特征是量化参数 γ 和去噪步数 t 都可以根据内容和目标比特率进行调整。这是通过一个神经网络 P_ϕ(y,λ)来实现的。直观地说，本方法学会了通过量化转换来摒弃可以在扩散过程中合成的信息。由于量化过程中引入的误差类似于添加噪音，而扩散模型在功能上是去噪模型，因此可以用来消除编码过程中引入的量化噪声。

参数预测网络

模型的量化参数 γ 和扩散模型执行的去噪步数 t 都是通过参数预测网络 P_ϕ(y,λ)得到的。网络的输入是潜像 y 和率失真控制参数 λ。

自适应量化操作 Q 是对 y 的每一个通道做仿射变换，参数为 γ ，然后做标准整数量化，过程如下：

γ 作为辅助信息传输，以便解码器在客户端进行逆变换。

时间步骤方面，与从纯噪声开始扩散的图像生成任务相反，压缩任务中从量化的潜影开始扩散过程，而潜影已经包含了内容的结构和语义信息。在这种情况下，在解码过程中执行整个去噪步骤既浪费又会导致图像过度平滑。因此，模型要学会预测能产生最佳解码图像的去噪扩散步骤子集。图2展示了解码器所执行的扩散步骤数量对解码图像质量的影响，步骤太少或太多分别会导致噪声图像或过度平滑图像。由于最佳去噪步数取决于潜像中的噪声量（也取决于量化程度），因此我们在参数估计模块中联合预测 t和 γ。这个网络采用了全卷积结构，下采样层和标准卷积层交替堆叠，滤波器数量随着深度的增加而增加。最后一层输出滤波器的数量减少，使之与估计的参数总数相对应，并对每个输出通道应用均值池化，为每个参数生成一个标量。每个卷积层之间使用 SiLU 激活，最后一个卷积层之后不使用激活。不过，时间步骤标量的预测中应用了 sigmoid 激活，以保证范围在 [0,1] 之间。为了使 P_ϕ 以目标比特率为条件，λ被扩展到与潜样本相同的空间维度，沿矢量线将它们串联起来，并沿信道维度进行串联，然后再用参数估计网络进行处理。

模型的损失函数对码率和失真进行了权衡：

训练时选择 Vimeo-90k 数据集，训练过程中将图像随机裁剪为 256×256 大小，用 1e-4 学习率训练 300,000 步。为了用一个模型实现多目标比特率的训练，在每次梯度更新时从 λ∈[1,5,10,20]中随机选择。此外本文对基础扩散模型进行了细微的调整，训练时只进行一次 DDIM 采样迭代，并直接使用 x^0作为完全去噪数据。在模型运行的低时间步长范围内，x^0 与真实的完全去噪数据 x₀ 之间的差异很小，足以满足参数估计模块的优化要求。在推理时，执行标准的迭代 DDIM 过程。此外，在扩散采样过程中，时间步骤 tt 被用来索引所需的预计算值数组。这种离散化方式阻碍了参数估计网络的优化，因此本文为每个所需值实现了连续函数，并在训练过程中使用预测的时间步对其进行评估。

测试时使用的数据集包括 Kodak、 CLIC2022 和 MS-COCO 30k，指标包括 PSNR、MS-SSIM、LPIPS 和 FID，对比的方法包括 HiFiC、ILLM、CDC和HFD。由于定性指标往往无法捕捉图像重建的感知质量，本文进一步进行了用户研究，以评估我们结果的视觉质量。研究设置为两选一强迫选择 (2AFC)，向每位参与者展示源图像和两种方法的重建图像，并要求他们选择自己喜欢的重建图像。最终用Elo评分系统来进行排名。

实验结果

图3展示了几种方法定量对比的结果。本文的方法在 FID 层面优于所有方法，在其他指标层面优于其他基于扩散模型的方法，并且与其他最好的生成式编码方法相接近。不出所料，本文的方法和其他生成式方法在 PSNR 测量中都会受到影响，因为这些方法更倾向于令人愉悦的重构，而不是细节的精确复制。

图4展示了用户研究的结果。本文的方法显著超过了其他方法。

图5展示了几种方法定性对比的结果。本方法可以在保持高度逼真性的同时，始终如一地重建精细的细节和可信的纹理。HFD 经常合成不正确的内容或产生平滑的重建。CDC 和 ILLM 会产生不自然的模糊或高频生成假象，即使它们使用的比特率是本方法的 2 倍。

模型效率方面，本文使用 RTX 3090 GPU 测试了几种方法的编解码时间。本方法处理一幅图像只需 3.49 秒，比 CDC 的 6.87 秒快近一倍。ILLM 处理一幅图像只需 0.27 秒。不过，值得注意的是，基于扩散模型的方法由于其迭代去噪的特性，一般比其他编解码器慢。由于采用了 Stable Diffusion 作为主干，本方法比 CDC 更复杂（分别含 1.3B 和 53.6M 参数），而 ILLM 包含 181.5M 个参数。但是本方法大部分参数都来自扩散模型，而需要训练的模块仅包含 36M 个参数。

与其他生成方法类似，本方法可以舍弃某些图像特征，同时在接收端合成类似信息。然而，其局限性在于，在特定情况下，这可能会导致重建不准确，例如弯曲直线或扭曲小物体的边界。这是本方法中基础模型的众所周知的问题，这可以归因于其 VAE 的特征维度相对较低。

结论

本文提出的基于潜在扩散模型的新型有损图像压缩编解码器可以在低比特率到超低比特率下生成逼真的图像重建，在感知指标和用户主观偏好方面均优于之前的生成式编解码器。通过将扩散模型的去噪能力与量化噪声的固有特性相结合，本文的方法预测了理想的去噪步骤数量，从而在一系列比特率下用单一模型生成令人愉悦的重构图像。与以前的扩散编解码器相比，本文的方案解码时间更快，而且由于重新使用了基础模型骨干，训练预算也低得多。未来可能开展的工作包括整合更高效的骨干模型，以及支持用户控制，在速率-失真-逼真度之间进行权衡。