AAAI 2023 Oral | 字节提出非对称图像重采样模型，JPEG、WebP上抗压缩性能领先 SOTA

字节跳动-火山引擎多媒体实验室针对图像重采样模型面向图像压缩的鲁棒性，设计了一种非对称的可逆重采样框架，提出新型图像重采样模型SAIN。

图像重采样 (Image Rescaling) 任务联合优化图像下采样和上采样操作，通过对图像分辨率的下降和还原，可以用于节省存储空间或传输带宽。在实际应用中，例如图集服务的多档位分发，下采样得到的低分辨率图像往往会进行有损压缩，而有损压缩往往导致现有算法的性能大幅下降。

近期，字节跳动-火山引擎多媒体实验室首次尝试了有损压缩下的图像重采样性能优化，设计了一种非对称可逆重采样框架，基于该框架下的两点观察，进一步提出了抗压缩图像重采样模型 SAIN。该研究将一组可逆网络模块解耦成重采样和压缩模拟两个部分，使用混合高斯分布建模分辨率下降和压缩失真造成的联合信息损失，结合可微的 JPEG 算子进行端到端训练，大幅提升了对于常见压缩算法的鲁棒性。

目前对于图像重采样的研究，SOTA 方法是基于可逆网络 (Invertible Network) 构建一个双射函数 (bijective function)，其正运算将高分辨率 (HR) 图像转换为低分辨率 (LR) 图像和一系列服从标准正态分布的隐变量，逆运算则随机采样隐变量，结合 LR 图像进行上采样还原。

由于可逆网络的特性，下采样和上采样算子保持高度的对称性 (symmetry)，这就导致受压缩的 LR 图像难以通过原本学得的上采样算子进行还原。为了增强对于有损压缩的鲁棒性，该研究提出了一个基于非对称可逆框架的抗压缩图像重采样模型 SAIN (Self-Asymmetric Invertible Network)。

SAIN 模型的核心创新点如下：

提出了非对称 (asymmetric) 式的可逆图像重采样框架，解决了先前方法中由于严格对称性导致的性能下降问题；提出了增强可逆模块 (E-InvBlock)，在共享大量参数和运算的前提下，增强模型拟合能力，同时建模压缩前和压缩后的两组 LR 图像，使模型能够通过逆运算进行压缩恢复和上采样。
构建可学习的混合高斯分布，建模分辨率下降和有损压缩造成的联合信息损失，通过重参数化技巧直接优化分布参数，更加符合隐变量的实际分布。

SAIN 模型进行了 JPEG 和 WebP 压缩下的性能验证，在多个公开数据集上的表现大幅领先 SOTA 模型，相关研究已经在AAAI 2023发表。

论文地址：https://arxiv.org/abs/2303.02353

代码链接：https://github.com/yang-jin-hai/SAIN

非对称重采样框架

为了提升抗压缩性能，该研究首先设计了一种非对称式的可逆图像重采样框架，提出了 baseline 方案 Dual-IRN 模型，深入分析了该方案的不足之处后，进一步优化提出了 SAIN 模型。如上图所示，Dual-IRN 模型包含两个分支，其中 D-IRN 和 U-IRN 为两组可逆网络，分别学习 HR 图像与压缩前/压缩后 LR 图像之间的双射。

在训练阶段，Dual-IRN 模型通过可微的 JPEG 算子传递两个分支之间的梯度。而在测试阶段，模型通过 D-IRN 进行下采样获得高质量的 LR 图像，经过现实环境中的真实压缩后，再通过具有压缩感知 (compression-aware) 的 U-IRN 完成压缩恢复和上采样。

这样的非对称式框架，使得上采样和下采样的算子避免了严格的可逆关系，从根源上解决了压缩算法破坏上下采样过程对称性带来的问题，相对于 SOTA 的对称式方案大幅提升了抗压缩性能。

随后，研究人员对 Dual-IRN 模型展开进一步的分析，观察到以下两个现象：

第一，衡量 D-IRN 和 U-IRN 两个分支中间层特征的 CKA 相似性。如上图 (b) 所示，D-IRN 最后一层的输出特征（即网络生成的高质量 LR 图像）与 U-IRN 浅层的输出特征具有高度的相似性，表明 U-IRN 的浅层行为更贴近采样损失的模拟，而深层行为更贴近压缩损失的模拟。
第二，统计 D-IRN 和 U-IRN 两个分支中间层隐变量的真实分布。如上图 (c)(d) 所示，不具压缩感知的 D-IRN 的隐变量整体上满足单峰的正态分布假设，而具有压缩感知的 U-IRN 的隐变量则呈现了多峰的形态，表明有损压缩带来的信息损失形式更加复杂。

