SlowGuess/NeurIPS_2023_OCR huggingface dataset:
next paragraph) to classifications for 100 epochs using an Adam optimizer with default parameters, with the learning rate decreasing by a factor of 10 based on plateauing training loss, with a patience of 5 epochs, and early stopping if the learning rate fell below  $10^{-8}$ .
One particular design choice that we made in finetuning predictors merits elaboration: shifting encodings. Rather than directly training predictors to map from encodings to predictions, we applied a simple linear transformation to the encodings before feeding them to the predictor. Specifically, we multiplied all tensor elements by 5 and increased them by 1. This simple linear transformation does not affect any relations between encodings except scale, and indeed we found that predictors could be successfully trained with this rescaling. Nevertheless, this linear transformation was important to provide a check against merely relying upon initialization conditions for good finetuning performance. In particular, we found that if we did not apply this linear transformation (i.e., pass the raw encodings to the predictor), predictors sometimes performed better than they should given the training data. For example, as a sanity check, we trained a predictor on a binary classification task, but only provided two positive datapoints and no negative data. In general, given this data, one would expect a trained predictor to only predict positive labels. However, we observed that in several cases, the predictor would achieve nearly perfect accuracy, including predicting negative labels for negative inputs. This surprising result disappeared when we simply shifted encodings, indicating that the particular initialization conditions of the predictor seemed to align well with pre-trained encoders. We wanted to measure the effect of data on finetuning performance, rather than just initialization conditions, but we note that this odd phenomenon of well-aligned initializations merits further investigation.
We ran 10 finetuning trials per model, which was important given the small amount of randomly-sampled finetuning data.
In the following subsections, we present the hyperparameters used for training different encoders in the different domains. In general, we used the following principles when choosing hyperparameters:
- For VQ-based methods, use a large enough codebook to have at least one element per class. Larger  $C$  are also acceptable, as tuning weights should decrease the effective codebook size.
- When annealing, use a small enough weight increment to generate smooth changes during training. Larger increments, however, speed up training.
- When annealing for larger  $n$  one can increase the annealing rate. Models with greater  $n$  tended to use more complex representations, so annealing could be extremely slow for small increments.
A.3.1 FashionMNIST
For FashionMNIST, we trained all models with batch size 64 for 200 epochs, using the hyperparameters specified in Table 1. The only differences across methods were which hyperparameters we annealed to penalize complexity. Other differences simply reflected differences in architecture (e.g., using a codebook for vector-quantization methods). Pre-training a single model for 200 epochs took approximately 5 minutes on a desktop computer with one NVIDIA 2080 GeForce RTX.
Table 2: Hyperparameters for CIFAR100 training.
ENCODER	LATENT DIM	C	n	λU	λI	λC0	λC INCR	λH0	λH INCR
β-VAE	32	NA	NA	10	10	0.01	0.1	0.0	0.0
VQ-VIBN	32	1000	1	10	10	0.01	0.5	0.0	0.0
VQ-VIBN	32	1000	2	10	10	0.01	0.5	0.0	0.0
VQ-VIBN	32	1000	4	10	10	0.01	0.5	0.0	0.0
VQ-VIBC	32	1000	1	10	10	0.0	0.0	0.001	0.04
VQ-VIBC	32	1000	2	10	10	0.0	0.0	0.001	0.08
VQ-VIBC	32	1000	4	10	10	0.0	0.0	0.001	0.12
Table 3: Hyperparameters for iNaturalist training.
ENCODER	LATENT DIM	C	n	λU	λI	λC0	λC INCR	λH0	λH INCR
β-VAE	32	NA	NA	10	10	0.0001	0.03	0.0	0.0
VQ-VIBN	32	2000	1	10	10	0.001	0.1	0.0	0.0
VQ-VIBN	32	2000	2	10	10	0.001	0.2	0.0	0.0
VQ-VIBN	32	2000	4	10	10	0.001	0.4	0.0	0.0
VQ-VIBC	32	2000	1	10	10	0.0	0.0	0.001	0.05
VQ-VIBC	32	2000	2	10	10	0.0	0.0	0.001	0.10
VQ-VIBC	32	2000	4	10	10	0.0	0.0	0.001	0.20
A.3.2 CIFAR100
For CIFAR100, we trained all models with batch size 256 for 400 epochs, using the hyperparameters specified in Table 2. As explained for FashionMNIST, the only substantial differences across architectures were architecture-specific terms that needed to be specified, and which terms were annealed to penalize complexity. Pre-training a single model for 400 epochs took approximately 10 minutes on a desktop computer with one NVIDIA 3080.
A.3.3 iNaturalist
For iNat, we trained all models with batch size 256, using the hyperparameters specified in Table 3. We trained  $\beta$ -VAE and VQ-VIBc models for 300 epochs, while we trained VQ-VIBN models for 600. We used more annealing epochs for VQ-VIBN simply because it seemed to need more epochs to eventually anneal to random chance. Likely, a larger annealing rate could also accomplish the desired effect, but initial experiments with faster annealing tended to induce rapid codebook collapse that did not generate the smooth spectrum of MSE values we desired. Pre-training a single model for 300 epochs took approximately 10 minutes on a desktop computer with one NVIDIA 3080 (and twice as long for 600 epochs).
B Further Finetuning Results
In the main paper, we included only a small number of graphs highlighting our key results. Here, we include further results that corroborate the main trends stated in the paper. Primarily, these plots include further experiments for varying the amount of finetuning data, as well as varying  $n$ , the number of codebook elements to combine into a latent representation. We further include details of validation-set experiments, wherein we evaluated a method for autonomously selecting the optimal encoder. Results are divided by domain: FashionMNIST, CIFAR100, and iNat.
B.1 FashionMNIST
Here, we include the finetuning results for the FashionMNIST domain for varying amounts of finetuning data, ranging over  $k \in [1,2,5,10,50]$ , and  $n$ , the number of quantized vectors to combine into a single representation. Results for each  $k$  are included in Figure 10.
(a)  $k = 1$
(b)  $k = 2$
(d)  $k = 10$
Figure 10: FashionMNIST finetuning results for varying  $k$ . As  $k$  increased, all models benefited. The data-efficiency of advantage of VQ-VIB $_C$  was most pronounced when using the least amount of data.
(c)  $k = 5$
(e)  $k = 50$
Table 4: Finetuning accuracy across domains, using encoders selected via validation set accuracy. For large validation  $(v)$  and training  $(k)$  set sizes, this method selected high-performance encoders, but for low  $k$  and  $v$ , it struggled. This motivates our human-in-the-loop framework in the low-data regime.
FASHIONMNIST	CIFAR100 2-WAY	iNAT 2-WAY
v	k=2	5	10	50	2	5	10	50	2	5	10	50
1	0.73	0.89	0.93	0.97	0.71	0.79	0.83	0.90	0.65	0.79	0.89	0.90
5	-	-	0.94	0.97	-	-	0.82	0.91	-	-	0.88	0.90
10	-	-	-	0.97	-	-	-	0.91	-	-	-	0.89
As expected, increasing the amount of finetuning data improved performance for all models, and the gap between all model types (VQ-VIB $_C$ , VQ-VIB $_N$ , and  $\beta$ -VAE) shrank. It is noteworthy, however, that a VQ-VIB $_C$  model, tuned to the right complexity level and trained with just one example per ternary class (Figure 10 a), achieved better accuracy than a  $\beta$ -VAE model trained with 50 examples per class (Figure 10 e). Further, for any fixed  $k$ , VQ-VIB $_C$  consistently outperformed VQ-VIB $_N$ , suggesting that many recent works that use VQ-VIB $_N$  could be improved by replacing the model type [27, 28, 12, 7]. Lastly, for both VQ-VIB $_N$  and VQ-VIB $_C$ , increasing  $n$  tended to support lower MSE but worse finetuning accuracy. This supports an intuition that combining more discrete representations starts to more densely fill the representation space, trending towards continuous representations.
We note briefly that  $\mathrm{VQ - VIB}_{\mathcal{N}}$ , both in this domain and others (explored in the next sections), typically failed to learn as complex representations as either  $\mathrm{VQ - VIB}_{\mathcal{C}}$  or  $\beta$ -VAEs. This is apparent given the limited range of MSE values for the  $\mathrm{VQ - VIB}_{\mathcal{N}}$  curves. We consistently struggled to make  $\mathrm{VQ - VIB}_{\mathcal{N}}$  learn as rich representations as for the other model types, which led to worse reconstructions and higher MSE values.
Lastly, Table 4 includes results from our validation experiments for all three experiment domains. Recall that we tested a method for autonomously selecting the best encoder, among the suite of encoders of different complexity levels, by measuring validation set accuracy. Table 4 includes results from such experiments for different validation set sizes  $(v)$  and different  $k$ . For a given  $k$  and  $v$ , we randomly sampled  $v$  datapoints per class label to be part of the validation set and used the remaining
data for finetuning. For large  $k$ , this method worked quite well by selecting high-accuracy encoders. However, for small  $k$  and  $v$ , validation set accuracy was a noisy proxy for model performance (because of the small validation set size), so performance tended to be suboptimal. For example, in the FashionMNIST domain, for  $k = 2$ ,  $v = 1$ , models achieved mean performance of  $73\%$ , lower than the over than  $80\%$  accuracy achieved by tuning to the right complexity (Figure 10 b). Overall, these validation set experiments confirm that, for large enough  $k$ , one may autonomously select optimal encoders, but for very small  $k$ , such autonomous methods fail.
B.2 CIFAR100
We found similar trends in the CIFAR100 to those in the FashionMNIST domain and plotted results in Figures 11 and 12 (for 2-way and 20-way finetuning tasks, respectively). In all experiments, VQ-VIB $_C$  outperformed both  $\beta$ -VAE and VQ-VIB $_N$ . In the 2-way finetuning example, we again found a peaked curve for VQ-VIB $_C$  finetuning accuracy as a function of MSE, indicating that tuning to the right complexity level induced the best accuracy. In the more complex 20-way classification task, however, we did not observe this peak.
This last result is unsurprising: the 20-way hierarchy in CIFAR100 is less semantically meaningful and likely less obvious in photos than the 2-way task of distinguishing living and non-living things. For example, two of the 20 categories are simply different sorts of vehicles. It would be extremely surprising for VQ-VIB $_{C}$  to learn such arbitrary groups automatically while compressing representations. Without learning the right groupings, VQ-VIB $_{C}$  cannot benefit from learning less complex representations.
(d)  $k = 10$
Figure 11: CIFAR100 2-way finetuning results for varying  $k$ .
(d)  $k = 10$
Figure 12: CIFAR100 20-way finetuning results for varying  $k$ .
B.3 iNaturalist
Lastly, we found similar trends in finetuning in the iNat domain, finetuned on a 3-way (Figure 13), 34-way (Figure 14), and 1010-way (Figure 15) finetuning task.
On the 3-way finetuning task (between animals, plants, and fungi), we observed similar peaking behavior as in earlier experiments, indicating yet again the importance of tuning to the right complexity. In addition, as in prior results, we found a similar trend that greater  $n$  tended to allow greater complexity (lower MSE) but induced worse finetuning performance. For example, in Figure 13 b, the orange line, corresponding to  $n = 1$  stays above and to the right of the green ( $n = 2$ ) and red ( $n = 4$ ) lines. Intuitively, this seems to indicate that the more combinatorial representations, with greater  $n$ , were somewhat of a midpoint between the continuous  $\beta$ -VAE representations and the discrete representations used by VQ-VIBc for  $n = 1$ .
Results from the 34-way finetuning followed similar patterns as before as well. Just as in CIFAR100 wherein we tested both a 2-way and 20-way finetuning task, this 34-way finetuning task for iNat showed that  $\mathrm{VQ - VIB}_C$  continued to outperform  $\mathrm{VQ - VIB}_{\mathcal{N}}$  and  $\beta$ -VAE for more complex finetuning tasks, although the performance gap shrank as  $k$  increased.
Most interestingly, perhaps, we conducted yet another iNat finetuning experiment, this time using the 1010 low-level labels that had originally been used during pre-training. As before, we used very small amounts of data in finetuning (e.g., for  $k = 1$ , only 1 example from each class, so 1010 labeled examples total). Results from those experiments are shown in Figure 15.
For small  $k$ , we again see that  $\mathrm{VQ - VIB}_{\mathcal{C}}$  outperforms other model types. For larger  $k$ , however, we see one of the limitations of  $\mathrm{VQ - VIB}_{\mathcal{C}}$ . Because the discrete encoders learned less complex representations than  $\beta$ -VAEs (as shown by the fact that they never reach lower MSE values), with enough finetuning data,  $\beta$ -VAEs are able to capture distinctions between classes that  $\mathrm{VQ - VIB}_{\mathcal{C}}$  models cannot. Thus, in the particular case of large amounts of finetuning data and complex finetuning tasks, more complex, continuous encoders continue to outperform our method.
(a)  $k = 1$
(d)  $k = 10$
Figure 13: iNat 3-way finetuning results for varying  $k$ .
(b)  $k = 2$
(c)  $k = 5$
(e)  $k = 50$
(d)  $k = 10$
Figure 14: iNat 34-way finetuning results for varying  $k$ .
(a)  $k = 1$
(b)  $k = 2$
(d)  $k = 10$
Figure 15: iNat 1010-way finetuning results for varying  $k$ .
(a) Epoch 40
(c) Epoch 70
(b) Epoch 60
(d) Epoch 199
Figure 16: The evolution of the distribution over prototypes during annealing for VQ-VIB $_C$  in the FashionMNIST domain. In early epochs, VQ-VIB $_C$  uses many prototypes, with a long-tailed distribution. Over the course of annealing the entropy, the probability distribution becomes more concentrated (b and c) before collapsing to a single prototype (d).
C Prototype Utilization: Further Visualizations
Here, we include some further visualizations that we omitted from the main paper due to space constraints. These visualization primarily illustrate the importance of entropy-regulated representation learning (for  $\mathrm{VQ - VIB}_{\mathcal{C}}$ ) vs. complexity-regulated (for  $\mathrm{VQ - VIB}_{\mathcal{N}}$ ).
Figures 16 and 17 shows the prototypes for VQ-VIB $_C$  and VQ-VIB $_N$ , respectively, in the FashionMNIST domain over the course of training. Each subfigure consists of a top row of decoded prototypes, with associated probabilities (frequency of use measured when passing through images from the test set) below. The 30 most frequent prototypes are visualized, or fewer prototypes if fewer were used.
There is an important trend in Figures 16 and 17: the entropy-based annealing for VQ-VIB $_{\mathcal{C}}$  caused models to use fewer prototypes, while the complexity-based annealing for VQ-VIB $_{\mathcal{N}}$  did not. At epoch 40, just as both methods begin annealing, VQ-VIB $_{\mathcal{C}}$  and VQ-VIB $_{\mathcal{N}}$  use a large number of prototypes, as seen by the long-tailed distributions. Over the course of annealing, however, VQ-VIB $_{\mathcal{C}}$  uses fewer prototypes, and merges images of different classes into the same prototype. Thus, the degenerate encoder at the end of annealing (epoch 199) uses just a single prototype to represent all possible inputs (Figure 16). At the same time, VQ-VIB $_{\mathcal{N}}$ , during annealing, does not use fewer prototypes. Rather, the complexity-penalization term seems to induce the model to make the mapping from input to prototype more stochastic (Figure 17). Thus, the degenerate VQ-VIB $_{\mathcal{N}}$  encoder uses many prototypes, each of which is blurry because it could correspond to any input.
Visualizations of decoded prototypes for the CIFAR100 domain is more challenging. In richer image domains, prototype-based methods often use training examples as prototypes [3, 18, 4], which can make it more difficult to understand when a single prototype represents more than one concept. Nevertheless, by visualizing the distribution over prototypes (without decoding them), we see the same pattern that  $\mathrm{VQ - VIB}_{\mathcal{C}}$  tends to learn to use fewer prototypes over the course of annealing than  $\mathrm{VQ - VIB}_{\mathcal{N}}$ . Snapshots of the categorical distributions for CIFAR100 are included in Figure 18.
D Ablation Study: Entropy vs. Complexity
Here, we present results motivating penalizing entropy, as opposed to complexity, in VQ-VIB $_{C}$ . Appendix C showed how annealing entropy in VQ-VIB $_{C}$  caused models to use fewer prototypes, whereas penalizing complexity in VQ-VIB $_{N}$  did not induce similar reductions in effective codebook
(b) Epoch 60
(d) Epoch 199
Figure 17: The evolution of the distribution over prototypes during annealing for  $\mathrm{VQ - VIB}_{\mathcal{N}}$  in the FashionMNIST domain. Unlike annealing entropy for  $\mathrm{VQ - VIB}_{\mathcal{C}}$ , annealing the complexity in  $\mathrm{VQ - VIB}_{\mathcal{N}}$  did not result in fewer prototypes being used. Instead, each prototype became blurrier, indicating that each prototype became more likely regardless of class.
(a) VQ-VIB<sub>C</sub> Epoch 50
(b) VQ-VIB Epoch 399
(c) VQ-VIB $_{N}$  Epoch 50
Figure 18: Categorical distribution over prototypes while annealing in the CIFAR100 domain at the start of annealing (Epoch 50) and at the end (Epoch 399) for VQ-VIB $_C$  (top row) and VQ-VIB $_N$  (bottom row). Annealing the entropy in VQ-VIB $_C$  caused the model to use fewer prototypes (note the degenerate categorical distribution over only one prototype at Epoch 399), whereas annealing complexity for VQ-VIB $_N$  did not cause similar concentration.
(d) VQ-VIB Epoch 399
(d)  $k = 10$
Figure 19: FashionMNIST finetuning results for varying  $k$ , comparing annealing by entropy (VQ-VIB $_C$ ) and annealing by complexity (VQ-VIB $_C$  Comp.). Annealing by complexity resulted in worse finetuning performance.
size. Further experiments corroborate our findings that penalizing entropy was the key to this difference in behavior.
We trained VQ-VIB $_C$  agents on the FashionMNIST task, using the same pre-training and finetuning procedures as in the main paper, with the only difference being that we annealed the complexity of representations instead of the entropy. Results from finetuning such models are included in Figure 19.
Figure 19 shows that annealing by entropy, as opposed to complexity, was the key factor in improving VQ-VIB $_{C}$  finetuning performance. The difference in performance when penalizing entropy vs. complexity closely matches the difference in performance between VQ-VIB $_{C}$  and VQ-VIB $_{N}$  examined in the main paper. Thus, the entropy-regularization term seems to explain much of the difference between VQ-VIB $_{C}$  and VQ-VIB $_{N}$ .
In subsequent experiments in the CIFAR100 and iNat domains, therefore, we tested whether penalizing the estimated entropy of VQ-VIB $_{\mathcal{N}}$  models matched VQ-VIB $_{c}$  results from the main paper. We note that Tucker et al. [27] advocate for a small positive  $\lambda_{H}$  to penalize entropy, but the authors also acknowledge that exactly computing this entropy is impossible given the VQ-VIB $_{\mathcal{N}}$  architecture.
Finetuning results for CIFAR100 (on the 2-way and 20-way finetuning tasks) and iNat (on the 3-way and 34-way finetuning tasks), for VQ-VIB $_C$  trained by varying  $\lambda_C$  and VQ-VIB $_N$  trained by varying  $\lambda_H$ , are included in Figures 20, 21, 22, and 23. Several important trends emerge from viewing these plots, especially compared to results from our main paper for VQ-VIB $_C$  controlled via  $\lambda_H$ .
First, finetuning performance is noisier using these models compared to results from the main text. This likely arises, for VQ-VIB $_N$  models, because increasing  $\lambda_H$  failed to consistently reduce the number of discrete representations used. Thus, for a given MSE value, different models used different numbers of representations, and therefore exhibited different finetuning performance.
Second, varying  $\lambda_{H}$ , instead of  $\lambda_{C}$ , seemed to somewhat improve VQ-VIB $_{\mathcal{N}}$  performance, but not as much as when varying  $\lambda_{H}$  for VQ-VIB $_{\mathcal{C}}$ , as presented in our main paper. For example, consider Figure 21 a. The best-performing model, VQ-VIB $_{\mathcal{N}}$ ,  $n = 1$ , peaks at finetuning accuracy of approximately 0.16, outperforming VQ-VIB $_{\mathcal{C}}$  models when varying  $\lambda_{C}$ . However, in Figure 12 a, we found that VQ-VIB $_{\mathcal{C}}$  models in the exact same setting achieved a mean accuracy of approximately 0.38: more than double the VQ-VIB $_{\mathcal{N}}$  performance. Thus, varying  $\lambda_{H}$  seemed to improve VQ-VIB $_{\mathcal{N}}$
(a)  $k = 1$
(d)  $k = 10$
Figure 20: Ablation study results for the CIFAR100 2-way finetuning task. Tuning  $\lambda_{C}$  for VQ-VIB $_{C}$  or  $\lambda_{H}$  for VQ-VIB $_{\mathcal{N}}$  led to worse results than in our main paper.
(a)  $k = 1$
(b)  $k = 2$
(d)  $k = 10$
Figure 21: Ablation study results for the CIFAR100 20-way finetuning task. Tuning  $\lambda_{C}$  for VQ-VIB $_{C}$  or  $\lambda_{H}$  for VQ-VIB $_{\mathcal{N}}$  led to worse results than in our main paper.
(d)  $k = 10$
Figure 22: Ablation study results for the iNat 3-way finetuning task. Tuning  $\lambda_{C}$  for VQ-VIB $_{C}$  or  $\lambda_{H}$  for VQ-VIB $_{N}$  led to worse results than in our main paper.
(a)  $k = 1$
(d)  $k = 10$
Figure 23: Ablation study results for the iNat 34-way finetuning task. Tuning  $\lambda_{C}$  for VQ-VIB $_{C}$  or  $\lambda_{H}$  for VQ-VIB $_{N}$  led to worse results than in our main paper.
performance somewhat, but  $\mathrm{VQ - VIB}_{\mathcal{C}}$  better supports penalizing entropy, and therefore achieves higher performance.
Third, varying  $\lambda_{C}$ , instead of  $\lambda_{H}$ , for VQ-VIB $_{C}$  worsened finetuning performance. Once again, by comparing finetuning performance for VQ-VIB $_{C}$  models in Figure 21 a (achieving a maximum accuracy around 0.14), to results from our main paper, we note the importance of penalizing the entropy of representations.
Thus, in general these ablation studies support many of the design decisions made in the main paper.
1. Varying  $\lambda_{H}$ , instead of  $\lambda_{C}$  for VQ-VIB $_{C}$  improves finetuning performance by decreasing the number of discrete representations used.
2. VQ-VIB $_{\mathcal{N}}$  benefits somewhat from penalizing entropy, but because it is architecturally unable to support exact calculations of entropy, we were unable to match VQ-VIB $_{\mathcal{N}}$  performance.
It is certainly possible that some optimal combination of  $\lambda_{H}$  and  $\lambda_{C}$  might further improve VQ-VIB $_{N}$  or VQ-VIB $_{C}$  performance; initial explorations of such combinations with fixed  $\lambda_{H}$  values while annealing  $\lambda_{C}$  did not yield obvious results. Most importantly, our current findings are enough to indicate that controlling the entropy of discrete representations appears important for data-efficient finetuning.
E User Study Results
Here, we include complete results from our user study. Table 5 includes accuracy rates for all model types and visualizations, for all three questions. For Questions 2 and 3, selecting the lowest-MSE (highest complexity) encoder was correct, so accuracy rates for all model and visualization types for these questions remained high. For Question 1, however, selecting the correct encoder required users to select non-maximally-complex representations; for this question, only users viewing prototypes of  $\mathrm{VQ - VIB}_C$  models performed above random chance.
Table 5: Accuracy rates by encoder type and visualization method for the three questions (reporting means and standard errors over 20 subjects). For the most challenging question (Question 1), visualizing VQ-VIB $_C$  prototypes supported the greatest accuracy.
Model	Viz.	Question 1	Question 2	Question 3
VQ-VIBc	MSE	0.10 (0.02)	0.70 (0.05)	0.85 (0.03)
VQ-VIBc	Proto	0.55 (0.06)	0.75 (0.04)	0.90 (0.02)
VQ-VIBN	MSE	0.10 (0.02)	0.80 (0.04)	0.85 (0.03)
VQ-VIBN	Proto	0.10 (0.02)	0.90 (0.02)	0.90 (0.02)