HybridFlow：将连续性融入掩码码本以实现极低比特率图像压缩

用户1324186

发布于 2024-05-20 13:57:04

1290

发布于 2024-05-20 13:57:04

引言

通过训练有素的神经网络，学习图像压缩（LIC）已被证明优于 VVC 或 JPEG2000 等传统方法。其整体思路是在编码器中将输入图像编码到潜空间中，压缩潜特征以减少传输比特，并在解码器中使用解压缩的潜特征重建输出图像。根据传输信息的类型，LIC 方法大致可分为两类。第一类已被广泛研究，其特点是传输连续的压缩特征图。原始潜特征经过经典的量化和熵编码，得到具有连续值的紧凑比特流，解码器恢复降级的潜特征进行重构。当比特率极低时，由于严重量化，恢复的潜特征质量很差，导致重建质量低，过于平滑，缺乏表现细节。

第二类方法的特点是传输整数索引，因为在图像复原任务中使用基于量化矢量的码本来学习一般图像先验的做法越来越流行。学习到的视觉码本经过预训练，可将潜在图像的分布离散到有限的离散集空间中。通过在编码器和解码器之间共享码本，编码器将潜在特征映射到码字索引，解码器通过使用整数索引检索码字特征来恢复近似的潜在特征，从而进行重建。从高质量图像中学到的高质量码本通常能确保高质量的感知重建，细节清晰丰富。然而，输出图像可能不忠于原始输入图像，例如，离散化的码本会消解微小的内容变化。捕捉视觉细节的大型码本或多个码本（每个码本都侧重于特定类别的表示）可以缓解这一问题，但要牺牲更高的比特率。因此，当比特率极低时，有限的码本大小会导致保真度较低的重建。

在本文中，作者提出了一种混合框架，它得益于上述两类技术的双流互补性，可同时实现极低比特率传输和高质量重建。从输入图像中生成两个并行流：一个是基于高质量码本的离散索引，它利用学习到的一般图像先验获得高感知重建质量；另一个是提供保真细节的极低比特率连续特征流。这两种流通过有效的桥接机制结合在一起，用于屏蔽token生成和校正像素解码。本文的贡献在于：

作者介绍了一种用于图像压缩的新型双流框架--HybridFlow，它甚至可以在极低的比特率（<0.05 bpp）下实现清晰、忠实的高质量图像重建，超越了以往的方法。
基于码本的离散流进一步引入了 “屏蔽预测 ”策略。受 MAGE 基于屏蔽token的Transformer架构的启发，只需从索引的一部分进行引导生成，不仅减少了传输索引，还实现了重建质量和比特率之间的可控权衡。
作者提出了一种桥接机制来合并两个信息流。连续特征被输入token解码器的交叉注意力模块，以指导基于码本特征的预测生成。同时，连续特征通过像素解码器旁的校正网络，纠正使用基于码本特征的像素解码过程的偏差。

作者通过几个基准数据集进行了实验，从定性和定量两个方面评估了方法的有效性。定性结果表明， “HybridFlow ”框架可以保持基于码本重构的高质量和清晰度，同时通过注入连续图像特征有效纠正像素级失真。在定量性能方面，与纯粹基于码本的方法相比，PSNR 平均提高了约 3.5dB，LPIPS 分数相同甚至更高；与纯粹基于连续特征的方法相比，LPIPS 分数显著提高（55.7%）。总体而言，“HybridFlow ”在可信度和清晰度之间、可信度和感知质量之间实现了平衡。

方法

图 1 给出了双流 HybridFlow 工作流程，用于以超低比特率（< 0.05bpp）进行高质量重建。

图 1：用于极端低比特率图像压缩的双流 “HybridFlow ”框架。

HybridFlow 图像压缩

基于码本的表示。 给定一个输入图像

x\in \mathbb{R}^{3\times H\times W}

，通过使用学习到的视觉码本

\mathbf{C}\in\mathbb{R}^{c\times n_z}

生成第一个数据流，作为离散索引

d\in\mathbb{R}^{\frac Hn\times\frac Wn}

。具体来说，𝑥 被 VQ-Encoder

\mathbf{E}_{\mathbf{VQ}}

编码成潜在表示

y_{VQ}\in\mathbb{R}^{c\times\frac{H}{n}\times\frac{W}{n}}

