【图像压缩】多层超先验模型《Coarse-to-Fine Hyper-Prior Modeling for Learned Image Compression》

第一篇中在建模部分表述不是很清晰，建模部分用的第二篇。

绪论

基于超先验和空间上下文的模型在建模长期依赖方面有局限性，不能很好地榨干图像中的空间冗余度。本文提出一个多层超先验模型，并使用由粗到细的框架，以进一步减少空间冗余。Signal Preserving Hyper Transforms是对隐层表示的更深入的分析变换，Information Aggregation Reconstruction 最大化利用多层边信息进行重建。

基于超先验和上下文模型存在的问题：当前待编码元素的概率值取决于少部分已编码元素，限制了长期条件概率估计的准确性；由于自回归模型需要顺序计算，难以并行加速；由于浅层网络限制，分析、合成变换无法在保证重建质量的同时最大程度减少空间冗余；超先验传输的信息没有被用于重建图像；超先验和上下文模型的交互和合作方式未得到很好的研究。

通过分析，上下文依赖可以建模为更高层的隐层表示。

动机

尽管使用GDN进行分析变换可以很大程度上减少像素间的空间冗余，从而将隐层表示建模为分解概率模型，但由于其感知上下文能力有限，仍然有可能忽略一些相关性。后来的工作使用上下文自适应的方法估计隐层表示 $\boldsymbol{X} = \{X_i\}$ 的联合分布 $P_X(X)$ 以建模上下文依赖， $Q_X(X;\theta)$ 表示熵模型是 $P_X(X)$ 的估计：

$Q_{X|Y} (X|Y)= \prod_i Q_i(X_i|X_{i-1}, X_{i-2}, \cdots, X_{i-m}, Y)$

其中，当计算当前 $X_i$ 的概率时， $X_{i-1}, X_{i-2}, \cdots, X_{i-m}$ 为之前计算过元素。存在问题：

1. 超先验形式的模型需保证 $Y$ 的维度与比特率低，否则 $Y$ 本身可能会包含冗余，致使压缩性能受损，在这种情况下，超先验可能不能提供足够的信息来建模条件概率，特别是对高比特率、高分辨率而言；

2. 引入上下文模型改善了这点，但上下文模型只考虑了前m个元素，因此不能建模长期依赖关系；实际应用中，m增大虽加强了建模能力但带来了高计算复杂度；只能逐元素计算，不能并行计算。因此不能很好地处理多尺度图像。

由粗到细超先验建模

在之前的超先验模型中，假设隐层表示的每个元素都是独立分布的，为了使假设性成立，需要设计更复杂的分析/合成变换来去除/重建冗余，这种分析合成变换包含太多参数、时间效率低。

未解决上述问题，引入多层条件框架：

$Q_X(X) = Q_{X,Y}(X,Y) = Q_Y(Y)Q_{X|Y}(X|Y)$

$Q_Y(Y) = Q_{Y,Z}(Y,Z) = Q_Z(Z)Q_{Y|Z}(Y|Z)$

实际上构成了一个由粗到细的框架，Z经过两层超分析变换，其维度进一步降低，更多冗余被去除，因此可以近似为分解概率：

$Q_Z(Z) = Q_Z(Z_1, Z_2, \cdots, Z_n)\approx \prod_i Q_{Z_i}(Z_i)$

由上面的式子， $Y$ 和 $X$ 的概率的分布可以被条件建模，之前的工作发现神经网络可以建模条件概率分布，通过训练， $Y$ 中可以嵌入被压缩图像的主要部分，因此联合分布可以近似分解如下：

$Q(X|Y) = Q(X_1, \cdots, X_n) \approx \prod_i Q_{X_i|Y}(X_i|Y)$

$Q(Y|Z) = Q(Y_1, \cdots, Y_n|Z) \approx \prod_i Q_{Y_i|Z}(Y_i|Z)$

之前层的表示被用作估计更高层分布的条件。虽然没有执行上下文条件，上下文条件被隐式地包含在从 $X$ 流向 $Y$ 的信息中，并在从 $Y$ 预测 $X$ 的过程中使用。不同于上下文模型以block为条件，由于由粗到细的结构设计，每个元素的概率估计使用到了更大区域的信息。

多层网络结构如上图所示。输入图像经过分析变换到隐层表示 $X$ ，旨在尽可能榨干冗余度，之后进行四舍五入量化。使用多层超分析与超合成变换构建一个由粗到细的框架。概率预测网络通过 $Y$ 计算 $Q_{X_i}(X_i = x_i)$ 来估计似然概率 $P_{X_i}(X_i = x_i)$ ，与超先验模型类似，将 $X$ 中每个元素的条件分布假设满足高斯分布，需要预测高斯分布的 $\mu$ 和 $\sigma$ 。因为隐层表示已经四舍五入到离散形式，所以隐层表示的似然可以计算如下：

$Q_{X_i|Y}(X_i=x_i|Y) = \phi(\frac{x_i + \frac{1}{2}-\mu_x}{\sigma_{x_i}}) - \phi(\frac{x_i - \frac{1}{2}-\mu_x}{\sigma_{x_i}})$

其中 $\phi$ 代表标准正太分布的累积分布函数， $\mu_{x_i}$ 和 $\sigma_{x_i}$ 从 $Y$ 中预测得到。从 $Z$ 预测 $Y$ 的分布的过程与之相同。将 $Z$ 假设为满足0均值的高斯分布， $Z$ 中每个元素的似然计算如下：

$Q_{Z_i}(Z_i = z_i) = \phi(\frac{z_i + \frac{1}{2}}{\sigma_{z_i}}) - \phi(\frac{z_i - \frac{1}{2}}{\sigma_{z_i}})$

