ELIC图像压缩算法理论基础及详细解析

信息量

感性上, 一个具体事件 发生所包含的信息量与 发生的概率 成反比, 即一个事件发生的概率越小, 它发生所包含的信息量越大。在获取信息量具体定义 的过程中, 我们发现信息量的定义 必须满足以下三个条件:

1. ;
2. 信息量与概率成反比, 即 ;
3. 两个独立事件 和 同时发生的信息量等于两者信息量之和, 和 同时发生的概率等于两者概率乘积, 即 ; 我们发现对数的性质刚好满足上述三个条件, 所以将事件 发生包含的信息量 定义为:

其中 一般取 2 , 这样信息量的单位就是 bit。

信息熵

用于衡量一个不确定事件 的不确定程度或预估包含的信息量, 也代表无损编码事件信息的最小平均编码长度。

假设不确定事件 存在多种可能的结果 , 每种结果产生对应的概率为 , 则 的信息熵 定义为:

KL 散度

也叫相对熵。用于衡量两个概率分布之间的差异。其值越大, 则两个概率分布的差异越大; 当两个概率分布完全相等时相对熵值为 0 。

散度和交叉熵的关系:

KL 散度和交叉熵都不具备对称性。

交叉嫡

在已知随机变量 的条件下随机变量 的不确定性。

我们训练模型常用的交叉熵损失就来源于此, 这里 就是真实分布, 就是预测分布。熵是真实值最小平均编码长度, 因此交叉熵 熵。

KL 散度可以被用于计算代价，而在特定情况下最小化 KL 散度等价于最小化交叉熵。而交叉熵的运算更简单，所以常用交叉熵来当做代价。

熵编码

摘编码即编码过程中按熵原理不丢失任何信息的编码。熵编码需要知道独立分布的待编码变量及其概率。常见的熵编码有: 香农(Shannon)编码、哈夫曼(Huffman)编码、算术编码 (Arithmetic Coding)、非对称数字系统编码(ANS)等。

图像压缩中的 和 一般用的算术编码或非对称数字系统编码, 详细原理这里不做详述。

自编码器

输入无标签数据从中提取有效低维特征, 并利用低维特征重构出输入数据。

自编码器包含编码器和解码器两部分, 编码器负责从输入数据 中提取有效低维特征 ,解码器根据 重构出输入数据 , 最后通过计算 与 的相似度作为损失函数监督训练。

不难发现，图像压缩和解压的过程本质就是自编码器结构，压缩就是提取图像中的低维重要特征，方便传输和存储，解压则是根据压缩特征恢复原始图像的过程。所以图像压缩算法的架构基于自编码器，而自编码器还应用于数据可视化分析、以图搜图等。

变分自编码器(VAE)

自编码器 (AE) 存在的问题:

自编码器本质上通过训练数据学到的是 “输入图片-特征编码-重构图片” 之间固定的映射关系。

如图 2 所示, 编码器输出为隐式变量 (特征编码) 代表每一维特征的编码值 , 代表隐式变量 的特征维度, 代表输入样本数, 代表 对应的特征编码空间,编码空间内包含所有样本的离散编码值以及其他无效的编码值, 每一维编码空间 代表提取到的一种有效特征, 这里假设该编码空间代表月亮的盈满特征。输入一张 “满月” 图片, 编码器会输出对应编码值并映射到编码空间 的位置 1 , 然后将位置 1 编码值输入给解码器就生成了对应的重构 “满月” 图片, 类似的操作适用于 “半月” 图片。因为这些类似分布的图片自编码器在训练中 “遇到过” , 并且将 “输入图片-特征编码-重构图片” 之间的映射关系存储在特征编码内部, 所以可以 “按图索骥” 重构出输入图片。

但是如果我们输入一张 “四分之一月” 的图片或者直接在编码空间取位置 1 和 3 之间的位置 2 编码值, 输入解码器是否能重构出输入图片呢? 此时解码器由于缺少对应 “特征编码 -重构图片” 对应的映射关系, 两个相邻编码值之间也不存在任何关联, 模型也无从知道位置 2 编码值的重构结果应该介于位置 1 和 3 编码值之间, 输出结果就会出现模糊或乱码。这也说明了自编码器存在泛化性差且不具备生成能力的缺陷。

对于自编码器 (AE), 上只有与训练样本对应位置的编码值才能用于解码生成对应的重构图片, 其他位置的编码则是无效的, 因为模型没有学到对应的映射关系。为了使得自编码器 (AE) 具备图像生成或强泛化性的能力, 变分自编码器 (VAE) 应运而生。

