On Disentangled Training for Nonlinear Transform in Learned Image Compression

TL;DR

This work tackles the slow training convergence of learned image compression (LIC) by reframing energy compaction into two controllable components: feature decorrelation and uneven energy modulation. It introduces a lightweight linear auxiliary transform (AuxT) with wavelet-based shortcuts (WLS) that acts in parallel with the nonlinear transforms, providing coarse decorrelation and energy modulation so the nonlinear transform can learn finer details more quickly. Across multiple LIC architectures and datasets, AuxT achieves substantial training-time reductions (often 2×) with competitive or improved RD performance, aided by a progressive, energy-preserving design and an orthogonality constraint. The approach offers a practical path to faster deployment and fine-tuning of LIC models in diverse image-synthesis and compression tasks.

Abstract

Learned image compression (LIC) has demonstrated superior rate-distortion (R-D) performance compared to traditional codecs, but is challenged by training inefficiency that could incur more than two weeks to train a state-of-the-art model from scratch. Existing LIC methods overlook the slow convergence caused by compacting energy in learning nonlinear transforms. In this paper, we first reveal that such energy compaction consists of two components, i.e., feature decorrelation and uneven energy modulation. On such basis, we propose a linear auxiliary transform (AuxT) to disentangle energy compaction in training nonlinear transforms. The proposed AuxT obtains coarse approximation to achieve efficient energy compaction such that distribution fitting with the nonlinear transforms can be simplified to fine details. We then develop wavelet-based linear shortcuts (WLSs) for AuxT that leverages wavelet-based downsampling and orthogonal linear projection for feature decorrelation and subband-aware scaling for

Figures (16)

• Figure 1: Illustration on the evolution of energy compaction and R-D loss during training of TCM-small with and without the proposed auxiliary transform (AuxT). The energy is computed by the $L_2$ norm. (a) Energy compaction of the top 10% of channels with highest energy. The numbers on the curve is the corresponding energy ratio relative to the total energy of all channels. (b) Convergence curves of the R-D loss. Both TCM-small liu2023learned and TCM-small+AuxT are trained with $\lambda$ set to 0.0483, and the energy and test R-D loss are averaged over the Kodak test set.
• Figure 2: Characteristics of energy distributions in latent representations for TCM-small liu2023learned, averaged over 24 images from the Kodak dataset. (a) Channel-wise energy distribution after convergence. We sort the channels of latent and split them into 10 groups. The energy is predominantly concentrated in the first group associated with low-frequency features, while other groups corresponding to high-frequency features carry significantly less energy. (b) Evolution of the energy distribution during the training process. The total energy of each group is plotted over the training process, with the dotted line representing the total energy of all channels. The numbers on the curve indicate the energy ratio of the group with the highest energy.
• Figure 3: Normalized histograms of pairwise channel similarities in the analysis transform for different training iterations. We rescale the $x$-axis for better visualization. Please Zoom in for a better view
• Figure 4: Maximum output intensity for each layer of the analysis transform $g_\textit{a}$ of TCM-small liu2023learned. The blue area is subsampling layer.
• Figure 5: (a) Overview of the proposed method. Without loss of generality, we adopt the nonlinear transform from TCM liu2023learned for illustration of our Auxiliary Transform (AuxT), while it can also integrate seamlessly with other LIC models. RBS and RBU denote the Residual Block with Stride and Residual Block Upsampling, respectively. Our method does not involve modifications to the entropy model, we omit the context model in this figure for simplicity. (b) Proposed wavelet-based linear shortrcut (WLS) for analysis transform, where DWT, subband-aware scaling and orthogonal linear projection (OLP) are performed sequentially. (c) iWLS for synthesis transform, which implements the inverse operation of WLS.

Theorems & Definitions (1)

• Lemma A.1