[论文解读] Compression with Bayesian Implicit Neural Representations

Zongyu Guo, Gergely Flamich+3 more|arXiv (Cornell University)|May 30, 2023

Advanced Data Compression Techniques被引用 6

一句话总结

本文提出 COMBINER，一种与模态无关的神经压缩方法，利用变分贝叶斯隐式神经表示（INRs）实现联合率失真优化。通过从变分后验中采样权重并使用相对熵编码压缩，COMBINER 直接优化 $β$-ELBO，无需昂贵的元学习循环，在图像和音频任务上均表现出色。

ABSTRACT

Many common types of data can be represented as functions that map coordinates to signal values, such as pixel locations to RGB values in the case of an image. Based on this view, data can be compressed by overfitting a compact neural network to its functional representation and then encoding the network weights. However, most current solutions for this are inefficient, as quantization to low-bit precision substantially degrades the reconstruction quality. To address this issue, we propose overfitting variational Bayesian neural networks to the data and compressing an approximate posterior weight sample using relative entropy coding instead of quantizing and entropy coding it. This strategy enables direct optimization of the rate-distortion performance by minimizing the $β$-ELBO, and target different rate-distortion trade-offs for a given network architecture by adjusting $β$. Moreover, we introduce an iterative algorithm for learning prior weight distributions and employ a progressive refinement process for the variational posterior that significantly enhances performance. Experiments show that our method achieves strong performance on image and audio compression while retaining simplicity.

研究动机与目标

• 解决现有基于 INR 的压缩方法存在的性能差距，这些方法依赖固定精度量化且缺乏端到端的码率-失真优化。
• 通过利用变分贝叶斯推断，实现隐式神经表示中码率与失真的联合优化。
• 开发一种支持自适应参数剪枝与激活的方法，提升压缩效率而不损失重建质量。
• 通过后验样本的相对熵编码，克服 INR 压缩中固定精度量化带来的局限性。
• 设计一种简单而有效的框架，无需针对模态调整网络结构，即可在图像和音频数据上实现泛化。

提出的方法

• 对 INR 权重上的变分后验分布 $q_{\mathbf{w}}$ 进行数据拟合，以概率表示替代点估计。
• 使用相对熵编码（REC）压缩采样得到的权重向量 $\mathbf{w}^* \sim q_{\mathbf{w}}$，其期望代价近似于 KL 散度 $D_{\mathrm{KL}}[q_{\mathbf{w}}\|p_{\mathbf{w}}]$。
• 通过调节 $\beta$ 最小化 $\beta$-ELBO 目标，直接优化码率-失真权衡。
• 提出一种迭代算法，用于学习网络权重上的最优先验分布 $p_{\mathbf{w}}$，以提升后验质量。
• 实施一种渐进式精炼策略：将权重划分为块，按顺序压缩并微调，以提升后验精度。
• 在每个块编码后，使用梯度下降优化剩余未压缩权重的后验，从而提升整体重建保真度。

实验结果

研究问题

• RQ1 与固定精度量化相比，变分贝叶斯推断是否能提升基于 INR 的压缩中的码率-失真性能？
• RQ2 学习 INR 权重上的先验分布如何影响压缩性能以及在不同数据模态间的泛化能力？
• RQ3 对变分后验的渐进式精炼在多大程度上提升了重建质量与码率-失真权衡？
• RQ4 一个简单、端到端可训练的基于 INR 的压缩框架，是否能在无需模态特定设计的情况下超越复杂的元学习基线？
• RQ5 迭代后验精炼与先验学习对图像与音频压缩中 PSNR 和比特率效率的影响如何？

主要发现

• COMBINER 在 CIFAR-10、Kodak 和 LibriSpeech 数据集上实现了出色的压缩性能，尽管结构简单，仍优于现有基于 INR 的方法。
• 渐进式精炼与迭代先验学习技术使低分辨率图像压缩任务的 PSNR 提升超过 4 dB。
• 在单张 A100 GPU 上编码 500 张 CIFAR-10 图像耗时 13–34 分钟，具体时间取决于码率，其中大部分时间用于后验推断与微调。
• 解码效率极高：单张 CIFAR-10 图像重建仅需 2–4 ms（GPU），Kodak 图像在 CPU 上解码时间低于 1 秒。
• 微调超过 3,000 步后收益递减，30k 步仅比 3k 步提升 0.3 dB PSNR，但耗时增加 90%。
• COMBINER 消除了最先进 INR 编解码器中使用的元学习循环，实现了无需数据分块的完整图像与完整音频 INR 训练。