COIN++: Neural Compression Across Modalities【论文笔记】

论文地址: https://arxiv.org/abs/2201.12904
代码地址: https://github.com/emiliendupont/coinpp

Abstract

神经压缩算法通常基于自动编码器，对于不同的数据模式需要特定的编解码算法。在本文中，我们提出了COIN++，这是一种神经压缩框架，可以无缝处理各种数据模式。我们的方法基于将数据转换为隐式神经表示，即将坐标（如像素位置）映射到特征（如RGB值）的神经函数。然后，我们存储应用于元学习基网络的调制作为数据的压缩码，而不是直接存储隐式神经表示的权重。我们进一步对这些调制进行量化和熵编码，从而获得较大的压缩增益，同时与基准相比将编码时间减少了两个数量级。我们通过压缩各种数据模式，从图像和音频到医疗和气候数据，实证了我们方法的有效性。

1. 介绍

据估计，每天都会有EB的数据被创建。这些数据由多种多样的数据模式组成，每种模式都可以从压缩中受益。然而，在神经压缩方面的大多数工作只关注图像和视频数据。在本文中，我们介绍了一种新的神经压缩方法，称为COIN++，适用于从图像和音频到医疗和气候数据的广泛数据模式（图1）。

图1：COIN++通过优化将各种数据模式转换为神经网络，然后将这些神经网络的参数存储为数据的压缩码。

大多数神经压缩算法基于自动编码器的。编码器将图像映射到潜在表示，该潜在表示被量化并熵编码到比特流中。然后将比特流传输到解码器中，被重建为图像。对编码器和解码器的参数进行训练，以共同最小化重建误差或失真，以及压缩码的长度或码率。为了获得良好的性能，这些算法在很大程度上依赖特定于图像的编码器和解码器架构。将这些模型应用于新的数据模式需要设计新的编码器和解码器，这有很大的挑战性。

最近有一种新的神经压缩框架，称为 COIN（Compression with Implicit Neural representations），它绕过了对专用编码器和解码器的需要。COIN 没有直接压缩图像，而是通过一个神经网络将图像的像素位置映射到RGB值，并将该网络的量化权重存储为图像的压缩码。虽然Dupont等人仅将 COIN 应用于图像，但它有望存储其他数据模式。实际上，将坐标（例如像素位置）映射到特征（例如 RGB 值）的神经网络，通常称为隐式神经表示（INR），现已用于表示符号距离函数、体素网格、3D 场景、温度场、视频、音频等。因此，将数据转换为 INR 并对其进行压缩的类似 COIN 的方法有希望构建适用于各种模式的灵活神经编解码器。

在本文中，我们识别并解决了 COIN 的几个关键问题，并提出了一种适用于多种模式的压缩算法，我们称之为COIN++。

更具体地说，我们确定了COIN的以下问题：

1. 编码速度慢：压缩单个图像可能需要一个小时。
2. 缺乏共享结构：由于每个图像都是独立压缩的，因此网络之间没有共享信息。
3. 性能远低于最先进的（SOTA）图像编解码器。

我们通过以下方式解决这些问题：

1. 与COIN的分钟或小时相比，使用元学习将编码时间从两个数量级以上减少到小于一秒。
2. 学习编码共享结构的基础网络，并对该网络应用调制来编码实例特定信息。
3. 量化和熵编码调制。我们的方法不经在压缩和速度方面都大大超过了COIN，而且在经过充分研究的模式（如图像）上，还部分缩小了与SOTA编解码器的差距。

而且，COIN++适用于传统方法无法使用的广泛数据模式，使其成为非标准领域神经压缩的一个有前景的工具。

2. 方法

