Test-time Adaptation for Image Compression with Distribution Regularization

Authors: Kecheng Chen, Pingping Zhang, Tiexin Qin, Shiqi Wang, Hong Yan, Haoliang Li
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: test-time adaptation, image compression, entropy coding

Abstract:
Current test- or compression-time adaptation image compression (TTA-IC) approaches, which leverage both latent and decoder refinements as a two-step adaptation scheme, have potentially enhanced the rate-distortion (R-D) performance of learned image compression models on cross-domain compression tasks, \textit{e.g.,} from natural to screen content images.  However, compared with the emergence of various decoder refinement variants, the latent refinement, as an inseparable ingredient, is barely
 tailored to cross-domain scenarios. To this end, we are interested in developing an advanced latent refinement method by extending the effective hybrid latent refinement (HLR) method, which is designed for \textit{in-domain} inference improvement but shows noticeable degradation of the rate cost in \textit{cross-domain} tasks. Specifically, we first provide theoretical analyses, in a cue of marginalization approximation from in- to cross-domain scenarios,  to uncover that the vanilla HLR suffers from an underlying mismatch between refined Gaussian conditional and hyperprior distributions, leading to deteriorated joint probability approximation of marginal distribution with increased rate consumption. To remedy this issue, we introduce a simple Bayesian approximation-endowed \textit{distribution regularization} to encourage learning a better joint probability approximation in a plug-and-play manner. Extensive experiments on six in- and cross-domain datasets demonstrate that our proposed method not only improves the R-D performance compared with other latent refinement counterparts, but also can be flexibly integrated into existing TTA-IC methods with incremental benefits.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730298440490

thread_type:
official_review

<thread_object>
creator: reviewer_5d5v

<content>
summary:
The author propose a novel test-time adaptation framework for image compression through directly refine the latent variables without altering any model parameters.

strengths:
This research topic of this work seems a bit a minority (image compression + test-time adaptation) and contains a lot of background knowledge. After hard attempts, I think I still can't give an accurate evaluation. It is recommended that AC find another highly relevant reviewer in order to make a more comprehensive judgment.

weaknesses:
None

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730455251631

thread_type:
official_review

<thread_object>
creator: reviewer_shhq

<content>
summary:
This article investigates test-time adaption image compression method on cross-domain compression tasks from the perspective of latent refinement, by extending hybrid latent refinement and introducing a simple Bayesian approximation-endowed distribution regularization. The authors provide theoretical analysis and experimental results to demonstrate the performance of this method. My detailed comments are as follows.

strengths:
1.	The authors propose their method from the perspective of latent refinement, which hasn’t been fully investigated. Compared with HLR, the method proposed by the author greatly improves the performance cross-domain compression tasks.
2.	The authors provide sufficient theoretical analysis on the method proposed in the article.

weaknesses:
1.	The experimental results of this plug-and-play method are relatively weak, which cannot illustrate the generalization ability of this method on other TTA-IC approaches. It would be better to apply this method to more baseline approaches.
2.	It would be better to report the image compression results on top of popular image reconstruction benchmarks, such as Set5, Set14, BSD100, Urban100, etc.
3.	This paper is hard to follow and the writing could be further improved. It would be better to explicitly highlight the key contributions of this paper.
4.	The compared methods BLR and HLR are too old since both are published before 2021. It would be better to compare with more recent work, such as [A].
[A] Downscaled representation matters: Improving image rescaling with collaborative downscaled images, ICCV 2023.

questions:
The compared methods are too old. More popular image reconstruction/compression datasets should be considered for comparisons.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730535934113

thread_type:
official_review

<thread_object>
creator: reviewer_7uht

<content>
summary:
Current Test-Time Adaptation Image Compression (TTA-IC) methods improve rate-distortion (R-D) performance on cross-domain tasks by refining both latent variables and decoders in a two-step process. However, latent refinement has not been specifically optimized for cross-domain scenarios. The authors conduct theoretical analyses to identify a mismatch between refined Gaussian conditionals and hyperprior distributions in the vanilla Hybrid Latent Refinement (HLR), leading to poorer joint probability approximations and higher rate consumption. To address this issue, they introduce a Bayesian approximation-based distribution regularization technique that enhances joint probability modeling in a plug-and-play manner. Extensive experiments across six in-domain and cross-domain datasets demonstrate that the proposed method outperforms existing latent refinement approaches in R-D performance.

strengths:
1. The paper effectively identifies the limitations of existing Hybrid Latent Refinement (HLR) methods in cross-domain image compression, particularly highlighting the mismatch between refined Gaussian conditionals and hyperprior distributions.

2. It provides a thorough theoretical analysis using marginalization approximation, establishing a solid foundation for the proposed improvements.

3. The introduction of a Bayesian approximation-based distribution regularization technique successfully addresses the identified mismatch, enhancing joint probability approximation. 

4. The method shows improvements in rate-distortion (R-D) performance compared to other latent refinement approaches, validating its effectiveness.

