Title: Compression with Bayesian Implicit Neural Representations
Venue: NeurIPS 2023
Accepted: yes
Keywords: Neural Compression

Implicit Neural Representation

Relative Entropy Coding

Bayesian Neural Network
Abstract: Many common types of data can be represented as functions that map coordinates to signal values, such as pixel locations to RGB values in the case of an image. Based on this view, data can be compressed by overfitting a compact neural network to its functional representation and then encoding the network weights. However, most current solutions for this are inefficient, as quantization to low-bit precision substantially degrades the reconstruction quality. To address this issue, we propose overfitting variational Bayesian neural networks to the data and compressing an approximate posterior weight sample using relative entropy coding instead of quantizing and entropy coding it. This strategy enables direct optimization of the rate-distortion performance by minimizing the $\beta$-ELBO, and target different rate-distortion trade-offs for a given network architecture by adjusting $\beta$. Moreover, we introduce an iterative algorithm for learning prior weight distributions and employ a progressive refinement process for the variational posterior that significantly enhances performance. Experiments show that our method achieves strong performance on image and audio compression while retaining simplicity.

Compression with Bayesian Implicit Neural Representations

Abstract

Many common types of data can be represented as functions that map coordinates to signal values, such as pixel locations to RGB values in the case of an image. Based on this view, data can be compressed by overfitting a compact neural network to its functional representation and then encoding the network weights. However, most current solutions for this are inefficient, as quantization to low-bit precision substantially degrades the reconstruction quality. To address this issue, we propose overfitting variational Bayesian neural networks to the data and compressing an approximate posterior weight sample using relative entropy coding instead of quantizing and entropy coding it. This strategy enables direct optimization of the rate-distortion performance by minimizing the β -ELBO, and target different rate-distortion trade-offs for a given network architecture by adjusting β . Moreover, we introduce an iterative algorithm for learning prior weight distributions and employ a progressive refinement process for the variational posterior that significantly enhances performance. Experiments show that our method achieves strong performance on image and audio compression while retaining simplicity. Our code is available at https://github.com/cambridge-mlg/combiner .

1 Introduction

With the celebrated development of deep learning, we have seen tremendous progress of neural data compression, particularly in the field of lossy image compression [1-4]. Taking inspiration from deep generative models, especially variational autoencoders (V AEs, [5]), neural image compression models have outperformed the best manually designed image compression schemes, in terms of both objective metrics, such as PSNR and MS-SSIM [6, 7] and perceptual quality [8, 9]. However, these methods' success is largely thanks to their elaborate architectures designed for a particular data modality. Unfortunately, this makes transferring their insights across data modalities challenging.

Arecent line of work [10-12] proposes to solve this issue by reformulating it as a model compression problem: we treat a single datum as a continuous signal that maps coordinates to values, to which we overfit a small neural network called its implicit neural representation (INR). While INRs were originally proposed in [13] to study structural relationships in the data, Dupont et al. [10] have demonstrated that we can also use them for compression by encoding their weights. Since the data

Gergely Flamich ˚

Jos´ e Miguel Hern´ andez-Lobato

…

…

…

…

…

…

…

…

…

…

…

…

Relative Entropy Coding

…

…

…

…

…

…

Transmit Index

Datum

Approximate Coding Cost

Distortion Measure

Variational Inference

+

Coordinates

Coordinates

-th sample from

is conceptualised as an abstract signal, INRs allow us to develop universal, modality-agnostic neural compression methods. However, despite their flexibility, current INR-based compression methods exhibit a substantial performance gap compared to modality-specific neural compression models. This discrepancy exists because these methods cannot optimize the compression cost directly and simply quantize the parameters to a fixed precision, as opposed to V AE-based methods that rely on expressive entropy models [2, 3, 14-17] for end-to-end joint rate-distortion optimization.

In this paper, we propose a simple yet general method to resolve this issue by extending INRs to the variational Bayesian setting, i.e., we overfit a variational posterior distribution q w over the weights w to the data, instead of a point estimate. Then, to compress the INRs, we use a relative entropy coding (REC) algorithm [18-20] to encode a weight sample w ' q w from the posterior. The average coding cost of REC algorithms is approximately D KL r q w } p w s , where p w is the prior over the weights. Therefore, the advantage of our method is that we can directly optimize the rate-distortion trade-off of our INR by minimising its negative β -ELBO [21], in a similar fashion to VAE-based methods [22, 2]. We dub our method Com pression with B ayesian I mplicit Ne ural R epresentations (COMBINER), and present a high-level description of it in Figure 1.

Wepropose and extensively evaluate two methodological improvements critical to enhancing COMBINER's performance further. First, we find that a good prior distribution over the weights is crucial for good performance in practice. Thus, we derive an iterative algorithm to learn the optimal weight prior when our INRs' variational posteriors are Gaussian. Second, adapting a technique from Havasi et al. [23], we randomly partition our weights into small blocks and compress our INRs progressively. Concretely, we encode a weight sample from one block at a time and perform a few gradient descent steps between the encoding steps to improve the posteriors over the remaining uncompressed weights. Our ablation studies show these techniques can improve PSNR performance by more than 4dB on low-resolution image compression.

Weevaluate COMBINER on the CIFAR-10 [24] and Kodak [25] image datasets and the LibriSpeech audio dataset [26], and show that it achieves strong performance despite being simpler than its competitors. In particular, COMBINER is not limited by the expensive meta-learning loop present in current state-of-the-art INR-based works [11, 12]. Thus we can directly optimize INRs on entire high-resolution images and audio files instead of splitting the data into chunks. As such, our INRs can capture dependencies across all the data, leading to significant performance gains.

To summarize, our contributions are as follows:

We propose variational Bayesian implicit neural representations for modality-agnostic data compression by encoding INR weight samples using relative entropy coding. We call our method Com pression with B ayesian I mplicit Ne ural R epresentations (COMBINER).

We propose an iterative algorithm to learn a prior distribution on the weights, and a progressive strategy to refine posteriors, both of which significantly improve performance.

We conduct experiments on the CIFAR-10, Kodak and LibriSpeech datasets, and show that COMBINER achieves strong performance despite being simpler than related methods.

2 Background and Motivation

In this section, we briefly review the three core ingredients of our method: implicit neural representations (INRs; [10]) and variational Bayesian neural networks (BNNs; [21]), which serve as the basis for our model of the data, and relative entropy coding, which we use to compress our model.

Implicit neural representations: We can conceptualise many types of data as continuous signals, such as images, audio and video. Based on neural networks' ability to approximate any continuous function arbitrarily well [27], Stanley [13] proposed to use neural networks to represent data. In practice, this involves treating a datum D as a point set, where each point corresponds to a coordinate-signal value pair p x , y q , and overfitting a small neural network f p x | w q , usually a multilayer perceptron (MLP) parameterised by weights w , which is then called the implicit neural representation (INR) of D . Recently, Dupont et al. [10] popularised INRs for lossy data compression by noting that compressing the INR's weights w amounts to compressing D . However, their method has a crucial shortcoming: they assume a uniform coding distribution over w , leading to a constant rate , and overfit the INR only using the distortion as the loss. Thus, unfortunately, they can only control the compression cost by varying the number of weights since they show that quantizing the weights to low precision significantly degrades performance. In this paper, we solve this issue using variational Bayesian neural networks, which we discuss next.

Variational Bayesian neural networks: Based on the minimum description length principle, we can explicitly control the network weights' compression cost by making them stochastic. Concretely, we introduce a prior p p w q (abbreviated as p w ) and a variational posterior q p w | D q (abbreviated as q w ) over the weights, in which case their information content is given by the Kullback-Leibler (KL) divergence D KL r q w } p w s , as shown in [28]. Therefore, for distortion measure ∆ and a coding budget of C bits, we can optimize the constrained objective