在文中，我们考虑压缩数据可以用坐标 $x\in \mathcal{X}$ 和 特征 $y\in \mathcal{Y}$ 来表示。例如，图像能够被一组像素坐标 $x=(x,y)in{\mathbb{R}}^2$ 和它们的对应 RGB 值 $y=(r,g,b)in\{0,1,...,255{\}}^3$ 来表示。同样的，MRI 扫描可以由一组 3D空间中的位置 $x=(x,y,z)$ 和一个强度值 $y\in {\mathbb{R}}^{+}$ 来表示。给定单个数据点作为坐标和特征对 $d=\{(x_i,y_i){\}}_{i=1}^n$ 的集合（例如，图像作为n个像素位置和RGB值的集合），$COIN$ 方法包括将带有参数 $\theta$ 的神经网络 $f_{\theta }:\mathcal{X}\to \mathcal{Y}$ 拟合到数据点，通过最小化函数：
$$\mathcal{L}(\theta ,d)=\sum _{i=1}^n\Vert f_{\theta }(x_i)-y_i{\Vert }_2\text{\hspace{1em}(1)}$$
权重 $\theta$ 被量化并存储作为数据点 $d$ 的压缩表示。 神经网络 $f_{\theta }$ 由 $SIREN$ 参数化，即具有正弦激活函数的 $MLP$ ，这对于拟合像自然图像这样的高频数据所必要的。具体而言， $SIZEN$ 层由应用于隐藏特征向量 $h\in {\mathbb{R}}^d$ 的元素 $sin$ 所定义：
$$SIREN(h)=sin(w_0(Wh+b))\text{\hspace{1em}(2)}$$
其中 $W\in {\mathbb{R}}^{d\times d}$ 是一个权重矩阵，$b\in {\mathbb{R}}^d$ 是一个偏差矩阵，$w_0\in {\mathbb{R}}^{+}$ 是一个正比例因子。

虽然这个方法很普遍，但是这里有几个关键性问题。首先压缩涉及最小化公式（1），编码速度会很慢。例如，在 1080Ti GPU 上压缩 Koadk 数据集中的单个图像需要将近一个小时。其次，每个数据点 $d$ 都配有单独的神经网络 $f_{\theta }$，因此数据点之间没有共享信息。这显然是次优的。例如，自然图像共享很多公共结构，不需要为每个单独的图像重复存储。在以下几节中，我们将展示我们提出的方法 COIN++，如何在保持COIN通用性的同时解决这些问题。

2.1 存储调制

虽然COIN将每个图像存储为一个单独的神经网络，但我们训练一个跨数据点共享的基网络，并对该基网络进行调制来参数化单个数据点。给定一个基网络，例如多层感知器（MLP)，我们使用FiLM层，通过应用元素尺度$\gamma \in {\mathbb{R}}^d$和位移$\beta \in {\mathbb{R}}^d$来调制网络的隐藏特征$h\in {\mathbb{R}}^d$。
$$FiLM(h)=\gamma \odot h+\beta \text{\hspace{1em}(3)}$$
在给定固定MLP的情况下，我们可以在每一层应用不同的尺度和移位来参数化神经网络族。每个神经网络函数由一组尺度和位移指定，这些尺度和位移统称为调制。最近，FiLM方法也被应用于INRs。Chan等人通过SIREN网络在生成对抗网络中参数化生成器，并通过应用调制$sin(\gamma \odot (Wh+b)+\beta )$于该网络来生成样本。类似地，Mehta等人通过使用比例因子$\alpha \odot sin(Wh+b)$参数化INR族。这两种方法都可以修改为使用一组低维潜在向量去映射调制，而不是直接应用调制。Chan等人将潜在向量映射到MLP的尺度和位移，而Mehta等人通过与基网络形状相同的MLP映射潜在向量，并使用该网络的隐藏激活作为调制。

我们对COIN++使用了类似的方法。存储应用于共享基网络的一组调制，而不是存储每个数据点的神经网络权重。更具体地说，给定一个基SIZEN网络，我们仅使用位移$\beta \in {\mathbb{R}}^d$作为调制:
$$sin(w_0(Wh+b+\beta ))\text{\hspace{1em}(4)}$$
实际上，我们根据经验发现，在MLP的每一层上，只使用移位与同时使用移位和尺度具有相同的性能，而且只需要一半的调制和一半的存储。但仅使用尺度会产生相当差的性能。为了进一步减少存储，我们使用了一个线性映射到调制的潜在向量，如图2所示。在稍微过多的符号中，我们也称该向量为调制或潜在调制。然后，我们将数据点$d$（例如图像）存储为一组（潜在）调制$\varphi$。为了解码数据点，我们只需在每个坐标$x$处计算调制的基本网络$f(\theta )$，
$$y=f_{\theta }(x;\varphi )\text{\hspace{1em}(5)}$$

