Towards In-context Scene Understanding
Authors: Ivana Balazevic
David Steiner
Nikhil Parthasarathy
Relja Arandjelovic
Olivier J Henaff
Venue: NeurIPS 2023 spotlight (2023)
Keywords: transfer learning
adaptation
self-supervised learning
contrastive learning
scene understanding
representation learning
in-context learning
vision transformers
TL;DR: Retrieval-augmented self-supervised models display fast and data-efficient adaptation to a range of scene understanding tasks.

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
Ivana Balažević* David Steiner" Nikhil Parthasarathy! Google DeepMind
Relja Arandjelovié Olivier J. Hénaff
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 assisting with any query. Computer vision, in contrast, has largely stayed the former regime: specialized decoders and finetuning protocols generally required to perform dense tasks such as semantic segmentation and depth estimation. In this work we explore simple mechanism for in-context learning of such scene understanding tasks: nearest neighbor retrieval from a prompt of annotated features. We propose new pretraining protocol--leveraging attention within and across images-which yields representations particularly useful in this regime. The resulting Humminggirrd 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.
Introduction
In natural language processing (NLP), the pretrain-fined finetune paradigm has long been the dominant way of acquiring domain-specific knowledge and adapting model's behavior particular task (e.g. question answering, natural language inference, summarization). More recently and predominantly due the increase in model and dataset sizes, large language models have exhibited impressive, task-agnostic emergent capabilities [T] 37 68), where a single model, given an appropriate prompt, can perform 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 31
which capture useful semantics [15] 22
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 classifi- cation [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-graid 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- graiied spatial
32
35,
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, expand the well-known non- parametric nearest neighbor (NN) retrieval method 9 [5 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.
'qqual contribution. 'Current affiliation: NYU CNS. work done while interning at Google DeepMind. Corresoondence to (balazevic, davidsteiner, henaff) @ google.com.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

We further show that the NN scene understanding capabilities of canonically pretrained vision trans- formers (such as MAE 30 and DINO [5]) vary greatly, despite similar finetuned performance. We find two pretraining components to yield reliable gains: (1) 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 image to summarize the (contextualized) spatial grid of features into a single image-level representation to be fed into the self-supervised objec- tive. We showcase the benefits of this approach in 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-purposs 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- non-paamettric methods combined with simple features such as SIFT
48 and HOG 21 saw early success in image classification
O shape matching 7
59 scene
recognition 166 76. and image parsing [44]. Exemplar-SVM. 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 669 74 shape their representations for this task, enabling aine-griied 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 722,, tracking abilities emerge after pretraining on a colorization task via retrieval from reference frames of a video. Retrieval has also been proposed as means of enriching the positive pairs used in eeff-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. 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 [7]] and MAML (26) learn to solve new classification and reinforcement learning tasks specified on the fly. Architectural innovations, such as image prompting 4 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 lask-specific weight updates and learned similarity functions. 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.
39 82 and adapter layers [27 [5)
contrast,
Self-super learning. Methodologically, our representation learning method is most similar to supervieee learning (SSL) techniques. Similarly to NLP, baaed SSL has witnessed great success in recent years, notably with the advent of contrastive methods [14 15| 16 23 [3).
self-distillation [2 18 28), and masked auto outo-encoding [30]. Due to their conceptual simplicity, we base our method on standard contrastive baselines such as SimCLR [66 and MoCo (II. Image-based SSL techniques have since been tailored to learning representations which transfer well to scene understanding tasks
K3
3 70 78), and although they have been shown to support zero-shot object
discovery [34 611. they generally still require task-specific decoders and end-to-ent finetuning.
2

---

Find nearest neighbors
Inpul prompi images
Query images
BB
Input prompt labels
Aggregate their labels
Prediclions
Figure I: 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 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 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 computed with our Hummingbirdd model.
Method
The following sections describe the retrieval- -asdd scene understanding decoding protocol (Section 3.1. followed by the contextual pretraining method (Section (.2) and the self-supervised (Section and supervised learning objectives (Section 3.4). We subscripts to differentiate between representations and superscripts æ to denote spatial locations within a representation.
3.3
3.1
Retrieval-based scene understanding
A general-pune purpose image representation should perform well across a variety of scene understanding tasks tut-of-the-box i.e. without modifying its parameters. 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 16 x 16 across all models this work). Given a prompt composed of training images from the downstream task and their corresponding
labels { y:), i I.... N, æ; E R x XC), j. our aim is to enable pretrained image encoder fa to make predictions about a new image æ from the test set. In tasks considered in this work, labels y:
are spatial maps either class labels y, E CHf xW' (c.g. for semantic segmentation, where the space of all classes) or scalars y; E E xW' (e.g. for monocular depth estimation).
We encode each prompt image into spatially flattened map K; = fe (xx,) E ER "f-WWx where a feature k/ ERP a spatial location is aligned with the local label created by averaging the pixel labels y of patch. We then sample subset of features and local labels for each image, which form the keys and values of the memory bank M [kk 12), N... N. jr~S} (see Appendix A. for details on the sampling distribution S). In the following, we do not distinguish between entries from different images, and use single integer j to index into the memory bank: M ((k), ((k,,)))3 f))jj=...... [M]]
Given test image we form representation fe() and use each spatial feature q as query to cross-attend over the memory bank with temperature J.. The cross-attention weights are then used to combine the corresponding labels and form a local prediction F::
(q) K 3
a' softmax(s') (s"),
i i- E a.
(1)
Equation| defines the cross-attention operation as CA( k', ()). The final prediction ÿ is simply the concatenation of local predictions i' upsampled to the original image size via bilinear interpolation. As result, nearest neighbor retrieval allows a simple image encoder to perform scene understanding tasks without any decoders or parameter adaplation (finetuning otherwise) to the downstream dataset. The mechanism is also entirely agnostic to the format of the labels, enabling to perform tasks diverse as semantic segmentation and depth estimation.
3

---

controtu hnd Notdken
aeterrion poit
MAg monnt
emooddr
corthol Lodin
memory cank
emooder moon pool
contraciv abjectre
VEW 2
poot
Figure 2: Hummingbird model components.
3.2 Contextual pretraining
Memory retrieval allows 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 in a similar manner, by enforcing its representation to be expressed examples. Over the course of training, we populate a memory bank Mp = {(k,, v;). i = 1,... [Mp]] with spatially averaged keys and values computed from training images æ; from previous batches:
combination of representations of nearby
H-W
h fo(xx) ERH.WxL
k. H. W.W.W
1
ERD. M
8; == oo(;)) E R.D.
where we use an MLP the value head Φo (see Appendix B for implementation details). We then form a representation q fa of a new training image æ and use each spatial feature q' to attend over the memory bank and compute an update &' =CA(q", k,, v,). Each feature "contextualized" as c² vee((1 >) प + A (T) ). where the is a linear layer and À weighting parameter. The contextualized image representation c= c=go(q, Mp) is simply the concatenation of local features et.
Note that the pretraining memory bank .Mp is discarded at test time and differs from the test time memory bank A described in Section 3.I, allowing for straightforward comparison of our representations fo( 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 locate the most distinctive part of image using a combination of attention pooling and contrastive learning. Following [16] 28), we construct different views of unlabeled images æ through random data augmentation # A (æ), 22 ~ A2(x), see Appendix (x2
C.I Each view is encoded as h; = =f(;)) and further contextualized with the mechanism described above as c; go go MMp) E fx (see Figure 2). Following [52], we compute attention pooled representations ē, R using masks m, derived from lightweight attention module ag, which we augment with an additional value head wa. Pooled features are then used to compute projections 2.:
H-W
= softmax((o(c:))
2=
Σ mf wol(c))
2
(o:)
Finally, following 20 28, 65), each view forms predictions qe (z,) about the other view's 's targets z5. which are computed with the same architecture and different set of weights & which vary more slowly (see Appendix C.2p. The online weights 0 are optimized using standard contrastive loss:
exp(go(z?)) =5))
Cs,,,(:; - -log
exp(px(z!) x5) + E. exp(gx(2!) z5)

---

Table 1: In-context scene understanding. All models pretrained source data in supervised or self-supervised manner, and applied to downstream datasets without modification. All downstream tasks performed using single mechanism, nearest neighbor retrieval. 'indicates reproduction of external work, all other models were evaluated using publiely available checkpoints.
Semantic segmentation
Depth pred.
Method
Encoder
Params (M)
Dataset
PASCAL
ADE20K t
NYUv2
Supervised
ViT-B
INIK
35.1
13.8
.913
DINO [15
VIT-B
INIK
55.9
21.8
MoCo-v3 20
VIT-B
86
INIK
37.2
14.6
.771
MAE 30
VIT-B
86
INIK
6.6
3.3
.981
LOCA 13
VIT-B
86
INIK
57.5
18.5
.880
Hummingbird
VITBB
INIK
70.5
28.3
.718
Hummingbird++
VTT-B
86
INIK
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
R3
VIT-B
86
IN22K
56.4
16.8
.829
Humminghird
VIT-B
IN22K
73.5
30.7
.706
Hummingbird++
VIT-B
IN22K
76.2
34.1
Comparison across architectures:
CLIP 55
NFNet-F6
ALIGN
57.2
25.0
.844
Supervised
NeXt-XL
[46
1300
IN22K
58.9
25.5
.791
Supervised
VIT-L
IN22K
65.8
26.1
.600
MAE (30
VIT-L
307
INIK
8.0
3.6
.934
1,OAA 13
VIT-I.
307
IN22K
59.5
17.6
.912
Humminghird
VIT-L
307
IN22K
76.9
35.0
.671
Humminghird++
VIT-I
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], explore the of label-supervision for learning representations that perform well in dense NN retrieval. While supervision is typically added with
linear classifier atop average pooled features 33 43 62), we instead it constrain contextual
pretraining and further align our training methodology with NN retrieval 669 74). Specifically, we
expand the memory bank Md include the labels: Mr {(k,, Us y:).
[MD and query with attention pooled features é; ER (see Equation 3 to form predictions ÿ. -A(C,, K; y;).
We then use the standard softmax cross entropy loss C Ws: y:)), which added the self-supervised objective of Equation 4 forms the total loss L = S + a(Cix + C(g), with supervised weight a. Note that the memory bank M, is only used during training and the added supervision relates global image classification task, not the downstream pixel-level tasks.
4 Experiments
We demonstrate the generality of Hummingbird representations through retrieval- based scene under- standing on several downstream tasks (Section 4.11 semantic segmentation on PASCAL VOC (2] and ADE20K 83 with mean IoU (mIOU) as metric, and monocular depth estimation
NYUv2 60 with root-mean-square error (RMSE) as metrie. 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-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-basee 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). refers the number of images fraction represents. All runs averaged over five different seeds, with standard deviation of the order of 0.04 O. 10% for NN and 0.36 1.35% for E2E FTon PASCAL/ MASCAL/ADE20K, ADE20K. Each method is trained with ViT-B ImageNet- Ik.
PASCAL1
ADE20K
Method
Decoder
1/128 (n=83)
1/64 (n=165)
1/128 (x==188
1/64
(w=316)
Supervised 67
В2Е FT
41.8
53.8
10.8
14.3
DINO [5
E2E FT
36.1
44.3
11.7
14.4
MoCo-v3
E2E FT
19.9
33.4
4.6
7.9
MAE 30
E2E FT
34.2
44.1
8.2
12.2
LOCA 113
E2E FT
40.1
53.9
11.2
15.5
Hummingbird
NN NN
50.5
57.2
11.7
15.1
Humminghird++
52.4
57.3
12.7
16.4
PASCAL semantic segmentation
ADE20K semantic segmentation
70
35
60
25
Mcon
40
Mcon
30
Linear + E2E FT
10-
Linear frozen
NN retrieval
5
10:
10
10a
10
10.
10a
10a
10*
Number of adaptation images
Number adaptation images
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*).
-2k (56]) scales, trained with supervision (Hummingbird+++) or without (Hummingbird). Figure shows example prediction made by Hummingbird with ViT-B encoder PASCAL VOC.
The top part of Table 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 Humminghird scales well with increasing the dataset size from ImageNet- to ImageNet-22K, which does not hold for all other methods (e.g. MAE, consistent with (5I)). Further, training with supervision generally beneficial, particularly for semantic segmentation. For an ablation on the impact of retrieval-basee supervision performance, see Appendix D.5
The bottom part of Table caontainn the best performing methods across architectures, showing 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.2 Data-efficient retrieval-based scene understanding
In addition adapting to downstream tasks with minimal (or ideally no) alterations to the model, second ingredient for in-context learning is adaptation given only a limited number examples.
We therefore evaluate the performance of Humminghing retrieval in the low-data regime, and compare 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 1~0000 images). For ADE20K the effect is less pronounced, however NN retrieval still exceeds end-to-end finetuning when given up 1 1/32 of the data (~-00 images). Hummingbird retrieval outperforms linear decoding

---

Table 3: Fast adaptation to new scene understanding tasks. After pretraining, models transferred to downstream tasks with the full dataset, but a small amount of computation: epoch. Models perform the task either with linear classifier (Frozen), end-to-end fine-tuning (E2E FT) or with our mechanism for in-context scene understanding (NN).
PASCAL
ADE20K
Method
Decoder
Frozen
E2E FT
Frozen
E2E FT
Supervised
Linear
61.5
66.3
27.6
15.1
DINO
15
Linear
54.9
64.0
25.6
23.4
MoCo-v3
20
Linear
41.2
4.8
14.6
3.2
MAE 30
Linear
20.1
42.5
8.3
7.9
LOCA 13
Linear
61.9
62.9
25.4
14.6
Humminghird
NN
70.5
Humminghird++
NN
72.1
30.5
PASCAL semantic segmentation
ADE20K semantic segmentation
0.7
Linear + E2E FT
0.25
Linear frozen
NN retrieval
0.20
Mean loU
0.4
Mean
0.15
0.10
0.2
0.1
200ms
15
10s
1min
10min
1h
200ms
1min
10min
1h
Adaptation time (1 TPU)
Adaptation time (1 TPU)
Figure 4: Adaptation time of Humminghird. The model is evaluated with retrieval- evaluation ("NN retrieval"), linear probing ("Linear frozen")) or full finetuning ("Linear + E2E FT*).
on top of the frozen encoder in all cases. These results show that given an appropriately designed encoder, NN retrieval provides data-efficient alternative end-to-end finetuning, and is strictly more expressive than linear decoding.
Second, verify the generality of these findings by comparing Humminngirrd retrieval to several other representation learning algorithms which transfer to the low-data regime with finetuning (Table [ see Appendix D.I 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 [5] for 1/128 of the data (158 images) and outperformed by LOCA for 1/64 of the data (316 images), whereas Hummingbinl++ outperforms all other models, demonstrating the benefit of retrieval-bait based supervision during pretraining. In summary, in the few-shot regime (c.g. < 100 images) relevant for in- in-coneext learning, Hummingbird retrieval provides compelling and robust alternative 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 the entire dataset. Yet even in this large-data regime, assistant systems must be quickly adaptable to new tasks. We thus evaluate 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 amount of computation given to NN retrieval by partially populating the memory bank with fractions of the dataset. Figure shows that minutes (I epoch through the downstream training set)
7

---

Table 4: Scene understanding with end-to-end finetuning. After pretraining, models are equipped with task-specific decoders and finetuned for that task on the entire downstream dataset. 'indicates results are taken from [30,, using UperNet 77 as the decoder. Results for all other baselines are taken from 13 and use the linear decoder from For ViT-L results, see Appendix [D.4
Fine-tuned accuracy (mloU)
Method
Encoder VYT-B
Dataset
PASCAL 1
ADE20K 7
Random
INIK
29.1
21.1
Supervised
ViT-B
IN1K
76.1
47.3
DINO 15
VIT-B
INIK
74.1
44.1
MoCo-v3
20
VYT-B
INIK
74.5
47.3*
BEiT
5
VIT-B
INIK+DALLE
54
47.1
MAE
30
VIT-B
INIK
75.0
48.1
LOCA
VIT-B
INIK
76.7
47.9
Hummingbird
VIT-B
INIK
80.0
44.9
Hummingbird++
VIT-B
INIK
81.2
44.9
Hummingbird
VIT-B
IN22K
81.6
46.9
Hummingbird++
VYT-B
IN22K
82.1
48.2
Table 5: Ablation of pretraining components. Effect of training with spatial attention pooling (as opposed to mean pooling or a [CLS] token) and memory contextualization ("Cont.") ") on performance. All models pretrained with ViT-B ImageNet-Ik.
Semantic segmentation
Depth pred.
Method
Pool.
Cont.
PASCAL f
ADE20K t
NYUv2 +
MoCL. 66
X
38.6
4.9
1.01
cont.
mean
55.6
15.3
.901
[CLS]
[CLS]
X
64.5
23.9
.741
[CLS] + cont.
[CLS]
V
65.6
25.1
.731
QK [52] cont.
QK att.
68.7
26.3
.728
QKV att.
QKV all.
68.0
27.4
.742
Hummingbird
QKV att.
70.5
28.3
.718
are sufficient to build performant NN decoder (70% mloU 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 tuming of learning rates, weight decay, and warm-up length. While linear classifier converges more quickly than finetuning, it saturates with significantly lower performance than NN retrieval (50% mloU PASCAL VOC, 20% ADE20K).
We also quantify these benefits in terms of relative convergence: on PASCAL VOC, NN retrieval reaches the performance of full finetuning after minutes rather than 3 hours, and the performance of a linear classifier in 2 seconds rather than minutes. For ADE20K, the speedups are smaller, but significant: 7 minutes rather than 30 minutes (relative to full finetuning), and 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 interactive setting.
Table 3 compares Hummingbird retrieval to other models (equipped with linear or end-to-end fine- tuned decoders) in the fast-adaptation regime (i.c. when given single pass over the full downstream dataset): Humminghing retrieval outperforms all other pretraining techniques and decoding mecha- nisms 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 comparison of fully finetuned Hummingbird with fully finetuned state-of- the-art models on the semantic segmentation task. We follow the finetuning protocol of MAE [30] and use UperNet (77) as a decoder. ssows that both Humminghinl and Hummingbird++ are competitive with state-of-the-ard when finetuned. Further analysis shows retrieval-based performance to be correlated with the finetuning performance (see Appendix (D2)). paving the way for using retrieval-baat baaed evaluation 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 We find attention pooling to yield superior performance compared to mean pooling or [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 the NN retrieval setup. We further find that modifying it ("QK att") with a value head ("QKV V att.") improves its performance across all tasks.
NN lookup time vs. memory bank length
200
Runtime (ms)
100
50
25
10
40,000
100,000
640, 000
2,500,000
10,240,000
Memory bank length
Figure 5: Effect of memory bank length on nearest neighbor lookup at inference time. Inference time is for a single image. Lookups were done on a single Nvidia A GPU.
PASCAL: Importance of pretraining memory
PASCAL: Importance of evaluation mernery
65-
Mcan loU
Mean 100
35
72
70-
supervised
downsamplc=True
68
supervised
45-
downsample-False
50000
109000
L50000
200000
250000
Training memory length
Evaluation memary length
Figure 6: Effecc of the pretraining (left) and evaluation (right) memory length on performance. All models were pretrained with ViT-B on ImageNet-22k. Left: Since the sarsevall-baedd supervised objective is only defined for memory banks of non-zero length, for the purpose of this ablation replace it with simple linear classifier when MM,l O. Right: For downsample=False, store representations of all patches into the memory bank. If downsample=True, we sample |M|/N patches per image (N is the length of the downstream training set), allowing for greater diversity.
Effect of evaluation memory length [M]. When transferring to downstream tasks with many training images (c.g. PASCAL VOC and ADE20K contain 10k and 20k images respectively, each image providing 100s of tokens), we see benefits of using large memory banks (c.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 22] 40) to limit cross attention (Equation T to small set of nearest neighbors for each query (e.g. k 30, Appendix| A.2 for details, where we find increasing & 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 million keys), the lookup cost is minimal (5 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 IMpl- In contrast to retrieval-based evaluation, we find contextual pretraining to be remarkably memory-efficient small memory banks (e.g. |M,]=40K,, see Figure 6 left for PASCAL VOC and Appendix [.3 for ADE20K) are sufficient to yield robust gains in retrieval-based scene understanding, adding relatively small computational overhead to training the representation (e.g. +22% for MMp = 40k). The module is agnostic to how the representation 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.
Conclusion
Inspired by impressive examples of in-context learning in language models, we investigate compo- nents necessary for in-context learning of dense scene understanding tasks in computer vision. To this end, we propose simple non paraameeettic nearest neighbor retrieval mechanism-wiic is agnostic to the downstream task and requires finetuning or specialized decoders- serve as a general-purpose decoder which we use to evaluate models on semantic segmentation and monocular depth estimation tasks. We further propose Hummingbird, 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 perform downstream tasks in a fast and data-efficient manner. By combining Hummingbirl 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 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 Hum- mingbird 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.
Acknowledgements
We thank Daniel Zoran, Andrew Zisserman, Evan Shelhamer and João Carreira for their thoughtful feedback, Skanda Koppula and Mathilde Caron for their assistance in reproducing baselines, and Aäron van den Oord and Oliver Vikbladh for fruitful discussions at the inception of the project.
10

---

References
[1] A. Aamodt and E. Plaza. Case-based reasoning: Foundational issues, methodological variations,
and system approaches. Al Communications, 7(1):39-59, 1994.
[2] J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lene, A. Mensch, K. Millican,
M. Reynolds, et al. Flamingo: A visual language model for few-shot learning. Advances in Neural Information Processing Systems, 2022.
[3] F. G. Ashby and W. T. Maddox. Human category learning, Annu. Rev. Psychol. 56:149-178.
2005.
[4] H. Bahng, A. Jahanian, S. Sankaranarayanan, and P. Isola. Exploring visual prompts for adapting
large-scale models. arXiv preprint vrXiv:2203.1727:: 2022.
[5] H. Bao, L. Dong, and F. Wei. BEÎT: BERT pre-training of image transformers. In International
Conference Learning Representations, 2022.
[6] A. Bar, Y. Gandelsman, T. Darrell, A. Globerson, and A. Efros. Visual prompting via image
inpainting. Advances in Neural Information Processing Systems, 2022.
[7] 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-52 2002.
[8] 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.
[9] O. Boiman, E. Shechtman, and M. Irani. In defense of heirest-neighbor based image classifi- cation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2008
[10] A. Brock, S. De. S. L. Smith, and K. Simonyan. hieghpperformncc large-scale image recognition without normalization. In International Conference on Machine Learning, 2021. [11] Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam,
G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in Neural
Information Processing Systems, 2020.
[12] 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.
[13] M. Caron, N. Houlsby, and Schmid. Location-aus aware self-supervised transformers. arXiv 2022.
preprint arXiv:
[14] 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.
[15] 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.
[16] 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.
[17] 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.
[18] X. Chen and K. He. Exploring simple siamese representation learning. In Proceedings of the
IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
[19] 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.
[20] 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.
[21] N. Dalal and B. Triggs. Histograms of oriented gradients for human detection. In Proceedings
of the IEEE Conference on Computer Vision and Puttern Recognition, 2005.
[22] 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 16x words: Transformers for
image recognition at scale. International Conference on Learning Representations, 2021. [23] 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.
[24] D. Dwibedi, Y. Aytar, J. Tompson, P. Sermanet, and A. Zisserman. With little help from my friends: Nearest-neighbor contrastive learning of visual representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
11

---

[25] 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 2015.
Vision,
[26] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, 2017.
[27] 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.
[28] 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.
[29] 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.
[30] 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 Computer Vision and Pattern Recognition, 2022.
[31] 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.
[32] 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.
[33] O. J. Hénaff, S. Koppula, J.-B. Alayrac, A. v. d. Oord, O. Vinyals, and Carreira. Efficient visual pretraining with contrastive detection. In Proceedings of the IEEE International Conference on
Computer Vision. 2021.
[34] O. J. Hénaff, S. Koppula, E. Shelhamer, D. Zoran, A. Jaegle, A. Zisserman, J. Carreira, and R. Arandjelović. Object discovery and representation networks. In Proceedings of the European Conference Computer Vision, 2022.
[35] O. 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.
[36] L. A. Hendricks and A. Nematzadeh. Probing image-language transformers for verb under- standing. In Findings of ACL, 2021.
[37] 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 computc-optimal large language models. In Advances in Neural Information Processing Systems, 2022.
[38] D. Homa, S. Sterling, and L. Trepel. Limitations of exemplar- generalization and the abstraction of categorical information. Journal of Experimental Psychology: Human Learning and Memory, 7(6):188, 1981
B. Hariharan, and S.-N. Lim. Visual prompt tuning. In Proceedings of the European Conference on Computer Vision, 2022. [40] J. Johnson, M. Douze, and H. Jégou. Billion-scale similarity search with GPUs. IEEE
[39] M. Jia, L. Tang. B.-C Chen, C. Cardie, S. Belongie,
Transactions Big Data, 7(3):535-5477 2019.
[41] 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.
[42] 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 Computer Vision, 2020.
[43] A. Krizhevsky, I. Sutskever, and G. E. Hinton. ImageNet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems, 2012.
[44] 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 Puttern Recognition, 2009.
[45] F. Liu, G. Emerson, and N. Collier. Visual spatial reasoning. arXiv preprint arXiv:2205.00336
2022.
[46] Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie, A ConvNet for the 2020s. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[47] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference
on Learning Representations, 2019.
12

---

[48] D. G. Lowe. Distinctive image features from scale-invarie invariant keypoints. International Journal of Computer Vision. 60:91-110, 2004.
[49] T. Malisiewicz, A. Gupta, and A. A. Efros. Ensemble of exemplar SVMMs for object detection and beyond. In Proceedings of the IEEE International Conference on Computer Vision, 2011.
[50] R. M. Nosofsky. Attention, similarity, and the identification categorization relationship. Journal
of Experimental Psychology: General, 1 115(1):39, 1986.
[51] M. Oquab, T. Darcet, T. Darcet, Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, P. Fermandez, D. Haziza, F. Massa, A. El-Nouby, et al. DINOv2: Learning robust visual features without supervision.
arXiv preprint
2023.
[52] 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.
[53] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agaawal, 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.
[54] 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. [55] S.-A. Rebuffi, H. Bilen, and A. Vedaldi. Learning multiple visual domains with residual adapters.
Advances in Neural Information Processing Systems, 2017.
[56] 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, (15(3):11-252) 2015.
[57] M. B. Sariyildiz, Y. Kalantidis, K. Alahari, and D. Larlus. No reason for no supervision:
Improved generalization in supervised models. In International Conference Learning Representations. 2023.
[58] R. C. Schank. Dynamic memory revisited. Cambridge University Press, 1999.
[59] 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.
[60] 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.
[61] 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.
[62] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image
recognition. In International Conference on Learning Reprexentations, 2015.
[63] R. Strudel, R. Garcia, Laptev, and C. Schmid. Segmenter: Transformer for semantic segmen- tation. In Proceedings of the IEEE International Conference on Computer Vision, 2021.
[64] 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.
[65] Y. Tian, O. J. Henaff, and A. van den Oord. Divide and contrast: suffrvuee learning from
uncurated data. In Proceedings of the IEFF/CVF International Conference on Computer Vision,
2021.
[66] 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(1 1958-1911970 2008
[67] H. Touvron, M. Cord, and H. Jégou. DeiT III: Revenge of the VIT. In Proceedings of the
European Conference on Computer Vision, 2022.
[68] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal,
E. Hambro, F. Azhar, et al. ILLMMa: Open and efficient foundation language models. arXiv
preprint arXiv:2302. 13971, 2023.
[69] H. Touvron, A. Sablayrolles, M. Douze, M. Cord, and H. Jégou. Grafit: Learning fine fine-grained image representations with coarse labels. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
[70] 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.
[71] O. Vinyals, C. Blundell, T. Lillicrap, K. Kavukeuoglu, and D. Wierstra. Matching networks for one shot learning. In Advances in Neural Information Processing Systems, 2016.
[72] 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.
13

---

[73] X. Wang, W. Wang, Y. Cao, C. Shen, and T. Huang. Images speak in images: A generalist painter for in-context visual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, 2023.
[74] Z. Wu, A. A. Efros, and S. X. Yu. Improving generalization via scalable neighborhood component analysis. In Proceedings of the European Conference on Computer Vision, 2018. [75] Z. Wu, Y. Xiong, S. X. Yu, and D. Lin. Unsupervised feature learning via non-parametric instance discrimination.
Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2.18.
[76] J. Xiao, J. Hays, K. A. Ehinger, A. Oliva, and A. Torralba. SUN database: Large-scale scene recognition from abbey to zoo. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, 2010.
[77] Xiao, Liu, Zhou, Y. Jiang, and Sun. Unified perceptual parsing for scene understanding.
In Proceedings of the European Conference on Computer Vision, 2018.
[78] E. Xie, J. Ding, W. Wang. X. Zhan, H. Xu, P. Sun, Z. Li, and P. Luo. Detco: Unsupervised contrastive learning for object detection. In Proceedings of the IEEE/CVF International
Conference Computer Vision, 2021.
[79] N. Xie, Lai, D. Doran, and A. Kadav. Visual entailment: A novel task for fine-grained image understanding. arXiv preprint arXiv:/901.06700
2019.
[80] J. Yu, Z. Wang, V. Vasudevan, L. Yeung, M. Seyedhosseini, and Y. Wu. CoCa: Contrastive captioners are image-text foundation models. In Transactions on Machine Learning Research, 2022.
[81] L. Yuan, D. Chen, Y.-.. 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 arXiv22111.11422,
2021.
[82] Y. Zhang, K. Zhou, and Z. Liu. What makes good examples for visual in-context learning? In Advances in Neural Information Processing Systems, 2023.
[83] B. Zhou, H. Zhao, X. Puig, T. Xiao, S. Fidler, A. Barriuso, and A. Torralba. Semantic understanding of scenes through the ADE20K dataset. International Journal of Computer Vision, 127:302-321, 2019.
[84] Zhou, C. Wei. H. Wang. W. Shen, Xie, A. Yuille, and Kong. iBOT: Image BERT pre-training with online tokenizer. International Conference on Learning Representations. 2022.
14

---

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 '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 spatial map k; = fo (x:;) &RH-W XxD, where a feature k/ E ERP at a given spatial location j is aligned with the local label 1/ created by averaging the pixel labels y in that patch. These features K; are then L2 normaazal
When the memory bank length is not large enough to accommodate all features for all images, it is necessary to subsample and only store subset of the features of each image. For concrete example using ADE20K, training set images have a resolution of 512 x 512 which when encoded by ViT-B/16 results in 32 32 grid of features (i.e. 1,024 features per image). To store every feature from each of ADE20K's i 20,120 training images would require a memory bank length of 20,120 32 32- 00,69,0040 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 IM
features to take per image as Nleeauuex_peer perriinnepper_ingee [D] =mmm angmentttiio 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, 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 refer to the label attached to the patch indexed by j in the image and let Loz = 1 if given class E P' and 0 otherwise. Then for each class we define Ke E, Leez (i.e. count of how many patches the class appeared in). We define a "class score" for each patch indexed by j as class_ caas_scoree = Eeee Ki. ceg Finally, we take the Pffeaturespper. per image from the spatial map k; with the lowest final scores using
(5) where x ~ Mo.1j- 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.
final score' (class_scoree xx) + (10% 1===)
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 Aiteatures_pez pic___prrree_eeeeeeeee image 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 Reauures per image from the list.
[M]
[D] We first randomly
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 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. we sequentially make predictions for each image in the evaluation set. Evaluation was done on 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 grid of features q fo (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 scores for the similarity that can be used as the attention logits. These scores
15

---

are then divided by temperature scaling value before having softmax applied to them to obtain the final attention values (see Equation
Throughout the paper, we use ScaNN in asymmetric hashing (AH) mode as opposed to brute-force negative impact on the evaluation from using approximate
mode. We find that there is little nearest neighbor search as opposed to brute-force exact search, despite the approximate search being several orders of magnitude faster. We use cosine similarity (Lz-normalized dot product) as a distance measure throughout this work. We also attempted some experiments using squared Euclidean distance and found it to have no benefits to performance for any of the models evaluated.
Table 6: NN retrieval hyperparameters. Note that training is involved with NN evaluation, hence there are no hyperparameters such as learning rates or training epochs.
Section 4.2
Everywhere else
M (Memory bank length)
20,480,000
10,240,000
k (nearest neighbors)
90
30
Temperature
.1
.02
Augmentation epochs
8
2
ScaNN dimensions Limensions_per_Dnee Bactoonssse__oeeteteee block
4
ScaNN num_leaver
512
512
ScaNN num jcaves_to_ssaetee
256
32
ScaNN reordering num neighbors
1800
120
Table 6 summarizes the hyperparameters used for NN evaluation throughout this work. For every 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 run and it is computationally fast to find nearest neighbors even with minimal approximations, hence we used slightly different setup in this regime. In general, we found nearest neighbor retrieval to be sur- prisingly robust to the choice of hyperparameters, with temperature and reordering tariig_uum__niihhhee being the most relevant performance. The same set of hyperparameters were used for the seman- tic 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 complete description of the meaning of the ScaNN hyperparame- ters, please see master/scann/ docs://///oo//ooeeeeeeee wttps//llloorittcoo .m
section except for Section 4.2
com/ gowglo-oogeooooooo
Table detailss 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.I When augmentations are enabled, the image is first scaled between the minimum and maximum scale factor, from which random crop is selected. Then photometric augmentations are applied independently with the probabilities and maximum intensities provided.
Table 7: Evaluation augmentations. Parameters used to augment the training dataset for semantic segmentation.
Parameter
Random crop probability
1.0
Minimum scale factor
0.5
Maximum scale factor
2.0
Brightness jittering probability
0.5
Contrast jittering probability
0.5
Saturation jittering probability
0.5
Hue jittering probability
0.5
Brightness adjustment max
0.1
Contrast adjustment max
0.1
Saturation adjustment max
0.1
Hue adjustment max
0.1
16

---

B Implementation details: contextual pretraining
The contextual pretraining module takes input batch of image representations (i.e. queries) = h = fo( E 1 DD from the VIT encoder fa. where / = 4096 is the batch size, H W
are the height and width of the spatial feature map and D 768 for VTT-B and D = 1024 for ViT-L is the feature dimension. Keys and values for the contextualization cross- attention operation are entries of the memory bank My {(k,, 02). i == 1=1..... [Mp]], where keys k; are taken from previous batches by spatially averaging h (see Equation 2 and values are obtained by applying two-layer MI.P Φo to the keys, where we use batch norm after the first layer and the hidden dimension is set to 4096. Each feature q' of the image representation is then updated as
c¹ - v00((1 - X) प + TT)) where the is a linear layer and || is the L2 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- -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- (256 TPUs with 600 entries for ImageNet- 22k), 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 I2 and Z2. The augmentations constructed as compositions of the following operations, each applied with a given probability:
1. random cropping: random patch of the image is selected, whose area is uniformly sampled in [0.08 A, AI, where AA 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 x 224 pixels using bicubic interpolation;
2. horizontal flipping:
3. color jittering: the brightness, contrast, saturation and hue are shifted by uniformly distributed offset;
4. color dropping: the RGB image is replaced by its grey-scale values;
5. gaussian blurring with a 23 23 square kernel and standard deviation uniformly sampled from 2.0];
6. solarization: point-wise color transformation x x 12000 + (1 - x)) 1 1220. with pixels in [0, 1].
The augmented images 11 and Ø2 result from augmentations sampled from distributions Ti and To respectively. These distributions apply the primitives described above with different probabilities and different magnitudes. Table specifies these parameters for the B BYL framework (28), which we adopt without modification.
Table 8: Pretraining augmentations. Parameters used to generate different views of a single image for contrastive pretraining.
Parameter
T
T2
Random crop probability
1.0
Flip probability
0.5
Color jittering probability
0.8
Color dropping probability
0.2
Brightness adjustment max
0.4
Contrast adjustment
0.4
Saturation adjustment max Hue adjustment max
0.2
0.1
Gaussian blurring probability
1.0
0.1
Solarization probability
0.0
0.2
17

---

C.2
Optimization
We pretrain the model for 300 epochs ImageNet-1K k or 100) epochs on ImageNet-22K using A 47 with a batch size of 4096, split across 128 Cloud TPU v3 workers for ImageNet- Ik and 256 Cloud TPU v3 workers for ImageNet-22k. Training VIT-B / ViT-L for 300 epochs ImageNet-Ik takes roughly 21 hours 53 hours, while 100 epochs on ImageNet-22k takes approximately 60 hours / 128 hours. We update the online parameters 0 with cosine learning
schedule with 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 decay rate of 0.99.
exponential moving average of the online parameters with
Following [16] the projections and predictions in Equation 4 Hare normalized and rescaled such that their norm equal to 1/ 1/vT 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 a 0.05 for all other experiments.
D Supplementary analysis
D.1 Data efficiency
In T we compared Humminghird with several leading representation learning techniques in the low-data regime. Here we provide the complete analysis from 1/128 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 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 linear decoder for all of the finetuned rows in each table.
For PASCAL VOC (Table 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 ImageNet- 1k) 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-wiis only nearest neighbors (no finetuning), Humminghird++ trained on ImageNet-224 with a VIT-L backbone beats all of the other finetuned variants for every dataset fraction. Hummingbird++ using a ViT-B and ImageNet-II Ik predictably lies in-between the other two models for every dataset fraction.
For ADE20K (Table I0,, the same general trends from above hold. Backbone and dataset scaling are once again beneficial as Hummingbird++ with ViT-L and ImageNet-221 22k training outperforms the other Hummingbird models, however this time the absolute performance relative the finetuned competition in the high-data regime is less favorable since the end-to-ent finetuned versions of other techniques start to outperform the nearest neighbors only ViT-I. Humminghird++ 1/16 of the data.
Table 9: PASCAL VOC data efficiency analysis. After pretraining, models are applied to down- stream tasks with the indicated fraction of the dataset size. Models perform the task either with end-to-end fine-tuning with a linear head (E2E FT) or with our mechanism for in-context scene understanding using nearest neighbors at evaluation time (NN). All fine-tuning runs are averaged over five different seeds. The metric reported is mean IoU (higher numbers are better). denotes models trained on ImageNet-22k: all other models were trained on ImageNet- 1k.
Fraction of dataset
Mcthod
Decoder
Backbone
1/128
1/64
1/32
1/16
1/8
1/4
1/2
1/1
DeiT-III
67
E2E FT
ViT-B
41.8
53.8
63.1
67.7
70.7
72.2
73.4
75.2
DINO 15
E2E FT
ViT-B
36.1
44.3
54.3
57.8
61.7
64.8
68.2
72.2
MoCo-v3
20
E2E FT
ViT-B
19.9
33.4
47.0
54.8
61.5
67.1
70.7
73.4
MAE
E2E
ViT-B
34.2
53.0
58.7
62.7
70.8
73.5
LOCA
13
E2E FT
ViT-B
40.1
53.9
63.1
67.8
70.7
72.8
74.4
75.5
Hummingbird
NN
VIT-B
50.5
57.2
60.1
62.6
64.3
65.9
68.9
71.8
Humminghird++
NN
VIT-B
52.4
57.3
61.5
64.6
66.2
67.9
70.5
73.2
Humminghird+4*
VIT-L
61.8
65.3
68.0
70.7
71.4
73.2
77.2
18

---

Table 10: ADE20K data efficiency analysis. After pretraining, models are applied to downstream tasks with the indicated fraction of the dataset size. Models perform the task either with end-to-end fine-tuning with a linear head (E2E FT) or with our mechanism for in-context scene understanding using nearest neighbors evaluation time (NN). All fine-tuning runs are averaged over five different seeds. The metric reported is mean IoU (higher numbers are better). The results for other techniques between 1 1/32 and 1/1 are sourced directly from [13, the rest are reproductions. denotes models trained on ImageNet-22k; all other models were trained on ImageNet-
Fraction of dataset
Method
Decoder
Backbone
1/128
1/64
1/32
1/16
1/8
1/4
DeiT-III
E2E
ViT-B
10.8
14.3
20.9
27.1
32.7
47.3
DINO
E2E FT
ViT-B
11.7
14.4
18.4
24.5
29.5
35.2
39.5
44.1
MoCo-v3 20
E2E FT
ViT-B
4.6
7.9
17.7
25.2
30.8
36.5
40.7
45.4
MAE
30
E2E FT
ViT-B
8.2
12.2
18.4
25.3
30.5
36.1
40.6
45.5
LOCA
13
E2E FT
VIT-B
11.2
15.5
22.2
30.0
34.4
39.1
42.8
47.9
Humminghird
ViT-B
11.7
15.1
17.3
20.0
22.3
24.9
Hummingbird++
NN
ViT-B
12.7
16.4
18.9
21.5
24.0
26.8
29.9
32.0
Hummingbird++*
NN
VIT-L
16.6
20.5
24.0
27.4
30.2
33.1
36.0
37.8
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 Hummingbirl models trained with different architectures (VTT-B vs ViT-L), datasets (ImageNet- 1k vs 2n2g), learning objectives self-superviied 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 decoding scheme is highly predictive of the other (Pearson's 's = 0.80 for PASCAL VOC, 0.89 for ADE20K). As such, in cases where NN retrieval underperforms end-to-end finetuning, it still be used powerful diagnostic tool. As illustrated in Section 4.3 evaluating with NN retrieval is much simpler and faster than with 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 10x less variable across runs and doesn't 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.
when using linear decoder. End-to-end finetuning
PASCAL: NN retrieval vs E2E FT
ADE20K: NN retrieval vs E2E FT
85
End-to-ent end fine- se-tuned
fine-tuned
84
50
83
82-
45
70
72
74
76
78
30
32
34
36
38
NN retrieval
NN retrieval
Figure 7: Relation between NN retrieval and end-to-end finetuning performance. We collect 14 models trained with different architectures, datasets, and learning objectives.
D.3 Effect of pretraining and evaluation memory length for ADE20K
We include the equivalent of Figure 6 on the ADE20K dataset Figure g Similar to what we observe for PASCAL VOC, we benefit from large memory banks evaluation. Since the ADE20K
19

---

training set is roughly 2x larger than that of PASCAL VOC, we also observe that sampling which features to store in the memory bank is more important than is for PASCAL VOC (see Appendix A. on the details of the sampling procedure). Similarly, at training time, ADF20K benefits from larger pretraining memory banks than PASCAL VOC, with performance plateauing for memory banks larger than 200,000. Thus, set the pretraining memory bank length to 1 153,600 in all our experiments (see Appendix for details on contextual pretraining).
ADE20K: Importance of pretraining memory
ADE20K: Importance of evaluation memery
33
25
Mean lou
Mean tou
31
20
supervised
downsample-True
supervised
downsample-Falss
50000
L50000
250000
10a
Training memory length
Evaluation memory length
Figure 8: Effect of the pretraining (left) and evaluation (right) memory length on performance of ADE20K. All models were pretrained with ViT-B on ImageNet- Left: Since the retrieval- supervised objective is only defined for memory banks of non-zero length, for the purpose of this ablation we replace it with a simple linear classifier when [Mpl =0, Right: For downsample-False store representations of all patches into the memory bank. If downsample=True we sample |MI/N patches per image (N the length of the downstream training set), allowing for greater memory bank diversity and thus superior performance than when 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 ViT-L is beneficial for both Humminghird and Hummiiggirdd++ where the suff-super 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.
Table 11: Impact of encoder size on finetuned performance. After pretraining, models are equipped with task-specific decoders and finetuned for that task on the entire downstream dataset.
Finetuned accuracy (mloU)
Method
Encoder
Dataset
PASCAL
ADE20K
Hummingbird
VIT-B
INIK
80.0
44,9
Humminghird
ViT-I
INIK
82.4
47.1
Hummingbird++
VIT-B
INIK
81.2
44,9
Hummingbird++
ViT-L
INIK
81.8
47.3
Hummingbird Hummingbird
VIT-B
IN22K
81.6
46,9
ViT-L
IN22K
84.1
50.8
Hummingbird++ Humminghird++
VIT-B VIT-L
IN22K IN22K
82.1
48.2
85.3
52.0
D.5 Importance of retrieval-based supervision for in-context scene understanding
We study the importance of retrieval-bacd supervised objective (see Section 3.4 on in-context scene understanding performance. We compare model trained purely with the retrieval-based supervised objective ("Sup"") with Hummingbird (purely self-supervised, i.e. "SSL") and Humminghird++ (both self-supervised and retrieval-basee supervised, i... "SSL Sup"). Results shown in Table 12 indicate
20

---

the necessity of the self-supervised objective for creating representations that are general and transfer well to downstream tasks.
Table 12:
Importance
ectrieval-based supervision in-context scene understanding. All
models are pretrained on source data and applied to downstream datasets without modification. All downstream tasks are performed using nearest neighbor retrieval. All models use VIT-B as an encoder.
Semantic segmentation
Depth pred.
Method
Ohjective
Dataset
PASCAL
ADF20K t
NYUv2 +
Supervised retrieval Hummingbird
Sup
INIK
56.9
22.9
.787
SSL
INIK
70.5
28.3
.718
Hummingbird++
SSL Sup
INIK
72.1
30.5
.738
Supervised retrieval
Sup
IN22K
69.3
30.3
.739
Hummingbird
SSL SSL + Sup
IN22K
73.5
30.7
.706
Hummingbird++
IN22K
76.2
34.1
.695
21


Reviews:
[{"review_id": "t5oXL7p3KU", "reviewer": "Reviewer_UEmV", "summary": "This 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. ", "questions": "1. 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. ", "limitations": "Yes. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. 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. ", "weaknesses": "1. 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\nWeaknesses\n1. We agree that further scaling Hummingbird (e.g. to ViT-H) would require scaling the dataset in tandem. For this, we could leverage standard extra-large datasets such as ALIGN, and adapt the supervised retrieval objective to the multimodal contrastive setting. Note that our retrieval-based supervised objective straightforwardly applies here: predictions for the textual modality would be linear combinations of textual embeddings of nearest neighbors, and vice-versa for the image modality.\n\n2. We agree that this would be an interesting set of experiments to run. Unfortunately, for these two sections we relied on an open-source framework to run the experiments for the baselines for DINO, MAE, MoCo-v3 etc. which had well-tested support for segmentation tasks but not depth related tasks, so collecting this data would require substantial new code beyond our team's current capacity in the rebuttal period.\n\nQuestions\n\n1. If this question refers to why we use mean-pooled hidden representations as opposed to attention-pooled ones as keys of the memory bank: we have experimented with attention pooling, which reduced performance slightly compared to mean pooling. We will add this ablation to the appendix.\n\n    If this question refers to the softmax in cross-attending over the memory bank: the cross-attention operation $\\text{CA} (q^i, k_j, v_j)$ (see L132) actually does contain a softmax (please see Eq 1 for the definition of the cross-attention operation) and Eq 2 describes how the keys and values of the memory bank at pretraining time are formed.\n\n2. We are working on the open source implementation and hope to have it ready by the camera-ready deadline. \n\n3. Indeed it does. Hummingbird ViT-L full finetuning results weren’t included in Table 4 due to space constraints but we will add the table below in the appendix. We can see that scaling the model size to ViT-L also helps in the finetuning regime, especially when scaling the dataset size accordingly. Note that this scaling study highlights the need for jointly scaling data and model size, in line with the discussion around Weakness #1. \t\t\n\n| Model | Backbone | Dataset | PASCAL | ADE20K | Difference ViT-B -> ViT-L\n|:----------------|:-------------:|:-------------:|:-------------:|:-------------:|:-------------:|\nHummingbird\t| ViT-B\t | IN1K \t| 80.0 | 44.9 | |\nHummingbird\t| ViT-L\t | IN1K \t| 82.4 | 47.1 |  +2.4%, +2.2% |\nHummingbird++ \t| ViT-B\t | IN1K\t| 81.2 | 44.9 | |\nHummingbird++\t| ViT-L\t | IN1K \t| 81.8 | 47.3 |  +0.6%, +2.4% |\nHummingbird \t| ViT-B\t | IN22K \t| 81.6 | 46.9 | |\nHummingbird \t| ViT-L\t | IN22K \t| 84.1 | 50.8  | +2.5%, +3.9% |\nHummingbird++ \t| ViT-B\t | IN22K  | 82.1 | 48.2 | |\nHummingbird++ \t| ViT-L\t | IN22K  | 85.3 | 52.0 | +3.2%, +3.8% |\n\n4. In Sections 4.2 and 4.3 we compare Hummingbird with nearest neighbor retrieval to Hummingbird with linear probing (i.e. frozen backbone finetuning) and find that retrieval consistently outperforms linear probing. The question then remains how Hummingbird retrieval and Hummingbird linear probing compare to linear probing for other techniques in Table 2 (full nearest neighbor and full fine-tuning comparisons are available in Tables 1 and 4, respectively). To answer this we re-ran the baselines on PASCAL VOC semantic segmentation and found that Hummingbird linear probing far exceeds MAE and MoCo-v3 but slightly lags behind DINO and LOCA in the linear probing regime, though surprisingly, Hummingbird paired with NN retrieval still exceeds all other techniques by a few points, even when they use linear probing rather than NN retrieval. \n\n    This raises an interesting area for future work - more rigorously understanding which properties of the different algorithms lead to good performance given NN decoding vs linear probing vs specialized decoders (e.g. UperNet). Indeed Shi et al. [1] explicitly showed a trade-off between \"label efficiency (the ability to learn an accurate classifier on top of the representation with a small amount of labeled data)\" and \"universality (usefulness across a wide range of downstream tasks)\" in the context of contrastive learning with linear probing.\n\n5. Preliminary experiments ​​with the supervised weight $\\alpha$ indicated lower values of $\\alpha$ to be necessary for strong downstream task performance (we set $\\alpha$ to 0.25 for the supervised ViT-B trained on ImageNet-22k and to 0.05 for all other experiments; see Appendix C.2). \n\n    The ablation below, which is in line with our preliminary findings, shows the results for self-supervised only (Hummingbird), supervised retrieval only, and both objectives (Hummingbird++) for a ViT-B trained on ImageNet-1k and ImageNet-22k. These results highlight the importance of the self-supervised objective in creating general representations that transfer well to downstream tasks. We will include this table in the appendix of the final version.\n\n| Model | Backbone | Dataset | PASCAL | ADE20K | NYUv2 |\n|:----------------|:-------------:|:-------------:|:-------------:|:-------------:|:-------------:|\nHummingbird (SSL)\t\t| ViT-B\t | IN1K \t| 70.5 | 28.3 | .718 |\nHummingbird++ (SSL + Sup)\t| ViT-B\t | IN1K\t| 72.1 | 30.5 | .738 |\nSupervised retrieval only (Sup)\t| ViT-B\t | IN1K \t| 56.9 | 22.9 | .787 |\nHummingbird (SSL)\t\t| ViT-B\t | IN22K \t| 73.5 | 30.7 | .706 |\nHummingbird++ (SSL + Sup)\t| ViT-B\t | IN22K\t| 76.2 | 34.1 | .695 |\nSupervised retrieval only (Sup)\t| ViT-B\t | IN22K \t| 69.3 | 30.3 | .739 |\n\n\n[1] Z. Shi, J. Chen, K. Li, J. Raghuram, X. Wu, Y. Liang and S. Jha. The Trade-off between Universality and Label Efficiency of Representations from Contrastive Learning. 2023."}, {"review_id": "LVsssff0MC", "reviewer": "Reviewer_32Tr", "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.", "questions": "- Why does the proposed contrastive learning objective enable the encoded features to be directly applied to downstream tasks without any parameter adaptation?\nThis 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper shows that non-parametric nearest neighbor retrieval method can be more effective than linear probing in existing self-supervised learning approaches.\n- The self-supervised pretraining does not necessitate either fine-tuning or specialized decoders.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\nWeaknesses\n\n* We demonstrate the generality of Hummingbird representations by evaluating the model on two different tasks (semantic segmentation and depth prediction), of which one is a pixel-level classification task and the other is a regression task. We agree with the reviewer that evaluating on other downstream tasks would be interesting and valuable, and leave this investigation to future work.\n\n* \n  * Low-data regime: In the low-data regime, we compare Hummingbird equipped with the general-purpose nearest neighbor decoder to fully finetuned baselines which can specialize to the task at hand. This puts these baselines at an advantage, particularly for complex datasets such as ADE20K. Despite this disadvantage, Hummingbird is only worse than one baseline (LOCA) in one setting (n=316), while outperforming it in the more challenging n=158 setting, and outperforms it convincingly in all settings on PASCAL (see Table 2).\n\n  * Full fine-tuning: As mentioned in the introduction (L32), the focus of this work is fast and effortless adaptation to downstream tasks (as these are the properties needed for \"in-context\" scene understanding). In this domain, Hummingbird is indeed superior to prior work. Since fully-finetuned performance wasn’t the focus of the work, we didn't spend as much time tuning the pretraining or fine-tuning hyperparameters here for every variant of Hummingbird, though this would likely help explain/close the gap. We agree that further improving fully-finetuned self-supervised Hummingbird's performance is an interesting avenue for future work as it would allow us to have a single framework be state-of-the-art across many disparate regimes. It is again noteworthy though that the weakest Hummingbird model still outperforms the rest on PASCAL.\n\nQuestions\n\n* Contrastive learning in general has indeed been shown to yield semantically relevant features (see e.g. SimCLR [1] and variants thereof), and our use of spatial attention pooling makes these features locally discriminative, preparing them for local semantic tasks such as semantic segmentation. The contextual pretraining of Section 3.2 also helps push the representations to have semantically meaningful nearest neighbors. Moreover, semantically relevant features have been shown to be relevant for geometric tasks such as depth prediction (see e.g. [2]) which has been confirmed by multiple works in SSL (e.g. [3,4]) as well as our empirical results.\n\n[1] 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.\n\n[2] ​​Hickson, S., Raveendran, K., Fathi, A., Murphy, K. and Essa, I. Floors are flat: Leveraging semantics for real-time surface normal prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019.\n\n[3] C Doersch, A Gupta, AA Efros. Unsupervised Visual Representation Learning by Context Prediction.  Proceedings of the IEEE International Conference on Computer Vision, 2015.\n\n[4] Grill, J. B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., ... & Valko, M. Bootstrap your own latent - A new approach to self-supervised learning. Advances in Neural Information Processing Systems, 2020.\n\n\n"}, {"review_id": "Oxfp9vtvo8", "reviewer": "Reviewer_hsMp", "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.", "questions": "1. In Tab. 3, why does the result of MoCo-v3 significantly drop for E2E FT?\n2. 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.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The proposed method addresses a promising regime for reducing task-specific training requirements with prompt images and labels.\n2. The framework is conceptually simple yet powerful, providing a concise and reasonable solution.\n3. 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.\n2. The retrieval process is inefficient compared to previous paradigms during inference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\nWeaknesses\n\n1. We focused on the pixel-level dense tasks in this paper due to their relative novelty in the nearest neighbors domain, however, the framework can be used in a straightforward manner for other tasks. For instance, image-level tasks such as classification are very easy to support by pooling the patch-level features for a prompt image and storing this pooled key in the memory bank. The query image can then pool its features and query the memory bank to derive its correct label at the image (rather than pixel/patch) level. For other use cases, nearest neighbor evaluation has historically been used for a wide variety of tasks, such as shape matching [6, 7, 56], scene recognition [62, 71], and image parsing [43, 47]. We discuss some of these examples in the Related Work section of the paper and leave the exploration of these other tasks with our model for future work.\n\n2. We find that thanks to efficient open-source approximate nearest-neighbor search libraries, we are able to keep the additional inference time necessitated by our approach to a minimum. We've attached a graph showing the overhead of the nearest neighbors lookup vs. the length of the memory bank on an A100 GPU in the attached 1-page PDF and will add this graph to the appendix of the final version. For small-to-medium sized memory banks (0 to 1 million keys), the look-up cost is minimal (<= 30ms), meaning the system is still fast enough to be used for real-time applications, e.g. segmenting videos at 30 frames per second. When scaling to very large memory banks (10 million+ keys), the scaling tends to be linear (each doubling of the memory bank size results in roughly a doubling of the inference time). However, the absolute performance is still reasonable: with a memory bank size of 10 million, the overhead from NN lookup is only 0.2 seconds for dense tasks, likely still suitable for most applications. One useful property of the nearest neighbors libraries is that they provide many options to trade-off precision of the NN search with inference time. In this paper we tuned the system for the best performance on the downstream tasks as .2 seconds per image was deemed fast enough for our purposes.\n\n\nQuestions\n\n1. In this “Fast adaptation” setting, the model is only allowed one pass over the downstream dataset during finetuning, equivalent to 1 epoch. This is different from the typical end-to-end fine-tuning regime in which the model performs multiple passes over the dataset. Note that different models generalize differently to this regime: MAE and MoCo-v3 are slower to adapt, whereas DINO is more robust, despite MAE and MoCo-v3 faring better when given a full compute budget. The fact that different models perform better at slow vs. fast adaptation constitutes an interesting research finding which we believe warrants further investigation. \n\n2. NN evaluation makes predictions as linear combinations of the labels of nearest neighbor examples. Contextual pretraining aligns the representations with this kind of evaluation by expressing an example’s representation as a linear combination of the representations of its nearest neighbors. By forcing a proportion (determined by the parameter $\\lambda$ in L133) of the representation to come from the nearest neighbors before making a prediction, the network learns representations that aren't just good at solving the self-supervised objective (Eq. 4), but good at doing so by attending to their neighbors.\n"}, {"review_id": "TTRJnhoAwO", "reviewer": "Reviewer_KiUe", "summary": "This 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.", "questions": "Please address concerns raised above. ", "limitations": "The 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?", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I 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", "weaknesses": "First, 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? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\nWeaknesses\n\n* (1) Thank you for the suggestion, we will add a figure showing model components of Hummingbird in the final version. (2) Nearest neighbors are patch-level, we will edit the figure to clarify. (3) The input prompt images are not curated and they contain various kinds of images (including e.g. horses, people, furniture, planes, buildings, etc.) and each patch from the query image attends over patches of all those images to select the nearest neighbors. In the case of the person on a horse query image, its nearest neighbors per the Hummingbird model are the patches from images of people on horses and not the furniture images. In other words, each patch of the query image (e.g. a patch of a horse) will have a similarity score with every patch in the memory bank (which contains many patches, some of horses, some of furniture, some of other things), and the label assigned to a horse patch in the query image will be the weighted average of the patches it is closest to in the memory bank.\n\n* L142-144 This claim follows naturally from the contrastive objective, which learns representations which maximally distinguish an example from others in the batch. When equipped with a spatial attention mechanism, this representation can in particular learn to select the most unique part of the image. This property has been found empirically in [50], which we use as the basis for our spatial attention pooling mechanism. Indeed in [50], the learned attention masks are shown to attend to the main object in an image.\n\n* L251-254 We've attached a graph showing the overhead of the nearest neighbors lookup vs. the length of the memory bank on an A100 GPU in the attached 1-page PDF. For small-to-medium sized memory banks (0 to 1 million keys), the look-up cost is minimal (<= 30ms). When scaling to very large memory banks (10 million+ keys), the scaling tends to be linear (each doubling of the memory bank size resulting in roughly a doubling of the inference time). However, the absolute performance is still decent: with a memory bank size of 10 million, the overhead from NN lookup is only .2 seconds for dense tasks, likely still suitable for most applications. One useful property of the nearest neighbors libraries is that they provide many options to trade-off precision of the NN search with inference time. In this paper we tuned the system for the best performance on the downstream tasks as .2 seconds per image was deemed fast enough for our purposes.\n\n* Cost of pre-training stage: The cost is similar to other contrastive pretraining frameworks. As mentioned in Section 5, contextual pretraining adds a relatively small computational overhead to training the representation (e.g. +22% for $|\\mathcal{M}_p| $= 40k relative to standard contrastive pretraining). Total computational requirements for the experiments are reported in section C.2 of the supplementary material.\n\n* We note that the focus of our work was _in-context_ learning of novel scene understanding tasks, which is why we didn’t investigate training on the downstream tasks themselves. All of our models are trained on ImageNet-1k / - 22k and adapted without any change in the parameters to semantic segmentation and depth prediction. Nevertheless this is an interesting question which has partially been addressed in concurrent work (e.g. Visual Token Matching [40]) which showed that such transfer across tasks is indeed possible. Our framework could straightforwardly be adapted to this multitask setting, by framing all downstream tasks as memory retrieval during pretraining rather than during evaluation only.\n\nLimitations\n\nThis is an important question, and we will add the following discussion to the appendix. We've attached a graph showing the overhead of the nearest neighbors lookup vs. the length of the memory bank on an A100 GPU in the attached 1-page PDF and we will also add it to the appendix of the final version. The graph shows that the inference time concerns are not as large as one might fear - even with a huge memory bank size of 20 million (far larger than the ideal few-shot/in-context setup), the inference time is less than half a second. For reasonably sized memory banks, even up to 1 million, the inference cost of the NN look-up is nearly free (less than 30 ms per image, making it fast enough to apply to tasks like video processing at 30 frames per second).\n\nThe added overhead for contextual pretraining is similarly manageable. Figure 4 shows that the optimal pretraining memory bank size is much smaller for pretraining (O(10^4)) than evaluation (O(10^6)). It is important to note that these two numbers (pretraining memory bank length and inference time memory bank length) have no relationship to one another, meaning we can pretrain with a memory bank length of X and evaluate with a memory bank length of Y. In Section 5, we show that contextual pretraining adds only a relatively small computational overhead to training the representation (e.g. +22% runtime for $|\\mathcal{M}_p|$ = 40k relative to standard contrastive pretraining).\n\nIt is also important to note that after the one-time cost of pretraining, Hummingbird can be applied to different downstream tasks without any finetuning, which saves on the computation needed to perform inference on a diverse set of tasks.\n"}, {"review_id": "tCA6ae3Ro1", "reviewer": "Reviewer_35vE", "summary": "The 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", "questions": "Major:\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", "limitations": "Please 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? \n", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. 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", "weaknesses": "Major:\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\n\n\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\nMajor\n\n1. \n   \n* “The features are retrieved independently of each other and there is no mechanism to consolidate the retrieved features into a whole scene output.” We agree with the reviewer that the neighbors are retrieved independently for each feature, but the features themselves are strongly dependent on one another (e.g. since they may correspond to the neighboring parts of an image), which then leads to coherent predictions, as seen in Fig 1. This happens despite the lack of a more complex consolidation mechanism, as is becoming standard given powerful encoder networks (see e.g. linear probing with DINO v2 [1], or per-feature object detection in OWLViT [2]).\n\n* “Correspondence - with the current approach, the authors assume there are spatial correspondences between features and their associated labels.” We acknowledge that certain challenges like image rotation and image flipping do not work well with this framework and will add this to the Limitations Section. Other advanced tasks such as image translation could be more amenable to this setup, by leveraging more complex mechanisms like \"representational algebra\" as proposed in [3] and as is common in NLP (e.g. [4]).\n\n * “Retrieval approach might not result in smooth outputs.” Indeed the discrete set of nearest neighbors will change as a function of the input. However the final prediction is a linear combination of their (continuous or discrete) labels, and the weighting will vary smoothly with the input representation. As a result, the NN predictions are as smooth as the underlying representation, particularly when using sufficient nearest neighbors (e.g. 30 or 90 in our experiments).\n\n* Style transfer: Somewhat surprisingly, NN (combined with a decoder) has recently been shown to be a useful component in SotA style transfer systems. See the diagram at the top of page 3 of Kolkin et al. ([5]) for a description of such a system. To do style transfer in our framework, the prompt would be composed of one or more style images, with the query image being the content image. However, as the reviewer correctly notes, with NN retrieval alone the image generated will not be perfectly smooth across patches. It would be necessary to add a decoding step as done in [5] (using our ViT backbone in place of VGG) to guide the content image towards the desired style.\n\n    We will include the summary of the points above in the Limitations Section and state more explicitly what the capabilities and limitations of the proposed framework are. Specifically, we note that our framework 1) relies on the dependencies between features to perform joint inference, and does not contain any notion of uncertainty when combining the predictions of multiple locations, 2) relies on a simple pairing of images and labels, and more complex relational reasoning (e.g. flipping labels) remains untested, and 3) relies on smooth features and dense memory banks to produce smooth predictions. We believe these points are the most important limitations of the approach.\n\n2. The model is pretrained with both attention across images (through contextual pretraining) as well as attention within an image (through spatial attention pooling). At test time, NN evaluation makes predictions by linearly combining the labels of nearest neighbor examples. At training time, Hummingbird’s contextual pretraining mechanism aligns the representations with this kind of evaluation by expressing them as a linear combination of the representations of nearest neighbor examples. Even though adding retrieval-based supervision (i.e. Hummingbird++) brings pretraining even closer to how the model is evaluated, we find the self-supervised component to be crucial for maintaining generality and transferability of learned representations across tasks (see answer to Q5 of reviewer UEmV for an ablation and a more detailed discussion on the role of supervision).\n\nMinor\n\n1. Thank you for the suggested references, we will cite them in the final version of the paper. To discuss each in more detail:\n    * [1] is based on choosing good images for the prompts to be used in Bar et al. (2022) \"Visual Prompting via Image Inpainting\", a paper we cite. [1]'s unsupervised method of choosing prompts picks images for the prompt such that they are most similar to the query image using a CLIP-based encoder. Intriguingly, their conclusion notes \"CLIP would not be sufficient… A model that can better balance spatial and semantic closedness in feature space would be more ideal for visual in-context learning.\" Our localized and semantic features seem ideal for this use-case, making exploring their combination an interesting possibility.\n\n    * For [2], in addition to the difference in pretraining objective (colorization vs. contrastive learning) and backbone (ViT vs. CNN), other key differences are that their memory bank is much smaller (\"four gray-scale video frames down-sampled to 256 × 256\") and that the \"model points within a single training example rather than across training examples\". Thus the segmentation they describe is more applicable to tracking frames within a single video, rather than being general purpose.\n\n2. We chose Hummingbird due to the model’s fast adaptation properties. Hummingbirds are known for their nimbleness (e.g. being able to fly both forward and backward and change directions quickly).\n\nOther\n* L109: Thank you, we will amend the inconsistency.\n* ViT-B (and ViT-L) patch size is 16x16 everywhere in the paper. We will include this in the paper.\n\n[1] Oquab, Maxime, et al. \"Dinov2: Learning robust visual features without supervision.\" 2023\n\n[2] Minderer, Matthias, et al. \"Simple open-vocabulary object detection.\" 2022\n\n[3] Eslami, SM Ali, et al. \"Neural scene representation and rendering.\" 2018\n\n[4] Mikolov, Tomas, et al. \"Efficient estimation of word representations in vector space.\" 2013\n\n[5] Kolkin, Nick et al. \"Neural Neighbor Style Transfer.\" 2022\n\n"}]

Author rebuttal:
Thank you for your time and effort in reviewing our work. 

Please see the attached PDF with a graph showing the overhead of the nearest neighbors lookup vs. the length of the memory bank on an A100 GPU.

Meta-review: This paper introduces Hummingbird, a novel representation learning method designed for in-context learning in dense prediction tasks. Hummingbird enhances self-supervised learning by incorporating a spatial attention-pooling mechanism and introducing attention between images, along with an optional retrieval objective during pre-training. During inference, it enables few-shot prediction through retrieval from a limited set of prompted images without any parameter updates. The effectiveness of this method is validated for semantic segmentation and depth estimation.

The paper initially received positive reviews with acceptance recommendations from all reviewers except one. The main concerns pointed out by reviewers related to the scalability of the approach due to retrieval complexity, the applicability of the approach to various scene understanding tasks, and the justifications for pre-training motivation.
The rebuttal effectively addressed these concerns, leading to a consensus to accept the paper after the authors' feedback.
The AC has thoroughly reviewed the submission and the discussions. The AC believes that the method makes a substantial contribution to adapting in-context learning for dense prediction tasks in visual recognition, potentially yielding a significant impact on the community. The proposed retrieval approach for prediction and the adapted pre-training appear to be sound. Consequently, the AC recommends acceptance but strongly advises the authors to include information regarding the scalability and applicability of their approach when preparing the final paper.
Decision: Accept (spotlight)