In practice, we introduce a slack variable β and optimize the Lagrangian dual, which yields:

with different settings of β corresponding to different coding budgets C . Thus, optimizing L β p D , q w , p w q is equivalent to directly optimizing the rate-distortion trade-off for a given rate C .

Relative entropy coding with A* coding: We will use relative entropy coding to directly encode a single random weight sample w ' q w instead of quantizing a point estimate and entropy coding it. This idea was first proposed by Havasi et al. [23] for model compression, who introduced minimal random coding (MRC) to encode a weight sample. In our paper, we use depth-limited, global-bound A* coding instead, to which we refer as A* coding hereafter for brevity's sake [29, 20]. We present it in Appendix A for completeness. A* coding is an importance sampling algorithm that draws 2 N ' X 2 D KL r q w } p w s` t \ independent samples w 1 , . . . , w N from the prior p w for some parameter t ě 0 , and computes their importance weights r i ' log ` q w p w i q L p w p w i q ˘ . Then, in a similar fashion to the Gumbel-max trick [31], it randomly perturbs the importance weights and selects the sample with the greatest perturbed weight. Unfortunately, this procedure returns an approximate sample with distribution ˜ q w . However, Theis and Yosri [32] have shown that the total variation distance ∥ q w ´ ˜ q w ∥ TV vanishes exponentially quickly as t Ñ 8 . Thus, t can be thought of as a free parameter of the algorithm that trades off compression rate for sample quality. Furthermore, A* coding is more efficient than MRC [23] in the following sense: Let N MRC and N A ˚ be the codes returned by MRC and A* coding, respectively, when given the same target and proposal distribution as input. Then, H r N A ˚ s ď H r N MRC s , hence using A* coding is always strictly more efficient [32].

3 Compression with Bayesian Implicit Neural Representations

Wenowintroduce our method, dubbed Com pression with B ayesian I mplicit Ne ural R epresentations (COMBINER). It extends INRs to the variational Bayesian setting by introducing a variational posterior q w over the network weights and fits INRs to the data D by minimizing Equation (2). Since

Algorithm 1 Learning the model prior

encoding the model weights is equivalent to compressing the data D , Equation (2) corresponds to jointly optimizing a given rate-distortion trade-off for the data. This is COMBINER's main advantage over other INR-based compression methods, which optimize the distortion only while keeping the rate fixed and cannot jointly optimize the rate-distortion. Moreover, another important difference is that we encode a random weight sample w ' q w from the weight posterior using A* coding [20] instead of quantizing the weights and entropy coding them. At a high level, COMBINER applies the model compression approach proposed by Havasi et al. [23] to encode variational Bayesian INRs, albeit with significant improvements which we discuss in Sections 3.1 and 3.2.

In this paper, we only consider networks with a diagonal Gaussian prior p w ' N p µ p , diag p σ p qq and posterior q w ' N p µ q , diag p σ q qq for mean and variance vectors µ p , µ q , σ p , σ q . Here, diag p v q denotes a diagonal matrix with v on the main diagonal. Following Havasi et al. [23], we optimize the variational parameters µ q and σ q using the local reparameterization trick [33] and, in Section 3.1, we derive an iterative algorithm to learn the prior parameters µ p and σ p .

3.1 Learning the Model Prior on the Training Set

To guarantee that COMBINER performs well in practice, it is critical that we find a good prior p w over the network weights, since it serves as the proposal distribution for A* coding and thus directly impacts the method's coding efficiency. To this end, in Algorithm 1 we describe an iterative algorithm to learn the prior parameters θ p ' t µ p , σ p u that minimize the average rate-distortion objective over some training data t D 1 , . . . , D M u :

In Equation (3) we write q p i q w ' N p µ p i q q , diag p σ p i q q qq , and p w ; θ p ' N p µ p , diag p σ p qq , explicitly denoting the prior's dependence on its parameters. Now, we propose a coordinate descent algorithm to minimize the objective in Equation (3), shown in Algorithm 1. To begin, we randomly initialize the model prior and the posteriors, and alternate the following two steps to optimize t q p i q w u and θ p :

Optimize the variational posteriors: We fix the prior parameters θ p and optimize the posteriors using the local reparameterization trick [33] with gradient descent. Note that, given θ p , optimizing s L β p θ p , t q p i q w uq can be split into M independent optimization problems, which we can perform in parallel:

Updating prior: We fix the posteriors t q p i q w u and update the model prior by computing θ p ' arg min θ s L β p θ p , t q p i q w uq . In the Gaussian case, this admits a closed-form solution:

We provide the full derivation of this procedure in Appendix B. Note that by the definition of coordinate descent, the value of s L β p θ p , t q p i q w uq decreases after each iteration, which ensures that our estimate of θ p converges to some optimum.

3.2 Compression with Posterior Refinement

Once the model prior is obtained using Algorithm 1, the sender uses the prior to train the variational posterior distribution for a specific test datum, as illustrated by Equation (2). To further improve the performance of INR compression, we also adopt a progressive posterior refinement strategy, a concept originally proposed in [23] for Bayesian model compression.

To motivate this strategy, we first consider the optimal weight posterior q ˚ w . Fixing the data D , trade-off parameter β and weight prior p w , q ˚ w is given by q ˚ w ' arg min q L β p D , q, p w q , where the minimization is performed over the set of all possible target distributions q . To compress D using our Bayesian INR, ideally we would like to encode a sample w ˚ ' q ˚ w , as it achieves optimal performance on average by definition. Unfortunately, finding q ˚ w is intractable in general, hence we restrict the search over the set of all factorized Gaussian distributions in practice, which yields a rather crude approximation. However, note that for compression, we only care about encoding a single, good quality sample using relative entropy coding. To achieve this, Havasi et al. [23] suggest partitioning the weight vector w into K blocks w 1: K ' t w 1 , . . . , w K u . For example, we might partition the weights per MLP layer with w i representing the weights on layer i , or into a preset number of random blocks; at the extremes, we could partition w per dimension, or we could just set K ' 1 for the trivial partition. Now, to obtain a good quality posterior sample given a partition w 1: K , we start with our crude posterior approximation and obtain

where each of the K minimization procedures takes place over the appropriate family of factorized Gaussian distributions. Then, we draw a sample w 1 ' q w 1 and refine the remaining approximation:

where L β p¨ | w 1 q indicates that w 1 is fixed during the optimization. We now draw w 2 ' q w 2 | w 1 to obtain the second chunk of our final sample. In total, we iterate the refinement procedure K times, progressively conditioning on more blocks, until we obtain our final sample w ' w 1: K . Note that already after the first step, the approximation becomes conditionally factorized Gaussian , which makes it far more flexible, and thus it approximates q ˚ w much better [18].

Combining the refinement procedure with compression: Above, we assumed that after each refinement step k , we draw the next weight block w k ' q w k | w 1: k ´ 1 . However, as suggested in [23], we can also extend the scheme to incorporate relative entropy coding, by encoding an approximate sample ˜ w k ' ˜ q w k | ˜ w 1: k ´ 1 with A* coding instead. This way, we actually feed two birds with one scone: the refinement process allows us to obtain a better overall approximate sample ˜ w by extending the variational family and by correcting for the occasional bad quality chunk ˜ w k at the same time, thus making COMBINER more robust in practice.

3.3 COMBINER in Practice

Given a partition w 1: K of the weight vector w , we use A* coding to encode a sample ˜ w k from each block. Let δ k ' D KL r q w k | ˜ w 1: k ´ 1 } p w k s represent the KL divergence in block k after the completion of the first k ´ 1 refinement steps, where we have already simulated and encoded samples from the first k ´ 1 blocks. As we discussed in Section 2, we need to simulate X 2 δ k ` t \ samples from the prior p w k to ensure that the sample ˜ w k encoded by A* coding has low bias. Therefore, for our method to be computationally tractable, it is important to ensure that there is no block with large divergence δ k . In fact, to guarantee that COMBINER's runtime is consistent, we would like the divergences across all blocks to be approximately equal, i.e., δ i « δ j for 0 ď i, j ď K . To this end, we set a bit-budget of κ bits per block and below we describe the to techniques we used to ensure δ k « κ for all k ' 1 , . . . , K . Unless stated otherwise, we set κ ' 16 bits and t ' 0 in our experiments.

