Paper: RECOMBINER: Robust and Enhanced Compression with Bayesian Implicit Neural Representations (ICLR-2024)

Reviewer comment (Evaluation):
Despite the improvements, overall INRs continue lag behind performance of competing methods on 3 out of 4 tasks.

Author response:
**INR performance gap and slow encoding:** While these are undoubtedly important observations, we argue that they are general limitations of INR-based approaches and aren’t weaknesses of RECOMBINER per se. We also note that VAE-based methods are much more mature than INRs and rely on specifically designed network architecture and powerful entropy models. Moreover, INRs offer unique advantages compared to VAE-based methods, such as data modality-agnosticism, lightweightness, and fast decoding times on even CPUs.

Reviewer comment (Evaluation):
INR performance comes at a compute cost penalty. This would be particularly large in the case of the video and image codec comparisons.

Author response:
**INR performance gap and slow encoding:** While these are undoubtedly important observations, we argue that they are general limitations of INR-based approaches and aren’t weaknesses of RECOMBINER per se. We also note that VAE-based methods are much more mature than INRs and rely on specifically designed network architecture and powerful entropy models. Moreover, INRs offer unique advantages compared to VAE-based methods, such as data modality-agnosticism, lightweightness, and fast decoding times on even CPUs.

Reviewer comment (Evaluation):
Many neural methods for compression (particularly for audio) rely on perceptual compression rather than rate-distortion as is experimented with in the present paper. The paper does not consider how RECOMBINER could be adapted to become a perceptual codec.

Author response:
**Perceptual compression:** We agree that this is a promising direction and are currently working on incorporating perceptual constraints into INR-based compression. However, we emphasize that all current INR-based compression methods focus on MSE rate-distortion performance. Therefore, we do not believe that this is one of our weaknesses. 

**Regarding the Reviewer’s questions:**
 1. Yes, one way to adapt RECOMBINER to a perceptual codec is to add an adversarial term in the loss function. This term can come from an extra discriminator.
 2. As discussed above, no INR-based method is currently designed for perceptual loss. But we will be happy to consider and compare with the work you mentioned once we successfully adapt RECOMBINER to a perceptual codec. We will also update our discussion section in our camera-ready version to note this future direction.
 3. For a $M$-dimensional modality, we can write the coordinate as $\mathbf{x} = [x_1, ..., x_M]^\top$. We define its Fourier embeddings $ \gamma(\mathbf{x}) \in \mathbb{R}^D$ as
\begin{align}
   \gamma(\mathbf{x}) = \Big[\cos(  a^{{0}/{d}} \pi x_1 ), &\dots, \cos(  a^{{0}/{d}} \pi x_M ), \\\\
   \cos(  a^{{1}/{d}} \pi x_1 ), &\dots, \cos(  a^{{1}/{d}} \pi x_M ), \\\\
   &\vdots\\\\
    \cos(  a^{{d}/{d}} \pi x_1 ), &\dots, \cos(  a^{{d}/{d}} \pi x_M ), \\\\
    \sin(  a^{{0}/{d}} \pi x_1 ), &\dots, \sin(  a^{{0}/{d}} \pi x_M ), \\\\
   \sin(  a^{{1}/{d}} \pi x_1 ), &\dots, \sin(  a^{{1}/{d}} \pi x_M ), \\\\
&\vdots\\\\
    \sin(  a^{{d}/{d}} \pi x_1 ), &\dots, \sin(  a^{{d}/{d}} \pi x_M )\Big]^\top\text{,}
\end{align}
For image, $M=2$; for audio and protein, $M=1$ (note that the backbone atoms of protein are located in a chain); for video, $M=3$. 
In our experiments, we follow COMBINER, setting $d = \lfloor\frac{D}{2M}\rfloor-1$ and $a=1024$.

Reviewer comment (Experiments):
Older image compression methods are used as baselines. The handcrafted baseline of choice is now VTM, and the baseline neural image compression is ELIC (He, 2022).

Author response:
We agree with the reviewer and have updated our paper with modifications:
 1) We have updated our paper to include VC-INR in our audio compression plot in Figure 3b and plan to include it for video compression in Figure 3c. Please see our response to Reviewer oN2Z for a detailed explanation of how we ensure a fair comparison between the methods.
 2) We added ELIC and VTM as baselines in Figure 3a. We find that they do not outperform RECOMBINER on CIFAR-10 at low bitrates, and our method is still SOTA in this region. Interestingly, VTM performs even worse than BPG on Cifar-10 at lower bitrates. This might be because VVC is designed for higher resolution. Thus, advanced techniques such as block partitioning and mode selection in VVC might be less effective when compressing low-resolution images. We also tested ELIC on CIFAR-10, but it did not work well on low-resolution images. Its performance was even worse than Balle 2018 and sometimes does not converge. Therefore, we omit the ELIC curve on CIFAR-10.

