Causal Context Adjustment Loss for Learned Image Compression

Authors: minghao  han
shiyin  jiang
shengxi  li
xin  deng
mai  xu
ce  zhu
shuhang  gu
hyper  encoder
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 10:59:28

Abstract: In recent years, learned image compression (LIC) technologies have surpassed conventional methods notably in terms of rate-distortion (RD) performance. Most present learned techniques are VAE-based with an autoregressive entropy model, which obviously promotes the RD performance by utilizing the decoded causal context. However, extant methods are highly dependent on the fixed hand-crafted causal context. The question of how to guide the auto-encoder to generate a more effective causal context benefit for the autoregressive entropy models is worth exploring. In this paper, we make the first attempt in investigating the way to explicitly adjust the causal context with our proposed Causal Context Adjustment loss (CCA-loss). By imposing the CCA-loss, we enable the neural network to spontaneously adjust important information into the early stage of the autoregressive entropy model. Furthermore, as transformer technology develops remarkably, variants of which have been adopted by many state-of-the-art (SOTA) LIC techniques. The existing computing devices have not adapted the calculation of the attention mechanism well, which leads to a burden on computation quantity and inference latency. To overcome it, we establish a convolutional neural network (CNN) image compression model and adopt the unevenly channel-wise grouped strategy for high efficiency. Ultimately, the proposed CNN-based LIC network trained with our Causal Context Adjustment loss attains a great trade-off between inference latency and rate-distortion performance.