基于以上分析，研究人员从多个方面对模型进行优化，最终得到的 SAIN 模型不仅网络参数量下降了将近一半，并且还取得了进一步的性能提升。

SAIN 模型细节

SAIN 模型的架构如上图所示，主要进行了以下四个方面的改进：

整体框架方面。 基于中间层特征的相似性，将一组可逆网络模块解耦成重采样和压缩模拟两个部分，形成自非对称 (self-asymmetric) 架构，避免使用两组完整的可逆网络。在测试阶段，使用正变换取得高质量的 LR 图像，先使用逆变换进行压缩恢复，再使用逆变换进行上采样。
网络结构方面。 基于压缩损失能够借助高频信息恢复的假设提出 E-InvBlock，在模块中增加一个加性变换，使得共享大量运算的前提下可以高效地建模压缩前/压缩后的两组 LR 图像。
信息损失建模方面。 基于隐变量的真实分布，提出使用可学习的混合高斯分布建模下采样和有损压缩带来的联合信息损失，通过重参数化技巧端到端优化分布参数。
目标函数方面。 设计了多个损失函数，用于约束网络的可逆性、提高重建精度，同时在损失函数中引入真实压缩操作以增强对真实压缩方案的鲁棒性。

实验与效果评估

评估数据集为 DIV2K 验证集和 Set5、Set14、BSD100 和 Urban100 四个标准测试集。

量化评估指标为：

PSNR：Peak Signal-to-Noise Ratio，峰值信噪比，反映重建图像与原始图像的均方误差，越高越好；
SSIM：Structural Similarity Image Measurement，衡量重建图像与原始图像的结构相似度，越高越好。

在表 1 和图 3 的对比实验中，SAIN 在所有数据集上的 PSNR 和 SSIM 分数都大幅领先 SOTA 的图像重采样模型。在比较低的 QF 下，现有方法普遍出现严重的性能下降，而 SAIN 模型仍然保持最优性能。

表 1 对比实验，在 DIV2K 数据集上比较不同 JPEG 压缩质量 (QF) 下的重建质量 ( PSNR / SSIM )。

图 3 对比实验，在四个标准测试集上比较不同 JPEG QF 下的重建质量 ( PSNR )。

在图 4 的可视化结果中，可以明显看出 SAIN 还原的 HR 图像更加清晰、准确。

图 4 不同方法在 JPEG 压缩下的可视化结果对比 (×4 倍率)

在表 2 的消融实验中，研究人员还比较了其他几种结合真实压缩进行训练的候选方案。这些候选方案相较于完全对称的现有模型 (IRN) 具有更高的抗压缩性能，但在参数量和准确率上仍不如 SAIN 模型。

表 2 针对整体框架和训练策略的消融实验

在图 5 的可视化结果中，研究人员对比了 WebP 压缩失真下不同图像重采样模型的重建结果。可以发现，SAIN 模型在 WebP 压缩方案下同样表现出了最高的重建分数，能够清晰准确地还原图像细节，证明了 SAIN 对于不同压缩方案的兼容性。

*图 5 不同方法在* *WebP* *压缩下的定性和定量对比 (×2 倍率)*

除此之外，该研究还针对混合高斯分布、E-InvBlock 和损失函数等部分进行了消融实验，证明了这些改进对于结果的正向贡献。

总结与展望

火山引擎多媒体实验室针对抗压缩图像重采样提出了一个基于非对称可逆框架的模型：SAIN。该模型包含重采样和压缩模拟两个部分，使用混合高斯分布建模分辨率下降和压缩失真造成的联合信息损失，结合可微的 JPEG 算子进行端到端训练，提出 E-InvBlock 来增强模型拟合能力，大幅提升了对于常见压缩算法的鲁棒性。

火山引擎多媒体实验室是字节跳动旗下的研究团队，致力于探索多媒体领域的前沿技术，参与国际标准化工作，其众多创新算法及软硬件解决方案已经广泛应用在抖音、西瓜视频等产品的多媒体业务，并向火山引擎的企业级客户提供技术服务。实验室成立以来，多篇论文入选国际顶会和旗舰期刊，并获得数项国际级技术赛事冠军、行业创新奖及最佳论文奖。

未来，研究团队会继续优化图像重采样模型在有损压缩下的性能，并且进一步探索抗压缩视频重采样、任意倍率重采样等更加复杂的应用情景。