Leveraging Second-Order Curvature for Efficient Learned Image Compression: Theory and Empirical Evidence

TL;DR

The paper tackles slow convergence and suboptimal rate–distortion trade-offs in learned image compression caused by gradient conflicts between rate and distortion. It proposes SOAP, a second-order quasi-Newton optimizer with Kronecker-structured preconditioning, as a drop-in replacement that accelerates training and improves RD performance across multiple LIC architectures. The authors provide theoretical analysis showing Newton preconditioning aligns gradients intra-step and inter-step, and they validate this with empirical metrics, including reduced steps/time and BD-Rate gains, plus a marked reduction in activation/latent outliers that enhances post-training quantization robustness. A key practical takeaway is that optimization strategy—specifically curvature information—offers a tangible lever to boost LIC efficiency and deployability without architectural changes. The work also presents a preliminary extension to learned video compression, suggesting broader applicability of curvature-aware optimization in high-dimensional compression tasks.

Abstract

Training learned image compression (LIC) models entails navigating a challenging optimization landscape defined by the fundamental trade-off between rate and distortion. Standard first-order optimizers, such as SGD and Adam, struggle with \emph{gradient conflicts} arising from competing objectives, leading to slow convergence and suboptimal rate-distortion performance. In this work, we demonstrate that a simple utilization of a second-order quasi-Newton optimizer, \textbf{SOAP}, dramatically improves both training efficiency and final performance across diverse LICs. Our theoretical and empirical analyses reveal that Newton preconditioning inherently resolves the intra-step and inter-step update conflicts intrinsic to the R-D objective, facilitating faster, more stable convergence. Beyond acceleration, we uncover a critical deployability benefit: second-order trained models exhibit significantly fewer activation and latent outliers. This substantially enhances robustness to post-training quantization. Together, these results establish second-order optimization, achievable as a seamless drop-in replacement of the imported optimizer, as a powerful, practical tool for advancing the efficiency and real-world readiness of LICs.

Figures (8)

• Figure 1: Comparison of Testing Loss: Epochs vs. R-D Loss for Various LICs. First 10 epochs are omitted for better visualization. The SOAP optimizer demonstrates significantly faster convergence compared to Adam across multiple LICs. Evaluation is performed on the Kodak dataset with $\lambda = 0.013$; the R-D loss is computed as $\lambda \cdot 255^2 \cdot \text{MSE} + \text{Bpp}$. Longer training period results are shown in Sec. \ref{['subsec:longtrain']}.
• Figure 2: Comparison of Testing Loss: Wall-Time vs. R-D Loss for Various LICs. Training with the SOAP optimizer leads to much faster and more stable convergence than Adam when comparing wall-clock time. Results are measured on the Kodak dataset with $\lambda = 0.013$. Longer training period results are shown in Sec. \ref{['subsec:longtrain']}.
• Figure 3: R-D curves of various methods.Please zoom in for more details.
• Figure 4: Evolution of intra-step and inter-step gradient scores for ELIC trained with Adam vs. SOAP. SOAP achieves high intra-step and inter-step scores, while Adam exhibits negative intra-step scores and oscillatory inter-step scores, highlighting SOAP’s ability to suppress gradient conflicts.
• Figure 5: Scaled deviation maps for ELIC latent representations. Each row shows the input image (left), latent scaled deviation with Adam (middle), and SOAP (right). SOAP consistently suppresses extreme values and yields lower maximum scaled deviation. (Best viewed zoomed in.)

Theorems & Definitions (18)

• Theorem 1: Newton approximation of SOAP
• Lemma 1: Inter-step alignment for Newton
• Proposition 1: Newton aligns component steps near the optimum
• Theorem 2: Conditional Newton approximation for SOAP
• proof
• proof
• Proposition 2: Local diagonal-preconditioner approximation
• proof
• Proposition 3
• proof