图2 COIN++ 架构：潜在调制 $\varphi$（绿色）映射到调制（蓝色），这些调制被添加到基本网络 $f(\theta )$（白色）的激活中，以参数化单个函数，可以在坐标 $x$ 处评估该函数以获得特征 $y$。

如 图3 所示。为了将一组调制 $\varphi$ 拟合到数据点 $d$，我们保持基网络的参数 $\theta$ 固定，并最小化：
$$\mathcal{L}(\theta ,\varphi ,d)=\sum _{i=1}^n\Vert f_{\theta }(x_i;\varphi )-y_i{\Vert }_2\text{\hspace{1em}(6)}$$

与 COIN 不同的是，每个数据点 $d$ 存储为单独的神经网络 $f_{\theta }$，COIN++ 只需要存储 $O(n)$ 调制，而不是 $O(n^2)$ 权重，其中 $n$ 是 MLP 的宽度。此外，这种方法允许我们在基网络中存储共享信息，并在调制中存储实例特定信息。例如，对于自然图像，基本网络对自然图像常见的结构进行编码，而调制存储重建单个图像所需的信息。

图3：通过应用调制 ${\varphi }^{(1)},{\varphi }^{(2)},{\varphi }^{(3)}$ 到一个基网络 $f_{\theta }$ ，我们获得了不同的函数，通过不同坐标下的评估函数，它们能够被解码为数据点 $d^{(1)},d^{(2)},d^{(3)}$ 。同时我们在这里展示了图像，同样的原理可以运用在一系列的数据模式中。

2.2 元学习调制

给定一个基网络 $f_{\theta }$ ，我们能够编码数据点 $d$ 通过最小化公式 6。然而我们始终面临着两个问题：1.我们需要学习基网络的权重 $\theta$，2.通过公式 6 编码数据点太慢了，需要数以千计的梯度下降迭代。COIN++ 使用元学习解决了这些问题。

最近，Sitzmann和Tancik等人表明，应用MAML到INRs中能够减少测试时的拟合时间到几个梯度步长。不进行随机初始化的梯度下降来最小化 $\mathcal{L}(\theta ,d)$，而是用元学习初始化 ${\theta }^{\ast }$，从而可以在几个梯度步骤内最小化 $\mathcal{L}(\theta ,d)$。更具体的，假设我们得到了由N个点 $\{d^{(j)}{\}}_{j=1}^N$ 组成的数据集，从初始化 $\theta$ 开始，数据点 $d^{(j)}$ 上的MAML内环步骤为：
$${\theta }^{(j)}=\theta -\alpha {▽}_{\theta }\mathcal{L}(\theta ,d^{(j)})\text{\hspace{1em}(7)}$$
其中 $\alpha$ 是内环学习率。然后，我们关注的是学习良好的初始化 ${\theta }^{\ast }$，以便在整个数据点集$\{d^{(j)}{\}}_{j=1}^N$ 上经过几个梯度步骤后，使得损失 $\mathcal{L}(\theta ,d^{(j)})$ 最小化。为了更新初始化 $\theta$ ，我们接着以外环学习率 $\beta$ 执行外环步骤，通过:
$$\theta \leftarrow \theta -\beta {▽}_{\theta }\sum _{j=1}^N\mathcal{L}({\theta }^{(j)},d^{(j)})\text{\hspace{1em}(8)}$$
在我们的例子中，MAML不能直接使用，因为在测试时，我们只拟合调制 $\varphi$，而不拟合共享参数 $\theta$。因此，我们需要元学习一个 $\theta$ 和 $\varphi$ 的初始化，以便在给定新数据点的情况下，可以快速计算调制 $\varphi$ ，同时保持 $\theta$ 恒定。实际上，我们只存储每个数据点的调制，并在所有数据点上共享参数 $\theta$。对于COIN++，内部循环的单个步骤由以下公式得出：
$${\theta }^{(j)}=\theta -\alpha {▽}_{\varphi }\mathcal{L}(\theta ,\varphi ,d^{(j)})\text{\hspace{1em}(9)}$$
其中 $\theta$保持不变。Zintgraf等人之前已经研究过在参数子集上执行内循环，称为 CAVIA。正如在 CAVIA中观察到的那样，元学习 $\varphi$ 的初始化是冗余的，因为它可以被吸收到基本网络权重 $\theta$ 的偏差参数中。因此，我们只需要元学习共享参数初始化 $\theta$。外部循环的更新规则由下式给出：
$$\theta \leftarrow \theta -\beta {▽}_{\theta }\sum _{j=1}^N\mathcal{L}(\theta ,{\varphi }^{(j)},d^{(j)})\text{\hspace{1em}(10)}$$
然后，内环更新调制 $\varphi$，而外环更新共享参数 $\theta$。该算法允许我们元学习基网络，以便可以轻松快速地拟合每组调制（见图4）。在实践中，我们发现，与COIN的数千个梯度步骤相比，我们只用3个梯度步骤就可以得到令人信服的结果。

