SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Mostafa Dehghani*, Basil Mustafa*, Josip Djolonga†, Jonathan Heek†, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, Avital Oliver, Piotr Padlewski, Alexey A. Gritsenko, Mario Lucic, Neil Houlsby
Google DeepMind {dehghani,basilm}@google.com
The ubiquitous and demonstrably suboptimal choice of resizing images to a fixed resolution before processing them with computer vision models has not yet been successfully challenged. However, models such as the Vision Transformer (ViT) offer flexible sequence-based modeling, and hence varying input sequence lengths. We take advantage of this with NaViT (Native Resolution ViT) which uses sequence packing during training to process inputs of arbitrary resolutions and aspect ratios. Alongside flexible model usage, we demonstrate improved training efficiency for large-scale supervised and contrastive image-text pretraining. NaViT can be efficiently transferred to standard tasks such as image and video classification, object detection, and semantic segmentation and leads to improved results on robustness and fairness benchmarks. At inference time, the input resolution flexibility can be used to smoothly navigate the test-time cost-performance trade-off. We believe that NaViT marks a departure from the standard, CNN-designed, input and modelling pipeline used by most computer vision models, and represents a promising direction for ViTs.
The simple, flexible and scalable nature of the Vision Transformer (ViT) [1] has rendered it an almost ubiquitous replacement to convolution based neural networks. Underpinning this model is a simple operation: splitting an image into patches, each of which is linearly projected to a token. Typically, input images are resized to a fixed square aspect ratio and then split into a fixed number of patches.
Recent works have explored alternatives to this paradigm: FlexiViT [2] supports multiple patch sizes within one architecture, enabling smooth variation of sequence length and thus compute cost. This is achieved via random sampling of a patch size at each training step and a resizing algorithm to allow the initial convolutional embedding to support multiple patch sizes. Pix2Struct [3] introduced an alternative patching approach which preserves the aspect ratio, which is particularly useful for tasks such as chart and document understanding.
We present an alternative, NaViT. Multiple patches from different images are packed in a single sequence—termed Patch n' Pack—which enables variable resolution while preserving the aspect
* Project lead. † Core contributor.
Figure 1: NaViT offers notable computational efficiency during pre-training (left) which carries over to downstream fine-tuning (middle). A single NaViT can be applied successfully to multiple resolutions (right), smoothly trading off performance and inference cost.
ratio (Figure 2). This is inspired by example packing in natural language processing, where multiple examples are packed into a single sequence to accommodate efficient training on variable length inputs.
Pooling Representations
Self-Attention
Data preprocessing
Figure 2: Example packing enables variable resolution images with preserved aspect ratio, reducing training time, improving performance and increasing flexibility. We show here the aspects of the data preprocessing and modelling that need to be modified to support Patch n' Pack. The position-wise operations in the network, such as MLPs, residual connections, and layer normalisations, do not need to be altered. Note that this diagram only showed the differences introduced by NaViT and other parts of the model architecture is identical to the Vision Transformer [1].
We demonstrate that: (i) Randomly sampling resolutions at training time significantly reduces training cost. (ii) NaViT results in high performance across a wide range of resolutions, enabling smooth cost-performance trade-off at inference time, and can be adapted with less cost to new tasks, (iii) Fixed batch shapes enabled by example packing lead to new research ideas, such as aspect-ratio preserving resolution-sampling, variable token dropping rates, and adaptive computation.
These observations have major practical implications. At a fixed computational budget NaViT consistently outperforms ViT. For instance, we match the performance of the top-performing ViT with  $4 \times$  less compute (Figure 1, left). We identify the substantial increase in the number of training examples processed within the allocated compute budget as the primary contributor to the improved performance over ViT (Appendix B.5) —example packing coupled with variable resolution inputs and variable token dropping enable NaViT-L/16 to process five times more images during training (Table 2). This improved efficiency extends to the fine-tuning process (Figure 1, middle). Furthermore, by exposing NaViT to multiple resolutions during both pre-training and fine-tuning, a single model demonstrates excellent performance when evaluated on various resolutions, significantly advantaging NaViT in terms of inference cost (Figure 1, right).
NaViT's training and adaptation efficiency, and flexible inference, presents a promising avenue for Vision Transformers. Patch n' Pack empowers computer vision systems to transcend limitations imposed by current data and modeling pipelines, enabling ideas that were previously restricted by the constraints of fixed batch shapes, unlocking new possibilities for innovation and advancement.
Figure 3: Height:width ratios of different datasets; most images are not squareish  $(>20\%$  deviation).
Deep neural networks are typically trained and run with batches of inputs. For efficient processing on the current hardware this implies fixed batch shapes which in turn imply fixed image sizes for computer vision applications. This coupled with architectural limitations historically associated with convolutional neural networks led to a practice of either resizing or padding images to a fixed size. Both of these have been shown to be flawed: the former harms performance and the latter is inefficient [3]. An analysis of aspect ratios in ImageNet [4], LVIS [5] and WebLI [6] as representative examples of classification, detection and web image datasets, respectively, shows that most images are typically not square (Figure 3).
In language modelling, it is common to bypass limitations of fixed sequence lengths via example packing: tokens from multiple distinct examples are combined in one sequence, which can significantly accelerate training of language models [7]. By treating images as sequences of patches (tokens), we show that Vision Transformers [1] can benefit from the same paradigm, which we call Patch n' Pack. Using this technique ViTs can be trained on images at their "native" resolution, and we name this approach NaViT.
2.1 Architectural changes
NaViT is built upon the original ViT, but in principle can use any ViT variant operating on a sequence of patches. To enable Patch n' Pack, we make the following architectural modifications.
Masked self attention and masked pooling. To prevent examples attending to each other, additional self-attention masks are introduced. Similarly, masked pooling on top of encoder aims to pool the token representations within each example, resulting in a single vector representation per example in the sequence. Figure 2 presents how the receptive filed of attention is controlled via masking.
Factorized & fractional positional embeddings. To handle arbitrary resolutions and aspect ratios, we revisit the position embeddings. Given square images of resolution  $R \times R$ , a vanilla ViT with patch size  $P$  learns 1-D positional embeddings of length  $(R / P)^2$  [1]. Linearly interpolating these embeddings is necessary to train or evaluate at higher resolution  $R$ .
Pix2struct [3] introduces learned 2D absolute positional embeddings, whereby positional embeddings of size  $\left[\max \mathrm{Len},\max \mathrm{Len}\right]$  are learned, and indexed with  $(x,y)$  coordinates of each patch. This enables variable aspect ratios, with resolutions of up to  $R = P\cdot \max \mathrm{Len}$ . However, every combination of  $(x,y)$  coordinates must be seen during training.
To support variable aspect ratios and readily extrapolate to unseen resolutions, we introduce factorized positional embeddings, where we decompose into separate embeddings  $\phi_x$  and  $\phi_y$  of  $x$  and  $y$  coordinates. These are then summed together (alternative combination strategies explored in Section 3.4). We consider two schema: absolute embeddings, where  $\phi(p): [0, \max\text{Len}] \to \mathbb{R}^D$  is a function of the absolute patch index, and fractional embeddings, where  $\phi(r): [0, 1] \to \mathbb{R}^D$  is a function of  $r = p/\text{side-length}$ , that is, the relative distance along the image. The latter provides positional embedding parameters independent of the image size, but partially obfuscates the original aspect ratio, which is then only implicit in the number of patches. We consider simple learned embeddings  $\phi$ , sinusoidal embeddings, and the learned Fourier positional embedding used by NeRF [8].
2.2 Training changes
Patch n' pack enables new techniques to be used during training of NaViT.
Continuous Token dropping. Token dropping (random omission of input patches during training) [9, 10] has been developed to accelerate training. However, typically the same proportion of tokens are dropped from all examples; packing enables continuous token dropping, whereby the token dropping rate can be varied per-image. This enables the benefits of faster throughput enabled by dropping while still seeing some complete images, reducing the train/inference discrepancy. Further, with packing, the
drop-distribution can vary throughout training, following some pre-defined schedule. In Section 3.3, we explore different schedules and the benefits of flexible token dropping.
Resolution sampling. NaViT can be trained using the original resolution of each image. Alternatively, the total number of pixels can be resampled while preserving aspect ratio. In vanilla ViT, there is a tension between greater throughput (training on smaller images), and greater performance (training on larger images, to enable high-resolution at evaluation time). Oftentimes, models are pre-trained at a smaller resolution and finetuned at a higher one [11]. NaViT is much more flexible; it allows mixed-resolution training by sampling from a distribution of image sizes, while retaining each images' original aspect ratio. This allows both higher throughput and exposure to large images, yielding substantial improved performance over equivalent ViTs (in terms of models size and training duration). Section 3.2 explores different sampling strategies, and variable resolution training for pre-training and finetuning.
2.3 Efficiency of NaViT
Here we discuss some implications of Patch n' Pack on the computational efficiency of NaViT.
Self attention cost. The  $\mathcal{O}(n^2)$  cost of attention is a natural concern when packing multiple images into longer sequences. Though many works aim to remove this quadratic scaling [12, 13], we demonstrate here that as the transformer hidden dimension is scaled, the attention becomes an increasingly smaller proportion of the overall cost, which encompasses the computation cost of the MLP as well. Figure 4 illustrates this trend, indicating a corresponding reduction in the overhead associated with packing examples. In addition to speed considerations, the memory cost of self-attention can pose a challenge for extremely long sequences. However, this challenge can be also addressed by employing memory-efficient methods [14, 15].
Figure 4: Overhead from extra attention due to packing, assuming 256 tokens per image; it diminishes with model scale.
Packing, and sequence-level padding. The final sequence lengths containing multiple examples must be fixed. We use a greedy packing approach discussed in Appendix B.3; there typically is no perfect combination of examples exactly adding up to the fixed length and padding tokens have to be used. One could for example dynamically choose the resolution or token dropping rate of the final example in a sequence to exactly fit the remaining tokens; however, we find typically less  $2\%$  of tokens are padding tokens, and thus the simple approach is sufficient.
Padding examples and the contrastive loss. Per-token losses are straightforward to implement with packed sequences. However, many computer vision models are trained with example-level losses, typically applied to a pooled representation. First, this requires modifications to the typical pooling heads to account for packing. Second, multiple pooled representations must be extracted from each sequence. Fixed batch shapes require an assumption that, from a batch of  $B$  sequences, we extract at most  $B \times E_{\mathrm{max}}$  pooled representations (i.e.  $E_{\mathrm{max}}$  examples per sequence). If a sequence contains more than  $E_{\mathrm{max}}$  images, the extra images will be dropped, wasting computation of the model's encoder. If a sequence has less than  $E_{\mathrm{max}}$  examples, then the loss will process lots of fake padding representations.
The latter is an issue for contrastive learning, where loss computation scales in time and memory  $\sim \mathcal{O}(n^2)$ . To avoid this, we used the chunked contrastive loss [16], which circumvents the need to gather all data points for the softmax by performing computations on local device subsets and efficiently accumulating the necessary statistics for global softmax normalization. This enable high values of  $E_{\mathrm{max}}$  (and thus efficient use of the model encoder), without being bottlenecked by the loss.
The base architecture we use for NaViT follows vanilla ViT [1], with the changes to enable packing, described in Section 2.1. In addition, we include small ViT improvements from previous works: query-key normalization and the omission of biases [17], and attention pooling [18].
We pre-train NaViT in two setups: classification training on JFT-4B [18] and contrastive language-image training [19] on WebLI [6]. Typically, for JFT, inception crop is applied pre-training [20, 1], and in both cases, images are resized to a square (distorting aspect ratio). Unless otherwise specified,
all NaViT models are pre-trained without these operations, and preserve aspect ratio. NaViT is implemented in JAX [21] using the FLAX library [22] and built within Scenic [23].
Classification pretraining. We pre-train NaViT with supervised classification objective, using a sigmoid cross-entropy loss, following the setup of [17] on JFT-4B [18]. Visual representations are evaluated following the linear evaluation protocol used for ViT [1], where 10 examples per class are used to train a linear classifier on top of frozen representations.
Contrastive pre-training. Alongside the image model, we train a text encoder with the same architectural modifications using the contrastive image-text loss [19, 24] (details in Appendix B.2). Packing also provides efficiency improvements on the text-tower, as text sequences do not need to be padded to a fixed lengths, which is the normal setup. The contrastive models are evaluated on zero-shot ImageNet classification and COCO image-text retrieval.
Note that throughout all the pre-training experiments and the transfer to downstream experiments, both the NaViT and baseline models are trained with compute-matched setup. This implies that during pretraining, the same amount of TPU time is used and for downstream tasks, the models are evaluated at the same effective resolution (i.e., the same number of tokens). To be precise, all downstream experiments utilized the top-rightmost points in Figure 1 (ViT-L/16 and NaViT-L/16).
3.1 Improved training efficiency and performance
Figure 1 illustrates the JFT pretraining performance of different NaViT models compared to compute-matched ViT baselines [25]. The experimental setup details are provided in Appendix B.1. NaViT consistently surpasses ViT in performance while using the same computational budget across different compute and parameter scales; for example, the performance of the top-performing ViT can be matched by a NaViT with four times less compute. Conversely, the computationally lightest NaViT in Figure 1 is five times more cost-effective than its equivalent ViT counterpart.
The NaViT models benefit from preserved aspect ratios and the ability to evaluate over many resolutions, but the chief contributor here is the significant increase in the number of training examples processed by NaViT within the allocated compute budget. This is achieved through the combination of sampling multiple variable-resolution examples and token dropping, leading to variable size images that are efficiently packed into a similar sequence length as the original model. We ablate these factors below.
3.2 Benefits of variable resolution
Here, we deep-dive the benefits of mixed-resolution training. Since we preserve the native aspect ratio, when we refer to "resolution" for NaViT, we mean "effective resolution". That is, images with the same area as a square image with a given resolution. For example, for NaViT a resolution of "128" has the same area of a square  $128 \times 128$  image, but could be  $64 \times 256$ , or  $170 \times 96$ , etc., and thus has the same inference cost as regular ViT on  $128 \times 128$  images.
Variable-resolution pre-training. Lower resolution images require fewer FLOPs to process and hence small resolutions (like 224) are used with fixed-resolution training. With fixed-resolution training, there is a trade-off between throughput and ability to process details and high-resolution images. With NaViT we can mix lower resolution images with large ones to get the best of both worlds.
Figure 5 shows a comparison between two NaViT variants trained at several different resolutions. Here, all trained for the same number of FLOPs. (1) Native aspect ratio, but fixed resolution  $R = R_{\max}$  for different chosen values of  $R_{\max}$ . (2) Variable resolution, where the resolution is distributed as  $R \sim \mathcal{U}(64, R_{\max})$ . Variable resolution models outperform models trained at only that resolution. Even in the best case for fixed resolution, where the train and evaluation resolutions are identical, variable resolution matches or outperforms fixed.
Variable-resolution finetuning. Prior works increase resolution late in pre-training or during finetuning, producing higher quality but more expensive models [1, 11]. We finetune NaViT and ViT models at different fixed resolutions, and additionally NaViT at variable resolutions. Figure 6 shows the results of fine-tuning pretrained ViT and NaViT on ImageNet-1k dataset. Performance gains during pretraining transfer well at all resolutions, but two phenomena are particularly interesting: First, NaViT finetuned with variable resolutions ("NaViT 64:512") is as good as a NaViT finetuned at a single resolution (and much better than single-resolution ViT), removing the need to pick a single downstream finetuning resolution. Second, NaViT finetuned at low resolution (64) still obtains good performance when
Figure 5: At fixed computational cost, sampling lower resolutions increases throughput, improves performance and enables better use of models at varied resolutions. NaViT-B/16 models trained with variable vs. fixed resolutions demonstrate the benefit of mixed resolution.
Figure 6: Variable-resolution finetuning, JFT B/16 models finetuned on ImageNet at various resolutions. Overall NaViT in all settings (blue, red), outperforms ViT (orange) Left: A single NaViT finetuned with variable resolutions (red) is as good as models tuned on only one resolution (blue). Right: Mixedresolution pretraining performs well at high resolution when finetuning at low resolution (left-hand end of blue curve).
evaluated at higher resolutions (Figure 6, right), enabling cheaper adaptation. This ability to perform cheap adaptation of a flexible pre-trained model corroborates findings in [2].
Resolution sampling strategies. Packing examples enables diverse resolution sampling strategies. We first consider whether to sample the target side length (average height/width,  $R$ ), or the target area (i.e. sequence length  $\propto R^2$ ). Sampling the side length from a uniform distribution biases towards lower sequence lengths, whereas sampling the area from a uniform distribution biases towards higher side lengths.
For each image, we sample  $u \sim \mathcal{D}$ , where  $\mathcal{D}$  is a distribution with support  $[-1,1]$ . We rescale  $u$  linearly to  $[64,384]$  for sampling side-lengths or  $[64^2,384^2]$  for sampling areas. We consider four distributions  $\mathcal{D}$ : uniform  $u \sim \mathcal{U}(-1,1)$ , truncated (to  $[-1,1]$ ) standard Normal  $u \sim \mathcal{N}_t(0,1)$ , and then two other Normals which bias towards lower resolution  $u \sim \mathcal{N}_t(-0.5,1)$  and higher resolution  $u \sim \mathcal{N}_t(0.5,1)$ . The results are shown in Figure 7. Here, the best resolution resampling strategy consistently performs over the default resolution. It is consistently better to
sample side-lengths (orange) directly as opposed to area (blue), and the best distribution is the truncated normal biasing towards lower values; both of these increase throughput by preferentially sampling smaller sequences.
Figure 7: Sampling side lengths directly with a bias towards lower resolutions gives overall best performance at a fixed computational budget.
3.3 Benefits of variable token dropping
Token dropping strategies.
We experimented with continuously sampled token dropping rates, and with resolution-dependent token dropping rates; both are explained in Appendix B.6.
Fig. 9a compares variable drop rates sampled from a Beta distribution to a constant drop rate, demonstrating consistent improvements from the former. Fig. 9b shows the use of a resolution dependent token dropping rate for models trained with  $R \sim \mathcal{U}(64,384)$  and dropping rates scaled between  $[0.5 - \delta, 0.5 + \delta] \propto R$ , which further improves over the beta distribution.
Figure 8: Time-varying token dropping rates improves performance and are easily done with Patch n' Pack.
(a) Constant vs. Beta-distributed token dropping rates.
(b) Resolution dependent token dropping sampled  $\in 0.5\pm \delta$
Figure 9: Continuous token dropping strategies enabled by sequence packing improves performance
Scheduled token dropping rates. Packing also enables easy
variation the token dropping rate during training. By changing the token dropping rate we can better tune the trade-off between number of images seen and information used per image, to maximize the final accuracy while keeping the total training cost constant. We varied the token dropping rate as a function of the number of images seen (details of the schedule in Appendix B.7). Fig. 8 demonstrates that further improvements are possible by reducing the token dropping rate during JFT pretraining of NaViT-B/16.
3.4 Positional embeddings
We evaluate our factorized embeddings introduced in Section 2.1, and their design choices. We are interested in both absolute performance, and extrapolation to resolutions outside the training regime. To test this, we train NaViT-B/16 models for  $200\mathrm{k}$  steps on JFT, with resolutions  $R \sim \mathcal{U}(160,352)$ . We evaluate performance at a range of resolutions, without modification of the embedding variables. We compare to a ViT-B/16 trained at fixed resolution 256 for the same amount of images seen, evaluated at new resolutions using standard interpolation of positional embeddings.
Figure 10a illustrates the disparity among different positional embedding methods when the model is evaluated under the "in-distribution" setup, where the resolutions of training and testing data are within a comparable range. Conversely, Figure 10b shows the model performance across out-of-distribution resolutions, revealing notable distinctions between the approaches.