2017-ICLR-End-to-end Optimized Image Compression

1. 摘要

本文^[1]提出了一种图像压缩框架，其由一个非线性分析变换、一个均匀量化器和一个非线性综合变换组成。这些变换都三层结构，每一层由一个卷积线性过滤器和非线性的激活函数构成。与大多数卷积神经网络不同的是，这些联合的非线性是用来实现一种局部增益控制，这种局部增益控制最初是来源于生物神经元领域的发现（局部增益控制已经成为生物神经元建模的一种方法）。作者使用随机梯度下降方法来联合优化整个模型的率失真性能。针对均匀标量量化器的目标函数不连续问题，作者提出了一种连续可微的近似量化，从而保证在优化过程中梯度能有效地反传。最后在测试集上，作者提出压缩模型展现出比 JPEG 和 JPEG2000 更好的性能。而且，该压缩模型在所有码率和所有图像上，MS-SSIM 性能都要好于 JPEG 和 JPEG2000。

2. 引言

数据压缩一直以来都是一个工程基础问题，自香农提出熵信息以来，数据压缩问题得到很好的研究。对于给定的离散数据，最优的压缩编码可以通过最小化编码熵来得到。而这个最优的压缩编码严重依赖于数据本身的概率分布，因此设计一个最优的压缩编码本质也是对数据分布进行建模的问题。然而，实际的编码必须具有有限的熵，因此连续值数据必须量化成一系列有限的离散值，而这会导致误差的引入。在这篇文章中，作者研究的便是有损压缩问题，针对有损压缩，我们必须平衡两个互斥的损失：离散编码的熵（率）和由量化引起的误差（失真）。因此最终要优化的目标函数即率失真函数。对于不同的压缩场景，比如数据存储或有限容量信道的数据传输，对不同的率失真 trade-off 要求也不同。

然而，联合优化率失真是很困难的。如果没有进一步的约束的话，这个问题的最优量化在高维空间是难以求解的。因此，大多数现有的图像压缩方法都是先通过一个线性变换将图像数据变换到合适的连续值表示，然后独立地量化这些连续值，最后使用一个无损的熵编码来编码量化得到的离散值。这个框架一般称为变换编码，因为变换在这个框架中起到决定性作用。比如，JPEG 使用 DCT 变换来编码像素块，而 JPEG2000 使用多尽度的正交小波变换来解耦图像块。通常，变换编码的三个组成部分：变换、量化、熵编码，往往是独立地进行优化的（而且往往是通过手动设置参数进行调整的）。

作者在本文中完善了他原先在另一篇文章^[2]中提出的一个基于非线性变换的端到端压缩框架，完善的框架如下图所示：

在他先前那篇文章^[2:1]中，他并没有详细介绍上述框架中的分析综合变换结构。在本文中，作者对此进行了进一步介绍，给出了该框架下的一个实际模型，并给出了充足的理论分析和实验来验证该模型的有效性。在此模型中，分析变换作者采用了级联线性卷积层和非线性层的结构，非线性层采用的是 GDN 模型（来源于作者的另一文章^[3]）；而综合变换则是采用了分析变换的近似逆结构。量化使用的是均匀的标量量化。

对于期望逼近的率失真曲线上的每个点，都直接使用梯度下降优化方法同时联合优化分析综合变换。对于均匀标量量化导致梯度几乎处处为 $0$ 的问题，作者给出了一种连续的近似均匀标量量化来替代原来的离散均匀量化。最后，作者对量化后的数据使用了一个熵编码得到实际的比特流，从而计算该模型实际的码率，以此来验证该模型的有效性。

3. 分析、综合以及感知变换的选择

大多数压缩方法都是建立在正交线性变换基础上，目的是为了降低数据间的相关性，从而简化后续的熵编码。但是线性变换输出之间的联合统计特性展现了很强的高阶依赖。这个问题可以通过使用局部增益控制操作（最早来源于生物神经元领域）来很好地解决，于是作者使用了他自己在另一篇文章^[3:1]中提出的 GDN 模型来替代线性变换（在那篇文章^[3:2]中作者已经验证了 GDN 具有很好的高斯化图像数据的能力）。