weaknesses:
1. The paper does not include a comparison of GPU memory usage and GFLOPs, which are crucial metrics for evaluating the complexity of various methods. Including these metrics would illustrate the trade-off between performance gains and computational complexity more clearly. It would be beneficial to present memory usage across different datasets, especially since DVI is introduced during optimization, making it reasonable to assess the additional complexity incurred.

2. In [1], the posterior distribution of $\hat{z}$ is assumed to be a uniform distribution. It is unclear whether directly inferring $\hat{z}$ through a hyper analysis transform in out-of-distribution (OOD) scenarios is the better choice? Because this approach not only avoids incurring a larger $\Delta H$ but also naturally matches its posterior distribution during pre-training. 

3. The paper introduces a regularization term based on HLR in TTA-IC. While this contribution is valuable, its impact on the community appears moderate.

[1] Ballé, J., Minnen, D., Singh, S., Hwang, S. J., & Johnston, N. (2018). Variational image compression with a scale hyperprior. arXiv preprint arXiv:1802.01436.

questions:
Will the authors make their code public to assist the community upon acceptance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730546982615

thread_type:
official_review

<thread_object>
creator: reviewer_gydx

<content>
summary:
The authors propose a new latent refinement method to enhance the rate-distortion (R-D) performance for test-time adaptation image compression. They theoretically analyze that existing methods require higher rate consumption to estimate the marginal probability of the latent variables in cross-domain scenarios. To address this, the authors incorporate a new distribution regularization term into the R-D objective, promoting a better approximation of the latent representation. Experimental results across six cross-domain datasets demonstrate the effectiveness of the proposed approach, showing its superiority over existing methods. Please see my detailed comments below.

strengths:
1. The authors offer a theoretical analysis highlighting limitations of existing methods, which informs their introduction of a distribution regularization technique. This technique effectively adapts the compression model to test data with domain shifts.
2. They provide detailed analysis results on the entropy curves of different probabilities about different methods. The results in Figure 4 demonstrate that the proposed method can lead to a good convergence of rate consumption and better R-D performance.
3. Experimental results demonstrate that the proposed method can be easily integrated with existing methods to enhance R-D performance across cross-domain image datasets.

weaknesses:
1. The test datasets used in the experiments contain a limited number of images, which may affect the robustness of the results. To strengthen the findings, please consider providing additional results on larger-scale datasets.
2. It would be helpful to include a preliminary section explaining the image compression process before the theoretical analysis of existing methods. This section could clarify the steps involved in obtaining the latent variable y and hyper-latent variable z (side information), enhancing readers' understanding of the compression model and the subsequent theoretical analysis.
3. In addition to the adaptation cost of different methods, it is better to provide the compression cost of the proposed methods, which would offer a more comprehensive view of its performance. 
4. The authors do not provide a limitation analysis of the proposed methods. For instance, the long adaptation time may limit the proposed methods to be used in real-time applications.
5. More discussion on recent test-time adaptation works would be beneficial to place this study in context. Relevant works to consider include, but are not limited to, [A-E].

[A] Test-Time Training Can Close the Natural Distribution Shift Performance Gap in Deep Learning Based Compressed Sensing, ICML 2022.

[B] Efficient Test-Time Model Adaptation without Forgetting, ICML 2022.

[C] Tent: Fully Test-Time Adaptation by Entropy Minimization, ICLR 2021.

[D] Efficient Test-Time Adaptation for Super-Resolution with Second-Order Degradation and Reconstruction, NeurIPS 2023.

[E] Test Time Adaptation for Blind Image Quality Assessment, ICCV 2023.

questions:
1. In Table 3, the results indicate that increasing the number of dropout layers can negatively impact image compression performance. Please add further discussion on how to effectively approximate the distribution regularization to avoid such performance degradation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652872645

thread_type:
official_review

<thread_object>
creator: reviewer_5frr

<content>
summary:
This paper addresses a latent refinement techqniue in test-time adaptation methods for image compression.
Especially, the paper focuses on methods that use a hyper-prior model, where the image is encoded into a latent variable $y$ and a hyper-latent variable $z$ in a hierarchical way.
The paper analyzes the issue of recent latent refinement methods not generalizing well in cross-domain settings, arguing the cause to a mismatch between the hyper-prior $p(z)$ and the gaussian conditional $p(y|z)$ used in refinement methods.
Based on this conjecture, the paper proposes a distribution regularization term that uses dropout variational inference, claiming it can improve probability approximation. 
Experimental results show that the proposed method continuously improves latent refinement and demonstrates effectiveness even when integrated into state-of-the-art methods that jointly utilize latent and decoder refinement.

strengths:
1. The proposed method can be applied without altering any model parameters, making optimization easy and immune to catastrophic forgetting.
2. The paper tries providing detailed discussions on the proposed method, both in a principled and empirical way.
3. The experimental results consistently demonstrate the effectiveness of the proposed method in cross-domain image compression across different settings: (1) with only latent refinement, (2) integrated with state-of-the-art TTA-IC methods, which further include decoder adaptation, and (3) on medical images.