First, we describe how we partition the weight vector based on the training data, to approximately enforce our budget on average. Note that we control COMBINER's rate-distortion trade-off by

varying β in its training loss in Equation (3). Thus, when we run Algorithm 1 to learn the prior, we also estimate the expected coding cost of the data given β as c β ' 1 M ř M i ' 1 D KL r q p i q w } p w s . Then, we set the number of blocks as K β,κ ' r c β { κ s and we partition the weight vector such that the average divergence s δ k of each block estimated on the training data matches the coding budget, i.e., s δ k « κ bits. Unfortunately, allocating individual weights to the blocks under this constraint is equivalent to the NP-hard bin packing problem [34]. However, we found that randomly permuting the weights and greedily assigning them using the next-fit bin packing algorithm [35] worked well in practice.

Relative entropy coding-aware fine-tuning: Assume we now wish to compress some data D , and we already selected the desired rate-distortion trade-off β , ran the prior learning procedure, fixed a bit budget κ for each block and partitioned the weight vector using the procedure from the previous paragraph. Despite our effort to set the blocks so that the average divergence s δ k « κ in each block on the training data, if we optimized the variational posterior q w using L β p D , q w , p w q , it is unlikely that the actual divergences δ k would match κ in each block. Therefore, we adapt the optimization procedure from [23], and we use a modified objective for each of the k posterior refinement steps:

where λ k : K ' t λ k , . . . , λ K u are slack variables, which we dynamically adjust during optimization. Roughly speaking, at each optimization step, we compute each δ i and increase its penalty term λ i if it exceeds the coding budget (i.e., δ i ą κ q and decrease the penalty term otherwise. See Appendix D for the detailed algorithm.

The comprehensive COMBINER pipeline: We now provide a brief summary of the entire COMBINER compression pipeline. To begin, given a dataset t D 1 , . . . , D M u , we select an appropriate INR architecture, and run the prior learning procedure (Algorithm 1) with different settings for β to obtain priors for a range of rate-distortion trade-offs.

To compress a new data point D , we select a prior with the desired rate-distortion trade-off and pick a blockwise coding budget κ . Then, we partition the weight vector w based on κ , and finally, we run the relative entropy coding-aware fine-tuning procedure from above, using A* coding to compress the weight blocks between the refinement steps to obtain the compressed representation of D .

4 Related Work

Neural Compression: Despite their short history, neural image compression methods' ratedistortion performance rapidly surpassed traditional image compression standards [16, 7, 9]. The current state-of-the-art methods follow a variational autoencoder (V AE) framework [2], optimizing the rate-distortion loss jointly. More recently, VAEs were also successfully applied to compressing other data modalities, such video [36] or point clouds [37]. However, mainstream methods quantize the latent variables produced by the encoder for transmission. Since the gradient of quantization is zero almost everywhere, learning the VAE encoder with standard back-propagation is not possible [38]. A popular solution [22] is to use additive uniform noise during training to approximate the quantization error, but it suffers from a train-test mismatch [39]. Relative entropy coding (REC) algorithms [19] eliminate this mismatch, as they can directly encode samples from the V AEs' latent posterior. Moreover, they bring unique advantages to compression with additional constraints, such as lossy compression with realism constraints [40, 41] and differentially private compression [42].

Compressing with INRs: INRs are parametric functional representations of data that offer many benefits over conventional grid-based representations, such as compactness and memory-efficiency [43-45]. Recently, compression with INRs has emerged as a new paradigm for neural compression [10], effective in compressing images [46], climate data [11], videos [47] and 3D scenes [48]. Usually, obtaining the INRs involves overfitting a neural network to a new signal, which is computationally costly [49]. Therefore, to ease the computational burden, some works [11, 46, 12] employ meta-learning loops [50] that largely reduce the fitting times during encoding. However, due to the expensive nature of the meta-learning process, these methods need to crop the data into patches to make training with second-order gradients practical. The biggest difficulty the current INR-based methods face is that quantizing the INR weights and activations can significantly degrade their performance, due to the brittle nature of the heavily overfitted parameters. Our method solves this issue

COMBINER (ours)

MSCN [12]

COIN++ [11]

COIN [10]

CVPR2020 [4]

ICLR2018 [2]

BPG

JPEG2000

0

1

2

3

4

5

6

7

18

22

26

30

34

38

42

Rate (bpp)

RGB PSNR (dB)

(a) CIFAR-10 dataset

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

22

24

26

28

30

32

Rate (bpp)

(b) Kodak dataset

by fitting a variational posterior over the parameters, from which we can encode samples directly using REC, eliminating the mismatch caused by quantization. Concurrent to our work, Schwarz et al. [51] introduced a method to learn a better coding distribution for the INR weights using a VAE, in a similar vein to our prior learning method in Algorithm 1. Their method achieves impressive performance on image and audio compression tasks, but is significantly more complex than our method: they run an expensive meta-learning procedure to learn the backbone architecture for their INRs and train a VAE to encode the INRs, making the already long training phase even longer.

5 Experiments

To assess COMBINER's performance across different data regimes and modalities, we conducted experiments compressing images from the low-resolution CIFAR-10 dataset [24], the highresolution Kodak dataset [25], and compressing audio from the LibriSpeech dataset [26]; the experiments and their results are described in Sections 5.1 and 5.2. Furthermore, in Section 5.3, we present analysis and ablation studies on COMBINER's ability to adaptively activate or prune the INR parameters, the effectiveness of its posterior refinement procedure and on the time complexity of its encoding procedure.

5.1 Image Compression

Datasets: Weconducted our image compression experiments on the CIFAR-10 [24] and Kodak [25] datasets. For the CIFAR-10 dataset, which contains 32 ˆ 32 pixel images, we randomly selected 2048 images from the training set for learning the model prior, and evaluated our model on all 10,000 images in the test set. For the high-resolution image compression experiments we use 512 randomly cropped 768 ˆ 512 pixel patches from the CLIC training set [52] to learn the model prior and tested on the Kodak images, which have matching resolution.

Models: Following previous methods [10-12], we utilize SIREN [43] as the network architecture. Input coordinates x are transformed into Fourier embeddings [44] before being fed into the MLP network, depicted as γ p x q in Figure 1. For the model structure, we experimentally find a 4-layer MLP with 16 hidden units per layer and 32 Fourier embeddings works well on CIFAR-10. When training on CLIC and testing on Kodak, we use models of different sizes to cover multiple rate points. We describe the model structure and other experimental settings in more detail in Appendix E. Remarkably, the networks utilized in our experiments are quite small. Our model for compressing CIFAR-10 images has only 1,123 parameters, and the larger model for compressing high-resolution Kodak images contains merely 21,563 parameters.

Performance: In Figure 2, we benchmark COMBINER's rate-distortion performance against classical codecs including JPEG2000 and BPG, and INR-based codecs including COIN [10], COIN++

10

20

30

38

40

42

44

46

48

Rate (kbps)

RGB PSNR (dB)

COMBINER

MP3

1

2

3

4

20

25

30

35

Rate (bpp)

RGB PSNR (dB)

COMBINER

Theoretically optimal

No fine-tuning

No prior learning

0

10k

20k

30k

28

.

6

28

.

7

28

.

8

28

.

9

Number of iterations

RGB PSNR (dB)

COMBINER

