SlowGuess/NeurIPS_2023_OCR huggingface dataset:
First, it is clear that the factorized approaches outperform both the baseline ViT and the Learned 2D embeddings from Pix2struct. The latter in particular struggles to generalize to higher resolution, likely because this requires an increasingly long tail of unseen  $(x,y)$  pairs. Factorized embeddings are best combined additively (as opposed to stacking or multiplying).
(a)
Figure 10: Factorized position embeddings improve generalization to new resolutions and aspect ratios. (a) Best (faded) and average accuracies (dark) across resolutions. (b) Accuracy normalized w.r.t. resolution 256.
3.5 Other aspects of NaViT's performance
Out of distribution generalization. We directly evaluate JFT-pretrained NaViT on downstream datasets, employing a label-map [26] from JFT-4B to ImageNet [4] and robustness-variants (ObjectNet [27] and ImageNet-A [28]). We compare the performance to a compute-matched ViT baseline.
Figure 11: Out of distribution evaluation of ViT and NaViT models that were matched for training compute. In addition to improved performance due to more images seen (see also Figure 1), NaViT performs much better on ImageNet-A that has many images with an extreme aspect ratio and important information outside the center crop (Appendix G). Same data as in Table 5.
Figure 11 shows that NaViT compares favorably both on ImageNet as well as datasets variants that were specifically designed to test out of distribution performance. It is interesting to note that NaViT performs much better on ImageNet-A, but ViT catches up on ObjectNet, even though both these datasets contain images that have extreme aspect ratios. We believe this is due to the aspect-preserving center crop that we apply to images for ImageNet-A and ObjectNet classification (same as in [17]), which is a useful prior for ObjectNet, but less so for ImageNet-A (see Appendix G for details). If no crop is applied to images and they're instead simply resized to the image resolution expected by the model (i.e., square for ViT, and aspect preserving resize to the same number of tokens for NaViT), then the observed difference is much larger (see Figure 21 in appendix).
Calibration. In addition to the in-depth accuracy analysis, we have also quantified the quality of the uncertainty computed by the model. In particular, for our ImageNet1K-finetuned models, we computed the expected calibration error [29] of the top prediction, as we vary the number of patches we assign per examples. We find that the calibration error remains very stable in the interval (0.045, 0.047) as we vary the number of patches per image in the range [128, 1024], without any post-hoc recalibration. We provide further details in the appendix Appendix F.
Inference trade-offs. Given the flexibility of the model, there are several viable approaches how one can maximize the aggregate accuracy under a given compute budget, which we quantify in terms of the latency measured on a Cloud TPUv3 chip. In an online inference setting, choosing an inference strategy translates to an algorithm how to allocate a fixed number of tokens per example. As we show in Fig. 1, NaViT offers much better trade-offs than ViT and it shows strong diminishing returns, s.t., even relatively few patches provide highly competitive results. In Appendix D we further study a cascading approach, which both provides Pareto optimal models and gives more precise trade-off opportunities.
Fairness signal annotation. We investigate annotating images with fairness-related signals, such as those pertaining to gender and ethnicity. Prior research has shown that metrics, such as group calibration, are susceptible to labeling inaccuracy, particularly for underrepresented groups. In addition, this problem persists even when accounting for label noise during training [30]. Thus, reducing the labeling error of fairness signals improves the reliability of bias mitigation and post-hoc auditing [31]. To explore whether NaViT can help in overcoming these challenges, we train annotators on FairFace [32] and CelebA [33] datasets as linear probes (i.e. using frozen features produced by NaViT or ViT), before comparing their accuracy.
First, NaViT provides representations of higher quality that improve the accuracy of fairness signal annotation, even when dealing with square images. Original images are of size  $448 \times 448$  in FairFace and  $178 \times 218$  in CelebA, and we resize them to area  $224^2$  while preserving aspect ratios in NaViT. Despite having the same sequence length, NaViT provides a higher prediction accuracy than ViT, as shown in Figure 12 (left). We verify statistical significance using the Wilcoxon signed-rank test test [34], which reveals that the improvement in NaViT is significant with  $p = 3 \times 10^{-4}$ .
Figure 12: Evaluating the accuracy of annotators trained on fairness-related signals using either NaViT-L/16 or ViT-L/16: Left: NaViT offers better representations that improve the accuracy of annotators. Right: Using native aspect ratios in NaViT results in a higher performance when compared to resizing images to squares.
Second, we apply inception-style cropping [35] with a fixed minimum area of  $50\%$ . This changes the aspect ratio but maintains native resolution. Similar to before, we resize all cropped images to area  $224 \times 224$ , either as square images or with native aspect ratios. Figure 12 (right) shows that native aspect ratios in NaViT improve accuracy, which is statistically significant at the  $95\%$  confidence level  $(p = 0.02)$ . Appendix H contains the full set of figures for other attributes in CelebA and Fairface.
3.6 Other downstream tasks
Semantic segmentation.
We finetune NaViT to semantic segmentation on ADE20k dataset [36], following the linear decoder protocol of Segmenter [37]. We use ViT-L/16 as baseline and compare its performance with that of NaViT-L/16. Both models are pretrained on JFT-4B [18] with comparable compute budget.
We experiment with different maximum resolution  $R_{max}$  at finetuning time: ViT takes in random square crops of resolution  $R_{max} \times R_{max}$  while NaViT's inputs are randomly resized (with preserved aspect ratio) so that the total number of pixels is  $R_{max}^2$ . This way, we have the same finetuning cost for the ViT and NaViT models. Note that following common practice, in order not to alter the ground truth segmentation maps, both models are evaluated at the native resolution [37, 38, 36]. This means that for ViT, the square predictions are resized from square back to native resolution. We observe in Figure 13 that NaViT outperforms ViT when transferred to semantic segmentation with the same maximum finetuning resolution  $R_{max}$ . Note that NaViT at  $R_{384}$  beats ViT at  $R_{512}$  while being twice as fast (see Appendix D). An advantage of NaViT over ViT is that it benefits from flexibility
of resolution during training and can ingest seamlessly square and non-square images.
Figure 13: NaViT transfers competitively to semantic segmentation. We transfer ViT-L/16 and NaViT-L/16 on ADE20k with fine-tuning at different resolutions. ViT consumes square images while NaViT preserves the aspect ratio of the original images, while maintaining the total number of pixels the same as ViT.
Object detection.
Native-resolution training may be especially beneficial for fine-grained tasks such as object detection, which require image understanding at many different spatial scales. We use compute matched NaViT and ViT models as backbones for OWL-ViT-L/14 object detectors [39], following the OWL-ViT protocol for training and evaluation. Results are presented in Table 1. The NaViT-based detector
performs significantly better on both common and unseen LVIS "rare" classes. These experiments used shorter pre-training than the original OWL-ViT and thus reach lower absolute performance; the relative difference nonetheless suggests that NaViT produces strong representations for fine-grained vision tasks.
Video Classification. Video processing with transformers is inherently challenging due to the heterogeneity of the underlying spatio-temporal signals. In practice, cumbersome protocols have to be
Table 1: NaViT improvements carry over to object detection.
ViT-L/14	NaViT-L/14
ImageNet zeroshot	68.3%	72.9%
LVIS AP	23.3%	28.3%
LVIS AP rare	17.2%	24.3%
established both for training and evaluation (e.g. spatial and temporal multi-cropping) [40]. NaViT alleviates some of these challenges by not only allowing training over different resolutions, but also over different temporal durations. We fine-tuned NaViT trained on JFT for Kinetics400 [41] classification by extracting three different spatio-temporal patches "tubelets". [42], extending the positional embedding to include the temporal dimension and initializing the embedding kernel using "central frame embedding" [40]. We posit that NaViT is an excellent starting point for multi-scale training due to the resolution diversity at training time. We observe that NaViT-L achieves competitive performance with ViViT-L  $(80.4\%)$  in approximately 6x less epochs, without multi-crop evaluation. Note that the Kinetics400 dataset used here contains less data than prior works [40].
Flexible Vision Transformers. FlexiViT [2] developed a novel kernel resizing approach which enables variable "resolution" via models which support multiple patch sizes. This is viewed as unifying multiple distinct models, and they study the relationship with distillation and neural architecture search. Pix2struct [3] supported variable aspect ratios with a novel positional embedding schema, and demonstrated significant efficiency and performance improvements for non-natural imagery such as chart and document understanding.
Multiscale Vision Transformers. Using feature maps at multiple spatial scales is a common approach to localization tasks such as segmentation and detection [43]. Many works developed Vision Transformers which do the same [44, 45], though some dispute the necessity for simple localization tasks [46]. NaViT does not build hierarchical representations using multiple scales; we believe combining the unique flexibility of our approach with the benefits of this modelling family is a promising avenue.
Accelerating training with mixed resolutions. Image modelling works considering resolution largely focus on accelerating pretraining with a fixed, low resolution [47, 48]. FixRes [11] is a popular technique whereby resolution is increased in a final stage of pretraining, and has been used in many subsequent works [19, 6]. PaLI [6], for example, successively increases the resolution of the vision backbone during its generative training. This approach's main downside is its irreversibility: compute cannot be scaled back by reducing resolution after tuning at the higher and more expensive resolution.
Multigrid training [49] is the most closely related work. The main idea is accelerate video modelling by processing large batches with "coarse" spatiotemporal resolution early in the training, followed by a "finer" resolution later. To this end, the authors apply a hierarchical grid sampling schedule coupled with appropriate learning rate scaling. In contrast, Patch n' Pack enables effortless incorporation of mixed resolutions without complex schedules or training pipelines.
Token dropping for improved efficiency. Research initially explored random token dropping [9, 10, 50]. Follow ups demonstrated benefits from considering structured [51] or "importance" based strategies [52, 53]. Better strategies will likely further boost performance, and Patch n' Pack sidesteps the fixed minibatch shape restriction which limited these works.
We have demonstrated that Patch n' Pack—the simple application of sequence packing to vision transformers—significantly improves training efficiency. The resultant NaViT models can be applied to many resolutions at inference time, and cheaply adapted to new tasks. We discuss limitations and impacts in Appendix A, but overall Patch n' Pack enables a wide variety of research previously hindered by the need for fixed batch shapes, including adaptive computation and new algorithms for improving training and inference efficiency.
[56] Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and tokenizer for neural text processing. In EMNLP, pages 66-71, November 2018.
[57] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[58] Foad Hamidi, Morgan Klaus Scheuerman, and Stacy M Branham. Gender recognition or gender reductionism? the social implications of embedded gender recognition systems. In Proceedings of the 2018 chi conference on human factors in computing systems, pages 1-13, 2018.
[59] Os Keyes. The misgendering machines: Trans/hci implications of automatic gender recognition. Proceedings of the ACM on human-computer interaction, 2(CSCW):1-22, 2018.
[60] Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. In NeurIPS, pages 23885-23899, 2021.
A Broader impacts & limitations
NaViT enables training of vision transformers on variable size inputs, which has a profound impact on advancing adaptive computation research. By training models to handle various input size, we can explore adaptive computation techniques that dynamically adjust the computational resources based on the specific requirements of a given input. This flexibility opens up new avenues for implementing ideas that aim at adjusting allocation of compute and improving efficiency in vision tasks per input. Furthermore, NaViT computational efficiency unlocks the potential for scaling up pre-training of vision models. With the ability to handle different resolutions, models can effectively tackle more complex and diverse visual data, allowing for the development of larger and more powerful vision models.
There is a wide range of applications that could benefit from a model capable of processing inputs of different resolutions, including OCR or document understanding with the use of vision models. Although we highlight the considerable advantages of employing NaViT and provide a comprehensive analysis of common computer vision tasks, we did not specifically investigate the benefits of NaViT in these particular applications. We consider this area as a priority for future research and follow-up work.
Implementation complexities The modelling changes shown in Figure 1 can complicate ViT; in particular, packing-aware changes to the loss function may be complex for more complicated usecases.
Computational overhead Done carelessly, significant extra cost can be introduced due to sequence packing (Discussed in Figure 4). NaViT can however be run with the same sequence length as a corresponding ViT model (e.g. ViT at resolution 384 with sequence length 576 vs NaViT at resolutions 64:384 with sequence length 576; the latter will fit on average  $\sim 2\times$  as many images).
Data augmentation aspects were not explored in this work; firstly, it is possible that token dropping and variable resolution can also act as data augmentations, which was not exhaustively studied here. Secondly, common data augmentations such as mixup and various random cropping schema assume square, constant resolution images, and analogues supporting variable aspect ratio and size were not developed here.
Scale; experiments were performed in a large-dataset regime, and future work should validate these results at smaller scale. This however necessitates research into data augmentation techniques as mentioned above.
B Training details
B.1 Classification pretraining
The experiments and ablations in the paper are with ViT-B/32, ViT-B/16, and ViT-L/16. We use a reciprocal square-root learning rate schedule, with linear warmup and cooldown, with a maximum value of  $8e - 4$ , and phases. We follow [18, 54] and use a higher weight decay of 3.0 on the head compared to the body's weight decay of 0.03 during upstream training to improve transfer to downstream tasks. For our experiments, we evaluated both NaViT and ViT models using configurations B/32, B/16, and L/16. Each ViT model was trained with varying compute budgets, with cooling down at different stages of training. We trained a corresponding NaViT model for each ViT size and computational budget, allowing us to perform "compute-matched" comparisons [25].
Table 2 presents the pretraining specifications for both ViT and NaViT models. During the pretraining phase, ViT models were trained using images of size  $224 \times 224$ . In contrast, NaViT models uniformly sampled a value, denoted as  $r$ , between 64 and 256 and resized the image to have a total of  $r^2$  pixels while preserving the aspect ratio. Although training NaViT models on native resolutions is possible, we empirically discovered that sampling a resolution provides greater control over maximizing the number of examples observed during pretaining within a fixed computational budget while maintaining performance across different resolutions in downstream tasks. Additionally, by controlling the resolution, we can ensure efficient packing by tuning the sequence length and limit padding to less than  $2\%$ .
Table 2: Pre-training details of ViT and NaViT with supervised classification.
Name	TPU Hours	Train Steps	Cooldown Steps	Sequence Length	Images Per Seq.	Batch Size	Training Images
ViT-B/32	7.9×101	1.0×105	1.0×104	49	1.0	≈4.0×103	4.0×108
1.9×102	2.5×105	5.0×104	49	1.0	≈4.0×103	1.0×109
3.9×102	5.0×105	1.0×105	49	1.0	≈4.0×103	2.0×109
7.9×102	1.0×106	1.0×105	49	1.0	≈4.0×103	4.0×109
ViT-B/16	2.6×102	1.0×105	1.0×104	196	1.0	≈4.0×103	4.0×108
6.6×102	2.5×105	5.0×104	196	1.0	≈4.0×103	1.0×109
1.3×103	5.0×105	1.0×105	196	1.0	≈4.0×103	2.0×109
2.6×103	1.0×106	1.0×105	196	1.0	≈4.0×103	4.0×109
ViT-L/16	5.4×102	1.0×105	1.0×104	196	1.0	≈4.0×103	4.0×108
1.3×103	2.5×105	5.0×104	196	1.0	≈4.0×103	1.0×109
2.7×103	5.0×105	1.0×105	196	1.0	≈4.0×103	2.0×109
5.4×103	1.0×106	1.0×105	196	1.0	≈4.0×103	4.0×109
NaViT-B/32	7.9×101	9.8×104	1.0×104	64	5.41	≈2.2×104	2.1×109
1.9×102	2.4×105	5.0×104	64	5.41	≈2.2×104	5.3×109
3.9×102	4.8×105	1.0×105	64	5.41	≈2.2×104	1.0×1010
7.9×102	9.7×105	1.0×105	64	5.41	≈2.2×104	2.1×1010
NaViT-B/16	2.6×102	9.3×104	1.0×104	256	4.87	≈1.9×104	1.8×109
6.6×102	2.3×105	5.0×104	256	4.88	≈1.9×104	4.6×109
1.3×103	4.6×105	1.0×105	256	4.88	≈1.9×104	9.2×109
2.6×103	9.2×105	1.0×105	256	4.88	≈1.9×104	1.8×1010
NaViT-L/16	5.4×102	9.7×104	1.0×104	256	4.88	≈1.9×104	1.9×109
1.3×103	2.4×105	5.0×104	256	4.87	≈1.9×104	4.8×109
2.7×103	4.8×105	1.0×105	256	4.87	≈1.9×104	9.6×109
5.4×103	9.6×105	1.0×105	256	4.88	≈1.9×104	1.9×1010
B.2 Contrastive pretraining
We use the 32000 token T5 [55] sentencepiece [56] tokenizer. By default, text sequences are truncated to a maximum length of 24. No token dropping is used for text. Models are trained under the same optimization regime as the classification models, but with a learning rate of  $3 \times 10^{-3}$ . Weight decay of  $1 \times 10^{-6}$  is applied consistently to all kernels in the model (no change for projection heads). By default, image side-lengths are sampled  $\sim \mathcal{U}(64, R_{\mathrm{max}})$ , and no other image augmentations are applied.
B.3 Packing algorithm
Packing of examples into sequences is done alongside batching. A simple greedy approach is used which adds examples to the first sequence with enough remaining space. Once no more examples can fit, sequences are filled with padding tokens, yielding the fixed sequence lengths needed for batched operations. Such simple packing algorithm can lead to a significant padding, depending on the distribution of length of inputs. There are several methods to address such limitations, like bin packing [7], which allows minimizing the padding. Here, in NaViT, since controlling the resolutions we sample, we can ensure efficient packing by tuning the sequence length and limit padding to less than  $2\%$ .
B.4 Pre-processing and augmentation
For both classification and contrastive pre-training, for the baseline models, we follow the original Vision Transformer paper [1], and follow-ups; models are trained on images with resolution  $224 \times 224$ , with inception crop followed by random horizontal flipping. For NaViT, we have no preprocessing besides resizing to the sampled resolution and token dropping.
Figure 14: Training on more images given a fixed computational budget, is a key factor contributing to the computational efficiency of NaViT compared to the standard ViT.
B.5 Sample efficiency with NaViT
Based on the experiments we ran, we believe a lot of the observed performance benefits are due to speeding up training (i.e., seeing more images in a given compute budget). This is evident in Figure 1, where relative performance gains diminish with longer training. We can also show this by plotting the number of images seen on the x-axis:
If compute cost was not a concern, for a given number of images seen, it is best to train on purely large resolution, with no token dropping. Asymptotically we therefore expect the token dropping and mixed resolutions to eventually hurt NaViT in unlimited compute regime. This assumes a large enough dataset that overfitting is not an issue. Eventually one may enter a regime where the token dropping and variable resolution is a useful data augmentation to prevent overfitting, which may alter results.
B.6 Sampling token dropping rates
Sampling with a beta distribution We use a parameterisation based on the mean  $d_{\mu}$  and standard deviation  $\sigma$ . We aim to sample dropout rate  $d \in [0.0, d_{\max}]$ , with some mean  $d_{\mu}$ .
Accordingly, we sample  $u \in [0,1] \sim \mathcal{B}(\alpha, \beta)$  and set drop rate  $d = u \times d_{\max}$ .  $\alpha$  and  $\beta$  are set such that the mean of  $u$  is  $u_{\mu} = \frac{d_{\mu}}{d_{\max}}$ . The maximum supported variance for a beta distribution of mean  $u_{\mu}$  is  $u_{\mu}(1 - u_{\mu})$ ; we pick by default a variance  $\sigma^2 = 0.3u_{\mu}(1 - u_{\mu})$ , which we found to work well in practice. The resultant distributions of token dropouts for different settings of  $d_{\mu}$  and  $d_{\max}$  are shown in Figure 15a.
(a) Beta-sampled token drop rates parameterised by the mean  $\mu$  and the max drop rate  $d_{\max}$
(b) Sampled resolution-dependent token drop rates
Sampling resolution-dependent dropping rates Given input data with sequence lengths ranging from  $s_{\min}$  to  $s_{\max}$ , we sample dropout rate  $d$  from a truncated normal distribution  $d \sim \mathcal{N}_{\mathrm{trunc}}(\mu, 0.02)$ , where samples more than two standard deviations away from  $\mu$  are rejected.
The mean of this distribution  $\mu$  is set according to the minimum and maximum token dropping rates  $d_{\min}$  and  $d_{\max}$ , and simply scales linearly with the sequence length  $s$  (such that  $s = s_{\min}$  has  $\mu = d_{\min}$  and  $s = s_{\max}$  has  $\mu = d_{\max}$ ).
Figure 15b shows example distributions of sampled drop rates given inputs with resolution  $R \sim \mathcal{U}(64,384)$ , and different values of  $d_{\min}$  and  $d_{\max}$ .
Figure 16: Decreasing the token dropping rate along training improves the ImageNet 10shot accuracy using the same pre-training resources.  $N$  is the total number of training examples seen with a fixed token dropping rate of  $\rho = \frac{1}{2}$ .
B.7 Scheduling token dropping rates
We experiment with a token dropping schedule which varies with total number of images seen. In particular, the rate applied for the  $n$ -th processed image during training is given by:
$$
\rho (n; \rho_ {\min }, \rho_ {\max }, \mu , \tau) = \rho_ {\min } + (\rho_ {\max } - \rho_ {\min }) \cdot \sigma \left(\frac {n - \mu}{\tau}\right), \tag {1}
$$