COIN: COmpression with Implicit Neural representations

论文地址:https://arxiv.org/abs/2103.03123
论文代码:https://github.com/EmilienDupont/coin

Abstract

我们提出了一种新的简单的图像压缩方法：不是存储图像每个像素的RGB值，而是存储过度拟合到图像的神经网络的权重。具体来说，为了编码图像，我们使用MLP将像素位置映射到RGB值。然后，我们量化并存储该MLP的权重作为图像的代码。为了解码图像，我们只需在每个像素位置评估MLP。我们发现，即使没有熵编码或学习权重分布，这种简单的方法在低比特率下也优于JPEG。虽然我们的框架还不能与最先进的压缩方法竞争，但我们表明，它具有各种吸引人的特性，这可能使其成为其他神经数据压缩方法的可行替代方法。

1 介绍

基于学习的图像压缩方法通常在自动编码器中运行[2、24、20]。发送端将数据输入编码器映射到离散的潜在码，再根据学习到的潜在分布熵编码为比特流。比特流被传输到接收器，被解码为潜在码，最终通过解码器重建图像。

在COIN中，用一个小的 MLP 网络映射像素位置到 RGB 值从而对图像进行拟合，然后将该 MLP 网络的权重 $\theta$ 作为编码进行图像的传输（图1）。由于自然图像中包含的高频信息，很难过拟合此类MLP（称为隐式神经表示）[3,34]，但最近的研究表明，可以通过使用正弦编码和激活函数来缓解这一问题[23,34,30]，其可以把低维坐标映射到高维。在这项工作中，我们把使用正弦激活 Sine 的 MLP，通常被称为 SIREN [30]，我们可以用很小的网络（8k）来拟合大图像（393k）。

我们在标准图像压缩任务上评估了我们的方法，并表明我们在低比特率下优于JPEG，即使没有熵编码或学习权重分布。由于隐式表示已成功应用于生成建模[10]，因此将我们的方法与学习的权重分布相结合可能会产生有前途的神经数据压缩新方法。此外，通过将我们的图像视为从像素位置到RGB值的函数，我们可以简单地通过以逐渐更高的分辨率评估我们的函数来执行渐进解码，这对于资源受限的接收设备特别有吸引力。

图1 基于隐式神经表示的压缩。我们使用神经网络将像素位置（x，y）映射到RGB值（通常称为隐式神经表示）来过拟合图像。然后，我们将该神经网络的权重θ量化到较低的比特宽度，并传输它们。

2 方法

在本节中，我们描述了 COmpressed Implicit Neural representations（COIN），这是我们提出的图像压缩方法。编码步骤包括将MLP拟合到图像，量化其权重并传输这些权重。在解码时，传输的MLP在所有像素位置进行评估以重建图像。

2.1 编码

用 $I$ 表示我们想要编码的图像，这样 $I[x,y]$ 返回像素位置 $(x,y)$ 处的 RGB 值。我们定义了一个函数 $f_{\theta }:{\mathbb{R}}^2\to {\mathbb{R}}^3$，其中参数 $\theta$ 将像素位置映射到图像中的 RGB 值，即 $f_{\theta }(x,y)=(r,g,b)$。然后，我们可以通过在某种失真度量下将 $f_{\theta }$ 过拟合到图像来对图像进行编码。在本文中，我们使用均方误差，得出以下优化问题
$$\underset{\theta }{\min }\sum _{x,y}\Vert f_{\theta }(x,y)-I[x,y]{\Vert }_2^2\text{\hspace{1em}(1)}$$
其中 $\Sigma$ 包含了所有像素位置。

选择 $f_{\theta }$ 的参数至关重要。事实上，即使使用大量参数，通过具有标准激活函数的MLP[34，30]来参数化 $f_{\theta }$ 也会导致欠拟合。这个问题可以通过多种方式克服，例如，通过傅立叶特征编码像素坐标[34]或使用正弦激活函数[30]。从经验上看，我们发现，对于给定的参数，后一种选择产生了更好的结果。

给定足够大的 MLP，最小化方程（1）很简单。然而，我们将 MLP 的参数 $\theta$ 存储为图像的压缩描述，所以限制权重的数量将提高压缩率。因此，目标是使用尽可能少的参数（即最大化速率）将 $f_{\theta }$ 拟合到 $I$（即，最小化失真）。然后，我们的方法有效地将数据压缩问题转化为模型压缩问=。