Main paper: 
 introduction: 
 The burgeoning quality of high-resolution photos is driving an increasing demand for advanced image storage and transmission technologies. Consequently, lossy image compression techniques have been growing extraordinarily fast in recent years. In parallel to conventional coding technologies such as JPEG [42] , BPG [6] , WebP [15] , VVC [40] , learned image compression (LIC) methods [3, 4, 10, 11, 17, 18, 20, 30, 35, 36, 47] emerge, achieving high peak signal-to-noise ratio (PSNR) and multiscale structural similarity (MS-SSIM) [44] while operating fairly fast. Their superior compression results over those of VVC demonstrate an enormous possibility that LIC technology would appear on par with the traditional ones in the near future.Learned lossy image compression methods are built upon a variational auto-encoder (VAE) framework proposed by Ballé et al. [4] . The VAE based LIC framework mainly comprises an auto-encoder and an entropy model. The auto-encoder conducts nonlinear transforms between the image space and the latent representation space; while, the entropy model minimizes the code length by estimating the probability distribution of latent representations. In comparison to the auto-encoder, which could borrow ideas from recent advances in network architecture design, the entropy model is a unique important component to LIC and has a vital influence on the final compression results.In the literature on LIC, the entropy model generally refers to a parameterized distribution model. In their seminal work [3] , Ballé et al. established the end-to-end rate-distortion minimization framework and showed that the smallest average code length of latent representation is given by the Shannon cross entropy between the actual marginal distribution and a learned entropy model. Since then, numerous entropy models have been investigated. One category of studies investigates advanced network architectures for accurately predicting the distribution of latent representations. Meanwhile, another line of research study a more fundamental perspective of the entropy model, i.e. conditional distribution modeling, to pursue a better rate-distortion trade-off. Taking side information (also termed as hyperprior) and decoded latent (also termed causal context) as conditions has become a prevailing strategy in state-of-the-art LIC models.In this paper, we advance conditional distribution modeling in the entropy model with our proposed causal context adjustment loss (CCA-loss). Existing works generally train LIC networks with a combination of the rate loss and the distortion loss. The conditional predictability of the representation is indirectly optimized, and the performance of entropy model highly relies on the hand-crafted causal context model, e.g. channel-wise [36] , checkerboard [18] and space-channel [17] context model. Our CCA-loss makes the first attempt on explicitly imposing loss to adjust the causal context, making the latter representation more accurately predicted by the previously decoded representations. To be more specific, considering a two stage autoregresive context model with hyperprior z, denote the latent representation to be decoded in the first and second stage as y 1 and y 2 ; in addition to minimizing the cross entropy loss for reducing bitstream, we introduce an auxiliary entropy model and a tailored context causal adjustment loss, which let y 2 can be accurately estimated by y 1 and z, while, at the same time, let y 2 can not be accurately estimated by merely z. In this vein, our CCA-loss explicitly guides the encoder to adjust important information into the early stage of the autoregressive entropy model, providing the LIC framework a more rational causal context sequence for entropy coding. As the codes in the early stages are enhanced with our CCA-loss, we further study the schedule of causal context transmission, and adopt an uneven channel dimension schedule for the pursuit of a better rate-distortion trade-off. The uneven channel schedule is also beneficial for reducing computational burden in the coding and decoding process, enabling our model to achieve state-of-the-art compression performance with less running time. Our contributions are summarized as follows:• We introduce causal context adjustment loss to explicitly adjust the causal context information, forcing the network to encode important information early and therefore improving the autoregressive prediction accuracy of the entropy model.• We adopt an uneven schedule of autoregressive causal context and a convolutional autoencoder architecture, delivering an efficient compression network which is easy to be implemented on modern deep learning platforms.• We evaluate our proposed compression network on various benchmark datasets, in which our method achieves better rate-distortion trade-offs towards the existing state-of-the-art methods, with more than 20% less compression latency.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: Our abstract and introduction accurately reflect our paper's contributions and scope.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related works: 
 Recent LIC studies broadly follow the seminal work of Ballé et al. [3] , which utilizes a VAE based framework for rate-distortion optimization. Generally, VAE-based LIC models comprise an auto-encoder and an entropy model. In this section, we review respective progresses in advanced auto-encoders and entropy models. Another considerable technique in LIC is the quantization method, however, as we did not dig into the details and simply followed the quantization method of [35] , we omit the review of quantization methods in this section.The auto-encoder plays the role of extracting a latent representation apt to be compressed in LIC framework. In their pioneering work, Ballé et al. [3] first proposed to use a generalized divisive normalization (GDN) [2] to transform the input image into latent space. The later works followed the same VAE framework for LIC but exploited the convolution neural network (CNN) architecture, which is easier to implement and train. Beyond the basic CNN auto-encoder, the introduction of more complex nonlinear transforms [9, 11, 33] and various architectures [14, 28, 45, 46] promotes RD performance. Recently, inspired by the great successes Transformers have made in other vision tasks, self-attention modules have been widely utilized for extracting latent representations. Embedding Transformer variants, for example ViT [12] , swin-Transformer [31] in the auto-encoder [32, 48, 49] enhances the RD performance. Moreover, Liu et al. [30] proposed a hybrid approach, combining conventional CNN and swin-Transformer. Although these Transformer-based auto-encoder could improve the RD performance by extracting better latent representations, the inference of transformer architecture has not been well optimized by the existing hardware, resulting in slow coding and decoding speed. In this paper, we borrow ideas from recent advances in image restoration [8] and adopt a CNN-based auto-encoder architecture. Thanks to our improved entropy model as well as the powerful NAF-block [8] , our LIC model could achieve state-of-the-art compression results with much less runtime than recent Transformer-based approaches.The entropy model plays a key role in LIC for minimizing the bitstream of latent representation. In the original work [3] , the probability distributions of the latent representation are modeled using a non-parametric, fully factorized density model. In order to improve the distribution predicting accuracy, Ballé et al. [4] introduced side information as a hyperprior latent variable and ultimately established the basic VAE architecture of LIC in the past decade. Beyond the hyperprior transmitted in the VAE-based LIC framework, newly proposed extra side information transmitted from encoder to decoder promotes the compression performance as well [20, 43] . Moreover, Duan et al. [13] explored the hierarchical VAE structure with multiple hyperpriors.In addition to the improvement on the side information, the introduction of causal context autoregression greatly promotes the RD performance of LIC, efficiently utilizing the information from decoded parts without a supererogatory amount of bits per pixel (bpp). Minnen et al. [35] proposed the first autoregressive structure, using the decoded spatial context to better estimate the current probability distribution. Numerous works attempt to establish an effective causal context for assistance in distribution estimations, such as channel-wise segmentation [36] , checkerboard [18] , unevenly grouping [17] . A very recent work [34] explored different strategies to selectively transmit tokens. However, a fixed hand-crafted causal context may not work well in diverse image distributions. In this work, we impose a loss to adjust the causal context in the training phase, allowing the network to achieve a more accurate probability estimation.An efficient network structure of the entropy model remains critical for achieving high RD performance [26, 27, 29] . Apart from the basic causal context and hyperprior entered into the entropy model, more references benefit RD performance [11, 16, 17, 38] . Just as how Transformer performs in auto-encoder, the advantages of integrating the entropy model with Transformer are unearthed quickly. Previous works applied various Transformer blocks [20, 23, 25, 30, 37] to enhance features before entropy estimation. Following our modified architecture in auto-encoder, we embed the NAF-block in the entropy model to improve estimation accuracy.3 Preliminary: Learned Image Compression with Variational Auto-Encoder Variational Auto-Encoder based Image Compression Framework. Ever since the VAE architecture was established [21] , learned lossy image compression techniques maintain the primary constituent structure [3] , including an auto-encoder to extract the latent representation for compression and an entropy model to assist in entropy coding. Given a source image vector x, the auto-encoder contains a parametric analysis transform g a to obtain the latent representation y from x and a parametric synthesis transform g s for reconstruction. y is then quantized to ŷ, the discrete coding symbol for storage. The probability distributions of ŷ are modeled using a factorized density model ψ as p ŷ|ψ ( ŷ|ψEQUATION))( ŷi ). As quantization introduces error, which is tolerated in the context of lossy compression, the optimization target approximates the true posterior p ŷ|x ( ŷ|x) with a neural network q( ŷ|x) as the expectation of their Kullback-Leibler (KL) divergence over the data distribution p x :EQUATIONThe former term refers to the image reconstruction distortion (measured by PSNR or MS-SSIM), and the latter term represents the bit-rate (expected code length). A hyperparameter λ is multiplied on the latter term, so that we can control the rate-distortion trade-off to obtain various compression rates.Entropy Model with Hyperprior. However, directly modeling ŷ with the factorized density model ψ is less than satisfactory, as the estimation of which is not accurate and out of correlation with the data distributions. To capture the spatial dependence among the elements of ŷ, the side information z is introduced [4] . z is generated by a hyper analysis transform h a from y, transmitted as a hyperprior latent feature to help predict the distributions of ŷ accurately. Similarly to y, z is quantized to ẑ in the same manner. The probability distributions of ẑ are calculated using a factorized density model ψ, to encode ẑ as p ẑ|ψ ( ẑ|ψEQUATION))( ẑi ). During the entropy coding process, ẑ would be entered into a hyper synthesis transform h s to acquire the estimations {µ, σ} i in normal distribution of each element ŷi . This course can be formulated asEQUATION2 ) ( ŷi ), with {µ, σ} = h s ( ẑ). The KL divergence in the basic VAE structure (Eq. 1) can be expanded as follows:E x∼px D KL [q ∥ p ŷ, ẑ|x ] = E x∼px E ŷ, ẑ∼q -log p x| ŷ (x| ŷ) -log p ŷ| ẑ ( ŷ| ẑ) -log p ẑ|ψ ( ẑ|ψ) . (2)On the other side of VAE, the parametric synthesis transform g s recovers the reconstructed image x from the decoded ŷ. Fig. 1 reveals the general basic structure.Autoregressive Entropy Model. In addition, an advanced architecture of the entropy model is the joint autoregression [35] , which soon develops into a more efficient channel-wise autoregression [36] . In the channel-wise autoregressive structure, the latent representation ŷ is grouped in the channel dimension and decoded in order. Thus, the second term of the KL divergence in hyperprior structure (Eq. 2) is expanded as:EQUATIONBesides the prior of ẑ, the estimation of the autoregressive entropy model conditions more on the causal context, that is, the model utilizes the information from the decoded parts (causal context) without introducing additional redundancy in information transmission. Therefore, the more effective the causal context, the stronger the performance of the autoregressive entropy model. Existing methods adopt various hand-crafted causal contexts to enhance it. We expect to establish a way that enables the network to adaptively adjust the causal context. Imposing a loss to explicitly adjust the causal context is a delicate way.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [NA] Justification: To our best knowledge, our proposed methods have no limitation. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: The data used from the Open Images datasets are publicly available. Since our codes have not been sorted and filed well, to avoid bringing the burden on reading a messy code and the details are presented enough for reproducing in our paper, we will release them when the codes are sorted well.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
