©作者 | 郭宗昱
单位 | 中国科学技术大学
来源 | 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 来进行熵编解码,所以在解码端我们的方法能取得非常快的解码速度,我们统计了实际的编解码时间如下表所示。
从上面两个表格可以看出来,我们提出的 COMBINER 方法编码时间较长,但是解码时间可以非常的短。考虑到 COMBINER 是第一个可以支持隐式表征压缩率失真联合优化的探索性框架,我们希望也相信未来会有一些后续进展可以弥补编码时间较长的问题。

结论

在本文中,我们提出了 COMBINER,一种新型的神经网络压缩方法。它首先将数据编码为变分贝叶斯隐式神经表示(INR),然后通过相对熵编码传递近似后验权重样本。与之前的基于 INR 的神经编解码器不同,COMBINER 支持联合率失真优化,因此可以适应性地激活和剪枝网络参数。

此外,我们还引入了一种迭代算法,用于学习网络权重的先验参数并逐步精细化变分后验。这些方法显著提高了 COMBINER 的率失真性能。特别是,COMBINER 在低分辨率和高分辨率的图像及音频压缩方面都实现了强大的性能,展示了其在不同数据模态和场景下的应用潜力。

然而,正如我们原文中所讨论的,COMBINER 也存在一些局限性。首先,尽管解码过程快速,但编码时间相对较长。优化变分后验分布需要数千次迭代,而逐步微调这些分布也非常耗时。其次,贝叶斯神经网络在本质上对初始化非常敏感。为了实现训练稳定性和优越的率失真性能,确定最佳初始化设置可能需要大量努力。尽管存在这些挑战,我们相信 COMBINER 为 INR 压缩的联合率失真优化铺平了道路。

参考文献
[1] Compression with Implicit Neural Representations. Dupont et al., ICLR 2021 workshop.
[2] Fast Relative Entropy Coding with A* Coding. Flamich et al., ICML 2022.
[3] Joint Autoregressive and Hierarchical Priors for Learned Image Compression. Minnel et al., NeurIPS 2018.
更多阅读
#投 稿 通 道#
 让你的文字被更多人看到 
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:[email protected] 
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
·
继续阅读
阅读原文