注意，GDN 本质也是一种归一化方法，但和经常在神经网络里看到的 BN 不是同一类型的归一化。BN 是针对一个 batch 里的数据进行归一化和尺度操作，且一旦神经网络训练完成后，BN 的尺度参数也固定了，此是就完全变成一个关于归一化数据的仿射变换。其数学公式描述如下：

$$\begin{align*} y_i = \gamma \hat{x}_i + \beta \\ \hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \varepsilon}} \\ \mu_B = \frac{1}{n} \sum_{i=1}^n x_i \\ \sigma_B = \frac{1}{n} \sum_{i=1}^n (x_i - \mu_B)^2 \end{align*} \tag{1}$$

而 GDN 的归一化和尺度操作对于图像所有位置的数据都是一致的，相当于一种空间自适应的归一化操作，且 GDN 是非线性的。其数学公式描述如下：

$$\begin{align*} \boldsymbol{y} = g(\boldsymbol{x};\boldsymbol{\theta}) && \text{ s.t. } &&& y_i = \frac{z_i}{(\beta_i + \sum_j \gamma_{ij} |z_j|^{\alpha_{ij}})^{\varepsilon_i}}\\ && \text{ and } &&& \boldsymbol{z} = \boldsymbol{H} \boldsymbol{x} \end{align*} \tag{2}$$

更为具体的，作者在本文中提出的模型的分析综合变换结构均由三阶结构组成，具体如下图所示：

• 分析变换：由三阶结构组成，每阶由一个卷积层+下采样层+ GDN 组成。设 $u_i^{(k)} (m, n)$ 表示第 $k$ 阶第 $i$ 通道上坐标为 $(m, n)$ 的数据，则 $u^{(0)} = \boldsymbol{x}, \boldsymbol{y} = u^{(3)}$。每一阶的三层结构用数学公式描述如下：

  $$\begin{align*} v_i^{(k)}(m, n) = \sum_{j} (h_{k, ij} * u_j^{(k)})(m, n) + c_{k, i} \\ w_i^{(k)}(m, n) = v_i^{(k)} (s_k m, s_k n) \\ u_i^{(k+1)}(m, n) = \frac{w_i^{(k)}(m, n)}{\left(\beta_{k, i} + \sum_j \gamma_{k, ij}\left(w_j^{(k)}(m, n)\right)^2\right)^{\frac{1}{2}}} \end{align*} \tag{3}$$

  其中，$h, c, \beta, \gamma$ 为分析变换的所有参数。

• 综合变换：由三阶结构组成，每阶由一个 IGDN + 上采样层 + 卷积层组成。设 $\hat{u}_i^{(k)}(m, n)$ 表示第 $k$ 阶第 $i$ 通道上坐标为 $(m, n)$ 的数据，则 $\hat{u}^{(0)} = \hat{\boldsymbol{x}}, \hat{\boldsymbol{y}} = \hat{u}^{(3)}$。每一阶的三层结构用数学公式描述如下：

  $$\begin{align*} \hat{w}_i^{(k)}(m, n) = \hat{u}_i^{(k)}(m, n) \cdot \left( \hat{\beta}_{k, i} + \sum_{j} \hat{\gamma}_{k, ij} \left(\hat{u}_j^{(k)}(m, n)\right)^2 \right)^{\frac{1}{2}} \\ \hat{v}_i^{(k)} (m, n) = \begin{cases} \hat{w}_i^{(k)} (m/\hat{s}_k, n/\hat{s}_k) & m/\hat{s}_k, n/\hat{s}_k \in \mathbb{Z} \\ 0 & \text{otherwise} \\ \end{cases} \\ \hat{u}_i^{(k+1)} (m, n) = \sum_j (\hat{h}_{k, ij} * \hat{v}_j^{(k)})(m, n) + \hat{c}_{k, i} \end{align*} \tag{4}$$

  其中，$\hat{h}, \hat{c}, \hat{\beta}, \hat{\gamma}$ 为综合变换的所有参数。

