AdaNIC:通过动态变换路由实现实用的神经图像压缩

近年来，基于神经网络的图像压缩（NIC）得到了积极的发展，与传统方法相比，其表现出了令人印象深刻的性能。然而，由于NIC基于学习的特性，产生的浮点运算（FLOP）数量比传统算法高出几个数量级。通过用基于神经网络的非线性变换代替传统的线性变换（即DCT、DST），尽管具有更好的表示能力，但推理计算成本将是巨大的。这种困境阻碍了NIC编解码器的实际部署，这需要一种有效的方法来减少计算开销，而不损害其性能优势。对此，作者提出了一种全新的技术框架，以实现高效图像编码。这个框架主要使用了三个新的技术：（1）根据变换容量图，空间自适应卷积和归一化算子能在不同的图像块上以不均匀的计算量进行卷积。（2）提出了Just-unpenalized model capacity（JUMC）方法，来寻找每个图像块的最佳变换容量。（3）通过轻量级路由代理模型来近似预测JUMC。这种方法在BD-Rate增加最小的情况下实现了高达40%的计算加速，具备应用潜力。

题目：AdaNIC: Towards Practical Neural Image Compression via Dynamic Transform Routing
作者：Lvfang Tao, Wei Gao, Ge Li, Chenhao Zhang
来源：ICCV 2023
文章链接：https://openaccess.thecvf.com/content/ICCV2023/html/Tao_AdaNIC_Towards_Practical_Neural_Image_Compression_via_Dynamic_Transform_Routing_ICCV_2023_paper.html
内容整理：李江川

引言

自动编码器的特定变体，即压缩自动编码器（CAE），已成为神经图像压缩中流行的架构选择。采用CAE学习图像信号的紧凑非线性表示取得了巨大成功，与现有的编解码器相比，产生了相当甚至更优的率失真性能。之前的研究工作已经证明，CAE的规模与图像质量或比特率高度相关。在这种情况下，经过充分研究的信道修剪方法可能适合复杂性缓解的需要。当使用信道修剪方法去除部分信道时，过度的信道修剪可能导致率失真性能严重下降。因此，静态的信道修剪方式可能不适合进一步的率失真复杂度优化。具体结果可见图1，对于三张不同的输入图像，直接将潜在变量的通道数由192裁剪为176。深色圆点代表了原始的率失真表现，浅色圆点代表裁剪后的率失真表现。可以看到，三张图像表现出了不同的下降趋势，但复杂度的降低是一致的。更进一步的，箭头代表不同图像块的率失真表现，可以发现，同一图像的不同图像块也会有不同的率失真下降趋势。因此，这种通道裁剪方法需要更细粒度的划分，而不仅仅是作用在整张图像上。此外，作者希望研究一种动态路由解决方案，以探索率失真和复杂度的联合优化。因为，在运行时使用内容自适应优化能实现最大的系统吞吐量。由于动态路由的作用空间被设计为样本或区域自适应，因此它可以无缝集成到其他可行的解决方案中，以加速神经非线性变换，从而产生静态轻量级模型，并通过联合优化提高其性能。这种动态路由方法在运行时做出编码决策，这类似于现代图像/视频编码标准通常采用的传统RDO过程或快速算法。这种运行时权衡可以带来更大的灵活性，从而通过定制行为实现更好的速率失真或复杂性权衡。

AdaNIC:通过动态变换路由实现实用的神经图像压缩 | ICCV 2023 — 图1 静态的信道修剪方法导致的率失真结果

提出的方法

总体流程

图2给出了算法的处理流程。对于一张输入图像，首先会将其分块，每个图像块单独由路由网络处理。路由网络采用了知识蒸馏的设计，存在一个老师网络和一个学生网络。训练时，老师和学生一起训练，实际压缩时，由轻量级的学生网络负责预测当前图像块的transform capacity map(TCM)。这个TCM就决定了当前图像块应该使用的通道数的多少。下面的空间自适应CAE就是负责压缩的主干网络，它能动态调整特征图的通道数，而调整的依据就是TCM。下面将分模块介绍各个部分。

空间自适应算子

能够实现裁剪通道数的CAE在这篇工作之前就已经被提出了，一般称其为SlimCAE。但SlimCAE实现的是整张特征图的通道调整，而本文实现的是部分特征图通道数的调整。这种调整方式可以用图3形象的表示。前面提到的TCM作为一个引导，通过空间自适应的卷积和归一化操作实现通道数裁剪。在网络训练时，训练的卷积核的数量是特征图支持的通道数的最大值。在这种空间自适应算子中，TCM负责选择使用的卷积核的数量，使用的卷积核少了，对应特征图的通道数自然也减少了。

这种空间自适应算子需要将特征图进行分块卷积，如果直接按照卷积块的空间大小进行划分，就会导致信息的丢失。因此，作者使用了有重叠的分块方案，图4是一个例子。这种分块方案在数学上是很容易证明的。

通过JUMC估计TCM

作者提出的JUMC方法是以率失真函数为依据的。一般的率失真函数可以写为：

这通常可以描述为R ~ D空间上的一条直线。作者使用了另外一个参数 μ 来控制斜率，并将两个率失真损失之差，作为路由网络的预测目标，这些定义可以写为：

作为基准的直线是使用最大通道数时，图像块的率失真性能（也就是上面的 L’_old）。JUMC会在一个可以接受的误差范围内搜索使用最少通道数的率失真点，并以这个通道数作为图像块压缩时的TCM。公式化为：

这个过程可以用图5表示，黑线是 L’_old，红线代表了可接受的最大率失真损失。对于位于红线内的率失真点，会选择最少的通道数作为JUMC的预测结果。

路由代理网络

路由代理网络采用了知识蒸馏的设计思想，同时训练一个老师网络和学生网络。老师网络采用比较复杂的网络设计，先由一个主干网络提取特征，然后使用两个MLP根据特征分别预测∆Cost和z_pred，这里设计了两个损失函数：

这里设计损失函数S是为了提高网络的预测准确率。总的损失函数是两者的加权和：

学生采用了更轻量级的设计，包括更精简的主干网络，以及一个MLP只预测z_pred。学生网络的损失函数采用了知识蒸馏的常见设计，包括交叉熵和KL散度：

实验

作者在多个不同质量级别和率失真性能的模型上进行了实验，表1给出了所训练模型的具体设置。这些参数代表了特征图的最大通道数，以及率失真损失的权衡。在后面的实验中，以C1到C5表示不同的通道数，而且是依次递增的。

根据之前的介绍，C5对应的模型的率失真性能就是L’_old，而作者认为可以接受的率失真误差范围是L’_old的1%。表2给出了具体的实验结果。结果都是在不同的q下结果的平均值。C1到C4代表了直接裁剪整张特征图的通道数，能够加速的时间和率失真损失。可以看到，代理网络的预测结果非常接近真实结果，而且在近乎无损的情况下，实现了远超C4的加速效果。