weaknesses:
Most importantly, as i understand, the theoretical analysis and subsequent derivation of the proposed method, which are claimed as major contribution in this paper, contain many unclear and misleading aspects. 
1. (L173, 193) The definition of "optimal probability representation" is not clearly defined, throughout the paper.
2. The entropy function, i.e., $H(X)=-\mathbb{E}_{x\sim p(x)}[\log p(x)]$, is interpreted in the form of log probability, and it seems that the expectation term is not considered in the analysis and derivation.
3. In the Proposition 1, assuming an optimal joint probability approximation of the true marginal distribution $p(y^*)=\int p(y,z)dz$, the proposition simply interprets $H(Y,Z) - H(Y^*) = H(Z|Y) + H(Y) - H(Y^*) = H(Z|Y)$ when $H(Y) = H(Y^*)$. However, the assumption $p(y^*)=\int p(y,z)dz$ may not induce equations (7,8), which can be misleading.
4. In the proof of Proposition 2, the authors argue that the entropy bottleneck $p(z_t)$ and the posterior $p(y_t|z_t)$ are learned from the source dataset (L212, source image correlated) and may not work effectively. While this makes sense, it is not a theoretical analysis, in my point of view.
5. For Corollary 1, it is unclear whether $z^*$ and $y^*$ in equation 13 correspond to in-domain or cross-domain cases. In the cross-domain case with $y^*_t$ and $z_t^*$ as in Proposition 2, equation 13 is as follows: $\triangle H = H(Y, Z) - H(Y^*_t) > H(Z^*|Y^*)$, which holds according to equations (11, 12) only if $H(Y^*_t)=H(Y^*)$ and $H(Z^*_t|Y^*_t)=H(Z^*|Y^*)$. However, there is no guarantee for the condition.
6. Importantly, in Section 3.3, the paper proposes distribution regularization using equation 13, interpreting the first three terms in equation 17 as corresponding to $H(Y|Z), H(Z), H(Z^*|Y^*)$ (the three terms at the front) in equation 13. However, it is unclear how $H(Z^*|Y^*)=\mathbb{E}_{y^*,z^*\sim p(y^*,z^*)}[p(z^*|y^*)]$ leads to $-\log p(z^m_t|y^m_t)$ in the optimization objective. Note that $H(Z^*|Y^*)$ includes expectation from $p(y^*,z^*)$.

In conclusion, the theoretical analysis in this paper seems to have many over claimed or misclaimed elements.

questions:
* Please address the weakness above.
* The writing is a bit difficult to understand, especially in the discussion parts of both the methodology and experiment sections
</content>
</thread_object>

</thread>
<thread>
creation date:
1733128004092

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thanks for your time

comment:
Dear Reviewers,

As the discussion period will be ended today, we would like to send a kind reminder about our latest responses for Reviewer 5FRR (regarding theoretical discussions) and Reviewer ShHq (regarding the concern of “insufficient comparisons”) .

We have addressed your concerns in detail via responses or incorporated your suggestions into our revisions. 
We are looking forward to your feedback.

Best wishes,

The authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper addresses test-time adaptation for learned image compression (TTA-IC) in the face of cross-domain inputs. It builds upon a latent-refinement approach originally designed for in-domain scenarios and identifies a key issue: a fundamental mismatch between the refined Gaussian conditionals and hyperprior distributions when adapting to out-of-distribution images. This mismatch leads to suboptimal joint probability approximations and increased bit rates. To remedy this, the authors introduce a Bayesian approximation-based distribution regularization, which better aligns latent distributions with their hyperpriors, resulting in improved rate-distortion (R-D) performance on cross-domain inputs.

Initial reviewer feedback was mixed. Major concerns included:

1. **Clarity and Accessibility:** The motivations and theoretical derivations were challenging to follow, especially for readers without substantial background in learned compression.
2. **Comparisons to General TTA Methods:** The paper lacked explicit connections and comparisons to broader test-time adaptation techniques.
3. **Computational Complexity:** Reviewers requested more details on computational overhead and complexity.
4. **Dataset Diversity:** The evaluation lacked experiments on a broader range of datasets.

During the rebuttal phase, the authors provided detailed clarifications, introduced additional experiments on more diverse datasets, and offered complexity analyses. These responses addressed most concerns, and by the end of the discussion, all reviewers agreed on a borderline acceptance.

The AC read paper, reviews, and discussions, and concurs with the reviewers. The additional clarifications and experiment results presented in the rebuttal adequately addresses the reviewers’ concerns. However, the AC also agrees with the reviewers that the paper can be benefited from a significant revision by improving its clarity and self-containedness; many of the mathematical derivations are either trivial or from the prior works, and can be greatly simplified  (e.g., Section 3.1.). Introducing more descriptions on background (e.g., what information the side-information is carrying and be used in decoding) can make the paper more easily accessable to non-experts of the domain. Overall, the AC believes that the paper provides a meaningful theoretical insight and demonstrates tangible performance improvements, hence recommends acceptance. The authors are encouraged to incorporate the clarifications and new results presented in the rebuttal into the camera-ready version for better clarity and accessibility.