causal context adjustment for efficient learned image compression: 
 In this section, we introduce the details of our LIC method. We firstly introduce our causal context adjustment (CCA) loss, which is able to explicitly push the encoder to encode important (in terms of information gain) representation at an earlier stage for better predicting the remaining representations. Subsequently, we introduce the implementation details of our efficient LIC method, including our CNN-based encoder and decoder architecture, unevenly grouped autoregressive schedule, light-weight entropy model, and overall loss function.As introduced in the previous section, exploiting the causal context from the hyperprior and the autoregressive framework to establish a conditional distribution task is the key to a state-of-the-art entropy model. While introducing conditions is undoubtedly beneficial for improving the accuracy of distribution estimation, existing works intuitively set up the context models, such as checkerboard context model and slice-based context model, and there still lack in-depth study on how to constitute the causal context rationally in the literature. More concretely, the rate and distortion loss reflect the prediction error given the causal context and the reconstruction error given the decoded representations, respectively; neither of them could explicitly affect the organization of the causal context. In this section, we introduce the CCA-loss, which explicitly encourages important information of the image to be encoded into earlier causal context, so as to enhance the predictability of the autoregressive entropy model. To the best of our knowledge, our work is the first attempt that introduces loss instead of intuitively adjusting the context architecture to improve the context model.To introduce our proposed Causal Context Adjustment (CCA) loss, we first revisit the hyperprior and the autoregressive entropy model. Without loss of generality, we consider a two-stage autoregressive context model. To encode the same latent representation, the cross entropy of the hyperprior model and the hyperprior + autoregressive model can be written as follows:EQUATIONEQUATIONwhere H H.P. and H H.P.+A.R. represent the cross entropy with hyperprior and with hyperprior + autoregressive estimation, respectively; q and p denotes the real distribution and the learned entropy model. According to Shannon information theory [39] , as more information is incorporated in the estimation of ŷ2 , H H.P.+A.R. is less than or equal to H H.P. . Moreover, the gap between H H.P. and H H.P.+A.R. is related to the amount of information ŷ1 could provide for estimating ŷ2 . Therefore, by calculating the following equation:EQUATIONwe could obtain the information gain introduced by causal context ŷ1 . The above analysis inspires us to explicitly optimize Eq. 6 to enhance ŷ1 , encouraging it to encode important information that helps to better estimate ŷ2 . Concretely, in addition to the original entropy model F ( ŷ1 , ẑ) → ŷ2 , we introduce an auxiliary entropy model, which only takes ẑ as input: G( ẑ) → ŷ2 . With the introduced auxiliary entropy model, our CCA-loss can be defined as follows:EQUATIONwhere p ŷ2| ẑ ( ŷ2 | ẑ) and p ŷ2| ẑ, ŷ1 ( ŷ2 | ẑ, ŷ1 ) are the estimated distributions of auxiliary and the major entropy model, respectively. The auxiliary and major entropy models are parameterized by two networks, i.e. G( ẑ) and F ( ŷ1 , ẑ). It should be noted that the auxiliary entropy model is only introduced in the training phase for better optimizing the causal context; in the testing phase, our model still uses F ( ŷ1 , ẑ) to compress the latent representation, our CCA-loss will not introduce additional computational burden for image compression.An illustration of the proposed CCA-loss can be found in Fig. 1 . Besides the above analysis in Eq. 4 to Eq. 7, a straightforward interpretation of our CCA-loss is enlarging the prediction gap between the major entropy model F ( ŷ1 , ẑ) and the auxiliary entropy model G( ẑ); so that the encoder is forced to adjust causal context ŷ1 and make it contain important information for conditional modeling. For autoregressive models with more than two stages, Eq. 7 can be extended easily. Denote the i-th stage entropy model as p ŷi| ẑ, ŷ<i ( ŷi | ẑ, ŷ<i ), which is parameterized with network F i ( ẑ, ŷ<i ); we introduce the corresponding auxiliary entropy model p ŷi| ẑ, ŷ<i-1 ( ŷi | ẑ, ŷ<i-1 ), which is parameterized with network G i ( ẑ, ŷ<i-1 ). The multi-stage CCA-loss can be defined as follows:L CCA = i E ŷ∼p ŷ|z E ẑ∼p ẑ|ψ -log p ŷi| ẑ, ŷ<i-1 ( ŷi | ẑ, ŷ<i-1 ) + log p ŷi| ẑ, ŷ<i ( ŷi | ẑ, ŷ<i ) . (8)With our proposed CCA-loss, the learned image compression model is able to spontaneously adjust the causal context, thereby promoting the rate-distortion performance.Network with CCA LossInspired by the recent work [8] , which designed a CNN-based nonlinear activation-free network to improve image restoration performance, we stack NAF-Blocks [8] in the analysis transform g a and the synthesis transform g s . Following the previous CNN-based model [9, 17] , we adopt the stacking residual blocks [19] in the auto-encoder transform for better nonlinearity. Due to the simplicity of the information that hyperprior z carries, there are only simple convolution layers for the hyper analyzer h a and synthesizer h s . Thanks to our convolutional architecture, our approach is much faster than recent LIC methods which generally adopt Transformer blocks to comprise the auto-encoder. Detailed architectures of our auto-encoder can be found in the Supplementary Materials.To establish a robust causal context and efficiently exploiting it in the autoregressive entropy models is the key to reaching state-of-the-art. The existing approaches generally constitute the causal context model intuitively. In this paper, we propose the causal context adjustment loss (CCA-loss), which compels the analysis transform to generate a more potent causal context, that is, the enhanced estimation gain of early-stage context towards the latter latent representation. Theoretically, our proposed CCA-loss is architecture-agnostic and can be utilized to train various encoders to adjust the causal context according to the given conditional modeling architecture. However, compared to the checkerboard context model that leverages adjacent spatial information as context, it is easier for our convolutional encoder to adjust information across feature channels. We therefore adopt a channel-wise grouped autoregressive architecture to design our entropy model. Furthermore, since our CCA-loss could explicitly adjust the significant information into the earlier channels, we explore an unevenly grouped strategy to take full advantage of the first several informative channels. On account of the accumulated contexts [ ŷ1 , ŷ2 , • • • , ŷi-1 ] as input to the autoregressive entropy model to predict ŷi , the unevenly grouped strategy also brings us advantages in the number of parameters and run time. Following our auto-encoder structure, we also utilize NAF-blocks [8] for a superior trade-off between accuracy and speed. For detailed network architectures of our entropy model as well as auxiliary entropy model, please refer to our Supplementary Materials. The comprehensive analysis of the benefits of the evenly and unevenly grouped strategies brought by our CCA-loss will be presented in the ablation study section.We follow the commonly used rate-distortion optimization framework to train our model. In addition to the rate losses R( ŷ), R( ẑ) and the distortion loss D( x, x), our proposed CCA-loss is introduced to explicitly adjust the causal context. The implementation of our CCA requires a group of auxiliary entropy models. In order to obtain feasible auxiliary entropy models, we further introduce auxiliary losses L Aux , which let the auxiliary model to estimate the same latent representation ŷ as the major entropy model. Therefore, the overall losses used for training our models are listed as follows:EQUATIONwe only use one parameter λ to adjust the compression rate. Detailed ablation studies about the introduced CCA-loss will be presented in our experimental section. 
