On Disentangled Training for Nonlinear Transform in Learned Image Compression

W1 : Typos and Phrasing Recommendation Thanks for your valuable suggestion. We have modified our manuscript accordingly to further enhance the quality of the manuscript.

W2 : Discussion with the contemporaneous work [R8] Thank you for providing this important contemporaneous work [R8], which is currently under review for TMLR. We have discussed [R8] in the RELATED WORK section in our revision. Our method differs from this work in both motivation and methodology. The core of this paper focuses on modelling the training dynamics of the parameters, reducing the training space dimension, and decreasing the number of active trainable parameters over time, thereby achieving lower training complexity. We notice that while this paper provides valuable insights, it does not deeply explore the underlying reasons for the slow convergence of LIC or incorporate the specific characteristics of LIC. In contrast, our method leverages these characteristics (i.e., energy compaction and uneven energy modulation). We make the first attempt to analyze the training process of LIC from the perspective of energy compaction and reveal the training challenges in the feature decorrelation and uneven energy modulation. I believe our approach, in combination with the methods in this paper, can complement each other and further accelerate convergence.

W3 : Related work for training acceleration Thank you for your valuable advice. We have added a subsection in the RELATED WORK section discussing the acceleration of training in other tasks. This will provide a broader context and demonstrate our awareness of related methodologies in the field.

[R8] Anonymous. Accelerating learned image compression with sensitivity-aware embedding and moving average. Submitted to Transactions on Machine Learning Research, 2024. URL https://openreview.net/forum?id=V1sChEsXZg. Under review

W2: Effect of entropy constraints on energy aggregation (continue)

We also have followed your suggestion to perform the progressive training strategy for inter-frame compression of DCVC [R6] (i.e., performing 200,000 reconstruction error training iterations followed by 400,000 joint training iterations), denoted as progressive training. The R-D performance results of TCM-small and ELIC are shown in the table below.

Table R9: Comparison with the progressive training of DCVC [R6] for ELIC. All the models are trained for total 600k iterations with $\lambda=0.0483$.

Table R10: Comparison with the progressive training of DCVC [R6] for TCM-small. All the models are trained for total 600k iterations with $\lambda=0.0483$.

We found that the progressive training approach results in a significant loss in rate-distortion performance and convergence speed. We speculate that this training strategy may be more suitable for inter-frame (P-frame) compression, which involves more complex modules such as motion compensation and motion vector compression. However, for intra-frame compression, joint R-D optimization performs better. In fact, DCVC also uses this R-D optimization for training intra-frame compression (see "Appendix C: Intra frame coding" of DCVC [R6]).

W3: Training time comparison with two-stage training strategy of the paper [R7]

Thank you for your suggestion. We notice in the paper [R7] that a total of 3.1M iterations were trained, where the first 2M iterations (about 65%) were used to train a small model with only the hyperprior and the remaining 1.1M iterations (about 35%) were used to train the model with a complex context model.

For comparison, we allocated 1.3M iterations for training only the hyperprior in the case of ELIC and TCM-small, with the remaining 0.7M iterations dedicated to adding the complex context model in our 2M training iterations setting. The calculated training time comparison is presented in the table below. Our method demonstrates a significant reduction in training time, achieving decreases of 47%, 48%, and 47% compared to the anchor model for ELIC, TCM-small, and TCM-large, respectively. In contrast, the two-stage training strategy [R7] achieves only 11%, 30%, and 21% training time reductions for ELIC, TCM-small, and TCM-large, respectively.

Table R11: Training time comparison with two-stage training strategy of [R7].

[R6] Li, Jiahao, Bin Li, and Yan Lu. "Deep contextual video compression." Advances in Neural Information Processing Systems 34 (2021): 18114-18125.

[R7] Zhu, Yinhao, Yang Yang, and Taco Cohen. "Transformer-based transform coding." International Conference on Learning Representations. 2022.

Highlighting our contributions

