Efficient Contextformer：利用空间通道窗口注意力机制实现可学习图像压缩中的快速上下文建模

来源：IEEE Transactions on Circuits and Systems for Video Technology 题目：Efficient Contextformer: Spatio-Channel Window Attention for Fast Context Modeling in Learned Image Compression 作者：A. Burakhan Koyuncu, Panqi Jia, Atanas Boev, Elena Alshina, Eckehard Steinbach 原文链接：https://ieeexplore.ieee.org/abstract/document/10453616 内容整理：刘潮磊 熵估计对学习型图像压缩的性能至关重要。研究表明，基于Transformer的熵模型在实现高压缩比方面具有重要意义，但代价是显著的计算开销。在这项工作中，我们引入了高效的上下文Transformer（eContextformer）——一种计算效率高的基于Transformer的自回归上下文模型，用于可学习图像压缩。eContextformer有效融合了按块、棋盘和按通道分组的技术进行并行上下文建模，并引入了移位窗口空间-通道注意机制。我们探索了更好的训练策略和架构设计，并引入了额外的复杂度优化。在解码过程中，所提出的优化技术动态缩放注意范围，并缓存先前的注意力计算，大幅减少模型和运行时的复杂度。与非并行方法相比，我们的方案模型复杂度降低约145倍，解码速度提高约210倍，并在Kodak、CLIC2020和Tecnick数据集上实现了更高的平均比特节省。此外，我们的上下文模型的低复杂度使得我们支持在线率失真算法，从而进一步提高压缩性能。相比于Versatile Video Coding (VVC) 测试模型 (VTM) 16.2的帧内编码，我们实现了高达17%的码率节省，并超越了各种基于学习的压缩模型。

引言

研究问题：对于可学习的图像压缩，如何构建一个计算高效的上下文模型。

SOTA 工作与所属团队：

• “Parallelized context modeling for faster image coding”, A. B. Koyuncu, K. Cui, A. Boev, and E. Steinbach
• “Checkerboard context model for efficient learned image compression”, D. He, Y. Zheng, B. Sun, Y. Wang, and H. Qin
• “Contextformer: A transformer with spatio-channel attention for context modeling in learned image compression”, AB Koyuncu, H Gao，ECCV 2022

动机：本文作者的前一篇文章构建了一个基于transformer的上下文模型，这篇文章采用对latent tensor分组、优化训练策略等方式降低模型复杂度。

贡献：

• 采用棋盘格加划分channel的方式，利用了通道之间的相关性
• 探究并行化策略
• 其他的复杂度优化策略

方法

Contextformer

架构：

作者前一篇文章构建的一个基于transformer的上下文模型架构如下，主体是Multi-Head Attention，和ViT类似：

图 1

对latent的划分：

1. 将通道划分成

段，一幅图像共有

个序列

1. 对于这些序列有两种处理方式：空间优先（sfo）、通道优先（cfo），如下图所示：

图 2

1. 为降低复杂度，使用3D sliding-window，窗口只在空间上滑动，在通道上延申，示意图如下：

图 3

Contextformer总体思想：

• 对latent划分
• sliding window attention，不同处理方式采用不同的mask

并行化策略探究（在Contextformer 基础上）

并行策略：棋盘格+patch-wise+channel-wise

图 4

参数设置：

1. 训练图像大小：256x256（latent tensor：16x16）
2. Spatial kernel size 𝐾：16、8（ 3D sliding-window 在空间上的大小）

：16、8、4、1（训练/测试时window的步长）

1. sfo/cfo（序列处理方式）

参数对：

1. 训练时全局attention：

1. Windows是否有重叠：

（无重叠）；

（有重叠）

探究结果：

图 5

表 1

表 2

结论：

• 训练时全局attention：sfo优于cfo
• 训练时窗口有overlap(

)：sfo优于cfo

本文的eContextformer(efficient Contextformer)

总体架构：

将Contextformer主体从ViT形式的sliding window替换成了swin-transformer形式

图 6

latent tensor分组处理流程：

图 7

eContextformer的改进：

相较于Contextformer，eContextformer的改进可以总结如下：

表 3

其他的复杂度优化策略

EGR(Efficient coding Group Rearrangement)：亦即上述对lantent tensor的分组策略。

SFG(Skipping First coding Group)：在对lantent tensor分组后，第一组只用hyperprior编码，不设置start token。

key-value caching：第n次迭代时需要对输入的

所有序列计算Attention：

由于前n-1次已经计算过：

于是将之前的结果保留，只需计算：

eContextformer架构细节

eContextformer各部分结构参数总结如下表所示：

表 4

实验设计与验证

实验设置

数据集：Train：Vimeo-90K Evaluate ：Kodak、CLIC2020 test dataset、Tecnick

训练细节：Image crops：256x256， Batch size：16， MSE作为失真指标

实验结果

在Kodak数据集上测试结果：

图 8

探究在高分辨率图像上的效果：

本文用COCO 2017的256×256/384×384 image crops微调，测试了在高分辨率图像上的效果，如下图所示，左图为Kodak数据集上测试结果，右图为Tecnick数据集上测试结果，可以看出，在高分辨率数据集Tecnick上有明显的提升。

图 9

在多个数据集上测试结果：

下图从左至右分别为Kodak、CLIC2020、Tecnick数据集上的测试结果：

图 10

模型复杂度和运行时间：各种复杂度优化策略的效果如下(Single pass是在编码端一次性计算整个

）：

表 5

模型参数量和计算复杂度如下：

表 6

编码和解码的时间如下：

表 7

消融实验：

表 8

视觉效果展示：

图 11

结论

这项工作引入了eContextformer——一个高效且快速的优化后的Contextformer。我们进行了大量的实验，以实现一个快速且低复杂度的上下文模型，这个模型还具有和SOTA相当的性能。值得注意的是，我们提供的算法优化进一步将复杂度降低了84%。为了缩小差距，我们还实验了编码端迭代算法。这进一步提高了率失真性能，并且其复杂度仍然低于最先进的视频压缩标准。本研究专注于提供一个高效的上下文模型架构，对于这种改进的变换留待未来工作进行探讨。