experimental settings: 
 Datasets. We follow the previous work [49] and train our models on the Open Images [24] dataset. Open Images Dataset contains 300k images with short edge no less than 256 pixels. For evaluation, three benchmarks, i.e., Kodak image set [22] , Tecnick test set [1] , and CLIC professional validation dataset [41] , are utilized to evaluate the proposed network.Implementation details. We set the channel of latent representation y as 320 and that of hyperprior z is set as 192. Following the previous works, we turn the quantization operation to ⌈y -µ⌋ instead of ⌈y⌋ and restore ŷ as ⌈y -µ⌋ + µ, which benefits the entropy models. We adopt the unevenly grouped strategy to segment the latent representation into 5 uneven slices. Our detailed unevenly grouped method and discussion on it can be found in the Supplementary Materials. Our experiments and evaluations are carried out on Intel Xeon Platinum 8375C and a single Nvidia RTX 4090 graphics card. We train our network with Adam optimizer. We randomly crop 256 × 256 sub-blocks from the Open Images dataset [24] with a batch size of 8. We optimize the network with the initial learning rate 1e -4 for 2M steps and then decrease the learning rate to 1e -5 for another 0.4M steps. The network is optimized with the MSE metric, which represents the distortion loss D in Eq. 9. For the MSE metric, the multipliers λ before rate loss are {0.3, 0.85, 1.8, 3.5, 7, 15}.Comparison methods and metrics. We compare our method with the hand-crafted coding standards VVC [40] , BPG [6] and WebP [15] and recent state-of-the-art methods [4, 11, 17, 30, 45, 49] .The results of hand-crafted methods and Ballé2018 [4] are based on the implementation from Com-pressAI [5] , while, the results of other methods are provided by the method authors. We mainly use PSNR to evaluate the image quality of compression results and use bits per pixel (bpp) value to indicate the compression ratio. The BD-rate [7] and runtime of several methods are also reported to comprehensively evaluate our model. Following the commonly used setting, we also compare the MS-SSIM metric on the Kodak dataset, the MS-SSIM optimized results by different methods are shown in our Supplementary Materials.Hyperprior 1 .We firstly conduct ablation experiments to validate the effectiveness of the proposed CCA-loss. In order to facilitate the analysis, we establish a tiny model to conduct our ablation experiments. We halve the channel number and stacking count of NAF-blocks [8] in our model and only adopt a three-stage autoregressive entropy model. We evaluate our CCA-loss on evenly grouped channel-wise autoregressive model as well as unevenly grouped channel-wise autoregressive model. The BDrates of different models are reported in Table 1 , without any additional computation in the testing phase, our CCA-loss could improve the evenly grouped and unevenly grouped models by a considerable margin. Especially for the case of unevenly grouped strategy, which adopts a more aggressive strategy and decodes less number of channels in the early stage, the enhancement brought by our CCA-loss is quite large. The phenomena reveals that utilizing small amount of significant information as the initial condition is beneficial for autoregressive entropy modeling, which is in line with the motivation of our paper.To further investigate the impact on information distributions of our proposed CCA-loss, we extend a visualization of the quantities of information (code length) in the hyperprior and latent representation. The averaged information distributed ratios on the Kodak testing images by different models are shown in Fig. 2 . As can be clearly found in the histogram, for entropy models with the same network architecture, our CCA-loss is able push the network to encode significant information at an earlier stage of the autoregressive model. Rate-Distortion Comparison. We evaluate the rate-distortion performance of our proposed models by drawing the rate-distortion curves. The distortion is assessed by PSNR while the rate is calculated by the bits per pixel (bpp). We first compare our proposed network with hand-crafted codec methods [6, 15, 40] and the LIC models that once reached state-of-the-art (SOTA) [4, 11, 17, 30, 45, 49] on the Kodak dataset. The result of the PSNR metric is presented in Fig. 3 (left), which demonstrates that our proposed methods could outperform other SOTA methods. The middle sub-figure and the right sub-figure in Fig. 3 are evaluated on the CLIC Professional Validation dataset and the Tecnick dataset, respectively. The SOTA results in various datasets show the generalization and robustness of our proposed model. Compression Latency. As described in our introduction, we established a convolutional compression model for the pursuit of efficient compression. In Table 2 , we present the coding latency, as well as the number of parameters and GFLOPs, by our proposed network and recent state-of-the-art methods [30, 48, 49] . The BD-rate values by different methods are also provided for reference, the anchor RD performance of which is set as the results of VVC (vtm-12.1) on Kodak dataset (BD-rate = 0%). As can be found in the table, our method achieves a better trade-off between compression performance and coding latency than the competing methods. With more than 20% less runtime, our model obtains about 2% BD-rate gain over [30] . Visualization Analysis. Our proposed learned image compression technology is capable of restoring the image details. Fig. 4 shows two sets of comparisons with the reconstruction of VVC [40] and a recent SOTA model [30] . The visualization results are produced at low bit-rates on the Kodak dataset [22] . The comparison of the reconstructed images demonstrates that our model restores more detailed and complicated textures than other methods. For example, we restore more sharp textures on the hat (kodim4), more details of the grassland (kodim19) and wrinkles on the sails (kodim10). 
