SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Relative entropy coding-aware fine-tuning: Assume we now wish to compress some data  $\mathcal{D}$ , and we already selected the desired rate-distortion trade-off  $\beta$ , ran the prior learning procedure, fixed a bit budget  $\kappa$  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  $\bar{\delta}_k \approx \kappa$  in each block on the training data, if we optimized the variational posterior  $q_{\mathbf{w}}$  using  $\mathcal{L}_{\beta}(\mathcal{D}, q_{\mathbf{w}}, p_{\mathbf{w}})$ , it is unlikely that the actual divergences  $\delta_k$  would match  $\kappa$  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:
$$
\mathcal {L} _ {\lambda_ {k: K}} \left(\mathcal {D}, q _ {\mathbf {w} | \tilde {\boldsymbol {w}} _ {1: k - 1}}, p _ {\mathbf {w}}\right) = \sum_ {(\boldsymbol {x}, \boldsymbol {y}) \in \mathcal {D}} \mathbb {E} _ {\boldsymbol {w} \sim q _ {\mathbf {w}}} [ \Delta (\boldsymbol {y}, f (\boldsymbol {x} \mid \boldsymbol {w}) ] + \sum_ {i = k} ^ {K} \lambda_ {i} \cdot \delta_ {i}, \tag {8}
$$
where  $\lambda_{k:K} = \{\lambda_k,\dots ,\lambda_K\}$  are slack variables, which we dynamically adjust during optimization. Roughly speaking, at each optimization step, we compute each  $\delta_{i}$  and increase its penalty term  $\lambda_{i}$  if it exceeds the coding budget (i.e.,  $\delta_i > \kappa$ ) 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  $\{\mathcal{D}_1,\dots ,\mathcal{D}_M\}$ , we select an appropriate INR architecture, and run the prior learning procedure (Algorithm 1) with different settings for  $\beta$  to obtain priors for a range of rate-distortion trade-offs.
To compress a new data point  $\mathcal{D}$ , we select a prior with the desired rate-distortion trade-off and pick a blockwise coding budget  $\kappa$ . Then, we partition the weight vector  $\mathbf{w}$  based on  $\kappa$ , and finally, we run the relative entropy coding-aware fine-tuning procedure from above, using  $\mathrm{A}^*$  coding to compress the weight blocks between the refinement steps to obtain the compressed representation of  $\mathcal{D}$ .
Neural Compression: Despite their short history, neural image compression methods' rate-distortion performance rapidly surpassed traditional image compression standards [16, 7, 9]. The current state-of-the-art methods follow a variational autoencoder (VAE) 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 VAEs' 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
Figure 2: Rate-distortion curves on two image datasets. In both figures, solid lines denote INR-based methods, dotted lines denote VAE-based methods and dashed lines denote classical methods. Examples of decoded Kodak images are provided in Appendix F.3
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.
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 high-resolution 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: We conducted our image compression experiments on the CIFAR-10 [24] and Kodak [25] datasets. For the CIFAR-10 dataset, which contains  $32 \times 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 \times 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  $\pmb{x}$  are transformed into Fourier embeddings [44] before being fed into the MLP network, depicted as  $\gamma(\pmb{x})$  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++
Figure 3: COMBINER's audio compression performance versus MP3 on the LibriSpeech dataset.
Figure 4: Ablation study on CIFAR-10 verifying the effectiveness of the fine-tuning and prior learning procedures.
Figure 5: COMBINER's performance improvement as a function of the number of finetuning steps.
[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 second-order 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  $16\mathrm{kHz}$  sampling rate. We train 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  $\mathrm{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  $\mathrm{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 audiios 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  $\beta = 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  $\mu_p$  and  $\sigma_p$  in the left column, the variational parameters of two distinct images in the second and third column and the KL divergence  $D_{\mathrm{KL}}[q_{\mathbf{w}}\| p_{\mathbf{w}}]$  in bits in the rightmost column. Since this layer incorporates 16 hidden units, the weight matrix of parameters has a  $17\times 16$  shape, where weights and bias are concatenated (the bias is represented by the last row). Interestingly, there are seven "active" columns within  $\sigma_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 CIFAR-10 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
Figure 6: Visualizations of the weight prior, posterior and information content of a variational INR trained on two CIFAR-10 images. We focus on the INR's weights connecting the first and second hidden layers. Each heatmap is  $17 \times 16$  because both layers have 16 hidden units and we concatenated the weights and the biases (last row). We write s.d. for standard deviation.
KL / bits, image 2
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_{\ell} = \mathcal{N}(0,\sigma_{\ell}I)$ , where  $p_{\ell}$  is the weight prior for the  $\ell$ th hidden layer. We learn the  $\sigma_{\ell}$ 's jointly with the posterior parameters by optimizing Equation (3) using gradient descent, and encode them at 32-bit precision alongside the  $\mathrm{A}^*$ -coded posterior weight samples. In the second ablation study, we omit the fine-tuning steps between encoding blocks with  $\mathrm{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 block's variational posterior between refinement steps instead of encoding an approximate sample with  $\mathrm{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  $\mathrm{A}^*$  coding yields, as there is a consistent  $0.8 - 1.3\mathrm{dB}$  discrepancy between COMBINER's and the theoretically optimal performance.
In Appendix C, we describe a further experiments we conducted to estimate how much each fine-tuning 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).
bit-rate	Encoding (500 images, GPU A100 80G)	Decoding (1 image, CPU)
Learning Posterior	REC + Fine-tuning	Total
0.91 bpp	~6 min	~13 min	2.06 ms
1.39 bpp	~9 min	~16 min	2.09 ms
2.28 bpp	~7 min	~14 min 30 s	~21 min 30 s	2.86 ms
3.50 bpp	~21 min 30 s	~28 min 30 s	3.82 ms
4.45 bpp	~27 min	~34 min	3.88 ms
Table 1: The encoding time and decoding time of COMBINER on CIFAR-10 dataset.
bit-rate	Encoding (1 image, GPU A100 80G)	Decoding (1 image, CPU)
Learning Posterior	REC + Fine-tuning	Total
0.07 bpp	~12 min 30 s	~21 min 30 s	348.42 ms
0.11 bpp	~9 min	~18 mins	~27 min	381.53 ms
0.13 bpp	~22 min	~31 min	405.38 ms
0.22 bpp	~50 min	~61 min	597.39 ms
0.29 bpp	~11 min	~68 min	~79 min	602.32 ms
Table 2: The encoding time and decoding time of COMBINER on Kodak dataset.
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 diminishing returns. In particular, while we used around  $30\mathrm{k}$  iterations in our other experiments, just using  $3\mathrm{k}$  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.
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.
ZG acknowledges funding from the Outstanding PhD Student Program at the University of Science and Technology of China. ZC is supported in part by National Natural Science Foundation of China under Grant U1908209, 62021001. GF acknowledges funding from DeepMind.
[39] Eirikur Agustsson and Lucas Theis. Universally quantized neural compression. In Advances in Neural Information Processing Systems, volume 33, 2020. 6
[40] Lucas Theis and Eirikur Agustsson. On the advantages of stochastic encoders. In Neural Compression Workshop at ICLR, 2021. 6
[41] Lucas Theis, Tim Salimans, Matthew D. Hoffman, and Fabian Mentzer. Lossy compression with Gaussian diffusion. arXiv preprint arXiv:2206.08889, 2022. 6
[42] 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
[43] 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
[44] 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
[45] 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
[46] Yannick Strümpler, 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
[47] 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
[48] Thomas Bird, Johannes Balle, 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
[49] 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
[50] 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
[51] 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
[52] 5th Workshop and Challenge on Learned Image Compression. http://compression.cc, 2022. 7, 17
[53] Marton Havasi. Advances in compression using probabilistic models. PhD thesis, University of Cambridge, 2021. 9
A Relative Entropy Coding with  $\mathbf{A}^*$  Coding
Algorithm 2 A* encoding
Require: Proposal distribution  $p_{\mathbf{w}}$  and target distribution  $q_{\mathbf{w}}$  Initialize:  $N,G_0,w^*,N^*,L\gets 2^{|C|},\infty ,\bot ,\bot , - \infty$  for  $i = 1,\dots ,N$  do  $\triangleright N$  samples from proposal distribution  $\boldsymbol {w}_i\sim p_{\mathbf{w}}$ $G_{i}\sim \mathrm{TruncGumbel}(G_{i - 1})$ $L_{i}\leftarrow G_{i} + \log (q_{\mathbf{w}}(\boldsymbol {w}_i) / p_{\mathbf{w}}(\boldsymbol {w}_i))$ $\triangleright$  Perturbed importance weight if  $L_{i}\leqslant L$  then  $L\gets L_i$ $\boldsymbol {w}^{*},N^{*}\gets \boldsymbol {w}_{i},i$  end if   
end for   
return  $\boldsymbol {w}^{*},N^{*}$ $\triangleright$  Transmit the index  $N^{*}$
Algorithm 3 A* decoding
Simulate  $\{\pmb {w}_i\} = \{\pmb {w}_1,\dots ,\pmb {w}_N\}$ $\triangleright$  Simulate  $N$  samples from  $p_{\mathbf{w}}$  with the shared seed   
Receive  $N^*$    
return  $\pmb {w}^{*}\gets \pmb{w}_{N^{*}}$ $\triangleright$  Receive the approximate posterior sample
Recall that we would like to communicate a sample from the variational posterior distribution  $q_{\mathbf{w}}$  using the proposal distribution  $p_{\mathbf{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  $\pmb{w}_1, \dots, \pmb{w}_N \sim p_{\mathbf{w}}$  from the proposal distribution  $p_{\mathbf{w}}$ , where  $N$  is a parameter we pick. Then, it computes the importance weights  $r(\pmb{w}_n) = q_{\mathbf{w}}(\pmb{w}_n) / p_{\mathbf{w}}(\pmb{w}_n)$ , and sequentially perturbs them with truncated Gumbel<sup>3</sup> noise:
$$
\tilde {r} _ {n} = r \left(\boldsymbol {w} _ {n}\right) + G _ {n}, \quad G _ {n} \sim \operatorname {T r u n c G u m b e l} \left(G _ {n - 1}\right), G _ {0} = \infty \tag {9}
$$
Then, it can be shown that by setting