[论文解读] Contextformer: A Transformer with Spatio-Channel Attention for Context Modeling in Learned Image Compression

A. Burakhan Koyuncu, Han Gao+4 more|arXiv (Cornell University)|Mar 4, 2022

Advanced Data Compression Techniques被引用 4

一句话总结

本文提出 Contextformer，一种基于 Transformer 的上下文模型，采用时空-通道注意力机制用于学习型图像压缩，可联合建模潜在空间中的空间与跨通道依赖关系。通过在现代压缩框架中替换标准上下文模型，Contextformer 在 VTM 16.2 的基础上实现了最高达 11% 的码率节省，并在 Kodak、CLIC2020 和 Tecnick 数据集上均优于现有学习型方法，在 PSNR 和 MS-SSIM 指标上表现更优。

ABSTRACT

Entropy modeling is a key component for high-performance image compression algorithms. Recent developments in autoregressive context modeling helped learning-based methods to surpass their classical counterparts. However, the performance of those models can be further improved due to the underexploited spatio-channel dependencies in latent space, and the suboptimal implementation of context adaptivity. Inspired by the adaptive characteristics of the transformers, we propose a transformer-based context model, named Contextformer, which generalizes the de facto standard attention mechanism to spatio-channel attention. We replace the context model of a modern compression framework with the Contextformer and test it on the widely used Kodak, CLIC2020, and Tecnick image datasets. Our experimental results show that the proposed model provides up to 11% rate savings compared to the standard Versatile Video Coding (VVC) Test Model (VTM) 16.2, and outperforms various learning-based models in terms of PSNR and MS-SSIM.

研究动机与目标

• 解决在学习型图像压缩的熵建模中，潜在空间内时空-通道依赖关系未被充分利用的问题。
• 通过用动态、基于注意力的机制替代非自适应的掩码卷积，提升上下文自适应能力。
• 设计一种高效的推理流水线，在不增加模型大小或无需微调的情况下保持高性能。
• 在率失真性能上超越当前最先进学习型压缩模型及 VVC 测试模型（VTM）16.2。

提出的方法

• 提出一种新颖的时空-通道注意力机制，将标准自注意力推广至联合建模潜在张量中的空间与通道依赖关系。
• 基于 Transformer 架构设计上下文模型 Contextformer，其中每个 token 会关注潜在表示中的空间与通道相关特征。
• 采用多头注意力机制，结合可学习的查询、键与值，以根据输入模式动态调整上下文信息。
• 应用两种算法优化：分块解码（BDS）与分段编码策略（SCS），通过支持并行处理降低编码器运行时间。
• 在解码器中采用波前解码（wavefront decoding）以实现并行化处理，减少解码时间，相比 3D 上下文模型实现 9 倍加速。
• 在保持与先前工作相同模型架构与参数量的前提下，通过架构创新与推理优化提升性能。

实验结果

研究问题

• RQ1 基于 Transformer 的上下文模型若能联合建模空间与跨通道依赖关系，是否可提升学习型图像压缩中的率失真性能？
• RQ2 在熵建模精度方面，时空-通道注意力相较于仅关注空间或仅关注通道的注意力机制表现如何？
• RQ3 能否在不损失性能或无需微调的前提下，降低基于注意力机制的上下文模型的计算开销？
• RQ4 编码顺序与通道分段对信息分布与模型性能有何影响？
• RQ5 与 VTM 16.2 及其他学习型模型相比，所提模型在 PSNR、MS-SSIM 与码率节省方面的表现如何？

主要发现

• 在 Kodak、CLIC2020 与 Tecnick 数据集上，Contextformer 相较于 VTM 16.2 实现了平均 6.9%–10.5% 的 BD-Rate 收益，即码率节省。
• 在所有评估的学习型基线方法中，该模型在 PSNR 与 MS-SSIM 指标上均表现更优，展现出卓越的率失真性能与感知质量。
• 采用通道优先编码策略（cfo）时，Contextformer 将超过 70% 的总信息量存储在前两个通道段中，表明其具备强大的跨通道依赖建模能力。
• 当通道段数为 4 时，模型在性能与复杂度之间达到最优平衡；进一步增加段数仅带来边际性能增益，但训练成本显著上升。
• 通过所提出的 BDS 与 SCS 优化，尽管像素数量增加 20 倍，从 Kodak（低分辨率）扩展到 4K 图像时，编码器运行时间仅增加 3 倍。
• 波前解码使解码速度相比 3D 上下文模型提升 9 倍，充分证明了并行化推理策略的有效性。