Balanced Rate-Distortion Optimization in Learned Image Compression

TL;DR

This paper introduces two adaptive gradient adjustment strategies for learned image compression that reformulate rate-distortion optimization as a multi-objective problem, improving balance and efficiency in model training.

Contribution

It proposes novel balanced R-D optimization methods based on multi-objective reformulation, enhancing training stability and compression performance in LIC models.

Findings

Achieves around 2% BD-Rate reduction

Improves balance between rate and distortion during training

Provides methods suitable for both training from scratch and fine-tuning

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.