在编码阶段, 变分自编码器 (VAE) 为了能利用训练数据构建编码空间任意位置 (比如图 2 中位置 2 编码）的映射关系, 不再直接学习 “输入图片-特征编码-重构图片” 的离散映射关系, 而是尝试将训练数据经编码器输出的编码值用覆盖整个编码空间 的高斯分布 (均值 , 方差 ) 代替。

在解码阶段, 由于现在编码器输出不是一个固定编码值 而是一个高斯分布, 就需要在这个高斯分布上随机采样, 以生成一个带噪声的编码值 用于解码。因为采样操作导致在训练时无法进行梯度传递, 所以利用重参数技巧引入采样输入 , 用来在编码器输出的高斯分布上采样一个编码值 。

这样做的关键在于解码阶段产生的编码值 是通过随机采样在高斯分布上产生的, 相比自动编码器 ( 的固定编码值 添加了偏移噪声 是从训练数据中学习到的。在训练过程中 的取值会以高斯分布概率出现在高斯分布的任意位置, 越靠近均值的位置几率越大。这样解码器就有机会学习到固定位置之外编码与重构图片的映射关系。

变分自编码器详细的正向传播过程见图 4。其中 代表图片样本数, 输入的经过卷积操作后的特征图维度假设为 5 , 均值为 , 方差为 。

在训练阶段, 除了和自编码器一样的重构损失之外, 为了方式 VAE 退化为 , 还需要对编码器输出的高斯分布进行约束。对于生成模型, 我们最终的目的是在训练完成之后丢掉编码器, 而通过在关于 的固定先验分布 (一般为标准正太分布) 手动采样得到编码后送入解码器生成全新的图片, 然后在编码阶段训练编码器让输出的高斯分布尽量接近 。这样做的好处是保证不同样本经编码器输出的均值约束在 0 附近, 防止模型在隐空间中的编码相互远离, 并鼓励尽可能多的返回分布发生 “重叠”, 从而满足预期的连续性和完整性条件。方差约束在 1 附近, 与重构误差构成对抗, 防止噪声在训练过程中被消解掉。

以上为变分自编码器的理论框架，在变分自编码器的损失函数推导之前，需要重点理解 这几个概率密度函数的定义及联系。

图 6 所示为解码器生成图片的概率论基础 (以下图文不区分大小写, 表示编码器参数 代表解码器参数) 。我们设计变分自编码器的目的是为了生成与输入图片类似的新图片,所以我们从解码器出发, 先验、后验区分也基于解码器。

是人为规定的编码空间的先验概率分布, 为了方便一般为标准正态分布, 即 为多维的话则 代表每一维编码空间的概率分布。训练完成后可根据 和解码器生成需要的新图片。

是生成图片单个像素值的概率分布。随机变量 代表像素值, 取值范围为 ,该分布是通过在 上对 连续多次采样送入解码器得到 , 然后对所有 求和生成的混合高斯分布。即 对应着所有采样 值下的多个先验高斯分布的混合, 该分布较为复杂, 见图 6 黑色曲线。因为我们想要生成图片与原始图片相似, 所以该分布在原始图片对应像素值处概率最大, 临近区域其次, 越远离该值的地方概率越小。

是用基于 或 采样到的 值生成关于 的先验概率密度函数,一般建模为高斯分布。图中每个绿色高斯分布曲线对应一个不同的采样 值。

是基于训练样本已知 值生成 的后验概率密度函数, 该过程即为编码过程, 与 互逆。VAE 的核心就是求解 。

在 VAE 推理阶段, 解码器 参数已经训练完成, 此时从标准正态分布 中采样得到具体 值输入解码器, 会输出以具体像素值为均值的高斯分布, 取均值作为最终像素值即可得到新图片。当上述采样的 值发微小偏移时, 理论上输出像素均值也应该接近偏移前输出的像素均值, 即满足连续性。简单来说, 训练好的参数构建起了 和 一一对应的关系。

在 VAE 训练阶段, 解码器 输出的目标像素均值已知, 的参数未知待学习,所以像素均值对应的具体 值也无从得知。相当于有了结果, 需要逆推出对应原因, 这就需要求解 , 即根据已知 值求解对应 的分布, 需注意该分布属于 的子集, 因为我们只有训练样本中特定的像素值 , 无法获取 对应的全域 。

在训练阶段, 样本像素值 和 已知, 分布均值已知, 和 未知, 以上所以分布都需要转化为以 为基础。

根据全概率公式得到 表达式:

根据贝叶斯定理和全概率公式得到 的表达式:

理论上可以通过 和 求得, 但是观察 的表达式, 分母中需要对 进行全域积分, 本身也是高维向量, 高维积分很难求解, 所以 分布复杂无法直接求解,只能选择变分推断的近似方法逼近。这里我们选择用神经网络去拟合一个函数 , 输入样本像素值 , 直接输出 的对应分布而不再需要知道 的分布。我们可以将 建模为简单的概率分布 (一般为高斯分布), 从而简化问题的处理难度。这样在训练阶段, 输入样本像素值 到编码器 生成对应 的高斯分布, 然后基于该分布采样一个具体 值送入解码器 得到该 值对应 的高斯分布。

所以我们需要 与 的分布尽量接近, 计算两者的 KL 散度:

由于输入 给定, 所以 是定值, 那么最小化 KL 散度就是最小化变分下限 ELBO。

上述 ELBO 推导过程中, 过程(2)即为 VAE 的优化目标: 第一项计算 和 散度并希望越小越好, 即编码器输出分布不断接近先验分布 。第二项计算 和 的交叉嫡, 是从原始数据中学习到的分布, 本质也是保证重构图片的分布和原始图片保持一致。

过程(1)是下一章节超先验变分图像压缩模型的理论基础, 其中第一项 一般为均匀噪声概率密度恒等于 1 所以可忽略。第二项与 VAE 相同, 计算 和 的交叉熵,当都为高斯分布时, 这一项就是原始图片和重构图片的 MSE 损失。第三项计算 和 的交叉熵, 代表使用熵模型 实际的最小编码长度。

超先验变分图像压缩模型

图像压缩一般用 指标衡量, 即压缩率, 即重构损失。在特定情况下, VAE 的损失刚好等价于 R-D 损失, 所以尝试用 VAE 架构来实现图像压缩, 详细推导见上一章节 VAE 损失函数的推导。

超先验图像压缩模型在 VAE 的基础上, 发现隐变量 之间存在空间依赖不是完全独立,这样通过熵模型编解码后的结果就不是最优, 所以又引入一套编解码结构获取隐变量 的空间关系。

压缩变换时: 输入图像 经过分析变换 后压缩成隐藏空间表示 经过量化编码为 ,同时捕捉 中元素服从高斯分布的 , 经过 变换生成 , 通过量化编码生成 ;

重构变换时: 通过对 解码后 变换生成重构 需要的 , 新的 经过 变换生成重构图像 。

根据 VAE 推导出损失函数定义:

式子第一项为 0 , 因为 为均匀分布。

第二项代表重构损失。

第三、四项代表相应交叉嫡, 代表编码长度损失。

ELIC

本文在基于超先验变分图像压缩模型架构的基础上, 提出对编码特征进行空间和通道两个维度融合的上下文结构 SCCTX，与超先验信息融合之后共同预测编码特征之间存在的依赖关系用于摘模型编解码。其中通道维度还进行不均匀分组并按照信息重要程度顺序存储,空间维度采用两阶段棋盘格方法加速。

改进 1

项目要求满足随机渐进压缩（随机丢包均能解压出图片）, 原版 ELIC 虽然也支持渐进压缩, 但是限制只能末端丢包, 不满足要求。为此我们进行第一项改进, 将 ELIC 的通道维度上下文中通道间的联系 “斩断” (图 9 中红色虚线), 这样通道间不存在依赖关系, 当前解压包就不需要前一个解压包信息才能解压。

改进 2

针对输出的压缩包大小不够均匀的问题, 将通道数全部改为 64 , 保证各个包的理论容量相同。在损失函数中对不同压缩包的 bpp 损失计算方差并加入损失函数监督训练, 结果有效保证了各个压缩包的大小均匀。

改进 3

针对 MSE-optimized 模型在 bpp 较小时出现的模糊问题, 在损失函数中引入 LPIPS 损失,该损失更加注重肉眼主观视觉效果, 最终在 bpp 变化不大的情况下解压图片的清晰度有较大提升。根据参考文献 5、6, 还可以引入 GAN-loss 来进一步增强重构图片的真实视觉效果。

参考文献

1. Ballé J, Minnen D, Singh S, et al. Variational image compression with a scale hyperprior[J]. arXiv preprint arXiv:1802.01436, 2018.
2. He D, Zheng Y, Sun B, et al. Checkerboard context model for efficient learned image compression[C]/Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 14771-14780.
3. Minnen D, Singh S. Channel-wise autoregressive entropy models for learned image compression[C]/2020 IEEE International Conference on Image Processing (ICIP). IEEE, 2020: 3339-3343.
4. He D, Yang Z, Peng W, et al. Elic: Efficient learned image compression with unevenly grouped space-channel contextual adaptive coding[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 5718-5727.
5. He D, Yang Z, Yu H, et al. PO-ELIC: Perception-Oriented Efficient Learned Image Coding [C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 1764-1769 
6. Agustsson E, Minnen D, Toderici G, et al. Multi-realism image compression with a conditional generator [C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 22324-22333.
7. 计算机视觉与深度学习 北京邮电大学 鲁鹏 清晰版合集 (完整版)
8. https://zhuanlan.zhihu.com/p/628604566
9. https://mp.weixin.qq.com/s/pvebLZa38BQrdNA4KGfCiw
10. https://zhuanlan.zhihu.com/p/144649293?utm_id=0
11. https://zhuanlan.zhihu.com/p/514477269
12. https://www.cnblogs.com/zblade/p/14338758.html
13. https://zhuanlan.zhihu.com/p/55791843
14. https://zhuanlan.zhihu.com/p/83865427
15. https://www.zhihu.com/question/513622870/answer/2985031683
16. https://www.jianshu.com/p/7d0c8fb026a4

公众号后台回复“数据集”获取100+深度学习各方向资源整理