SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
N ^ {*} = \underset {n \in [ 1: N ]} {\arg \max } \tilde {r} _ {n}, \tag {10}
$$
we have that  $\mathbf{w}_{N^*} \sim \tilde{q}_{\mathbf{w}}$  is approximately distributed according to the target, i.e.  $\tilde{q}_{\mathbf{w}} \approx q_{\mathbf{w}}$ . More precisely, we have the following result:
Lemma A.1 (Bound on the total variation between  $\tilde{q}_{\mathbf{w}}$  and  $q_{\mathbf{w}}$  (Lemma D.1 in [32])). Let us set the number of proposal samples simulated by Algorithm 2 to  $N = 2^{D_{\mathrm{KL}}[q_{\mathbf{w}} \| p_{\mathbf{w}}] + t}$  for some parameter  $t \geqslant 0$ . Let  $\tilde{q}_{\mathbf{w}}$  denote the approximate distribution of the algorithm's output for this choice of  $N$ . Then,
$$
D _ {T V} \left(q _ {\mathbf {w}}, \tilde {q} _ {\mathbf {w}}\right) \leqslant 4 \epsilon , \tag {11}
$$
$$
\epsilon = \left(2 ^ {- t / 4} + 2 \sqrt {\mathbb {P} _ {Z \sim q _ {\mathbf {w}}} \left[ \log_ {2} r (Z) \geqslant D _ {\mathrm {K L}} [ Q \| P ] + t / 2 \right]}\right) ^ {1 / 2}. \tag {12}
$$
This result essentially tells us that we should draw at least around  $2^{D_{\mathrm{KL}}[q_{\mathrm{w}} \| p_{\mathrm{w}}]}$  samples to ensure low sample bias, and beyond this, the bias decreases exponentially quickly as  $t \to \infty$ . However,
3The PDF of a standard Gumbel random variable truncated to  $(-\infty, b)$  is given by  $\operatorname{TruncGumbel}(x \mid b) = \mathbf{1}[x \leqslant b] \cdot \exp(-x - \exp(-x) + \exp(-b))$ .
note that the number of samples we need also increases exponentially quickly with  $t$ . In practice, we observed that when  $D_{\mathrm{KL}}[q_{\mathbf{w}} \| p_{\mathbf{w}}]$  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_{\mathbf{w}}$ . Then we select the sample  $\boldsymbol{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  $\{\mathcal{D}_i\} = \{\mathcal{D}_1,\mathcal{D}_2,\dots,\mathcal{D}_M\}$ , we fit a variational distribution  $q_{\mathbf{w}}^{(i)}$  to represent each of the  $\mathcal{D}_i$ s. To do this, we minimize the loss (abbreviated as  $\mathcal{L}$  later)
$$
\begin{array}{l} \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) (13) \\ = \frac {1}{M} \sum_ {i = 1} ^ {M} \left\{\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}} \left[ q _ {\mathbf {w}} \| p _ {\mathbf {w}; \boldsymbol {\theta} _ {p}} \right] \right\}. (14) \\ \end{array}
$$
Now calculate the derivative w.r.t. the prior distribution parameter  $p_{\mathbf{w};\boldsymbol{\theta}_p}$
$$
\frac {\partial \mathcal {L}}{\partial \boldsymbol {\theta} _ {p}} = \frac {1}{M} \sum_ {i = 1} ^ {M} \frac {\partial D _ {\mathrm {K L}} \left[ q _ {\mathbf {w}} \| p _ {\mathbf {w} , \boldsymbol {\theta} _ {p}} \right]}{\partial \boldsymbol {\theta} _ {p}} \tag {15}
$$
Considering we choose factorized Gaussian as variational distributions, the KL divergence is
$$
\begin{array}{l} D _ {\mathrm {K L}} \left[ q _ {\mathbf {w}} ^ {(i)} \| p _ {\mathbf {w}, \boldsymbol {\theta} _ {p}} \right] = D _ {\mathrm {K L}} \left[ \mathcal {N} \left(\boldsymbol {\mu} _ {i}, \operatorname {d i a g} \left(\boldsymbol {\sigma} _ {i}\right)\right) \| \mathcal {N} \left(\boldsymbol {\mu} _ {i}, \operatorname {d i a g} \left(\boldsymbol {\sigma} _ {i}\right)\right) \right] (16) \\ = \frac {1}{2} \log \frac {\sigma_ {p}}{\sigma_ {q} ^ {(i)}} + \frac {\sigma_ {q} ^ {(i)} + \left(\boldsymbol {\mu} _ {q} ^ {(i)} - \boldsymbol {\mu} _ {p}\right) ^ {2}}{\sigma_ {p}} - \frac {1}{2} (17) \\ \end{array}
$$
To compute the analytical solution, let
$$
\frac {\partial \mathcal {L}}{\partial \boldsymbol {\theta} _ {p}} = \frac {1}{M} \sum_ {i = 1} ^ {M} \frac {\partial D _ {\mathrm {K L}} \left[ q _ {\mathbf {w}} \| p _ {\mathbf {w} , \boldsymbol {\theta} _ {p}} \right]}{\partial \boldsymbol {\theta} _ {p}} = 0. \tag {18}
$$
Note here  $\sigma$  refers to variance rather than standard deviation. The above equation is equivalent to
$$
\frac {\partial \mathcal {L}}{\partial \boldsymbol {\mu} _ {p}} = \sum_ {i = 1} ^ {M} \frac {\boldsymbol {\mu} _ {p} - \boldsymbol {\mu} _ {q} ^ {(i)}}{\boldsymbol {\sigma} _ {p}} = 0, \tag {19}
$$
$$
\frac {\partial \mathcal {L}}{\partial \pmb {\sigma} _ {p}} = \sum_ {i = 1} ^ {M} [ \frac {1}{\pmb {\sigma} _ {p}} - \frac {\pmb {\sigma} _ {q} ^ {(i)} + (\pmb {\mu} _ {q} ^ {(i)} - \pmb {\mu} _ {p}) ^ {2}}{\pmb {\sigma} _ {p} ^ {2}} ] = 0.
$$
We finally can solve these equations and get
$$
\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 {20}
$$
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.
Figure 7: The approximated PSNR value changes as the fine-tuning process goes on.
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  $\beta$
When learning the model prior, the value of  $\beta$  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  $\mathbf{w}_{1:K} = \{\mathbf{w}_1,\dots ,\mathbf{w}_K\}$  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  $\kappa = 16$  bits, we adjust the value of  $\beta$  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  $[\kappa - 0.4, \kappa]$  as buffer area where we do not change the value of  $\lambda_{k}$ . Here we only adjust  $\lambda_{k}$  every 15 iterations to avoid frequent changes at the initial training stage.
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 \times 10^{-6}$ .
Algorithm 4 Dynamic  $\beta$  adjustment for optimizing the posteriors
Require:  $\beta ,\mathbf{w}_{1:K} = \{\mathbf{w}_1,\dots ,\mathbf{w}_K\}$    
Initialize:  $\lambda_{k} = \beta ,k = 1,\dots ,K$    
Initialize: variational posterior  $q_{\mathbf{w}_k},k = 1,\dots ,K$    
for  $i\gets$  NumberIter do   
 $\delta_{k} = D_{\mathrm{KL}}[q_{\mathbf{w}_{k}}\| p_{\mathbf{w}_{k}}],k = 1,\dots ,K$ $q_{\mathbf{w}_{1:K}}\gets$  VariationalUpdate  $(\mathcal{L}_{\lambda_{1:K}})$ $\triangleright$ $\mathcal{L}_{\lambda_1:K}$  is defined in Equation 8 in the main text if (i mod 15)  $= 0$  then if  $\delta_k > \kappa$  then  $\lambda_{k} = \lambda_{k}\cdot 1.05$  end if if  $\delta_{k} <   \kappa -0.4$  then  $\lambda_{k} = \lambda_{k} / 1.05$  end if end if   
