SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zongyu Guo*
University of Science and
Technology of China
guozy@mail.ustc.edu.cn
Gergely Flamich*
gf332@cam.ac.uk
Jiajun He
jh2383@cam.ac.uk
Zhibo Chen
chenzhibo@ustc.edu.cn
Jose Miguel Hernandez-Lobato
jmh233@cam.ac.uk
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. Our code is available at https://github.com/cambridge-mlg/combiner.
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 (VAEs, [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.
A recent 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
Figure 1: Framework overview of COMBINER. It first encodes a datum  $\mathcal{D}$  into Bayesian implicit neural representations, as variational posterior distribution  $q_{\mathbf{w}}$ . Then an approximate posterior sample  $\boldsymbol{w}^*$  is communicated from the sender to the receiver using relative entropy coding.
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 VAE-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_{\mathbf{w}}$  over the weights  $\mathbf{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  $\mathbf{w} \sim q_{\mathbf{w}}$  from the posterior. The average coding cost of REC algorithms is approximately  $D_{\mathrm{KL}}[q_{\mathbf{w}} \| p_{\mathbf{w}}]$ , where  $p_{\mathbf{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  $\beta$ -ELBO [21], in a similar fashion to VAE-based methods [22, 2]. We dub our method Compression with Bayesian Implicit Neural Representations (COMBINER), and present a high-level description of it in Figure 1.
We propose 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.
We evaluate 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.
- 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 Compression with Bayesian Implicit Neural Representations (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  $\mathcal{D}$  as a point set, where each point corresponds to a coordinate-signal value pair  $(\mathbf{x},\mathbf{y})$ , and overfitting a small neural network  $f(\mathbf{x} \mid \mathbf{w})$ , usually a multilayer perceptron (MLP) parameterised by weights  $\mathbf{w}$ , which is then called the implicit neural representation (INR) of  $\mathcal{D}$ . Recently, Dupont et al. [10] popularised INRs for lossy data compression by noting that compressing the INR's weights  $\mathbf{w}$  amounts to compressing  $\mathcal{D}$ . However, their method has a crucial shortcoming: they assume a uniform coding distribution over  $\mathbf{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(\mathbf{w})$  (abbreviated as  $p_{\mathbf{w}}$ ) and a variational posterior  $q(\mathbf{w}|\mathcal{D})$  (abbreviated as  $q_{\mathbf{w}}$ ) over the weights, in which case their information content is given by the Kullback-Leibler (KL) divergence  $D_{\mathrm{KL}}[q_{\mathbf{w}}\| p_{\mathbf{w}}]$ , as shown in [28]. Therefore, for distortion measure  $\Delta$  and a coding budget of  $C$  bits, we can optimize the constrained objective
$$
\min  _ {q _ {\mathbf {w}}} \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}) ], \quad \text {s u b j e c t t o} D _ {\mathrm {K L}} \left[ q _ {\mathbf {w}} \| p _ {\mathbf {w}} \right] \leqslant C. \tag {1}
$$
In practice, we introduce a slack variable  $\beta$  and optimize the Lagrangian dual, which yields:
$$
\mathcal {L} _ {\beta} (\mathcal {D}, q _ {\mathbf {w}}, p _ {\mathbf {w}}) = \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}) ] + \beta \cdot D _ {\mathrm {K L}} [ q _ {\mathbf {w}} \| p _ {\mathbf {w}} ] + \text {c o n s t .}, \tag {2}
$$
with different settings of  $\beta$  corresponding to different coding budgets  $C$ . Thus, optimizing  $\mathcal{L}_{\beta}(\mathcal{D}, q_{\mathbf{w}}, p_{\mathbf{w}})$  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 \sim q_{\mathbf{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  $N = \left[2^{D_{\mathrm{KL}}[q_{\mathbf{w}} \| p_{\mathbf{w}}] + t}\right]$  independent samples  $\boldsymbol{w}_1, \ldots, \boldsymbol{w}_N$  from the prior  $p_{\mathbf{w}}$  for some parameter  $t \geqslant 0$ , and computes their importance weights  $r_i = \log \left(q_{\mathbf{w}}(\boldsymbol{w}_i) / p_{\mathbf{w}}(\boldsymbol{w}_i)\right)$ . 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  $\tilde{q}_{\mathbf{w}}$ . However, Theis and Yosri [32] have shown that the total variation distance  $\| q_{\mathbf{w}} - \tilde{q}_{\mathbf{w}} \|_{\mathrm{TV}}$  vanishes exponentially quickly as  $t \to \infty$ . 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_{\mathrm{MRC}}$  and  $N_{\mathrm{A*}}$  be the codes returned by MRC and A* coding, respectively, when given the same target and proposal distribution as input. Then,  $\mathbb{H}[N_{\mathrm{A*}}] \leqslant \mathbb{H}[N_{\mathrm{MRC}}}$ , hence using A* coding is always strictly more efficient [32].
3 Compression with Bayesian Implicit Neural Representations
We now introduce our method, dubbed Compression with Bayesian Implicit Neural Representations (COMBINER). It extends INRs to the variational Bayesian setting by introducing a variational posterior  $q_{\mathbf{w}}$  over the network weights and fits INRs to the data  $\mathcal{D}$  by minimizing Equation (2). Since
In practice, we use quasi-random number generation with multi-dimensional Sobol sequences [30] to simulate our random variables to ensure that they cover the sample space as evenly as possible.
Algorithm 1 Learning the model prior
Require: Training data  $\{\mathcal{D}_i\} = \{\mathcal{D}_1,\mathcal{D}_2,\dots,\mathcal{D}_M\}$  Initialize: The model posteriors  $q_{\mathbf{w}}^{(i)} = \mathcal{N}(\pmb {\mu}_i,\mathrm{diag}(\pmb {\sigma}_i))$  of every training datum  $\mathcal{D}_i$  Initialize: The model priors  $p_{\mathbf{w};\pmb{\theta}_p} = \mathcal{N}(\pmb {\mu}_p,\mathrm{diag}(\pmb {\sigma}_p))$  repeat until convergence for  $i\gets 1$  to  $M$  do  $\{q_{\mathbf{w}}^{(i)}\} \leftarrow \arg \min_{\{q_{\mathbf{w}}^{(i)}\}}\mathcal{L}(\pmb {\theta}_p,\{q_{\mathbf{w}}^{(i)}\})$  Gradient descent for optimizing posteriors end for  $\theta_p\gets \arg \min_{\theta_p}\mathcal{L}(\theta_p,\{q_{\mathbf{w}}^{(i)}\})$  Closed-form solution in Equation (5) end repeat Return  $p_{\mathbf{w};\pmb{\theta}_p} = \mathcal{N}(\pmb {\mu}_p,\mathrm{diag}(\pmb {\sigma}_p))$
encoding the model weights is equivalent to compressing the data  $\mathcal{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\sim q_{\mathrm{w}}$  from the weight posterior using  $\mathrm{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_{\mathbf{w}} = \mathcal{N}(\pmb{\mu}_p, \mathrm{diag}(\pmb{\sigma}_p))$  and posterior  $q_{\mathbf{w}} = \mathcal{N}(\pmb{\mu}_q, \mathrm{diag}(\pmb{\sigma}_q))$  for mean and variance vectors  $\pmb{\mu}_p, \pmb{\mu}_q, \pmb{\sigma}_p, \pmb{\sigma}_q$ . Here,  $\mathrm{diag}(\mathbf{v})$  denotes a diagonal matrix with  $\mathbf{v}$  on the main diagonal. Following Havasi et al. [23], we optimize the variational parameters  $\pmb{\mu}_q$  and  $\pmb{\sigma}_q$  using the local reparameterization trick [33] and, in Section 3.1, we derive an iterative algorithm to learn the prior parameters  $\pmb{\mu}_p$  and  $\pmb{\sigma}_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_{\mathbf{w}}$  over the network weights, since it serves as the proposal distribution for  $\mathrm{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  $\pmb{\theta}_{p} = \{\pmb{\mu}_{p}, \pmb{\sigma}_{p}\}$  that minimize the average rate-distortion objective over some training data  $\{\mathcal{D}_1, \dots, \mathcal{D}_M\}$ :
$$
\bar {\mathcal {L}} _ {\beta} \left(\boldsymbol {\theta} _ {p}, \left\{q _ {\mathbf {w}} ^ {(i)} \right\}\right) = \frac {1}{M} \sum_ {i = 1} ^ {M} \mathcal {L} _ {\beta} \left(\mathcal {D} _ {i}, q _ {\mathbf {w}} ^ {(i)}, p _ {\mathbf {w};} \boldsymbol {\theta} _ {p}\right). \tag {3}
$$
In Equation (3) we write  $q_{\mathbf{w}}^{(i)} = \mathcal{N}(\pmb{\mu}_q^{(i)}, \mathrm{diag}(\pmb{\sigma}_q^{(i)}))$ , and  $p_{\mathbf{w};\pmb{\theta}_p} = \mathcal{N}(\pmb{\mu}_p, \mathrm{diag}(\pmb{\sigma}_p))$ , 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  $\{q_{\mathbf{w}}^{(i)}\}$  and  $\pmb{\theta}_p$ :
1. Optimize the variational posteriors: We fix the prior parameters  $\theta_p$  and optimize the posteriors using the local reparameterization trick [33] with gradient descent. Note that, given  $\theta_p$ , optimizing  $\bar{\mathcal{L}}_{\beta}(\theta_p, \{q_{\mathbf{w}}^{(i)}\})$  can be split into  $M$  independent optimization problems, which we can perform in parallel:
$$
\text {f o r e a c h} i = 1, \dots , M: \quad q _ {\mathbf {w}} ^ {(i)} = \underset {q} {\arg \min } \mathcal {L} _ {\beta} \left(\mathcal {D} _ {i}, q, p _ {\mathbf {w}; \boldsymbol {\theta} _ {p}}\right). \tag {4}
$$
2. Updating prior: We fix the posteriors  $\{q_{\mathbf{w}}^{(i)}\}$  and update the model prior by computing  $\theta_p = \arg \min_\pmb{\theta} \bar{\mathcal{L}}_\beta(\pmb{\theta}_p, \{q_{\mathbf{w}}^{(i)}\})$ . In the Gaussian case, this admits a closed-form solution:
$$
\boldsymbol {\mu} _ {p} = \frac {1}{M} \sum_ {i = 1} ^ {M} \boldsymbol {\mu} _ {q} ^ {(i)}, \quad \boldsymbol {\sigma} _ {p} = \frac {1}{M} \sum_ {i = 1} ^ {M} \left[ \boldsymbol {\sigma} _ {q} ^ {(i)} + \left(\boldsymbol {\mu} _ {q} ^ {(i)} - \boldsymbol {\mu} _ {p}\right) ^ {2} \right]. \tag {5}
$$
We provide the full derivation of this procedure in Appendix B. Note that by the definition of coordinate descent, the value of  $\bar{\mathcal{L}}_{\beta}(\pmb{\theta}_p, \{q_{\mathbf{w}}^{(i)}\})$  decreases after each iteration, which ensures that our estimate of  $\pmb{\theta}_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_{\mathbf{w}}^{*}$ . Fixing the data  $\mathcal{D}$ , trade-off parameter  $\beta$  and weight prior  $p_{\mathbf{w}}$ ,  $q_{\mathbf{w}}^{*}$  is given by  $q_{\mathbf{w}}^{*} = \arg \min_q \mathcal{L}_\beta(\mathcal{D}, q, p_{\mathbf{w}})$ , where the minimization is performed over the set of all possible target distributions  $q$ . To compress  $\mathcal{D}$  using our Bayesian INR, ideally we would like to encode a sample  $\boldsymbol{w}^{*} \sim q_{\mathbf{w}}^{*}$ , as it achieves optimal performance on average by definition. Unfortunately, finding  $q_{\mathbf{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  $\mathbf{w}$  into  $K$  blocks  $\mathbf{w}_{1:K} = \{\mathbf{w}_1, \dots, \mathbf{w}_K\}$ . For example, we might partition the weights per MLP layer with  $\mathbf{w}_i$  representing the weights on layer  $i$ , or into a preset number of random blocks; at the extremes, we could partition  $\mathbf{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  $\mathbf{w}_{1:K}$ , we start with our crude posterior approximation and obtain
$$
q _ {\mathbf {w}} = q _ {\mathbf {w} _ {1}} \times \dots \times q _ {\mathbf {w} _ {K}} = \underset {q _ {1}, \dots , q _ {K}} {\arg \min } \mathcal {L} _ {\beta} (\mathcal {D}, q _ {1} \times \dots \times q _ {K}, p _ {\mathbf {w}}), \tag {6}
$$
where each of the  $K$  minimization procedures takes place over the appropriate family of factorized Gaussian distributions. Then, we draw a sample  $\boldsymbol{w}_1 \sim q_{\mathbf{w}_1}$  and refine the remaining approximation:
$$
q _ {\mathbf {w} \mid \boldsymbol {w} _ {1}} = q _ {\mathbf {w} _ {2} \mid \boldsymbol {w} _ {1}} \times \dots \times q _ {\mathbf {w} _ {K} \mid \boldsymbol {w} _ {1}} = \underset {q _ {2}, \dots , q _ {K}} {\arg \min } \mathcal {L} _ {\beta} (\mathcal {D}, q _ {2} \times \dots \times q _ {K}, p _ {\mathbf {w}} \mid \boldsymbol {w} _ {1}), \tag {7}
$$
where  $\mathcal{L}_{\beta}(\cdot \mid \boldsymbol{w}_1)$  indicates that  $\boldsymbol{w}_1$  is fixed during the optimization. We now draw  $\boldsymbol{w}_2 \sim q_{\mathbf{w}_2 \mid \boldsymbol{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  $\boldsymbol{w} = \boldsymbol{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_{\mathbf{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  $\boldsymbol{w}_k \sim q_{\mathbf{w}_k|\boldsymbol{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  $\tilde{\boldsymbol{w}}_k \sim \tilde{q}_{\mathbf{w}_k|\tilde{\boldsymbol{w}}_{1:k-1}}$  with  $\mathrm{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  $\tilde{\boldsymbol{w}}$  by extending the variational family and by correcting for the occasional bad quality chunk  $\tilde{\boldsymbol{w}}_k$  at the same time, thus making COMBINER more robust in practice.
3.3 COMBINER in Practice
Given a partition  $\mathbf{w}_{1:K}$  of the weight vector  $\mathbf{w}$ , we use  $\mathrm{A}^*$  coding to encode a sample  $\tilde{w}_k$  from each block. Let  $\delta_k = D_{\mathrm{KL}}[q_{\mathbf{w}_k}|\tilde{w}_{1:k-1}\|p_{\mathbf{w}_k}]$  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  $\left\lfloor 2^{\delta_k + t}\right\rfloor$  samples from the prior  $p_{\mathbf{w}_k}$  to ensure that the sample  $\tilde{w}_k$  encoded by  $\mathrm{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  $\delta_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.,  $\delta_i \approx \delta_j$  for  $0 \leqslant i,j \leqslant K$ . To this end, we set a bit-budget of  $\kappa$  bits per block and below we describe the to techniques we used to ensure  $\delta_k \approx \kappa$  for all  $k = 1,\dots,K$ . Unless stated otherwise, we set  $\kappa = 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  $\beta$  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  $\beta$  as  $c_{\beta} = \frac{1}{M}\sum_{i=1}^{M}D_{\mathrm{KL}}[q_{\mathbf{w}}^{(i)}\| p_{\mathbf{w}}]$ . Then, we set the number of blocks as  $K_{\beta,\kappa} = \lceil c_{\beta}/\kappa \rceil$  and we partition the weight vector such that the average divergence  $\bar{\delta}_k$  of each block estimated on the training data matches the coding budget, i.e.,  $\bar{\delta}_k \approx \kappa$  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.