SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ivana Balažević* David Steiner* Nikhil Parthasarathy† Relja Arandjelović Olivier J. Henaff Google DeepMind
In-context learning—the ability to configure a model's behavior with different prompts—has revolutionized the field of natural language processing, alleviating the need for task-specific models and paving the way for generalist models capable of assisting with any query. Computer vision, in contrast, has largely stayed in the former regime: specialized decoders and finetuning protocols are generally required to perform dense tasks such as semantic segmentation and depth estimation. In this work we explore a simple mechanism for in-context learning of such scene understanding tasks: nearest neighbor retrieval from a prompt of annotated features. We propose a new pretraining protocol—leveraging attention within and across images—which yields representations particularly useful in this regime. The resulting Hummingbird model, suitably prompted, performs various scene understanding tasks without modification while approaching the performance of specialists that have been finetuned for each task. Moreover, Hummingbird can be configured to perform new tasks much more efficiently than finetuned models, raising the possibility of scene understanding in the interactive assistant regime.
In natural language processing (NLP), the pretrain-finetune paradigm has long been the dominant way of acquiring domain-specific knowledge and adapting a model's behavior to a particular task (e.g. question answering, natural language inference, summarization). More recently and predominantly due to the increase in model and dataset sizes, large language models have exhibited impressive, task-agnostic emergent capabilities [11, 37, 68], where a single model, given an appropriate prompt, can perform a wide range of downstream tasks without any change in its parameters.
While large-scale supervised and self-supervised pretraining in vision has yielded powerful encoders which capture useful semantics [15, 22, 31, 32, 35, 42, 43], applying these representations to solve downstream tasks has typically required bespoke decoders and end-to-end finetuning. The most readily applicable representations are trained for image-text alignment, enabling zero-shot classification [53] and image-based dialogue [2, 19, 80, 81], however these models are inherently limited by the coarseness of natural language outputs. Attempts have been made at casting fine-grained tasks (e.g. detection) as language modeling [17], but dense scene understanding tasks requiring millions of outputs do not lend themselves to this format. Indeed, deficiencies in fine-grained spatial understanding have been well documented in visual language models [36, 45, 64, 79].
In this work, we investigate the components required for in-context learning of scene understanding tasks, which we characterize along three axes: generality, data efficiency, and fast adaptation. To this end, we expand the well-known non-parametric nearest neighbor (NN) retrieval method [7, 9, 15, 75] to support a variety of dense scene understanding tasks. This retrieval-based decoding mechanism has the advantage of requiring no task-specific parameters or finetuning, thus enabling effortless adaption of standard encoders (e.g. ResNet [32] or ViT [22]) to any dense task of interest, as well as faster research iteration by allowing for simpler and more efficient model selection during pretraining.
*Equal contribution. †Current affiliation: NYU CNS, work done while interning at Google DeepMind. Correspondence to {balazevic, davidsteiner, henaff} @google.com.
We further show that the NN scene understanding capabilities of canonically pretrained vision transformers (such as MAE [30] and DINO [15]) vary greatly, despite similar finetuned performance. We find two pretraining components to yield reliable gains: (1) a simple modification to the standard self-supervised pretraining protocol, termed contextual pretraining, which performs attention across images by updating the spatial representation of each image with features retrieved from a memory bank, and (2) a spatial attention pooling mechanism (as opposed to the more standard mean pooling or the [CLS] token), which computes attention within an image to summarize the (contextualized) spatial grid of features into a single image-level representation to be fed into the self-supervised objective. We showcase the benefits of this approach in a standard contrastive framework, demonstrating large gains in NN scene understanding over prior pretraining methods.
Finally we find that our model, named Hummingbird due to its fast adaptation properties: (1) yields general-purpose representations which perform well in non-parametric semantic segmentation and monocular depth estimation using NN retrieval, (2) approaches the performance of fully finetuned models on some tasks, and (3) is more data-efficient and faster to adapt to new tasks when equipped with NN retrieval, compared to other pretraining methods and decoding mechanisms. By adapting quickly and efficiently to new tasks specified on the fly, Hummingbird raises the possibility of vision systems providing general-purpose assistants with in-context scene understanding.
Retrieval-based perception. Non-parametric evaluation has a long history with roots in the exemplar theory of human cognition [3, 38, 50] and case-based theories of artificial intelligence [1, 58]. In computer vision, non-parametric methods combined with simple features such as SIFT [48] and HOG [21] saw early success in image classification [9], shape matching [7, 8, 59], scene recognition [66, 76], and image parsing [44]. Exemplar-SVMs [49] showcased the versatility of non-parametric methods by retrieving arbitrary meta-data (such as segmentations, geometry, even 3D models) from training examples. We leverage these insights with modern architectures and training paradigms coupled with dense retrieval.
Retrieval-based training. To improve retrieval-based performance at test time, retrieval-based classifiers [69, 74] shape their representations for this task, enabling fine-grained classification from coarse supervision. While not explicitly training for it, DINO [15] witnessed NN classification abilities emerge from self-supervised training of vision transformers, enabling global retrieval tasks such as landmark recognition and copy detection. In [72], tracking abilities emerge after pretraining on a colorization task via retrieval from reference frames of a video. Retrieval has also been proposed as a means of enriching the positive pairs used in self-supervised contrastive learning [24]. These works differ from ours in that they encode and retrieve global representations of entire images, in contrast to the local inferences required by dense scene understanding tasks.
Fast adaptation. A number of methods have tackled the problem of adapting to newly specified tasks, most often from the perspective of meta-learning. For example, matching networks [71] and MAML [26] learn to solve new classification and reinforcement learning tasks specified on the fly. Architectural innovations, such as image prompting [4, 39, 82] and adapter layers [27, 55] have also facilitated transfer to new image recognition tasks. While fast adaptation to dense scene understanding tasks has been less studied, image inpainting [6, 73] and VTM [41] have made progress in this direction, particularly in the low-data regime. These approaches differ from ours in that they achieve fast adaptation by training on related dense tasks and (in the case of VTM) adapt to downstream tasks with task-specific weight updates and learned similarity functions. In contrast, we maintain the simplicity of pure retrieval-based approaches by adapting to new downstream tasks without modifying any model parameters, and the generality of self-supervised approaches by learning representations from generic pretraining data with no dense annotations.
Self-supervised learning. Methodologically, our representation learning method is most similar to self-supervised learning (SSL) techniques. Similarly to NLP, image-based SSL has witnessed great success in recent years, notably with the advent of contrastive methods [14, 15, 16, 23, 33], self-distillation [12, 18, 28], and masked auto-encoding [30]. Due to their conceptual simplicity, we base our method on standard contrastive baselines such as SimCLR [16] and MoCo [31]. Image-based SSL techniques have since been tailored to learning representations which transfer well to scene understanding tasks [13, 33, 70, 78], and although they have been shown to support zero-shot object discovery [34, 61], they generally still require task-specific decoders and end-to-end finetuning.
Figure 1: In-context scene understanding with nearest neighbor retrieval. On the left, we provide the system with a "prompt" of annotated images. On the right, we ask the system to describe new query images. The network computes dense features for each location and uses them to query features computed from the prompt. The labels associated with the nearest prompt features are then aggregated to make predictions about the query. Note that the system makes no assumptions about the nature of the labels, and as such can be used to solve a variety of different scene understanding tasks in-context. The nearest neighbors and predictions in this example are computed with our Hummingbird model.
The following sections describe the retrieval-based scene understanding decoding protocol (Section 3.1), followed by the contextual pretraining method (Section 3.2) and the self-supervised (Section 3.3) and supervised learning objectives (Section 3.4). We use subscripts  $\boldsymbol{x}_i$  to differentiate between representations and superscripts  $\boldsymbol{x}^j$  to denote spatial locations within a representation.
3.1 Retrieval-based scene understanding
A general-purpose image representation should perform well across a variety of scene understanding tasks out-of-the-box, i.e. without modifying its parameters. To test whether a representation satisfies this condition, we extend the standard image-level nearest neighbor (NN) retrieval [7, 9] decoding mechanism to dense, patch-level retrieval (with patch size set to  $16 \times 16$  across all models in this work). Given a prompt composed of training images from the downstream task and their corresponding labels  $\{(\pmb{x}_i, \pmb{y}_i), i = 1, \dots, N, \pmb{x}_i \in \mathbb{R}^{H' \times W' \times C}\}$ , our aim is to enable a pretrained image encoder  $f_{\theta}$  to make predictions about a new image  $\pmb{x}$  from the test set. In tasks considered in this work, labels  $\pmb{y}_i$  are spatial maps of either class labels  $\pmb{y}_i \in \mathbb{C}^{H' \times W'}$  (e.g. for semantic segmentation, where  $\mathbb{C}$  is the space of all classes) or scalars  $\pmb{y}_i \in \mathbb{R}^{H' \times W'}$  (e.g. for monocular depth estimation).
We encode each prompt image into a spatially flattened 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 spatial location  $j$  is aligned with the local label  $l_i^j$  created by averaging the pixel labels  $\pmb{y}_i^j$  of a patch. We then sample a subset of features and local labels for each image, which form the keys and values of the memory bank  $\mathcal{M} = \{(k_i^j, l_i^j), i = 1, \dots, N, j \sim \mathcal{S}\}$  (see Appendix A.1 for details on the sampling distribution  $S$ ). In the following, we do not distinguish between entries from different images, and use a single integer  $j$  to index into the memory bank:  $\mathcal{M} = \{(k^j, l^j), j = 1, \dots, |\mathcal{M}|\}$ .
Given a test image  $\pmb{x}$ , we form a representation  $\pmb{q} = f_{\theta}(\pmb{x})$  and use each spatial feature  $\pmb{q}^i$  as a query to cross-attend over the memory bank with temperature  $\beta$ . The cross-attention weights are then used to combine the corresponding labels and form a local prediction  $\hat{l}^i$ :
$$
s ^ {i, j} = \frac {1}{\beta} \frac {\left\langle \boldsymbol {q} ^ {i} , \boldsymbol {k} ^ {j} \right\rangle}{\| \boldsymbol {q} ^ {i} \| \| \boldsymbol {k} ^ {j} \|}, \quad \boldsymbol {a} ^ {i} = \underset {j} {\operatorname {s o f t m a x}} (\boldsymbol {s} ^ {i}), \quad \hat {\boldsymbol {l}} ^ {i} = \sum_ {j} a ^ {i, j} \boldsymbol {l} ^ {j}. \tag {1}
$$
Equation 1 defines the cross-attention operation as  $\hat{l}^i = \mathrm{CA}(q^i, k^j, l^j)$ . The final prediction  $\hat{y}$  is simply the concatenation of local predictions  $\hat{l}^i$  upsampled to the original image size via bilinear interpolation. As a result, nearest neighbor retrieval allows a simple image encoder to perform scene understanding tasks without any decoders or parameter adaptation (finetuning or otherwise) to the downstream dataset. The mechanism is also entirely agnostic to the format of the labels, enabling it to perform tasks as diverse as semantic segmentation and depth estimation.
Figure 2: Hummingbird model components.
3.2 Contextual pretraining
Memory retrieval allows an image encoder to perform various tasks by combining labels of nearby examples. To ensure that a model will perform well in this regime, we propose to train it in a similar manner, by enforcing its representation to be expressed as a combination of representations of nearby examples. Over the course of training, we populate a memory bank  $\mathcal{M}_p = \{(k_i,\boldsymbol {v}_i),i = 1,\dots,|\mathcal{M}_p|\}$  with spatially averaged keys and values computed from training images  $\boldsymbol{x}_i$  from previous batches:
$$
\boldsymbol {h} _ {i} = f _ {\theta} \left(\boldsymbol {x} _ {i}\right) \in \mathbb {R} ^ {H \cdot W \times D}, \quad \boldsymbol {k} _ {i} = \frac {1}{H \cdot W} \sum_ {j = 1} ^ {H \cdot W} \boldsymbol {h} _ {i} ^ {j} \in \mathbb {R} ^ {D}, \quad \boldsymbol {v} _ {i} = \phi_ {\theta} \left(\boldsymbol {k} _ {i}\right) \in \mathbb {R} ^ {D}, \tag {2}
$$
where we use an MLP as the value head  $\phi_{\theta}$  (see Appendix B for implementation details). We then form a representation  $q = f_{\theta}(\pmb{x})$  of a new training image  $\pmb{x}$  and use each spatial feature  $\pmb{q}^i$  to attend over the memory bank and compute an update  $\hat{\pmb{v}}^i = \mathrm{CA}(\pmb{q}^i, \pmb{k}_j, \pmb{v}_j)$ . Each feature is "contextualized" as  $\pmb{c}^i = \psi_{\theta}((1 - \lambda)\frac{\pmb{q}^i}{\|\pmb{q}^i\|} + \lambda \frac{\hat{\pmb{v}}^i}{\|\hat{\pmb{v}}^i\|})$ , where  $\psi_{\theta}$  is a linear layer and  $\lambda$  a weighting parameter. The contextualized image representation  $c = g_{\theta}(\pmb{q}, \mathcal{M}_p)$  is simply the concatenation of local features  $c^i$ .
Note that the pretraining memory bank  $\mathcal{M}_p$  is discarded at test time and differs from the test time memory bank  $\mathcal{M}$  described in Section 3.1, allowing for straightforward comparison of our representations  $f_{\theta}(\boldsymbol{x})$  to those trained without the memory bank.
3.3 Self-supervised objective
While contextual pretraining updates representations by attending across images, we hypothesize that learning to attend within images will also enable fine-grained predictions required by dense tasks. To that end, we train representations to locate the most distinctive part of an image using a combination of attention pooling and contrastive learning. Following [16, 28], we construct different views of unlabeled images  $\pmb{x}$  through random data augmentation  $\pmb{x}_1\sim \mathcal{A}_1(\pmb {x}),\pmb{x}_2\sim \mathcal{A}_2(\pmb {x})$  , see Appendix C.1. Each view is encoded as  $h_i = f_\theta (x_i)$  and further contextualized with the mechanism described above as  $c_{i} = g_{\theta}(h_{i},\mathcal{M}_{p})\in \mathbb{R}^{H\cdot W\times D}$  (see Figure 2). Following [52], we compute attention pooled representations  $\hat{\pmb{c}}_i\in \mathbb{R}^D$  using masks  $\pmb{m}_{i}$  derived from a lightweight attention module  $a_{\theta}$  , which we augment with an additional value head  $\omega_{\theta}$  . Pooled features are then used to compute projections  $\pmb{z}_i^\theta$  ..
$$
\boldsymbol {m} _ {i} = \underset {j} {\operatorname {s o f t m a x}} \left(a _ {\theta} \left(\boldsymbol {c} _ {i}\right)\right), \quad \hat {\boldsymbol {c}} _ {i} = \sum_ {j = 1} ^ {H \cdot W} m _ {l} ^ {j} \omega_ {\theta} \left(\boldsymbol {c} _ {i} ^ {j}\right), \quad \boldsymbol {z} _ {i} ^ {\theta} = p _ {\theta} (\hat {\boldsymbol {c}} _ {i}). \tag {3}
$$
Finally, following [20, 28, 65], each view forms predictions  $q_{\theta}(z_i^\theta)$  about the other view's targets  $z_j^\xi$ , which are computed with the same architecture and a different set of weights  $\xi$  which vary more slowly (see Appendix C.2). The online weights  $\theta$  are optimized using a standard contrastive loss:
$$
\mathcal {L} _ {\mathrm {S S L}} ^ {i j} (\theta ; \xi) = - \log \frac {\exp \left(q _ {\theta} \left(\boldsymbol {z} _ {i} ^ {\theta}\right) \cdot \boldsymbol {z} _ {j} ^ {\xi}\right)}{\exp \left(q _ {\theta} \left(\boldsymbol {z} _ {i} ^ {\theta}\right) \cdot \boldsymbol {z} _ {j} ^ {\xi}\right) + \sum_ {k} \exp \left(q _ {\theta} \left(\boldsymbol {z} _ {i} ^ {\theta}\right) \cdot \boldsymbol {z} _ {k} ^ {\xi}\right)}. \tag {4}
$$
Table 1: In-context scene understanding. All models are pretrained on source data in a supervised or self-supervised manner, and applied to downstream datasets without modification. All downstream tasks are performed using a single mechanism, nearest neighbor retrieval.  $\dagger$  indicates our reproduction of external work, all other models were evaluated using publicly available checkpoints.
Method	Encoder	Params (M)	Dataset	Semantic segmentation	Depth pred.
PASCAL ↑	ADE20K ↑	NYUv2 ↓
Supervised†	ViT-B	86	IN1K	35.1	13.8	.913
DINO [15]	ViT-B	86	IN1K	55.9	21.8	.793
MoCo-v3 [20]	ViT-B	86	IN1K	37.2	14.6	.771
MAE [30]	ViT-B	86	IN1K	6.6	3.3	.981
LOCA [13]	ViT-B	86	IN1K	57.5	18.5	.880
Hummingbird	ViT-B	86	IN1K	70.5	28.3	.718
Hummingbird++	ViT-B	86	IN1K	72.1	30.5	.738
Supervised†	ViT-B	86	IN22K	63.5	28.0	1.07
MAE† [30]	ViT-B	86	IN22K	9.8	4.2	.968
LOCA [13]	ViT-B	86	IN22K	56.4	16.8	.829
Hummingbird	ViT-B	86	IN22K	73.5	30.7	.706
Hummingbird++	ViT-B	86	IN22K	76.2	34.1	.695
Comparison across architectures:
CLIP† [53]	NFNet-F6 [10]	438	ALIGN	57.2	25.0	.844
Supervised†	NeXt-XL [46]	1300	IN22K	58.9	25.5	.791
Supervised†	ViT-L	307	IN22K	65.8	26.1	.860
MAE [30]	ViT-L	307	IN1K	8.0	3.6	.934
LOCA [13]	ViT-L	307	IN22K	59.5	17.6	.912
Hummingbird	ViT-L	307	IN22K	76.9	35.0	.671
Hummingbird++	ViT-L	307	IN22K	77.3	35.8	.671
3.4 Retrieval-based supervised objective
Given the availability of large labeled datasets, and noting that correctly designed supervision does not necessarily hurt generalization [57], we explore the use of label-supervision for learning representations that perform well in dense NN retrieval. While supervision is typically added with a linear classifier atop average pooled features [32, 43, 62], we instead use it to constrain contextual pretraining and further align our training methodology with NN retrieval [69, 74]. Specifically, we expand the memory bank  $\mathcal{M}_p$  to include the labels:  $\mathcal{M}_p' = \{(k_i, v_i, y_i), i = 1, \dots, |\mathcal{M}_p'|\}$  and query it with attention pooled features  $\hat{\pmb{c}}_i \in \mathbb{R}^D$  (see Equation 3) to form predictions  $\hat{y}_i = \mathrm{CA}(\hat{c}_i, k_j, y_j)$ . We then use the softmax cross entropy loss  $\mathcal{L}_{\mathrm{CE}}^i(\hat{y}_i, y_i)$ , which added to the self-supervised objective of Equation 4, forms the total loss  $\mathcal{L}^{ij} = \mathcal{L}_{\mathrm{SSL}}^{ij} + \alpha (\mathcal{L}_{\mathrm{CE}}^i + \mathcal{L}_{\mathrm{CE}}^j)$ , with supervised weight  $\alpha$ . Note that the memory bank  $\mathcal{M}_p'$  is only used during training and the added supervision relates to a global image classification task, not the downstream pixel-level tasks.
We demonstrate the generality of Hummingbird representations through retrieval-based scene understanding on several downstream tasks (Section 4.1): semantic segmentation on PASCAL VOC [25] and ADE20K [83] with mean IoU (mIOU) as metric, and monocular depth estimation on NYUv2 [60] with root-mean-square error (RMSE) as metric. We further show that, in the low-data regime (Section 4.2) and when looking at adaptation speed (Section 4.3), Hummingbird with NN retrieval outperforms other pretraining techniques and decoding mechanisms, including end-to-end finetuning. Section 4.4 compares the performance of fully finetuned Hummingbird with prior work.
4.1 Retrieval-based scene understanding
We consider the performance of learned representations in the retrieval-based scene understanding setup described in Section 3.1 across architecture (ViT-B and ViT-L) and dataset (ImageNet-1k and
Table 2: Data-efficient scene understanding. After pretraining, models are adapted to downstream tasks on small amounts of data with end-to-end fine-tuning with a linear head (E2E FT) or with nearest neighbor retrieval (NN).  $n$  refers to the number of images a fraction represents. All runs are averaged over five different seeds, with standard deviation of the order of  ${0.04}/{0.10}\%$  for NN and  ${0.36}/{1.35}\%$  for E2E FTon PASCAL / ADE20K. Each method is trained with ViT-B on ImageNet-1k.
Method	Decoder	PASCAL ↑	ADE20K ↑
1/128 (n=83)	1/64 (n=165)	1/128 (n=158)	1/64 (n=316)
Supervised [67]	E2E FT	41.8	53.8	10.8	14.3
DINO [15]	E2E FT	36.1	44.3	11.7	14.4
MoCo-v3 [20]	E2E FT	19.9	33.4	4.6	7.9
MAE [30]	E2E FT	34.2	44.1	8.2	12.2
LOCA [13]	E2E FT	40.1	53.9	11.2	15.5
Hummingbird	NN	50.5	57.2	11.7	15.1
Hummingbird++	NN	52.4	57.3	12.7	16.4
Figure 3: Data efficiency of Hummingbird. The model is evaluated with retrieval-based evaluation ("NN retrieval"), linear probing ("Linear + frozen"), or full finetuning ("Linear + E2E FT").
-22k [56]) scales, trained with supervision (Hummingbird++) or without (Hummingbird). Figure 1 shows an example prediction made by Hummingbird with a ViT-B encoder on PASCAL VOC.
The top part of Table 1 shows an apples-to-apples comparison of Hummingbird with existing methods for pretraining ViT-B encoders, where it outperforms all baselines by a large margin. We also note that Hummingbird scales well with increasing the dataset size from ImageNet-1k to ImageNet-22k, which does not hold for all other methods (e.g. MAE, consistent with [51]). Further, training with supervision is generally beneficial, particularly for semantic segmentation. For an ablation on the impact of retrieval-based supervision on performance, see Appendix D.5.
The bottom part of Table 1 contains the best performing methods across architectures, showing a performance increase for Hummingbird with encoder size. Note that results achieved by Hummingbird retrieval on PASCAL VOC and ADE20K, without any finetuning, approach the performance of methods fully finetuned on each of those tasks with specialized decoders (see Table 4).
4.2 Data-efficient retrieval-based scene understanding
In addition to adapting to downstream tasks with minimal (or ideally no) alterations to the model, a second ingredient for in-context learning is adaptation given only a limited number of examples.
We therefore evaluate the performance of Hummingbird retrieval in the low-data regime, and compare it with other decoding techniques: linear probing and end-to-end finetuning (Figure 3). For PASCAL VOC, NN retrieval outperforms the end-to-end finetuning for up to  $1/8$  of the data ( $\sim 1300$  images). For ADE20K the effect is less pronounced, however NN retrieval still exceeds end-to-end finetuning when given up to  $1/32$  of the data ( $\sim 600$  images). Hummingbird retrieval outperforms linear decoding