end for   
return  $q_{\mathbf{w}_k},\lambda_k,k = 1,\dots ,K$
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  $\{10^{-7}, 10^{-8}, 4 \times 10^{-8}\}$ . 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 \times 10^{-9}$ . We empirically adjust the variance initialization from the set  $\{4 \times 10^{-6}, 4 \times 10^{-10}\}$  and find they can affect the converged bitrate and achieve good performance. In particular, the posterior variance is initialized as  $4 \times 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 \times 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  $\beta = \{10^{-7},3\times 10^{-8},10^{-8},10^{-9}\}$ . 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\times 10^{-9}$ . The settings for optimizing and fine-tuning posterior of a test datum are the same as the experiments on Kodak dataset.
CIFAR-10	Kodak	LibriSpeech
Smaller Model	Larger Model
Network Structure
number of MLP layer	4	6	7	6
hidden unit	16	48	56	48
Fourier embedding	32	64	96	64
number of parameters	1123	12675	21563	12675
Learning Model Prior from Training Data
number of training data	2048	512	512	1024
epoch number	128	96	96	128
learning rate	0.0002	0.0001	0.0001	0.0002
iteration / epoch (except the first epoch)	100	200	200	100
iteration number in the first epoch	250	500	500	250
initialization of posterior variance	\( 9 \times 10^{-6} \)	\( 4 \times 10^{-6} \)	\( 4 \times 10^{-6}, 4 \times 10^{-10} \)	\( 4 \times 10^{-9} \)
\( \beta \)	\( 2 \times 10^{-5}, 5 \times 10^{-6}, 2 \times 10^{-6} \)\( 1 \times 10^{-6}, 5 \times 10^{-7} \)	\( 10^{-7}, 10^{-8}, 4 \times 10^{-8} \)	\( 4 \times 10^{-6} \)	\( 10^{-7}, 3 \times 10^{-8} \)\( 10^{-8}, 10^{-9} \)
Optimize the Posterior of a Test Datum
iteration number	25000	25000	25000	25000
learning rate	0.0002	0.0001	0.0001	0.0002
training with 1/4 the points (pixels)	✘	✓	✓	✘
number of group (KL budget = 16 bits / group)	(58, 89, 146, 224, 285)	(1729, 2962, 3264)	(5503, 7176)	(1005, 2924, 4575, 6289)
bitrate, (bpp for images, Kbps for audios)	(0.91, 1.39, 2.28, 3.50, 4.45)	(0.070, 0.110, 0.132)	(0.224, 0.293)	(5.36, 15.59, 24.40, 33.54)
PSNR, dB	(0.91, 1.39, 2.28, 3.50, 4.45)	(0.070, 0.110, 0.132)	(0.224, 0.293)	(5.36, 15.59, 24.40, 33.54)
Table 3: Hyper parameters in our experiments.
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 SmallChunks
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 \times 200$ . We use the same
Figure 8: Impact of the number of training data.
Figure 9: Compressing audios.
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.
Ground Truth
Figure 10: Decoded Kodak images.
0.0703 bpp, 23.02 dB
0.0703 bpp, 29.73 dB
0.2928 bpp, 25.43 dB
0.2928 bpp, 33.59 dB