CVPR2020:基于GMM和Attention的图像压缩的论文学习

在这里插入图片描述
论文地址：https://arxiv.org/pdf/2002.01657.pdf
代码地址：https://github.com/ZhengxueCheng/Learned-Image-Compression-with-GMM-and-Attention

在这里插入图片描述
现在基于深度学习的图像压缩方法有着快速的发展同时也有着令人满意的结果，但是在和现行压缩标准之间仍然存在性能差距，尤其是在广泛使用的PSNR度量方面。
本文探讨了最近的深度学习图像压缩算法中的存在的冗余问题，作者提出了建议使用离散的高斯混合模型来对分布进行参数化，从而消除了剩余的冗余以实现准确的熵模型，从而直接导致所需的编码位数更少。此外，网络架构中采用了注意力模块的简化版本，注意力模块可以使学习的模型更加关注复杂区域，从而以中等训练复杂度提高我们的编码性能。
在这里插入图片描述

baseline 框架具体来自于Balle论文《End-to-end optimized image compression》
Hyperprior（超先验模型）框架同样来自于google实验室Balle出品《Variational Image Compression With A Scale Hyperprior》
Joint 是基于Hyperprior之上优化的，由Minnen发表的《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》
本文提出的在 c) 的基础上，优化了对于潜在的特征层的建模模型，从单高斯分布的形式优化为混合高斯分布的形式。

离散混合高斯

在这里插入图片描述

Baseline模型，一个基准模型，类似于auto-encoder自编码器的一个过程，就是原始图像输入映射到一个latent space，一个潜在空间的过程，便于量化和熵编码，然后再将信息进行解码的过程。
Hyperprior模型是Balle根据Baseline改进而提出的，在前期熵编码过程中，我们需要知道编码的tensor张量的具体值以及所对应的概率，常规编码比如算术编码（Arithmetic coding），霍夫曼编码（Huffman coding），他的编码原理就是对源符号进行编码，出现概率高的字母使用较短的编码，反之出现概率低的则使用较长的编码，这样编码之后的熵就会变小，占的储存空间就会变小。
但是对于第一种模型，也就是baseline来说，他的编码的值和他所对应的概率虽然是可以训练的，但是只有在训练过程中他是可以改变的，训练完成后，他的这个概率就确定好了，无论你输入图片x是什么样的，它所对应的概率都是固定的，就会产生很大的冗余。
Hyperprior是一个双层结构，我们假设他的（latent representation）隐藏层是独立同分布的，然后通过知道他的这个隐层的特征值服从一个什么分布，例如拉普拉斯分布，高斯分布等等就可以计算出他的理论熵。这里的第一种模型baseline是离散的，一个点对应一个点的概率，而hyperprior则是隐层服从具体的概率公式，训练时它可以是连续的，但在编码时也可以进行离散化。
在这里隐层特征y是在z的基础上进行估计的，其中y是服从一个均值为0，方差为δ的单高斯分布进行估计的，隐层的概率分布参数是可以根据外层去学习的。
Joint相较于前者，高斯分布概率公式中均值不为零了，也变成了一个可学习的参数，此外又引入了一个context model 进行自回归模型（PixelCNN）的预测，与带有均值与方差的高斯分布联合对隐层特征y进行估计。
本文作者认为，单高斯并不能很好的估计隐层的特征分布，所以提出了一个混合高斯的模型，因为高斯本身是处理连续数值的，但是隐层特征量化后是离散的，所以采用离散的混合高斯模型，包括权重、均值、方差三个参数，都是可以学习的。

Hyperprior模型，Joint模型和作者提出的熵估计模型(离散高斯混合模型)的熵估计的可视化结果：
在这里插入图片描述
推理阶段高斯混合：

简化注意力模型

在论文中，增加了简易注意力机制模块，通过1x1的卷积核增强通道间的信息交流。并且通过sigmoid函数生成重要性掩膜，使得网络更注重复杂纹理区域的特征，然后作者简化掉了Non-local block，这个过程中包含了矩阵乘法，运算过程很费时间。
在这里插入图片描述
简化注意力模型：

RB结构：

性能结果

在整体的网络中，对比《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》中的模型，该整体框架将1个5x5的卷积核分解成2个3x3的卷积核，这种卷积核的变换在保持感受野范围的同时又减少了参数量，有一定的性能增益。
在这里插入图片描述

在PSNR性能指标上已经获得了和VVC相近的性能效果，MS-SSIM指标则传统的编码器效果好像比起端到端的压缩方式一直以来都差了一点，总的来说离散高斯混合模型在较低的计算cost的给整个编码框架带来了比较大的性能增益。
在这里插入图片描述
结论：作者提出了一种使用离散高斯混合似然和注意模块的学习图像压缩方法。通过对现有压缩技术剩余冗余的研究，我们发现单参数化模型不能实现任意似然，限制了熵模型的精度。
因此，我们使用离散化的高斯混合似然来实现更灵活和准确的熵模型，以提高性能。此外，我们在我们的网络架构中使用了一个简化的、中等复杂度的注意力模块，以达到较高的编码效率。
实验结果表明，与现有的学习压缩方法和编码标准HEVC、JPEG2000、JPEG相比，本文提出的方法达到了最先进的性能，并且实现了与下一代PSNR压缩标准VVC相当的性能。