4. 优化非线性变换编码模型

4.1 模型优化

要优化本文提出的压缩模型，便是要优化率失真目标函数：$R + \lambda D$，$\lambda$ 是用来控制率失真之间的 trade-off。作者已经给出了分析综合变换的结构，Fig 1 框架中还有 $g_p$ 和量化没有确定。对于量化部分，作者这里没有考虑试图在图像数据上找到最优的量化器，而是直接使用固定的均匀标量量化器。通过设计一个合适的熵编码，实际的码率只比理论的熵略大。作者给出的优化目标函数如下：

$$L[g_a, g_s, P_{\boldsymbol{q}}] = -\mathbb{E}[\log_2{P_{\boldsymbol{q}}}] + \lambda \mathbb{E}[d(\boldsymbol{z}, \hat{\boldsymbol{z}})] \tag{5}$$

其中，期望值在训练时是使用 batch 的均值来近似的。不失一般性，作者假定所用的标量量化器为均匀的，且量化区间为 $1$。

对于量化区间不为 $1$ 的，我们可以通过调整分析变换和综合变换的系数来缩放量化区间到 $1$；对于非均匀量化，我们可以通过压缩扩展来将其转变为均匀量化。

从而量化过程用数学公式描述如下：

$$\hat{y}_i = q_i = \mathrm{round}(y_i) \tag{6}$$

其中 $i$ 遍历所有的数据向量的所有下标（包括坐标位置和通道）。$\hat{y}_i$ 的概率密度由一系列的离散累积概率给出（见 Fig 2），如下公式所示：

$$P_{q_i}(n) = \int_{n-\frac{1}{2}}^{n+\frac{1}{2}} p_{y_i}(t) dt, \qquad \forall n \in \mathbb{Z} \tag{7}$$

需要注意的是，式 $(5)$ 右边的两项都依赖于量化后的值。然而式 $(6)$ 的导数几乎处处为 $0$，无法直接用梯度下降算法来反传梯度进行优化。因此作者提出了离散标量量化一种连续的近似，即将离散量化替换成加性 $U(0, 1)$ 均匀噪声 $\Delta \boldsymbol{y}$。这样，式 $(6)$ 就变成如下公式：

$$\tilde{\boldsymbol{y}} = \boldsymbol{y} + \Delta \boldsymbol{y} \tag{8}$$

这样做有两个好处：

• 第一，$\tilde{\boldsymbol{y}}$ 的概率密度经过 $\hat{\boldsymbol{y}}$ 的所有离散累积概率，如 Fig 2 所示，用数学公式描述如下：

  $$p_{\tilde{\boldsymbol{y}}}(\boldsymbol{n}) = P_{\boldsymbol{q}}(\boldsymbol{n}), \qquad \forall \boldsymbol{n} \in \mathbb{Z}^M$$

  这说明 $\tilde{\boldsymbol{y}}$ 的可微熵的确可以作为 $\boldsymbol{q}$ 的熵的近似。

• 第二，独立的均匀噪声以其边缘矩近似量化误差，因此往往被用来作为量化误差的模型^[4]。

这部分涉及到对量化器的建模与分析，可能需要看看文章^[4:1]。

通过将离散标量量化替换为连续的近似标量量化，目标损失函数式 $(5)$ 最终可写成如下形式：

$$L(\boldsymbol{\theta}, \boldsymbol{\phi}) = \mathbb{E}_{\boldsymbol{x}, \Delta \boldsymbol{y}} \left[ -\sum_i \log_2{p_{\tilde{y}_i}} \left(g_a(\boldsymbol{x}; \boldsymbol{\phi}) + \Delta \boldsymbol{y}; \boldsymbol{\psi}^{(i)}\right) + \lambda d\Big(g_p\big(g_s(g_a(\boldsymbol{x}; \boldsymbol{\theta}) + \Delta \boldsymbol{y}; \boldsymbol{\theta})\big), g_p(\boldsymbol{x})\Big) \right] \tag{9}$$