[11], and MSCN [12]. Additionally, we include results from VAE-based codecs such as ICLR2018 [2] and CVPR2020 [4] for reference. We observe that COMBINER exhibits competitive performance on the CIFAR-10 dataset, on par with COIN++ and marginally lower than MSCN. Furthermore, our proposed method achieves impressive performance on the Kodak dataset, surpassing JPEG2000 and other INR-based codecs. This superior performance is in part due to our method not requiring an expensive meta-learning loop [11, 46, 12], which would involve computing secondorder gradients during training. Since we avoid this cost, we can compress the whole high-resolution image using a single MLP network, thus the model can capture global patterns in the image.

5.2 Audio Compression

To demonstrate the effectiveness of COMBINER for compressing data in other modalities, we also conduct experiments on audio data. Since our method does not need to compute the second-order gradient during training, we can directly compress a long audio segment with a single INR model. We evaluate our method on LibriSpeech [26], a speech dataset recorded at a 16kHz sampling rate. Wetrain the model prior with 3-second chunks of audio, with 48000 samples per chunk. The detailed experimental setup is described in Appendix E. Due to COMBINER's time-consuming encoding process, we restrict our evaluation to 24 randomly selected audio chunks from the test set. Since we lack COIN++ statistics for this subset of 24 audio chunks, we only compare our method with MP3 (implemented using the ffmpeg package), which has been shown to be much better than COIN++ on the complete test set [11]. Figure 3 shows that COMBINER outperforms MP3 at low bitrate points, which verifies its effectiveness in audio compression. We also conducted another group of experiments where the audios are cropped into shorter chunks, which we describe in Appendix F.2.

5.3 Analysis, Ablation Study and Time Complexity

Model Visualizations: To provide more insight into COMBINER's behavior, we visualize its parameters and information content on the second hidden layer of two small 4-layer models trained on two CIFAR-10 images with β ' 10 ´ 5 . We use the KL in bits as an estimate of their coding cost, and do not encode the weights with A* coding or perform fine-tuning.

In Figure 6, we visualize the learned model prior parameters µ p and σ p in the left column, the variational parameters of two distinct images in the second and third column and the KL divergence D KL r q w } p w s in bits in the rightmost column. Since this layer incorporates 16 hidden units, the weight matrix of parameters has a 17 ˆ 16 shape, where weights and bias are concatenated (the bias is represented by the last row). Interestingly, there are seven 'active' columns within σ p , indicating that only seven hidden units of this layer would be activated for signal representation at this rate point. For instance, when representing image 1 that is randomly selected from the CIFAR10 test set, four columns are activated for representation. This activation is evident in the four blue columns within the KL map, which require a few bits to transmit the sample of the posterior distribution. Similarly, three hidden units are engaged in the representation of image 2. As their variational Gaussian distributions have close to zero variance, the posterior distributions at these

Prior mean

Posterior mean, image 1

Posterior mean, image 2

KL

/ bits, image 1

Prior s.d.

Posterior s.d., image 1

Posterior s.d., image 2

KL

/ bits, image 2

0.030

0.030

0.030

3.0

2.5

0.015

0.015

0.015

2.0

0.000

0.000

0.000

1.5

1.0

0.015

0.015

0.015

0.5

-0.030

-0.030

0.030

0.0

0.016

0.012

0.012

3.0

2.5

0.012

0.009

0.009

2.0

0.008

0.006

0.006

1.5

1.0

0.004

0.003

0.003

0.5

0.000

0.000

0.000

0.0

activated columns basically approach a Dirac delta distribution. In summary, by optimizing the rate-distortion objective, our proposed method can adaptively activate or prune network parameters.

Ablation Studies: We conducted ablation studies on the CIFAR-10 dataset to verify the effectiveness of learning the model prior (Section 3.1) and posterior fine-tuning (Section 3.2). In the first ablation study, instead of learning the prior parameters, we follow the methodology of Havasi ([53], p. 73) and use a layer-wise zero-mean isotropic Gaussian prior p ℓ ' N p 0 , σ ℓ I q , where p ℓ is the weight prior for the ℓ th hidden layer. We learn the σ ℓ 's jointly with the posterior parameters by optimizing Equation (3) using gradient descent, and encode them at 32-bit precision alongside the A*-coded posterior weight samples. In the second ablation study, we omit the fine-tuning steps between encoding blocks with A* coding, i.e. we never correct for bad quality approximate samples. In both experiments, we compress each block using 16 bits. Finally, as a reference, we also compare with the theoretically optimal scenario: we draw an exact sample from each blocks's variational posterior between refinement steps instead of encoding an approximate sample with A* coding, and estimate the sample's codelength with the block's KL divergence.

We compare the results of these experiments with our proposed pipeline (Section 3.3) using the above mentioned techniques in Figure 4. We find that both the prior learning and posterior refinement contribute significantly to COMBINER's performance. In particular, fine-tuning the posteriors is more effective at higher bitrates, while prior learning increases yields a consistent 4dB in gain in PSNR across all bitrates. Finally, we see that fine-tuning cannot completely compensate for the occasional bad approximate samples that A* coding yields, as there is a consistent 0.8 - 1.3dB discrepancy between COMBINER's and the theoretically optimal performance.

In Appendix C, we describe a further experiments we conducted to estimate how much each finetuning step contributes to the PSNR gain between compressing two blocks. The results are shown in Figure 7, which demonstrate that quality of the encoded approximate posterior sample doesn't just monotonically increase with each fine-tuning step, see Appendix C for an explanation.

Time Complexity: COMBINER's encoding procedure is slow, as it requires several thousand gradient descent steps to infer the parameters of the INR's weight posterior, and thousands more for the progressive fine-tuning. To get a better understanding of COMBINER's practical time complexity, we evaluate its coding time on both the CIFAR-10 and Kodak datasets at different rates and report our findings in Tables 1 and 2. We find that it can take between 13 minutes (0.91 bpp) to 34 minutes (4.45 bpp) to encode 500 CIFAR-10 images in parallel with a single A100 GPU, including posterior inference (7 minutes) and progressive fine-tuning. Note, that the fine-tuning takes longer for higher bitrates, as the weights are partitioned into more groups as each weight has higher individual information content. To compress high-resolution images from the Kodak dataset, the encoding time varies between 21.5 minutes (0.070 bpp) and 79 minutes (0.293 bpp).

To assess the effect of the fine-tuning procedure's length, we randomly selected a CIFAR-10 image and encoded it using the whole COMBINER pipeline, but varied the number of fine-tuning steps between 2148 and 30260; we report the results of our experiment in Figure 5. We find that running the fine-tuning process beyond a certain point has diminshing returns. In particular, while we used around 30k iterations in our other experiments, just using 3k iterations would sacrifice a mere 0.3 dB in the reconstruction quality, while saving 90% on the original tuning time.

On the other hand, COMBINER has fast decoding speed, since once we decode the compressed weight sample, we can reconstruct the data with a single forward pass through the network at each coordinate, which can be easily parallelized. Specifically, the decoding time of a single CIFAR-10 image is between 2 ms and 4 ms using an A100 GPU, and less than 1 second for a Kodak image.

6 Conclusion and Limitations

In this paper, we proposed COMBINER, a new neural compression approach that first encodes data as variational Bayesian implicit neural representations and then communicates an approximate posterior weight sample using relative entropy coding. Unlike previous INR-based neural codecs, COMBINER supports joint rate-distortion optimization and thus can adaptively activate and prune the network parameters. Moreover, we introduced an iterative algorithm for learning the prior parameters on the network weights and progressively refining the variational posterior. Our ablation studies show that these methods significantly enhance the COMBINER's rate-distortion performance. Finally, COMBINER achieves strong compression performance on low and high-resolution image and audio compression, showcasing its potential across different data regimes and modalities.

