[论文解读] Learned Image Compression with Mixed Transformer-CNN Architectures

Jinming Liu, Heming Sun+1 more|arXiv (Cornell University)|Mar 27, 2023

Advanced Data Compression Techniques被引用 11

一句话总结

本文提出并行 Transformer-CNN 混合（TCM）块和基于 Swin-transformer 的注意力（SWAtten）的参数高效熵模型，在 Kodak、Tecnick 和 CLIC 数据集上实现相比 VVC 的最先进的速率-失真性能。

ABSTRACT

Learned image compression (LIC) methods have exhibited promising progress and superior rate-distortion performance compared with classical image compression standards. Most existing LIC methods are Convolutional Neural Networks-based (CNN-based) or Transformer-based, which have different advantages. Exploiting both advantages is a point worth exploring, which has two challenges: 1) how to effectively fuse the two methods? 2) how to achieve higher performance with a suitable complexity? In this paper, we propose an efficient parallel Transformer-CNN Mixture (TCM) block with a controllable complexity to incorporate the local modeling ability of CNN and the non-local modeling ability of transformers to improve the overall architecture of image compression models. Besides, inspired by the recent progress of entropy estimation models and attention modules, we propose a channel-wise entropy model with parameter-efficient swin-transformer-based attention (SWAtten) modules by using channel squeezing. Experimental results demonstrate our proposed method achieves state-of-the-art rate-distortion performances on three different resolution datasets (i.e., Kodak, Tecnick, CLIC Professional Validation) compared to existing LIC methods. The code is at https://github.com/jmliu206/LIC_TCM.

研究动机与目标

• 激励将局部的 CNN 建模与非局部的 transformer 建模相结合，以提升 LIC 的 RD 性能。
• 设计一个高效的并行 TCM 块，将 CNN 与 transformer 特征通过可控复杂度进行融合。
• 开发一个带有 SWAtten 和通道压缩的逐通道熵模型，在降低参数的同时维持性能。
• 在多个数据集（Kodak、Tecnick、CLIC）上展示最先进的 RD 性能并分析复杂度。
• 提供消融研究以理解 TCM 和 SWAtten 对 RD 增益的贡献。

提出的方法

• 提出一个并行的 Transformer-CNN 混合（TCM）块，将特征分成 CNN 路径和 transformer 路径，通过 1x1 卷积和拼接进行融合，并使用两阶段的 Swin Transformer（窗口化和移位窗口）进行非本地建模。
• 在主路径残差块之后以及超先验路径中加入 TCM 块，以聚合局部与非局部信息。
• 引入带有参数高效 SWAtten 模块的逐通道自回归熵模型，使用通道压缩来减少熵模型的输入通道数。
• 将熵模型的切片数量从 10 减少到 5，以平衡速度和 RD 性能，并在 SWAtten 中将通道压缩应用到 128 通道。
• 以 RD 目标函数（RD 损失）进行训练，结合潜在变量 y 与 z 的熵项以及失真项（MSE 或 MS-SSIM），并使用拉格朗日乘子 λ。
• 提供一个带有超先验和基于上下文的细化的 RD 框架，用于估计 y 的每个切片的高斯参数。

实验结果

研究问题

• RQ1 在给定复杂度下，并行 Transformer-CNN 混合（TCM）块相较于仅 transformer 的 LIC 或仅 CNN 的模型，是否能提升 RD 性能？
• RQ2 在熵模型中，带通道压缩的参数高效 SWAtten 模块能否在不大幅增加模型规模的情况下提供 RD 增益？
• RQ3 在主路径和超先验路径中都应用 TCM 块在学习型图像压缩中有哪些 RD 优势？
• RQ4 在标准 LIC 基准（Kodak、Tecnick、CLIC）上，与 VVC（VTM-12.1）在 PSNR 和 MS-SSIM 指标下相比，所提出的方法表现如何？
• RQ5 减少熵切片数量和使用通道压缩对 RD 和计算效率有何影响？

主要发现

• 该方法在 Kodak、Tecnick 和 CLIC 数据集上实现了最先进的 RD 性能，分别比 VVC（VTM-12.1）在 BD-rate 上高出 12.30%、13.71%、11.85%。
• 与现有 SOTA 方法相比，在相同比特率下，大型模型在 Kodak 上实现约 0.4 dB 的 PSNR 和 0.5 dB 的 MS-SSIM 增益。
• SWAtten 模块降低信息损失，在代表性 Kodak 样本中，将缩放后的像素平均绝对偏差 εs 从 0.451/0.422 降低到 0.389/0.365。
• 在 SWAtten 内的通道压缩策略在保持有竞争力的 BD-rate 性能的同时实现了参数和 FLOP 的减少。
• 一个 Transformer-CNN 混合（TCM）块在 RD 上优于仅 Transformer 或仅 CNN 的基线，证明了结合局部和非局部建模的好处。