RECOMBINER: Robust and Enhanced Compression with Bayesian Implicit Neural Representations

We thank the reviewer for their review and suggestions to improve our paper. Below, we respond to the reviewer’s concerns. We are happy to discuss any further concerns the reviewer might have. However, should we have addressed the reviewer’s concerns and questions, we kindly invite them to raise their score.

Some of the tasks seem a little bit selective in terms of baselines.

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.

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.

Errors in protein data: As shown at https://www.rcsb.org/stats/distribution-resolution, most structures have a resolution around 2Å, and only a very small minority have a resolution below 1Å. Thus, our errors are still much less than the machine resolution; hence, we believe they should be insignificant for any downstream task.

Furthermore, unlike the baselines, our method is highly extensible. For example, it can be easily extended to capture the rotation/translation-invariant by modifying the loss function. It can also take the protein’s uncertainty into account with ease. As another example, it is known that the alpha-helix region of the protein has lower uncertainty than coils. We can capture this information by appropriately weighting the loss function, which ensures that the compression budget is allocated to the more important regions, such as the alpha-helix. In our paper, our main goal for this experiment is to demonstrate RECOMBINER’s modality-agnosticism and leave these specific extensions for future work.

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

$

\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{,}

$

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$.

We thank the reviewer for their constructive feedback on our paper. We are delighted that the reviewer appreciates our technical contributions and values our work. We respond to the reviewer’s concerns below. We are also happy to discuss any further questions the reviewer might have.

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.

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].

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.

References

[1] Martens, J., & Grosse, R. (2015, June). Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning (pp. 2408-2417). PMLR

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.

Practicality. We agree this is a limitation of our work, as discussed in the limitation section. However, we highlight that our method’s decoding time is fast. Here, we provide the encoding time of COMBINER and RECOMBINER on Kodak for comparison. As we also mentioned in our paper, we can significantly reduce the coding time by reducing the sample size, with a minor impact on the performance. Therefore, we provide the reduced encoding time in this Table. Note that RECOMBINER is generally faster than COMBINER on Kodak due to our patching and parallelizing strategy.

As for COIN++, the author only reported the time complexity on CIFAR-10. Their reported time complexity is based on one image. While we follow COMBINER, compress 500 images in one batch. To fairly compare with them, we also report the average coding time per image. Since COIN++ is based on MAML, it only requires a handful of gradient descent steps during test time, and as such, its encoding complexity is lower them ours. The authors of VC-INR do not report runtime results, but since their method is also based on MAML, we expect their coding complexity to be lower than ours.

We thank the reviewer for their detailed and insightful review. We are delighted that the reviewer found our paper well-written and recognized our method outperforms COMBINER and that we performed extensive empirical evaluation. We respond to the reviewer's concerns below. Should the reviewer find our answers satisfactory, we kindly invite them to consider raising their score.

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.

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.

We thank again the reviewer for their effort during the reviewing process. We believe we have addressed the stated concerns with our response, and we would like to ask the reviewer if they think this is the case as well. If the response is affirmative and based on the higher ratings provided by the other reviewers, we would like to ask the reviewer if they are happy to increase their score.

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.

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).

We are happy to answer any further questions the reviewer might have. If we answered the reviewer's questions adequately, we kindly invite the reviewer to consider raising their score.

Dear Reviewer oN2Z,

The authors have submitted the response and also requested to discuss it with you.

Could you please read their response and give your feedback? Thanks.

Best,

AC