2.3 调制的patch、量化和熵编码

大规模数据的patch ：虽然基网络的元学习允许我们将新数据点快速编码为调制，但训练过程成本高，因为 MAML 必须通过内环获取梯度。对于大数据点（如高分辨率图像或MRI扫描），这可能会变得非常昂贵。虽然存在一阶近似，但我们发现它们严重阻碍了性能。相反，为了减少内存使用，我们在训练期间将数据点分成4个随机patch。例如，对于大规模图像，我们在32×32 patch上进行训练。在训练时，我们学习一个基本网络，这样调制可以很容易地匹配patch。在测试时，我们将一个新图像分割成块，并计算每个块的调制。然后，图像由所有patch的调制集表示（见 图4 ）。我们对其他数据模式使用了类似的方法，例如，将MRI扫描分割为3D补丁。

图4：（左）我们元学习基本网络的参数 ${\theta }^{\ast }$，使得调制 $\varphi$可以很容易地在几个梯度步骤中拟合。（右）在训练期间，我们随机抽取补丁，而在测试时，我们将数据点划分为补丁，并将调制适应每个补丁。

量化：虽然 COIN 将神经网络权重从32位量化到16位以减少存储，但超出此范围的量化严重阻碍了性能[15]。相反，我们发现调制可以令人惊讶地量化。在元学习期间，调制由32位浮点表示。为了将这些量化到较短的比特宽度，我们只需使用均匀量化。我们首先将调制剪裁在其平均值的3个标准差内。然后，我们将该间隔分割为 $2^b$个大小相等的存储单元（其中 $b$是位的数量）。值得注意的是，我们发现将比特数从32减少到5（即将符号数从 $10^9$个以上减少到32个）只会导致重建精度的小幅度降低。然后，简单地应用均匀量化将压缩提高了6倍，重建质量成本很低。

熵编码

几乎所有编解码器的核心组件都是熵编码，它允许使用例如算术编码对量化代码进行无损压缩。这依赖于量化码的分布模型。与量化一样，我们使用一种非常简单的方法来建模此分布：我们计算每个量化调制值的频率，并使用此分布进行算术编码。在我们的实验中，这在重建质量方面没有成本的情况下减少了 8-15%的存储。虽然这种简单的熵编码方案工作良好，但我们期望更复杂的方法能够显著提高性能，这是未来工作的一个令人振奋的方向。

3.相关工作

3.1 神经压缩

正如Ballé等人最初介绍的那样，学习的压缩方法通常基于联合最小化比率和失真的自动编码器。Ballé等人通过添加超先验来扩展这一点，而Mentzer、Minnen、Lee等人使用自回归模型来改进熵编码。Cheng等人通过为潜在代码的分布添加注意力和高斯混合模型来提高熵模型的准确性，而Xie等人使用可逆卷积层来进一步提高性能。虽然其中大多数是在传统的失真度量（如MSE或SSIM）上优化的，但其他工作探索了使用生成对抗网络优化感知度量。神经压缩也被应用于视频和音频。