conclusion: 
 In this work, we explore the approach to adjust the causal context, which enables a superior channelwise autoregressive model and more accurate estimation in probability distributions. By imposing the Causal Context Adjustment loss (CCA-loss) and the unevenly channel-wise grouped strategy on our proposed CNN-based model, we achieve state-of-the-art rate-distortion performance. Thanks to the advantages of convolutional neural network, our discussed unevenly grouped schedule and the training method by the proposed CCA-loss, our learned image compression model maintains a great trade-off between compression latency and RD performance. Furthermore, since we did not dive into the information redistributed phenomenon brought by the unevenly grouped strategy and CCA-loss training in this paper, the issue of the laws about the information distributed among the latent representation to be compressed is still worth investigating in the future.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We set up a special experimental setting section and provide more details in the Supplementary Materials.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: For learning-based image processing methods, it is difficult to assess the error bars. However, a massive dataset drawn from natural images is utilized to perform the experiments and evaluate the results. The large number of experiments ensures the statistical significance of our experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We provide the compute resources in our experimental setting section.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our paper has no societal impact of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We are sure that the creators or original owners of assets (e.g., code, data, models), used in the paper, are properly credited and are the license and terms of use explicitly mentioned and properly respected. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Our paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Architecture of transform networks As introduced in our main paper, our compression framework is adopt the VAE framework proposed by Ballé et al. [3] , and use the same strategy of hyperprior [4] and autoregressive entropy model [35] . Generally, the transform network comprise an analyzer g a and a synthesizer g s , which play the role of feature extraction and image reconstruction. For extracting side information, another pair of hyper analyzer h a and hyper synthesizer h s is used to extracting and reconstructing the hyperprior variable z. The detailed network architecture of the above components can be found in Table 3 . The 5×5 convolution and the 5×5 transposed convolution are utilized to downsample and upsample the feature maps, respectively. Following the previous works [9, 17] , we adopt the commonly used residual blocks and the newly proposed NAF-Blocks to establish the analyzer and synthesizer. While, due to the simplicity of the information that hyperprior z carries, there are only simple convolution layers for the hyper analyzer h a and hyper synthesizer h s . The dimension numbers dim 0 , dim 1 and dim 2 in the table are set as 192, 224 and 256, respectively.Our proposed entropy model utilizes the NAF-block as well (see Fig. 5 ). The stacking NAF-blocks can enhance the concatenated features input to the entropy model, in order to obtain a more accurate estimation of the latent representation. The dimension of the latent representation in NAF-block is set as 224. Please note that we do not conduct the special training strategy like [8] for the simple channel attention (SCA) in the NAF-block, on account of no performance loss caused by this. Following the previous work [36] , we append the latent residual prediction (LRP) to restore the error introduced by the quantization operation.For our auxiliary entropy model, the only difference is that the input removes the previous one slice, that is, replace ŷ<i with ŷ<i-1 in Fig. 5 .To take full advantage of our proposed CCA-loss, we adopt the unevenly grouped strategy proposed by He et al. [17] in our method. In this part, We dive deeper into the specific grouped method and the advantages it brings. For the channel-wise autoregressive entropy models, the input to them is accumulated as the decoding progresses to the latter slices.Efficiency. For the i-th stage entropy model, the causal context containsEQUATIONCi , where H × W is the spatial size of the latent representation and C i denotes the channel number of ŷi . For the overall n-stage entropy model, the shape of the total causal context is written asn 2 R H×W × i-1 1Ci , which can be expanded as n-1 1 R H×W ×(n-i)Ci . From this expression, we could see that the former slices are reused more times, leading to more parameters and latency. Thus, the unevenly grouped strategy could benefit the model in complexity, as Table 1 shows in the ablation study section. Rate-Distortion Trade-off. In this part, we analyze the effects of selecting different grouped schedules. Inspired by previous work [34] , we parameterize the unevenly grouped strategy via a power schedule, i.e.,C(i) = N k,n • i k ,where k denotes the steepness of the increasing slices and N normalizes the n-stage autoregressive slices in sum of M . For selecting best schedule for our model, we train compression models with different grouping hyperparameters with the loss function in Eq. 9 (including our CCA-loss). We evaluate grouping strategies with different k values. The rate-distortion trade-offs by different models are shown in Fig. 6 , the RD curve achieved by our selected schedule (i.e. k = 1.7) is presented for reference. As can be seen in the figure, the setting of k = 1.7 achieves the best rate-distortion performance, which we select for our ultimate SOTA model.For higher MS-SSIM performance to adapt the real eyesight, we also produce the model of the MS-SSIM optimized objective. The distortion loss is replaced by 1 -MS-SSIM( x) and the multiplier λ before the rate loss are set as {0.2, 0.65, 1.5, 3.2, 6, 15}. The comparison of the rate-distortion curves with previous works is released in Fig. 7 . We compare the reconstruction results on kodim20 (Fig. 8 ) and kodim24 (Fig. 9 ) of our model with those of Liu et al. [30] and several hand-crafted methods, i.e., VVC [40] , Webp [15] , JPEG [42] . NeurIPS Paper ChecklistJustification: We are sure that research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: While there were some reservations in the initial reviews, the rebuttal addressed most of the reviewers' concerns. All the reviewers unanimously voted for acceptance in their final rating. The AC agrees with the reviewers' recommendations and supports acceptance of the paper. The AC recommends including additional evaluations performed during the discussion of the paper (either in the appendix or main paper, space permitting.).