COMBINER has several limitations. First, as discussed in Section 5.3, while its decoding process is fast, its encoding time is considerably longer. Optimizing the variational posterior distributions requires thousands of iterations, and progressively fine-tuning them is also time-consuming. Second, Bayesian neural networks are inherently sensitive to initialization [21]. Identifying the optimal initialization setting for achieving training stability and superior rate-distortion performance may require considerable effort. Despite these challenges, we believe COMBINER paves the way for joint rate-distortion optimization of INRs for compression.

References

Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Husz´ ar. Lossy image compression with compressive autoencoders. In International Conference on Learning Representations , 2017. 1

Johannes Ball´ e, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. Variational image compression with a scale hyperprior. In International Conference on Learning Representations , 2018. 2, 6, 7, 8

David Minnen, Johannes Ball´ e, and George Toderici. Joint autoregressive and hierarchical priors for learned image compression. In Advances in Neural Information Processing Systems , 2018. 2

Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Learned image compression with discretized Gaussian mixture likelihoods and attention modules. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2020. 1, 7, 8

Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations , 2014. 1

Dailan He, Ziming Yang, Weikun Peng, Rui Ma, Hongwei Qin, and Yan Wang. ELIC: efficient learned image compression with unevenly grouped space-channel contextual adaptive coding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2022. 1

Jinming Liu, Heming Sun, and Jiro Katto. Learned image compression with mixed transformer-CNN architectures. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2023. 1, 6

Eirikur Agustsson, Michael Tschannen, Fabian Mentzer, Radu Timofte, and Luc Van Gool. Generative adversarial networks for extreme learned image compression. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 221-231, 2019. 1

Fabian Mentzer, George Toderici, Michael Tschannen, and Eirikur Agustsson. High-fidelity generative image compression. In Advances in Neural Information Processing Systems , volume 33, pages 1191311924, 2020. 1, 6

Emilien Dupont, Adam Golinski, Milad Alizadeh, Yee Whye Teh, and Arnaud Doucet. COIN: Compression with implicit neural representations. In Neural Compression: From Information Theory to Applications-Workshop@ ICLR , 2021. 1, 3, 6, 7

Emilien Dupont, Hrushikesh Loya, Milad Alizadeh, Adam Golinski, Y Whye Teh, and Arnaud Doucet. COIN++: Neural compression across modalities. Transactions on Machine Learning Research , 2022(11), 2022. 2, 6, 7, 8, 15

Jonathan Richard Schwarz and Yee Whye Teh. Meta-learning sparse compression networks. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. 1, 2, 6, 7, 8, 15

Kenneth O Stanley. Compositional pattern producing networks: A novel abstraction of development. Genetic programming and evolvable machines , 8(2):131-162, 2007. 1, 3

David Minnen and Saurabh Singh. Channel-wise autoregressive entropy models for learned image compression. In 2020 IEEE International Conference on Image Processing (ICIP) . IEEE, 2020. 2

Dailan He, Yaoyan Zheng, Baocheng Sun, Yan Wang, and Hongwei Qin. Checkerboard context model for efficient learned image compression. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2021.

Zongyu Guo, Zhizheng Zhang, Runsen Feng, and Zhibo Chen. Causal contextual prediction for learned image compression. IEEE Transactions on Circuits and Systems for Video Technology , 32(4), 2021. 6

Ahmet Burakhan Koyuncu, Han Gao, Atanas Boev, Georgii Gaikov, Elena Alshina, and Eckehard G. Steinbach. Contextformer: A transformer with spatio-channel attention for context modeling in learned image compression. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XIX . Springer, 2022. 2

Marton Havasi, Jasper Snoek, Dustin Tran, Jonathan Gordon, and Jos´ e Miguel Hern´ andez-Lobato. Refining the variational posterior through iterative optimization. In Bayesian Deep Learning Workshop @ NeurIPS , 2019. 2, 5

Gergely Flamich, Marton Havasi, and Jos´ e Miguel Hern´ andez-Lobato. Compressing images by encoding their latent representations with relative entropy coding. In Advances in Neural Information Processing Systems , volume 33, pages 16131-16141, 2020. 6

Gergely Flamich, Stratis Markou, and Jos´ e Miguel Hern´ andez-Lobato. Fast relative entropy coding with A* coding. In International Conference on Machine Learning , pages 6548-6577. PMLR, 2022. 2, 3, 4, 14

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In International Conference on Machine Learning , pages 1613-1622. PMLR, 2015. 2, 3, 10, 17

Johannes Ball´ e, Valero Laparra, and Eero P. Simoncelli. End-to-end optimized image compression. In International Conference on Learning Representations , 2017. 2, 6

Marton Havasi, Robert Peharz, and Jos´ e Miguel Hern´ andez-Lobato. Minimal random code learning: Getting bits back from compressed model parameters. In International Conference on Learning Representations , 2019. 2, 3, 4, 5, 6, 16

Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, Canadian Institute for Advanced Research, 2009. 2, 7

Eastman Kodak. Kodak Lossless True Color Image Suite (PhotoCD PCD0992). http://r0k.us/ graphics/kodak/ , 1993. 2, 7

Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: An ASR corpus based on public domain audio books. In ICASSP , pages 5206-5210. IEEE, 2015. 2, 7, 8, 17

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989. 3

Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the sixth annual conference on Computational learning theory , pages 5-13, 1993. 3

Chris J. Maddison, Daniel Tarlow, and Tom Minka. A* sampling. In Advances in Neural Information Processing Systems , volume 27, 2014. 3

Il'ya Meerovich Sobol'. On the distribution of points in a cube and the approximate evaluation of integrals. Zhurnal Vychislitel'noi Matematiki i Matematicheskoi Fiziki , 7(4):784-802, 1967. 3

George Papandreou and Alan L Yuille. Perturb-and-map random fields: Using discrete optimization to learn and sample from energy models. In 2011 International Conference on Computer Vision , pages 193-200. IEEE, 2011. 3

Lucas Theis and Noureldin A. Yosri. Algorithms for the communication of samples. In International Conference on Machine Learning , pages 21308-21328. PMLR, 2022. 3, 14

Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. In Advances in Neural Information Processing Systems , volume 28, 2015. 4, 16

Silvano Martello and Paolo Toth. Bin-packing problem. Knapsack problems: Algorithms and computer implementations , pages 221-245, 1990. 6

David S Johnson. Near-optimal bin packing algorithms . PhD thesis, Massachusetts Institute of Technology, 1973. 6

Guo Lu, Wanli Ouyang, Dong Xu, Xiaoyun Zhang, Chunlei Cai, and Zhiyong Gao. DVC: an end-to-end deep video compression framework. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11006-11015, 2019. 6

Yun He, Xinlin Ren, Danhang Tang, Yinda Zhang, Xiangyang Xue, and Yanwei Fu. Density-preserving deep point cloud compression. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2333-2342, 2022. 6

Zongyu Guo, Zhizheng Zhang, Runsen Feng, and Zhibo Chen. Soft then hard: Rethinking the quantization in neural image compression. In International Conference on Machine Learning . PMLR, 2021. 6

Eirikur Agustsson and Lucas Theis. Universally quantized neural compression. In Advances in Neural Information Processing Systems , volume 33, 2020. 6

Lucas Theis and Eirikur Agustsson. On the advantages of stochastic encoders. In Neural Compression Workshop at ICLR , 2021. 6

Lucas Theis, Tim Salimans, Matthew D. Hoffman, and Fabian Mentzer. Lossy compression with Gaussian diffusion. arXiv preprint arXiv:2206.08889 , 2022. 6

Abhin Shah, Wei-Ning Chen, Johannes Balle, Peter Kairouz, and Lucas Theis. Optimal compression of locally differentially private mechanisms. In International Conference on Artificial Intelligence and Statistics , 2022. 6