，并进一步映射成索引映射 𝑑。每个条目向量

y_{ij}\in y_{VQ} (i=1,\ldots,\frac{H}{n},j=1,\ldots,\frac{W}{n})

被映射到最接近的码字

c_{ij}\in\mathbb{C}

，其码字索引为

d_{ij}\in[0,n_z)

。在实践中，为了获得超低比特率（

n_{z}=1024

，在本文的实验中，离散数据流的比特率约为 0.02 bpp），使用了相对较小的单个码本。

屏蔽模块。 为了进一步减少离散数据流中的传输比特，作者没有像之前的研究那样直接传输 “𝑑 ”索引，而是选择性地只传输 “𝑑 ”的一部分。如图 2 所示，利用配备了预定义、结构化和便于压缩的（如 1_4 屏蔽和 1_9 屏蔽）的屏蔽模块，传输经过屏蔽的 𝑑′ 而不是 𝑑。所选的屏蔽方案决定了剩余信息的比例，从而提供更高效的压缩率。

图 2：码字索引的候选屏蔽方案（以 8x8 索引显示）。各方案的平均屏蔽率分别为 50%、75%、90.3% 和 93.75%。

连续域表示。 通过使用基于连续特征的 LIC 方法，将输入

x\in\mathbb{R}^{3\times H\times W}

作为连续潜在特征

y_{c}\in\mathbb{R}^{f\times\frac{H}{m}\times\frac{W}{m}}

，生成第二个数据流。本文使用的 MLIC 方法具有最先进的低比特率压缩性能。一般来说，以前的方法能提供合理重构的最低比特率约为 0.1 bpp。为了获得超低比特率，在将原始图像输入 MLIC 流水线之前，首先对其进行 4 倍降采样。在本文的实验中，连续数据流的比特率约为 0.025 bpp。

屏蔽预测器。 在解码器方面，基于token的Transformer预测器 T 用于从接收到的屏蔽索引

d^{\prime}

恢复完整的索引

\hat{d}

。受 MAGE 的启发，作者为该预测器采用了编码器-解码器Transformer结构。此外，利用连续域 MLIC 解码过程，可以恢复连续域的潜在特征

\hat{y}_{c}

，并将其输入屏蔽预测器，以指导缺失token的生成。这种想法类似于音频到文本的翻译，在这种翻译中，音频信息通过交叉注意力用于指导文本token的生成。作者在转换解码器的每个解码块中插入一个交叉注意力模块，将连续域潜特征输入该模块，作为指导来辅助离散token生成。

像素解码器。 为了合并双数据流以实现高质量重建，在 VQ-Decoder

\mathrm{D}_{\mathrm{VQ}}^{\cdot}

的同时引入了复制的像素解码器。从恢复的

\hat{d}

中，矢量量化的潜在

\hat{y}_{VQ}

可以从包含

\hat{d}

所指示的相应编码词的码本 C 中检索出来。在以前的研究中 [21, 35]，VQ-Decoder仅根据

\hat{y}_{VQ}

重建输出

\hat{x}

。在本文的方法中，复制的解码器充当校正网络，将恢复的连续域潜在

\hat{y}_{c}

作为输入，并将重复像素解码器每个上采样层的解码表示发送到 VQ-Decoder的相应上采样层，以纠正偏差。这种纠偏将重要的保真信息从连续域传递到重建的图像中，同时提供高感知质量和高保真。

复杂性感知动态屏蔽

图像区域的细节复杂度参差不齐。根据不同区域的复杂度调整比特分配，可以进一步提高压缩效率。本文使用多个指标来分析图像区域的复杂度得分，包括熵、对比度、颜色多样性（直方图熵）和空间频率（傅里叶域），类似于 ClassSR 及其进一步扩展。对训练数据中的每个指标进行归一化处理，使其范围在 [0,1] 之间，然后对归一化指标求平均值，得出最终的复杂度得分。然后，根据图像复杂度得分，为每个图像区域设置一个动态屏蔽方案。对于具有丰富复杂细节的区域，可采用低屏蔽率，以保留复杂信息；而对于简单区域，则可采用高屏蔽率，以极大地减少比特。

训练过程

为了有效地训练所提出的 HybridFlow 框架，训练过程分为三个阶段，如图 3 所示。

