论文解读-Variational Image Compression With a Scale Hyperprior

论文解读-Variational Image Compression With a Scale Hyperprior

本文发表在ICLR2018，作者Balle沿袭了ICLR2017上的工作，增加了可以提取边信息的Hyperprior，进一步消除了图像在latent space上的空间相关性，同时改进了此前的分段线性函数拟合边缘概率分布的做法。本文所提出的方法在PSNR上接近HEVC BPG，而在SSIM上则超过了HEVC BPG。

Hyper Prior

需要解决的问题

在此前的框架中，解决微分熵不可导的操作是设计了一个entropy model $p_{\hat{y}}(\hat{y})$以逼近图像$x$在latent space上的实际边缘概率分布$m(\hat{y})$，并尝试降低压缩后的码率，即
$$R={\mathbb{E}}_{\hat{y}\sim m}[-{\log }_2{p}_{\hat{y}}(\hat{y})]=-\sum _{i}m(\widehat{y_i}){\log }_2{p}_{\widehat{y_i}}(\widehat{y_i})$$
当模型预测的分布$p_{\hat{y}}(\hat{y})$和实际分布$m(\hat{y})$一致时，该交叉熵才能准确朝最小化的方向优化，否则可能导致模型收敛到次优点。

出于使模型预测的概率分布更加接近于特定图片的概率分布（在latent space上），作者以HEVC通过引入块划分表示边信息而取得较JPEG更好的压缩效果为例（作者认为块划分结构体现了边的信息），猜测通过引入边信息可以使模型学习到更好的概率分布以用于熵编码实现更高的压缩效率，即传递边信息所需的码率较将$p_{\hat{y}}$刻画得更加精细而言更少。

边信息本质上也可看做是entropy model的一种先验，因此作者在ICLR2017的可分离先验（factorized prior）上增加了可捕获在latent representation的不同尺度上元素的空间相关性的超先验（Hyperprior）。

作者进行了预实验，可视化了图片在latent space上的特征图

其中，从左到右分别是原图、factorized prior处理后的latent representation、Hyperprior预测得到的标准差、图二经去除标准差后的分布。可见，factorized prior确实没有完全学习到部分纹理和边缘信息，从而不能完全去除离散信号的空间耦合，这一方面并不利于模型学习概率分布（因为不如完全高斯化后的图四更容易学习到概率分布），从而不容易找准最佳的梯度下降方向；另一方面，我认为，从信息论的角度思考，均匀分布的图片能够实现更高的压缩率。

带有Hyperprior的端到端压缩框架

作者提出的带有Hyperprior的端到端压缩框架如下图所示

作者沿袭了此前的基于代理函数近似量化过程的操作，并证明了优化先验概率q和真实的边缘概率p的KL散度等价于优化码率和失真，具体证明过程可参考ICLR2017中的证明端到端的压缩方法在某种约束下等价于变分自编码器。上图中的U|Q表示训练时施加均匀噪声，测试时进行量化操作。

具体操作是在latent representation $\hat{y}$上又进行了一次编解码，得到$\hat{z}$，作者对该过程的formulation如下所示：
$$\begin{array}{r}{\mathbb{E}}_{\mathbf{x}\sim p_{\mathbf{x}}}{D}_{\mathrm{K}\mathrm{L}}\left[q\Vert p_{\tilde{\mathbf{y}},\tilde{\mathbf{z}}\mid \mathbf{x}}\right]={\mathbb{E}}_{\mathbf{x}\sim p_{\mathbf{x}}}{\mathbb{E}}_{\tilde{\mathbf{y}},\tilde{\mathbf{z}}\sim q}[\log q(\tilde{\mathbf{y}},\tilde{\mathbf{z}}\mid \mathbf{x})-\log p_{\mathbf{x}\mid \tilde{\mathbf{y}}}(\mathbf{x}\mid \tilde{\mathbf{y}})\\ -\log p_{\tilde{\mathbf{y}}\mid \tilde{\mathbf{z}}}(\tilde{\mathbf{y}}\mid \tilde{\mathbf{z}})-\log p_{\tilde{\mathbf{z}}}(\tilde{\mathbf{z}})]+\text{ const. }\end{array}$$
其中第四项即代表边信息。

网络架构如下所示

其中左半边对应ICLR2017中的factorized model，右半边对应Hyperprior model，AE和AD表示算数熵编码器和熵解码器，卷积层的标识为“卷积核数×Height×Width×Stride”，而N和M需要根据$\lambda$决定，对于5个低$\lambda$值，$N=128M=192$，对于3个高$\lambda$值，$N=192M=320$，附录2中对此进行了探讨。

$\hat{\sigma }$表示标准差的空间分布，用于$\hat{y}$的压缩和重建。

实验

本文取消了感知域变换，分别以PSNR和MS-SSIM为失真指标进行训练，发现训练和测试时损失函数是否一致会显著影响本文提出的端到端图像压缩的效果，作者在两类评价指标上均证明了这点