Vincent Sitzmann, Julien N. P. Martel, Alexander W. Bergman, David B. Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. In Advances in Neural Information Processing Systems , volume 33, pages 7462-7473, 2020. 6, 7

Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. In Advances in Neural Information Processing Systems , volume 33, pages 7537-7547, 2020. 7

Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. NeRF: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM , 65(1):99-106, 2021. 6

Yannick Str¨ umpler, Janis Postels, Ren Yang, Luc Van Gool, and Federico Tombari. Implicit neural representations for image compression. In ECCV (26) , volume 13686 of Lecture Notes in Computer Science , pages 74-91. Springer, 2022. 6, 8

Hao Chen, Bo He, Hanyu Wang, Yixuan Ren, Ser-Nam Lim, and Abhinav Shrivastava. NeRV: Neural representations for videos. In Advances in Neural Information Processing Systems , volume 34, pages 21557-21568, 2021. 6

Thomas Bird, Johannes Ball´ e, Saurabh Singh, and Philip A. Chou. 3D scene compression through entropy penalized neural representation functions. In 2021 Picture Coding Symposium (PCS) , pages 1-5. IEEE, 2021. 6

Matthew Tancik, Ben Mildenhall, Terrance Wang, Divi Schmidt, Pratul P. Srinivasan, Jonathan T. Barron, and Ren Ng. Learned initializations for optimizing coordinate-based neural representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2846-2855, 2021. 6

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning , pages 1126-1135. PMLR, 2017. 6

Jonathan Richard Schwarz, Jihoon Tack, Yee Whye Teh, Jaeho Lee, and Jinwoo Shin. Modality-agnostic variational compression of implicit neural representations. In International Conference on Machine Learning , 2023. 7

5th Workshop and Challenge on Learned Image Compression. http://compression.cc , 2022. 7, 17

Marton Havasi. Advances in compression using probabilistic models . PhD thesis, University of Cambridge, 2021. 9

A Relative Entropy Coding with A* Coding

Algorithm 2 A* encoding

Require: Proposal distribution p w and target distribution q w . Initialize : N,G 0 , w ˚ , N ˚ , L Ð 2 | C | , 8 , K , K , ´8 for i ' 1 , . . . , N do Ź N samples from proposal distribution w i ' p w G i ' TruncGumbel p G i ´ 1 q L i Ð G i ` log p q w p w i q{ p w p w i qq Ź Perturbed importance weight if L i ď L then L Ð L i w ˚ , N ˚ Ð w i , i end if end for return w ˚ , N ˚ Ź Transmit the index N ˚

Recall that we would like to communicate a sample from the variational posterior distribution q w using the proposal distribution p w . In our experiments, we used global-bound depth-limited A* coding to achieve this [20]. We describe the encoding procedure in Algorithm 2 and the decoding procedure in Algorithm 3. For brevity, we refer to this particular variant of the algorithm as A* coding for the rest of the appendix.

A* coding is an importance sampler that draws N samples w 1 , . . . , w N ' p w from the proposal distribution p w , where N is a parameter we pick. Then, it computes the importance weights r p w n q ' q w p w n q{ p w p w n q , and sequentially perturbs them with truncated Gumbel 3 noise:

Then, it can be shown that by setting where

This result essentially tells us that we should draw at least around 2 D KL r q w } p w s samples to ensure low sample bias, and beyond this, the bias decreases exponentially quickly as t Ñ 8 . However,

we have that w N ˚ ' ˜ q w is approximately distributed according to the target, i.e. ˜ q w « q w . More preciesly, we have the following result:

Lemma A.1 (Bound on the total variation between ˜ q w and q w (Lemma D.1 in [32])) . Let us set the number of proposal samples simulated by Algorithm 2 to N ' 2 D KL r q w } p w s` t for some parameter t ě 0 . Let ˜ q w denote the approximate distribution of the algorithm's output for this choice of N . Then,

note that the number of samples we need also increases exponentially quickly with t . In practice, we observed that when D KL r q w } p w s is sufficiently large (around 16-20 bits), setting t ' 0 already gave good results. To encode N ˚ , we built an empirical distribution over indices using our training datasets and used it for entropy coding to find the optimal variable-length code for the index.

In short, on the encoder side, N random samples are obtained from the proposal distribution p w . Then we select the sample w i and transmit its index N ˚ that has the greatest perturbed importance weight. On the decoder side, those N random samples can be simulated with the same seed held by the encoder. The decoder only needs to find the sample with the index N ˚ . Therefore, the decoding process of our method is very fast.

B Closed-Form Solution for Updating Model Prior

In this section, we derive the analytic expressions for the prior parameter updates in our iterative prior learning procedure when both the prior and the posterior are Gaussian distributions. Given a set of training data t D i u ' t D 1 , D 2 , ..., D M u , we fit a variational distribution q p i q w to represent each of the D i s. To do this, we minimize the loss (abbreviated as L later)

Now calculate the derivative w.r.t. the prior distribution parameter p w ; θ p ,

Considering we choose factorized Gaussian as variational distributions, the KL divergence is

To compute the analytical solution, let

Note here σ refers to variance rather than standard deviation. The above equation is equivalent to

We finally can solve these equations and get

as the result of Equation (5) in our main text. In short, this closed-form solution provides an efficient way to update the model prior from a bunch of variational posteriors. It makes our method simple in practice, unlike some previous methods [11, 12] that require expensive meta-learning loops.

0

50

100

150

200

Number of Compressed Groups

30.0

30.5

31.0

31.5

32.0

32.5

33.0

33.5

34.0

PSNR (dB)

REC with fine-tuning

REC without fine-tuning

Fine-tuning

REC of Group 5

Fine-tuning

REC of Group 6

Fine-tuning

31.64

31.66

31.68

31.70

31.72

31.74

Fine-tuning

REC of Group 217

Fine-tuning

REC of Group 218

Fine-tuning

33.90

33.92

33.94

33.96

33.98

34.00

C The Approximated PSNR Changes As Fine-tuning Goes On

We compressed some of the parameters using A* coding, directly sampled the rest from the posterior distributions, and used their corresponding KL divergence to estimate the coding cost. At the same time, we can obtain the approximated PSNR value by using the posterior samples to estimate the decoding quality. As shown in Figure 7, the PSNR tends to increase as the fine-tuning process goes on. However, it tends to drop when the fine-tuning process is nearing completion. This phenomenon occurs because, at the initial fine-tuning stage, the fine-tuning gain is more than the loss from A* coding, as many uncompressed groups can be fine-tuned to correct the errors of A* coding. But when the fine-tuning process nears completion, there are fewer uncompressed groups which could compensate for the bad sample of A* coding. Therefore, the general PSNR curve tends to decrease when it approaches the end of fine-tuning. This figure shows that while A* coding's sample results may have a distance to the accurate posterior, our proposed progressive fine-tuning strategy effectively mitigates most of these discrepancies.

D Dynamic Adjustment of β

When learning the model prior, the value of β that controlling the rate-distortion trade-off is defined in advance to train the model prior at a specific bitrate point. After obtaining the model prior, we will first partition the network parameters into K groups w 1: K ' t w 1 , . . . , w K u according to the average approximate coding cost of training data, as described in Section 3.3. Now for training the variational posterior for a given test datum, to ensure the coding cost of each group is close to κ ' 16 bits, we adjust the value of β dynamically when optimizing the posteriors. The detailed algorithm is illustrated here in Algorithm 4.

The algorithm is improved from Havasi et al. [23] to stabilize training, in the way that we set an interval r κ ´ 0 . 4 , κ s as buffer area where we do not change the value of λ k . Here we only adjust λ k every 15 iterations to avoid frequent changes at the initial training stage.

E Experiment Details

We introduce the experimental settings here and summarize the settings in Table 3.

