NeurIPS 2023 Spotlight | 通过贝叶斯隐式表征实现数据压缩

国际科技财经移民娱乐民生时事体育

Bendi新闻

11月前

©作者 | 郭宗昱

单位 | 中国科学技术大学

来源 | IMCL研习社

论文标题：

Compression with Bayesian Implicit Neural Representations

论文链接：

https://arxiv.org/pdf/2305.19185.pdf

在数字世界里，很多常见的数据类型都可以看作是一种特殊的连续函数形式，表征为隐式表征。例如，图像可以表征为将每个像素位置映射到对应 RGB 色值的连续函数。从这个角度出发，我们可以通过让一个紧凑的神经网络拟合这种函数表示形式，接着对网络权重进行编码，来实现数据压缩。

然而，目前大多数的方法在这一过程中效率并不高。原因在于，将网络权重量化到低比特精度会严重影响重构质量。

为了解决这个问题，我们提出了一个新颖的思路：利用变分贝叶斯神经网络对数据进行连续函数拟合，然后通过相对熵编码对近似后验权重样本进行压缩，而不是采用传统的量化和熵编码方法。这种策略使我们能够通过最小化-证据下界（-ELBO）直接优化压缩的率失真性能，并且通过调整值，针对给定的网络架构实现不同的率失真权衡。

此外，我们还引入了一种迭代算法来学习先验权重分布，并采用了一种逐步精细化的变分后验处理过程，显著提升了性能。实验表明，我们的方法在图像和音频压缩方面都取得了出色的性能，同时保持了方法的简洁性。

论文简介

近年来，深度学习的发展带来了数据压缩领域的显著发展，特别是在有损图像压缩领域。受到深度生成模型，尤其是变分自编码器（VAE）的启发，深度图像压缩模型在客观指标（如 PSNR 和 MS-SSIM）和感知质量上均超越了最佳手工设计的图像压缩方案。然而，这些方法的成功很大程度上归功于它们为特定数据模态设计的复杂架构。

然而，这使得将它们转移来压缩不同数据模态变得具有挑战性。为了解决这一问题，近期的研究 [1] 提出了将单个数据视为映射坐标到值的连续信号，并通过过拟合小型神经网络，即隐式神经表示（Implicit Neural Representations, INRs），来压缩它们的权重。尽管 INR 在灵活性上有优势，但现有的基于 INR 的压缩方法与针对特定模态的神经压缩模型相比仍存在显著的性能差距。

在这篇论文中，我们提出了一种简单却通用的方法，通过将 INR 扩展到变分贝叶斯形式 (variational Bayesian INRs) 来解决这一问题。我们将数据比如单张图像拟合成一个变分后验分布 (variational posterior) 的贝叶斯网络 (Bayesian Neural Network)，而不是一个点估计。

然后，我们使用相对熵编码（Relative Entropy Coding, REC）算法 [2] 来压缩 INR 的分布的一个采样。这样，我们可以直接优化 INR 的 rate-distortion tradeoff 这一有损压缩中最重要的目的。我们将这套新的框架称为 COMBINER (Compression with Bayesian Implicit Neural Representations)，如下图所示。

基于这套我们新提出的框架，我们还引入了两个重要的技术改进。首先，我们提出了一种学习网络权重先验分布 (model prior) 的迭代算法，给出了通过一组后验分布来估计先验分布的解析解。其次，我们还提出了一种逐步细化后验分布的渐进策略，可以不断修正相对熵编码所带来的损失。具体方法如下一节所述。

方法介绍

(1) A* Coding / Relative Entropy Coding

为了实现码率失真联合优化，我们将数据的信息编码 (encode) 为贝叶斯网络的高斯参数分布，但是我们知道分布的无损传输是非常困难的，比如对于高斯分布我们需要分别用足够的精度传输分布参数 mean 和 scale。

事实上，我们可以通过相对熵编码（Relative Entropy Coding, REC），在这里也是 A* Coding [2]，来实现分布的一个近似采样 (a sample from the posterior distribution) 的传输。A* Coding 的具体算法流程如下：

可以看到 A* Coding 需要在编解码端维护一个相同的随机数发生器，以及共享一个对于不同压缩数据都一样的 model prior。这充分说明了 model prior 对于网络的重要性。需要指出的是，另一方面 model prior 还直接影响了 rate-distortion tradeoff 中的 rate 值。

(2) Model Prior

为了得到尽可能好的 model prior，我们提出了一种迭代式的策略配合解析解来高效的学习。给定一组图像，我们可以先随机初始化一个 model prior，然后“稍微”用这个 model prior 来优化一下这组图像中每一张图像对应的 model posterior，然后可以用这里算法中提到的 Equation (5) 来用解析解更新一下 model prior，然后再迭代式的优化 posterior。

如上图所示，我们在这里也给出了如何从一组 posterior distribution 中来更新 prior distribution。

整个学习 model prior 的算法如下所示，我们通常在实验中会进行 64 到 256 轮迭代，每轮迭代会梯度下降优化约 30-100 次 model posterior distribution，然后会利用上 (5) 式来更新一次 model prior。

(3) Posterior Refinement

类似于基于 VAE 编码的自回归上下文模型 [3]，我们提出了在编码时候进行自回归微调，我们将 INR 网络参数按照固定的编码 budget 分为了数组，然后按照顺序一个一个的进行组的 A* Coding。

在编码完一个组之后，我们对剩下的所有未编码的参数按照率失真损失函数进行微调，这样可以一方面弥补前面实际编码参数组偶尔出现的坏的样本，另一方面可以通过微调后面的组来更好表征原始待压缩信号。通过我们提出的 Posterior Refinement 策略，我们可以在小数据集上大大提升压缩的率失真性能，但是会使得编码时间大大增长。

实验结果

我们首先在图像压缩任务上在小分辨率的 CIFAR 数据集上（如下图左）和大分辨率的 Kodak 上（如下图右）分别做了实验。实验结果与之前方法的比较如下图所示。

在上图中，我们比较了三类方法，分别有，传统编码方法包括 BPG，JPEG2000，在图中用虚线来表示。基于 VAE 的编码方法包括 ICLR2018 ，CVPR2020，在图中用点线来表示。基于隐式表征的编码方法，包括 COIN [1]，COIN++，MSCN 和我们提出的 COMBINER，在图中用实线来表示。

可以看到我们提出的 COMBINER 与之前的隐式表征压缩方法相比取得了最优的性能，但是仍然与现在流行的复杂的 VAE 框架下的图像编码方法的率失真性能有一些差距，特别是在高码率端。另一方面来说，基于隐式表征的图像编码方法在低码率端展现出超过 VAE 编码方法的潜能。

由于我们的方法在实际编解码图片时，编码的过程分为两步，分别是学习给定信号的对应隐式表征网络的后验分布，然后是渐进式的微调来提高编码的性能，所以可以想象编码的时间复杂度相较于之前的方法有不小的增加。另一方面来说，由于我们采用相对熵编解码 REC 来进行熵编解码，所以在解码端我们的方法能取得非常快的解码速度，我们统计了实际的编解码时间如下表所示。