[论文解读] Checkerboard Context Model for Efficient Learned Image Compression

Dailan He, Yaoyan Zheng+3 more|arXiv (Cornell University)|Mar 29, 2021

Advanced Data Compression Techniques参考文献 32被引用 16

一句话总结

本文提出了一种可并行化的棋盘状上下文模型（CCM），用于学习型图像压缩，通过将解码顺序重新组织为两轮处理，实现了高效、高速的解码。与传统自回归上下文模型必须严格串行处理不同，CCM 使用棋盘形状的卷积，允许上下文特征的并行计算，与最先进模型相比，在码率-失真性能上损失可忽略不计的同时，实现了超过40倍的加速。

ABSTRACT

For learned image compression, the autoregressive context model is proved effective in improving the rate-distortion (RD) performance. Because it helps remove spatial redundancies among latent representations. However, the decoding process must be done in a strict scan order, which breaks the parallelization. We propose a parallelizable checkerboard context model (CCM) to solve the problem. Our two-pass checkerboard context calculation eliminates such limitations on spatial locations by re-organizing the decoding order. Speeding up the decoding process more than 40 times in our experiments, it achieves significantly improved computational efficiency with almost the same rate-distortion performance. To the best of our knowledge, this is the first exploration on parallelization-friendly spatial context model for learned image compression.

研究动机与目标

• 为解决学习型图像压缩中自回归上下文模型存在的计算效率低下问题，此类模型依赖严格串行解码，阻碍了实时部署。
• 开发一种上下文建模方法，在保持高码率-失真性能的同时，实现在解码过程中完全并行化。
• 提供一种即插即用的替代方案，无需修改模型架构或容量，即可替代现有的串行上下文模型。
• 通过克服自回归上下文建模的解码瓶颈，实现高性能学习型编解码器的实际部署。

提出的方法

• 提出一种棋盘形状的卷积，将潜在特征划分为两个交错的集合（锚点和非锚点），以实现上下文计算的并行化。
• 设计一种两轮解码方案：首先解码所有锚点位置，然后并行计算所有非锚点位置的上下文。
• 使用掩码卷积确保因果性，即在上下文建模过程中仅关注已解码（可见）的邻居。
• 采用与标准自回归模型相同的上下文模型结构，但通过重新排序计算过程，消除顺序依赖。
• 通过仅替换上下文建模模块，与现有学习型图像压缩框架保持兼容。
• 采用标准的熵编码和超先验建模，主自编码器和超先验组件无需任何修改。

实验结果

研究问题

• RQ1 能否在不牺牲码率-失真性能的前提下，使学习型图像压缩的空间上下文模型实现完全并行化？
• RQ2 何种解码策略能够实现在自回归模型中高效并行计算上下文特征？
• RQ3 通过重新排序潜在表示的计算顺序，解码速度最多可提升多少？
• RQ4 所提出的方法是否可作为现有串行上下文模型在最先进压缩架构中的可行即插即用替代方案？

主要发现

• 所提出的棋盘状上下文模型在并行硬件（Nvidia TITAN XP）上实现了解码速度超过40倍的提升，相比传统串行上下文模型。
• 该模型保持了几乎相同的码率-失真性能，与原始串行模型相比，PSNR和MS-SSIM仅出现微小下降。
• 在Kodak数据集上，基于CCM的模型在Minnen2018和Cheng2020的BDBR指标上分别优于无上下文的超先验基线17.0%和27.4%。
• 该方法在多个数据集上表现稳健，在Kodak和Tecnick数据集上均表现出一致的性能增益，证实了其泛化能力。
• 两轮解码方案成功解耦了上下文计算，实现了完全并行化，同时保持了因果性和模型精度。
• 该方法与现有SOTA架构兼容，可作为插件式替换直接部署，无需修改架构。