为了减小模型尺寸，我们考虑了两种方法：架构搜索和权重量化。更具体地说，我们在MLP的宽度和层数上执行超参数扫描，并将权重从32位精度量化为16位精度，这足以在低比特率下优于JPEG标准。然而，我们认为，更复杂的架构搜索方法[11]，尤其是模型压缩[36、13、37]将进一步改善结果。

2.2 解码

给定存储的量化权重 $\theta$，解码可以简单的由在每个像素位置进行评估函数 $f_{\theta }$ 以组成，以重建图像。这种解码方法为我们提供了额外的灵活性：我们可以逐步解码图像，例如，首先对部分图像或低分辨率图像进行解码，只需评估不同像素位置的函数。使用基于自动编码器的方法很难以这种方式对图像进行部分解码，这显示了COIN方法的进一步优势。

3 相关工作

隐式神经表征 ：用神经网络表示数据最初是由[32]提出的，但最近3D视觉界的兴趣激增[27、26、6]。受深度体素表示的内存需求呈立方体增长[25，31，9]的推动，提出了隐式表示来紧凑地编码高分辨率信号[23，34，30]。虽然用于表示图像的 MLP 通常具有相对较少的参数[10]，但我们更进一步，并表明通过仔细选择 MLP 的架构并量化权重，我们可以使用比存储 RGB 值占用更少空间的 MLP 来拟合图像。

神经数据压缩 ：学习的图像压缩方法通常基于分层变分自动编码器[2，24，20]，为了熵编码的目的，将学习的先验和潜在变量离散化。与潜变量模型文献[14、19、16、22]中的工作类似，一些工作[5、12、39]试图通过在使用分期推理网络的基础上执行基于梯度的迭代优化步骤来缩小分期差距[8]。此外，[39]还通过每个实例的推理时间优化，识别并尝试缩小潜在变量量化产生的离散化差距。[38]进一步考虑模型的持续优化思想：它们执行解码器的每实例微调，并随潜在代码传输量化的解码器参数更新，从而提高率失真性能。在本文中，我们从每个实例优化的角度采取了一种不同的、甚至更极端的方法：我们优化 MLP 以过拟合单个图像，并传输其权重作为图像的压缩描述。

模型压缩 ：众所周知，数据和模型压缩问题密切相关，但COIN明确地将数据压缩问题转化为模型压缩问题。关于模型压缩，有大量文献[36、21、13、37、18、15]，这些文献可能用于改善COIN的性能。

4 EXPERIMENTS

我们在柯达图像数据集[17]上进行了实验，该数据集由 24 幅尺寸为 768×512 的图像组成。我们将我们的模型与三个基于自动编码器的神经压缩基线（我们称之为 BMS [2]、MBT [24] 和 CST [7]）进行比较。我们还与 JPEG、JPEG2000、BPG 和 VTM 图像编解码器进行了比较。为了对我们的模型进行基准测试，我们使用了 CompressAI 库 [4] 和其中提供的预训练模型。我们在 PyTorch [28] 中实现了我们的模型，并在单个 RTX2080Ti GPU 上进行了所有实验。

率-失真图。为了确定给定参数预算（以每像素比特或 bpp 为单位）下的最佳模型架构，我们首先要为表示图像的 MLP 找到有效的深度和宽度组合。例如，对于使用 16 位权重的 0.3bpp，有效网络包括宽度为 28 的 10 层 MLP 和宽度为 34 的 7 层 MLP，以此类推。然后，我们使用贝叶斯优化法对单张图像上的学习率和有效架构进行超参数搜索，从而选出最佳架构（我们发现架构搜索的结果可以很好地应用于其他图像）。由此产生的模型以 32 位精度在数据集中的每幅图像上进行训练，并在训练后转换为 16 位精度。我们注意到，在训练后将权重精度从 32 位降低到 16 位几乎不会导致失真增加，但将权重精度进一步降低到 8 位会导致大量失真，超过了将 bpp 减半所带来的好处。

图2 显示了这一过程在不同比特率水平下的结果。可以看出，在低比特率下，即使不使用熵编码，我们的模型也比 JPEG 有所改进。虽然我们的方法与目前最先进的压缩方法还相差甚远，但我们相信，这种简单方法的性能对于未来在这一方向的工作大有可为。

[图 2：柯达数据集的比率失真图。] [图 3：0.3bpp 时的模型尺寸。]