Reviewer comment (Writing):
The quality of this paper needs to be comprehensively improved, whereby many typos exist. For example, "have a neural network memorize the data (Stanley, 2007) and encode the network weights instead.".

Author response:
**Typo:** We thank the reviewer for their feedback and have clarified the sentence they highlighted. We will also carefully re-read our manuscript and correct any typos.

Reviewer comment (Theory):
For the linear reparameterization module, why A^[l] is updated during the training stage? Also the authors claim that the block-wise diagonal matrix operates as good as the full covariance matrix. Is any verification on this, from either theoretically or emperically?

Author response:
**Linear reparameterization:** The matrix $A$ controls the correlation between the network weights. This correlation will depend on the bitrate and the data modality and thus cannot be handcrafted in advance. Hence, we learn it from the training data. To better understand what matrix $A$ has learned, we provide the visualization of $A$ in Figure 7(c)(d) in the Appendix. By updating in the training period, we can see that $A$ is specifically tailored for different bitrates: some columns are pruned out (i.e., almost 0), and some are activated. Note that before the training, we randomly initialize the matrix $A$. So, this visualization illustrates what $A$ has learned by updating. 

We found that models trained with full covariance matrices performed very similarly to the ones using layer-wise block diagonal matrices in practice but were much more expensive to train/evaluate. We also note that the layer-wise block diagonal matrix is a standard approximation widely used for neural networks. For example, KFAC also assumes independence between layers [1].

Reviewer comment (Theory):
For the learnt positional embeddings, I am a bit confused on using additional position cues. Since x_i also includes the coordinate information, why using z_i helps to address the global representation challenge?

Author response:
**Positional encoding:** At a high level, the $z_i$s just allow more flexibility for the model and can be thought of as “zeroth-level bias.” We visualize the learned positional encodings for a Kodak image in Figure 7a in the Appendix to illustrate their effect.

Reviewer comment (Experiments):
**R-D tradeoff on image.** While it is definitely good to see that recombiner outperforms combiner, the performance does not clearly outperform VC-INR, an even older baseline. I do not see this as a very big drawback, but this observation makes it quite questionable whether combiner-like approach has any great potential in the long run.

Author response:
**R-D tradeoff on image.**
We appreciate that RECOMBINER’s potential compared to other methods is not clear at first sight from the figures. However, we believe that RECOMBINER is more promising than previous INR-based methods because it does the "right thing" and jointly optimizes the rate-distortion tradeoff. Non-variational INR compression methods cannot do this. Even VC-INR, the most advanced technique, follows a two-stage training pattern: 
 1) Over a training set, it fits an INR to each datum using only the distortion as the objective. Thus, it obtains a new dataset of INR weights. 
 2) Then, it fits a variational autoencoder to the INR weight dataset inspired by Balle et al.'s model. 

Regarding VC-INR in particular, we would like to highlight three aspects of the work and contrast them with RECOMBINER:
 1. From a theoretical perspective, VC-INR's performance is capped by the abovementioned first stage of training. We don't necessarily believe that this cap is a big problem in practice, but it's still worth noting that RE/COMBINER does not have this limitation.
 2. We can view VC-INR's second training stage as building a complex entropy model for the weights, essentially equivalent to RECOMBINER's weight prior. Thus, considering that VC-INR's entropy model is given by a VAE while RECOMBINER's prior is multivariate Gaussian, it is quite encouraging that the two methods' performances match closely. Furthermore, this indicates that designing more elaborate weight priors for RECOMBINER (e.g., using normalizing flows) could significantly boost its performance.
 3. As a more subtle point, in private communication with the authors of VC-INR, the authors highlighted to us that to obtain each point on VC-INR's rate-distortion curve, they trained several different INRs with varying sizes but equal rate constraints and always compressed the best-performing INR. They noted that this selection process was especially important to ensure good performance in the low-to-medium bpp range. In contrast, in our paper, we fixed a single architecture for all our experiments, even across data modalities, and trained and compressed only a single model to obtain each point. We could likely get significantly improved curves if we followed a similar "grid search" over architectures at each rate point.

To further add to the future potential of the COMBINER-like approach, we also highlight that training RE/COMBINER is much more stable than training VC-INR, as it does not involve meta-learning. The authors of VC-INR note in their paper: “In line with prior work, we find the direct optimisation of non-linear networks via Meta-Learning to be unstable and under-performing. As low-rank parameterisations are also known to suffer from stability issues, the direct use yield unsatisfactory results.” Hence, they need to introduce several techniques to stabilize training. In contrast, RECOMBINER’s training is robust and stable without requiring stabilization techniques. To illustrate this, we added a plot of RECOMBINER’s training curve in Figure 10 in Appendix D.4. As the figure shows, there is a single dip in the training curve, which is due to the dynamic beta-adjustment procedure we describe in Section 3.4. After beta adjusts to its appropriate level early in the training, the training curves essentially monotonically increase without any dips.