其中，向量 $\boldsymbol{\psi}^{(i)}$ 参数化了 $p_{\tilde{y}_i}$ 分段线性近似（见下文熵模型），在训练时和参数 $\boldsymbol{\theta}, \boldsymbol{\phi}$ 一同进行优化。

4.2 和 VAE 模型的关系

作者从经典的率失真优化问题推出最终目标函数 $(9)$。他发现，一旦过渡到连续的损失函数，优化问题便类似于拟合图像生成模型遇到的问题，更具体地说是可以转化为处理 VAE 模型^[5]时遇到的问题。在变分贝叶斯推断中，随机变量 $x$ 的一个观测集合以及生成模型 $p_{x | y}(x|y)$ 是给定的，我们需要找到一个后验模型 $p_{y|x}(y|x)$，但这个后验模型往往无法求出解析解。前人文章^[5:1]中采用的方法是使用一个概率密度 $q(y|x)$ 来近似，然后最小化 $q(y|x)$ 和 $p_{y|x}(y|x)$ 之间的 KL 散度：

$$\begin{align*} D_{KL}\left[q \Vert p_{y|x}\right] & = \mathbb{E}_{y \sim q} \log{q(y|x)} - \mathbb{E}_{y \sim q} \log{p_{y|x}(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)} + C \tag{10} \end{align*}$$

其中，常数 $C = \mathbb{E}_{y \sim q} \log{p_x(x)}$。

如果使用 MSE 损失，并分别定义生成模型为如下形式：

$$p_{\boldsymbol{x}|\tilde{\boldsymbol{y}}}(\boldsymbol{x}|\tilde{\boldsymbol{y}}; \lambda, \boldsymbol{\theta}) = \mathcal{N}(\boldsymbol{x}; g_s(\tilde{\boldsymbol{y}}; \boldsymbol{\theta}), (2\lambda)^{-1} \boldsymbol{1}) \\ p_{\tilde{\boldsymbol{y}}} (\tilde{\boldsymbol{y}}; \boldsymbol{\psi}^{(0)}, \boldsymbol{\psi}^{(1)}, \cdots) = \prod_i p_{\tilde{y}_i}(\tilde{y}_i; \boldsymbol{\psi}^{(i)})$$

近似的后验模型为如下形式：

$$q(\tilde{\boldsymbol{y}} | \boldsymbol{x} ; \boldsymbol{\phi})=\prod \mathcal{U}\left(\tilde{y}_{i} ; y_{i}, 1\right) \quad \text { with } \boldsymbol{y}=g_{a}(\boldsymbol{x} ; \boldsymbol{\phi})$$

其中，$\mathcal{U}(\tilde{y}_i; y_i, 1)$ 表示 $[y_i - 0.5, y_i + 0.5]$ 区间上的均匀密度。则式 $(10)$ 中的第一项为常数，第二项对应失真，第三项对应码率。可视化表示如下：

注意，如果使用了感知变换 $g_p$ 或者非欧拉距离度量，则 $p_{\boldsymbol{x} | \tilde{\boldsymbol{y}}}$ 就不再是高斯模型了，此时式 $(10)$ 就可能不再是变分自编码器了，因为此时失真项不再是规范化的密度模型。

对任何仿射可逆的感知变换 $g_p$ 和平移不变的度量 $d$，$p_{\boldsymbol{x} | \tilde{\boldsymbol{y}}}$ 可表示为如下形式：

$$p_{\boldsymbol{x} \mid \tilde{\boldsymbol{y}}}(\boldsymbol{x} \mid \tilde{\boldsymbol{y}} ; \lambda, \boldsymbol{\theta})=\frac{1}{Z(\lambda)} \exp \Big(-\lambda d\big(g_{p}\left(g_{s}(\tilde{\boldsymbol{y}} ; \boldsymbol{\theta})\right), g_{p}(\boldsymbol{x})\big)\Big) \tag{11}$$

