[论文解读] The Devil Is in the Details: Window-based Attention for Image Compression

Renjie Zou, Chunfeng Song+1 more|arXiv (Cornell University)|Mar 16, 2022

Advanced Data Compression Techniques被引用 8

一句话总结

本文提出一种基于窗口的局部注意力机制，通过提升局部纹理重建能力来增强学习图像压缩，解决了CNN在捕捉非重复纹理方面的固有局限。该方法可无缝集成至CNN与Transformer架构中，在率失真性能方面达到当前最优水平，尤其在MS-SSIM优化下，显著提升了视觉质量与细节保留能力。

ABSTRACT

Learned image compression methods have exhibited superior rate-distortion performance than classical image compression standards. Most existing learned image compression models are based on Convolutional Neural Networks (CNNs). Despite great contributions, a main drawback of CNN based model is that its structure is not designed for capturing local redundancy, especially the non-repetitive textures, which severely affects the reconstruction quality. Therefore, how to make full use of both global structure and local texture becomes the core problem for learning-based image compression. Inspired by recent progresses of Vision Transformer (ViT) and Swin Transformer, we found that combining the local-aware attention mechanism with the global-related feature learning could meet the expectation in image compression. In this paper, we first extensively study the effects of multiple kinds of attention mechanisms for local features learning, then introduce a more straightforward yet effective window-based local attention block. The proposed window-based attention is very flexible which could work as a plug-and-play component to enhance CNN and Transformer models. Moreover, we propose a novel Symmetrical TransFormer (STF) framework with absolute transformer blocks in the down-sampling encoder and up-sampling decoder. Extensive experimental evaluations have shown that the proposed method is effective and outperforms the state-of-the-art methods. The code is publicly available at https://github.com/Googolxx/STF.

研究动机与目标

• 解决基于CNN的端到端图像压缩在重建非重复局部纹理方面的固有局限。
• 探究局部感知注意力机制是否在图像压缩任务中优于全局注意力。
• 设计一种灵活、即插即用的注意力模块，以增强基于CNN与Transformer的压缩模型。
• 提出一种新型对称Transformer（STF）框架，其编码器与解码器均采用绝对Transformer模块，以提升重建保真度。
• 通过优化高对比度与细节区域的比特分配，提升感知质量，而无需完全依赖PSNR或MS-SSIM指标。

提出的方法

• 提出一种基于窗口的注意力机制（WAM），通过在局部空间窗口内计算自注意力，以捕捉短程空间相关性。
• 将WAM作为即插即用模块集成至现有的CNN与Transformer架构中，包括STF框架。
• 设计对称Transformer（STF），在编码器与解码器中均采用相同的绝对Transformer模块，以实现对称的特征学习。
• 在固定空间窗口内使用通道-wise 多头自注意力，以在保持计算效率的同时保留局部纹理细节。
• 采用SGM-based 信道模型实现高效的比特流编码，并在Transformer模块中使用层归一化。
• 采用率失真损失进行端到端训练，通过调整超参数以同时优化MSE与MS-SSIM目标。

实验结果

研究问题

• RQ1 与全局注意力相比，局部感知注意力机制是否能提升学习图像压缩中非重复纹理的重建效果？
• RQ2 在率失真性能与视觉质量方面，基于窗口的注意力机制与非局部注意力相比表现如何？
• RQ3 基于窗口的注意力模块在多大程度上可泛化并成功集成至基于CNN与Transformer的压缩模型中？
• RQ4 与非对称或标准Transformer设计相比，采用对称Transformer模块的所提对称Transformer（STF）框架是否能提升重建保真度？
• RQ5 在不同优化目标（如MSE与MS-SSIM）下，所提方法的性能在感知质量方面有何差异？

主要发现

• 在λ=0.0035时，与基线CNN模型（Minnen2020）相比，WAM将PSNR提升0.16 dB，同时将bpp降低0.003。
• 在λ=0.0130时，WAM增强模型达到34.10 dB PSNR与0.448 bpp，优于基线模型（33.87 dB，0.454 bpp）与非局部注意力变体（33.95 dB，0.467 bpp）。
• 采用窗口注意力的STF模型在视觉质量方面表现更优，尤其在文本与边缘等高对比度区域，如定性对比所示。
• 消融实验表明，WAM在不同λ值下均能持续提升性能，证明其鲁棒性与有效性。
• 所提方法实现了当前最优的率失真性能，尤其在MS-SSIM优化下取得显著提升，表明其感知质量更优。
• 基于窗口的注意力机制兼容CNN与Transformer架构，作为即插即用组件展现出强大的泛化能力。