注意 $\sigma_{z_i}$ 是网络中的一个可训练参数（？经过训练网络可以得到的参数），一个通道的所有隐层表示共享一个相同的 $\sigma$ ，不同的通道间 $\sigma$ 相互独立。

由信息论，编码X需要的最小比特率为实际分布 $P(X|Y)$ 与估计分布 $Q_{X|Y}\sim N(\mu_x, \sigma_x)$ 的交叉熵：

$R = H(Q) + D_{KL}(P||Q) = E_{P_{X|Y}}[-log(Q)]$

优化目标函数为 $L_{RD} = R + \lambda D$ 。为了加速多层网络收敛，引入信息保真损失。该损失鼓励高层表示 $Y$ 保留 $X$ 中的关键信息，列式如下：

$\mathop{min}\limits_{Y,\theta} L_{if} = \left \| F(Y, \theta) - X \right \|_2$

其中 $F$ 是没有激活层的一层卷积层，信息保真损失使得 $\mu$ 和 $\sigma$ 的估计更准确。

可以保留信号的超变换（Signal Preserving Hyper Transform）

为了对图像由粗到细建模，关键是在执行更复杂的合成分析变换的同时保留信息。一方面，与自然图像相比，隐层表示的相关性已经大大降低，而使用大卷积核的卷积依赖于局部相关性建模；另一方面，之前的超变换使用跨步卷积和Relu，跨步卷积进行下采样，Relu禁用了某些为负的神经元。由于需要限制这些卷积层的维度以确保隐层表示的逐渐分解，原来的超变换丢失了很多信息。

原来的超变换存在两个问题：1）原始分析变换保持维度不变，下采样特征图，降低了隐层表示的尺寸 2）在分析变换的开始和合成变换的结尾使用大卷积+Relu导致丢失尚未变换的信息，限制了模型能力。

可以保留信号的超变换网络结构如表1，在第一层使用相对较小的卷积，使用的是线性激活函数，之后全部使用 $1\times1$ 卷积。第一层增加了原始表示的通道数，结合后面的非线性层，维度的拓展在提供非线性建模的同时保留了信息。使用space-to-depth操作reshape张量，使得空间上相邻的元素分布在不同通道的相同位置，通过这种方法，后面的 $1\times1$ 卷积可以通过非线性变换减少空间冗余。在网络最后一层，减少通道数使得表示更紧凑。对称设计超合成变换。

space-to-depth与depth-to-space操作实现代码如下建模，图示：

# Space-to-depth & depth-to-space module
# same to TensorFlow implementations
class Space2Depth(nn.Module):def __init__(self, r):super(Space2Depth, self).__init__()self.r = rdef forward(self, x):r = self.rb, c, h, w = x.size()out_c = c * (r**2)out_h = h//2out_w = w//2x_view = x.view(b, c, out_h, r, out_w, r)x_prime = x_view.permute(0, 3, 5, 1, 2, 4).contiguous().view(b, out_c, out_h, out_w)return x_primeclass Depth2Space(nn.Module):def __init__(self, r):super(Depth2Space, self).__init__()self.r = rdef forward(self, x):r = self.rb, c, h, w = x.size()out_c = c // (r**2)out_h = h * 2out_w = w * 2x_view = x.view(b, r, r, out_c, h, w)x_prime = x_view.permute(0, 3, 4, 1, 5, 2).contiguous().view(b, out_c, out_h, out_w)return x_prime

聚合信息重建图像

不同于超先验模型，为了最大化利用比特流中数据，不仅使用边信息估计分布，而且用它来辅助重构图像。信息聚合解码网络如图2：

融合不同尺度的表示有助于图像重构，高阶表示提供全局信息，其他的提供图像细节信息。融合过程在小空间分辨率上执行以避免高计算复杂度。

实现细节

概率估计子网络

如图1所示，为了进行算术编码需要使用概率估计子网络对量化后的隐层表示进行概率估计，对于尺度最小的最内层超先验假设分布满足零均值高斯分布，第一层超先验和隐层表示都采用有均值和标准差的高斯分布。

为了充分利用超先验，设计了一个有更合适采样区域的预测子网络。图3所示子网络中，以被预测位置为中心采样一个 $5\times5$ 的区域，每一个采样区域经过多层卷积估计预测位置的概率值。最后一层使用全连接网络，输出预测的均值和标准差。局部卷积在所有空间位置上共享卷积，超隐层表示被密集采样，在卷积操作过程中没有信息丢失。

训练建模

网络分为三个阶段训练。第一阶段，添加一个上下文模型，训练主合成、分析变换以获得良好的图像重建。第二阶段，去掉上下文模型，添加超变换、合成网络，端到端训练。最后，将原始的合成变换最后一层换成信息聚合网络，对整个网络训练。

前向传播的量化使用四舍五入，反向传播的时候看做 $y=x$ 导数为1，训练估计表示的熵的时候，向原来的隐层表示添加一个噪声 $U(-\frac{1}{2}, \frac{1}{2})$ 提供更好的拟合效果（超先验模型中的方法）。

实验结果

在实验过程中，首先发现如果采用原来的超分析、合成则很难构建一个多层超先验模型，添加的超先验层几乎捕捉不到用来建模上层概率分布的边信息，这可能是因为前向过程中由于大卷积和ReLU层丢失了信息。通过引入可以保留信息的超变换构建的多层超先验模型比超先验模型表现要好，添加信息聚合模块之后性能进一步提升。

该模型在高分辨率数据集Tecnick数据集上表现的更出众，这是因为这些图像包含更多的长期依赖，由粗到细模型可以很好的解决。

下图计算加速比证明了该模型在并行能力上的优越性：