1-end2end_optimized_image_compression

论文解读 END-To-END optimized image compression

本文发表在2017年的ICLR上，主要贡献是提出了一套包括分析变换（analysis　transformation）、均匀量化器（uniform quantizer）和综合变换（synthesis　transformation）的端到端图像压缩框架，并使用代理函数的形式将量化和离散熵等不可导的步骤松弛为可导的步骤，作者同时证明，在特定条件下，本文的松弛化后的损失函数在形式上类似于变分自编码器（variable autoencoder）的似然。尽管本文所提出的端到端图像压缩编码器是基于MSE优化的，但是在MS-SSIM等客观评价指标上均优于JPEG和JPEG2000，尤其在低码率的情况下，端到端的图像压缩较基于正交线性变换的压缩方法在人眼感知上更加友好。

端到端图像压缩框架

其中，$x$和$\hat{x}$分别是原图像和重建图像，$g_a$、$g_s$和$g_p$分别表示分析变换（编码）、综合变换（解码）和感知域变换，在具体实现中，$g_a$和$g_s$在结构上是相反的，如下图所示，$g_a$主要由卷积层和GDN层（Generalized Divisive Normalization）交替拼接而成，而$g_s$正好相反，图中注明了每层需要优化的参数量，而分析变换和综合变换各自所需要优化的参数集合记为$\varphi$和$\theta$，对应$y=g_a(x;\theta )$和$\hat{x}=g_s(\hat{y};\theta )$。

该压缩框架的具体过程是，原始图像$x$经过$g_a$变换到code space上的连续向量$y$，经量化后得到可传输的离散码流$q$；在接收端$q$作为输入$\hat{y}$，经过$g_s$变换到data space的$x$作为重建图像。

在实际操作中，量化操作是不可导的，因此在训练过程中，并没有进行量化，而是使用松弛的损失函数模拟量化操作的误差，具体操作是使用$\tilde{y}=y+\Delta y$代替反量化后的数据，这样做有两个考虑：

1. 当$\Delta y\in U(0,1)$，即噪声服从均匀分布，此时$\tilde{y}$和$q$在整数点上的概率密度是一样的，即连续变量$\tilde{y}$的概率密度可作为离散码流$q$的概率的插值函数，如下图所示
   

   推导如下：
   $$p_{\widetilde{y_i}}(n)=p_{y_i}\ast rect(n)={\int }_{n-\frac{1}{2}}^{n+\frac{1}{2}}{p}_{y_i}(t)d\text{t}=P_{q_i}(n)$$
   这样我们就可以使用均匀噪声模拟量化造成的误差$\Delta y$，而且该过程是可导的！（补充，此处的$i$表示不同的channel）

2. 使用均匀噪声模拟量化操作的误差已经有先例（Gray and Neuhoff, 1998）

因此，我们解决了量化操作不可导的问题，可以计算重建图像和原始图像在感知域上的差异（这里简单采用了MSE），使失真的计算过程变得可导，对应于式（7）的第二项
$$\begin{array}{rl}E[d(z,\hat{z})]& =E_{x,\Delta y}[d(g_p(\tilde{x}),g_p(x))]\\ & =E_{x,\Delta y}[d(g_p(g_s(g_a(x;\varphi );\theta )),g_p(x))]\end{array}$$
此外，计算码率的过程也是不可导的——本文将离散熵作为码率的指标，但是统计每个符号的操作是不可导的，因此作者使用线性分段函数的形式估计${\tilde{y}}_i$的分布，对应的参数是${\psi }_i$，具体操作应该是基于普通的梯度下降算法最大化似然函数：
$$L_{\psi }({\psi }^{(0)},{\psi }^{(1)},...)=-E_{\widetilde{y_i}}\sum _i{p}_{\widetilde{y_i}}(\widetilde{y_i};{\psi }^{(i)})$$

概率和似然在表达式上几乎没有区别，但是前者是参数已知，后者是参数未知。而基于梯度下降以实现最大化似然的过程也是类似的试探操作，改动参数如果能使得当前的概率更加符合实际情况，则认为参数的改动是合理的。

到此，本文已经解决了最棘手的两个问题：量化和计算熵的不可导，可以实现端到端的训练了。

损失函数

本文的损失函数分为码率和失真两部分，并通过$\lambda$平衡二者的权重，形式如下：
$$\begin{array}{rl}L(\theta ,\varphi )& =-\mathbb{E}[lo{g}_2{P}_q]+\lambda \mathbb{E}[d(z,\tilde{z})]\\ & ={\mathbb{E}}_{x,\Delta y}[-\sum _{i}lo{g}_2{p}_{\widetilde{y_i}}(g_a(x;\varphi )+\Delta y;{\psi }^{(i)})+\lambda d(g_p(x),g_p(g_s(g_a(x;\varphi )+\Delta y;\theta )))]\\ & ={\mathbb{E}}_x[-\sum _i{p}_{{\tilde{y}}_i}(g_a(x;\varphi )+\Delta y;{\psi }^{(i)})lo{g}_2{p}_{\widetilde{y_i}}(g_a(x;\varphi )+\Delta y;{\psi }^{(i)})+\lambda d(g_p(x),g_p(g_s(g_a(x;\varphi )+\Delta y;\theta )))]\end{array}$$
此外，作者解释从data space到perceptual space的变换$g_p$可以是固定的（a fixed transform），因为并没有需要优化的参数。