E.1 CIFAR-10

We use a 4-layer MLP with 16 hidden units and 32 Fourier embeddings for the CIFAR-10 dataset. The model prior is trained with 128 epochs to ensure convergence. Here, the term 'epoch' is used to refer to optimizing the posteriors and updating the prior in the Algorithm 1 in the main text. For each epoch, the posteriors of all 2048 training data are optimized for 100 iterations using the local reparameterization trick [33], except the first epoch that contains 250 iterations. We use the Adam optimizer with learning rate 0.0002. The posterior variances are initialized as 9 ˆ 10 ´ 6 .

After obtaining the model prior, given a specific test CIFAR-10 image to be compressed, the posterior of this image is optimized for 25000 iterations, with the same optimizer. When we finally progressively compress and fine-tune the posterior, the posteriors of the uncompressed parameter groups are fine-tuned for 15 iterations with the same optimizer once a previous group is compressed.

E.2 Kodak

For Kodak dataset, since training on high-resolution image takes much longer time, the model prior is learned using fewer training data, i.e., only 512 cropped CLIC images [52]. We also reduce the learning rate of the Adam optimizer to 0.0001 to stabilize training. In each epoch, the posterior of each image is trained for 200 iterations, except the first epoch that contains 500 iterations. We also reduce the total epoch number to 96 which is empirically enough to learn the model prior.

We use two models with different capacity for compressing high-resolution Kodak images. The smaller model is a 6-layer SIREN with 48 hidden units and 64 Fourier embeddings. This model is used to get the three low-bitrate points in Figure 2b in our main text, where the corresponding beta is set as t 10 ´ 7 , 10 ´ 8 , 4 ˆ 10 ´ 8 u . Another larger model comprises a 7-layer MLP with 56 hidden units and 96 Fourier embeddings, which is used for evaluation at the two relatively higher bitrate points in Figure 2b in our main text. The betas of these two models have the same value 2 ˆ 10 ´ 9 . We empirically adjust the variance initialization from the set t 4 ˆ 10 ´ 6 , 4 ˆ 10 ´ 10 u and find they can affect the converged bitrate and achieve good performance. In particular, the posterior variance is initialized as 4 ˆ 10 ´ 10 to reach the highest bitrate point in the rate-distortion curve. The posterior variance of other bitrate-points on Kodak dataset are all initialized as 4 ˆ 10 ´ 6 .

Important note: It required significant empirical effort to find the optimal parameter settings we described above, hence our note in the Conclusion and Limitations section that Bayesian neural networks are inherently sensitive to initialization [21].

E.3 LibriSpeech

We randomly crop 1024 audio samples from LibriSpeech 'train-clean-100' set [26] for learning the model prior and randomly crop 24 test samples from 'test-clean' set for evaluation. The model structure is the same as the small model used for compressing Kodak images. We evaluate on four bitrate points by setting β ' t 10 ´ 7 , 3 ˆ 10 ´ 8 , 10 ´ 8 , 10 ´ 9 u . There are 128 epochs, and each epoch has 100 iterations with learning rate as 0.0002. The first epoch has 250 iterations. In addition, the posterior variance is initialized as 4 ˆ 10 ´ 9 . The settings for optimizing and fine-tuning posterior of a test datum are the same as the experiments on Kodak dataset.

F Supplementary Experimental Results

F.1 Number of Training Samples

Since the model prior is learned from a few training data, the number of training data may influence the quality of the learned model prior. We train the model prior with a different number of training images from the CIFAR-10 training set and evaluate the performance on 100 randomly selected test images from the CIFAR-10 test set. Surprisingly, as shown in Figure 8, we found that even merely 16 training images can help to learn a good model prior. Considering the randomness of training and testing, the performance on this test subset is almost the same when the number of training data exceeds 16. This demonstrates that the model prior is quite robust and generalizes well to test data. In our final experiments, the number of training samples is set to 2048 (on CIFAR-10 dataset) to ensure the prior converges to a good optimum.

F.2 Compressing Audios with Small Chunks

The proposed approach does not need to compute the second-order gradient during training, which helps to learn the model prior of the entire datum. Hence, compression with a single Bayesian INR network helps to fully capture the global dependencies of a datum. That is the reason for our strong performance on Kodak and LibriSpeech datasets. Here, we also conduct a group of experiment to compare the influence of cropping audios into chunks. Unlike the experimental setting in our main text that compresses every 3-second audio ( 1 ˆ 48000 ) with a single MLP network, here we try to crop all the 24 audios into small chunks, each of the chunk has the shape of 1 ˆ 200 . Weuse the same

2 2

2 3

2 4

2 5

2 6

2 7

2 8

2 9

2 10

2 11

Number of training samples

26.5

27.0

27.5

28.0

28.5

PSNR(dB)

COMBINER

0

10

20

30

40

Kbps

28

32

36

40

44

48

PSNR(dB)

COMBINER - 1 x 48000

COMBINER - 1 x 200

MP3

network used for compressing CIFAR-10 images for our experiments here. As shown in Figure 9, if we do not compress the audio as an entire entity, the performance will drops for around 5 dB. It demonstrates the importance of compressing with a single MLP network to capture the inherent redundancies within the entire data.

F.3 Additional Figures

We provide some examples of the decoded Kodak images in Figure 10.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n## Summary\nThe authors propose an approximated correlation communication approach to compress Bayesian INF. Two practical considerations, such as prior fitting and posterior refinement are proposed. It achieves an R-D performance comparable to the SOTA INF method on image compression. Moreover, the audio compression is also considered.\n\n========================================\n## Post rebuttal\nThe authors have addressed my concern wrt to the experiments, in a quite unexpected way. I raise my rating to accept.\n\nSTRENGTHS:\n## Strength\n* The idea of using correlation communication to compress image INR is novel. The compression of bayesian INR is also novel and likely to have impact outside compression community. Furthermore, it also considers audio compression and achieves reasonable results.\n\nWEAKNESSES:\n## Weakness\n* To the best of my knowledge, there is no algorithm that communicates correlation using finite number of samples. The adaptative reject sampling [Harsha 2010, The Communication Complexity of Correlation] and Possion Functional Representation [Li 2018, Strong Functional Representation Lemma and Applications to Coding Theorems] both requires infinite number of samples in the worst case to communicate exact posterior. As [Theis 2021, Algorithms for the Communication of Samples] has shown in Fig. 2, the expected code length of relative entropy coding [Flamich 2020], even if converge expotentially fast, still have quite some distance to the true posterior. And the expected code length of [Flamich 2020] is also inferior to [Harsha 2010] [Li 2018], which is $D + \\log D + \\log\\log D +,... $ The [Flamich 2020] does not emphasis this gap, while it is supposed to negatively effect the performance of the proposed approach. And this inaccurate posterior gap also exists for A* coding [Flamich 2022]. Then my question is, how bad such posterior approximation it is? Or to say, if we assume that we can use infinite sample optimal PFR to achieve exact posterior, what is the R-D curve like?\n* Another weakness is that the advantage of the proposed approach over MSCN [Schwarz 2022, Meta-Learning Sparse Compression Networks] is unclear. From the perspective of R-D performance, the proposed approach is only comparable to MSCN on Kodak and is slightly outperformed by MSCN on CIFAR. However, the encoding time of the proposed approach seems to be quite long due to the progressive posterior refinement. As neither the author nor original MSCN paper report encoding time of MSCN, I can not evaluate the advantage of the proposed approach pver MSCN. Or to say, if the authors can show that the run time of propoased approach is significantly faster than MSCN, this paper is more acceptable. This is especially true as the authors claim that their approach is simpler (L54,66,253). They need evdience, such as faster runtime, to support their claim.\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nThe limtation is discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-24"}

