On Disentangled Training for Nonlinear Transform in Learned Image Compression
Authors: Han Li
Shaohui Li
Wenrui Dai
Maida Cao
Nuowen Kan
Chenglin Li
Junni Zou
Hongkai Xiong
Venue: ICLR 2025 Spotlight (2025)
Keywords: learned image compression， training efficiency， auxiliary transform
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

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, \emph{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 uneven energy modulation. AuxT is lightweight and plug-and-play to be integrated into diverse LIC models to address the slow convergence issue. Experimental results demonstrate that the proposed approach can accelerate training of LIC models by 2  times and simultaneously achieves an average 1\% BD-rate reduction. To our best knowledge, this is one of the first successful attempt that can significantly improve the convergence of LIC with comparable or superior rate-distortion performance.

Published as a conference paper at ICLR 2025
ON DISENTANGLED TRAINING FOR NONLINEAR TRANSFORM IN LEARNED IMAGE COMPRESSION
Han Li', Shaohui Li Wenrui Dai' Maida , Kan', Chenglin Li', Junni Zou', Hongkai Xiong' Shanghai Jiao Tong University. "Tsinghua Shenzhen International Graduate School, Tsinghua University (qingshi9974, da ee edu, 1 lishaohuiee ses.. ttttahhhuiee...tsshhaaaaa twwnggui..nieeee www.ghuaa.dd....o edu.cn
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 ofate-of-the-ar 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 (WI.Ss) for AuxT that leverages wavelet-based downsampling and orthogonal linear projection for feature decorrelation and subband-aware scaling for uneven energy modulation. AuxT is lightweight and plug-and-play to integrated into diverse LIC models to address the slow convergence issue. Experimental results demonstrate that the proposed approach can accelerate training of LIC models by 2 times and simultancously achieves average 1% BD-rate reduction. To our best knowledge, this is one of the first successful attempt that can significantly improve the convergence of LIC with comparable superior rate-distortion performance. Code will be relcased at httpt:t::/// /tEtps://githüa github......ggggggg com com./gingsh13974//Mu
INTRODUCTION
Recent advances in learned image compression (LIC) (Ballé et al. 2018 Minnen et al. 2018
Cheng et al., 2020 Liu et al. 2023 Li et al. 2024a) have been attracting increasing attention. LIC usually employs a pair of nonlinear analysis and synthesis transforms (Ballé et al. 2020 to achieve mapping between input images and their latent representations. The nonlinear transforms are jointly optimized along with uniform quantizer and entropy model to render superior rate-distortion (R-D) performance than traditional handcrafted image codecs.
Despite their promising R-D performance, existing LIC methods suffer from slow convergence training. We take TCM (Liu et al. 2023), the recent state-of-the-are IJC model, as an example. TCM requires more than 2 million training iterations that take more than 15 days on NVIDIA GeForce RTX 4090 GPU to train a model for a single R-D point, and needs 6 models at different rates guarantee satisfactory R-D performance over the whole rate region. Moreover, low-efficiency training hinders fast fine-tuning for varying image datasets and downstream tasks or adaptation to variable target bitrates. This issue limits the practical application of LIC models.
In fact, efficient training of LIC models has not been thoroughly explored, despite the properties of nonlinear transforms have been studied. Previous studies (Duan et al. 2022; He al., 2022]
Li et al. 2024c) find that the analysis transform in LIC exhibits energy compaction property (i.e., the energy of latent representation concentrates in several dominant channels), and further observe two characteristics of nonlinear transforms to realize energy compaction (i.e., feature decorrelation and uneven energy modulation). However, they ignore to consider the evolution of energy compaction during training. In this paper, further investigate this problem, and reveal that models struggle to improve the efficiency of energy compaction through the training process. As
"Corresponding authors: Wenrui Dai; Shaohui Li.

---

Published as a conference paper at ICLR 2025
TCM amal
7Th smail AvrT
1.9
1.8
Energy
Test R-0
1.7
1.6-
10
1.4
1.422
1.414
AuxT 1.3
Eorations
Iterations (M)
(a) Energy compaction vs. iterations
(b) Test R-D loss vs. iterations
Figure 1: Illustration 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 L2 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 (Liu et al.) 2023) and TCM-small+ANu À set to 0.0483, and the energy and test R-D loss are averaged over the Kodak test set.
are trained with
demonstrated in Figure la (blue line), during training, an LIC model could use only 10% channels of the latent representation to preserve 87% of its energy in the first 0.5 million iterations, but would require extra 1.5 million iterations to further increase the energy by 6%. This suggests the necessity to ease energy compaction training nonlinear transform and thus improve training efficiency.
Motivated by these facts, we propose to disentangle the learning process of energy compaction for nonlinear transform with a simple yet effective linear auxiliary transform (AuxT). AuxT works as a bypass transform in addition the nonlinear transform to achieve coarse approximation of feature decorrelation and uneven energy modulation. Such a design enables the nonlinear transform to focus on fitting the distribution of fine details and consequently yield faster convergence. Specifically, AuxT contains several wavelet-based linear shortcuts (WLSs) with each comprising a wavelet-based downsampling and orthogonal linear projection for feature decorrelation and subband-aware scaling for uneven energy modulation.
The auxiliary transform lightweight and plug-and-play, which be integerated seamlessly with existing LIC models. Figure [ presents an instance of incorporating AuxT with TCM (specifically, the small version). Compared vanilla TCM, TCM with AuxT rapidly reduces the channel-wise correlations of latent representations and evidently enhances energy compaction with identical training iterations. Consequently, the training process of LIC is significantly accelerated by introducing the lightweight AuxT.
To our best knowledge, this is one of the first successful attempt for LIC that significantly accelerates the convergence of training while achieving comparable or superior R-D performance. The contributions of this paper are summarized as below.
We interpret the training process of LIC from the perspective of energy compaction property, and reveal the low efficiency of existing nonlinear transforms achieving feature decorrelation and uneven energy modulation for energy compaction.
We design novel method to accelerate LIC training. Specifically, we propose lightweight and plug-and-play auxiliary transform (AuxT) to achieve energy compaction with coarse feature decorrelation and uneven energy modulation.
We demonstrate superior efficiency of AuxT on diverse LIC models, achieving approxi- mately 40% to 70% reduction in training time with an average BD-rate reduction of 1.3%.
2 PRELIMINARY: LEARNED IMAGE COMPRESSION
Learned image compression models typically consist of two fundamental components: nonlinear transforms and entropy model. The nonlinear transforms include a nonlinear analysis transform
2

---

Published as a conference paper at ICLR 2025
334
10"
107
Chennel group index
10
Energy
Energy
10
10a
10r
107
192
256
320
0.00
0.25
0.50
0.75
1,00
1.25
1.50
1.75
2.00
Sorted channel index
Iterations (M)
(b)
Figure 2: Characteristics energy distributions in latent representations for TCM-small (Liu al, 2023). 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 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.
ga(:;0,) 0.) with trainable parameters 0₂ at the encoder side and a nonlinear synthesis transform I. 0,) with 0, at the decoder side. (-(-) transforms the input image into latent representation = (x; 0,), which is then discretized to ÿ = Qyy using uniform quantization Q(-), whereas g. (-) obtains the reconstructed image & ga ( 0a) from the quantized latent ğ.
Gaussians with mean and scale a,, where the mean and scale estimated by the entropy model using side information and contextual information Φ The estimated probability is further employed in entropy coding that outputs bitstream with an average length calculated by
The probability of ý is usually estimated with
= Empp, [- logg Pg + Eaap, [- logg P2 (2)].
Consequently, LIC models optimized using the R-D loss, which incorporates Lagrangian multiplier À that controls the trade-off between the rate and the distortion D, expressed as:
LRD R A D.
different bitrates are realized using different values of A. Please refer to Appendix A.I for more related works on LIC.
According to Equation 2
ENERGY BASED INTERPRETATION OF LIC TRAINING
Understanding the training process of LIC crucial to addressing the slow convergence issue. In this section, we begin by re-examining the results of previous work on the mmery-ccoppaction property of nonlinear transforms (Duan al. 2022, He 2022 Li 2024c). and then show the two
reasons why the energy-compaction property is hard to leamn
Traditional transform coding usually leverages orthogonal linear transforms such as Discrete Cosine Transform (DCT), Discrete Wavelet Transform (DWT), Karhunen-Loève Transform (KLT), and multiscale geometric analysis to aggregate most energy of the image into few components. And the energy-compaction property has a significant contribution to the compression performance. In an optimized LIC model, the majority of energy in the latent representation concentrated a few dominant channels, while the other channels contain considerably less energy (as illustrated in Figure ( According to (Li et al. 2024c), the dominant channels holding high energy are low- frequency channels, whereas the others high-frequency channels.
3

---

Published as a conference paper at ICLR 2025
For clear demonstration, categorize these latent representation channels into 10 groups according to their average energy, and visualize the evolution of the energy ratio (i.e., the ratio of energy in the total energy) for each group during the training process. Figure 2 shows that the LIC model is trained to gradually achieve energy compaction. Specifically, we obtain the following two observations.
(i) High-energy channels exhibit a trend of increasing energy ratio, but the growth of both energy and energy ratio slow.
(ii) Low-energy channels are not stable in energy ratios, and might adversely affect the training stability of LJC.
These observations highlight the challenge of slow energy compaction in LIC training. To address this challenge, we further identify two concrete reasons behind it.
Feature Decorrelation. Traditional transform cod- ing inherently obtains energy-compacting representation through orthogonal linear transforms that are well designed for decorrelating frequency components in structured manner. III. learns nonlinear transforms to adapt to the source distribution of natural images and also achieves decorrelation in latent representation. In Figure 3 we estimate the average similarity between each two channels based on their guided backpropagation patterns (Springenberg et al., 2015) on the Kodak dataset throughout the whole training process. The pairwise channel similarity suggests the correlation of any two channels of latent representations is gradually reduced during the training process. However, unlike traditional transform coding, LIC requires numerous iterations in training to achieve ideal decorrelation, since orthogonality not guaranteed for the analysis transform.
Convergence
12%
Frequency
Iterations (MI
6%
4%
2%
1974taaton
0.00
1.00
Correlation
Figure Normalized histograms of pairwise channel similarities in the analysis transform for different training iterations. We rescale the x-axis for better visualization. Please Zoom for a better view
Uneven Energy Modulation. Traditional transform coding commonly applies different quantization steps on the transform coefficients of different frequencies. The low-frequency coefficients usually adopt smaller quantization steps than high-frequency coefficients to reduce average distortion. On the contrary, LIC models employ uniform quantizers with fixed quantization step of for all the latent representation. Low- frequency channels that inherently carry higher energy require additional amplitude amplification (i.e., uneven energy modulation) to achieve finer quantization. Thus, the analysis transform gradually learns to amplify low- frequency channels, and overall energy in the latent representation increases as training proceeds. Figure 4 shows that the max intensity of latent representations is significantly scaled up from value of in the source space to over 100 in the latent space. However, the scaling-up primarily occurs in the last layer of neural network realizing the analysis transform. This causes unbalanced gradients that could slow down the training process.
140
120
Max Amplitude
100
80
6
10
Layer Depth
Figure Maximum output intensity for each layer of the analysis transform 9a of TCM-small (Liu et al.) (023). The blue area is subsampling layer.
In light of these findings. propose method that disentangles the training process of the nonlinear transforms into feature decorrelation and uneven energy modulation to accelcrate training.
4 AUXILIARY TRANSFORM FOR DISENTANGLED TRAINING
We introduce auxiliary transform (AuxT) disentangle the learning process of the nonlinear transform such that the difficulty in achieving energy compaction in training could be reduced and the convergence of training could be accelerated. The proposed AuxT works in parallel with
4

---

Published as a conference paper at ICLR 2025
AuxT
L.I
Subhand Boaand-Awarr Amarre
WLS
W1.
WISS
ME.S
DWT
LII
Scalling
Concat
TCM Bltck
TCM Block
HL
JIBS 12
1R 1/
RBS 12
Conv3rc
HH
OLP
TCM
Entropy Coding
Ex
1
(b) WLS
8x
sihpel Conehel Fz
RBH 12
TCM Block
REA 12
Bleck
12
LL
TCM
Subhand-dAuu
Spint
Subhaa
HH
OLP
DATX
JALS
AuXT
(a) Overview of proposed method
(c) IWLS
Figure 5: (a) Overview of the proposed method. Without loss of generality, we adopt the nonlinear transform from TCM (Liu et al., 2023) for illustration of our Auxiliary Transform (AuxT), while 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 involve modifications to the entropy model, we omit the context model in this figure for simplicity. (b) Proposed wavelet-based linear shortreut (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.
the nonlinear transform and efficiently provides a complementary latent representation with coarse decorrelation and energy modulation. This allows the nonlinear transform to pay more attention to learning latent representations for fine details, ultimately achieving higher compression performance with significantly reduced overall training effort. AuxT is achieved with multiple wavelet-based linear shortcuts (WLSs), where each WLS consists of three components a sequence, i.e., wavelet- based downsampling, subband-awer -ware scaling, and orthogonal linear projection (OLP).
4.1 WAVELET-BASED LINEAR SHORTCUT (WLS)
Wavelet-based down-sampling is employed to achieve feature decorrelation. Besides, it achieves spatial down-sampling align with the common design of LIC models. Let P E H be the input of WLS, where H, W, and C' denote the height, width, and channel respectively. The input input P is down-sampled by the Discrete Wavelet Transform (DWT). For simplicity, we use the Haar wavelet decompose into four wavelet subbands PLL R x xc PLH E RR f x xc
PHL E R x XC and Рии E R # x W. xC, The filters to produce the four subbands are
1 fLL = 2
fLH 2 2 1
1 1
-
1 [ -| (3)
JHL
fHE
2
2
where fLL is the low-frequency filter, and fLH, FHL. and fH# are the high- figh-frequency filters.
Subband-awar scaling. Considering lower-frequency channels require greater energy amplifica- tion (i.e., uneven energy modulation) than higher-frequeny channels, we design subband-awe aware scaling method to adaptively amplify the energy of wavelet subbands. Specifically, each WLS employs four learnable scaling vectors, 8u1, 8th, 8h1, E CRC to scale the corresponding subbands PLL,, PLH, PHLL, and РиН, respectively. The scaled subbands are then concatenated along the channel dimension to form E R
W x4C
Concat Pp C e P 2 8th , Phi O "hl Pash A tha))
(4)
where @ denotes Channel-wiss multiplication. The values of Sil, 8th, Shi, Sha are initialized with 1, 0.5, 0.5, O, respectively, reflecting uneven energy modulation for different subbands.
Orthogonal Linear Projection (OLP). Channel projection is considered to ensure that the output channels align with the intermediate features of the nonlinear transform. However, an unconstrained
5

---

Published as a conference paper at ICLR 2025
linear projection layer is inferior in channel-wise decorrelation and causes excessive training time, due to the lack of orthogonality guarantees. To address this issue, we develop OLP by applying an orthogonality constraint || [ W I to a 1 x convolution layer W R4CxD OLP projects P to obtain the final output of WLS, i.e., E I2 x xD, Please refer to Appendix A.5 for more details about the orthogonality constraint.
4.2 OVERALL ARCHITECTURE WITH MULTI-STAGE SHORTCUTS
As illustrated in Figure 5 apply four stacked WLS in the analysis transform and four stacked iWLS in the synthesis transform. For the analysis transform, the output of each WLS will be added the output of each down-sampling layer of the nonlinear transform Go. The output of the last WLS (i.e., P(ima)) will be added the output of analysis transform F to form the final latent representation y that is subsequently quantized and encoded. For the synthesis transform, iWLS implements the inverse operation of WL.S in the analysis transform, where the DWT is replaced by the inverse DWT (IDWT) and the channel-wise multiplication in Equation 4 is replaced by channel-wise division to achieve corresponding energy reduction.
Progressive Energy Modulation. The orthogonality of the DWT and our OLP ensures that they act as energy-preserving transforms. In this structure, multiple subband-aware scaling operations the multi-stage architecture can progressively and efficiently scales energy, which can avoid dramatic changes in energy and amplitude, leading to a more stable training process.
Loss Function. The proposed architecture is optimized in end-to-end fashion using combination of the R-D loss LRD Equation 2 and orthogonality regularization loss Corth for OLP. The overall loss function is formulated as
Loverall LRD + North Lorth, where
Lorth
Σ W|lWT
w I
(5)
WEW
where W is the set of the weight matrix for all OLPs and Aorth the regularization parameter.
EXPERIMENTS
5.1 EXPERIMENTAL. SETUP
Training. We apply our AuxT to several mainstream LIC models to show its effectiveness. including mb201 8mcan (Minnen et al., 2018), ELIC (He et al. 2022). STF (Zou et al. 2022), and TCM (Liu al. 2023). All the models are trained on the ImageNet-Ik (Degg et al. 2009) dataset and optimized using Adam optimizer (Kingma & Ba 2015). We set the batch size to 16 for convolution- LIC models (Minnen al. 2018 He al.) 2022) and 8 for transformer based LIC models (Zou ct al. (2022); Liu et (2023) We train the models without our AuxT for 0.6M and 2M iterations respectively, and train the models with our AuxT for 0.6M and 1M iterations, respectively. The learning rate is initialized as 10-4 and is decayed by factor of 10 after 0.55M iterations for 0.6M training iterations scenario, after 0.9M iterations for IM training iterations scenario, and after 1.8M iterations for 2M training iterations scenario.
Two kinds of quality metrics, i.e., mean square error (MSE) and multi-scale structural similarity (MS-SSIM), are used to measure the distortion DD. The Lagrangian multiplier X in the R-D loss used for training MSE-optimized models are (0.0025 0.0035, 0.0067, 0.0130, 0.0250, 0.0483), and those for MS-SSIM-optinize models are (2.40, 4.58, 8.73, 16.64, 31.73, 60.50). The orthogonal regularization weight Aorth is 0. Experiments are performed on NVIDIA GeForce RTX 4090 GPU and Intel Xcon Platinum 8260 CPU.
Evaluation. We adopt three benchmark datasets, i.e., Kodak image set (Kodak) 1993) with 24 images of 768 x 512 pixels, Tecnick testset (Asuni & Giachetti] 2014) with 100 images of 1200 x 1200 pixels, and CLIC Professional Validation dataset (CLIC) 202T) with 41 images of at most 2K resolution, for evaluations. We use both PSNR and MS-SSIM to measure the distortion, and bits per pixel (BPP) to evaluate bitrates. We present the PSNR- BPP results evaluated on Kodak in Table I and provide the results of other dataset and MS- MS-SSIM metric in Appendix A.7

---

Published as a conference paper at ICLR 2025
Table 1: Performance comparison of the proposed method applied various LIC anchor models on Kodak. We report the GMACs calculated using x x52 input images for complexity comparison and GPU hours for training a single R-D point model for time comparison. BD-rate computed from PSNR-BPP curves by comparing with the anchor VTM-18.0 is adopted as the quantitative metric.
The relative values represent the comparison with the anchor model trained for 2M iterations
of Iterations
GMAC's
Training time
#Params (M)
BD-rate
Model
(M)
Enc.
Dec.
(GPU hours)
(%)
Convolution bause nonlinear transforms
mb1201 Smean
Minnen et al. 2018
0.6
43
17.6
34.3
mh201 8mean
Minnen et al.
2018
2,0
44
43
50
17.6
25.5
0bb2018mean
Aux
0.6
49
48
18 (-64%)
18.6
26.4 (+0.9)
1hx2018mean
AuxT
1.0
49
48
(-40%)
18.6
22.6 (-2.9)
ELIC (He ct al.
2022
0.6
132
130
43
33.8
40.5
ELIC et al.
2.0
132
130
143
33.8
4.5
ELIC + AuxT
0.6
137
135
46 (-68%)
34.8
4.0 (+0.5)
ELIC AuxT
1.0
137
135
76 (-47%)
34.8
-5.7 (-1.2)
Transformer-buse nonlinear transforms
STF Zou et al.
2022
0.6
143
161
35
99.8
4.6
STF Zou et al.
2022
2.0
143
161
116
99.8
-3,2
STF
Aux
0.6
144
162
36 (-68%)
100.6
-3.2
STF AuxT
1.0
144
162
(-47%)
100.6
-5.7 (-2.5)
TCM-small (Luu et al.
2023
0.6
112
148
72
45.2
-0.1
TCM-small Liu
2023
2.0
112
148
240
45.2
-5.3
TCM-small + A
0.6
114
150
75
(-68%)
45.8
4.8 (+0.5)
TCM-small T
1.0
114
150
125 (-48%)
45.8
-6.0 (-0.7)
TCM-large (Liu et al.
2023
0.6
315
449
76.6
4.1
TCM-large (Lu et al.
2023
2.0
315
449
330
76.6
-10.6
TCM-large Aux1
0.6
324
4S7
105 (-68%)
78.2
-9.7 (40.9)
TCM-large + AuxT
1.0
324
457
175 (-47%)
78.2
-11.3 (+0.7)
5.2 EXPERIMENT RESULTS
Table presents a thorough comparison of the proposed AuxT applied to diverse LIC models. Each anchor LIC model can achieve good R-D performance when trained for long iterations. However, when trained for only 0.6M iterations, these anchor models perform substantially worse due to their slow convergence. By integrating AuxT, we consistently achieve a significant acceleration on training. For instance, after 0.6M iterations, ATX+ achieves comparable BD-rate performance to STF trained for 2M iterations, reducing training time by 68%. Furthermore, training STF+AuxT for IM iterations yields a 2.5% BD-rate improvement over the anchor model trained for 2M iterations, with only half the training time. Since AuxT is lightweight, it only bring few additional parameter complexity (about IM) and has little increase on inference complexity (1~10 GMACs), In addition, Figure [bb shows the convergence curve for TCM-smal (Liu et al., (202)). it clearly demonstrates that AuxT not only accelerates convergence but also enhances stability. We further compare the inference time in Appendix A.6 and provide the detailed R-D curves in Appendix A.8
5.3 ANALYSIS ON FEATURE DECORRELATION AND ENERGY MODULATION
We analyze the energy and decorrelation behavior of our proposed auxiliary transform along with analysis transform to better understand the disentangling effect of proposed method.
Energy distribution among channels of both branches. Figure 6a illustrates the energy distribution among channels for both the nonlinear transform and AuxT when the model converges. We observe the following key points:
(1) The simple auxiliary transform can effectively capture the high-energy lower-frequency repre- sentations. In contrast, the complicated nonlinear transform focuses and contributes more on the intricate low-energy nigher-frequenee representations.
(2) By integrating AuxT, the nonlinear transform exhibits a more balanced energy distribution across channels, which facilitates stable train when the architecture is sophisticated.
7

---

Published as a conference paper at ICLR 2025
001
72
10"
120 Charnal
110
CX6 Farationa (91
banatora
(a)
(b)
(c)
Visualization of the energy distribution of proposed auxiliary transform and the nonlinear transform. (a) Energy distribution among channels different branches, with the dotted line indicating the total energy for both branches. (b) Evolution of energy during training of the two transforms, where the dotted line represents the total energy for both branches over time. (c) Evolution of energy distribution during training, channels are grouped same Figure 2 The dotted line presents the total energy for all channels. Please Zoom in for a better view.
Figure 6:
Evolution of the energy of both branches. Figure 6b shows the evolution of the total energy for both branches during training. AuxT and the nonlinear transform have similar energy at the beginning of the training process. As the training progresses, the energy of the auxiliary transform gradually increases, ultimately exhibiting significant 31.1x x energy amplification. energy of the nonlinear transform exhibits a more stable trend, with only sudden energy reduction in the early stages. After this drop. only gradually exhibits a 1.5x energy amplification. This phenomenon further demonstrates that uneven energy modulation is primarily achieved by AuxT, which alleviates the training burden on the nonlinear transform.
contrast, the
Evolution of the energy of different Channels. Figure 6c illustrates the evolution of total energy across different channels, where the grouping strategy is identical to Figure [ We observe the following phenomena. (1) AuxT significantly enhances energy compaction. The top 10% high-
energy channels account for 96% of the total energy in the latent representation after 1M training iterations, compared 93% in the anchor model trained for 2M iterations (see Figure 2b). (2) Low-energy high-frequency channels exhibit a more stable energy evolution, showing almost no change in the latter half of training, compared to the unstable trend observed in the anchor model (see Figure 2b). Since these channels are primarily realized by the nonlinear transform, this further demonstrates that AuxT contributes the stable training of the nonlinear transform.
efficiency of both branches. Figure 7 shows the pairwise channel similarity for both AuxT and the nonlinear transform. We observe that the nonlinear transform exhibits a strong decorrelation ability due to its sophisticated design, with most pairwise channel similarities being less than 0.01. Additionally, AuxT alone also demonstrates decent level of decorrelation.
Decorrelation
18%
total
not inear transform
ausilary transfomm
13%
Frequency
When integrating both branches, the overall decorrelation further enhanced, leading even greater decorrelation efficiency and compression performance.
8%
5%
2%
5.4 ABLATION STUDIES
0%
We perform ablation studies further evaluate the effectiveness of our proposed AuxT. Experiments are conducted TCM-small (Liu et al. (202)), and each model is trained for 0.6M training iterations.
0.00
0.02
1.00
Correlation
Figure Normalized histograms of pairwise channel similarities for differ- ent branches. We rescale the z-axis for better visualization.
Effect of the components of WLS. We first evaluate the effectiveness of each WLS component through the following experiments: [A1] removing the orthogonal constraint for OLP; [A2] removing the subband-aware scaling; [A3] replacing DWT with an average pooling layer; and [A4] replacing DWT with convolutional layer with a stride of 2. For [A3] and [A4], also remove the subband-aware scaling since DWT is longer used.
8

---

Published as a conference paper at ICLR 2025
Aither bartioomeeee
PSMR
(a) Effect the each component of WLS
(b) Effect the base of wavelet
2.5
31
32
33
34
35
36
37
32
32
33
34
35
36
PSNR [68]
PSAR (6B)
(c) Effect AuxT in analysis and synthesis
(d) Effect the linearity of AuxT
Figure 8: Ablation studies. We use the TCM-small the anchor model and calculate the rate saving relatives Please Zoom for a better view.
As shown in Figure 8a removing the subband-aware scaling results in a notable degradation in performance. However, replacing OLP or DWT with non-orthogonal layers leads to an even more significant performance drop, highlighting the importance of orthogonal linear transforms in WLS for stable energy modulation and enhanced feature decorrelation. Additionally, we observe that frequency decomposition provided by DWT is essential and cannot be replaced by simple low- frequency filter such as pooling.
Effeet of the bases of wavelet. [A3] and [A4] highlight the importance of DWT in WLS. We further evaluate the effect of different wavelet bases in DWT. As shown Figure 8b employing more complex wavelets like Daubechies wavelet db4 and biorthogonal wavelet bior2.2 leads to performance degradation, indicating that the simple Haar wavelet is better.
Effect of AuxT in analysis and synthesis. We evaluate the effectiveness of AuxT in the analysis transform and the synthesis transform Is- Figure [cc indicates that AuxT in Js yields trivial gain on convergence. The primary benefits come from implementing AuxT in Ja- highlighting the importance of enhancing the feature decorrelation and uneven energy modulation for Ja-
Effect of the linearity of AuxT. Unlike the nonlinear transform used in LIC, our AuxT is purely linear transform without any nonlinear operators. We demonstrate the necessity of this linearity by introducing a nonlinear layer after the OLP. As shown Figure [&d] adding either ReLU (Goort et al., (011)) or GDN (Ballé et al., 2016) significantly hinders the performance of AuxT. This is because the nonlinearity causes energy attenuation and disrupts orthogonality, which negatively impacts effective energy compaction. This demonstrates that simple linear module can be more effective for energy compaction. More ablation studies can be found Appendix A.3
CONCLUSION
In this paper, we propose AuxT, lightweight and plug-and-play linear auxiliary transform that works as a bypass transform alongside the nonlinear transform to address the slow convergence issue for existing learned image compression (LIC) models. AuxT leverages multiple wavelet- based linear shortcuts to achieve coarse approximations for feature decorrelation and uneven energy modulation, allowing the nonlinear transform to learn simplified distributions, thus accelerating the training without sacrificing performance. Empirical results demonstrate that AuxT is effectively integrated into several mainstream LIC models and speeds up their training process by times, with average BDD-rate reduction of 1%.

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work was supported in part by the National Natural Science Foundation of China under Grant 62320106003, Grant 62371288, Grant 62301299, Grant 62401357, Grant 62431017, Grant 62401366, Grant 62120106007, Grant 62125109, Grant U24A20251, and in part by the Program of Shanghai Science and Technology Innovation Project under Grant 24BC3200800.
REFERENCES
Anonymous.
Accelerating learned image compression with sensitivity-awarc embedding and moving average. Submitted to Transactions on Machine Learning Research, 2024. URL
https: //pper////eeeeeeeew net Setr//ooommiid_i_.
Under review.
Nicola Asuni and Andrea Giachetti. TESTIMAGES: a large-scale archive for testing visual devices and basic image processing algorithms. In STAG: Smart Tools and Apps for Graphics, pp. 63-70, 2014.
Johannes Ballé, Valero Laparra, and Ecro P Simoncelli. Density modeling of images using generalized normalization transformation. In The Fourth International Conference on Learning Representations, 2016.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. In The Fifth International Conference on Learning Representations, 2017.
Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. Variational image compression with a scale hyperprior. In The Sixth International Conference on Learning Representations, 2018.
Johannes Ballé, Philip A Chou, David Minnen, Saurabh Singh, Nick Johnston, Eirikur Agustsson, Sung Jin Hwang, and George Toderici. Nonlinear transform coding. IEEE Journal of Selected Topics in Signal Processing, 15(2):339-353 2020.
Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can we gain more from orthogonality regularizations in training deep networks? In Advances in Neural Information Processing Systems 31, 2018
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Energy compaction-based image compression using convolutional autoencoder. IEEE Transactions on Multimedia, 22(4):860-873, 2019a.
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Learning image and video
compression through spatial-temporal energy compaction. In 2019 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 10071-10080, 2019b.
Zhengxue Cheng. Heming Sun, Masaru Takeuchi, and Jiro Katto. Learned image compression with discretized Gaussian mixture likelihoods and attention modules. In 2020 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 7939-7948. 2020.
CLIC. Workshop and challenge on learned image compression. In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Jia Deng, Wei Dong. Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In 2009 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009.
Zhihao Duan, Ming Lu, Zhan Ma. and Fengqing Zhu. Opening the black box of learned image coders. In 2022 Picture Coding Symposium (PCS), pp. 73-77, 2022.
Haisheng Fu, Jie Liang, Zhenman Fang, Jingning Han, Feng Liang, and Guohe Zhang. Weconvene: Learned image compression with wavelet-domain convolution and entropy model. In 18th European Conference on Computer Vision (ECCV), pp. 37-53, 2024.
10

---

Published as a conference paper at ICLR 2025
Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the 14th International Conference on Artificial Intelligence and Statistics, pp- 315-323, 2011.
Priya Goyal, Piotr Dolkár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint arXiv:/706.000 (Ttv:/700.0677 2017.
Frithjof Gressmann, Zach Eaton-Ron Rosen, and Carlo Luschi. Improving neural network training in low dimensional random bases. In Advances in Neural Information Processing Systems 33, pp. 12140-121555 2020.
Zongyu Guo, Zhizheng Zhang, Runsen Feng, and Zhibo Chen. Causal contextual prediction for learned image compression. IEEE Transactions on Circuits and Systems for Video Technology, (2(4):2329-23444 2021.
Dailan He, Ziming Yang, Weikun Peng. Rui Ma, Hongwei Qin, and Yan Wang. ELIC: Efficient learned image compression with unevenly grouped space-channel contextual adaptive coding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
5718-5727, 2022.
Feng Hong.
Yueming Lyu, Jiangchao Yao, Ya Zhang, Ivor Tsang, and Yanfeng Wang. Diversified batch selection for training acceleration. In Proceedings of the 41st International Conference on Machine Learning, pp. 18648-18667 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. The Third International Conference on Learning Representations, 2015.
Eastman Kodak. Kodak lossless true color image suite (photoed pcd0992), 1993.
A Burakhan Koyuncu, Han Gao, and Eckehard Steinbach. Contextformer: A transformer with spatio-channel attention for context modeling in learned image compression. In /7th European Conference Computer Vision (ECCV), pp. 447-463, 2022.
Quoc Le, Alexandre Karpenko, Jiquan Ngiam, and Andrew Ng. Ica with reconstruction cost for efficient overcomplete feature learning. Advances neural information processing systems, 24, 2011.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In The Sixth International Conference on Learning Representations, 2018.
Han Li, Shaohui Li, Wenrui Dai, Chenglin Li, Junni Zou, and Hongkai Xiong. Frequency- aware transformer for learned image compression. In The Twelfth International Conference on Learning Representations, 2024a. URL htttps://openree/
n mete//oorrmmm.ddHGDDD 1d==KGgDDTuw?
Han Li, Shaohui Li, Shuangrui Ding, Wenrui Dai, Maida Cao, Chenglin Li, Junni Zou, and Hongkai Xiong. Image compression for machine and human vision with spatial-frequency adaptation. In 18th European Conference on Computer Vision (ECCV), pp. 382-399, 2024b.
Shaohui Li, Wenrui Dai, Yimian Fang, Ziyang Zheng, Wen Fei, Hongkai Xiong, and Wei Zhang. Revisiting learned image compression with statistical measurement of latent representations. IEEE Transactions on Circuits and Systems for Video Technology, 34(4):2891-290) 2024c.
Tao Li, Lei Tan, Zhehao Huang, Qinghua Tao, Yipeng Liu, and Xiaolin Huang. Low dimensional trajectory hypothesis is true: DNNs can be trained in tiny subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(3):3411-3420, 2022.
Jinming Liu, Heming Sun, and Jiro Katto. Improving multiple machine vision tasks in the compressed domain. In 2022 26th International Conference on Pattern Recognition (ICPR), pp- 331-337, 2022.
11

---

Published as a conference paper at ICLR 2025
Jinming Liu, Heming Sun, and Jiro Katto. Learned image compression with mixed transformer- CNN architectures. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 14388-14397, 2023.
Jinming Liu, Ruoyu Feng, Yunpeng Qi, Qiuyu Chen, Zhibo Chen, Wenjun Zeng, and Xin Jin. Rate-dissortion-ccnniito controllable versatile neural image compression. In 18th European Conference on Computer Vision (ECCV), pp. 329-348. 2024.
Ming Lu, Peiyao Guo, Huiqing Shi, Chuntong Cao, and Zhan Ma. Transformer-bawe image compression. In 2022 Data Compression Conference (DCC), pp. 469-469, 2022.
Haichuan Ma, Dong Liu, Ning Yan, Houqiang Li, and Feng Wu. End-to-end optimized versatile image compression with wavelet-like transform. IEEE Transactions on Pattern Analysis and Machine Intelligence, 14(3):224--1266) 2020.
David Minnen and Saurabh Singh. Channel-wise autoregressive entropy models for learned image compression. In 2020 IEEE International Conference on Image Processing (ICIP), pp. 3339- 3343, 2020.
David Minnen, Johannes Ballé, and George Toderici. Joint autoregressive and hierarchical priors for learned image compression. In Advances in Neural Information Processing Systems 31, pp- 10771-10780, 2018.
Dipti Mishra, Satish Kumar Singh, and Rajat Kumar Singh. Wavelet-based deep auto encoder- decoder (WDAED))based image compression. IEEE Transactions Circuits and Systems for Video Technology, 31(4): 2020.
Yurii Nesterov. A method for solving the convex programming problem with convergence rate O(1/k2). Soviet Mathematics Doklady, 27:372-376, 1983.
Yichen Qian, Xiuyu Sun, Ming Lin, Zhiyu Tan, and Rong Jin. Entroformer: A transformer- based entropy model for learned image compression. In The Tenth International Conference on Learning Representations, 2022.
Jost Tobias Springenberg, Alexey Dosovitskiy, Thomas Brox, and Martin Riedmiller. Striving for simplicity: The all convolutional net. In The Third International Conference on Leanring Representationss 2015.
Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. Proceedings of the 34th International Conference Machine Learning, pp. 3570-3578, 2017.
Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural networks. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11505-11515, 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint Xivv:/700 038888 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training BERT in 76 minutes. The Eighth International Conference Learning Representations, 2020.
Renjie Zou, Chunfeng Song, and Zhaoxiang Zhang. The devil is in the details: Window-basee attention for image compression. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 17471-17480, 2022.
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 RELATED WORK
Learned Image Compression. Learned image compression has significant advancements in recent years due to its impressive rate-distortion (R-D) performance. Unlike traditional transform coding methods that use linear transforms to decorrelate input images, LIC employs nonlinear transforms to achieve a more compact latent representation. Early works (Ballé et al.) 2017]] 2018 Minnen et al. 2018; Minnen & Singh (020) utilized convolutional neural networks (CNNs) to develop these nonlinear transforms, focusing on capturing local patterns in images. Later approaches incorporated non-local blocks (Cheng et al. 2020) to enhance the ability to capture global correlations. More recently, advancements have leveraged attention-basee mechanisms and
transformer architectures (Lu et al.) 2022 Zou et al.] 2022; Liu et al. [023 Li et al. 2024a) for
LIC, underscoring the idea that increasing complexity and capabilities can lead to more expressive nonlinear transforms.
Another key aspect of LIC research involves designing more effective entropy models. Earlier works (Ballé et al. 2017; 2018) introduced factorized and hyperprior models to boost compression efficiency. More recent efforts (Minnen et al. 2018 Minnen & Singh, 2020 Koyuncu et al. [022] Qian et al. 2022) have incorporated spatial or channel-wise autoregression into entropy models, further enhancing their performance.
Recent works (Liu et al., 2022 2024 Li et al. 2024b) in LIC have focused on developing image codec tailored for machine perception rather than solely for human perception, aiming enhance the performance of downstream machine vision tasks while reducing storage and transmission costs.
Despite these advancements, many existing LIC methods suffer from long training times, requiring millions of iterations to converge effectively. This prolonged training duration poses significant challenges, hindering the real-world applications.
Energy Compaction Property in LIC. Energy compaction is a fundamental characteristic in traditional transform coding, indicating that most of the signal's energy is concentrated into a few components. In the field of LIC, several studies have explored this concept. For instance, penalty
Cheng et al. (2019b) and[Ch et al. (2019a) introduced a spatial energy compaction- in the loss function, encouraging energy concentration within few channels. He et al. (2022) observed that the analysis transform in LIC inherently exhibits an information compaction property, where limited number of channels carry significantly more average energy. Additionally, Li cct al. (2024c) further demonstrated that higher-energy channels are typically characterized by low- frequency components, while lower-energy channels tend to carry high-frequency information. In our work, we find that the energy compaction process is closely linked to the slow training issue of LIC and propose an auxiliary transform address this issue.
Decorrelation for LIC. Existing LIC methods often struggle with achieving better decorrelation ability to further remove the redundancies in the latent representation. The main solution is to leverage more complex neural network modules to enhance the overall representation ability of
nonlinear transforms (Ballé et al. 2017 2018 Minnen et al. 2018, Minnen & Singh 2020 Lu et al.) 2022] Zou et al. 2022] Liu et al.
2023 Li et al. (024a). In addition, [allé et al. (2016) proposed the Generalized Divisive Normalization (GDN) layer helps to decorrelate information
across channels. [uo et al. (2021) further explored the cross-channe relationships of the latents and achieve better channel-wise decorrelation by the proposed causal context prediction module.
Training Acceleration for Neural Networks. The acceleration of neural network training has been a heated topic in deep learning research.
One key approach to speeding up training is the development of advanced optimizers. In addition to widely-used adaptive learning rate methods like Adam (Kingma Ba 2015) and accelerated schemes such as Nesterov momentum Nesterov 1983), several recent approaches have emerged to further enhance training efficiency. For example, Goyal et al. (017) introduced a scaling rule to adjust the learning rate, enabling faster training when using large miniebatches. Furthermore, You ct al. (2017) and You et al. (2020) proposed layer wiiee-wise&i wise adaptive learning rates, which allow for the scaling of batch size, ultimately reducing training time.
13

---

Published as a conference paper at ICLR 2025
Other research has focused training neural networks in lower-dimensional subspaces and employing batch selection to improve efficiency. For example, ei al. (2018) proposed the idea of training in a reduced random subspace measure the intrinsic dimension of the loss objective. Later work by Gressmann et al. (020) enhanced this approach by considering the layer-wise structure and re-drawing the random bases at each step, leading improved training performance. [Li et al. (2022) proposed a low-dimensional trajectory hypothesis, which extracts subspaces from historical training dynamics, significantly improving the dimensionality efficiency of neural network training. In addition, Hong et al. (2024) introduced Diversified Batch Selection (DivBS), veference-modee free method that efficienlly selects diverse and representative samples for machine learning training.
In the filed of learned image compression, concurrent work (Anonymous) 2024) also focuses on training acceleration. This work is inspired by prior research on subspace training (Li et al. 2018 Gressmann et al., 2020; Li et al. 2022) and focuses on modeling the training dynamics of the parameters of LIC model. reduces the training space dimension and decreases the number of active trainable parameters over time, thereby achieving lower training complexity for LIC models. However, Anonymous (2024) does not thoroughly investigate the underlying reasons of slow convergence in LIC models, nor does it incorporate the specific characteristics of LIC. In contrast, our method leverages these unique characteristics (i.e., energy compaction and uneven energy modulation) and, for the first time, analyzes the training process of LIC from the perspective of energy compaction. We believe that our approach, in combination with the methods presented in Anonymous (2024), can complement each other and further accelerate convergence.
A.2 DISCUSSION ON THE DWT IN LEARNED IMAGE COMPRESSION.
In end-to-end learned image compression, many works have introduced DWT into the nonlinear transform to enhance its decorrelation ability. Mishra et al. (2020) use DWT to preprocess the input image and separately encode its wavelet subbands. Ma et al. (2020) design wavelet-like nonlinear transform using the lifting scheme. Fu et al. (2024) replace the down-sampling and up-sampling layers of LIC with DWT and IDWT, respectively, to better remove frequency-domain correlation. We highlight that our method differs from these methods as we do not modify the nonlinear transform itself, but instead introduce bypass auxiliary transform issue. To further demonstrate the superiority of our method, we conduct experiments to test whether adopting DWT directly in the nonlinear transform can accelerate convergence. Specifically, we replace the sub-sampling and up-sampling layers of TCM with DWT and IDWT, respectively and append Residual Block before and after the DWT and IDWT layers to maintain the parameter count.
address the convergence
2.0
TCM-small (45
TCM small replace all the sub./up. with DWT/IDWT (51.5M) TCM-small replace all the sub./up. with WLS/WLS (52.9M)
1.9
AuxT (45.8M)
Test R-D Loss
1.8
1.7
1.5
1.4
0.0
0.2
0.4
0.6
Iterations (M)
Figure 9: Convergence curves of the test R-D loss. X is 0.0483.
Figure 9 shows the convergence curves different methods. We observe that replacing all the sub-sampling and up-sampling layers with wavelet yields slight performance gain compared to
14

---

Published as a conference paper at ICLR 2025
2.0
TCM-small Our AuxT
1.9
TCM-small Energy Compartion Penalty
Test R-D Loss
1.8
1.7
1.6
1.5
1.4
0.0
0.2
0.4
0.6
Iterations (M)
Figure 10: Convergence curves of the test R-D loss. À is set 0.0483.
Table 2: Effect of the orthogonality regularization weight (Aorth) on the R-D performance. The TCM-small+AuxT model trained for 0.6M iterations by À set as 0.0483. and the R-D losses are evaluated on the Kodak dataset.
Aorth
Rate (bpp)
Distortion (PSNR)
Test R-D Loss
0.8958
37.726
1,447
0.01
0.8890
37.816
1,429
0.1
0.8865
37.834
1.424
0.8950
37.742
1.445
anchor model, but it also results in higher parameter complexity (51.5M vs. the original 45.2M). In addition, replacing all the sub-sampling and up-sampling layers with proposed WLS and iWLS leads to performance drop. This may be because these additional modules (subband-aware scaling, OLP) and orthognoal constraints hinder the nonlinear transforms from effectively learning latent representations for fine details. This further highlights the importance of the additional parallel branch design accelerating convergence.
A.3 ADDITIONAL ABLATION STUDIES
Compared with Energy Compaction Penalty. Cheng et (2019b) and Cheng (2019a) introduce a spatial energy compaction-based penalty to encourage the nonlinear transform of Learned Image Compression (LIC) exhibit strong energy concentration properties. We also apply this penalty in the TCM-small (Liu et al., 2023), with the convergence curve shown in Figure IO] Our observations indicate that this spatial energy compaction- compaction, with the top 10% of energy channels accounting for 99% of the total energy, compared 96% for TCM-small + AuxT and 90% for TCM-small. However, this strong penalty leads to degradation in rate-distortion (R-D) performance and the training progress is still unstable. In contrast, our AuxT enhances energy compaction in structured manner, which benefits both convergence and overall R-D performance.
penalty does enhance energy
Effect of Aorth- We further evaluate the effect of the weight of orthogonality regularization loss. We train the Aul- with different value of Aorth for 0.6M iterations and the R-D losses evaluated on Kodak dataset are listed in the Table The results indicate the the best choice of Aorth is 0.1.
15

---

Published as a conference paper at ICLR 2025
140
205
120
0.8
100
Max Amplitude
Energy
0.4
40
20
0.0
Layer Index
Layer Index
(a) TCM-small
(b) TCM-small
1.0
120
0.8
100
Max Amplituide
80
0.6
nadincer ouorie berafum
traniim
0.4
0.0
Layer
Layer
(c) TCM-small+ AuxT
(d) TCM-small AuxT
Figure 11: (a) Maximum output intensity for each layer the analysis transform Ja of TCM-small (b)Average energy for each layer of the analysis transform Ia of TCM-small (c) Maximum output intensity for each layer of the analysis transform 9a and auxiliary transform of TCM-small+Auu (d) Average energy for each layer of the analysis transform Ia and auxiliary transform of TCM- Auu
A.4 ENERGY MODULATION BEHAVIOR
Figure [1 shows the effect of our subband-aware scaling which can progressively achieve energy modulation, thus avoiding dramatic changes in energy and amplitude, leading to a more stable training prcoess.
A.5 DETAILS FOR ORTHOGONAL CONSTRAINT.
The orthogonality concept has been widely explored in the filed of deep learning (Le et al.) 2011 Vorontsov al. 2017 Bansal et al. 2018 Wang al. 2020), as implies energy preservation and encourages filter diversity. In our work. we use the orthogonal kernel regularization to force our orthogonal linear projection W E 1CDD as orthogonal layer. Specifically. if 4C > D. W is an undercomplete matrix. and the identify matrix. Conversely, 4C' < D. W becomes an overcomplete matrix, and column orthogonality is encouraged using the penalty WWT I' In practical implementation, we use the row orthogonal regularizer for both cases, since the row orthogonality and column orthogonality equivalent in the mean squared error (MSE).
orthogonal regularizer ||WT I| where I R DXD
Lemma A.1 The row orthogonality and column orthogonality are equivalent in the mean squared error (MSE), i.e., ||W 11 WWT r' + U, where U is constant.
16

---

Published as a conference paper at ICLR 2025
Table 3: Comparison on the coding complexity of the proposed method applied to various LIC anchor models. The entropy coding time is excluded.
Model
Enc. Inference Time (ms)
Dec. Inference Time (ms)
Convolution-baul based nonlinear transforms
mb1201 8mean
35.3
11.8
mbt2018mean AuxT
35.4
13.2
ELIC
74.2
53.5
ELIC + AuxT
76.7
55.6
Transformer-based nonlinear transforms
STF
104.4
110.3
STF + AuxT
106.0
111.6
TCM-small
185.2
183.8
TCM-smal AuxT
188.6
186.4
TCM-large
216.4
215.6
TCM- AuxT
222.1
219.6
The following proof is provided in the supplementary material Le al. (001). We would like to present it here for the reader's convenience.
Proof. For W MxN arbitrary matrix, we denote |WTW IN Ia 1/: as La. L, -IWW W IN -tr [W W IN) ( W - II)]
L, and WWT
"r(W"WW' W) 2tr(W"") + tr(xN)
-tr(WW7 WWT) 2tr(WWTT) + tr( +N M
(6)
-tr -tr[WW WWT 2WW +IM] +N-
-tr [WWW (WW -I]] +N-
-IW M + +N- M
=Le +U
where U N M.
A.6 COMPARISON ON THE CODING COMPLEXITY
We compare the coding complexity of the proposed AuxT the proposed method applied to various LIC anchor models. The coding complexity measured by inference latency during encoding and decoding process, where the entropy coding time is excluded. The experiments show that the increase on the additional latency time caused by our AuxT can be ignored.
A.7 ADDITIONAL RATE RESULTS
Table 4 provide the BD-rate results for MS-SSIM metric evaluated Kodak dataset (Kodak 1993) and Table provide the BD-rate results for PSNR metric evaluated on the CLIC (CLIC 2021 and Tecnick (Asuni & Giachetti] 2014) dataset.
A.8 ADDITIONAL R-D CURVES.
As shown in Figure 12 Figure B Figure 14 Figure 15 and Figure 16 present the detailed
rate-distortion (R-D) curves obtained by integrating our proposed auxiliary transform (AuxT) with various LIC anchor models different training iterations.
17

---

Published as a conference paper at ICLR 2025
Table Comparison on the performance of the proposed method applied TCM (Liu al.
2023
The BD-rate is computed from MS-SSIM-BPP curves evaluated on the Kodak (Kodak as the quantitative metric with VTM-18.0 as the anchor.
1993 dataset
# of Iterations
BD-rate (%) with MS-SSIM
Model
(M)
Kodak (Kodak 1993)
Convolution- nonlinear transforms
TCM-small
2.0
-44.6
TCM-small + AuxT
1.0
-45.7 (-1.1)
TCM-large
2.0
-48.3
TCM-large + AuxT
1.0
-49.0 (-0.7)
Table 5: Comparison on the performance of the proposed method applied various LIC anchor models. The BD-rate is computed from PSNR-BPP curves evaluated on the CLIC (CLIC) 2021) and Tecnick (Asuni & Giachetti 2014 dataset as the quantitative metric with VTM-18.0 as the anchor.
# of Iterations
BD-rate (%) with PSNR
Model
(M)
Tecnick (Asuni Giachetti
2014
CLIC
CLIC 2021
Convolution-uused nonlinear transforms
mbt2018ment
2.0
21.9
26.4
mbt201 AuxT
1.0
20.6 (-1.3)
24.2(-2.2)
ELIC
2.0
-7.6
-2.0
ELIC AuxT
1.0
-8.2 (-0.6)
-3.2 (-1.2)
Transformer-bass nonlinear transforms
STF
2.0
-5.6
-1.6
STF AuxT
1.0
-6.6 (-1.0)
-3.7 (-2.1)
TCM-small
2.0
-5,9
-3.2
TCM-small + AuxT
1.0
-6.6 (-0.7)
-4.0 (-0.8)
TCM-large
2.0
-11.4
-8.0
TCM-large AuxT
1.0
-11.9 (-0.5)
-8.4 (-0.4)
18

---

Published as a conference paper at ICLR 2025
mbt2018mean (2M iterations)
mbt2018mean (0.6M iterations)
mbt2018meen-Auxx (1M steps)
Bree2008mman+Auxx (0.6M iterations)
36
34
PSNR
32
30
0.2
0.4
1.0
Bits Per Pixel
Figure 12: R-D curve for mnbt2018ment
ELIC (2M iterations)
ELIC (0.6M iterations)
ELIC+AuXTI AuxT(1M iterations)
ELIC+ AuxT (0.6M iterations)
36
34
PSNR
32
30
0.2
0.4
0.6
1.0
Bits Per Pixel
Figure 13: R-D curve for ELIC
19

---

Published as a conference paper at ICLR 2025
STF (2M iterations)
STF (0.6M iterations)
STF AuxT (1M Iterationn)
STF AuxT (0.6M iterations)
36
34
PSNR
32
30
0.2
0.4
1.0
Bits Per Pixel
Figure 14: R-D curve for STF
TCM-small (2M iterations)
TCM-small (0.6M iterations)
TCM-small +AuxT (1M iterations)
RCM-small+Auxx (0.6M Rerations]r
36
PSNR
34
32
30
0.2
0.4
0.6
0.8
Bits Per Pixel
Figure 15: R-D curve for TCM-small
20

---

Published as a conference paper at ICLR 2025
38
TCM-large (2M iterations)
TCM-arge (0.6M iterations)
TCM-lrge+AAx" (1M iterations)
TCM-largeAAuTT (0.6M iiteration iterationsir
PSNR
34
32
30
0.2
0.4
0.6
0.8
Bits Per Pixel
Figure 16: R-D curve for TCM-large
21


Reviews:
[{"review_id": "paTywocnQb", "reviewer": "Reviewer_DW5J", "summary": "This paper introduces a linear auxiliary transform (AuxT) to address the slow convergence issue in learned image compression (LIC). Specifically, this paper analyzes and designs the method based on feature decorrelation and uneven energy modulation. Experimental results demonstrate that, compared to traditional LIC methods, AuxT accelerates model training by 2 to 3 times and achieves an average 1% BD-rate reduction, improving convergence and rate-distortion performance.", "questions": "1. The input to the WLAS in the proposed method is either the coded image x or the output from the previous layer of WLAS. I t would be interesting to compare this with using the intermediate features from the analysis transform at the corresponding resolution as input for the 2nd, 3rd, and 4th layers of WLAS. Similarly, for iWLAS, the input could be the intermediate features from the synthesis transform at the corresponding resolution.\n2. Are $s_{ll}, s_{lh}, s_{hl}, s_{hh}$ trainable? Are they shared across all WLAS and iWLAS?\n3. Is there an explanation for why employing more complex wavelets, such as the Daubechies wavelet db4 and biorthogonal wavelet bior2.2, performs worse than the simpler Haar wavelet?\n4. In the ablation study “Effect of the linearity of AuxT,” it is mentioned that the 1x1 convolution layer W in OLP is placed before the added non-linear (ReLU or GDN) layer in comparison. However, it is unclear whether the orthogonality constraint is applied before or after the non-linear layer.\n5. In Fig.1, it would be better to explain how energy is measured.\n6. Cross-ref of figures mentioned in Section 5.4 should be wrong.\n7. In Fig. 8, [A4] DWT -> Conv should be [A3] DWT -> Conv.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-written.\n2. This paper introduces a valuable new research direction. Accelerating convergence in training learned image codecs is indeed an underexplored area and is both novel and important for this field. \n3. The idea of analyzing the training process from the perspective of energy compaction is interesting.\n4. The experiment results are extensive and complete, showing the effectiveness of the proposed method.", "weaknesses": "1. This paper appears to focus solely on cases where learned image codecs do not incorporate a context model. However, including a context model could potentially enhance coding performance by leveraging correlations between latents. However, this paper primarily aims to decorrelate. It remains unclear whether the proposed method is applicable in scenarios with a context model.\n2. The motivation of adding additional branch is unclear. Is it possible to merge the operation of the proposed method (DWT / IDWT + Sub-band-aware scaling + OLP + orthogonality constraint) to $g_a$ and $g_s$?\n3. In Section A.2, this paper compares the proposed method with two alternatives that substitutes the sub-sampling and up-sampling layers with DWT and IDWT to highlight the effectiveness of introducing a bypass auxiliary transform. However, the comparison is not entirely fair, as the proposed method also includes OLP layers and orthogonality constraints.\n4. The appendix of this paper compares the proposed method with an energy compaction penalty. However, the paper “Towards Efficient Image Compression Without Autoregressive Models (NeurIPS'23)” also addresses latent decorrelation for learned image compression, but this paper does not discuss or compare it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "skTEdwuGmy", "reviewer": "Reviewer_xyeS", "summary": "This paper identifies the slow convergence in training Learned Image Compression (LIC) models due to energy compaction issues and proposes a Linear Auxiliary Transform (AuxT) to address this. AuxT uses Wavelet-based Linear Shortcuts (WLSs) to improve feature decorrelation and energy modulation. Integrating AuxT into LIC models can accelerate training by 2-3 times and achieve a 1% reduction in BD-rate, improving both efficiency and performance.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) This paper explores the problem of slow convergence in LIC. Through a large number of experiments, the author designed a universal plug-in AuxT to solve this problem to a certain extent. I think the paper is interesting and novel.\n\n(2) AuxT shows very strong generalization, achieving a significant reduction in training time and a small increase in performance in several classic LIC architectures (mbt2018,ELIC,TCM, etc.).", "weaknesses": "(1) In Subband-aware scaling of wavelet-based linear shortrcut (WLS), the authors appear to adopt simple linear scaling (reduction coefficients 1,0.5,0.5,0) for the four subbands LL,LH,HL,HH to modulate. Can you give further reasons for the use of these reduction coefficients, that is, it is proved through experiments that the reduction coefficients set in this way are optimal?\n\n(2) The derivation of Loss Function seems to be a little short, I am not sure what $L_{orth}$ includes. Can you give the formula? Are they added on top of the original Rate-Distortion loss?\n\n(3) In Table 1, what does Trainning time (-64%), (-40%) mean and to whom is it compared?\n\n(4) TCM is a relatively new model in the field, and it seems that AuxT has relatively little performance improvement when applied on TCM. Can you further analyze the reason?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JPW0jiFmkY", "reviewer": "Reviewer_5Rti", "summary": "The aim of this paper is to accelerate the convergence of learned image compression models. The authors observe the phenomena of energy compaction and decorrelation in learned image compression training. A portion of the channels in the model takes up most of the energy, and the energy compaction becomes more pronounced as the number of iteration steps increases. For a learned model, it usually requires many iterations to achieve energy compaction. This paper proposes the use of a wavelet transform based bypass transform module, to achieve coarse approximation of feature decorrelation and uneven energy modulation, which enhances the convergence speed of the model.", "questions": "Please refer to weakness.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The experimental results demonstrate the effectiveness of the bypassed wavelet transform module, and the model with the wavelet transform module requires only 1M iterations to achieve comparable rate-distortion performance compared to the original model with 2M iterations.\n\n2) The authors demonstrate the validity of their approach on most of the existing models (e.g. ELIC, Minnen’18, STF, TCM).", "weaknesses": "1) The introduction of AuxT will introduce additional hyperparameters ($\\lambda_{orth}$), will the optimal $\\lambda_{orth}$ be different for different models? According to Table 2 in the Supplementary Material, different $\\lambda_{orth}$ will have a large impact on the performance, how to choose the optimal $\\lambda_{orth}$?\n\n2) In the training of image compression models, is the difficulty in aggregating energy due to the presence of entropy constraints? In the presence of entropy constraints, in order to avoid large bit rates, the model would automatically reduce the rate of information aggregation to match the entropy model's ability to estimate redundancy? In fact, in the training of the DCVC series [1], only reconstruction loss training would be performed first, and then entropy constraints would be introduced later, and I intuitively think that this approach would accelerate the training of the model. Please compare with this approach (e.g. performing 200,000 reconstruction error training followed by 400000 joint training). \n\n3) Also, there are currently methods that will train a small model (e.g. only hyperprior) [2] and then reload the trained transform module and then add a more complex entropy model, in terms of total time, how does this compare to the training time of the method that introduces AuxT (e.g. TCM & ELIC)?\n\nReferences\n\n[1] $\\mathcal{L}\\_{reconstruction}$ and $\\mathcal{L}\\_{contextual coding}$ for inter compression model in Table 4 in https://arxiv.org/pdf/2109.15047\n\n[2] Sec B.1 TRAINING in https://openreview.net/pdf?id=IDwN6xjHnK8", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ymXdWbRENv", "reviewer": "Reviewer_faRP", "summary": "Overall, this paper presents a linear auxiliary transform (AuxT) to disentangle energy compaction in training nonlinear transforms. By analyzing the characteristics of the LIC framework, the authors find the factors that caused the slow convergence of the LIC. Their methods accelerate the training of LIC models by around 2 times with compatible performance.", "questions": "Please refer to the weaknesses section. I am open to increasing my rating if the authors further enhance the quality of their manuscript.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The authors analyze the training process of LIC from the perspective of energy compaction and highlight the inefficiency of existing nonlinear transforms in achieving effective feature decorrelation and uneven energy modulation necessary for optimal energy compaction.\n\n2. The authors propose an auxiliary transform (AuxT) specifically designed to enhance energy compaction by enabling coarse feature decorrelation and uneven energy modulation.\n\n3. Experimental results demonstrate that their method can reduce training time by approximately 45%.\n\nOverall, this is an interesting and insightful paper. The authors propose a straightforward method that results in significant improvements.", "weaknesses": "1. Typos: For example, in the sentence: \"and further observe two characteristics of nonlinear transforms to realize energy (i.e., feature decorrelation and uneven energy modulation),\" there seems to be a missing word. It would be clearer to say \"realize energy compaction\" to fully capture the intended meaning.\n\n2. Phrasing Recommendation: (1). Instead of using the phrase \"this is the first successful attempt\", which may not be entirely accurate, consider revising the sentence to: \n\n\"To the best of our knowledge, this is one of the first successful attempts in LIC that significantly accelerates the convergence of training while achieving comparable or superior R-D performance.\"\n\nAdditionally, it would be beneficial to include the earlier contemporaneous work that also explored training acceleration for LIC in the RELATED WORK section. You could reference:\n\"Accelerating Learned Image Compression with Sensitivity-aware Embedding and Moving Average\" (https://openreview.net/forum?id=V1sChEsXZg).\n\n(2). \"Experimental results demonstrate that the proposed approach can accelerate the training of LIC models by 2 to 3 times while simultaneously achieving an average 1% BD-rate reduction.\"\n\nHowever, based on Table 1, when the training time is reduced to approximately one-third of the original, there is no observed BD-rate reduction. I suggest revising the statement to: \"can accelerate the training of LIC models by 2 times while simultaneously achieving an average 1% BD-rate reduction.\"\n\n\n3. It would be important to add a subsection in RELATED WORK discussing the acceleration of training in other tasks. This could provide a broader context and demonstrate awareness of related methodologies in the field. You could discuss approaches from other domains that have addressed similar challenges in accelerating the training process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)