We emphasize that our key contribution lies in being the first to analyze the convergence issue of LIC from the perspective of energy compaction and introducing an auxiliary bypass transform to address this issue. It is important to note that our approach differs significantly from existing works that simply modify the loss function for progressive training [R6] or adopt multi-stage training [R7]. Instead, we delve deeply into the characteristics of LIC and propose a novel architecture specifically designed to tackle the convergence challenge and successfully achieve 2$\times$ acceleration with 1% BD-rate reduction. We believe our analysis and exploration of the convergence issue in LIC can shed light on future research direction of LIC.

W1: How to choose the optimal $\lambda_{orth}$.

Thank you for your suggestion. We selected the optimal value of $\lambda_{orth} = 0.1$ based on experimental results on TCM in our previously submitted manuscript and directly applied it to other models. We further validate the impact of different $\lambda_{orth}$ values on ELIC and mbt2018-mean. As shown in the table below, $\lambda_{orth} = 0.1$ is demonstrated to consistently obtain the best performance. This indicates that the choice of $\lambda_{orth}$ is not sensitive to the model and can be adapted to different networks.

Table R6: Effect on the $\lambda_{orth}$ for ELIC. Models are trained for 600k iterations with $\lambda=0.0483$.

Table R7: Effect on the $\lambda_{orth}$ for mbt2018mean. Models are trained for 600k iterations with $\lambda=0.0483$.

W2: Effect of entropy constraints on energy aggregation

In fact, the presence of entropy constraints does promote more significant energy aggregation (i.e., energy compaction) rather than reducing the rate of aggregation. This is because entropy is lower when energy is concentrated in fewer channels, while the entropy is maximized when the energy is evenly distributed across all channels. In this way, entropy constraints can force the model to concentrate the information in low-frequency channels.

We have conducted experiments to further verify it. Without entropy constraints, the top 10% high-energy channels only account for 26.1% of the total latent energy after 300k iterations. However, with entropy constraints, they account for 66.3% after 300k iterations, while our method with entropy constraints achieves 94.4%. This demonstrates significantly more efficient energy aggregation.

Table R8: Energy ratio of top 10% channels with highest energy at different training stages.

W1: Scaling factor

In fact, $1, 0.5, 0.5, 0$ represent the initial values of the scaling factors $s_{LL}$, $s_{LH}$, $s_{HL}$, and $s_{HH}$, respectively. These scaling factors are indeed learnable parameters, and are optimized with the entire network. The choice of initial values is motivated by the observed phenomenon of uneven energy modulation, where low-frequency components require more significant energy amplification than high-frequency components. Indeed, the initial value does not have significant impact on the overall performance. We have performed ablation studies with initial values set uniformly to 1 or randomly initialized. As shown in the table below, the overall performance is not sensitive to the initial value, and our uneven initialization is able to obtain a lowest test R-D loss.

Table 5: Effect of the initialized value of scaling factors. All models are trained for 600k iterations with $\lambda=0.0483$.

W2: Loss Function

$L_{orth}$ represents the sum of orthogonality constraint terms across all OLPs, formulated as:

$L_{orth}=\sum_{\boldsymbol{W} \in \mathcal{W}} \left\|| \boldsymbol{W}^\top \boldsymbol{W} - \boldsymbol{I} \right\||_F^2,$

where $\boldsymbol{W}$ is the weight matrix of an OLP layer and $\mathcal{W}$ is the set of weight matrices for all OLPs. The orthogonality constraint $\left\||\boldsymbol{W}^\top\boldsymbol{W}-\boldsymbol{I}\right\||_F^2$ was defined in Section 4.1 in the previously submitted manuscript to enforce the weight $\boldsymbol{W}$ to be an orthogonal matrix [R5].

This term $L_{orth}$ is added to the original R-D loss function defined in Equation (2) of the manuscript and the overall loss function for our method is :

$L_{overall} = L_{RD}+\lambda_{orth}L_{orth} = \mathcal{R} +\lambda \mathcal{D}+\lambda_{orth}L_{orth}.$

We have updated the our manuscript to clarify it.