其中，$Z(\lambda)$ 规范化了式 $(11)$ 对应的概率模型。虽然这个变换编码的压缩模型很类似 VAE 模型， 但作者还是强调了一下二者不的不同之处：

• 首先，VAE 是用来处理连续数据的，而变换编码的模型是用来处理离散数据的；
• 其次，生成模型的目标在于最小化给定模型下数据集全的微分熵，也即意味着最大化率失真函数 $(5)$ 中的参数 $\lambda$；而变换编码模型的目标在于求得码率和失真间的一个最好的 trade-off。因此变换编码模型需要优化各种 $\lambda$ 值对应的场景，而生成模型的目标是优化当 $\lambda = \infty$ 时对应的场景，此时相当于是一种无损压缩模型，见图 Figure 2。
• 最后，变换编码中可能使用感知变换来针对人眼视觉进行优化，但在生成模型中一般不会加入一个这样感知变换度量。

5. 熵模型

5.1 训练熵模型

虽然作者提出了一种离散量化的近似连续量化操作，但即使使用连续的量化操作，仍然不知道量化的数据所服从的分布 $p_{\tilde{\boldsymbol{y}}}$，我们只能得到量化后的数据 $\tilde{\boldsymbol{y}}$，因此还是无法直接计算出目标损失函数中的码率项，也就无法优化损失函数。为此，作者又提出了一种近拟计算码率的方法，即为每个 $p_{\tilde{y}_i}$ 设定一个可学习的参数向量 $\boldsymbol{\psi}^{(i)}$，其中的元素看成是从分布 $p_{\tilde{y}_i}$ 中采样得到的，采样密度为每单位长度采样 $10$ 个点。然后通过使用 SGD 最大化如下负期望似然来优化参数 $\boldsymbol{\psi}^{(i)}$：

$$L_{\boldsymbol{\psi}}(\boldsymbol{\psi}^{(0)}, \boldsymbol{\psi}^{(1)}, \ldots)=-\mathbb{E}_{\tilde{\boldsymbol{y}}} \sum_{i} p_{\tilde{y}_{i}}(\tilde{y}_{i} ; \boldsymbol{\psi}^{(i)}) \tag{12}$$

在计算 $p_{\tilde{y}_{i}}(\tilde{y}_{i} ; \boldsymbol{\psi}^{(i)})$ 时，将 $p_{\tilde{y}_i}$ 表达成分段线性函数，过参数向量 $\boldsymbol{\psi}^{(i)}$ 中的那些点。然后将 $\tilde{y}_i$ 代入即可计算得到 $p_{\tilde{y}_{i}}(\tilde{y}_{i} ; \boldsymbol{\psi}^{(i)})$。整个近似过程都是可微的，故能够使用梯度下降来优化 $\boldsymbol{\psi}^{(i)}$ 和码率。

5.2 实际熵模型

为了得到实际的压缩效果，作者在文章中提出了一种用于实际编码的熵模型，其主要基于 H.264/AVC 中使用的 CABAC 压缩算法^[6]。该熵模型主要编码的信息包含：

• 图像的大小：两个 $16$ 比特的整数（不使用 CABAC 算法）；
• 图像类型：灰度图或 RGB 图，$1$ 比特表示（不使用 CABAC 算法）；
• 模型参数：使用一个 $16$ 比特的整数，来表示分析综合变换的模型参数 $\boldsymbol{\lambda}$ 的索引（一般训练好模型参数后，会发给编解码端）；
• 量化后的图像数据 $\boldsymbol{q}$：将 $\boldsymbol{q}$ 在空间上光栅扫描并在通道上依次迭代处理形成一串数据流，再使用 CABAC 算法进行压缩。