图 3：提出的框架的训练流程。

预训练。 基于连续特征的 LIC 编码器使用预先训练好的 MLIC 编码器来生成连续潜特征和超先验。用于生成基于码本的潜特征的基于码本的 VQ-Encoder和学习的视觉码本以及 VQ-Decoder使用了预先训练好的 VQGAN 模型。这些预训练模块旨在计算双流潜在特征，分别强调高质量重构和高保真重构。

Transformer预测器训练。 上一阶段的预训练模块被冻结，在这一阶段训练Transformer预测器。对于Transformer预测器中的每个解码器块，选择性地训练交叉注意力模块、MLP 模块和负责映射到token logits的最外层 MLP 结构，同时保持与自注意力相关的预训练参数不变。作者沿用了 MAGE 的训练设计，即随机屏蔽部分token序列。让

表示从 VQGAN 码本输出的扁平化token（索引），

M_{b}

表示随机生成的二进制屏蔽，

Y_{m}

表示所有被屏蔽的token，

Y_{r}

表示剩余的未被屏蔽的token。Transformer预测器经过训练，可以准确预测被屏蔽的token。预测的屏蔽token的概率分布与相应的真实分布之间的损失函数形式化为

L_{prediction}=-E(\sum\log p(m_i|Y_r)) \quad (1)

其中，

m_i\in Y_m

是预测的屏蔽token。根据之前的 MIM 研究，作者只计算了屏蔽比例的损失，以获得更好的模型容量。

像素解码器训练。 前面提到的模块被冻结，作者在这一阶段训练像素解码器。首先，从预先训练好的 VQGAN 解码器中复制了一个相同的像素解码器，作为复制的解码器的校正网络。然后，通过最小化原始输入

和重建

\hat{x}

之间的像素级损失，对重复解码器进行微调：

L_{pixel\_distortion}=w_1*L1(x,\hat x)+w_2*L_{perceputal}(x,\hat x) \quad (2)

w_{i}

为损失权重，

为 L1_loss，

L_{𝑝𝑒𝑟𝑐𝑒𝑝𝑡𝑢𝑎𝑙}

为 AlexNet 产生的感知损失。预先训练好的 VQGAN 模型中的 VQ-Decoder 被冻结，这样像素级损失就能使重复解码器在不过度影响基于码本的表示所获得的感知质量的情况下，提高像素保真度。

实验结果

实验设置

数据集。 为了充分利用双流系统的建模能力，作者使用 ImageNet 进行了训练，其中包含 100 多万张不同的图像。为了进行性能评估并与现有的几种方法进行公平比较，作者在Kodak、CLIC 2020 测试集和 Tecnick数据集上进行了测试。

模型配置。 对于基于连续特征的数据流，作者使用了预先训练好的质量最低的 MLIC 模型（𝜆 = 0.0018）。对于基于码本的数据流，使用了预先训练好的 VQGAN，以及 MAGE 的Transformer屏蔽预测器。具体来说，为了与预先训练好的 MAGE 模型的默认 256 长度的token索引输入相匹配，作者将 256x256x3 的图像片段输入到本文的系统中。这确保了扁平化索引的长度为 256。附加的复杂度感知模块使用三个复杂度得分阈值从三个屏蔽方案中进行选择：1_9 屏蔽方案用于复杂度得分小于 0.24 的简单区域，1_2 屏蔽方案用于复杂度得分大于 0.77 的复杂区域，1_4 屏蔽方案用于介于两者之间的中等区域。

屏蔽预测推理。 与之前使用随机抽样来创建新内容的基于token的Transformer不同，本文的解码依赖于最大logits来消除随机性，从而稳定生成以达到压缩目的。此外，与 MAGE 和 MaskGIT 所使用的迭代解码不同，作者插入的交叉注意力模块允许在测试过程中一次前向传递就能进行准确预测，无需多次分步恢复。这些变化减少了屏蔽token预测的随机性，提高了解码效率。

桥接机制的有效性

作者首先从两个方面验证了所提出的桥接机制的有效性：利用连续域超先验信息帮助Transformer预测器恢复原始索引图，以及利用连续潜特征的保真度信息通过复制解码器帮助像素解码。