Reviews: review: paper_summary: paper_summary: The paper proposes an auxiliary neural network in conjunction with a loss  function during training to improve the better contextual modeling of the data during training phase. The advantage of the proposed technique is that the auxiliary network is not required during inference and hence the complexity is minimal. The work conducts a series of experiments that supports the claims.
strength_weakness: strength_weakness: * There is no strong weakness
questions: questions: * Could you elaborate on MAC operations versus FLOPS?
limitations: limitations: No discussion on limitations
paper_summary: paper_summary: This work proposes a novel causal context adjustment loss to explicitly guide the encoder in prioritizing important information at the early stage of the autoregressive entropy model, which is both interesting and significant compared to the implicit modeling in ELIC. The loss is designed based on the entropy loss between different conditional channel-wise transformed latent slices. Extensive experiments have demonstrated the effectiveness of this design.
strength_weakness: strength_weakness: There are some weaknesses that need to be addressed:

1. In Table 1, why is the performance only compared with anchor BPG instead of VVC? Whether this strategy is effective with stronger foundational compression codecs? It is suggested to validate your method with stronger codecs and compare the performance with VVC as shown in Table 2.
2. The abstract could be reorganized for a better illustration of the contribution. For example, you should highlight the explicit guidance for the encoder to adjust important information in the early stage of the autoregressive entropy model.
3. It would be better to compare with ELIC in lines 58-60.
questions: questions: Whether this loss can be combined with the spatial context modeling strategy like "checkerboard"?
limitations: limitations: No limitation analysis is provided in this work.
paper_summary: paper_summary: This paper proposed Causal Context Adjustment loss (CCA-loss) to explicitly adjust the causal context. However, this paper also proposes an efficient image compression model, which seems to be irrelevant to its main contribution, namely CCA loss. CCA loss and efficient architectures are orthogonal, and there is no obvious correlation between the two, which makes this paper rather scattered.
strength_weakness: strength_weakness: 1. CCA Loss should assist more network architectures to prove its effectiveness. Using NAFNet alone is not enough.
2. This paper lacks analysis. Is the causal context organization of the model better after using CCA Loss? Some visual analysis can be given, for example, after using CCA Loss, can the model notice areas that the model without CCA Loss did not notice?
questions: questions: Why use MS-SSIM instead of SSIM?
limitations: limitations: Please see the weaknesses and questions.
paper_summary: paper_summary: The paper presents a novel approach to learned image compression (LIC) by introducing a Causal Context Adjustment loss (CCA-loss). This method aims to improve the rate-distortion (RD) performance of autoregressive entropy models used in LIC. The proposed approach allows the neural network to adjust the causal context dynamically, enhancing the accuracy of latent representation estimations. The paper also leverages a convolutional neural network (CNN) based model and an unevenly channel-wise grouped strategy to achieve a balance between inference latency and RD performance. Experimental results on benchmark datasets demonstrate that the proposed method outperforms existing state-of-the-art LIC techniques in both RD performance and computational efficiency.
strength_weakness: strength_weakness: 1) Limited Exploration of Context Models: The paper mentions the potential for further investigation into the organization of causal contexts, indicating that current methods are somewhat intuitive and may benefit from more structured approaches.
2) Complexity in Implementation: While the proposed method is efficient, implementing the uneven channel-wise grouped strategy and the CCA-loss function may pose challenges for practitioners not well-versed in deep learning techniques.
questions: questions: 1) Further Exploration of Causal Context Models:
a) Question: The paper mentions that current methods for organizing causal contexts are somewhat intuitive. Can you elaborate on the potential for more structured approaches in organizing causal contexts? Are there any preliminary ideas or experiments that could guide future research in this direction?
b) Suggestion: Consider conducting a more detailed analysis or ablation study on different causal context models. This would help in understanding the strengths and limitations of various approaches and could provide a more comprehensive justification for the chosen method.
2) Effectiveness of CCA-loss Across Different Models:
a) Question: The paper demonstrates the effectiveness of CCA-loss within the proposed model. Have you tested the applicability of CCA-loss with other LIC architectures or models? How does it perform in different settings?
b) Suggestion: Extending the evaluation of CCA-loss to other architectures could demonstrate its versatility and robustness. Including results from such experiments would show whether the benefits of CCA-loss are generalizable across different models.
limitations: limitations: 1) The authors mention that current methods for organizing causal contexts are somewhat intuitive and that further exploration is needed. However, they do not delve deeply into the potential limitations this may impose on their results.
2) The implementation complexity of the proposed method, particularly the uneven channel-wise grouped strategy and CCA-loss function, is not explicitly discussed as a limitation.

