Balanced Rate-Distortion Optimization in Learned Image Compression

TL;DR

This paper tackles imbalance in rate-distortion optimization for learned image compression by formulating LIC R-D as a multi-objective problem. It introduces two balanced strategies: Solution 1, a coarse-to-fine gradient-descent method for training LIC models from scratch, and Solution 2, an analytic quadratic-programming approach for fine-tuning pre-trained models. Both methods adapt gradient updates via adaptive weights on the rate and distortion objectives to achieve balanced progress, yielding about a 2% BD-Rate improvement with modest training-time overhead. The work demonstrates consistent R-D gains across multiple LIC architectures and datasets, while maintaining unchanged inference, and lays groundwork for more efficient MOO-based optimization in learned compression.

Abstract

Learned image compression (LIC) using deep learning architectures has seen significant advancements, yet standard rate-distortion (R-D) optimization often encounters imbalanced updates due to diverse gradients of the rate and distortion objectives. This imbalance can lead to suboptimal optimization, where one objective dominates, thereby reducing overall compression efficiency. To address this challenge, we reformulate R-D optimization as a multi-objective optimization (MOO) problem and introduce two balanced R-D optimization strategies that adaptively adjust gradient updates to achieve more equitable improvements in both rate and distortion. The first proposed strategy utilizes a coarse-to-fine gradient descent approach along standard R-D optimization trajectories, making it particularly suitable for training LIC models from scratch. The second proposed strategy analytically addresses the reformulated optimization as a quadratic programming problem with an equality constraint, which is ideal for fine-tuning existing models. Experimental results demonstrate that both proposed methods enhance the R-D performance of LIC models, achieving around a 2\% BD-Rate reduction with acceptable additional training cost, leading to a more balanced and efficient optimization process. Code will be available at https://gitlab.com/viper-purdue/Balanced-RD.

Figures (3)

• Figure 1: Comparison of loss trends and improvement speeds for the M&S hyperprior minnen2018joint model with $\lambda = 0.013$ from epoch 10 to 150, where Distortion = MSE $\times \lambda \times 255^2$. The first 10 epochs are omitted for improved readability. (a) Testing loss trend for the standard R-D optimization versus the proposed balanced R-D optimization. The balanced approach demonstrates more stable and smoother loss improvements, showing simultaneous and consistent reductions in both distortion and bits-per-pixel (bpp) metrics. In contrast, the standard method focuses on reducing distortion, with the bpp loss only gradually increasing in this period. (b) Comparison of loss improvement speeds, Eq. \ref{['eq:speed']}, for standard and balanced R-D optimizations. The balanced R-D optimization yields a more consistent, less volatile improvement speed, outperforming the standard approach by achieving steady and reliable convergence across epochs. This highlights the advantage of the balanced approach in optimizing both objectives cohesively and effectively.
• Figure 2: R-D curves of various methods. Please zoom in for more details.
• Figure 3: Ablation experiments on proposed methods. M&S Hyperprior model, $\lambda$ = 0.013. Kodak dataset.