[Paper Review] 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 Methods13 references4 citations

TL;DR

CAE-ADMM proposes a novel lossy image compression method that replaces traditional entropy estimation with ADMM-based pruning to implicitly optimize bitrate by promoting sparsity in the latent code. The approach achieves superior SSIM and MS-SSIM performance compared to original CAE and traditional codecs like JPEG and JPEG 2000, with competitive inference speed and a 17.65% zero-element ratio in the latent representation after pruning.

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.

Motivation & Objective

• To address the non-differentiability of bitrate in deep learning-based lossy image compression by replacing entropy estimation with a sparsity-promoting alternative.
• To improve compression efficiency and reconstruction quality without relying on parameterized entropy estimators or complex context models.
• To explore whether ADMM-based pruning can implicitly optimize the entropy rate while maintaining or improving perceptual image quality.
• To achieve competitive inference speed while outperforming traditional codecs and the original CAE in structural similarity metrics.

Proposed method

• The method employs an encoder-decoder architecture with convolutional layers and sub-pixel deconvolution for up-sampling, using a differentiable quantizer defined as $ Q(t) = ig floor t ig floor + ext{Bernoulli}(t - ig floor t ig floor) $.
• Instead of using an entropy estimator, the bitrate is implicitly optimized via ADMM-based pruning that enforces sparsity in the quantized latent code $ ilde{\mathbf{z}} $, reducing the number of non-zero elements.
• The optimization problem is formulated as minimizing distortion plus a penalty on the cardinality of the latent code, $ R(\tilde{\mathbf{z}}) = \text{card}(\tilde{\mathbf{z}}) $, which is solved using the ADMM framework.
• The ADMM module alternates between updating the latent code, dual variables, and slack variables to enforce sparsity while maintaining reconstruction fidelity.
• A two-stage training schedule is used: first, a warm-up phase with MSE loss only, followed by a pruning phase where 10% of the smallest-magnitude elements are set to zero.
• The final encoder layer is modified to allow fine-grained control over the bitrate by adjusting the number of retained elements.

Experimental results

Research questions

• RQ1 Can ADMM-based pruning effectively replace entropy estimation in deep learning-based image compression while maintaining or improving rate-distortion performance?
• RQ2 Does promoting sparsity in the latent code through pruning implicitly optimize the entropy rate and thus reduce bitrate without explicit modeling?
• RQ3 How does the performance of CAE-ADMM compare to the original CAE and traditional codecs (JPEG, JPEG 2000) in terms of SSIM and MS-SSIM at comparable bitrates?
• RQ4 To what extent does the pruning process affect the distribution of non-zero elements in the latent code, and does this correlate with improved compression efficiency?
• RQ5 What is the inference speed of CAE-ADMM relative to traditional codecs, and can it maintain high quality at low bitrates?

Key findings

• CAE-ADMM outperforms the original CAE and traditional codecs in both SSIM and MS-SSIM on the Kodak PhotoCD dataset, with a 17.65% zero-element ratio in the latent code after pruning.
• At a target bitrate of approximately 0.3 bpp, CAE-ADMM maintains high-quality reconstruction, while JPEG degrades significantly, demonstrating robustness at low bitrates.
• The model achieves an average SSIM of $ 0.871 \pm 0.003 $ at $ 0.597 $ bpp, outperforming JPEG ($ 0.828 \pm 0.006 $) and JPEG 2000 ($ 0.793 \pm 0.013 $) in perceptual quality.
• The inference speed of CAE-ADMM is $ 0.140 \pm 0.008 $ seconds per image, which is acceptable and comparable to JPEG (0.033 s) and JPEG 2000 (0.177 s).
• The ablation study confirms that the ADMM pruning module significantly improves performance, as evidenced by the increase in zero-element ratio from 7.80% to 17.65% and the corresponding bitrate reduction from 1.684 to 1.257 bpp.
• Visual comparisons on image kodim21 show that CAE-ADMM preserves fine details better than JPEG and JPEG 2000 at low bitrates, with minimal blocking artifacts.