W3: Unclear illustration of Table 1

We apologize for the unclear explanation. The relative values in Table 1 of the previously submitted manuscript compare the results with the anchor model trained for 2M iterations (the second row of each section). We have updated the revision to improve clarity.

W4: Performance gain for TCM

This is because TCM, as a relatively new state-of-the-art model, has already achieved a superior performance. As a result, further improving performance on this anchor model becomes increasingly challenging, leading to relatively limited gains. Additionally, it is important to note that our method was trained for only 1 million iterations, while the anchor model was trained for 2 million iterations. If training of our method were to continue, the gains could increase further.

[R5] Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11505–11515, 2020.

Q3: Why is a simple wavelet more effective?

The Haar wavelet is the simplest wavelet that is both orthogonal and symmetric, ensuring a linear phase property, which minimizes phase distortion and improves compression performance [R3, R4]. In contrast, the db4 wavelet lacks the symmetry of the Haar wavelet, potentially introducing phase distortion and necessitating the use of discontinuous periodic boundary extensions, which can introduce high-frequency artifacts in the reconstructed image [R5]. Although the bior2.2 wavelet is symmetric, it is not orthogonal and results in a loss of overall orthogonality when applied in our AuxT. This causes degraded performance. Moreover, the Haar wavelet is more computation-efficient, since it has a short support and requires less processing power and time compared to wavelets with longer supports, such as db4 and bior2.2.

Q4: Effect of the linearity of AuxT

The used orthogonality constraint cannot regularize the nonlinear activation layer, as it is designed to constrain the weights $\boldsymbol{W}$ of the OLP, meaning it is implemented before the non-linear layer. However, the non-linear layer disrupts the overall orthogonality and linearity, which may cause degraded performance and slower convergence speed.

Q5: Figure 1 and other typos

Thank you for your kind reminder. The energy of each channel is computed using the $L_2$ norm in Fig.1. We have updated the caption of Fig.1 and corrected the typos you mentioned.

[R2] Soman, Anand K., P. P. Vaidyanathan, and Truong Q. Nguyen. "Linear phase paraunitary filter banks: Theory, factorizations and designs." IEEE transactions on signal processing 41.12 (1993): 3480-3496.

[R3] Stanhill, David, and Yehoshua Y. Zeevi. "Two-dimensional orthogonal filter banks and wavelets with linear phase." IEEE transactions on signal processing 46.1 (1998): 183-190.

[R4] Belzer, Benjamin, J-M. Lina, and John Villasenor. "Complex, linear-phase filters for efficient image coding." IEEE Transactions on Signal Processing 43.10 (1995): 2425-2427.

W4: Comparison with the correlation loss of NIPS2023

The proposed method differs from the correlation loss [R1] in both objectives and methodology. [R1] aims to enhance the decorrelating ability of nonlinear transforms without relying on a context model. It designs a metric for evaluating the neighborhood correlations of latent representations and uses it as a penalty in the loss function. Different from [R1], we focus on accelerating the convergence of training nonlinear transforms by learning coarse decorrelated representations with an easy-to-learn auxiliary transform. It is worth mentioning that the correlation loss and our method are orthogonal and can be combined to complementary each other to further enhance the network's decorrelating ability and accelerate convergence. We have discussed this paper in the revised appendices.

We have attempted to compare the proposed method with [R1]. However, the codes of the method developed in [R1] are not yet publicly available and we suffer from unstable convergence (i.e., NaN loss in training) when reproducing it. We strictly followed the methodology and settings described in the original paper, as well as the hyperparameters provided in the appendix. In detail, we set the training iterations to 600K for comparison and applied the correlation loss [R1] to diverse models (e.g., TCM-small and mbt2018-mean) and conducted experiments on different bitrate points (i.e., $\lambda= 0.0018$ for low bitrate and $\lambda = 0.0483$ for high bitrate). We consistently observed a training collapse with the loss becoming NaN. Moreover, prior to the collapse, we did not observe any significant performance improvements compared to the anchor model. The comparison between the proposed method (before collapse) and the anchor are depicted in the following tables.