值得注意的是，上式的优化列表中并没有$\psi$的出现，作者在附录6.1的末尾解释了$\psi$的优化过程，每次反向传播时更新下面的负似然函数：
$$L_{\psi }({\psi }^{(0)},{\psi }^{(1)},...)=-{\mathbb{E}}_{\tilde{y}}\sum _i{p}_{{\tilde{y}}_i}({\tilde{y}}_i;{\psi }^{(i)})$$
而且，每10⁶ 次梯度更新后，更新分段线性函数$p_{{\tilde{y}}_i}$的值。

理论推导：与变分自编码器的关系

变分自编码器的优化目标如下所示：
$$\begin{array}{rl}& D_{KL}(q(y|x)||p_{y|x})\\ & ={\mathbb{E}}_{y\sim q}\log q(y|x)-{\mathbb{E}}_{y\sim q}\log p_{y|x}\\ & ={\mathbb{E}}_{y\sim q}\log q(y|x)-{\mathbb{E}}_{y\sim q}\log p_{x|y}(x|y)+{\mathbb{E}}_{y\sim q}\log p_y(y)-{\mathbb{E}}_{y\sim q}\log p_x(x)\end{array}$$
显然，最后一项为常数，作者提出，如果满足以下三个条件，前文的优化目标类似于变分自编码器的优化目标：
$$p_{\mathbf{x}\mid \tilde{\mathbf{y}}}(\mathbf{x}\mid \tilde{\mathbf{y}};\lambda ,\mathbf{\theta })=\mathcal{N}\left(\mathbf{x};g_s(\tilde{\mathbf{y}};\mathbf{\theta }),(2\lambda {)}^{-1}1\right)$$

$$p_{\tilde{\mathbf{y}}}\left(\tilde{\mathbf{y}};{\mathbf{\psi }}^{(0)},{\mathbf{\psi }}^{(1)},\dots \right)=\prod _i{p}_{{\tilde{y}}_i}\left({\tilde{y}}_i;{\mathbf{\psi }}^{(i)}\right)$$

$$q(\tilde{\mathbf{y}}\mid \mathbf{x};\mathbf{\varphi })=\prod _i\mathcal{U}\left({\tilde{y}}_i;y_i,1\right)\text{ with }\mathbf{y}=g_a(\mathbf{x};\mathbf{\varphi })$$

其中，满足条件一需要避免使用感知域变换$g_p$且失真度量$d$必须是二范数，满足条件二意味着$\tilde{y}$各通道相互独立，条件三则对应proxy function。

此时，（6）中的第一项和第四项都是常数，第二项对应失真，第三项对应码率。

尽管本文提出的端到端优化框架和变分自编码器（VAE）有诸多相似之处，但是两者仍然存在本质的区别：

1. VAE是基于连续值优化的，而实际压缩和传输是在离散域进行的，而本文则在训练时使用离散熵模拟码率，在推理时使用熵编码计算码率，更加符合实际情况；
2. VAE生成模型在优化时过于关注数据的微小波动，而本文的模型则可以实现失真-码率的平衡，忽略不重要的细节；而且在VAE生成模型的中间结果进行离散化，这可能产生意料之外的结果，但是也有一些研究（van den Oord, Kalchbrenner, and Kavukcuoglu, 2016 ）
3. 从生成模型的视角看，尽管（7）对应于端到端模型在较简单的度量形式下的失真，但是这种特殊形式的概率密度可能导致模型的选择不够普遍，从而难以达到最优。

实验结果

1. 松弛后的结果和实际结果的对比

作者在训练数据集上，分别比较了采用代理函数和量化引起的误差，以及微分熵和实际中的离散熵，结果如下所示

可见，使用均匀噪声替代量化操作引起的误差是很成功的，但是使用分段线性函数估计实际的$p_{\tilde{y}}$存在一定的偏差，但是也要看到随着量化程度越精细，偏差得到了缓解。作者给出的建议是，$\lambda$仅作为平衡各项的参数，可以在推理时修正这种偏差即可，即使用$\lambda +\Delta \lambda$。

2. 与JPEG和JPEG2000的对比

作者首先在客观评价指标 MS-SSIM（结构相似性，一种偏向人类感知的度量）和PSNR上与JPEG和JPEG2000进行了对比，结果如下所示

可见，本文提出的方法基本超越了JPEG和JPEG2000。

作者也给出了可视化的结果对比，同一码率下的结果对比，从上到下依次为JPEG、Ours和JPEG2000：

从主观感受上，本文所提出方法的压缩效果在低码率下的感受是明显好于此前传统编码的压缩效果。

作者还给出了不同码率下三种方法的压缩效果：

这个结果相当惊艳，端到端的压缩效果在不同码率上的效果十分类似，将对人眼不重要的细节成功“挑选”了出来，作者将其归功于端到端的学习方式在生物学意义上与大脑视觉皮层对图片的编码类似。

在附录中，作者还介绍了在模型推理时使用的经典熵编码算法CABAC（算术编码的一种形式），此处略。

压缩效果在不同码率上的效果十分类似，将对人眼不重要的细节成功“挑选”了出来，作者将其归功于端到端的学习方式在生物学意义上与大脑视觉皮层对图片的编码类似。

在附录中，作者还介绍了在模型推理时使用的经典熵编码算法CABAC（算术编码的一种形式），此处略。