从图中可以发现：

1. 训练和测试时的评价指标是否一致会显著影响模型评价效果
2. 就通用的PSNR指标而言，当前Hyperprior model接近于HEVC BPG的性能，但是仍有不足

下图中，从上到下分别为在较低码率下的基于SSIM和MSE的重建图像，其中篮框表示SSIM优于MSE的区域，而红框表示MSE优于SSIM的区域

可见，SSIM倾向于低对比度区域的重建质量，但是语义上的相关性应该更多被分配到文字等较为突出的物体上，而基于MSE的训练目标可以更好的匹配这一点。作者强调，两种metric在捕捉语义上都不够优秀。

作者还给出了边信息所占码率和总码率的对比，如下图所示，可见编码边信息增加的码率很少

讨论

1. 单变量非参数概率密度模型

ICLR2017中Balle等人采用分段线性函数拟合$m(\hat{y})$，在操作中这存在两个问题：

1. 非零概率值必须是有限且提前知道
2. 对于现有的自动微分器，实现该操作并不容易

因此本文基于概率密度的积分实现了新的非参数概率密度模型。

首先定义概率密度$p$和累积概率密度$c$，并被分解为以下形式
$$c=f_K\circ f_{K-1}\cdots f_1$$

$$p=f_K^{\prime }\cdot f_{K-1}^{\prime }\cdots f_1^{\prime }$$

其中，$f_K$为矩阵${\mathbb{R}}^{d_k}\to {\mathbb{R}}^{r_k}$，而$f_K^{\prime }$为$f_K$的一阶偏导（雅各比矩阵），$\circ$表示矩阵乘法。此处，为保证概率密度服从定义（位于0-1之间的非负数），要求雅各比矩阵的元素为非负数。

$f_k$被设计为如下的形式
$$f_k(\mathbf{x})=g_k\left({\mathbf{H}}^{(k)}\mathbf{x}+{\mathbf{b}}^{(k)}\right)1\le k<K$$

$$f_k(\mathbf{x})=\text{sigmoid}\left({\mathbf{H}}^{(k)}\mathbf{x}+{\mathbf{b}}^{(k)}\right)k=K$$

其中的$g_k$被设计为如下形式
$$g_k(\mathbf{x})=\mathbf{x}+{\mathbf{a}}^{(k)}\odot \tanh (\mathbf{x})$$
其中，$\odot$表示elementwise乘法，$a^{(k)}$的正负可用来控制是否伸缩零值附近的空间。

整体的运算如下所示
$$\begin{array}{rl}{f}_k^{\prime }(\mathbf{x})& =\mathrm{diag}{g}_k^{\prime }\left({\mathbf{H}}^{(k)}\mathbf{x}+{\mathbf{b}}^{(k)}\right)\cdot {\mathbf{H}}^{(k)}\\ g_k^{\prime }(\mathbf{x})& =1+{\mathbf{a}}^{(k)}\odot {\tanh }^{\prime }(\mathbf{x})\\ f_K^{\prime }(\mathbf{x})& ={\mathrm{sigmoid}}^{\prime }\left({\mathbf{H}}^{(K)}\mathbf{x}+{\mathbf{b}}^{(K)}\right)\cdot {\mathbf{H}}^{(K)}\end{array}$$
此处为限制雅各比矩阵非负以及$a^{(k)}$大于-1，需要进行参数重整：
$$\begin{array}{rl}{\mathbf{H}}^{(k)}& =\mathrm{softplus}\left({\hat{\mathbf{H}}}^{(k)}\right)\\ {\mathbf{a}}^{(k)}& =\tanh \left({\hat{\mathbf{a}}}^{(k)}\right)\end{array}$$
下图展示了本文使用三层的非参数概率密度模型$p$拟合一个混合高斯分布，并实现了很好的拟合效果，灰色线条展示了拟合的收敛过程

在本文中，始终设置K=4，并实现了和分段线性函数一样好的性能，同时归于基于自动微分器实现该操作更加友好。

2. 模型性能瓶颈与$\lambda$的关系

作者探究了$\lambda$与模型表达性能的关系，此处通过改变卷积核的个数N和M以改变模型表达性能（实验中保证N=M），发现对于确定码率，增加模型容量可提高重建质量，但是存在一个上限，例如当码率等于0.5时，卷积核个数从64提升为128会显著提升性能，但是再提升卷积核个数并不能提升性能；增大$\lambda$，应该换用容量更大的模型。

作者还发现，$M>N$这种更“宽”的结构可实现更好的性能。

当码率等于0.5时，卷积核个数从64提升为128会显著提升性能，但是再提升卷积核个数并不能提升性能；增大$\lambda$，应该换用容量更大的模型。

作者还发现，$M>N$这种更“宽”的结构可实现更好的性能。