SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 3: Fast adaptation to new scene understanding tasks. After pretraining, models are transferred to downstream tasks with the full dataset, but a small amount of computation: 1 epoch. Models perform the task either with a linear classifier (Frozen), end-to-end fine-tuning (E2E FT) or with our mechanism for in-context scene understanding (NN).
Method	Decoder	PASCAL ↑	ADE20K ↑
Frozen	E2E FT	Frozen	E2E FT
Supervised [67]	Linear	61.5	66.3	27.6	15.1
DINO [15]	Linear	54.9	64.0	25.6	23.4
MoCo-v3 [20]	Linear	41.2	4.8	14.6	3.2
MAE [30]	Linear	20.1	42.5	8.3	7.9
LOCA [13]	Linear	61.9	62.9	25.4	14.6
Hummingbird	NN	70.5	28.3
Hummingbird++	NN	72.1	30.5
Figure 4: Adaptation time of Hummingbird. The model is evaluated with retrieval-based evaluation ("NN retrieval"), linear probing ("Linear + frozen"), or full finetuning ("Linear + E2E FT").
on top of the frozen encoder in all cases. These results show that given an appropriately designed encoder, NN retrieval provides a data-efficient alternative to end-to-end finetuning, and is strictly more expressive than linear decoding.
Second, we verify the generality of these findings by comparing Hummingbird retrieval to several other representation learning algorithms which transfer to the low-data regime with finetuning (Table 2, see Appendix D.1 for higher-data regime and additional analysis). For PASCAL VOC, Hummingbird with the NN retrieval decoder outperforms the end-to-end finetuned version of all other techniques for both 1/128 (83 images) and 1/64 (165 images) of the data, which holds for both the purely self-supervised Hummingbird and its supervised variant. For ADE20K, Hummingbird is competitive with DINO [15] for 1/128 of the data (158 images) and outperformed by LOCA for 1/64 of the data (316 images), whereas Hummingbird++ outperforms all other models, demonstrating the benefit of retrieval-based supervision during pretraining. In summary, in the few-shot regime (e.g.  $\leq 100$  images) relevant for in-context learning, Hummingbird retrieval provides a compelling and robust alternative to end-to-end finetuning.
4.3 Fast adaptation to downstream tasks
While Hummingbird retrieval displays useful data-efficiency properties relative to fully finetuned methods, finetuning yields better performance when given access to the entire dataset. Yet even in this large-data regime, assistant systems must be quickly adaptable to new tasks. We thus evaluate the amount of computation required to reach good performance with the different decoding schemes from Section 4.2. All decoders are given the full training set and varying compute budgets. We titrate the amount of computation given to NN retrieval by partially populating the memory bank with fractions of the dataset. Figure 4 shows that 5 minutes (1 epoch through the downstream training set)
Table 4: Scene understanding with end-to-end finetuning. After pretraining, models are equipped with task-specific decoders and finetuned for that task on the entire downstream dataset.  $\dagger$  indicates results are taken from [30], using UperNet [77] as the decoder. Results for all other baselines are taken from [13] and use the linear decoder from [63]. For ViT-L results, see Appendix D.4.
Method	Encoder	Dataset	Fine-tuned accuracy (mIoU)
PASCAL ↑	ADE20K ↑
Random	ViT-B	IN1K	29.1	21.1
Supervised [67]	ViT-B	IN1K	76.1	47.3
DINO [15]	ViT-B	IN1K	74.1	44.1
MoCo-v3 [20]	ViT-B	IN1K	74.5	47.3†
BEiT [5]	ViT-B	IN1K+DALLE [54]	-	47.1†
MAE [30]	ViT-B	IN1K	75.0	48.1†
LOCA [13]	ViT-B	IN1K	76.7	47.9
Hummingbird	ViT-B	IN1K	80.0	44.9
Hummingbird++	ViT-B	IN1K	81.2	44.9
Hummingbird	ViT-B	IN22K	81.6	46.9
Hummingbird++	ViT-B	IN22K	82.1	48.2
Table 5: Ablation of pretraining components. Effect of training with spatial attention pooling (as opposed to mean pooling or a [CLS] token) and memory contextualization ("Cont.") on performance. All models were pretrained with ViT-B on ImageNet-1k.
Method	Pool.	Cont.	Semantic segmentation	Depth pred.
PASCAL ↑	ADE20K ↑	NYUv2 ↓
MoCLR [65]	mean	X	38.6	4.9	1.01
+ cont.	mean	✓	55.6	15.3	.901
+ [CLS]	[CLS]	X	64.5	23.9	.741
+ [CLS] + cont.	[CLS]	✓	65.6	25.1	.731
+ QK att. [52] + cont.	QK att.	✓	68.7	26.3	.728
+ QKV att.	QKV att.	X	68.0	27.4	.742
Hummingbird	QKV att.	✓	70.5	28.3	.718
are sufficient to build a performant NN decoder (70% mIoU on PASCAL VOC, 28% on ADE20K). In contrast, given the same amount of time, end-to-end finetuning still exhibits performance near chance, despite benefitting from hyperparameter tuning of learning rates, weight decay, and warm-up length. While a linear classifier converges more quickly than finetuning, it saturates with a significantly lower performance than NN retrieval (50% mIoU on PASCAL VOC, 20% on ADE20K).
We also quantify these benefits in terms of relative convergence: on PASCAL VOC, NN retrieval reaches the performance of full finetuning after 3 minutes rather than 3 hours, and the performance of a linear classifier in 2 seconds rather than 3 minutes. For ADE20K, the speedups are smaller, but significant: 7 minutes rather than 30 minutes (relative to full finetuning), and 1 minute rather than 30 minutes (relative to the linear classifier). By making substantial gains in this near-real-time use case, we believe NN retrieval lays the groundwork for scene understanding in an interactive setting.
Table 3 compares Hummingbird retrieval to other models (equipped with linear or end-to-end fine-tuned decoders) in the fast-adaptation regime (i.e. when given a single pass over the full downstream dataset): Hummingbird retrieval outperforms all other pretraining techniques and decoding mechanisms on both PASCAL VOC and ADE20K.
4.4 Fully finetuned scene understanding
Although the primary focus of this work is on fast and effortless adaption to downstream tasks, for completeness, we include a comparison of fully finetuned Hummingbird with fully finetuned state-of-the-art models on the semantic segmentation task. We follow the finetuning protocol of MAE [30] and use UperNet [77] as a decoder. Table 4 shows that both Hummingbird and Hummingbird++ are competitive with state-of-the-art when finetuned. Further analysis shows retrieval-based performance to be correlated with the finetuning performance (see Appendix D.2), paving the way for using retrieval-based evaluation as a model selection tool during training.
Ablating the pretraining components. We perform an ablation of pretraining components required for adaptation to downstream tasks through NN retrieval in Table 5. We find attention pooling to yield superior performance compared to mean pooling or a [CLS] token. Both contextual pretraining and attention pooling separately lead to large performance improvements over a baseline MoCLR [65] model and best results are achieved when combining the two. Note that although spatial attention pooling was initially introduced in the context of video understanding [52], this work is the first to show its utility for downstream task adaptation in the NN retrieval setup. We further find that modifying it ("QK att") with a value head ("QKV att.") improves its performance across all tasks.
Figure 5: Effect of memory bank length on nearest neighbor lookup at inference time. Inference time is for a single image. Lookups were done on a single Nvidia A100 GPU.
Figure 6: Effect of the pretraining (left) and evaluation (right) memory length on performance. All models were pretrained with ViT-B on ImageNet-22k. Left: Since the retrieval-based supervised objective is only defined for memory banks of non-zero length, for the purpose of this ablation we replace it with a simple linear classifier when  $|\mathcal{M}_p| = 0$ . Right: For downsample=True, we store representations of all patches into the memory bank. If downsample=True, we sample  $|\mathcal{M}| / N$  patches per image ( $N$  is the length of the downstream training set), allowing for greater diversity.
Effect of evaluation memory length  $|\mathcal{M}|$ . When transferring to downstream tasks with many training images (e.g. PASCAL VOC and ADE20K contain  $\sim 10\mathrm{k}$  and  $\sim 20\mathrm{k}$  images respectively, each image providing 100s of tokens), we see benefits of using large memory banks (e.g.  $|\mathcal{M}|$  of the order of 1-10 million tokens, see Figure 6, right). Since this makes the cross-attention operation computationally intractable, we leverage powerful libraries for approximate NN search [29, 40] to limit cross-attention (Equation 1) to a small set of nearest neighbors for each query (e.g.  $k = 30$ , see Appendix A.2 for details, where we find increasing  $k$  not to have significant impact on performance).
Figure 5 shows the relationship between evaluation memory length and the cost of the nearest neighbor lookup at inference time. For small-to-medium sized memory banks (0 to 1 million keys), the lookup cost is minimal ( $\leq 30$  ms), meaning the system is still fast enough to be used for real-time
applications, such as segmenting videos at 30 frames per second. When scaling to very large memory banks of 10 million keys or more, the scaling tends to be linear. However, the absolute performance is likely still suitable for most applications: with a memory bank size of 10 million, the overhead from NN lookup is only 0.2 seconds for dense tasks.
Effect of pretraining memory length  $|\mathcal{M}_p|$ . In contrast to retrieval-based evaluation, we find contextual pretraining to be remarkably memory-efficient: small memory banks (e.g.  $|\mathcal{M}_p| = 40\mathrm{k}$ , see Figure 6, left for PASCAL VOC and Appendix D.3 for ADE20K) are sufficient to yield robust gains in retrieval-based scene understanding, adding a relatively small computational overhead to training the representation (e.g.  $+22\%$  for  $|\mathcal{M}_p| = 40\mathrm{k}$ ). The module is agnostic to how the representation is trained and it benefits both self-supervised and supervised pretraining. Note that contextual pretraining is only present at training time and does not affect inference speed, and that pretraining and evaluation memory length are fully decoupled, allowing us to set them independently.
Inspired by impressive examples of in-context learning in language models, we investigate components necessary for in-context learning of dense scene understanding tasks in computer vision. To this end, we propose a simple non-parametric nearest neighbor retrieval mechanism—which is agnostic to the downstream task and requires no finetuning or specialized decoders—to serve as a general-purpose decoder which we use to evaluate models on semantic segmentation and monocular depth estimation tasks. We further propose Hummingbird, a pretraining method which benefits from attention across images (through contextual pretraining) and within an image (through spatial attention pooling) to produce image representations that can be easily configured to perform downstream tasks in a fast and data-efficient manner. By combining Hummingbird as the encoder with NN retrieval as the decoder, we take an important step towards in-context learning for dense vision tasks.
7 Broader Impact and Limitations
Broader impact. In laying the groundwork for scene understanding methods to be used in the interactive regime, our work could potentially benefit general-purpose assistants that are seeing rapid adoption. While these may enable a host of beneficial applications, they suffer from the biases and potential harms associated with visual language models and large language models more generally.
Limitations. Despite offering large relative improvements compared to finetuning and linear classification in the low-data regime, the absolute performance of Hummingbird when given less than 100 examples in the prompt is still far from perfect. To truly match the in-context learning abilities displayed in NLP, we would ideally need good performance from a handful of examples.
Further, given that retrieval-based scene understanding is task-agnostic, we leave expanding Hummingbird nearest neighbor evaluation to other tasks (e.g. object detection) to future work. Certain tasks, such as image rotation or image flipping, are not currently amenable to our framework as we assume a spatial correspondence between features and their labels. Explicitly post-processing to smooth outputs across patches was also not explored in this paper. Although the set of nearest neighbors and their weights vary smoothly as a function of the query image representation, it should be noted that this smoothness is dependent on the input prompt, since the memory bank needs to be sufficiently diverse and dense to allow a linear combination of neighbors to be expressive enough to cover all possible labels.
Finally, while we have showcased the benefits of attention within and across images in a contrastive framework, we defer adding them to more recent approaches using advanced data curation [51] and self-distillation [15, 84] to future work.
We thank Daniel Zoran, Andrew Zisserman, Evan Shelhamer and João Carreira for their thoughtful feedback, Skanda Koppula and Mathilde Caron for their assistance in reproducing baselines, and Aïron van den Oord and Oliver Vikbladh for fruitful discussions at the inception of the project.
[73] X. Wang, W. Wang, Y. Cao, C. Shen, and T. Huang. Images speak in images: A generalist painter for in-context visual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[74] Z. Wu, A. A. Efros, and S. X. Yu. Improving generalization via scalable neighborhood component analysis. In Proceedings of the European Conference on Computer Vision, 2018.
[75] Z. Wu, Y. Xiong, S. X. Yu, and D. Lin. Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018.
[76] J. Xiao, J. Hays, K. A. Ehinger, A. Oliva, and A. Torralba. SUN database: Large-scale scene recognition from abbey to zoo. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2010.
[77] T. Xiao, Y. Liu, B. Zhou, Y. Jiang, and J. Sun. Unified perceptual parsing for scene understanding. In Proceedings of the European Conference on Computer Vision, 2018.
[78] E. Xie, J. Ding, W. Wang, X. Zhan, H. Xu, P. Sun, Z. Li, and P. Luo. Detco: Unsupervised contrastive learning for object detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
[79] N. Xie, F. Lai, D. Doran, and A. Kadav. Visual entailment: A novel task for fine-grained image understanding. arXiv preprint arXiv:1901.06706, 2019.
[80] J. Yu, Z. Wang, V. Vasudevan, L. Yeung, M. Seyedhosseini, and Y. Wu. CoCa: Contrastive captioners are image-text foundation models. In Transactions on Machine Learning Research, 2022.
[81] L. Yuan, D. Chen, Y.-L. Chen, N. Codella, X. Dai, J. Gao, H. Hu, X. Huang, B. Li, C. Li, et al. Florence: A new foundation model for computer vision. arXiv preprint arXiv:2111.11432, 2021.
[82] Y. Zhang, K. Zhou, and Z. Liu. What makes good examples for visual in-context learning? In Advances in Neural Information Processing Systems, 2023.
[83] B. Zhou, H. Zhao, X. Puig, T. Xiao, S. Fidler, A. Barriuso, and A. Torralba. Semantic understanding of scenes through the ADE20K dataset. International Journal of Computer Vision, 127:302-321, 2019.
[84] J. Zhou, C. Wei, H. Wang, W. Shen, C. Xie, A. Yuille, and T. Kong. iBOT: Image BERT pre-training with online tokenizer. International Conference on Learning Representations, 2022.
A Implementation details: retrieval based scene understanding
A.1 How to store memories?
The memory bank only needs to be calculated once per dataset and can then be re-used for each of the images in the evaluation set. To populate the memory bank, each image in the dataset's training set (i.e. the "prompt") is encoded using the frozen backbone of the pretrained network to evaluate. We encode each of the training set images into a spatial map  $\pmb{k}_i = f_\theta(\pmb{x}_i) \in \mathbb{R}^{H \cdot W \times D}$ , where a feature  $\pmb{k}_i^j \in \mathbb{R}^D$  at a given spatial location  $j$  is aligned with the local label  $\pmb{l}_i^j$  created by averaging the pixel labels  $\pmb{y}_i^j$  in that patch. These features  $\pmb{k}_i$  are then  $L_2$ -normalized.
When the memory bank length is not large enough to accommodate all features for all images, it is necessary to subsample and only store a subset of the features of each image. For a concrete example using ADE20K, training set images have a resolution of  $512 \times 512$  which when encoded by a ViT-B/16 results in a  $32 \times 32$  grid of features (i.e. 1,024 features per image). To store every feature from each of ADE20K's 20,120 training images would require a memory bank length of  $20,120 \times 32 \times 32 = 20,695,040$ . When using data augmentation to increase the number of training images, the required length is even higher.
Our subsampling strategy for semantic segmentation works as follows. We define the number of features to take per image as  $n_{\text{features\_per\_image}} = \frac{|\mathcal{M}|}{|\mathcal{D}| * \text{num\_augmentation\_epochs}}$  where  $|\mathcal{D}|$  refers to the number of images in the training dataset. We thus sample the same number of features for each training image. Rather than sampling this number of features per image from the grid uniformly, we attempt to sample the most salient features using a simple strategy: upweighting patches containing class labels that appear less frequently in the image. Following the notation of Section 3.1, let  $l^j$  refer to the label attached to the patch indexed by  $j$  in the image and let  $\mathbb{1}_{c \in l^j} = 1$  if a given class  $c \in l^j$  and 0 otherwise. Then for each class  $c$  we define  $\kappa_c = \sum_j \mathbb{1}_{c \in l^j}$  (i.e., a count of how many patches the class  $c$  appeared in). We define a "class score" for each patch indexed by  $j$  as class_score  $= \sum_{c \in \mathcal{C}} \kappa_c \cdot \mathbb{1}_{c \in l^j}$ . Finally, we take the  $n_{\text{features\_per\_image}}$  from the spatial map  $k_i$  with the lowest final scores using
$$
\text {f i n a l} \quad \text {s c o r e} ^ {j} = \left(\text {c l a s s} \quad \text {s c o r e} ^ {j} \cdot x\right) + \left(1 0 ^ {6} \cdot \mathbb {1} _ {l ^ {j} = \emptyset}\right) \tag {5}
$$
where  $x \sim \mathcal{U}_{[0,1]}$ . The first term introduces some stochasticity into the sampling process and the second term deprioritizes locations that have no class label. The chosen features serve as the memory bank keys and their associated labels are the memory bank values.
The subsampling strategy used for depth estimation is simpler since there are no classes involved. We opted not to use data augmentation for this task making  $n_{\text{features\_per\_image}} = \frac{|\mathcal{M}|}{|\mathcal{D}|}$ . We first randomly order each patch in the image, then place all patches that contain no valid pixel labels after any patch with valid pixel labels, and then take the first  $n_{\text{features\_per\_image}}$  from the list.
There are many possible alternative strategies for sampling the most salient patches within an image in the event that the memory bank length cannot fit every feature from every image. We leave exploration of these possibly better sampling strategies for future work because in general we found this technique to perform well and wanted to show that nearest neighbor evaluation does not require complicated, hand-crafted strategies but rather works well out of the box with a simple heuristic calculated per image. For a complete listing of the hyperparameters involved in building and retrieving from the memory bank, see Appendix A.2.
A.2 How to recall memories?
After the memory bank has been populated as described in Appendix A.1, we sequentially make predictions for each image in the evaluation set. Evaluation was done on a single Nvidia A100 GPU per downstream task and takes approximately 15 minutes for PASCAL VOC, 25 minutes for ADE20K, and 30 minutes for NYUv2. Each image  $\pmb{x}$  is encoded as a grid of features  $\pmb{q} = f_{\theta}(\pmb{x})$  and each of the features from this grid will serve as the query that we will look up the nearest neighbors for. We use the open-source ScaNN library [29] to perform the approximate nearest neighbor search efficiently. ScaNN natively provides the functionality to return both the top-k nearest neighbors for a given query as well as scores for the similarity that can be used as the attention logits. These scores
are then divided by a temperature scaling value before having a softmax applied to them to obtain the final attention values (see Equation 1).
Throughout the paper, we use ScaNN in asymmetric hashing (AH) mode as opposed to brute-force mode. We find that there is little to no negative impact on the evaluation from using approximate nearest neighbor search as opposed to a brute-force exact search, despite the approximate search being several orders of magnitude faster. We use cosine similarity ( $L_{2}$ -normalized dot product) as a distance measure throughout this work. We also attempted some experiments using squared Euclidean distance and found it to have no benefits to performance for any of the models evaluated.
Table 6: NN retrieval hyperparameters. Note that no training is involved with NN evaluation, hence there are no hyperparameters such as learning rates or training epochs.
Section 4.2	Everywhere else
|M| (Memory bank length)	20,480,000	10,240,000
k (nearest neighbors)	90	30
Temperature	.1	.02
Augmentation epochs	8	2
ScaNN dimensions_per_block	4	4
ScaNN num_leaves	512	512
ScaNN num_leaves_to_search	256	32
ScaNN reordering_num_neighbors	1800	120
Table 6 summarizes the hyperparameters used for NN evaluation throughout this work. For every section except for Section 4.2, we use a flat set of hyperparameters detailed in the "Everywhere else" column of Table 6. Because Section 4.2 is concerned with small subsets of the data (i.e. training on the order of hundreds of images), hyperparameter sweeps are extremely cheap to run and it is computationally fast to find nearest neighbors even with minimal approximations, hence we used a slightly different setup in this regime. In general, we found nearest neighbor retrieval to be surprisingly robust to the choice of hyperparameters, with temperature and reordering_num_neighbors being the most relevant to performance. The same set of hyperparameters were used for the semantic segmentation tasks (PASCAL VOC and ADE20K) as for the monocular depth estimation task (NYUv2), with the exception of the number of augmentation epochs (we did not use augmentations for depth estimation). For a complete description of the meaning of the ScaNN hyperparameters, please see https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md.
Table 7 details the parameters used for augmenting the training dataset for semantic segmentation tasks. Note that the augmentations used to augment the training set when evaluating downstream tasks differ from the augmentations used for creating different views of the same image during contrastive pretraining described in Appendix C.1. When augmentations are enabled, the image is first scaled between the minimum and maximum scale factor, from which a random crop is selected. Then photometric augmentations are applied independently with the probabilities and maximum intensities provided.
Table 7: Evaluation augmentations. Parameters used to augment the training dataset for semantic segmentation.
Parameter
Random crop probability	1.0
Minimum scale factor	0.5
Maximum scale factor	2.0
Brightness jittering probability	0.5
Contrast jittering probability	0.5
Saturation jittering probability	0.5
Hue jittering probability	0.5
Brightness adjustment max	0.1
Contrast adjustment max	0.1
Saturation adjustment max	0.1
Hue adjustment max	0.1
B Implementation details: contextual pretraining