Reviewer comment (Experiments):
**Comparison on audio/video.** I wonder how the RECOMBINER compare with VC-INR and COIN++ on audio/video datasets. The paper says that it does not compare with these baselines, because "the works use different data splits." I do not think this is a good excuse to not compare with these baselines. It seems evident that recombiner outperforms combiner, but in the end, we would like to understand whether the combiner-like approach is indeed a useful framework, when compared with other INR-based paradigms (and furthermore, VAE-based ones).

Author response:
**Comparison on audio/video.**

*Regarding COIN++:* We omitted COIN++ in the audio comparison, as on the reported bitrates, it performs worse than MP3, which RECOMBINER already outperforms significantly. Moreover, the authors of COIN++ do not report experiments on video data; hence, it is missing from that comparison plot.

*Regarding VC-INR:* We agree with the reviewer and have changed our view. To ensure a fair comparison, we started evaluating RECOMBINER on the same test set on which VC-INR was evaluated. We will include the results in the camera-ready version of the paper. To get an idea of the performance, we included VC-INR in an indicative comparison plot in the updated manuscript on audio comparison in Figure 3b to include VC-INR. As we can see, the two methods have similar performance, with RECOMBINER even slightly outperforming VC-INR at a higher bitrate.

Furthermore, we reached out to the authors of VC-INR to share the precise details of their train/test split on the video data so that we can provide a fair comparison to their method. We will add this comparison in the final version of the paper.

Reviewer comment (Experiments):
**Limited range of bitrate in comparison---problems in extending to higher bitrates?.** Comparing with baselines, the range of bitrates considered is considerably smaller. For Kodak, VC-INR compares on the bitrate up to 3.5, while this work only considers up to 1.2. For videos, the range is again up to 1.2, while VC-INR does it to over 4. The same for the audio. I wonder why the authors made this choice. Does this mean that recombiner has difficulties in training in high bpp regime?

Author response:
**Limited range of bitrate in comparison---problems in extending to higher bitrates?**
The reason for the limited range for the bitrate is quite prosaic: other than VC-INR, all other INR works focus on the smaller range highlighted by the reviewer. However, we emphasize that there are absolutely no issues with training RECOMBINER at higher bitrates. To demonstrate this point, we added Figures 11 and 12 in the appendix that compare the methods on an extended range on Kodak. We added a further high-bitrate point to the audio comparison in Figure 3b. As we can see, RECOMBINER remains competitive with VC-INR in this regime as well.

Reviewer comment (Novelty):
**(minor) Limited Impact of Technical Contributions.** Given the lukewarm performance gain over competing paradigms, it would be great if the proposed technical innovations are very novel or have a wider applicability to other fields of machine learning. I am not sure if this is the case; the linear reparameterization and learned positional encodings are either not very new or highly specialized to the context of recombiner. I do appreciate the technicality of the hierarchical Bayes part, but some components are somewhat mysterious to me (the strange importance of random permutation) and the idea may not have a wider impact outside this specific context.

Author response:
**Limited Impact of Technical Contributions.**

*Linear reparameterization:* We agree that a linear reparameterization of the weights, in general, is not new. However, to the best of our knowledge,  the type of linear reparameterization we propose, i.e., using the same linear transform for the prior and the posterior, is novel in the context of Bayesian neural networks. It is particularly well-suited for any further work utilizing variational INRs, as the number of its variational parameters only grows linearly in the number of weights instead of quadratically (which would be the case if we learned a linear transform separately for each INR). 

*Random permutations:* We need the random weight permutation because we encode an INR’s weights sequentially, block-by-block. Ideally, when we infer the INR weight posterior, the KL divergence of the posterior from the prior in each block would conform to the bit budget we set (in the paper, we set a budget of 16 bits per block following COMBINER). However, this constraint is only enforced on average, and it is unlikely that each block will conform to the budget after inference. Thus, randomly permuting the weights allows us to distribute the KL better among the blocks.

Regarding the specific random permutation technique we propose in Section 3.3 for patch-INRs: As we break an image up into patches, in addition to distributing the KL (information content) as evenly as possible, an additional aspect here is that we want to be able to encode the patch-INRs in parallel. With this in mind, the permutation technique we propose is essentially the simplest one that allows us to achieve these two goals simultaneously.

Reviewer comment (Evaluation):
The figure 5c is very interesting, in the sense that using the tricks without random permutation is almost the worst among all choices (brown). Do you have any explanation?

Author response:
**The figure 5c is very interesting, in the sense that using the tricks without random permutation is almost the worst among all choices (brown). Do you have any explanation?**
As we mentioned in our previous point, randomly permutating the weights is crucial to ensure that block-wise relative entropy coding can work well. Especially regarding our patch-INRs, it ensures information content is spread as evenly as possible. Otherwise, we might spend too many bits on simple parts of the image (e.g., the sky), leaving insufficient bits for complex parts (e.g., one containing a person).