端到端图像压缩《Checkerboard Context Model for Efficient Learned Image Compression》

论文地址：http://arxiv.org/abs/2103.15306

##代码地址：https://github.com/JiangWeibeta/Checkerboard-Context-Model-for-Efficient-Learned-Image-Compression

商汤出品，这两年商汤在端到端图像压缩领域展露头脚，综合实力很强，2022的CLIC夺得冠军。

一 简介

       提出一种两阶段的Checkerboard Context Model，重新定义了解码的顺序，代码非官方部署。

二 内容

2.1 目前方法的缺陷

        整体框架是如上图，但是本文主要针对上图中Contex Model模型进行优化，论文joint2019以及cheng2020都采用了基于Mask Convolution的自回归模型，自回归模型能够极大地提高模型的性能，但是由于严格的解码顺序，会导致解码速度虽特征图的空域大小呈指数增加，为了探索熵模型并行化的可能，本文提出了一种两阶段的上下文模式。
       如上图，所有的小格子表示量化后的潜在表示点$\hat{y}$，蓝色点表示在解码过程中，已经解码的特征点，红色表示当前正在解码的点，黄色表示解码当前点时用到的已解码的点。之前的工作均采用(a) (b) 图捕获上下文信息，(a)和(b)仅卷积核的大小不同，无本质差别，本文则采用基于©(d)的卷积核的上下文模型。

2.2 探索上下文“距离”的影响

       论文joint2019中实验了几种简单的卷积核作为Context Model的卷积候选，有结论即5x5 Mask Convolution相比于7x7和3x3尺寸的卷积核性能更优，我们可以理解5x5比3x3卷积核的性能更好，但是7x7的效果会差于5x5则有点奇怪，为此，作者研究了对当前解码点不同距离的特征点对当前解码点的码率贡献。在训练阶段，通过对5x5的卷积核进行随机采样，即采用Random Mask Convolution，如下图所示，在sampling space中，对5x5的卷积核进行随机的掩码，具体的随机参数论文中没有披露，暂时认为是25个空域位置的权重有0.5的概率会被掩蔽置0。因此，权重在使用不同掩码的上下文模型之间共享，以便训练的随机掩码模型可用于在推理过程中评估任意掩码。 训练后，为了测试特定掩码模式的性能，只需将该掩码作为 M 输入上下文模型，然后随机掩码模型具有具有固定掩码的上下文模型。通过掩盖逐个空域位置的权重邻近像素的距离对码率点的影响，对码率的节省率如下：

$R_0$ 表示不采用已经解码的信息作为上下文，$R_M$表示随机掩码。

       图(a) 表示距离周围像素对于当前像素的码率节省的贡献，通过由图(b) 所示的码率掩码获取的，测试的时候每次对于Random Mask Convolution，会白色部分的权重全部掩盖，进保留一个空域位置的权重，以此测试该位置对于节省码率的贡献，得到最终24个位置的贡献程度，可以得出结论，其空域距离越近，相关性越高，对码率的影响也越大。

2.3 Checkerboard 上下文模型

       网络整体结构和之前Google的论文一致，在上下文模块进行修改，将上下文模块调整为两阶段的上下文模式。


图5表示了具体的两阶段的Checkerboard的具体流程。

1. 图5中，黄色的图表示为图4中，Hyper Decoer的输出，这个输出经过$g_{ep}$的变换过程得到如${\Phi }_{anchor}$显示的蓝灰色位置，可以看出每个蓝色的位置都有一定的间隔，$g_{ep}$输出的就是对应这部分位置的潜在表示$y$的概率密度参数，例如高斯分布中的$\mu$和$\sigma$，在得到这些位置的概率密度的参数后，就可以对这些蓝色的点进行熵解码，解码后得到${\hat{y}}_{anchor}$，可以看出，${\hat{y}}_{anchor}$中橙色的点和蓝色的点位置是一致的，这部分的解码有两个特点，i） 不会利用的已经解码点的信息。ii）近解码了一部分的位置，对于灰色位置的特征点仍然没有解码。
2. 第二阶段的解码过程则利用第一阶段得到的解码点作为上下文信息，即如图5所示的${\hat{y}}_{anchor}$，利用第一阶段得到的解码点，以及$g_{ep}$的输出信息，获取灰色点的概率密度参数，得到参数后，对灰色的点进行解码，得到${\hat{y}}_{non-anchor}$，最后组合${\hat{y}}_{anchor}$和${\hat{y}}_{non-anchor}$点，得到最终完整的$\hat{y}$，这就是Checkboard Context的全过程。

三 性能

3.1 失真性能

可以从Cheng2020的性能（绿色）和cheng2020 用了checkerboard的性能对比看出，checkerboard的方式基本不会导致模型的性能下降，并且在高码率点的时候，反而超过了anchor的性能。

3.2 时间复杂度性能

对比采用Mask Convolution的Context Model，Checkerboard的Context Model大幅度加速了解码速度，提高了接近50倍，而且如果编码的图像分辨率越高，则可以推测两种Context的解码时间倍数会相差更多。

四 结论

       很棒的工作，后续也有很多人follow这篇论文对上下文进一步地增强，并且之前的Mask Convolution从代码层面上解码端比较难实现，Checkerboard能够有效地保持base的RD性能，并且极大地降低了解码的复杂度，在代码层面上也比较容易复现。
        文章从空域角度加速了解码端的Context Model，对比$论{文}^1$采用通道域的自回归模型，在本质上其实是一致的，特征图是可以在通道域和空域相互转换的，例如可以采用卷积核控制步长得到预设特征图的分辨率和通道数量，完成简单的转换，或者类似超分辨领域的PixelShuffle层，可以将四个通道的特征图拼接为一个，完成空域分辨率的两倍放大。所以换个角度，是不是可以把Chekerboard Context理解为一种两阶段的通道自回归模型？