[论文解读] CAE-ADMM: Implicit Bitrate Optimization via ADMM-based Pruning in Compressive Autoencoders

Haimeng Zhao, Peiyuan Liao|arXiv (Cornell University)|Jan 22, 2019

Image and Signal Denoising Methods参考文献 13被引用 4

一句话总结

CAE-ADMM 提出了一种新颖的有损图像压缩方法，用基于 ADMM 的剪枝替代传统的熵估计，通过在潜在码中促进稀疏性来隐式优化比特率。该方法在 SSIM 和 MS-SSIM 性能上优于原始 CAE 及 JPEG 和 JPEG 2000 等传统编码器，在推理速度上具有竞争力，且剪枝后潜在表示中的零元素比例达到 17.65%。

ABSTRACT

We introduce ADMM-pruned Compressive AutoEncoder (CAE-ADMM) that uses Alternative Direction Method of Multipliers (ADMM) to optimize the trade-off between distortion and efficiency of lossy image compression. Specifically, ADMM in our method is to promote sparsity to implicitly optimize the bitrate, different from entropy estimators used in the previous research. The experiments on public datasets show that our method outperforms the original CAE and some traditional codecs in terms of SSIM/MS-SSIM metrics, at reasonable inference speed.

研究动机与目标

• 通过用一种促进稀疏性的替代方法替代熵估计，解决基于深度学习的有损图像压缩中比特率不可微的问题。
• 在不依赖参数化熵估计器或复杂上下文模型的前提下，提升压缩效率和重建质量。
• 探究基于 ADMM 的剪枝是否能通过隐式优化熵速率来减少比特率，同时保持或提升感知图像质量。
• 在保持与传统编码器相当或更优的结构相似性指标的同时，实现具有竞争力的推理速度。

提出的方法

• 该方法采用编码器-解码器架构，结合卷积层和子像素转置卷积进行上采样，使用可微分量化器 $ Q(t) = \big\rfloor t \big\rfloor + \text{Bernoulli}(t - \big\rfloor t \big\rfloor) $。
• 不使用熵估计器，而是通过基于 ADMM 的剪枝隐式优化比特率，强制对量化后的潜在码 $ \tilde{\mathbf{z}} $ 实现稀疏性，从而减少非零元素数量。
• 将优化问题表述为最小化失真加上对潜在码基数的惩罚，$ R(\tilde{\mathbf{z}}) = \text{card}(\tilde{\mathbf{z}}) $，并使用 ADMM 框架求解。
• ADMM 模块通过交替更新潜在码、对偶变量和松弛变量，在保持重建保真度的同时强制实现稀疏性。
• 采用两阶段训练策略：首先进行仅使用 MSE 损失的预热阶段，随后进入剪枝阶段，将最小幅值的 10% 元素置零。
• 最终的编码器层经过修改，以允许通过调整保留元素数量来实现对比特率的细粒度控制。

实验结果

研究问题

• RQ1 基于 ADMM 的剪枝是否能有效替代基于深度学习的图像压缩中的熵估计，同时保持或提升率失真性能？
• RQ2 通过剪枝在潜在码中促进稀疏性，是否能隐式优化熵速率，从而在无需显式建模的情况下降低比特率？
• RQ3 在相近比特率下，CAE-ADMM 在 SSIM 和 MS-SSIM 指标上与原始 CAE 及传统编码器（JPEG、JPEG 2000）相比表现如何？
• RQ4 剪枝过程在多大程度上影响了潜在码中非零元素的分布，这种分布变化是否与压缩效率的提升相关？
• RQ5 CAE-ADMM 的推理速度与传统编码器相比如何，是否能在低比特率下保持高质量？

主要发现

• 在 Kodak PhotoCD 数据集上，CAE-ADMM 在 SSIM 和 MS-SSIM 上均优于原始 CAE 及传统编码器，剪枝后潜在码中的零元素比例达到 17.65%。
• 在目标比特率为约 0.3 bpp 时，CAE-ADMM 仍能保持高质量重建，而 JPEG 显著退化，表明其在低比特率下具有强鲁棒性。
• 在 0.597 bpp 时，CAE-ADMM 的平均 SSIM 达到 $ 0.871 \pm 0.003 $，优于 JPEG ($ 0.828 \pm 0.006 $) 和 JPEG 2000 ($ 0.793 \pm 0.013 $)，在感知质量上表现更优。
• CAE-ADMM 的推理速度为 $ 0.140 \pm 0.008 $ 秒/图像，可接受且与 JPEG（0.033 秒）和 JPEG 2000（0.177 秒）相当。
• 消融实验表明，ADMM 剪枝模块显著提升了性能，表现为零元素比例从 7.80% 提升至 17.65%，比特率相应从 1.684 降低至 1.257 bpp。
• 在图像 kodim21 上的视觉对比显示，CAE-ADMM 在低比特率下比 JPEG 和 JPEG 2000 更好地保留了细节，且阻塞伪影极少。