对于量化后的图像数据 $\boldsymbol{q}$ 部分，由于 CABAC 模型是用来处理二进制数据的，因此还需要将 $\boldsymbol{q}$ 转为二进制决策。对每个 $q_i$，二分决策其和众数的关系：

每次往左决策，即 $+0$；往右决策，即 $+1$。最终得到落在区间 $[q_{i, \min}, q_{i, \max}]$ 的所有 $q_i$ 值的二进制编码，再使用 CABAC 进行压缩。对于不在 $[q_{i, \min}, q_{i, \max}]$ 的 $q_i$，值，使用指数哥伦布码进行压缩（也是借荐于 H.264）。作者在 Kodak 数据集上进行了实验，验证了这种基于上下文的自适应压缩算法是要好于传统的确定性压缩算法的。

6. 实验

作者对比了传统的 JPEG 和 JPEG2000 算法，它们的率失真曲线如下：

可以看到，作者提出的端到端模型是优于传统的 JPEG 和 JPEG2000 算法的。更多可视化结果可以到作者这篇文章的主页上查阅：End-to-end Optimized Image Compression 。

7. 总结

首先，作者提到自己提出的方法相比于传统的 JPEG 和 JPEG2000 压缩效果更好，并且这还只是在数据域（data space）上使用 MSE 损失函数优化效果。根据前面作者的分析，如果将 MSE 替换成一个在人眼感知域（perceptual space）上更优的损失函数，比如 NLP^[7]，效果应该会好。

其次，作者强调在本文中使用非线性分析综合变换对 GDN/IGDN^[3:3] 展现出了很好率失真性能。

最后，作者提到，率失真目标函数可以看作是一般的无监督学习/概率密度估计问题的一个实例。而且，由于变换编码涉及到将图像数据映射到一系列的离散值上，相当于分类问题的一种形式。因此作者说可以考虑自己的工作是否对其它有监督学习问题的优化有帮助。作者举了一个例子，即其在本文中引入的可微近似量化，其可以看作是加了一种噪声扰动，可能对增强模型的鲁棒性有所帮助，在一定程度上能避免模型受到对抗攻击而导致灾难性的失效。

关于这篇文章中出现的理论细节，作者在他 2016 年的文章^[2:2]中有很详细的推导，也可以看我对那篇文章的笔记：2016-PCS-End-to-end optimization of nonlinear transform codes for perceptual quality。

附录

---

1. Ballé, J., Laparra, V., & Simoncelli, E. P. (2017). End-to-end optimized image compression. In 5th International Conference on Learning Representations, ICLR 2017. ↩︎

2. Ballé, J., Laparra, V., & Simoncelli, E. P. (2016, December). End-to-end optimization of nonlinear transform codes for perceptual quality. In 2016 Picture Coding Symposium (PCS) (pp. 1-5). IEEE. ↩︎ ↩︎ ↩︎

3. Ballé, J., Laparra, V., & Simoncelli, E. P. (2016, January). Density modeling of images using a generalized normalization transformation. In 4th International Conference on Learning Representations, ICLR 2016. ↩︎ ↩︎ ↩︎ ↩︎

4. Gray, Robert M. and David L. Neuhoff (1998). “Quantization”. In: IEEE Transactions on Information Theory 44.6. DOI: 10.1109/18.720541. ↩︎ ↩︎

5. Kingma, D. P. & Welling, M. (2014). Auto-Encoding Variational Bayes. 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, . ↩︎ ↩︎

6. D. Marpe, H. Schwarz and T. Wiegand, “Context-based adaptive binary arithmetic coding in the H.264/AVC video compression standard,” in IEEE Transactions on Circuits and Systems for Video Technology, vol. 13, no. 7, pp. 620-636, July 2003, doi: 10.1109/TCSVT.2003.815173. ↩︎

7. Laparra, V., Ballé, J., Berardino, A., & Simoncelli, E. P. (2016). Perceptual image quality assessment using a normalized Laplacian pyramid. Electronic Imaging, 2016(16), 1-6. ↩︎