Table R2: Comparison with correlation loss [R1] using $\lambda=0.0483$. Training with the correlation loss collapses at 410k iterations.

Table R3: Comparison with correlation loss [R1] using $\lambda=0.0018$. Training with the correlation loss collapses at 100k iterations.

[R1] Ali, Muhammad Salman, et al. "Towards efficient image compression without autoregressive models." Advances in Neural Information Processing Systems, 36 (2024).

Q1: Input of WLS

Thank you for your valuable advice. We have performed experiments where the inputs of WLS/i-WLS were replaced with the intermediate features of $g_a$ or $g_s$, denoted as AuxT*. This modification means that WLSs no longer form a standalone bypass transform by sequentially stacking but instead act as shortcuts at each intermediate stage of the nonlinear transform separately. The results are shown in the table below. We find that this modification causes degraded rate-distortion performance. The results further demonstrate the necessity of introducing an additional branch for accelerated convergence of LIC.

Table R4: Effect of the input of WLS. All the models are trained for 600k iterations with $\lambda=0.0483$.

Q2: Are scaling factors trainable?

Yes, the scaling factors are learnable. Each WLS/iWLS employs a set of scaling factors and they are jointly optimized with other parts of the model. We have updated Section 4.1 in the revised manuscript to clarify it.

W1: On the context model

Thank you for your advice. It is noted that we did evaluate the proposed method on LIC models with diverse context models, as shown in Table 1 of the previously submitted manuscript. The context models include the space-channel context model in ELIC, the channel-wise context model in STF, and the transformer-based channel-wise context model in TCM. We found that the proposed method benefits all these methods and achieves BD-rate reduction from 0.7% to 2.5% when incorporated into these models. We have revised the manuscript to clarify the compatibility of the proposed method with diverse context models.

W2: Reason for introducing an additional branch

The additional branch helps accelerate convergence in training LIC models with nonlinear transforms. The main branch (i.e., nonlinear transform) and the additional branch (i.e., the proposed AuxT) can work in parallel to extract fine-grained and coarse information from the images. However, this cannot be achieved when AuxT is employed as part of a sequential model. Additionally, we have conducted a study on merging the proposed WLS/iWLS into the nonlinear transform by substituting the sub-sampling and up-sampling layers of TCM with the proposed WLS/iWLS (i.e., DWT/IDWT + Sub-band-aware scaling + OLP + orthogonality constraint). The model variants used in the experiments are as follows:

• (a) TCM-small (anchor model)
• (b) Substituting all the sub-sampling and up-sampling layers of TCM-small with DWT/IDWT (single branch)
• (c) Substituting all the sub-sampling and up-sampling layers of TCM-small with the proposed WLS/iWLS (single branch)
• (d) TCM-small + AuxT (two parallel branches of our method)

The R-D results are shown in the table below and the convergence curves are provided in Fig.9 of the revised manuscript.

Table R1: Effect on merging WLS/iWLS to $g_a$ and $g_s$. All the models are trained for 600k iterations with $\lambda=0.0483$.

Compared to introducing additional branches, merging WLS/iWLS into $g_a$ and $g_s$ causes degraded R-D performance. This may occur because these additional modules and constraints hinder the nonlinear transforms from effectively learning latent representations for fine details. This further highlights the importance of the additional parallel branch design in accelerating convergence.

W3: Substituting the sub-sampling and up-sampling layers with DWT and IDWT

Thank you for your valuable suggestion. As introduced in the response to W2, we have conducted an additional study on substituting the sub-sampling and up-sampling layers of TCM with the full wavelet-based shortcuts (i.e., DWT/IDWT + Sub-band-aware scaling + OLP + orthogonality constraint). The results indicate that incorporating sub-band-aware scaling, OLP, and orthogonality constraints leads to a performance drop compared to using only DWT/IDWT.

Thank you for your detailed response addressing my questions. I will hold the same score.