3.2 隐式神经表示和压缩

除COIN外，最近的几项工作也探索了INR在压缩中的应用。Davies等人使用神经网络对3D形状进行编码，并表明与传统的抽取网格相比，这可以减少内存使用。Chen等人用卷积神经网络表示视频，卷积神经网络以时间索引作为输入，并在视频中输出一帧。通过剪枝、量化和熵编码该网络的权重，作者实现了接近标准视频编解码器的压缩性能。Lee等人对INR进行了元学习稀疏和参数有效初始化，并表明这可以减少以给定重建质量存储图像所需的参数数量，尽管它还不能与JPEG等图像编解码器竞争。Lu等人、Isik等人探讨了INR在体积压缩中的应用。两个并行工作还使用函数表示进行图像和视频压缩。Strümpler等人元学习MLP初始化，然后量化和熵编码适合图像的MLP权重，从而比COIN获得更大的性能增益。然而，他们的方法仍然需要在测试时进行数万次迭代才能完全收敛，同时性能不如JPEG2000等图像编解码器。Zhang等人使用INR（量化和熵编码）压缩视频中的帧，同时学习流扭曲以模拟帧之间的差异。虽然性能仍然落后于标准视频编解码器，但视频基准测试的结果很有希望。据我们所知，这些工作都没有考虑INR在跨数据模式构建统一的压缩框架。

4. 实验

我们评估了四种数据模式：图像、音频、医疗数据和气候数据。我们在PyTorch中实现所有模型，并在单个GPU上进行训练。我们使用SGD作为内环，学习速率为1e-2，Adam作为外环，学习速率为1e-6或3e-6。我们规范化坐标 x，使其位于[−1，1]，特征 y位于[0，1]。再现所有结果所需的完整实验细节见附录。我们使用压缩和真实数据之间的均方误差来训练COIN++。按照标准，我们使用 PSNR（dB）测量重建性能（或失真），其定义为 $PSNR=-10lo{g}_{10}(MSE)$。我们以每像素比特数（bpp）来衡量压缩数据的大小（或速率），bpp由 ${\frac{numberofbits}{numberofpixels}}^2$和音频的千比特/秒（kpbs）给出。我们根据大量基线对COIN++进行基准测试，包括标准图像编解码器-JPEG、JPEG2000、BPG和VTM-基于自动编码器的神经压缩-BMS、MBT和CST-标准音频编解码器-MP3和COIN。为了清晰起见，我们对不同的编解码器使用一致的颜色，并用实线绘制学习的编解码器，用虚线绘制标准编解码器。

4.1 图像 CIFAR10

我们使用128、256、384、512、768和1024个潜在调制在CIFAR10上训练COIN++。如 图5 所示，COIN++远远优于COIN、JPEG和JPEG2000，同时部分缩小了与BPG的差距，尤其是在低比特率下。据我们所知，这是INRs压缩首次优于JPEG2000等图像编解码器。COIN++和SOTA编解码器（BMS，CST）之间的剩余差距可能是由于熵编码：我们使用第2.3节中描述的简单方案，而BMS和CST使用深度生成模型。我们假设对调制使用深度熵编码将显著减少或缩小这一差距。图5 显示了我们的模型和BPG之间的定性比较，以突出使用COIN++获得的压缩伪影的类型。为了彻底分析和评估COIN++的每个组成部分，我们进行了大量的消融实验。

（图5: (左)CIFAR10上的率失真图。(右)对相似重建质量的模型的压缩伪影进行定性比较。COIN++在3.29 bpp时达到32.4dB，而BPG在1.88 bpp时达到31.9dB。）

量化比特宽度

将调制量化到较低的比特宽度会产生更多的压缩码，但代价是重建精度。为了理解这两者之间的权衡，我们在 图6a 中显示了从3位量化到8位时的率失真图。可以看出，最佳比特宽度低得惊人：5比特在低比特率下是最佳的，而6比特在高比特率下是最佳的。通过量化调制获得的定性伪影如附录中的图13所示。