Comments: title: Response to authors
Response to authors
comment: Thanks for the responses of authors. I will keep my score.
The author answered my doubts. I think the method proposed in this paper achieves excellent compression performance and contributes to the field of compression. I will keep my score.
Thank you for your response and the insightful suggestions in enhancing our paper. We will add these supplementary experiments to the appendix in our revised manuscript.
Thanks for the responses from the authors. This response solved the issues I raised, but I suggest that the authors add these supplementary experiments to the main text or appendix to enhance the paper's solidity. Overall, I intend to up my rating to 5.

Also, authors could use LAM[1] to obtain clearer visualization results, further enhancing the persuasiveness of this paper.

[1] Interpreting Super-Resolution Networks with Local Attribution Maps, CVPR 2021.
We sincerely appreciate your response. The differences between the large and small models lie in the encoder-decoder transform as well as entropy model. The large model has 64.89M parameters and 615.93G FLOPs, while the small model has 22.16M parameters and 150.55G FLOPs. Specifically, compared to the large model, we have removed the residual blocks in both the encoder and decoder, and reduced the number of channels in the encoder and decoder from [192, 224, 256] to [128, 128, 128]. Additionally, the NAF-block stacks in the auto-encoder and entropy models contain 4 blocks for the large model and 2 blocks for the small model. The autoregressive slices are set to 5 for the large model and 3 for the small model. The stronger encoder-decoder in the large model produces a more robust latent representation, and the larger entropy model more accurately estimates the probability distribution, leading to better rate-distortion performance, which accounts for the performance gap between the two models. We will release the code and checkpoints for both models after completing the code cleanup. 

| Model | #Params | FLOPs(G) | Channel number  | NAF-blocks | AR slices | Residual blocks |
| :---: | :-----: | :------: | :-------------: | :-------: | :-------: | :-------------: |
| large | 64.89M  |  615.93  | [192, 224, 256] |     4     |     5     |  $\checkmark$   |
| small | 22.16M  |  150.55  | [128, 128, 128] |     2     |     3     |        X        |
Thanks for your responses from the authors. Can you clarify the differences between your large model and your small model? Generally, a simple model size cannot bring such performance improvement in compression actually. What's the improvement are from? Transform? Quantization? Entropy Coding?