连续性辅助屏蔽预测器。 作者比较了本文提出的Transformer预测器和原始的预训练 MAGE，看它们是否能根据部分真实索引恢复原始索引图。作者测试了几种屏蔽方案，并将恢复的索引图直接输入预训练的 VQGAN 解码器进行图像复原。如图 4 所示，与真实值相比，原始 MAGE 的预测偏差随着屏蔽比例的增加而增大。相比之下，本文的Transformer预测器却能获得高度稳定的预测结果，并忠实于原始图像的特征，这要归功于交叉注意力模块引入的连续流全局视觉线索。

图 4：屏蔽预测器的效果。第一行到最后一行分别是 1_4、1_9、1_16 的屏蔽方案。

连续性辅助像素解码器。 作者比较了提出的连续性融合像素解码器和原始的预训练 VQGAN 解码器。在这项测试中，索引图没有被屏蔽，而是直接输入到像素解码器中，以单独评估所提出的像素解码器的有效性。如图 5 所示，由于码字特征数量有限，仅根据基于码本的潜在特征解码像素可能无法很好地与一般像素分布保持一致，尤其是在人脸或文本等敏感区域。连续特征包含有关原始分布的重要保真信息，可在差异敏感区域为像素解码器提供保真校正。因此，生成的图像可以更好地匹配原始图像内容，同时保持较高的感知质量。

图 5：连续性辅助像素解码器的效果。红色方框强调了复制解码器利用连续域信息有效纠正基于码本的偏差的特定区域。

复杂性感知屏蔽的有效性

对复杂程度不同的模块采用了不同的屏蔽方案。也就是说，在保持相似图像质量的前提下，将更多的传输位分配给高复杂度区域（通常是感知敏感区域），而将较少的位用于低复杂度区域。如图 6 所示，根据分割得到的图像片段的复杂度，大量具有简单特征的区块被定义为 “简单”，而对特征敏感的区域，如人脸和复杂的服装图案，则被归类为 “困难”（“复杂”）。复杂性感知模块为进一步减少传输比特提供了一种简单而有效的方法，例如，与统一的 1_4 遮罩计划相比，平均可减少 12.5% bpp，使本文的框架能够更好地服务于极端低比特率的图像压缩场景。

图 6：复杂性感知动态屏蔽调度的效果。红色、绿色和蓝色分别代表 “容易”、“中等 ”和 “困难”。最右边一列描述了使用动态屏蔽的输出图像质量（PSNR），以及与 1_4 均匀遮罩相比减少的 bpp 百分比。

性能比较

比较方法。 为了证明提出的双流 HybirdFlow 图像压缩框架在极低比特率情况下的优势，作者将本文的工作与单流 VQGAN 和单流 MLIC （作为基于 SOTA 连续特征的 LIC）进行了比较。这种比较有效地显示了双流融合所带来的均衡的整体性能提升。此外，作者还将本文的方法与微调 VQGAN 压缩方法、另一种基于码本的 LIC 方法以及 SOTA 传统压缩方法 VVC 进行了比较。

评估指标。 在评估低比特率图像压缩模型方面，先前的研究通常要么测量像素级差异（PSNR），要么测量感知差异（LPIPS），很少同时测量这两种差异。实际上，PSNR 强调的是人眼观看时的视觉相似度，而 LPIPS 则倾向于强调局部图像质量，包括清晰度和具体细节。实际上，这两个指标对于图像压缩任务都至关重要。在本实验中，作者发现即使是在极低比特率的图像压缩环境下，本文提出的双流结构也能很好地平衡这两个指标。可以保持图像的清晰度和细节表现力，而不会明显偏离原始图像。

定性比较。 在本实验中，屏蔽模块采用了固定的 1_4 屏蔽方案，以保持模型压缩质量与比特率之间的一致关系，从而进行公平比较。如图 7 所示，与单流 MLIC 和 VQGAN 相比，本文的方法在极低比特率情况下的图像重建质量具有显著优势。MLIC 的重建结果显示出明显的油画般的模糊，并经常出现异常噪声。基于 VQGAN 的重建结果虽然整体清晰，但在细节敏感区域，特别是图像斑块的边缘，像素偏差明显，导致像素不连贯，严重影响视觉感受。此外，对于包含文本信息的图像，VQGAN 的有限特征会造成文本内容的局部失真。相比之下，本文的方法在视觉上平衡了清晰度和保真度，既保证了图像重建的清晰度，又在很大程度上减少了模糊和噪点，还通过来自连续流的互补结构信息纠正了 VQGAN 中的像素失真。