量化 COIN 和COIN++

我们比较了 图6b 中COIN和COIN+量化导致的峰值信噪比下降。可以看出，调制是显著可量化的：当直接量化 COIN权重时，性能在14位左右显著下降，而量化调制即使在使用5位时也会产生较小的峰值信噪比下降。然而，如 图6c 所示，对于较大的模型，量化后的峰值信噪比下降更大。

熵编码

熵编码。附录中的图12显示了全精度、量化和熵编码调制的率失真图。可以看出，量化和熵编码都显著提高了性能。

编码时间

编码时间。图6d 显示了CIFAR10上COIN++、COIN和BPG的平均编码时间（硬件细节见附录B.1）。可以看出，COIN++压缩图像的速度比COIN快300倍，同时实现了更好的4倍压缩率。注意，这些结果是通过单独压缩每个图像获得的。当使用成批图像时，当使用10个内部循环步骤时，我们可以在4分钟内压缩整个CIFAR10测试集（10k图像）（当使用3个步骤时，压缩时间仅为一分钟多）。此外，如附录中的图14所示，COIN++只需要3个梯度步骤就可以达到与COIN在10000个步骤中相同的性能，同时使用4×更少的存储空间。

4.3 patch压缩

为了评估第2.3节中的patch方法，并证明COIN++可以扩展到大数据（尽管在性能上有成本），我们在图像、音频和MRI数据上测试了我们的模型。

大规模图像，koadk

Kodak数据集包含24幅大小为768×512的大比例尺图像。为了训练模型，我们使用Vimeo90k数据集[64]中的随机32×32patch，其中包含大小为448×256的154k幅图像。在评估时，每个Kodak图像被分割成384个32×32面片，分别进行压缩。由于我们没有对图像的全局结构建模，因此，与不需要patch的情况相比，我们预计性能会显著下降。如图8所示，COIN++的性能确实有所下降，但在低比特率下仍优于COIN和JPEG。我们希望通过对图像的全局结构进行建模（例如，两片蓝天几乎完全相同，但当前设置中未利用信息冗余），可以极大地改善这一点，但请留待将来的工作。

5. 总结，限制和未来工作

总结：我们介绍了COIN++，这是第一个适用于多种模式的神经编解码器。与COIN相比，我们的框架在压缩和编码时间方面都显著提高了性能，同时与JPEG等成熟的编解码器竞争。虽然COIN++不能匹配SOTA编解码器的性能，我们希望我们的工作将有助于扩大神经压缩适用的领域范围。

限制：COIN++的主要缺点是，由于MAML需要二阶梯度，训练模型是内存密集型的。这反过来限制了可伸缩性，并要求我们对大数据使用补丁。设计有效的一阶近似或完全避开元学习将缓解这些问题。此外，训练coin++有时会不稳定，尽管模型通常会从损失的不稳定中恢复(参见附录中的图11)。此外，还有一些我们的框架不能处理的常见模式，如文本或表格数据，因为这些不容易表示为连续函数。最后，COIN++仍然落后于SOTA编解码器。然而，我们相信未来的工作有几个有趣的方向来缩小这一差距。

未来的工作：在其目前的形式中，COIN++使用非常基本的量化和熵编码方法——对这两个步骤使用更复杂的技术可能会带来很大的性能提升。事实上，最近对函数分布建模的成功[52,3,57,16]表明，使用深度生成模型来学习熵编码调制的分布可以获得很大的收益。同样，更好的训练后量化[43,29]或量化感知训练[26,17]也会提高性能。更普遍的是，模型压缩文献中有大量的方法可以应用于COIN++[11,31]。对于大尺度数据，对patch进行全局结构建模而不是对patch进行单独编码和熵编码将是一件有趣的事情。此外，INRs领域正在迅速发展，这些进展也有可能改善COIN++。例如，Martel等人[35]使用自适应补丁将INRs缩放到10亿像素图像——这样的输入分区类似于[7]中使用的可变大小块。此外，使用更好的激活函数[50]来增加PSNR和平衡模型[21]来减少内存使用是未来令人兴奋的研究方向。