SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The contextual pretraining module takes as input a batch of image representations (i.e. queries)  $\mathbf{q} = \mathbf{h} = f_{\theta}(\mathbf{x}) \in \mathbb{R}^{B \times H \cdot W \times D}$  from the ViT encoder  $f_{\theta}$ , where  $B = 4096$  is the batch size,  $H = W = 14$  are the height and width of the spatial feature map and  $D = 768$  for ViT-B and  $D = 1024$  for ViT-L is the feature dimension. Keys and values for the contextualization crossattention operation are entries of the memory bank  $\mathcal{M}_p = \{(k_i, v_i), i = 1, \dots, |\mathcal{M}_p|\}$ , where keys  $k_i$  are taken from previous batches by spatially averaging  $h$  (see Equation 2) and values  $v_i$  are obtained by applying a two-layer MLP  $\phi_{\theta}$  to the keys, where we use batch norm after the first layer and the hidden dimension is set to 4096. Each feature  $q^i$  of the image representation is then updated as  $\boldsymbol{c}^i = \psi_{\theta}((1 - \lambda)\frac{\boldsymbol{q}^i}{\|\boldsymbol{q}^i\|} + \lambda\frac{\hat{\boldsymbol{v}}^i}{\|\hat{\boldsymbol{v}}^i\|})$ , where  $\psi_{\theta}$  is a linear layer and  $\| \boldsymbol{x} \|$  is the  $L_2$  norm. Preliminary analysis showed  $\lambda = 0.2$  to work well across datasets, so we use it for all our experiments, with higher values  $\lambda \geq 0.5$  degrading performance.
We populate the memory bank with all batch entries of ImageNet-1k / -22k at each step, using the representations from the target network. The memory bank is spread across 128 Cloud TPU v3 workers with 1200 entries on each TPU for ImageNet-1k (256 TPUs with 600 entries for 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  $x_{1}$  and  $x_{2}$ . The augmentations are constructed as compositions of the following operations, each applied with a given probability:
1. random cropping: a random patch of the image is selected, whose area is uniformly sampled in  $[0.08 \cdot \mathcal{A}, \mathcal{A}]$ , where  $\mathcal{A}$  is the area of the original image, and whose aspect ratio is logarithmically sampled in  $[3/4, 4/3]$ . The patch is then resized to  $224 \times 224$  pixels using bicubic interpolation;
2. horizontal flipping;
3. color jittering: the brightness, contrast, saturation and hue are shifted by a uniformly distributed offset;
4. color dropping: the RGB image is replaced by its grey-scale values;
5. gaussian blurring with a  $23 \times 23$  square kernel and a standard deviation uniformly sampled from [0.1, 2.0];
6. solarization: a point-wise color transformation  $x \mapsto x \cdot \mathbb{1}_{x < 0.5} + (1 - x) \cdot \mathbb{1}_{x \geq 0.5}$  with pixels  $x$  in [0, 1].
The augmented images  $x_{1}$  and  $x_{2}$  result from augmentations sampled from distributions  $\mathcal{T}_1$  and  $\mathcal{T}_2$  respectively. These distributions apply the primitives described above with different probabilities and different magnitudes. Table 8 specifies these parameters for the BYOL framework [28], which we adopt without modification.
Table 8: Pretraining augmentations. Parameters used to generate different views of a single image for contrastive pretraining.
Parameter	T1	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 max	0.4
Saturation adjustment max	0.2
Hue adjustment max	0.1
Gaussian blurring probability	1.0	0.1
Solarization probability	0.0	0.2
C.2 Optimization
We pretrain the model for 300 epochs on ImageNet-1k or 100 epochs on ImageNet-22k using AdamW [47] with a batch size of 4096, split across 128 Cloud TPU v3 workers for ImageNet-1k and 256 Cloud TPU v3 workers for ImageNet-22k. Training a ViT-B / ViT-L for 300 epochs on ImageNet-1k takes roughly 21 hours / 53 hours, while 100 epochs on ImageNet-22k takes approximately 60 hours / 128 hours. We update the online parameters  $\theta$  with a cosine learning rate schedule with a base learning rate of 0.001, weight decay of 0.1 and gradient clipping with a maximum norm of 1. We update the target parameters  $\xi$  as an exponential moving average of the online parameters with a decay rate of 0.99.
Following [16] the projections and predictions in Equation 4 are normalized and rescaled such that their norm is equal to  $1 / \sqrt{\tau}$  where the contrastive loss temperature  $\tau$  is equal to 0.1. When using additional supervision we set the supervised loss weight  $\alpha$  to 0.25 for the supervised ViT-B trained on ImageNet-22k and  $\alpha = 0.05$  for all other experiments.
D Supplementary analysis
D.1 Data efficiency
In Table 2 we compared Hummingbird with several leading representation learning techniques in the low-data regime. Here we provide the complete analysis from  $1/128$  to  $100\%$  of the data, as well as results for our ViT-L model trained on ImageNet-22k to show the scaling properties of Hummingbird. Note that there is a difference between the experiments run here and those found in Section 4.4 of the main paper; that section uses an UperNet [77] decoder and this section uses a linear decoder for all of the finetuned rows in each table.
For PASCAL VOC (Table 9), Hummingbird performs very well not only in the low-data regime but in the full-data regime, with the apples-to-apples comparison (ViT-B self-supervised on 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—with only nearest neighbors (no finetuning), Hummingbird++ trained on ImageNet-22k with a ViT-L backbone beats all of the other finetuned variants for every dataset fraction. Hummingbird++ using a ViT-B and ImageNet-1k predictably lies in-between the other two models for every dataset fraction.
For ADE20K (Table 10), the same general trends from above hold. Backbone and dataset scaling are once again beneficial as Hummingbird++ with ViT-L and ImageNet-22k training outperforms the other Hummingbird models, however this time the absolute performance relative to the finetuned competition in the high-data regime is less favorable since the end-to-end finetuned versions of other techniques start to outperform the nearest neighbors only ViT-L Hummingbird++ at  $1/16$  of the data.
Table 9: PASCAL VOC 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 at evaluation time (NN). All fine-tuning runs are averaged over five different seeds. The metric reported is mean IoU (higher numbers are better).  $\dagger$  denotes models trained on ImageNet-22k; all other models were trained on ImageNet-1k.
Fraction of dataset
Method	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 [30]	E2E FT	ViT-B	34.2	44.1	53.0	58.7	62.7	67.4	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
Hummingbird++	NN	ViT-B	52.4	57.3	61.5	64.6	66.2	67.9	70.5	73.2
Hummingbird++†	NN	ViT-L	61.8	65.3	68.0	70.7	71.4	73.2	75.3	77.2
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 at 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/32$  and  $1/1$  are sourced directly from [13], the rest are reproductions.  $\dagger$  denotes models trained on ImageNet-22k; all other models were trained on ImageNet-1k.
Fraction of dataset
Method	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	10.8	14.3	20.9	27.1	32.7	38.3	42.0	47.3
DINO [15]	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
Hummingbird	NN	ViT-B	11.7	15.1	17.3	20.0	22.3	24.9	27.9	29.6
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 Hummingbird models trained with different architectures (ViT-B vs ViT-L), datasets (ImageNet-1k vs ImageNet-22k), learning objectives (self-supervised or with additional supervision), and training lengths. Figure 7 plots the performance of these models when equipped with NN retrieval decoders (x-axis) and fully-finetuned UperNet decoders (y-axis). For both PASCAL VOC and ADE20K semantic segmentation, performance using one decoding scheme is highly predictive of the other (Pearson's  $\rho = 0.80$  for PASCAL VOC,  $\rho = 0.89$  for ADE20K). As such, even in cases where NN retrieval underperforms end-to-end finetuning, it can still be used as a powerful diagnostic tool. As illustrated in Section 4.3, evaluating with NN retrieval is much simpler and faster than with end-to-end finetuning, even when using a linear decoder. End-to-end finetuning often requires sweeping over optimization hyperparameters and averaging across multiple seeds, making it unsuitable for online evaluation, whereas NN retrieval is  $10\times$  less variable across runs and doesn't require any hyperparameter sweeps. As such NN retrieval can be used as an online evaluation that is highly predictive of performance obtained with more expensive finetuning protocols.
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 in Figure 8. Similar to what we observe for PASCAL VOC, we benefit from large memory banks at evaluation. Since the ADE20K
training set is roughly  $2 \times$  larger than that of PASCAL VOC, we also observe that sampling which features to store in the memory bank is more important than it is for PASCAL VOC (see Appendix A.1 on the details of the sampling procedure). Similarly, at training time, ADE20K benefits from larger pretraining memory banks than PASCAL VOC, with performance plateauing for memory banks larger than 200,000. Thus, we set the pretraining memory bank length to 153,600 in all our experiments (see Appendix B for details on contextual pretraining).
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-22k. Left: Since the retrieval-based supervised objective is only defined for memory banks of non-zero length, for the purpose of this ablation we replace it with a simple linear classifier when  $|\mathcal{M}_p| = 0$ . Right: For downsample=True, we store representations of all patches into the memory bank. If downsample=True, we sample  $|\mathcal{M}| / N$  patches per image ( $N$  is the length of the downstream training set), allowing for greater 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 to ViT-L is beneficial for both Hummingbird and Hummingbird++, where the self-supervised Hummingbird benefits slightly more from model scaling than its supervised counterpart. Note that this scaling study highlights the need for jointly scaling data and model size, as the best performing model overall is the one with ViT-L as an encoder trained on ImageNet-22k.
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.
Method	Encoder	Dataset	Finetuned accuracy (mIoU)
PASCAL ↑	ADE20K ↑
Hummingbird	ViT-B	IN1K	80.0	44.9
Hummingbird	ViT-L	IN1K	82.4	47.1
Hummingbird++	ViT-B	IN1K	81.2	44.9
Hummingbird++	ViT-L	IN1K	81.8	47.3
Hummingbird	ViT-B	IN22K	81.6	46.9
Hummingbird	ViT-L	IN22K	84.1	50.8
Hummingbird++	ViT-B	IN22K	82.1	48.2
Hummingbird++	ViT-L	IN22K	85.3	52.0
D.5 Importance of retrieval-based supervision for in-context scene understanding
We study the importance of retrieval-based supervised objective (see Section 3.4) on in-context scene understanding performance. We compare a model trained purely with the retrieval-based supervised objective ("Sup") with Hummingbird (purely self-supervised, i.e. "SSL") and Hummingbird++ (both self-supervised and retrieval-based supervised, i.e. "SSL + Sup"). Results shown in Table 12 indicate
the necessity of the self-supervised objective for creating representations that are general and transfer well to downstream tasks.
Table 12: Importance of retrieval-based supervision for 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.
Method	Objective	Dataset	Semantic segmentation	Depth pred.
PASCAL↑	ADE20K↑	NYUv2 ↓
Supervised retrieval	Sup	IN1K	56.9	22.9	.787
Hummingbird	SSL	IN1K	70.5	28.3	.718
Hummingbird++	SSL + Sup	IN1K	72.1	30.5	.738
Supervised retrieval	Sup	IN22K	69.3	30.3	.739
Hummingbird	SSL	IN22K	73.5	30.7	.706
Hummingbird++	SSL + Sup	IN22K	76.2	34.1	.695