图 7：基于单流的 VQGAN、MLIC 和提出的双流结果的定性可视化。

定量比较。 本文的方法通过可变的屏蔽方案提供了不同的压缩率，从大约 0.025 到 0.065 bpp 不等。最低的质量对应于屏蔽所有码本引，最高的质量对应于完全不屏蔽。如图 8 所示，就像素级 PSNR 而言，VVC 等传统方法以及 Cheng2020 和 MLIC 等基于单流连续特征的 LIC 方法通常比基于单流码本的 LIC 方法表现更好。然而，它们在感知质量方面的表现（如 LPIPS）却很差。它们倾向于优先考虑整体像素的相似性，生成大块、无细节、模糊的斑块，从而牺牲了图像的清晰度和精细度。相比之下，基于单流码本的 LIC 方法的 PSNR 性能普遍较差（在相同比特率下比传统方法低约 4.5 dB）。这证实了仅基于学习的码本进行重构会带来明显的像素级失真。然而，由于像素解码器依赖于高质量的码本特征，重构相当稳定，从而获得了更好的 LPIPS。本文的方法结合了两种方法的优点，在感知 LPIPS 和像素级 PSNR 之间取得了良好的平衡。与基于单流码本的 LIC 相比， PSNR 曲线与传统方法非常相似，平均 PSNR 提高了 3.5 dB，LPIPS 甚至进一步提高。与只关注 PSNR 性能的传统方法相比，虽然 PSNR 较低，但重建图像的感知质量和清晰度明显更好，在三个测试集中，LPIPS 平均提高了 55.7%。这些结果表明，对于极端的低比特率图像压缩场景，HybridFlow 方法为实际应用提供了良好的均衡图像重建质量。

图 8：Kodak、CLIC2020 和 Tecnick 数据集的量化结果。(PSNR 越高越好，LPIPS 越低越好）

图像分割的边界效应。 为了降低图像压缩过程的内存和计算要求，通常的做法是将大图像分割成较小的块，如 256 × 256，以便进行高效的单独处理。这种策略通常会造成像素连续性的差异，导致潜在特征连续性的差异，从而在重建图像中出现明显的块边界。传统上，需要额外的后处理模块（如平滑网络）来缓解这一问题。然而，在比特率极低的情况下，仅使用基于码本的压缩方法（如 VQGAN ）会产生严重的边界效应，这是由于使用有限的编码字数在视觉空间中产生了剧烈的量化效应，很难通过后处理简单解决（如图 7 中放大的图像补丁所示）。相比之下，本文的方法通过使用连续潜特征来校正像素解码过程，可以自适应地调整饱和度偏差，纠正图像内部的结构潜特征，并在不使用任何显式后处理模块的情况下大幅降低边界效应。

总结

在本文中，作者提出了为超低比特率图像压缩量身定制的双流 HybridFlow 框架。通过将连续域特征整合到离散域表示中，可以在超低比特率的同时提供高感知质量和高保真的重建图像。作者还选择性地屏蔽索引图，以进一步降低信息率。具体来说，引入了基于token的Transformer和交叉注意力模块，有效地结合了连续域的指导，能够根据部分屏蔽的索引预测完整的索引图，并保持对原始分布的保真度。作者还通过校正网络将连续域特征注入像素解码器，从而减少重建中的像素级失真，实现高感知质量和高保真。最后，作者还提供了一个可选的复杂性感知模块，为不同的图像片段选择不同的屏蔽方案，从而更有效地分配有限的比特。实验结果证明了本文的方法在各种数据集上的鲁棒性，与现有的基于码本的 LIC 方法相比，本文的方法显著提高了 PSNR，LPIPS 与之相近甚至更好，与基于连续特征的 LIC 方法相比，本文的方法显著提高了 LPIPS。本文的方法提供了一个通用的双流 LIC 框架，可在连续特征域和基于码本的特征域之间搭建桥梁，从而进一步推动超低比特率图像压缩的研究。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2024-05-16，如有侵权请联系 cloudcommunity@tencent.com 删除

框架