SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 5: CIFAR100 finetuning, using 1 (a) or 50 (b) training examples for living vs. non-living things. With small amounts of data VQ-VIBc benefits from tuning to the right complexity level. When finetuning on the 20-way classification task (c) VQ-VIBc continues to outperform other architectures.
(a) iNat 3-way  $k = 1$
Figure 6: iNat finetuning results for the plant-animal-fungus classification task. Using only 1 finetuning example (a), VQ-VIB $_C$  performance improved as complexity decreased, and outperformed  $\beta$ -VAE and VQ-VIB $_N$ . Using 50 finetuning examples, performance plateaued, but VQ-VIB $_C$  continued to outperform other methods (b). VQ-VIB $_C$  performance shifts smoothly for varying  $k$  (c).
(b) iNat 3-way  $k = 50$
(c) iNat VQ-VIB  $\mathcal{C}$  varying  $k$
performance (see Figure 10 in Appendix B). We also tested methods for selecting encoders via validation-set accuracy and found similar trends. In finetuning, we held out a subset of the data for assessing finetuning accuracy and selected the best-performing encoder via validation set accuracy. For  $k \leq 5$ , this validation-set approach did not consistently converge to optimal performance, but, for sufficiently large  $k$ , it did. Further results from this approach are included in Table 4 and Appendix B. Generally, we found that for large enough  $k$ , the importance of tuning to the right complexity decreases, and several methods can select optimal encoders; for very small  $k$ , however, autonomous methods are suboptimal.
This illustrative FashionMNIST use case demonstrates many of the important trends we explore in our later experiments: tuning VQ-VIB $_C$  representations to the "right" complexity was important for optimal performance for small  $k$ , and VQ-VIB $_C$  generally outperformed other encoders for few-shot finetuning.
4.4 CIFAR100 and iNat
In this section, we present results from computational experiments in the more challenging CIFAR100 and iNat domains. As before, in the smallest data regime (small  $k$ ), using less complex representations afforded greater efficiency benefits, and  $\mathrm{VQ - VIB}_{\mathcal{C}}$  outperformed  $\beta$ -VAE and  $\mathrm{VQ - VIB}_{\mathcal{N}}$ .
Figure 5 a and b show finetuning accuracy on the binary classification task for CIFAR100 on identifying living vs. non-living things. With only 1 datapoint per class (Figure 5 a),  $\beta$ -VAE performance remained effectively flat at random chance: barely exceeding  $50\%$ . However, for VQ-VIBc, for MSE  $< 0.35$ , performance increased as MSE increased (linear regression slope was positive  $p < 0.001$ ), up to a  $70\%$  accuracy rate, before worsening as MSE increased further. Increasing  $k$  to 50 (Figure 5 b) shrank the gap between the encoder architectures, and flattened the improvements previously observed, but VQ-VIBc continued to outperform  $\beta$ -VAE and VQ-
(a) PCA of high-complexity VQ-VIB $_C$ . Note the orange bird in the pine tree.
(b) PCA of low-complexity  $\mathrm{VQ - VIB}_C$  Note the bird in its nest on the left.
Figure 7: 2D PCA of the VQ-VIB $_{C}$  latent space for iNat, at different complexity levels. Each gray point represents the continuous encoding output by the encoder; the colorful points represent the learnable prototypes. For visualization purposes, the closest training image to each prototype is included in the diagrams. In the high-complexity case, the model has learned many distinct prototypes, representing concepts like birds or insects. In the low-complexity case, VQ-VIB $_{C}$  uses only two prototypes, which capture a plant-animal distinction the aligns closely with the iNat hierarchy.
$\mathrm{VIB}_{\mathcal{N}}$  models. When finetuned on the 20-way classification task, the same trends of  $\mathrm{VQ - VIB}_C$  outperforming other architectures held (Figure 5 c), although the finetuning accuracy decreased monotonically as MSE increased, rather than peaking at a specific complexity. Plots for a larger range of  $k$ , and for varying  $n$ , for both finetuning tasks, are included in Appendix B.
Similar trends held in the iNat dataset when finetuned on the 3-way animal-plant-fungus task, as depicted in Figure 6. For small  $k$ , VQ-VIB $_C$  finetuning performance initially improved as the encoder learned more compressed representations (Figure 6 a). For VQ-VIB $_C$ ,  $k = 1$ , the linear correlation between MSE and finetuning accuracy is positive ( $p < 0.001$ ) for MSE  $< 0.55$ . Further analysis of finetuning performance for VQ-VIB $_C$ , for varying  $k$ , shows a smooth change in behavior as  $k$  increases: simultaneously improving overall performance, and benefiting less from compressed representations (Figure 6 c). This indicates that VQ-VIB $_C$  is most advantageous in a low-data regime. Similar results for finetuning on the 34-way finetuning task are included in Appendix B. Visualization via 2D principle component analysis (PCA) confirms our intuition of how VQ-VIB $_C$  models support few-shot learning for iNat. At lower complexity levels, a VQ-VIB $_C$  encoder used a decreasing number of prototypes to represent increasingly abstract concepts (Figure 7).
As in the FashionMNIST domain, we evaluated autonomous methods for selecting the optimal encoder and found that, while they succeeded for sufficiently large  $k$ , they struggled in a low-data regime. For  $k = 1$ , for example, the heuristic of choosing the most complex encoder was suboptimal (e.g., see Figure 5 a and Figure 6 a). At the same time, for  $k = 50$ , this heuristic worked quite well (e.g., Figure 5 b and Figure 6 b). Lastly, selecting models via validation set accuracy similarly worked well for large enough  $k$ , but struggled for  $k \leq 5$  (see Appendix B).
Finally, we briefly note some limitations of VQ-VIB $_C$  and our finetuning method. In the results discussed so far, we found consistent advantages in using VQ-VIB $_C$  and low-complexity representations. However, as the amount of finetuning data increases, more complex representation methods like  $\beta$ -VAEs outperform VQ-VIB $_C$ . A more complete discussion of this phenomenon is included in Appendix B. Between our main results and these limitations, we find that the data-efficiency benefits of using VQ-VIB $_C$  in finetuning are greatest in data-poor and low-complexity settings.
5 Human-in-the-Loop Selection of Task-Appropriate Representations
Given our main motivation of a human-in-the-loop framework for selecting task-appropriate abstractions (recall Figure 1), we conducted a human-participant study to evaluate whether users could select the highest-performing task-appropriate models given visualizations of prototypes as task abstractions. This is important because, in order to take full advantage of our setup, users must be able to select the appropriate task-level representation so the neural network best learns the finetuning task in a sample efficient manner. We asked users to select the optimal encoder for a given task, given a visualization of encoder prototypes. A positive result would show that, given a spectrum of pre-trained VQ-VIB $_c$  encoders, a human user wishing to deploy a task-appropriate model could specify the right encoder for a given task, and thus fine-tune the model in a sample efficient manner.
Figure 8: Participants were asked to select the optimal encoder for the FashionMNIST finetuning task based on MSE values (left) or decoded representations (middle). Using only MSE values, users often incorrectly selected option b), whereas when viewing prototypes, users were able to correctly identify the appropriate abstraction level, and thus the optimal encoder (c).
5.1 User Study
Our human experiment was performed using the pre-trained models from the FashionMNIST domain in Section 4.3. Users were told that a robot was good at sorting clothing into 10 distinct categories (i.e., the 10 normal categories in FashionMNIST), but that in their case they should consider a different set of categories. Each user was randomly assigned a model type (VQ-VIB $_C$  or VQ-VIB $_\mathcal{N}$ ) and a visualization type (prototypes or a plot of MSE during training). In three questions (explained below), users were shown different groupings of the 10 FashionMNIST classes and asked to select which of four model checkpoints they thought best represented the groupings. Overall, this between-subjects setup allowed us to compare the effects of different visualization methods on user accuracy in selecting the optimal representation.
For our main question, in which users were asked to select an encoder that could sort FashionMNIST items into the three finetuning categories (tops, shoes, and other), our null hypothesis,  $\mathbf{H_0}$ , was "Users viewing  $VQ-VIB_{\mathcal{C}}$  prototypes will select the optimal encoder as often as those viewing MSE scores". Our alternative hypothesis,  $\mathbf{H_1}$ , was "Viewing  $VQ-VIB_{\mathcal{C}}$  prototypes improves users' ability to select the optimal encoder compared to just viewing MSE scores."
Participants. We crowd-sourced 20 participants per group from Prolific.com (N=80). The sample was balanced to have an even number of male and female participants. All participants were native English speakers above the age of 18 and resided in the U.S., the U.K., or Ireland. Given estimates of survey duration, calculated in a pilot study, we paid participants according to an estimated $12 per hour wage. The study received IRB approval from MIT.
Materials. Each user was presented three questions, corresponding to different groupings of the 10 FashionMNIST classes: in Question 1 classes were grouped according to the finetuning labels in our computational experiments (tops, shoes, and other), in Question 2, there were 10 distinct classes corresponding to the 10 FashionMNIST classes, and in Question 3, classes were grouped according to an unintuitive 3-way grouping (e.g., Pullover, Dress, Sneaker were one class). Despite the numbering, we randomized the order of questions for each participant. In each question, users were asked to select which of four model checkpoints they thought best represented groupings for that specific question. The four models corresponded to checkpoints taken 1) near the start of training, 2) at the minimum MSE value, 3) at an intermediate MSE value, and 4) at the end of training. We primarily focused on results for Question 1, as it was the only question for which selecting the most complex encoder was not optimal. Responses were categorized as correct if users picked the optimal encoder (option "c" in Figure 8) and otherwise incorrect (but during the study, users were not told whether they were correct).
Results from our user study supported our hypothesis. As shown in Figure 8, for Question 1, corresponding to the 3-way grouping from our computational experiments, users who viewed VQ-VIB $_C$  prototypes selected the optimal encoder  $55\%$  of the time, compared to  $10\%$  accuracy when
viewing MSE scores (prototype accuracy was significantly greater at  $p < 0.01$  for a Fisher's exact test). Interestingly, users who viewed VQ-VIB $_{\mathcal{N}}$  models never achieved high performance (binomial test for non-random chance at  $p > 0.2$ ), suggesting an important advantage of the VQ-VIB $_{\mathcal{C}}$  architecture. Thus, VQ-VIB $_{\mathcal{C}}$  prototypes was the only visualization method that supported greater-than-random chance performance in selecting optimal abstractions.
Further analysis, using responses to Questions 2 and 3 as well, highlighted the benefit afforded by VQ-VIB $_C$  prototypes. Overall, we observed high accuracy rates for Questions 2 and 3 regardless of visualization method; this is unsurprising given that the correct behavior for both questions was the selecting most complex encoder. We fitted a Mixed Linear Effects Model (MLEM) to the survey results, predicting user accuracy as a function of model visualization and question, grouped by participant. Using Wilkinson notation [30], our model was:  $Acc. \sim Q + M + Q : M + (1|Participant)$  where  $Q$  represents the categorical variable for question and  $M$  represents the categorical variable for model type and visualization (VQ-VIB $_C$  Prototypes, VQ-VIB $_C$  MSE, VQ-VIB $_N$  Prototypes, etc.). We grouped results by participant to model random-intercept effects of some participants being more accurate than others. Full results for accuracy rates and the MLEM model are included in Appendix E.
Overall, we found a significant  $(p < 0.05)$  positive interaction effect between Question 1 and visualization of VQ-VIB $_C$  prototypes. The only other significant effect we found was a negative effect for Question 1. Together, these results show that 1) Question 1 was harder than the other two questions and 2) viewing VQ-VIB $_C$  prototypes mitigated the effects of increased difficulty for Question 1. This provides crucial support for our hypothesis, showing that viewing VQ-VIB $_C$  prototypes enabled participants to select the optimal encoder more often.
Lastly, the combination of all results helps rule out several alternative hypotheses for how participants selected encoders based on VQ-VIB $_C$  prototypes. Users did not merely select the most complex encoder; otherwise, they would have performed poorly on Question 1. Similarly, users did not simply use the fewest number of prototypes greater than the number of finetuning classes; otherwise, they would have performed worse on Question 3. Thus, our results suggest that users interpreted VQ-VIB $_C$  prototypes as desired in intelligently selecting task-appropriate representations.
6 Contributions
We proposed a human-in-the-loop framework for selecting task-appropriate representations for data-efficient adaptation to new tasks. Discrete information bottleneck methods provide a principled way to generate representations along a spectrum of low to high complexity; we tested methods from prior literature and a novel architecture,  $\mathrm{VQ - VIB}_C$ , for generating such representations.
We found that controlling the complexity of representations was important for data-efficient fine-tuning: overly-complex representations required more training examples, but overly-simplified representations failed to capture important information. In computational studies, we showed the importance of learning low-entropy representations and that  $\mathrm{VQ - VIB}_{\mathcal{C}}$  tends to outperform other discrete information bottleneck methods. In a human study, we found that human partners are able to identify optimal complexity levels, indicating a promising direction for human-in-the-loop training.
Broader Impact: Generally, we hope that this work supports broader and better human-AI collaboration, supporting human-selected individualized models for particular use cases; however, we recognize that our current prototype-inspection framework may be limited and that care must be taken in future work to verify that representations along the complexity spectrum correspond to desired human concepts. In future work, we hope to explore models that simultaneously support different complexity levels instead of training separate models for different levels.
We thank members of the Interactive Robotics Group for helpful feedback and discussions. Andi Peng is supported by the NSF Graduate Research Fellowship and Open Philanthropy. Mycal Tucker is supported by an Amazon Alexa Science Hub Fellowship.
Figure 9: In experiments, we used a common feature-extractor  $(F)$ , decoder  $(D)$ , and predictor  $(P)$  network backbone while replacing different encoder heads  $(E)$ .
A.1 Pretraining
In pre-training (before the finetuning with a small number of examples on a cruder task), we used the following setup. In general, the overall network architecture comprised a feature extractor, an encoder head, a decoder, and a predictor, as depicted in Figure 9.
In all experiments, the predictor was parametrized as a two-layer feedforward neural network with hidden dimension 128 and a ReLU activation. The last layer's dimension depended upon the exact prediction task (e.g., 10 neurons for FashionMNIST, 100 for CIFAR100, and 1010 for iNat) and used a softmax activation.
The feature extractors and decoders varied by domain. For FashionMNIST, the feature extractor used 3 2D convolution layers, followed by one fully connected layer. The decoder used two linear layers, followed by 3 inverse convolution layers. We again emphasize that the code for these models is available in our codebase, linked to at the beginning of this section.
For CIFAR100 and iNat, we pre-processed the images to extract the 512-dimensional activations from the penultimate layer of a ResNet18 pretrained on ImageNet [8]. These features were used as inputs to the feature extractor ( $x$  in Figure 9). For both CIFAR100 and iNat, the feature extractor used two linear layers, with a ReLU activation after the first layer, which had hidden dimension 128. The decoder was used to reconstruct the 512-dimension outputs of the ResNet18, using 3 fully-connected layers of dimension 128, 256, and 512, with ReLU activations between layers.
The different encoder heads were  $\beta$ -VAE, VQ-VIB $_C$ , and VQ-VIB $_N$  models.  $\beta$ -VAE models used two linear layers, branching off the output of the feature extractor, to generate  $\mu$  and  $\sigma$  from which to sample a continuous latent variable. VQ-VIB $_C$  directly passed the output of the feature extractor into the vector quantization layer, from which the discrete latent representations were sampled, as described in the main paper. In VQ-VIB $_N$ , the output of the feature extractor was passed through two linear layers to generate a  $\mu$  and a  $\sigma$  (exactly as in the  $\beta$ -VAE case) before the sampled continuous representation was discretized via vector quantization. Across experiments, the only differences among encoder heads that could arise were due to different latent dimensions (although we fixed it to 32 for all experiments) or, for VQ-VIB $_C$  and VQ-VIB $_N$ , the number of elements in the learnable codebook or  $n$ , the number of quantized vectors to combine into a latent representation.
In the main paper, we discussed the VQ-VIB $_C$  training loss (Equation 1), maximizing utility, minimizing reconstruction loss, and minimizing the entropy of the categorical distribution over codebook elements. A strict generalization of Equation 1, in which a variational bound on the complexity of representations is also penalized, is included in Equation 2:
$$
\begin{array}{l} \max  \quad \lambda_ {U} \mathbb {E} [ U (x, y) ] - \lambda_ {I} \mathbb {E} [ \| x - \hat {x} \| ^ {2} ] \\ - \lambda_ {H} \mathbb {E} \left[ \sum_ {i \in [ 1, n ]} H \left(\mathbb {P} \left(\zeta \mid h _ {i} (x)\right) \right. \right] \tag {2} \\ - \lambda_ {C} \mathbb {E} \left[ D _ {\mathrm {K L}} \left[ \mathbb {P} (\zeta | h (x)) \| \mathcal {U} (C) \right] \right] \\ - \| \mathbf {s g} [ h _ {i} (x) ] - \zeta_ {i} (x) \| ^ {2} - \alpha \| h _ {i} (x) - \mathbf {s g} [ \zeta_ {i} (x) ] \| ^ {2} \\ \end{array}
$$
Equation 2 differs from Equation 1 via the third line, penalizing the KL divergence between the conditional categorical distribution over codebook elements and a uniform distribution over the  $C$  elements. This provides a variational bound on  $I(X,Z)$ , dubbed the complexity of representations in prior literature [36, 27]. In our main experiments, we set  $\lambda_C = 0$  and vary  $\lambda_H$ ; ablation studies in which we varied  $\lambda_C$  instead of  $\lambda_H$  are included in Appendix D and confirm that controlling the entropy of representations supported better finetuning accuracy.
In training  $\beta$ -VAEs, we trained to maximize the function described in Equation 3, where  $\mu(x)$  and  $\sigma(x)$  represent the  $\mu$  and  $\sigma$  parameters output by the encoder.
$$
\begin{array}{l} \max  \quad \lambda_ {U} \mathbb {E} [ U (x, y) ] - \lambda_ {I} \mathbb {E} [ \| x - \hat {x} \| ^ {2} ] \\ - \lambda_ {C} \mathbb {E} \left[ D _ {\mathrm {K L}} \left[ \mathcal {N} (\mu (x), \sigma (x)) \| \mathcal {N} (0, 1) \right] \right] \tag {3} \\ \end{array}
$$
Equation 3 trains agents to maximize classification accuracy, minimize MSE, and minimize the complexity of representations. The scalar weight  $\lambda_{C}$  can be viewed as a Lagrange multiplier, constraining how much information can be encoded in representations. This equation is closely related to Equation 2 but, given the continuous nature of encodings in  $\beta$ -VAE, we could not penalize the entropy of a categorical distribution.
In training VQ-VIB $_{\mathcal{N}}$ , we used the training objective proposed by Tucker et al. [27], which closely resembles the training loss for VQ-VIB $_{\mathcal{C}}$  and is shown in Equation 4 (and closely matches Equation 2). We use the same notation as for the  $\beta$ -VAE and VQ-VIB $_{\mathcal{C}}$  models.
$$
\begin{array}{l} \max  \quad \lambda_ {U} \mathbb {E} [ U (x, y) ] - \lambda_ {I} \mathbb {E} [ \| x - \hat {x} \| ^ {2} ] \\ - \lambda_ {H} \mathbb {E} \left[ \hat {H} \left(\mathbb {P} (\zeta | \mu (x))\right) \right] \tag {4} \\ - \lambda_ {C} \mathbb {E} \left[ D _ {\mathrm {K L}} \left[ \mathcal {N} (\mu (x), \sigma (x)) \| \mathcal {N} (0, 1) \right] \right] \\ - \| \mathbf {s g} [ h _ {i} (x) ] - \zeta_ {i} (x) \| ^ {2} - \alpha \| h _ {i} (x) - \mathbf {s g} [ \zeta_ {i} (x) ] \| ^ {2} \\ \end{array}
$$
The two key differences between Equation 4 and Equation 2 (used for training VQ-VIB $_C$ ) are bounds on entropy and complexity (on the second and third lines of Equation 4). Just as for  $\beta$ -VAEs, VQ-VIB $_N$  models use a KL divergence loss to regulate the complexity of representations. Increasing  $\lambda_C$ , as we did while annealing complexity in experiments, decreases the amount of encoded information. However, as shown in our results, simply increasing  $\lambda_C$  does not ensure that VQ-VIB $_N$  models will use fewer discrete representations. (For visualizations of this effect, see Appendix C.) Tucker et al. [27] advocate for using a small positive  $\lambda_H$  to penalize the estimated entropy over codebook elements. We explore varying  $\lambda_H$  in Appendix D and find some benefits relative to our main experiments, in which we set  $\lambda_H = 0$ . However, given that VQ-VIB $_N$  only supports an approximation of the entropy term, we find that controlling the entropy for VQ-VIB $_N$  is not as effective as controlling the entropy for VQ-VIB $_C$ .
A.2 Finetuning
In finetuning, we loaded pretrained frozen encoders and trained new predictor models to map from encodings to downstream predictions.
For a finetuning task with  $M$  distinct classes, and a "duplication factor,"  $k$ , for how many examples of each class to train with, we randomly selected  $k * M$  datapoints to train with. For example, when finetuning on the binary CIFAR100 task of living vs. non-living things,  $M = 2$ , so we loaded 2 total datapoints for  $k = 1, 4$  datapoints for  $k = 2$ , etc.. For each input in the finetuning dataset, we generated an encoding by passing through the encoder once. This generated a new dataset of encodings and labels, which we used the train the predictor. (Note that this approach is distinct from passing the input through the encoder many times; given stochastic encoders, which we used, the same input could result in many different encodings. Here, we assumed a limited budget of encodings.)
Predictor neural networks were instantiated as feedforward networks with 4 fully connected layers, with hidden dimension 256 and ReLU activations, and trained to map from shifted encodings (see
Table 1: Hyperparameters for FashionMNIST training.
ENCODER	LATENT DIM	C	n	λU	λI	λC0	λC INCR	λH0	λH INCR
β-VAE	32	NA	NA	10	10	0.01	0.5	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.2
VQ-VIBC	32	1000	2	10	10	0.0	0.0	0.001	0.4
VQ-VIBC	32	1000	4	10	10	0.0	0.0	0.001	0.8