comments: SUMMARY:
This paper proposes overfitting variational Bayesian neural networks to the data and compressing an approximate posterior weight sample using relative entropy coding, which enables direct optimization of the rate-distortion performance by minimizing the $\beta$-ELBO. Moreover, an iterative algorithm for learning prior weight distributions is introduced and a progressive refinement process for the variational posterior is employed for improved performance. Experiments were conducted on image and audio compression.

STRENGTHS:
- the work encodes data with variational Bayesian implicit neural representations, which enables direct optimization of the rate-distortion performance by minimizing the $\beta$-ELBO. 
- the work presents an iterative algorithm for learning prior weight distributions and a progressive refinement process for the variational posterior for improved performance.

WEAKNESSES:
The presentation can be improved, e.g.,
- $\theta_p$ indicates $\mu_p, \sigma_p$, which should be clearly described.
- Some typos, e.g., Line 183: represent -> represents; Line 190: the to -> the; Line 209: the its -> its; Line 296: subscript "2".
- Title of [5] is missed, "Auto-Encoding Variational Bayes". Similar problem in [35]. And [35] actually repeats [34].

QUESTIONS:
- It would be nice to compare the proposed method with the concurrent work [52[ published in ICML 2023.
- Will the encoding time be the obstacle for applying he porposed method to video compression or 3D data compression?
- Code availability helps reproducibility although enough implementation details are provided.

LIMITATIONS:
The authors adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper improves INR-based (image) compression by introducing majorly two techniques: 1) a relative entropy coding based model compressing framework as an alternative to commonly used quantization - entropy coding pipeline; 2) a semi-amortized approach to train the model prior which is similar to beta-VAE and enables RD tradeoff with relative entropy coding. With the proposed novel pipeline, one can adopt COIN-like INR compression with better RD performance because it gets rid of quantization errors.\n\nSTRENGTHS:\nThe paper is well-written with dedicated figures and easy to follow. Overall I like the story of adopting beta-VAE/REC in INR coding methods. It is intuitive that given a dataset, all data neural representations share some common knowledge, and a potential amortized approach tends to exist. Thus, a meta-model describing the prior model distribution is soundable. This re-connects compression with VAEs, though a different way from feature-based learned image compression methods. And this re-connection makes a finer rate control of INR more promising. I believe this work contributes to society and can inspire future studies.\n\nWEAKNESSES:\nMy major concern is about practicality. Seems that to train a model prior, we should first train many model posteriors. As also discussed by the authors, this training makes the entire training (encoding) time extremely slow. Is it possible to adopt this model posterior to further speed up the encoding of samples out of the training set? i.e. we may imagine efficiently finetuning a new INR (posterior) from the obtained prior. If this finetuning requires less time to convergence, the enlarged encoding time may be amortized.\n\nAnother issue for me is the somewhat marginal performance improvement. Is it worth costing such a large encoding time for the model prior, instead of simply training a COIN/COIN++ model longer to compensate for the quantization error? It would be better to report the converging speed of the models e.g. something similar to PSNR-iteration curves.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nThe authors have discussed the limitations, which is convincing.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of lossy data compression (evaluation is on images and audio) using implicit neural representations (INRs). In this approach, a neural network is designed that maps coordinates (e.g., x,y locations in an image or time for audio) to samples (RGB values or audio amplitude). Then the overfit network itself is the representation that is stored or transmitted since the data can be approximately recovered by running the network over appropriate coordinates.\n\nINRs for data compression are not new, but this paper includes two key novelties: (1) a variational Bayesian approach is adopted, which (2) allows the method to jointly optimize for rate (the entropy of the neural network parameters) and distortion (the quality of the reconstructed data). Previous INRs for compression used less powerful formulations and separately optimized for rate and distortion, which typically leads to worse overall performance (and the empirical evaluation in this paper demonstrates that here).\n\nWithin this setup, the authors describe an iterative algorithm for learning the prior over network weights, and they describe a progressive refinement method that improves performance by dividing the network parameters into blocks and conditionally coding each block given previous blocks.\n\nThe method is evaluated on images (CIFAR-10 and Kodak) and audio (LibriSpeech). In both cases, the method is shown to outperform previous INR-based models. For images, it does not outperform the best VAE-based compression methods.\n\nSTRENGTHS:\nOriginality: the problem is not new and jointly optimizing rate and distortion is quite common in neural compression in general, but I have not seen it done before for INRs. Nor have I seen an approach that uses variational Bayes and relative entropy coding with INRs.\n\nQuality and clarity are both excellent.\n\nI think the significance of this paper is quite high for the neural compression community and especially researchers looking at INRs. VAE-based (often called \"nonlinear transform coding\" or NTC) is the dominant approach right now, but INRs are quite interesting and there's a growing literature using such methods for compression. This work is significant because it does the right thing (joint optimization of rate and distortion), which has not been done before with INRs, and because it sets a new SOTA for compression with this approach.\n\nWEAKNESSES:\n1. Obviously the paper would be stronger if the empirical results were better. Focusing on Fig. 2b (RD curves on the Kodak image set), COMBINER trails CVPR2020 by more than 2dB at 0.2 bpp (a huge gap), and CVPR2020 is no longer a SOTA approach. That said, COMBINER is SOTA for INR-based methods (as far as I know), which is an important result, and I think research can get stuck in a local minimum if we reviewers require new approaches to outperform established ones too soon in the research cycle, e.g., Fig. 2b shows that COMBINER slightly outperforms the VAE-based method from ICLR2018.\n\n2. Neural model compression is a large subfield so some discussion or comparison with a model compression method (beyond simple weight quantization) seems like an important baseline that's missing.\n\nQUESTIONS:\nSection 5.1 (\"Models\") discusses the model architecture saying that it was empirically determined. How sensitive are the results to the base model architecture, and how easily can the encoding process prune (reduce parameter entropy to near zero) large parts of a model architecture that is much larger than the RD optimal model?\n\nLIMITATIONS:\nadequately addressed", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new method for compressing general signals, by using Variational Bayesian implicit neural representations. It proposes an algorithm for learning a prior distribution over the implicit representation weights, as well as a pipeline for inferring the posterior distribution corresponding to every given signal to be encoded. The authors employ several tricks (e.g. Relative Entropy Coding) borrowed from previous works, that are meant to increase compression efficiency.\n\nThe method is evaluated over two image datasets (namely the CIFAR10 dataset of very small images and the Kodak dataset containing 24 larger images) and one audio dataset, while comparing performance with several image compression methods and the classic MP3 audio compression method, respectively.\n\nSTRENGTHS:\nGenerally speaking, the paper is presented well - despite the fact that compression is not my main expertise, I was able to follow and understand most of it.\nThe combination of the different ideas used in this method is interesting and appealing.\nLimitations, as well as concurrent work, are discussed candidly.\n\nWEAKNESSES:\nBesides the method's complexity, which could (at least partially) be attributed to my lack of expertise in this field, I found two major flows:\n\nPerformance:\nThe performance curves presented in Fig. 2 and 5 in the paper do not indicate an advantage of the proposed method over many of the existing methods, in terms of the rate-distortion trade-off. In the case of audio (Fig. 5), some advantage is shown only over the classic MP3 compression, and only for part of the kbps range.\n\nSpeed/Complexity:\nThe encoding process of this method is long and expensive. Besides being a big limitation (as noted by the authors) and hence a weakness of the method compared to some of the other ones, I'm missing some detailed comparison of these aspects.\n\nA third, relatively more minor flow has to do with the experimental setup itself, as experiments were conducted over only two small image datasets and one audio dataset, and compared only with MP3 in the latter case.\n\nQUESTIONS:\nAs I'm not entirely familiar with this field, the authors are welcome to direct my attention to certain points they feel that I overlooked, and I'll gladly consider increasing my initial rating.\n\nLIMITATIONS:\nLimitations were candidly discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}