[图 4：对柯达数据集中的图像 15 进行模型训练。最大 COIN 表示在训练过程中任何时候达到的最大 PSNR。] [图 5：采用不同架构的相同大小（0.3bpp）网络的最大 PSNR 图。]

模型大小。与其他大多数神经数据压缩算法不同，我们的方法在测试时不需要解码器。事实上，虽然这类方法中代表压缩图像的潜码很小，但解码器模型却很大（通常比未压缩图像大得多）。因此，解码设备所需的内存也很大。在我们的例子中，我们只需要解码器一侧的一个（非常小的）MLP 的权重，因此内存需求要小很多个数量级。如 图3 所示，在 0.3bpp 的情况下，我们的方法需要 14kB 的内存，而其他基线方法需要 10MB 到 40MB 的内存。

动态编码优化。我们在 图4 中展示了过拟合过程的一个示例。可以看出，COIN 在 15k 次迭代后的表现优于 JPEG，并在 15k 次迭代后继续提高。

图4 ：柯达数据集中图像15的模型训练。Max COIN表示训练过程中任何一点达到的最大峰值信噪比。

架构选择。图5 显示了大小为 0.3bpp 的各种有效架构的性能。可以看出，压缩质量取决于架构的选择，不同的 bpp 值有不同的最佳架构，详见附录 A。

5 SCOPE, LIMITATIONS AND FUTURE WORK

局限性。我们的方法的主要局限是编码速度较慢，因为我们必须解决每个编码图像的优化问题。然而，在一对多媒体分发的环境中，例如在 Netflix [1]，为向许多接收者分发压缩内容而预先支付大量计算成本是一种标准做法。不过，元学习（meta-learning）[29, 33] 或摊销推理（amortized inference）[35, 40]方法可以避免这一限制。此外，在解码时，我们需要在每个像素位置对网络进行评估，以解码整个图像。然而，这种计算的并行化程度很低，只需对所有像素进行一次前向传递即可。最后，我们的方法比最先进的压缩方法性能更差。不过，我们相信有几个方向有望缩小这一差距。

未来工作。最近在隐式表征生成建模方面的工作[10]表明，学习函数权重的分布可以为我们的方法带来显著的压缩收益。此外，探索元学习或其他摊销方法以加快编码速度也是未来工作的一个重要方向[29, 33]。完善表示图像的函数架构（例如通过神经架构搜索或剪枝）是另一个有前途的途径。在本文中，我们只是将权重转换为半精度，但通过使用更先进的模型压缩技术，可能会大幅提高性能[13, 37, 15]。最后，由于隐式表示法可将任意坐标映射到任意特征[34, 30, 10]，因此将我们的方法应用于不同类型的数据（如视频或音频）将会非常有趣。

6 CONCLUSION

在本文中，我们提出了 COIN，一种通过对像素拟合神经网络并存储所得模型权重来压缩图像的新方法。我们通过实验证明，即使不使用熵编码，这种简单的方法也能在低比特率下优于 JPEG。我们希望，在这一领域的进一步工作将带来一类新的神经数据压缩方法。

A EXPERIMENTAL DETAILS

所有模型均使用 Adam 进行了 50k 次迭代训练。我们使用了 2 个输入维度（对应 (x; y) 坐标）和 3 个输出维度（对应 RGB 值）的 MLP。坐标归一化为 [-1; 1]，RGB 值归一化为 [0; 1]。除最后一层外，我们在每一层都使用了正弦非线性，并使用了 [30] 中描述的初始化。我们使用的学习率为 2e-4。下面我们将介绍每个 bpp 级别的架构。

• 0.07bpp. Number of layers: 5, width of layers: 20.
• 0.15bpp. Number of layers: 5, width of layers: 30.
• 0.3bpp. Number of layers: 10, width of layers: 28.
• 0.6bpp. Number of layers: 10, width of layers: 40.
• 1.2bpp. Number of layers: 13, width of layers: 49.

B ADDITIONAL RESULTS

在 图6 中，我们绘制了柯达数据集中所有图像在 0.3bpp 时 COIN 和 JPEG 的性能（因为 COIN 和 JPEG 在此比特率下的性能相似）。可以看出，失真值紧随其后：COIN 难以编码的图像，JPEG 也难以编码。

[图 6：柯达数据集中 COIN 和 JPEG 所有图像的 PSNR 直方图。]