SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Andi Peng*  
MIT
Mycal Tucker* MIT
Eoin M. Kenny  
MIT
Noga Zaslavsky  
UC Irvine
Pulkit Agrawal MIT
Julie A. Shah  
MIT
Neural networks often learn task-specific latent representations that fail to generalize to novel settings or tasks. Conversely, humans learn discrete representations (i.e., concepts or words) at a variety of abstraction levels (e.g., "bird" vs. "sparrow") and deploy the appropriate abstraction based on task. Inspired by this, we train neural models to generate a spectrum of discrete representations and control the complexity of the representations (roughly, how many bits are allocated for encoding inputs) by tuning the entropy of the distribution over representations. In finetuning experiments, using only a small number of labeled examples for a new task, we show that (1) tuning the representation to a task-appropriate complexity level supports the highest finetuning performance, and (2) in a human-participant study, users were able to identify the appropriate complexity level for a downstream task using visualizations of discrete representations. Our results indicate a promising direction for rapid model finetuning by leveraging human insight.
Neural networks learn implicit representations tailored to specific training tasks, but such representations, or abstractions, often fail to generalize to distinct test tasks. One approach to mitigating such generalization failures is based on the Information Bottleneck (IB) method [26, 1, 25, 24], which provides a framework for controlling how much information is passed through the network, effectively limiting its representational complexity. Unfortunately, it is difficult to know  $a$ -priori how much information to retain for optimal task performance.
Humans, however, when given a task, are remarkably adept at understanding which abstraction to deploy [35]. Consider an expert birdwatcher who has learned fine-grained classes of birds such as "sparrows," "goldfinches," and "kiwis." When accompanied by other experts, a birdwatcher knows to deploy their "most complex" abstraction for classifying birds; meanwhile, when at home with a 6-year old child, a birdwatcher knows to deploy a far simpler representation to communicate about a "red" vs. "yellow" bird. In other words, given a task, humans naturally adopt the right abstraction level that meaningfully captures task-relevant information and enables rapid learning [11, 5]. Inspired by humans, we seek to train neural nets along an axis of representational complexity and allow end users to select the desired complexity level to support data-efficient finetuning.
In this paper, we introduce a human-in-the-loop framework, depicted in Figure 1, for pre-training and finetuning complexity-regulated neural representations. Within this framework, we first generate a spectrum of complexity-controlled representations by training discrete information bottleneck methods [29] on a pre-training task. Second, we allow a human to specify a finetuning task, unknown
$^{1}$ Code available at github.com/mycal-tucker/human-guided-abstractions.
Figure 1: Our human-in-the-loop framework for pre-training and finetuning, illustrated for a bird-identification example. In pre-training, we generate a spectrum of encoders using representations from low to high complexity (e.g., just two crude categories to fine-grained species classifications). In fine-tuning, based on a desired task, a human selects an appropriate model for finetuning (e.g., crude categories for a child learning colors and fine-grained categories for a birder).
a priori, and select and finetune a pre-trained model. Given that humans specify the finetuning task, and may have to provide finetuning labels, few-shot adaptation is important.
In computational experiments, we find that finetuning performance is non-monotonically linked to representation complexity: representations that are too complex are data-inefficient, and representations that are too simple fail to capture important information. In a user study, we show that humans, given a desired finetuning task, can select high-performance models from a set of pre-trained models at different complexity levels. For example, as in Figure 1, a child performing a low-complexity task might select a low-complexity representation. More generally, while our computational experiments establish that finetuning performance is a function of model complexity, our study shows that humans can select (near) optimal complexity levels given a task for model finetuning. Lastly, the choice of neural architecture significantly affects finetuning efficiency: on one task, for example, our best-performing architecture, tuned to the right complexity, achieves better performance than other standard encoding methods with  $50 \times$  more finetuning data.
Our findings suggest that automatically constructing complexity-regulated abstractions during pretraining and then providing a diverse spectrum for human use for finetuning is a promising direction for human-in-the-loop few-shot adaptation. In summary, our contributions are (1) introducing a human-in-the-loop framework for automatically generating a spectrum of complexity-regulated abstractions for fast adaptation, (2) establishing that finetuning performance is a function of representation complexity, and (3) demonstrating the utility of our human-in-the-loop framework in a user study.
2.1 Abstraction in Human Cognition
There is substantial evidence that suggests that much of human learning, perception, communication, and cognition may be understood as compression of relevant information [31, 36, 11]. For example, fast human learning can be enabled by merging two or more instances of statistical patterns into one when appropriate [20]. Moreover, the simplicity of these patterns appears key to supporting their predictive power on downstream tasks [2]. This connection between compression and prediction provides an elegant explanation for why human brains have evolved to be such efficient compressors of information and experience [17]. Furthermore, visual abstractions have been found to prioritize functional properties, i.e. downstream task use, at the expense of visual fidelity, i.e. image reconstruction [11]. This suggests that different abstractions are constructed and deployed conditioned on tasks. Inspired by this, we seek to train neural networks to output visual abstractions that are also functionally useful to human users on a diverse range of downstream tasks.
2.2 Discrete Information Bottleneck
In our work, we build upon and compare to methods from prior research in discrete information bottlenecks. Generally, such work seeks to generate complexity-limited representations (roughly, limiting the number of bits about the input in a representation) using a finite set of representations (dubbed quantized vectors or prototypes). Recent work [27, 28, 7, 12] has approached this problem by combining ideas from the Variational Information Bottleneck [1, 9] with Vector Quantization (VQ) [29]. In such works, an encoder network outputs parameters to a normal distribution, from which a continuous latent variable is sampled, and the sample is discretized to the closest element of a learnable codebook. By penalizing the KL divergence between the normal and a fixed prior (typically a unit normal), one can limit the complexity of representations. We refer to this family of approaches as the Vector Quantized Variational Information Bottleneck - Normal  $(\mathrm{VQ - VIB}_{\mathcal{N}})$ . Other work proposes a different sampling mechanism before vector quantization, but experimental evaluation of these methods remains limited [22, 32].
In our work, we propose that complexity-constrained discrete representation learning can be used to generate a meaningful variety of encoders for a human-in-the-loop finetuning process. We use methods from prior work, and we propose a novel combination of entropy regularization and categorical sampling that, in experiments, supports the best finetuning performance.
3.1 Problem Formulation and Human-in-the-loop Framework
We consider a pre-training and finetuning problem, wherein a model is first trained on a pre-training dataset and must be rapidly adapted to a distinct finetuning task. Unlike classic meta-learning frameworks [6, 21, 19], we do not assume access to a distribution of tasks.
In pre-training, we assume access to a dataset of inputs and pre-training outputs,  $(X,Y_{p})$  (e.g., images of birds and species labels). That is,  $x\sim \mathbb{P}(X)$ ,  $y\sim \mathbb{P}_p(Y|X)$ . In finetuning, we assume access to a task-specific dataset with similarly-drawn inputs but novel task labels:  $(X,Y_{t})$ , for  $x\sim \mathbb{P}(X)$ ,  $y\sim \mathbb{P}_{t}(Y|X)$  (e.g., for the depicted girl's finetuning task, the color of the bird). Lastly, we assume that the pre-training labels are a sufficient statistic for the task-specific labels:  $I(X;Y_t) = I(Y_p;Y_t)$ . Intuitively, this states that the pre-training objective must include relevant information for the finetuning task. This is trivially satisfied with a reconstruction loss (i.e.,  $Y_{p} = X$ ) or fine-grained classifications (e.g., pre-training on exact species, but finetuning on groups of species).
Without information about the finetuning task, it is difficult, a priori, to identify how to pre-train a model to perform optimally on the finetuning task. Regularization methods like information bottleneck, for example, depend upon insight on the downstream task to specify the right level of regularization [1]. Rather than automatically identifying the right model, therefore, we propose a human-in-the-loop framework, depicted in Figure 1. Within our framework, we seek to generate a suite of pre-trained neural net encoders such that an end user can identify which one uses abstractions that support high performance for their desired task.
3.2 Technical Approach
Within our human-in-the-loop framework, it is critical to generate a "good" set of encoders from which the human chooses. This set must exhibit important variation (such that some encoders are better than others) and human-interpretability (such that humans can select the better encoders). We propose and validate in experiments that complexity-controlled discrete representation learning mechanisms achieve these desiderata. First, discrete representations can support human interpretability (via visualizations of the finite set of representations) [15, 18]. Second, controlling the complexity of representations is a meaningful axis of variation for encoders that likely supports human-specified tasks [34].
Figure 2: A VQ-VIB $_C$  encoder maps an input,  $x$ , to a representation,  $h$  in  $\mathbb{R}^Z$ , which is divided into  $n$  sub-representations,  $h_i$ . Each  $h_i$  is discretized stochastically by sampling based on distance to each quantized vector in the codebook,  $\zeta$ . Lastly, the sampled quantized vectors are concatenated for the full latent representation,  $z$ .
Figure 3: Different encoders train within an encoder-decoder-predictor framework via utility, entropy, and reconstruction losses [adapted from 27].
3.3 Neural Architecture Improvements
To generate a spectrum of encoders using representations at different complexity levels, we extend prior methods in discrete information bottleneck. We briefly present a novel neural method, which we dub the Vector-Quantized Variational Information Bottleneck - Categorical, or VQ-VIB $_C$ .
In VQ-VIB $_C$ , we combine information-theoretic losses from Tucker et al. [27]'s Vector Quantized Variational Information Bottleneck method (which we dub VQ-VIB $_N$ ) to emphasize that latent representations are drawn from a normal distribution) with categorical sampling mechanisms similar to Roy et al. [22] and Wu and Flierl [32]. Our VQ-VIB $_C$  encoder architecture is depicted in Figure 2. The encoder is parametrized by a feature extractor (in our cases, a standard feedforward neural network), an integer,  $n$ , representing how many quantized vectors to combine into a single latent representation, and  $C$  learnable quantized vectors  $\zeta^i; i \in [1, C], \zeta^i \in \mathbb{R}^{Z/n}$ . Using this architecture, the encoder is characterized as a function mapping from an input to a distribution over latent representations:  $q(z|x); x \in \mathbb{R}^X, z \in \mathbb{R}^Z$ . Concretely, the VQ-VIB $_C$  architecture deterministically maps from  $x$  to a hidden representation,  $h \in \mathbb{R}^Z$ . That hidden representation is divided into  $n$  vectors of equal size:  $h = [h_1, h_2, \dots, h_n]$ . Each  $h_i$  is then probabilistically quantized by sampling a quantized vector according to the L2 distance from  $h_i$  to each quantized vector:  $\mathbb{P}(z_i = \zeta^j | h_i) \propto e^{-||h_i - \zeta^j||^2}$ . (One can differentiate through sampling from this categorical distribution via the gumbel-softmax trick [13, 16].) Lastly, these  $n$  sampled discrete representations, each of which is dubbed  $z_i$ , are concatenated to form the latent representation,  $z = [z_1, z_2, \dots, z_n]$ .
We train VQ-VIB $_C$  via a combination of losses introduced by Tucker et al. [27]. We assume the VQ-VIB $_C$  encoder is trained with a decoder and a predictor, as depicted in Figure 3. Given a utility function,  $U$ , VQ-VIB $_C$  is trained to maximize the objective function in Equation 1:
$$
\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 (\mathbb {P} (\zeta | h _ {i} (x)) \right] \tag {1} \\ - \| \mathbf {s g} [ h _ {i} (x) ] - \zeta_ {i} (x) \| ^ {2} - \alpha \| h _ {i} (x) - \mathbf {s g} [ \zeta_ {i} (x) ] \| ^ {2} \\ \end{array}
$$
This objective trades off, in order in the equation: 1) maximizing the expected utility (e.g., cross-entropy loss), 2) minimizing the expected reconstruction loss (here, MSE), 3) minimizing the entropy of the distribution over codebook elements, and 4) minimizing a clustering loss from prior literature, encouraging encodings and quantized vectors to cluster (and, as in prior art, we leave  $\alpha = 0.25$  in all experiments) [29]. Here, sg represents the stop-gradient operation.
We include a more thorough discussion of this loss function, and comparisons to losses and architectures proposed in prior literature, in Appendix A. We emphasize that, while we propose some modifications to neural architectures, the primary contributions of this work are not about VQ-VIB $_C$  but rather: 1) our human-in-the-loop framework and 2) the recognition that discrete information bottleneck methods support high performance within this framework. In experiments, we found that our VQ-VIB $_C$  method performed better than methods from prior art, so we include this technical section to inform future researchers about modest architectural changes that support better performance.
4 Assessing Representational Complexity's Impact on Finetuning
We first present results for computational experiments in three different visual classification domains, indicating that finetuning performance varies as a function of representation complexity in a non-monotonic way. All experiments consisted of pre-training and finetuning phases. First, in pre-training on a low-level task, we generated a spectrum of models at varying complexity levels by varying loss hyperparameters. Second, we used the encoders from the pre-trained models and trained predictors to map from encodings to downstream predictions. Jointly, these steps enabled us to assess the importance of complexity-controlled representations for data-efficient finetuning. Overall, we found that, for small amounts of finetuning data, tuning representations to the right complexity was important, but with large amounts of data, any sufficiently-complex encoder performed similarly.
4.1 Domains
We trained agents on three image classification datasets: FashionMNIST [33], CIFAR100 [14], and iNaturalist 2019 (iNat) [10]. For FashionMNIST, we used a two-level hierarchy, grouping the 10 low-level classes into 3 higher-level classes (top = T-shirt/top, Pullover, Coat, and Shirt; shoes = Sandal, Sneaker, and Ankle Boot; other = Trouser, Dress, Bag). The CIFAR100 and iNat datasets have more extensive hierarchical structures, as detailed by, e.g., Sainte Fare Garnot and Landrieu [23]. For both datasets, we considered two levels of increasing crudeness in the hierarchy. For CIFAR100, we used both a 20-way division (each class consisting of 5 low-level classes) and a binary division for alive vs. non-alive objects. For iNat, we used a 34-way division (representing categories like "butterflies" and "mushrooms") and a 3-way division for plants, animals, or fungi. Thus, all domains were characterized by semantically-meaningful hierarchies.
4.2 Pretraining and Finetuning
Our experiments comprised two phases: pre-training and finetuning. In pre-training, we trained an encoder, decoder, and predictor on a low-level task. For example, for iNat, this consisted of classifying a photograph among 1010 species. During pre-training, after convergence to high accuracy, we decreased the complexity of representations by tuning a hyperparameter until representations were uninformative. For our  $\beta$ -VAE and VQ-VIB $_{\mathcal{N}}$  baselines, we increased  $\lambda_{C}$ , a scalar weight penalizing the KL divergence of the conditional Gaussian (generated by the encoder) from a unit Gaussian. (See Appendix A for details of training losses for  $\beta$ -VAE and VQ-VIB $_{\mathcal{N}}$ , including  $\lambda_{C}$ . In general, larger values of  $\lambda_{C}$  led to less complex representations and greater MSE values for reconstructions.) For VQ-VIB $_{\mathcal{C}}$ , we increased  $\lambda_{H}$ , the scalar weight penalizing the entropy of the categorical distribution over quantized vectors. In Appendix D, we examine other combinations of losses and architectures to control complexity, but found that tuning entropy for VQ-VIB $_{\mathcal{C}}$  achieved the best results. In general, to address some conflicting definitions of complexity in prior literature, we distinguish between an entropy loss, as introduced for our VQ-VIB $_{\mathcal{C}}$  model, and a complexity loss, where we use the definition of complexity as  $I(X;Z)$ .
In finetuning, we trained new predictor networks to map from encodings to classifications, using a small amount of training data. Using an encoder saved during pre-training, we generated encodings from inputs. New predictors were trained on classification tasks given the supervised (encoding, label) data. We generated  $k$  encodings for each class in the supervised dataset, and report results for different  $k$ . By varying  $k$  and which encoder was used to generate encodings (from high to low complexity), we could investigate the effect of encoder complexity on data-efficiency in finetuning. In all experiments, we pre-trained 5 models and ran 10 finetuning trials for each encoder. Further details on pre-training and finetuning are included in Appendix A.
4.3 Illustrative Example: FashionMNIST
We illustrate the high-level trends from our computational results using the FashionMNIST dataset. We trained models on the 10-way classification task to high accuracy (typically around  $90\%$ ) and decreased representational complexity until the model was accurate  $10\%$  of the time (random chance).
Depictions of the learned quantized vectors for  $\mathrm{VQ - VIB}_{\mathcal{C}}$  are included in Figure 4 at high and low complexities. At high complexity,  $\mathrm{VQ - VIB}_{\mathcal{C}}$  used a large number of distinct prototypes, including multiple prototypes per class (e.g., for two different types of handbags). This high complexity
(a) VQ-VIBc; high complexity
(b) VQ-VIB; low complexity
(c) Finetuning accuracy vs. MSE.
Figure 4: 2D PCA of VQ-VIB $_C$  latent representations for FashionMNIST at high (a) or low (b) complexity. Colorful points and images are prototypes, scaled in proportion to frequency of use. At high complexity, VQ-VIB $_C$  uses a large number of prototypes, including multiple per class (e.g., two types of handbags); at low complexity, classes merge (e.g. different types of shirts). c) When finetuned on a 3-way classification task with 1 example per class, VQ-VIB $_C$  outperformed  $\beta$ -VAE and VQ-VIB $_N$ , and less complex representations (greater MSE) improved model scores.
supported low MSE reconstruction loss. As we decreased the representational complexity,  $\mathrm{VQ - VIB}_C$  used fewer prototypes that represented more abstract concepts (Figure 4 b). For example, five distinct classes (T-shirt/top, Pullover, Coat, Shirt, and Dress) were merged into a single prototype, increasing MSE. Further examples of prototype evolution are in Appendix C.
In finetuning experiments, we loaded VQ-VIB $_C$  encoders across a range of complexities and finetuned a predictor on the 3-way classification task described earlier (tops, shoes, or other). Using just one example per class ( $k = 1$ ), VQ-VIB $_C$  models were more accurate than  $\beta$ -VAE or VQ-VIB $_N$  models, as shown in Figure 4 c. The  $x$  axis displays each encoder's MSE, which is a proxy for the inverse of the complexity of representations: more complex representations capture more information and generate better reconstructions (lower MSE). The  $y$  axis shows the finetuned predictor's accuracy on the 3-way task. While VQ-VIB $_C$  peak performance is roughly  $70\%$  accuracy (remarkably high, given  $k = 1$ ),  $\beta$ -VAE and VQ-VIB $_N$  performance peaks at  $40\%$  and  $55\%$ , respectively.
Beyond comparing VQ-VIB $_{C}$  to other models, Figure 4 c shows the importance of task-appropriate complexity levels. At low MSE values, VQ-VIB $_{C}$  uses many different prototypes, which impedes finetuning data efficiency. However, as VQ-VIB $_{C}$  learns more abstract representations (which increases MSE), finetuning performance improves (the correlation between MSE and accuracy is positive for  $\mathrm{MSE} < 0.055$  ( $p < 0.03$ )). However, past an MSE value around 0.055, VQ-VIB $_{C}$  lacks the representational capacity to distinguish between images that should be classified differently, and performance worsens.
Here, we discussed finetuning results for 1 labeled example per class  $(k = 1)$  and for  $n = 1$ , the number of quantized vectors to combine into representations, but results and analysis for  $k \in [1,2,5,10,50]$ ,  $n \in [1,2,4]$ , and for  $\beta$ -VAE, VQ-VIB $_{\mathcal{N}}$ , and VQ-VIB $_{\mathcal{C}}$  models are included in Appendix B. For all  $k$ , we found that VQ-VIB $_{\mathcal{C}}$  outperformed  $\beta$ -VAE and VQ-VIB $_{\mathcal{N}}$ . In fact, VQ-VIB $_{\mathcal{C}}$  with  $k = 1$  outperformed  $\beta$ -VAE for  $k = 50$ , indicating important architectural benefits. Increasing  $n$  supported higher complexity (lower MSE) but worse fine-tuning performance; intuitively, many discrete representations approximate continuous representations, which have worse sample complexity. This is an important limitation of combinatorial codebooks that others propose [28, 12].
A series of ablation studies confirm the importance of the VQ-VIB $_{C}$  architecture and annealing entropy, instead of complexity, for efficient finetuning (Appendix D, for FashionMNIST and other domains). In particular, we found that for VQ-VIB $_{C}$  and VQ-VIB $_{N}$ , penalizing entropy led to greater finetuning accuracy than when penalizing complexity, and VQ-VIB $_{C}$  outperformed VQ-VIB $_{N}$  when both were trained via entropy regularization. In other words, penalizing entropy enabled better finetuning performance, and VQ-VIB $_{C}$  supported better entropy regularization.
Lastly, given the importance of complexity on finetuning performance, we considered two methods for autonomously selecting the optimal complexity level. For low-data regimes ( $k \leq 5$ ), the simple heuristic of choosing the most complex encoder was clearly suboptimal; however, as the amount of finetuning data increased (e.g.,  $k = 50$ ), the most complex encoders achieved near-optimal
(a) CIFAR100 2-way  $k = 1$
(b) CIFAR100 2-way  $k = 50$
(c) CIFAR100 20-way  $k = 1$