【图像压缩】自回归模型《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》

1. 绪论

第一个在PSNR和MS-SSIM都优于BPG的学习模型。引入自回归模型改善熵模型，虽然自回归模型计算很慢，但作者发现在图像压缩领域，自回归模型与多层先验模型互补，能够比之前的模型更能挖掘隐层表示的概率结构。

训练目标如下：

$R+\lambda \cdot D = \underbrace{\mathbb{E}_{x\sim p_x}[-log_2\ p_{\hat{y}}(\left \lfloor f(x) \right \rceil)]}_{rate} +\lambda \cdot \underbrace{\mathbb{E}_{x\sim p_x}[d(x, g(\left \lfloor f(x) \right \rceil))]}_{distortion}$

其中， $\lambda$ 是拉格朗日乘子平衡压缩率和失真， $p_x$ 是自然图像的未知的分布， $\left \lfloor \cdot \right \rceil$ 代表量化操作。 $y=f(x)$ 代表encoder， $\hat{y} = \left \lfloor y \right \rceil$ 将隐层量化， $p_{\hat{y}}$ 是离散熵模型， $\hat{x} = g(\hat{y})$ 是decoder， $\hat{x}$ 代表重构图像。rate项对应隐层边缘分布与先验概率（熵模型）之间的交叉熵，当两分布相同的时候，rate项最小。对于distortion项，当 $d(x, \hat{x})$ 代表MSE的时候，可以解释为VAE模型，当使用其他失真指标如MS-SSIM的时候，可以看做一个能量函数。

本篇工作基于超先验模型，在超先验模型中，关键在于引入超先验作为边信息导入比特流，使得decoder可以共享熵模型。通过这种方式，熵模型是图像相关和空间自适应的，使得模型更准确。这说明深度学习的优化方法足以在辅助信息的大小和更精确的熵模型之间学习一个有效的平衡。

本篇工作，将均值0只有标准差参数的高斯模型拓展为高斯混合模型，同时添加了自回归分支。专注于不同先验形式的潜力，本文章针对计算复杂度优化。

2. 模型细节

模型结构如下图所示：

将熵模型建模成为以超先验和超解码器、上下文模型参数、熵模型参数（ $\theta_{hd}, \theta_{cm}, \theta_{ep}$ ）为条件的高斯分布：

$p_{\hat{y}}(\hat{y}|\hat{z},\theta_{hd}, \theta_{cm}, \theta_{ep}) = \prod _i (\boldsymbol{N}(\mu_i, \sigma_i^2)*\boldsymbol{U}(-\frac{1}{2}, \frac{1}{2}))(\hat{y}_i) \\ with\ \mu_i, \sigma_i = g_{ep}(\psi, \phi_i; \theta_{ep}), \psi = g_h(\hat{z};\theta_{hd}), and\ \phi_i = g_{cm}(\hat{y}<i;\theta_{cm})$

超先验模型依旧使用非参数、全分解概率模型Entropy Bottleneck，虽然改善超先验的熵模型比如累计几层本文的上下文模型，这可以提高超先验的压缩率，但是 $\hat{z}$ 占总比特流的很小一部分，考虑到网络性能的原因没有讨论。将优化目标拓展到包含超先验：

$R+\lambda \cdot D = \underbrace{ \mathbb{E}_{x\sim p_x}[-log_2\ p_{\hat{y}}(\hat{y})]}_{rate(latents)} + \underbrace{\mathbb{E}_{x\sim p_x}[-log_2\ p_{\hat{z}}(\hat{z})]}_{rate(hyper-latents)} +\lambda \cdot \underbrace{\mathbb{E}_{x\sim p_x}[d(x, g(\left \lfloor f(x) \right \rceil))]}_{distortion}$

本质上自回归模型的预测 $\phi_i = g_{cm}(\hat{y}<i;\theta_{cm})$ 应该以所有的潜层表示为条件，但实现的时候使用的是有限上下文（5x5卷积），实现方式与PixelCNN中相似。PixelCNN中自回归模型被用来生成图像，但本质上其作用是拟合数据分布，这里用在图像压缩中拟合隐层分布。实现细节：使用一层掩码卷积，padding=2，掩码作用是屏蔽没有预测的信号，如下图所示。训练的时候，执行正常的卷积操作，最终得到的有效参数为绿色覆盖的部分；推理（解压缩）的时候，先将预测结果 $p_{ctx}$ 初始化为0，然后从左到右、从上到下依次对隐层每个pixel进行预测，每次预测后将得到的结果写回 $p_{ctx}$ 。

之后将 $p_{ctx}$ 和 $h_s(\hat{z})$ cat起来，送入熵模型参数网络得到高斯分布的参数。

此时熵模型参数网络同样应受到限制，因为它无法访问来自当前潜层元素之外的上下文模型的预测，简单起见使用1x1卷积。

2. 网络细节与限制

编码器最后一层对应自编码器中的瓶颈，通过率失真函数来平衡，我们的模型通过生成相同的隐层值并将其概率设置为1来忽略某些特定的通道，这种模型的灵活性使得能够将瓶颈层设置得比需要的更大，让模型确定产生最佳的通道数。同时发现，当比特率变高的时候如果瓶颈处通道数太少的话会阻碍率失真表现，太多的话则不会损失压缩性能。

限制解码器最后一层通道数为3，熵模型输出通道数为192*2，因为要预测均值和标准差。上下文模型和超解码器模型不受限制，在实验中将他们的通道数设置为瓶颈层的两倍。

3. 实验结果

对于上下文模型，图4比较了单个压缩点下文件大小相对增长。baseline使用5x5掩码卷积、高斯分布熵模型。绿色条显示把高斯模型替换为逻辑回归模型几乎没没有影响，替换为拉普拉斯分布则降低很多性能。蓝条比较不同上下文配置，无论是使用3x3 7x7还是多加几层上下文性能都会差于baseline。紫色条显示使用严格限制后的上下文模型的效果，只考虑前一行的1个或者3个邻居，虽然性能有所下降，但可以提高并行性。

5 讨论

关于超先验模型与自回归模型互补：首先，从超先验的角度出发，我们看到对于相同的超网络架构，熵模型的改进需要更多的边信息，但边信息增加了总压缩文件的大小，两者形成trade-off。相比之下，将自回归引入先验不会导致编码率降低，因为预测仅基于因果上下文，即已经解码的隐层元素。同样，从自回归模型的角度来看，我们希望消除一些不确定性，但不能仅仅依靠上下文完成消除，超先验可以“展望未来”，因为它是压缩比特流的一部分，并且解码器可以共享。因此，超先验可以学习存储减少自回归模型中的不确定性所需的信息，同时避免可以从上下文中准确预测的信息。

图6显示了三种变体的内部机制，第一列为提取隐层熵值最高的通道，第二列为预测的均值（基于VAE的模型比基于自回归的模型更容易产生模糊的图像），第三列为预测误差（隐层减去均值），第四列为预测尺度，最后两列为熵的分布。可以看出本文模型预测误差和预测尺度更小。

自回归模型因为串行的缘故，不能并行优化，耗时。没有为本文实现完整的解码器，而是只比较香农熵。

未来研究的另一个方向可能是通过在严格的多层先验中引入更多的复杂性来完全避免因果关系问题。