Title: Towards In-context Scene Understanding
Venue: NeurIPS 2023
Accepted: yes
Keywords: transfer learning

adaptation

self-supervised learning

contrastive learning

scene understanding

representation learning

in-context learning

vision transformers
Abstract: 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.

Towards In-context Scene Understanding

Abstract

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.

1 Introduction

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.

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.

2 Related Work

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.

Find nearest neighbors

Aggregate their labels

Input prompt images

Input prompt labels

Query images

Predictions

3 Method

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 x i to differentiate between representations and superscripts 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 × 16 across all models in this work). Given a prompt composed of training images from the downstream task and their corresponding labels { ( x i , y i ) , i = 1 , ..., N, x i ∈ R H ′ × W ′ × C } , our aim is to enable a pretrained image encoder f θ to make predictions about a new image x from the test set. In tasks considered in this work, labels y i are spatial maps of either class labels y i ∈ C H ′ × W ′ (e.g. for semantic segmentation, where C is the space of all classes) or scalars y i ∈ R H ′ × W ′ (e.g. for monocular depth estimation).

We encode each prompt image into a spatially flattened map k i = f θ ( x i ) ∈ R H · W × D , where a feature k j i ∈ R D at a spatial location j is aligned with the local label l j i created by averaging the pixel labels y j i 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 M = { ( k j i , l j i ) , i =1 , ..., N, j ∼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: M = { ( k j , l j ) , j =1 , ..., |M|} .

Given a test image x , we form a representation q = f θ ( x ) and use each spatial feature q i as a query to cross-attend over the memory bank with temperature β . The cross-attention weights are then used to combine the corresponding labels and form a local prediction ˆ l i :

Equation 1 defines the cross-attention operation as ˆ l i = CA ( q i , k j , l j ) . The final prediction ˆ y is simply the concatenation of local predictions ˆ 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.

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 M p = { ( k i , v i ) , i = 1 , ..., |M p |} with spatially averaged keys and values computed from training images x i from previous batches:

where we use an MLP as the value head ϕ θ (see Appendix B for implementation details). We then form a representation q = f θ ( x ) of a new training image x and use each spatial feature q i to attend over the memory bank and compute an update ˆ v i = CA ( q i , k , v )

j j . Each feature is 'contextualized' as c i = ψ θ ((1 -λ ) q i ∥ q i ∥ + λ ˆ v i ∥ ˆ v i ∥ ) , where ψ θ is a linear layer and λ a weighting parameter. The contextualized image representation c = g θ ( q , M p ) is simply the concatenation of local features c i .

Note that the pretraining memory bank M p is discarded at test time and differs from the test time memory bank M described in Section 3.1, allowing for straightforward comparison of our representations f θ ( 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 x through random data augmentation x 1 ∼ A 1 ( x ) , x 2 ∼ A 2 ( x ) , see Appendix C.1. Each view is encoded as h i = f θ ( x i ) and further contextualized with the mechanism described above as c i = g θ ( h i , M p ) ∈ R H · W × D (see Figure 2). Following [52], we compute attention pooled representations ˆ c i ∈ R D using masks m i derived from a lightweight attention module a θ , which we augment with an additional value head ω θ . Pooled features are then used to compute projections z θ i :

Finally, following [20, 28, 65], each view forms predictions q θ ( z θ i ) about the other view's targets z ξ j , which are computed with the same architecture and a different set of weights ξ which vary more slowly (see Appendix C.2). The online weights θ are optimized using a standard contrastive loss:

contrastive

encoder

attention

pool

encoder

objective

encoder +

mean pool

contextualize

augment

augment

contextualize

view 1

view 2

hiddens

contextualized

hiddens

memory bank

attention

pool

downstream tasks

Confidential - DeepMind

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 M p to include the labels: M ′ p = { ( k i , v i , y i ) , i = 1 , ..., |M ′ p |} and query it with attention pooled features ˆ c i ∈ R D (see Equation 3) to form predictions ˆ y i = CA ( ˆ c i , k j , y j ) . We then use the standard softmax cross entropy loss L i CE ( ˆ y i , y i ) , which added to the self-supervised objective of Equation 4, forms the total loss L ij = L ij SSL + α ( L i CE + L j CE ) , with supervised weight α . Note that the memory bank M ′ p is only used during training and the added supervision relates to a global image classification task, not the downstream pixel-level tasks.

4 Experiments

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

10 1

10 2

10 3

10 4

Number of adaptation images

10

20

30

40

50

60

70

PASCAL semantic segmentation

Linear + E2E FT

Linear + frozen

NN retrieval

10 1

10 2

10 3

10 4

Number of adaptation images

5

10

15

20

25

30

35

40

ADE20K semantic segmentation

-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 ( ∼ 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 ( ∼ 600 images). Hummingbird retrieval outperforms linear decoding

200ms

1s

10s

1min

10min

1h

Adaptation time (1 TPU)

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

PASCAL semantic segmentation

200ms

1s

10s

1min

10min

1h

Adaptation time (1 TPU)

0.00

0.05

0.10

0.15

0.20

0.25

ADE20K semantic segmentation

Linear + E2E FT

Linear + frozen

NN retrieval

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. ≤ 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)

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 finetuned 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-ofthe-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.

5 Analysis

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.

40,000

160,000

640,000

2,560,000

10,240,000

Memory bank length

10

25

50

100

200

400

Runtime (ms)

NN lookup time vs. memory bank length

0

50000

100000

150000

200000

250000

300000

Training memory length

69

70

71

72

73

74

75

76

PASCAL: Importance of pretraining memory

self-supervised

supervised

10 5

10 6

10 7

Evaluation memory length

45

50

55

60

65

70

PASCAL: Importance of evaluation memory

downsample=True

downsample=False

Effect of evaluation memory length |M| . When transferring to downstream tasks with many training images (e.g. PASCAL VOC and ADE20K contain ∼ 10 k and ∼ 20 k images respectively, each image providing 100s of tokens), we see benefits of using large memory banks (e.g. |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 ( ≤ 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 |M p | . In contrast to retrieval-based evaluation, we find contextual pretraining to be remarkably memory-efficient: small memory banks (e.g. |M p | = 40k, 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 |M p | = 40k). 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.

6 Conclusion

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.

References

A. Aamodt and E. Plaza. Case-based reasoning: Foundational issues, methodological variations, and system approaches. AI Communications , 7(1):39-59, 1994.

J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y . Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: A visual language model for few-shot learning. Advances in Neural Information Processing Systems , 2022.

F. G. Ashby and W. T. Maddox. Human category learning. Annu. Rev. Psychol. , 56:149-178, 2005.

H. Bahng, A. Jahanian, S. Sankaranarayanan, and P. Isola. Exploring visual prompts for adapting large-scale models. arXiv preprint arXiv:2203.17274 , 2022.

H. Bao, L. Dong, and F. Wei. BEiT: BERT pre-training of image transformers. In International Conference on Learning Representations , 2022.

A. Bar, Y. Gandelsman, T. Darrell, A. Globerson, and A. Efros. Visual prompting via image inpainting. Advances in Neural Information Processing Systems , 2022.

S. Belongie, J. Malik, and J. Puzicha. Shape matching and object recognition using shape contexts. IEEE Transactions on Pattern Analysis and Machine Intelligence , 24(4):509-522, 2002.

A. C. Berg, T. L. Berg, and J. Malik. Shape matching and object recognition using low distortion correspondences. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2005.

O. Boiman, E. Shechtman, and M. Irani. In defense of nearest-neighbor based image classification. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2008.

A. Brock, S. De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization. In International Conference on Machine Learning , 2021.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems , 2020.

M. Caron, P. Bojanowski, A. Joulin, and M. Douze. Deep clustering for unsupervised learning of visual features. In Proceedings of the European Conference on Computer Vision , 2018.

M. Caron, N. Houlsby, and C. Schmid. Location-aware self-supervised transformers. arXiv preprint arXiv:2212.02400 , 2022.

M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems , 2020.

M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , 2021.

T. Chen, S. Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning , 2020.

T. Chen, S. Saxena, L. Li, D. J. Fleet, and G. Hinton. Pix2Seq: A language modeling framework for object detection. In International Conference on Learning Representations , 2022.

X. Chen and K. He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2021.

X. Chen, X. Wang, S. Changpinyo, A. Piergiovanni, P. Padlewski, D. Salz, S. Goodman, A. Grycner, B. Mustafa, L. Beyer, et al. PaLI: A jointly-scaled multilingual language-image model. In International Conference on Learning Representations , 2023.

X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , 2021.

N. Dalal and B. Triggs. Histograms of oriented gradients for human detection. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2005.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021.

A. Dosovitskiy, J. T. Springenberg, M. Riedmiller, and T. Brox. Discriminative unsupervised feature learning with convolutional neural networks. In Advances in Neural Information Processing Systems , 2014.

D. Dwibedi, Y. Aytar, J. Tompson, P. Sermanet, and A. Zisserman. With a little help from my friends: Nearest-neighbor contrastive learning of visual representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision , 2021.

M. Everingham, S. A. Eslami, L. Van Gool, C. K. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes challenge: A retrospective. International Journal of Computer Vision , 111(1):98-136, 2015.

C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning , 2017.

J. Frankle, D. J. Schwab, and A. S. Morcos. Training batchnorm and only batchnorm: On the expressive power of random features in CNNs. In International Conference on Learning Representations , 2021.

J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in Neural Information Processing Systems , 2020.

R. Guo, P. Sun, E. Lindgren, Q. Geng, D. Simcha, F. Chern, and S. Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning , 2020.

K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2022.

K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2020.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition , 2016.

O. J. Hénaff, S. Koppula, J.-B. Alayrac, A. v. d. Oord, O. Vinyals, and J. Carreira. Efficient visual pretraining with contrastive detection. In Proceedings of the IEEE International Conference on Computer Vision , 2021.

O. J. Hénaff, S. Koppula, E. Shelhamer, D. Zoran, A. Jaegle, A. Zisserman, J. Carreira, and R. Arandjelovi´ c. Object discovery and representation networks. In Proceedings of the European Conference on Computer Vision , 2022.

O. J. Hénaff, A. Srinivas, J. De Fauw, A. Razavi, C. Doersch, S. Eslami, and A. v. d. Oord. Data-efficient image recognition with contrastive predictive coding. International Conference on Machine Learning , 2020.

L. A. Hendricks and A. Nematzadeh. Probing image-language transformers for verb understanding. In Findings of ACL , 2021.

J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. In Advances in Neural Information Processing Systems , 2022.

D. Homa, S. Sterling, and L. Trepel. Limitations of exemplar-based generalization and the abstraction of categorical information. Journal of Experimental Psychology: Human Learning and Memory , 7(6):418, 1981.

M. Jia, L. Tang, B.-C. Chen, C. Cardie, S. Belongie, B. Hariharan, and S.-N. Lim. Visual prompt tuning. In Proceedings of the European Conference on Computer Vision , 2022.

J. Johnson, M. Douze, and H. Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data , 7(3):535-547, 2019.

D. Kim, J. Kim, S. Cho, C. Luo, and S. Hong. Universal few-shot learning of dense prediction tasks with visual token matching. In International Conference on Learning Representations , 2023.

A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung, S. Gelly, and N. Houlsby. Big transfer (BiT): General visual representation learning. In Proceedings of the European Conference on Computer Vision , 2020.

A. Krizhevsky, I. Sutskever, and G. E. Hinton. ImageNet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems , 2012.

C. Liu, J. Yuen, and A. Torralba. Nonparametric scene parsing: Label transfer via dense scene alignment. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2009.

F. Liu, G. Emerson, and N. Collier. Visual spatial reasoning. arXiv preprint arXiv:2205.00363 , 2022.

Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A ConvNet for the 2020s. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2022.

I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019.

D. G. Lowe. Distinctive image features from scale-invariant keypoints. International Journal of Computer Vision , 60:91-110, 2004.

T. Malisiewicz, A. Gupta, and A. A. Efros. Ensemble of exemplar-SVMs for object detection and beyond. In Proceedings of the IEEE International Conference on Computer Vision , 2011.

R. M. Nosofsky. Attention, similarity, and the identification-categorization relationship. Journal of Experimental Psychology: General , 115(1):39, 1986.

M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza, F. Massa, A. El-Nouby, et al. DINOv2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193 , 2023.

N. Parthasarathy, S. Eslami, J. Carreira, and O. J. Hénaff. Self-supervised video pretraining yields strong image representations. In Advances in Neural Information Processing Systems , 2023.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning , 2021.

A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , 2021.

S.-A. Rebuffi, H. Bilen, and A. Vedaldi. Learning multiple visual domains with residual adapters. Advances in Neural Information Processing Systems , 2017.

O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision , 115(3):211-252, 2015.

M. B. Sariyildiz, Y. Kalantidis, K. Alahari, and D. Larlus. No reason for no supervision: Improved generalization in supervised models. In International Conference on Learning Representations , 2023.

R. C. Schank. Dynamic memory revisited . Cambridge University Press, 1999.

E. Shechtman and M. Irani. Matching local self-similarities across images and videos. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2007.

N. Silberman, D. Hoiem, P. Kohli, and R. Fergus. Indoor segmentation and support inference from rgbd images. Proceedings of the European Conference on Computer Vision , 2012.

O. Siméoni, G. Puy, H. V. Vo, S. Roburin, S. Gidaris, A. Bursuc, P. Pérez, R. Marlet, and J. Ponce. Localizing objects with self-supervised transformers and no labels. Proceedings of the British Machine Vision Conference , 2021.

K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations , 2015.

R. Strudel, R. Garcia, I. Laptev, and C. Schmid. Segmenter: Transformer for semantic segmentation. In Proceedings of the IEEE International Conference on Computer Vision , 2021.

T. Thrush, R. Jiang, M. Bartolo, A. Singh, A. Williams, D. Kiela, and C. Ross. Winoground: Probing vision and language models for visio-linguistic compositionality. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2022.

Y. Tian, O. J. Henaff, and A. van den Oord. Divide and contrast: Self-supervised learning from uncurated data. In Proceedings of the IEEE/CVF International Conference on Computer Vision , 2021.

A. Torralba, R. Fergus, and W. T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence , 30(11):1958-1970, 2008.

H. Touvron, M. Cord, and H. Jégou. DeiT III: Revenge of the ViT. In Proceedings of the European Conference on Computer Vision , 2022.

H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. LLaMa: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

H. Touvron, A. Sablayrolles, M. Douze, M. Cord, and H. Jégou. Grafit: Learning fine-grained image representations with coarse labels. In Proceedings of the IEEE/CVF International Conference on Computer Vision , 2021.

W. Van Gansbeke, S. Vandenhende, S. Georgoulis, and L. Van Gool. Unsupervised semantic segmentation by contrasting object mask proposals. In Proceedings of the IEEE International Conference on Computer Vision , 2021.

O. Vinyals, C. Blundell, T. Lillicrap, K. Kavukcuoglu, and D. Wierstra. Matching networks for one shot learning. In Advances in Neural Information Processing Systems , 2016.

C. Vondrick, A. Shrivastava, A. Fathi, S. Guadarrama, and K. Murphy. Tracking emerges by colorizing videos. In Proceedings of the European Conference on Computer Vision , 2018.

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.

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.

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.

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.

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.

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.

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.

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.

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.

Y. Zhang, K. Zhou, and Z. Liu. What makes good examples for visual in-context learning? In Advances in Neural Information Processing Systems , 2023.

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.

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 k i = f θ ( x i ) ∈ R H · W × D , where a feature k j i ∈ R D at a given spatial location j is aligned with the local label l j i created by averaging the pixel labels y j i in that patch. These features 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 × 512 which when encoded by a ViT-B/16 results in a 32 × 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 × 32 × 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 features_per_image = |M| |D|∗ num_augmentation_epochs where |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 ✶ c ∈ l j = 1 if a given class c ∈ l j and 0 otherwise. Then for each class c we define κ c = ∑ j ✶ c ∈ 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 j = ∑ c ∈C κ c · ✶ c ∈ l j . Finally, we take the n features_per_image from the spatial map k i with the lowest final scores using final_score j = ( class_score j · x ) + (10 6 · ✶ l j = ∅ ) (5) where x ∼ 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 features_per_image = |M| |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 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 x is encoded as a grid of features q = f θ ( 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 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.

B Implementation details: contextual pretraining

The contextual pretraining module takes as input a batch of image representations (i.e. queries) q = h = f θ ( x ) ∈ R B × H · W × D from the ViT encoder f θ , where B = 4096 is the batch size, H = W = 14 are the height and width of the spatial feature map and D = 768 for ViT-B and D = 1024 for ViT-L is the feature dimension. Keys and values for the contextualization crossattention operation are entries of the memory bank M p = { ( k i , v i ) , i =1 , ..., |M p |} , where keys k i are taken from previous batches by spatially averaging h (see Equation 2) and values v i are obtained by applying a two-layer MLP ϕ θ to the keys, where we use batch norm after the first layer and the hidden dimension is set to 4096. Each feature q i of the image representation is then updated as c i = ψ θ ((1 -λ ) q i ∥ q i ∥ + λ ˆ v i ∥ ˆ v i ∥ ) , where ψ θ is a linear layer and ∥ x ∥ is the L 2 norm. Preliminary analysis showed λ =0 . 2 to work well across datasets, so we use it for all our experiments, with higher values λ ≥ 0 . 5 degrading performance.

We populate the memory bank with all batch entries of ImageNet-1k / -22k at each step, using the representations from the target network. The memory bank is spread across 128 Cloud TPU v3 workers with 1200 entries on each TPU for ImageNet-1k (256 TPUs with 600 entries for ImageNet22k), resulting in total memory length of 153 , 600 .

C Implementation details: self-supervised pretraining

C.1 Data augmentation

Each image is randomly augmented twice, resulting in two views x 1 and x 2 . The augmentations are constructed as compositions of the following operations, each applied with a given probability:

random cropping: a random patch of the image is selected, whose area is uniformly sampled in [0 . 08 · A , A ] , where A is the area of the original image, and whose aspect ratio is logarithmically sampled in [3 / 4 , 4 / 3] . The patch is then resized to 224 × 224 pixels using bicubic interpolation;

horizontal flipping;

color jittering: the brightness, contrast, saturation and hue are shifted by a uniformly distributed offset;

color dropping: the RGB image is replaced by its grey-scale values;

gaussian blurring with a 23 × 23 square kernel and a standard deviation uniformly sampled from [0 . 1 , 2 . 0] ;

solarization: a point-wise color transformation x ↦→ x · ✶ x< 0 . 5 + (1 -x ) · ✶ x ≥ 0 . 5 with pixels x in [0 , 1] .

The augmented images x 1 and x 2 result from augmentations sampled from distributions T 1 and T 2 respectively. These distributions apply the primitives described above with different probabilities and different magnitudes. Table 8 specifies these parameters for the BYOL framework [28], which we adopt without modification.

C.2 Optimization

We pretrain the model for 300 epochs on ImageNet-1k or 100 epochs on ImageNet-22k using AdamW [47] with a batch size of 4096, split across 128 Cloud TPU v3 workers for ImageNet-1k and 256 Cloud TPU v3 workers for ImageNet-22k. Training a ViT-B / ViT-L for 300 epochs on ImageNet-1k takes roughly 21 hours / 53 hours, while 100 epochs on ImageNet-22k takes approximately 60 hours / 128 hours. We update the online parameters θ with a cosine learning rate schedule with a base learning rate of 0.001, weight decay of 0.1 and gradient clipping with a maximum norm of 1. We update the target parameters ξ as an exponential moving average of the online parameters with a decay rate of 0.99.

Following [16] the projections and predictions in Equation 4 are normalized and rescaled such that their norm is equal to 1 / √ τ where the contrastive loss temperature τ is equal to 0.1. When using additional supervision we set the supervised loss weight α to 0.25 for the supervised ViT-B trained on ImageNet-22k and α =0 . 05 for all other experiments.

D Supplementary analysis

D.1 Data efficiency

In Table 2 we compared Hummingbird with several leading representation learning techniques in the low-data regime. Here we provide the complete analysis from 1 / 128 to 100% of the data, as well as results for our ViT-L model trained on ImageNet-22k to show the scaling properties of Hummingbird . Note that there is a difference between the experiments run here and those found in Section 4.4 of the main paper; that section uses an UperNet [77] decoder and this section uses a linear decoder for all of the finetuned rows in each table.

For PASCAL VOC (Table 9), Hummingbird performs very well not only in the low-data regime but in the full-data regime, with the apples-to-apples comparison (ViT-B self-supervised on ImageNet1k) competitive with all other techniques even as the dataset fraction increases. This table also demonstrates the clear benefit of supervision as well as model-size and dataset size scaling-with only nearest neighbors (no finetuning), Hummingbird++ trained on ImageNet-22k with a ViT-L backbone beats all of the other finetuned variants for every dataset fraction. Hummingbird++ using a ViT-B and ImageNet-1k predictably lies in-between the other two models for every dataset fraction.

For ADE20K (Table 10), the same general trends from above hold. Backbone and dataset scaling are once again beneficial as Hummingbird++ with ViT-L and ImageNet-22k training outperforms the other Hummingbird models, however this time the absolute performance relative to the finetuned competition in the high-data regime is less favorable since the end-to-end finetuned versions of other techniques start to outperform the nearest neighbors only ViT-L Hummingbird++ at 1 / 16 of the data.

D.2 Correlation of NN retrieval and finetuning performance

In this section, we study the relation between NN retrieval performance and end-to-end finetuning. To that end, we collect 14 Hummingbird models trained with different architectures (ViT-B vs ViT-L), datasets (ImageNet-1k vs ImageNet-22k), learning objectives (self-supervised or with additional supervision), and training lengths. Figure 7 plots the performance of these models when equipped with NN retrieval decoders (x-axis) and fully-finetuned UperNet decoders (y-axis). For both PASCAL VOC and ADE20K semantic segmentation, performance using one decoding scheme is highly predictive of the other (Pearson's ρ = 0 . 80 for PASCAL VOC, ρ = 0 . 89 for ADE20K). As such, even in cases where NN retrieval underperforms end-to-end finetuning, it can still be used as a powerful diagnostic tool. As illustrated in Section 4.3, evaluating with NN retrieval is much simpler and faster than with end-to-end finetuning, even when using a linear decoder. End-to-end finetuning often requires sweeping over optimization hyperparameters and averaging across multiple seeds, making it unsuitable for online evaluation, whereas NN retrieval is 10 × less variable across runs and doesn't require any hyperparameter sweeps. As such NN retrieval can be used as an online evaluation that is highly predictive of performance obtained with more expensive finetuning protocols.

70

72

74

76

78

80

81

82

83

84

85

End-to-end fine-tuned

PASCAL: NN retrieval vs E2E FT

30

32

34

36

38

45

46

47

48

49

50

51

52

End-to-end fine-tuned

ADE20K: NN retrieval vs E2E FT

NN retrieval

NN retrieval

D.3 Effect of pretraining and evaluation memory length for ADE20K

We include the equivalent of Figure 6 on the ADE20K dataset in Figure 8. Similar to what we observe for PASCAL VOC, we benefit from large memory banks at evaluation. Since the ADE20K

training set is roughly 2 × larger than that of PASCAL VOC, we also observe that sampling which features to store in the memory bank is more important than it is for PASCAL VOC (see Appendix A.1 on the details of the sampling procedure). Similarly, at training time, ADE20K benefits from larger pretraining memory banks than PASCAL VOC, with performance plateauing for memory banks larger than 200 , 000 . Thus, we set the pretraining memory bank length to 153 , 600 in all our experiments (see Appendix B for details on contextual pretraining).

0

50000

100000

150000

200000

250000

300000

Training memory length

28

29

30

31

32

33

34

ADE20K: Importance of pretraining memory

self-supervised

supervised

10 5

10 6

10 7

Evaluation memory length

10

15

20

25

30

ADE20K: Importance of evaluation memory

downsample=True

downsample=False

D.4 Impact of encoder size on finetuned performance

We investigate the impact of encoder size in the finetuning regime in Table 11. We find that scaling the encoder from ViT-B to ViT-L is beneficial for both Hummingbird and Hummingbird++ , where the self-supervised Hummingbird benefits slightly more from model scaling than its supervised counterpart. Note that this scaling study highlights the need for jointly scaling data and model size, as the best performing model overall is the one with ViT-L as an encoder trained on ImageNet-22k.

D.5 Importance of retrieval-based supervision for in-context scene understanding

We study the importance of retrieval-based supervised objective (see Section 3.4) on in-context scene understanding performance. We compare a model trained purely with the retrieval-based supervised objective ('Sup') with Hummingbird (purely self-supervised, i.e. 'SSL') and Hummingbird++ (both self-supervised and retrieval-based supervised, i.e. 'SSL + Sup'). Results shown in Table 12 indicate

the necessity of the self-supervised objective for creating representations that are general and transfer well to downstream tasks.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a self-supervised approach to adapt models learned via SSL to in-context learning via retrieval. To obtain the best performance, the authors propose to add attention pooling to encourage features attention (even though within an image) and supervised training for retrieval. Through an extensive number of ablations and empirical studies, the authors show the features produced via Hummingbird outperform other SSL approaches on segmentation and depth estimation.\n\nSTRENGTHS:\n1. The paper is well written and the presentation is very good. \n2. The idea is to perform in-context via retrieval is very simple\n3. The empirical results are effective and impressive compared to the baseline\n\nWEAKNESSES:\nMajor:\n1. There are a few potential limitations with the current approach to in-context learning. Since the main idea of in-context learning is the adaption to new tasks, it is important to explicitly mention what can and what cannot be done with Hummingbird. \n2. With Hummingbird (no retrieval loss) there is a discrepancy between train time to test time\n\nMinor:\n1. Some missing RW. \n2. L109 - slight inconsistency in notation going from HxWxD then using a single index for spatial location\nRating\n3. What is the ViT-B patch size used? can you add it to Table 1/4?\n\nQUESTIONS:\nMajor:\n1. There are a few potential limitations with the current approach to in-context learning. Since the main idea of in-context learning is the adaption to new tasks, it is important to explicitly mention what can and what cannot be done with Hummingbird. \n\n* The title implies that scene understanding, however, the retrieved features might actually be very local and specific. The features are retrieved independently of each other and there is no mechanism to consolidate the retrieved features into a whole scene output.\n* Correspondence - with the current approach, the authors assume there are spatial correspondences between features and their associated labels. This is the case for segmentation or depth estimation but this is not necessarily the case for every possible task. E.g, if the task is to translate, crop and resize, rotate, or flip the image, etc, this assumption breaks.\n* Retrieval approach might not result in smooth outputs. This is ok for segmentation but how would this work for style transfer? \n\nI hope that the authors can explicitly clarify these in the final manuscript. On the one hand, the author's main claim is on segmentation and depth estimation, but on the other hand, in-context learning is very broad and therefore I think it should be stated explicitly what can be achieved and what not under this framework. I see this as the most important point, and I'll be willing to increase my score if this is addressed. \n\n2. Discrepancy between train time to test time (attending within the image to attention across examples). I assume that with an end-to-end objective (perhaps similar to [2]), the gap between Hummingbird to Hummingbird++ can be diminished. \n\n\nMinor:\n1. Some missing RW.\n* Recently retrieval-based a retrieval-based approach was proposed for choosing visual prompts [1]. To be fair, this work might be parallel research.\n* Conceptually, the idea here is similar to a (now old) paper [2]. E.g., learning to attend to different features to colorize, which can later be used for segmentation. I think it is at least worth citing.\n2. Why Hummingbird? Is there any connection from Hummingbird to the idea/method I'm missing?\n\n[1] Zhang et al. \"What Makes Good Examples for Visual In-Context Learning?\". 2023.\n\n[2] Vondrick et al. \"Tracking Emerges by Colorizing Videos\", ECCV 2018.\n\nLIMITATIONS:\nPlease consider adding the following limitations (also stated above):\n* The title implies that scene understanding, however, the retrieved features might actually be very local and specific. The features are retrieved independently of each other and there is no mechanism to consolidate the retrieved features into a whole scene output.\n* Correspondence - with the current approach, the authors assume there are spatial correspondences between features and their associated labels. This is the case for segmentation or depth estimation but this is not necessarily the case for every possible task. E.g, if the task is to translate, crop and resize, rotate, or flip the image, etc, this assumption breaks.\n* Retrieval approach might not result in smooth outputs. This is ok for segmentation but how would this work for style transfer?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates in-context learning in computer vision. It proposes the Hummingbird model, which uses a pretraining protocol and a nearest neighbor retrieval mechanism. This model avoids the need for task-specific parameters or finetuning, offering fast and data-efficient adaptation to various scene understanding tasks, including semantic segmentation and depth estimation (at the current stage). The pretraining protocol utilizes a unique combination of attention across and within images, providing robust image representations. Tested on common benchmarks of semantic segmentation and monocular depth estimation, the model's performance approaches the capabilities of specialist models finetuned for each task, suggesting potential for in-context scene understanding.\n\nSTRENGTHS:\nI believe the studied direction, in-context scene understanding, is very important, and this work takes step towards it. The proposed nearest neighbor retrieval framework reminds me of nearest neighbor baselines in few-shot recognition [1,2]. Once the learned representation is powerful enough, a simple NN-baseline can have comparable or even higher performance than SOTA meta-learning algorithm. It is interesting to see such simple framework being re-proposed and shown strong performance in in-context scene understanding. \n\nAll the used techniques (cross-attention, patch-level attention, self-contrastive objective) are sound and composed in a reasonable way. Their effectiveness are strongly proven in Table 5. Overall, the paper is easy to follow. \n\n\n\n\n\n[1] Large-Scale Few-Shot Learning: Knowledge Transfer With Class Hierarchy\n[2] SimpleShot: Revisiting Nearest-Neighbor Classification for Few-Shot Learning\n\nWEAKNESSES:\nFirst, I am unsure if Figure 1 is the best design: (1) I cannot know the components used in Hummingbird; (2) does the nearest neighbors (on the top) refers to image-level or patch-level or pixel level? (3) does query image is people-horse the nearest neighbor also contains furniture images?\n\nL 142-144 said: \"we train representations to locate the most distinctive part of an image\". Does this claim has any supportive evidence? \n\nL 251-254 shows the author adopts a NN search way to let cross-attention to a small set. What's the relationship between the NN search time and memory length. \n\nWhat's the cost for the pre-training stage? Will this be a huge commitment?\n\n(1) If Hummingbird pre-training with semantic segmentation and fast adopt to monocular depth estimation, how will performance change? (2) If Hummingbird pre-training with monocular depth estimation and fast adopt to semantic segmentation, how will performance change? (3) if Hummingbird pre-training with both semantic segmentation and monocular depth estimation, will this model outperforms than the models pre-trained with one type of data?\n\nQUESTIONS:\nPlease address concerns raised above.\n\nLIMITATIONS:\nThe author has discussed broader impact and limitations in Section 7 (L 273~285). I personally feel one potential limitation was not discussed: training time and inference time. The training/inference time could be huge if adopting large memory. Also, what's the relationship between training/inference time and memory size?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper presents a novel in-context learning paradigm based on NN retrieval, which can be adapted for downstream dense vision tasks without requiring fine-tuning or specialized decoders.

STRENGTHS:
1. The proposed method addresses a promising regime for reducing task-specific training requirements with prompt images and labels.
2. The framework is conceptually simple yet powerful, providing a concise and reasonable solution.
3. The methodology is clearly presented, and the paper is well-organized and easy to follow.

WEAKNESSES:
1. The current paradigm seems to only be applicable to pixel-wise prediction tasks. As it is unclear how the in-context learning paradigm could be extended to other vision tasks, its application may be limited.
2. The retrieval process is inefficient compared to previous paradigms during inference.

QUESTIONS:
1. In Tab. 3, why does the result of MoCo-v3 significantly drop for E2E FT?
2. In Sec. 3.2, why contextual pretraining can learn meaningful representations while "enforcing its representation to be expressed as a combination of representations of nearby examples", even without supervision? I don't understand its intuition.

LIMITATIONS:
The authors have discussed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
The paper presents a non-parametric nearest neighbor retrieval method designed to efficiently adapt a pretrained model for dense scene understanding tasks. The fundamental elements include a self-supervised pretraining method, which incorporates attention across images, and a spatial attention-pooling mechanism. The conducted experiments demonstrate improved data efficiency and training speed compared to fully-finetuned methods, while maintaining competitive performance.

STRENGTHS:
- The paper shows that non-parametric nearest neighbor retrieval method can be more effective than linear probing in existing self-supervised learning approaches.
- The self-supervised pretraining does not necessitate either fine-tuning or specialized decoders.
- The proposed method significantly outperforms existing self-supervised baselines in downstream tasks and demonstrates improved data efficiency.

WEAKNESSES:
- The paper only consider two downstream tasks (semantic segmentation and depth estimation) in the experiments, despite aiming for general-purpose representations.
- When dealing with a complex dataset like ADE20K, the training approach without any supervision (Hummingbird) exhibits lower performance than some baselines, particularly when the data is scarce (n=158, 316) or when fine-tuning is applied.

QUESTIONS:
- Why does the proposed contrastive learning objective enable the encoded features to be directly applied to downstream tasks without any parameter adaptation?
This seems to align well with semantic tasks like semantic segmentation, but can you explain its applicability to geometric tasks like depth estimation?

LIMITATIONS:
The authors have adequately addressed the limitations and broader impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis submission proposed a new representation and pre-training method, named Hummingbird for its fast adaption ability. Hummingbird could perform inference simply by nearest neighbor retrieval from a prompt of annotated features. Unlike the previous pre-training and then fine-tuning/linear probing protocol, the new representation could predict dense outputs by computing the cross attention (dot product similarity with propagation) between the testing query image and the memory bank. The pre-training process could be fully self-supervised or mixed with class label supervision. The new representation of Hummingbird outperforms strong representations like DINO, MAE on nearest neighbor evaluation of semantic segmentation and depth estimation benchmarks by a large margin. Hummingbird also shows great data efficiency in low-data regime.\n\nSTRENGTHS:\n1. I like the idea of nearest neighbor propagation very much and I am very delighted that authors make it works. With the nearest neighbor propagation protocol, a single model could adapt to any vision task without fine-tuning. It's also very interesting that Hummingbird uses patches as the memory bank instead of the whole image, which saves a lot of storage and computation. \n2. The fast adaption ability of Hummingbird is very surprisingly good. I like Figure 3 very much, and it shows with just 10s on V100, Hummingbird could achieve very good results already, which is nearly real-time. \n3. It's very interesting to see the proposed  self-supervised objective works very well with retrieval-based supervised objective. The self-supervised Hummingbird already achieves amazing results and retrieval-based supervised objective could further boost it.\n\nWEAKNESSES:\n1. I am a little bit concerned about the scalability of Hummingbird. I truly appreciate that authors provide ViT-L version of Hummingbird and Hummingbird++. But if we are going to scale Hummingbird to ViT-H level, would ImageNet-22k be sufficient enough? Is there any more gain from scaling up the model? Or if we follow the self-supervised objective only, how much gain could we expect if there is more data?\n2. Depth results in \"Fast adaptation to downstream tasks\" and \"Fully finetuned scene understanding\" sections are missing. I would also be very interesting to see the adaption ability on tasks other than semantic segmentation.\n\nQUESTIONS:\n1. I found that there is no softmax operation in contextual pretraining Eqn 2 but there is softmax in Eqn 3. Is there any explanation for this design?  \n2. Would the work be open-sourced? I think this work would be highly impactful for the community and it would be quite beneficial if authors would like to release the implementation and checkpoints. \n3. Is Hummingbird ViT-L fully fine-tuning yields higher accuracy than ViT-B? \n4. Since Hummingbird is very good at nearest neighbor, it would also be interest to see the linear probing results and comparison. \n5. It would be interesting to see whether retrieval-based supervised objective works solely without self-supervised objective.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}