Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Authors: Mostafa Dehghani
Basil Mustafa
Josip Djolonga
Jonathan Heek
Matthias Minderer
Mathilde Caron
Andreas Peter Steiner
Joan Puigcerver
Robert Geirhos
Ibrahim Alabdulmohsin
Avital Oliver
Piotr Padlewski
Alexey A. Gritsenko
Mario Lucic
Neil Houlsby
Venue: NeurIPS 2023 poster (2023)
Keywords: Vision Transformer
variable aspect ratio
flexible inference
efficient training
TL;DR: NaViT challenges the conventional practice of resizing images to a fixed resolution  and utilizes sequence packing to handle inputs of arbitrary resolutions and aspect ratios, offering flexible model usage and improved training efficiency.

Abstract: 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 NaViTmarks a departure from the standard, CNN-designed, input and modelling pipeline used by most computer vision models, and represents a promising direction for ViTs.

Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Mostafa Dehghani', Basil Mustafa', Josip Djolonga', Jonathan Heeki, 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
1mm11mmaaoooo {dehghani, basi 1m) @googlooooo...oooooooooooooooeoe com
Abstract
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- 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. Vi can be effi- ciently transferred to standard tasks such 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 sost-performance trade-off. We believe that ViT marks departure from the standard, CNN-designed, input and modelling pipeline used by most computer vision models, and represents promising direction for ViTs.
Introduction
The simple, flexible and scalable nature of the Vision Transformer (ViT) [ has rendered an almost ubiquitous replacement to convolution based neural networks. Underpinning this model is simple operation: splitting an image into patches, each of which is linearly projected to a token. Typically, input images resized to fixed square aspect ratio and then split into fixed number of patches.
Recent works have explored alternatives to this paradigm: Flexi [2] supports multiple patch sizes within one architecture, enabling smooth variation of sequence length and thus compute cost. This achieved via random sampling of patch size at each training step and a resizing algorithm to allow the initial convolutional embedding support multiple patch sizes. Pix2Struct [3] introduced an alternative patching approach which preserves the aspect ratio, which is particularly useful for tasks such chart and document understanding.
We present an alternative, NaVIT. Multiple patches from different images are packed in a single sequence-termed Patch n' Pack- enables variable resolution while preserving the aspect
Project lead. (Core contributor.
37th Conference Neural Information Processing Systems (NeurIPS 2023)

---

Pre-training lincar
Finctuned acc pretrain
Finetuned acc inference cost
faster
88%
90%
80%
75%
86%
80%
Model
Accuracy
Ap
VIT
70%
84%
NaVIT
65%
4.9% fuster
82%
Version
Br32
60%
80%
B/16
L/16
V1T1 10
10
10a
10
10'
10"
Pro-training TPU chips hours
Pre-training TPU chips hours
Inferemoc cost (ms/image)
Figure VaT offers notable computational efficiency during pre-training (left) which carries down- stream ane-tuning (middle). A single NaVIT can applied successfully to multiple resolutions (right), smoothly trading off performance and inference cost.
This is inspired by example packing in natural language processing, where multiple examples are packed into single sequence to accommodate efficient training on variable length inputs.
ratio (Figure 2.
Figure Example packing enables variable resolution images with preserved as- pect ratio, reducing train- ing time, improving perfor- mance and increasing flexi- bility. We show here the as- pects of the data preprocess- ing and modelling that need be modified to support Patch Pack. The position- wise operations in the net- work, such as MLPs, residual connections, and layer malisations, do not need to be altered. Note that this dia- gram only showed the differ- ences introduced by NaVIT and other parts of the model architecture identical to the Vision Transformer [I]
Pooling Representations
Pad examples are masked out from loss computation
are from pooling
Self-Attention
token's receptive field
eecample
Pad tokens are masked out from attention
query
key
Imaaee
Image
Image
Padi
Padied
Data preprocessing
Image1
Imaget
Imaget
Imaget
Imsqged
Imsged
Imaqee
Imaget
Imagee
Inputs
Patchify
Token drop
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-performanee 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-sampae samplinn,. 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 4x less compute [ left). We identify the substantial increase examples processed within the allocated compute budget the primary contributor to the improved performance over ViT (Appendix B.5) example packing coupled with variable resolution inputs and variable token dropping enable VaViT-L/16 to process five times more images during training (Table This improved efficiency extends to the fine-tuning process (Figure middle). Furthermore, by exposing VaT to multiple resolutions during both pre-training and fine-tuning, single model demonstrates excellent performance when evaluated on various resolutions, significantly advantaging Na in terms of inference cost (Figure [ right).
number of training
ViT'''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.

---

ImageNet (85.%% noa-square)
LVIS (92.2% uon-square)
WebLI (57.3% non-square)
Figure 3: Height:idth ra- tios of different datasets; most images are not square- ish (> 20% deviation).
Density
11-22 2m
A1 A0
50 2k
K- 23v
A1 1V Aspect ratio
to 24
A- 2
A w 2A
Method
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 vi- sion applications. This coupled with architectural limitations historically associated with convolutional neural networks led to practice of either resizing padding images to fixed size. Both of these have been shown to flawed: the former harms performance and the latter is inefficient [ An analysis of aspect ratios in ImageNet 4 LVIS [S] and WebLI [6] as representative examples of classification, de- tection and web image datasets, respectively, shows that most images are typivally not square (Figure
In language modelling, it is common to bypass limitations of fixed sequence lengths via example packing: tokens from multiple examins are combined in one sequence, which can significantly accelerate training of language models 77.. By treating images as sequences of patches (tokens), we show that Vision Transformers [ 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 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 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 Rt a vanilla ViT with patch size P llarns I-D positional embeddings of length (² [ Linearly interpolating these embeddings necessary train evaluate at higher resolution R.
Pix2struct 3 introduces learned 2D absolute positional embeddings, whereby positional embeddings of size mmaxLen,aaaaaaaamaa variable aspect ratios, with resolutions of up to R= P. maxLen. However, every combination of (z.y) coordinates must be seen during training.
are learned, and indexed with (I.g) coordinates of each patch. This enables
To support variable aspect ratios and readily extrapolate to unseen resolutions, we introduce factorized positional embeddings, where we decompose into separate embeddings &z and d, of and coordinates. These are then summed together (alternative combination strategies explored in R is
Section 3.4). We consider two schema: absolute embeddings, where Φ(p): [0, function the absolute patch index, and fractional embeddings, where ( R is function of r p/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 Iearned embeddings d sinusoidal embeddings, and the learned Fourier positional embedding used by NeRF g
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)
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
3

---

coop-distribttion 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. Na 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 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 smaller resolution and finetuned at a higher one [1]] ViV much more flexible; 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- vre-training and finetuning.
2.3 Efficiency of Na ViT
Here we discuss some implications of Patch n' Pack on the computational efficiency of Na
Self attention cost. The (n²) ) cost of attention is natu- ral concern when packing multiple images into longer se- quences. Though many works aim to remove this quadratic scaling 12 [3), we demonstrate here that as the transformer hidden dimension is scaled, the attention becomes in- creasingly smaller proportion of the the overall cost, which encompasses the computation cost of the MLP as well. Figure this trend, indicating a corresponding reduction in the overhead associated with packing exam- ples. In addition to speed considerations, the memory cost of self-attention can pose challenge for extremely long sequences. However, this challenge can be also addressed by employing memory-efficient methods [5).
Attention overheud
50%
packed
40%
pucked
pucked
30%
packed
20%
10%
0%%
1024
2048
4096
6144
Model dimension
Figure 4: Overhead from extra attention due to packing, assuming 256 tokens per image; it diminishes with model scale.
Packing, and sequencc-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 be used. One could for example dynamically choose the resolution or token dropping rate of the final example in a sequence 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 requires an assumption that, from a batch of B sequences, we extract at most x Emax pooled representations (i.e. Emax examples per sequence). If sequence contains more than Emax images, the extra images will be dropped, wasting computation of the model's ''s encoder. If a sequence has less than Fmax 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 On²). To avoid this, we used the chunked contrastive loss [16), which cireumvents 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 Emax (and thus efficient use of the model encoder), without being bottlenecked by the loss.
3 Experiments
The base architecture we use for NaVIT follows vanilla VIT [.. with the changes to enable packing, described in Section 2. 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 JFT-4B [18] and contrastive language- image training [19] on WebL.I (6. Typically, for JFT, inception crop is applied pre-training 20 1),
and in both cases, images resized to 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 [2]] 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 [. where 10 examples per class are used to train linear classifier on top of frozen representations.
Contrastive pre-training. Alongside the image model, we train a text encoder with the same architec- tural modifications using the contrastive image-text loss [9 24 (details in Appendix [8.2]. Packing also provides efficiency improvements on the text-tower, 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- pre-training experiments and the transfer to downstream experiments, both the VIT and baseline models are trained with cenpute-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 T (ViT-L/16 and NaViT-L/11).
T
3.1 Improved training efficiency and performance
Figure Tiatreeillsttaate the JFT pretraining performance of different NaVIT models compared to compute- matched ViT baselines 225.. The experimental setup details are provided in Appendix [..] ViT consistently surpasses VIT in performance while using the same computational budget across different compute and parameter scales; for example, the performance of the top- top-performing ViT can be matched by a VaV with four times less compute. Conversely, the computationally lightest Va in Figure is five times more cost-effective than its equivalent ViT counterpart.
The VTT 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 Na within the allocated compute budget. This is achieved through the combination of sampling multiple variable-resolutin examples and token dropping. leading to variable size images that are efficiently packed into similar sequence length 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 Na we mean "effective resolution", That is, images with the same area as square image with a given resolution. For example, for ViT a resolution of "228" has the same area of a square 128 x 128 image, but could be 64 256, or 170 x 96, etc., and thus has the same inference cost as regular VIT on 128 x 128 images.
Variaile-resolution pre-training. Lower resolution images require fewer FLOPs to process and hence small resolutions (like 224) are used with exed-resolution training. With fixed-resolution training, there is trade-off between throughput and ability to process details and high-resolution images. With ViT we can mix lower resolution images with large ones to get the best both worlds.
Figure showss 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 Rmax for different chosen values of Rmax. (2) Variable resolution, where the resolution is distributed as R (-2((64,,mma))) 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 finetun- ing. producing higher quality but more expensive models [ We finetune NaViT and ViT models at different fixed resolutions, and additionally Na at variable resolutions. Figure the results of fine-tuning pretrained ViT and Na on ImageNet- kk dataset. Performance gains during pretraining transfer well at resolutions, but two phenomena are particularly interesting: First, VTT finetuned with variable resolutions ("N 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 single downstream finetuning resolution. Second, NaVIT finetuned at low resolution (64) still obtains good performance when
5

---

79%
Best ImageNet 10shot
R0S
Figure At fixed computa- tional cost, sampling lower res- olutions increases throughput, improves performance and en- ables better use of models varied resolutions. NaViT- B/16 models trained with vari- able vs. fixed resolutions demonstrate the benefit of mixed resolution.
75%
60%
70%
Top Lcccurcy
40
766
513
2/6 164
77%
30%
80%
75
60R
---
VanooNee mns -tiph Mani
70%
Rem
401
128
361
512
164
Max trin resolutions Recs
Roux
Rayi
Ryw
Eval Rn
across R 512)
Figure 6: duriablc-resolttin finetuning, JFT B/16 models finetuned ImageNet various resolutions. Overall NaViT in all settings (blue, red), outper- forms ViT (orange) Left: A single NaViT finetuned with variable resolutions (red) is as good as models tuned on only one resolution (blue). Right: Mixed- resolution pretraining performs well at high resolu- tion when finetuning at low resolution (left-hand end blue curve).
InsageNet Top-1 Acsuraey
85%
808
75%
70%
NaVIT
N&VIT per
60%
VIT (one
VIT fone NaVIT-64512
NaVIT-64-512
558
256
Si2
256 Fine tune resolution Rp.
Fine tune resalation Rg.
evaluated at higher resolutions (Figure right), enabling cheaper adaptation. This ability to perform cheap adaptation of a flexible pre-trained model corroborates findings 2
Resolution sampling strategies. Packing examples enables diverse resolution sampling strategies. We first consider whether to sample the target side length (average height/wich R),, or the target area (i.e. sequence length R²). Sampling the side length from a uniform distribution biases towards lower sequence lengths, whereas sampling the area from uniform distribution biases towards higher side lengths.
85.25
nontial
53.45
distebuttoon
normal
#=
02:45
63.55
uniform
For each image, we sample u~ D, where D is distribution with support [-1,1]. We rescale u linearly to [64,384] for sampling side-lengths or [64" 3389] for sampling areas. We consider four distributions uniform u ~ (,,,,, truncated (to [-1,1) standard Normal ~ (1 1), and then two other Normals which bias towards lower resolution (=..55, and higher resolution (0.5,, 1). The results shown in Figure 7 Here, the best resolution resampling strategy consistently per- forms over the default resolution. It is consistently better to sample side-lengths (orange) directly as opposed to area (blue), and the best distribution the truncated normal biasing towards lower values; both of these increase throughput by preferentially sampling smaller sequences.
55%
58%
60%
62%
65%
68%
ImageNet zero Accunicy
Sumple side length
Notive ze5
Sumple 1823
Sampling side lengths di- rectly with a bias towards lower reso- lutions gives overall best performance a fixed computational budget.
Figure 7:
3.3 Benefits of variable token dropping
Token dropping strategies.
Drop schedule
We experimented with continuously sampled token dropping rates, and with resolution -eoouuiinnddeeedentt token dropping rates; both are explained in Appendix B.66 [366
Increasing
72.2%
Docrcasing (fest)
76.3%
Decreasing (slow)
76.5%
Fig. 9a compares variable drop rates sampled from Beta distribution constant drop rate, demonstrating consistent improvements from the former. Fig. 19b] shows the use of resolution dependent token dropping rate for models trained with R 2/664 384) and dropping rates scaled between [0.5 8,0.5 + 6] x R. which further improves over the beta distribution.
70.0%
80.0%
Top-1 accuracy
Figure 8: Time-varying token dropping rates improves performance and are ily done with Patch n' Pack.

---

79.0%
Resolution besoluzion-depeniaee
Top-1 Accurncy
Top-1 accuracy
Figure 9; Continuous to- ken dropping strategies abled by sequence packing improves performance
78.5%
78%
drop fixed for every image
betn: B image
76%
0.25
0.50
0.75
fixed
heta
8-0.25 6=04
Mean token drop rate
Token drop strategy
Constant vs. Beta-distributee to- ken dropping rates.
(b) Resolution dependent token drop- ping sampled 60.5+6
Scheduled token dropping rates. Packing also enables easy
variation the token dropping rate during training. By changing the token dropping rate 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 function of the number of images seen (details of the schedule in Appendix (B.)). F @ddeonsttatess 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 Na ViT-B/16 models for 200k steps on JFT, with resolutions R~~2(1600,32)) We evaluate performance at 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 T0a (lusttats 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 within a comparable range. Conversely, Figure shows the model performance across out-of of-distributtio resolutions, revealing notable distinctions between the approaches.
First, it clear that the factorized approaches outperform both the baseline ViT and the Learned 2D embeddings from Pix2struct. The latter in particular struggles to generalize higher resolution, likely because this requires an increasingly long of unseen (æ.y) pairs. Factorized embeddings are best combined additively (as opposed to stacking or multiplying).
Leumed ID (VIT)
53.8%
69.3%
100%
Learned 1D (Fix2siruc)
Factorizdd
Absolute coordinates
64.9%
71.5%
Relinive nccuney
Factoried (stack)
64.8%
71.5%
Paotorized OX)
64.2%
71.3%
60%
71.4%
Fourier
64.7%
71.3%
40%
Pactorized (+)
64.8%
71.4%
Train resolutions
Lexmed gpactstmut)
Simmoidal
Practioual coordinates
Packoized (+)
63 63.9%
70.8%
Foutter (fraciioaa)
0%
75%
256
640
ImageNet 10shot accuracy
Evaluatio resolution
(a)
(b)
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 ViT'ss performance
Out of distribution generalization. We directly evaluate JFT-pretrained NaViT on downstream datasets, employing label-map (66 from JFT-4B to ImageNet and robustness-variant -ariinnnss--visnsss (Object- Net (27] and ImageNet-A (28)). We compare the performance to compute- cotphee VIT baseline.

---

VIT
NaVIT
* 1//6
1/66
ImageNNet
InageNet-A
Object
75%
70%
55%
60%
70%
50%
50%
6.%
40%
45%
60%
30%
40%
20%
35%
53%
10
10
10a
10'
10 Training TPU chip hours
Training TPU chip hours
Training TPU chip hours
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 [igure[[]. NaViT performs much better on ImageNet-A A that has many images with an extreme aspect ratio and important information outside the center crop (Appendix G Same data in Table
Figure Figure[[ssowss shows that ViT compares favorably both on ImageNet as well as datasets variants that were specifically designed to test out of distribution performance. is interesting to note that Vi 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 ppeecttpreserving center crop that we apply to images for ImageNet-A and ObjectNet classification (same as in [17]), which is useful prior for ObjectNet,, but less for ImageNet-A (see Appendix for details). If no crop is applied to images and they'r re instead simply resized to the image resolution expected by the model (i.e., square for ViT, and aspect preserving resize the same number of tokens for Na ViT), then the observed difference is much larger (see Figure 21 in appendix).
Calibration. In addition to the in-depth accuracy analysis, we have also quantfied the quality of the uncertainty computed by the model. In particular, for our ImageNet MoeeNeelK-finnn K-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.04) as vary the number of patches per image in the range [128, 1024]. without any post-hoc recalibration. We provide further details the appendix Appendix
Inference trade-offs. Given the flexibility of the model, there are several viable approaches how one can maximize the aggregate accuracy under 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. [ ViT 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 [e further study cascading approach, which both provides Pareto optimal models and gives more precise trade-off opportunities.
Fairness signal annotation. We investigate annotating images with fairnes--related signals, such as those pertaining to gender and ethnicity. Prior research has shown that metrics, such 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 ViT can help in overcoming these challenges, train annotators on FairFace (32] and CelebA 33 datasets as lincar probes (i.e. using frozen features produced by
ViVIT or ViT), before comparing their accuracy.
First, ViT 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 x 448 in FairFace and 178 x 218 in CelebA, and we resize them to area 224" while preserving aspect ratios NaV. Despite having the same sequence length, NiV provides a higher prediction accuracy than ViT, as shown in T2 (left). We verify statistical significance using the Wilcoxon signed-rank test test 34), which reveals that the improvement in Na ViT is significant with p=3x x 10

---

furrfcee gender
faifcce ethnicty
faitace / gender
eetniciity
occmaccy
Madel VIT
0.0
Square
NAVIT
Natvee
100
250
500
1000
100)
250
500
1000
100
500
1000
100
2.00
500
1900
Pratsining Staps (K)
Protraining Steps (K)
Protraiming Stops (K)
Proteiniing Stups (K)
Figure 12: Evaluating the accuracy of annotators trained faimees--related signals using either NaViT-L/16 ViT-L/16: Left: NaVIT offers better representations that improve the accuracy annotators. Right: Using native aspect ratios in NaViT results in a higher performance when compared to resizing images to squares.
with 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 x 224, either square images with native aspect ratios. Figure 12 (right) shows that native aspect ratios in VaVIT improve accuracy, which is statistically significant at the 95% confidence level 0p=0.. Appendix H contains the full set of figures for other attributes in CelebA and Fairface.
Second, we apply inception-style cropping 35
3.6
Other downstream tasks
Semantic segmentation.
We finetune NaVIT to semantic segmentation on ADE20k following the linear decoder protocol of Seg-
dataset 36
Evalhuted natire resolations
menter 37 We use ViT-L/16 as baseline and compare its performance with that of NaViT-L/16. Both models are pre- trained on JFT-4B [18] with comparable compute budget.
mdoU nouU
30-
We experiment with different maximum resolution Riaz finetuning time: VIT takes in random square crops of resolution Riaaz Reaa while NaVIT's inputs randomly resized (with preserved aspect ratio) so that the total number of pixels is R This way, have the same finetuning cost for the ViT and NaViT models. Note that following common practice, in order alter the ground truth segmentation maps, both models evaluated at the native resolution 33 $7 36). This means that for ViT. the square predictions are resized from square back to native resolution. We observe in Figure 13 that VaViT outperforms ViT when transferred to semantic segmentation with the same maximum finetuning resolution Rmar. Note that ViT at R384 beats ViT at R512 while being twice fast (see Appendix(), An advantage of VaV over VIT is that it benefits from flexibility of resolution during training and can ingest seamlessly square and non- hon-square images.
VIT (squie R
NaVIT tres R
364 Max Finctuning Resolution
Figure 13: NaViT transfers competi- tively semantic segmentation. We transfer ViT-L/16 and VaViL-L/16 on ADE20k with fine-tuning at different resolutions. ViT consumes square im- ages while ViT preserves the aspect ratio of the original images, while main- taning the total number pixels the same as VIT.
Object detection.
Native-resolution training may be especially ben- eficial for fine-grained tasks such object detec- tion, which require image understanding at many different spatial scales. We use compute matched Na 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 | The NaViT-based detector
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%
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 het- erogeneity of the underlying spatio-temporal signals. In practice, cumbersome protocols have to be

---

established both for training and evaluation (e.g. spatial and temporal multi-cropping [40] Va alleviates some of these challenges by not only allowing training over different resolutions, but also over different temporal durations. We fine-tuned Va trained on JFT for Kinetics400 [4]] classifica- tion by extracting three different apatio-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 Na NaViT'a is an excellent starting point for multi-scale training due to the resolution diversity at training time. We observe that Va 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
4 Related work
Flexible Vision Transformers. FlexiViT developed 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 [ supported variable aspect ratios with 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 localization tasks such as segmentation and detection 43 Many works developed Vision Transform- which do the same 45), though some dispute the necessity for simple localization tasks [46]. VaV 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 promising avenue.
Accelerating training with mixed resolutions. Image modelling works considering resolution largely focus on accelerating pretraining with a fixed, low resolution [47 FixRes [M] is a popular technique whereby resolution is increased in 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 's main downside is its irteversibility: 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 the training, followed by "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 training pipelines.
Token dropping for improved efficiency. Research initially explored random token dropping 9 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.
Conclusions and future work
We have demonstrated that Patch n' Pack-the simple application of sequence packing to vision transformers siiermeer---innifiiantyy 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) AA] but overall Patch n' Pack enables 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.
10

---

References
[1] Alexey Dosovitskiy. Lucas Beyer. Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Ncil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scalc. In ICLR, 2021.
[2] Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alahdulmohsin, and Filip Pavetic. Flexivit: One model for all patch sizes. In CVPR, 2023.
[3] Kenton Lee, Mandar Joshi, Iulia Ture, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming- visual language understanding, 2022.
Ming-Wei Wei Chang, and Kristina Toutanova. Pix2struct: Screenshot parsing as pretraining for
[4] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, pages 248-255, 2009.
[5] Agrim Gupta, Piotr Dollár, and Ross B. Girshick. LVIS: A dataset for large vocabulary instance segmentation. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR, 2019.
[6] Xi Chen, Xiao Wang. Soravit Changpinyo, A. J. Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, Alexander Kolesnikov, Joan Puigcerver, Nan Ding, Keran Rong, Hassan Akbari, Gaurav Mishra, Linting Xue, Ashish Thapliyal, James Bradbury, Weicheng Kuo, Mojtaba Seyedhosseini, Chao Jia, Burcu Karagol Ayan, Carlos Riquelme, Andreas Steiner, Anelia Angelova, Xiaohua Zhai, Neil Houlsby, and Radu Soricut. PaLI: A jointly-scaled multilingual language- model. arXiv preprint arXiv:2209.0677 2022.
[7] Mario Michael Krell, Matej Kosec, Sergio P Perez, and Andrew William Fitzgibbon. Efficient sequence packing without cross contammmttiinnnnn Accelerating large language models without impacting performance. 2021.
[8] Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Syxtems 2020, NeurIPS 2020, 2020.
[9] Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang. Shih-Fu Fu Chang. Yin Cui, and Boqing Gong. Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text. Advances in Neural Information Processing Systems, 34:24206-24222 2021.
[10] Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-image pre-training via masking, 2023.
[11] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou. Fixing the train-test resolution discrep- ancy. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, 2019.
[12] Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. Efficient transformers: A survey, ACM Computing Surveys, 55(6):1 28, 2022.
[13] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:201 1.0:::20/110000000 1.04006, 2020.
[14] Markus N Rabe and Charles Staats. Self-attention does not need O(n²) memory. arXiv preprint arXiv:2112.05682, 2021.
[15] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io awareness. In NewrlPS, 2022.
[16] Basil Mustafa, Josip Djolonga, and Mostafa Dehghani. On efficient losses for distributed contrastive learning. arXiv preprint, 2023.

---

[17] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer. Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme, Matthias Minderer, Joan Puigoerver, Utku Evei, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetié, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters, In International Conference on Machine Learning, 2023.
[18] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, pages 12104-12115 13, 2022.
[19] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICMI., pages 8748-8763, 2021.
[20] Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big Transfer (BiT): General visual representation learning, In ECCV, pages 491-507, 2020.
[21] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris I Leary, Dougal Maclaurin, George Necula, Adam Paszke. Jake VanderPlas, Skye Wanderman-Milnn and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[22] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zec. Flax: A neural network library and ecosystem for JAX, 2020.
[23] Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab, Matthias Minderer, and Yi Tay. Scenie: A jax library for computer vision research and beyond. In CVPR, 2022.
[24] Yuhao Zhang, Hang Jiang, Yasuhide Miura, Christopher D. Manning, and Curtis P. Langlotz. Contrastive learning of medical visual representations from paired images and text. In Proceedings of the Machine Learning for Healthcare Conference, MLHC, 2022.
[25] Mostafa Dehghani, Anurag Arnab, Lucas Beyer, Ashish Vaswani, and Yi Tay. The efficiency misnomer. arXiv preprint arXiv:2110.12884 2021.
[26] Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Korblithh, Rebecca Roclofs, Raphael Gontijo Lopes, Hannanch Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Ro- bust fine-tuning of zero-shot models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 7949-7961. IEFF, 2022.
[27] Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenen- baum, and Boris Katz. ObjectNet: A large-scale bias-controlles dataset for pushing the limits of object recognition models. In NewrlPS, pages 9448 9458, 2019.
[28] Dan Hendrycks, Kevin Zhao, Steven Basart, Steinhardt, and Dawn Song. Natural adversarial examples. In CVPR, pages 5262-15271. 2021.
[29] Jeremy Nixon, Michacl W Dusenberry, Linchuan Zhang, Ghassen Jerfel, and Dustin Tran. Mcasuring calibration in deep learning. In CVPR Workshops, 2019.
[30] Julius Adebayo. Melissa Hall, Bowen Yu, and Bobbie Chern. Quantifying and mitigating the impact of label errors on model disparity metries. In ICLR, 2023.
[31] Inioluwa Deborah Raji and Joy Buolamwini. Actionable auditing: Investigating the impact of publicly naming biased performance results of commercial ai products. In Proceedings of the 2019 AAA//ACM Conference on AI, Ethics, and Society, pages 429-435, 2019.
[32] Kimmo Kärkkäänen and Jungseock Joo. Joo. Fairface: Face attribute dataset for balanced race, gender, and age. 2019.
[33] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning f'ace attributes in the wild. In ICCV, 2015.
[34] Frank Wilcoxon. Individual comparisons by ranking methods. Springer. 1992.
[35] Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In CVPR, 2016.
12

---

[36] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler. Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In CVPR, 2017.
[37] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In ICCV, 2021.
[38] Mathilde Caron, Ncil Houlsby, and Cordelia Schmid. Location- awarc self-supervvsed transformers for semantic segmentation. arXiv preprimt arXiv:2212.02400, 2022.
[39] Matthias Minderer, Alexey Gritsenko, Austin Stone, Maxim Neumann. Dirk Weissenborn, Alexey Dosovit- skiy, Aravindh Mahendran, Anurag Arna,, Mostafa Dehghani, Zhuoran Shen, et al. Simple open- vocabblayy object detection with vision transformers, arXiv preprint arXiv:2205.062300 2022.
[40] Anurag Amnb, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. ViVIT: A video vision transformer. In CVPR, 2021.
[41] Will Kay. Joao Carreira, Karen Simonyan. Brian Zhang, Chloe Hillier, Sudheendra Vjjayanaraiimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, al. The kinetics human action video dataset. arXiv preprint arXiv:1705.06950. 2017.
[42] Piergiovanni, Weicheng Kuo, and Anelia Angelova, Rethinking video vits: Sparse video tubes for joint image and video learning, 2022.
[43] Olaf Ronncherger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab. Joachim Hornegger, William M. Wells III, and Alejandro F. Frangi, editors, Medical Image Computing and Cumputer-sssisted Intervention MICCAI, 2015.
[44] Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In 2021 IEEE/CVF International Conference Computer Vision, ICCV, 2021.
[45] Yanghao Li, Chao-Yuan Wu, Haoqi Fan, Karttikeya Mangalam, Bo Xiong. Jitendra Malik, and Christoph Feichtenhofer. Mvity2: Improved multiscale vision transformers for classification and detection. In CVPR, 2022.
[46] Wuyang Chen, Xianzhi Du, Fan Yang, Lucas Beyer, Xiaohua Zhai, Tsung- Yi Lin, Huizhong Chen, Jing Li, Xiaodan Song. Zhangyang Wang, and Denny Zhou. A simple single-scale vision transformer for object detection and instance segmentation. In Shai Avidan, Gabriel J. Brostow, Moustapha Cissé, Giovanni Maria Farinella, and Tal Hassner, editors, Computer Vision ECCV 2022, 2022.
[47] Hugo Touvron, Matthieu Cord, and Hervé Jégou. DeiT III: Revenge of the ViT. In ECCV, 2022.
[48]
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning. Yue Cao, Zheng Zhang, Li Dong, al. Swin Transformer V2: Scaling Up Capacity and Resolution. CVPR, 2022.
[49] Chao-Yuaa Yuan Wo, Ross B. Girshick, Kaiming He, Christoph Feichtenhofer, and Philipp Krähenbühl. multigrid method for efficiently training video models. CVPR, pages 150-159. Computer Vision Foundation/ / IEEE, 2020.
[50] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In CVPR, pages 15979-15988 IEEE, 2022.
[51] Wuyang Chen, Wei Huang, Xianzhi Du, Xiaodan Song, Zhangyang Wang, and Denny Zhou, Auto-scaling vision transformers without training. In The Tenth International Conference on Learning Representations, ICLR, 2022.
[52] Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffmann. Token merging: Your but faster. CoRR, abs/2210. 0.4901 2022.
[53] Hongxu Yin, Arash Vahdat, Jose M. Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In CVPR, 2022.
[54] Samira Abnar, Mostafa Dehghami, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training, arXiv preprini 1rXiv:2110.00095, 2021.
[55] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. arXiv preprint ArXiv: 2019.
13

---

[56] Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing, EMNLP, pages 66-71, November 2018.
[57] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz 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 gender reduc- tionism? social implications embedded gender recognition systems. Proceedings of the 2018 chi conference on human factors in computing systems, pages 1-13, 1-,1,2015 2018.
[59] Keyes. The misgendering machines: Tran//tci automatic gender recognition. Proceedings of the ACM on computce interaction, 21S 2018.
[60] Robert Geirhos, Kantharaju Narayanappa. Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. In NewrIPS, pages 23885 23885-23899, 2021.
14

---

A Broader impacts & limitations
Broader Impacts
NaVIT enables training of vision transformers on variable size inputs, which has 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.
Limitations
There is 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 ViT in these particular applications. We consider this area as priority for future research and follow-up work.
Implementation complexities The modelling changes shown in Figure I can complicate ViT; in particular, packing-awan 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). NaV can however be run with the same sequence length a corresponding ViT model (c.g. VIT at resolution 384 with sequence length 576 vs VTT at resolutions 64:384 with sequence length 576; the latter will fit on average 2x 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- 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 the paper are with ViT-B/32, VIT-B/16, and VIT-L/16. We use reciprocal square-root learning rate schedule, with linear warmup and cooldown, with a maximum value of 4, and phases. We follow [18 (54) and use higher weight decay of 3.0 on the head compared to the body's 's weight decay of 0.03 during upstream training improve transfer to downstream tasks. For our experiments, we evaluated both ViT 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 ViV model for each ViT size and computational budget, allowing to perform "compute-matched"" matched" comparisons (25).
Table the pretraining specifications for both ViT and VaV models. During the pretraining phase, ViT models were trained using images of size 224 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 pixels while preserving the aspect ratio. Although training Na ViT 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 fixed computational budget while maintaining perfor- mance 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%.
15

---

Table Pre-training details of VIT and NaViT with supervised classification.
Name
TPU Hours
Train Steps
Cooldown
Sequence Length
Images Per Seq.
Batch Size
Training Images
Steps
x.9 10'
1.0x
1.0x 10*
49
1.0
#4.00 x 10
4.0 10
1.9x 10
2.5x
5.0 10*
49
1.0
xx.00 10
1.0x 10
ViT-B/32
3.9x 10
5.0x
1.0x 10
49
1.0
x4.0 10
2.0 109
7.9x10²
1.0
1.0x 10
49
1.0
x4.0010*
4.0x
2.6x10²
1.0x10$
1.0x 104
196
1.0
#4.00
4.0x
6.6x102
2
5.0x 10
196
1.0
10.
1.0x 10
VIT-B/16
1,3
5.0
1.0x 10a
196
1.0
x4.0010*
2.0x 109
206x
100
1.0x 10
196
1.0
10
1.0%
3.4x102
100
1.0 104
196
1.0
x4.0
1.0
1.3x
2.5x10
5.0 10*
196
1.0
=4.0010
1.0
VIT-L/16
2.7x 103
5.0x105
1.0 105
196
1.0
100
2.0x109
5,4x10a
100
1.0 10
196
#x.00 1.0
1.0%
7.9 x00'
9.8x10
1.0x 104
64
5,41
#22x
2.1x
1.9 102
2.4x
5.0 101
64
5.41
1/222101
5.3 109
HaViT-B/32
4.8 3.9
1.0x 10S
64
5,41
#2.2x 104
1.0
9.7x 1.9
1.0x 10
64
5,41
2.1x1010 20
2.6x102
9.3x
1.0 104
256
4.87
10
1.8x109
6.6x10²
2.3x10
5.0 104
256
4.88
~1.9100
4.6
NaViT-B/16
1.3x 103
4.6x105
1.0% 105
256
4.88
101
9.2
2.6x103
9.2x10$
10
256
4.88
10
1.8x10
1.2
9.4
1.0
256
4.88
~1.9100
1.9 109
2.4x 103
5.0 104
256
4.87
10*
4.8 109
NaViT-L/16
2.7x 103
4.8
1.0x 105
256
4.87
104
9.6 10
1.4
9.6x
1.0x 10
256
4.88
x119110
1.9x1020
B.2 Contrastive pretraining
use the 32000 token [55] sentencepiece (56] tokenizer. By default, text sequences are truncated to 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 x 10-1 Weight decay of 1x 10 is applied consistently all kernels in the model (no change for projection heads). By default, image side-lengths are sampled (64, and no other image augmentations are applied.
B.3 Packing algorithm
Packing of examples into sequences is done alongside batching, A simple greedy approach used examples
which adds examples the first sequence with enough remaining space. Once can fit, sequences are filled with padding tokens, yielding the fixed sequence lengths needed for batched operations. Such simple packing algorithm can lead 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 NaV 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 x 224, with inception crop followed by random horizontal flipping. For NaViT, we have no preprocessing besides resizing to the sampled resolution and token dropping.
16

---

165-
228
45x bod
Venke
NAVI
Trarkg TPC clp hoar
Tekex
teags
Figure 14: Training more images given fixed computational budget, is key factor contributing the computational efficiency of NaViT compared to to standard VIT.
B.5
Sample efliciency with VaV
Based on the experiments ran, believe lot of the observed performance benefits due speeding up training (i.e., seeing more images in a given compute budget). This evident Figure] [ 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 concern, for given number of images seen, it is best to train purely large resolution, with no token dropping. Asymptotically we therefore expect the token dropping and mixed resolutions eventually hurt NaViT in unlimited compute regime. This assumes large enough dataset that overfitting is not an issue. Eventually one may enter a regime where the token dropping and variable resolution is useful data augmentation 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, and standard deviation We aim to sample dropout rate d E m0.0,dmax], with some mean dy..
Accordingly, we sample E [0.1] ~ B(a,3) and set drop rate d= x dmax- a and | are set such that the mean of u is Uh omax The maximum supported variance for a beta distribution of mean Up un (1.112u4) u,);; we pick by default variance o² 0.3u,, (1 Un,), which we found to work well in practice. The resultant distributions of token dropouts for different settings of dμ and dmax are shown Figure 15a
dp
25%
drop
rate
rate
5%
20%
ACNULE
4%
Porcent
1.5%
3%
10%
2%
5%
1%
0%
25%
50%
75%
100%
25%
50%
75%
100%
25% 50%
75%
100%
0%
25%
50%
drop
drop
drop
drop rate
(a) Beta-sampled token drop rates parameterised by the mean yu and the max (b) Sampled resolution resolution-depeeddett token drop rates
drop rate dmax
Sampling resolution-dependee dropping rates Given input data with sequence lengths rang- ing from Smin to Smax. we sample dropout rate from a truncated normal distribution (2...0.) where samples more than two standard deviations away from lr are rejected.
The mean of this distribution μ is set according to the minimum and maximum token dropping rates dmin and dmax, and simply scales linearly with the sequence length s (such that s= =minn has μ d and Smax has μ dmax-
Figure 15b shows example distributions of sampled drop rates given inputs with resolution R U 64,384) and different values of dmin and dmax.
17

---

0.75
Total images soen
Avenge drop
Image Ner 10shat
75%
0.50
70%
200
400
200
100
200 &
400
Total training steps
Total training steps
training sleps
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 fixed token dropping rate of p=
B.7 Scheduling token dropping rates
We experiment with 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:
((n:Pmin-PPaase H1,T) = Punin (Pmax Pmiin) )
where a represents the sigmoid function; Pwiin Pmax control the minimum and maximum dropping rate applied; and μ and r control the shape of the schedule. We experimented with both increasing (r > 0) and decreasing (r schedules. all cases set Pauin 0.2 and Pmux = 0.8. Figure T6] shows that, by decreasing the dropping rate throughout training, one can improve the final accuracy, at fixed training cost. Conversely, increasing the token dropping rate harms performance.
C Model information
C.1 Positional embeddings
Extending ViTs to variable input sizes necessitates rethinking positional embeddings added to every token after embedding. We considered several variants of positional embeddings, and evaluated them based on (1) the best performance model using them achieve within training distribution of input sizes; and based on (2) how well these models perform when evaluated image sizes outside of the training distribution. Results and discussion of these experiments can be found Section 3.4
Broadly, considered positional embeddings that varied along three axes: (1) whether they were learned, parametric or fixed; (2) whether they were absolute or fractional; and (3) whether they are factorized.
Aboolate cooediastes
Proctiooa coondaates seppett wwwa pelapereeeg
Practional cocedisates inpert we pneering
Image credit: Matthew Henry burst shopify.ph...... httppiiy.cooooooommoo http://hooooooooo/ wtwp:oss//ooottsnnnnnno....
Figure 17: We use two different views of the same image, of resolutions 96x 128 and 224 160, and demonstrate different coordinate systems when using patch size 32.
Absolute and fractional coordinates A natural way of indexing token within an image is to select priori maximum possible image side length (shared for width and height) maxLen, and to assign to token integer coordinates (x.y) based on their original location within the image. Embedding coordinates defined this way allow models consume images with resolutions up to R= P- max Len. However, when learned absolute coordinate embeddings are considered, extreme values of and must also observed during training, which necessitates training on images with varied aspect ratios and limits models generalisation.
18

---

To alleviate the necessity of observing extreme aspect ratios and image size during learning of positional embeddings, we also consider fractional coordinates, which are normalized to the actual size of the input image and are obtained by dividing the absolute coordinates ar and y above by the number number of columns and rows respectively, i.e. the corresponding side length. Doing this allows the model to observe extreme token coordinates during training, which intuitively should help with generalization to higher resolutions. However, this is accomplished at the cost of obfuscating the input images aspect ratio.
Factorized embeddings We further consider whether coordinates and y should be embedded independently or jointly. In case of independent embedding, the two coordinates æ and are embedded independently, and their embeddings are combined via addition or by stacking. For joint embeddings and embedding for each position (z.y) is obtained directly.
Learned, parametric and fixed positional embeddings Finally, we also explored the relative benefits of fixed, learned and parametric embeddings. For fixed embeddings followed [57] and used sinusoidal positional embeddings, and learned embeddings were implemented as in []]
For parametric positional embeddings we followed g and used Fourier embeddings. Specifically, coordinates (a.y) were mapped using single linear layer before applying sin and cos activations to them, and stacking the results to obtained the positional embeddings.
Experiments Because not all combinations of the above embedding choices are equally promising or natural, we experimented only with subset of them shown in Table and Figure T0a
Table 3: Classification of positional embedding experiments from Figure
10a
Name
Coordinates
Type
Factorized
Learned 1D (ViT)
Absolute
Learned
No, position in flatted token sequence
Learned 2D (Pix2struct)
Absolute
Learned
No
Factorized abs. (+)
Absolute
Learned
Yes, sum
Factorized abs. (stack)
Absolute
Learned
Yes, stack
Factorized abs. (x)
Absolute
Learned
Yes, product
Fourier abs.
Absolute
Parametric
No
Sinusoidal abs.
Absolute
Fixed
Yes, stack
Factorized frac. (+)
Fractional
Learned
Yes, sum
Fourier frac.
Fractional
Parametric
No
Sinusoidal frac.
Fractional
Fixed
Yes, stack
In sinusoidal and factorised embeddings experiments with fractional coordinates fractional coordinate
embeddings were obtained from absolute coordinate embeddings via bilinear interpolation.
D Inference strategies
We performed various experiments to measure model quality for given runtime cost. The runtime can be tuned by changing the number of processed patches, or by using choosing different size of the model.
We firstly looked at how model quality changes in respect to decreasing area of the image compared to native resolution, presented in Figure 8a We observed that ImageNet 4 model retains most of the quality down to 40% of the image size. After that, the quality drastically decreases. On the other hand, increasing the size of the image have a diminishing return in quality. This can be directly compared with random token dropping as an alternative to resizing, which showed to be very ineffective way to decrease number of patches during inference Figure T8b
Please note that this highly depends on the native resolution of the images in the dataset e.g. dataset with twice big images than ImageNet can probably be safely resized to 20% of area.
better way to quantify the performance is by giving constant compute budget corresponding to number of patches. Figure 19a[shows that resizing the image (while preserving aspect ratio) to 256 tokens retains most the quality (within 0.3%). This corresponds to 256x256 area (given patch size of 16). At 128 tokens (181 1x181 181 area) the quality difference reaches 1 1% and drops significantly after that.
19

---

Imoganee accunity
NAVIT-LJAS
6734
NeVIT-L
65-
NaVIT-B/16
dnp
NaVi--Lono opscale)
0.75
0.55
0.75 Image
seggenos leugh
(b)
(c)
Figure 18: (a) The effect of resizing the image. (b) Dropping random tokens is ineffective way to decrease number of patches compared resizing the image. Data from NaViT-L/16, (c) Given number of patches compute budget, is beneficial upscale the image.
Imagenel accuracy
88
roduction
1%
Imagenet accuracy
88.0
256
512
640
768
896
1024
87-
1%
87.5
86
roduction
128
540
NaViT-L/16
18
512
-----
NaViT-L/16
NaViT-B/I6
86.5
NaVIT-B/I
64128
256
384
512
640
768
896
1024
sequence length
average runtime in per image
(b)
Figure 19: (a) Quality on ImageNet in respect to number of patches (sequence length). (b) Runtime of models compared the accuracy ImageNet.
Here we also resized the image past its native resolution in case it already fit the given sequence length budget. We observed that it is beneficial to resize the image to the given sequence length past the native resolution to keep monotonic increase in quality, which is showed on Figure T8c
Figure 19b presents the runtime of NaViT-L/16 and NaViT-B/16 for different sequence lengths. We can see that Na ViT-L//1 at sequence length 128 is as fast as Na ViT-//66 with sequence length 512, while having almost 1% difference in quality.
E
Cascades
Another strategy to be more compute efficient would be to assign more tokens (and thus FLOPs) to the examples deemed hard by the model. This in particular interesting for bulk inference workloads, where one can amortize over large datasets and where only the total inference time matters.
124
$8.0
256
128
Cascade (ViT-B)
ImageNet
87.0
ViT-B/16
Cascade (VIT-L)
36.0
ViT-L/16
85.5
Pareto frontier
i's
20
25
Average inference time in milliseconds
Figure 20: Performance model cascade versus the average inference time. The labels select points denote the number of tokens at that scale.
20

---

To evaluate the feasibility of this approach, we consider two sequence lengths n1 and ny with respective inference times 1 and tz. Then, we (i) send all examples though the model with n, tokens, and send only the E (0,1) ) deemed hardest (those with the smallest maximum probability) to the model with n2 tokens. To have an input almost exactly fit into 11 or n2 tokens we perform and aspect ratio preserving resize. Hence, the total amortized inference time per-example is & 422 while the accuracy obtained by combining the accuracy of the first model on the 1 1 most-confident fraction of the data, and the performance of the more expensive model the remaining data. By considering several pairs of models and varying a we obtain the plot in Figure 20 As we can see this strategy is indeed useful and provides not only the best performing models given compute budgets, but because is a real parameter one can obtain very fine-gen mne-grained trade-offs.
F
Calibration
To evaluate behaviour of the predicted uncertainties with scale, we compute the calibration error of -B sized maageNet-finetunee model (the -L model performs similarly). Note that these models were trained with sigmoid loss, i.e., the 1000 labels were predicted independently without enforcing that the probabilities should sum to 1. As we varied the sequence of tokens per example between 128 and 1024, we obtained very stable calibration errors (top-I, using & and 30 buckets, i.c., the settings from 29 which we present in T
Table 4: Expected calibration error on ImageNet- IK with varying sequence lengths.
Sequence Length
128
256
384
512
640
768
1024
Calibration Error
0.047
0.046
0.048
0.047
0.047
0.046
0.045
G Out of distribution evaluation
For VIT, we apply the "Crop" strategy from namely an of the image for ObjectNet and ImageNet-A, and square resize followed by 87.5% central crop for the other datasets. We also apply a simple "Resize" strategy that does not crop the images. For ViT, both the "Crop" and the "Resize" strategy do an aspect preserving resize of the target images.
crop of the central 75%
VIT
NaVeT
BA2
-X- B/16
.W. L/16
ImaggeNet
ImageNet-A
ObjectNet
50%
75%
60%
70%
50%
65%
40%
35%
30%
30%
25%
55%
20%
20%
10%
15%
1012
10
10
1015
10
10
Training TPU chip hours
Training TPU chip hours
Training TPU chip hours
Figure 21: Same evaluation as in Figure 11 but without any special preprocessing of the images before the evaluation. Employing simple resize (square for VIT, aspect preserving for Na ViT) results in much better performance on datasets that have images with an extreme aspect ratio. Same data in table T
H Fairness Signal Annotation
In Figure 22 we demonstrate that using native image resolution improves the performance of fairness signal annotation. Prior research has shown that metries, such as group calibration, are vulnerable to labeling errors, particularly for underrepresented groups. Moreover, this problem persists even when accounting for label noise during training (0). Thus, reducing the labeling error of fairness signals has the potential of improving the reliability of bias mitigation and post-hoc auditing [31]. Nevertheless,
21

---

Table 5: Detailed results of out evaluation of pretrained models with a label label-map (sce Section 3.5 in Figure and Figure 21
Same data as
ImageNet
ImageNet-A
ObjectNet
VIT NaViT
VIT
NaVIT
VIT
NaViT
Compute
Crop
B/32
7.9 10*
54.4
63.5
14.7
32.7
32.9
41.2
1.9 x 102
60.2
65.6
22.0
37.0
38.0
43.5
3.9 10
63.7
67.2
26.6
40.1
41.7
45.1
7.9 10
65.5
68.2
30.7
42.5
44.2
45.7
B/16
2.6 10
66.4
68.5
37.3
52.3
47.2
6.6 10
61.7
66.0
27.0
44.2
41.3
45.9
x 102
70.3
71.2
48.8
57.0
52.8
50.7
2.6 10
68.7
70.1
43.5
55.1
50.5
49.8
L/16
x 10
51.5
63.5
53.3
1.3x 10
66.4
71.1
39.2
58.6
47.6
52.1
2.7 10
73.9
75.1
60.4
68.9
57.7
57.9
5.4 10
73.0
74.6
57.6
67.9
56.2
57.1
Resize
BB/2
7.9 10"
51.7
64.0
12.6
26.7
15.9
31.6
1.9 102
57.9
66.2
17.2
30.1
20.0
33.8
3.9 10
61.6
67.4
20.7
32.6
23.5
34.6
7.9 10
63.6
68.6
22.8
35.0
25.3
36.1
B/16
2.6 102
65.4
69.2
27.4
43.9
28.7
38.7
6.6 10
60.4
66.5
20.4
36.8
23.6
35,3
1.3x 10
69.5
72.5
36.2
51.5
34.2
42.1
2.6 10
67.7
71.0
32.5
48.6
32.0
40.5
L/16
x.4 102
70.0
73.9
38.5
59.9
35.3
45.6
1.3x 10
65.1
71.5
28.2
51.6
29.3
41.5
2.7 10
73.2
76.0
47.3
65.5
39.8
48.8
5.4 102
72.3
75.3
44.3
64.1
38.8
48.2
celeba/ gender
Gatrface ago
fairface ctuncitt
80
Model
VIT
98
100
25/0
500
1000
93
100
250
500
1000
500
1000
Protraining Steps
Protraining Steps
Pretraining
Pretmuning Steps
crleta/ gender
Eairficce / gender
fairface age
ethnicity
Squar
11:
01:0
Native
100
250
300
1000
100
250
500
1000
100
500
3000
100
2.03
340
2000
Pretraining Steps (K)
Pretraining Steps (K)
Pretruising Steps (K)
Pretraining
Steps (K)
Figure 22: Summary results of evaluating the accuracy of annotators trained on fairness-relatee signals using either NaViT-L/I6 VIT-L/16. TOP: NaVIT offers better representations that improve accuracy annotators. воттом: Using native aspect ratios in NaViT results a higher performance when compared to resizing images to squares.
we emphasize that while NaViT improves the annotation accuracy in these tasks, care must be taken such situations since classifiers can be inaccurate and lead broad categorization of people that misidentifies real identities. We encourage readers to delve into the comprehensive work outlining such potential risks, e.g. 58 15,, for further insight. In assessing the technical capabilities of Na VTT, our intent is not to promote or encourage their application in inappropriate contexts. Rather, our objective only to illuminate these technical findings for scenarios where they may be considered beneficial, such when measuring the level of diversity in dataset or auditing/mitigating biases in predictive models. We strongly advocate for responsible Al use, maintaining that the benefits of technological advancements should not overshadow the importance of user safety and privacy. Al tools, including
22

---

ours, should always be deployed judiciously, with a keen awareness of potential risks and a commitment to avoiding harm.
I Evaluation on model-vs-human OOD datasets on different resolutions
Just like NaV,T, human visual perception works across flexible aspect ratios and resolutions (just imagine how strange the world would look like if we could only see it through a 224 x 224 pixel window!). We investigate how the ability cope with variable resolutions affects performance on poodel-ss-human", a benchmark of 17 challenging datasets [60] For this purpose, we replicate the setup from Figure but instead of evaluating ImageNet accuracy, we evaluate OOD accuracy on the model-vs-human benchmark.
rokelimn
Accunncy (30)
3
Accurbi
60
Ricuracy (S)
90
NaVIT lone per res)
18
16
80-
NJVTT-64552
64
128
224
384
512
64
128
224
364
512
04
128
224
384
512
04
128
224
384
552
Finedlume resolutioo Rg
Fire-tuee resolution lon Ry
Fine-lure resoiition R.
Fine-lure resolation Rr.
sighoppsss
Accuviy (N)
Accurang 193)
fccuratty 190
Accuracy (9)
20-
128 274
184
512
14 128 22t
512
65
128
224
184
512
128
224
104
512
Fine tunc resolution R.
Fine tune resolation Rs
Fine tunc resolution R.
Fine tune resolution R.
phose scrambiing
eccolani
ctdblonili
Alcuracy (N4)
Accuras Mi
Aecuunery (he
Aucuracy 1961
60.0
60-
128
224
resolution Re
384
512
64
128 224
Fite-ture resolution Re stylized
364
512
Fine tume resolation
354
64 128
Finc resolution
224
384
512
seatch
sthowelle
wage
70-
Recurracy (8)
88
Accuracy
Accuraty
60-
384
224
502
resolution
Fire-lune resoldion
resolution
Fine-lame resolution Ry
Figure 23: OOD accuracy on "model-vs-human datasets across different fine-tuming resolutions. A single Na VIT model trained on varying resolutions (red) performs roughly on par as fine-tuning one ViT model per test resolution (blue). The ViT haseline (orange) is mostly worse than NaVIT models for lower resolutions and mostly better for higher resolutions.
This corresponds to testing JFT B/16 models finetuned on ImageNet at various resolutions. The test dataset has a fixed 224 224 square resolution; thus we resize the test images to fit each model's fine- tuning resolution. Note that using square images has been standard practice designed for convolutional networks, but NaViT models no longer require square input, thus existing benchmarks are not tailored to those new possibilities. For datasets with multiple difficulty levels (such as different levels of blur), we average performance across levels while excluding levels that are too easy (not OOD) or too hard (human performance is close to chance), which follows the approach of (author?) [60 as explained in their "Appendix G Benchmark scores".
To ensure a fair comparison, we use models that are compute-man motchee for pretraining, and during fine-tuning, compute and data are identical for all models. The results of our comparison are shown in Figure 23 Overall, a single ViT model trained on varying resolutions (red) performs roughly on par with fine-tuning one ViT model per test resolution (blue).
The VIT baseline (orange) is mostly worse than ViV models for lower resolutions and mostly better for higher resolutions. It may be worth noting that ViT models have bit of an advantage in this comparison since they are fine-tuned on square images, whereas NaViT models are fine-tuned on flexible resolution images (preserving the image's s aspect ratio) that have the same number of pixels, while not necessarily being square.
For the purpose of our comparison, we exclude the "cue-conflie confeectnict" dataset from the OOD evaluation, since there is no objective ground truth class in the case of images with texture- texture-shape cue conflict.
23


Reviews:
[{"review_id": "OKb46fuhuC", "reviewer": "Reviewer_WGZP", "summary": "The authors propose to use example packing to train ViTs, where training examples of various lengths are packed into a single sequence. This requires a few straightforward architecture changes, including modified attention masking, pooling, and positional embeddings. This scheme allows for some interesting ideas, such as variable image resolutions during training, variable token drop rates, and adaptive inference time computation.\n\nExtensive experiments in the paper show that NaViT results in more efficient training and finetuning (in terms of TPU hours) and a better compute-performance tradeoff at inference time. Further analysis shows that mixed resolution training is beneficial to model performance, that the time-varying token drop ratio allowed by the model can improve results when the correct schedule is used, and that the proposed factorized positional embeddings can offer very good generalization to aspect ratios and resolutions unseen during training. Finally, additional experiments show promising behavior with respect to calibration, fairness, object detection, and video classification.", "questions": "See weaknesses section, thank you.", "limitations": "I think the \"limitations\" section from the appendix could be more forthcoming and candid. It mostly just says \"NaViT\" is a great idea, but we didn't get to apply it to all the tasks we wanted to. It could benefit from a more honest discussion of the method's limitations, such as compute overhead from packing or additional architectural changes needed to support example packing.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Overall, this is a very comprehensive paper. The number of experiments and the aspects of the proposed models performance evaluated is impressive. The authors touch on a large number of relevant properties of the model, including calibration stability, positional embedding ablations, and out of distribution performance. Additionally, results are quite positive, showing promising results on image classification, object detection, video classification, and training efficiency. Finally, the appendix is thorough, and gives enough details to accurately reproduce all experiments from what I can see.", "weaknesses": "According to appendix section B.1, classification experiments seem to be \"compute-matched.\" Are there any experiments analyzing the asymptotic performance between ViT and NaViT, or can the authors discuss this? The experiments in the paper seem inconclusive on this point. \n\nThe question of compute matching also applies to downstream experiments. For example, in the fairness analysis, are compute matched ViT and NaViT being compared? If so then the conclusion from Appendix H that \"native image resolution improves the performance of fairness signal annotation\" may not hold.\n\nNaViT seems to benefit from training at the original aspect ratio, and at variable resolutions. However, could either of these benefits be achieved through scale and crop data augmentation? What data augmentation is used for ViT? Additionally, I was under the impression that  image stretching or shearing could be useful as a data augmentation. Could the authors please discuss this?\n\nWhere are the contrastive results? I might have overlooked something, but I can't find zero-shot imagenet or COCO image-text retrieval results as discussed on line 154.\n\nMinor typos:\n- Figure 2 typo: \"image 2\" is repeated twice in the \"data processing\" part\n- Figure 9 is before Figure 7 and 8\n- It seems that a lot of (if not all) appendix section references are incorrect. Examples of these typos occur in line 152, line 158, line 339, and line 120.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Firstly, many thanks for your time and detailed feedback. It was very useful and constructive, and we made multiple improvements to the manuscript based on your comments. Please see below for some detailed responses to the weaknesses you pointed out.\n\n1. `Compute matching and asymptotic performance`\n\nThis is a good point; based on current experiments, we believe a lot of the observed performance benefits are due to speeding up training with NaViT (i.e. seeing more images in less time).\nWe introduce [this figure](https://i.imgur.com/Rpn4GBU.png) which augments Figure 1. In the center, we can show that per-token seen, NaViT trains much faster than ViT (and per TPU-hour, left), and this is because Patch 'n' Pack enables NaViT to see tokens from more unique images (figure, right).\n\nComparison of asymptotic performance can be difficult. First, when training on these large pre-training on large datasets, it is prohibitively expensive to train to convergence, so more training always yields improvements. Therefore it is common for foundation models to compare at a fixed (often large) training budget as we show above. Even at the largest budget ( a large 2000+ TPU chip hours), NaViT is still far ahead.\n\nFurther, in the infinite compute limit, for a given number of images seen, it is best to train on purely large resolution, without token dropping. This is not a setting that we explore since it substantially slows down training. Asymptotically, we therefore expect the token dropping and mixed resolutions to eventually hurt NaViT, although we believe that it would be very expensive to approach this regime. \nWe have added a new section (Appendix B.5) to the appendix with this figure, discussing asymptotic performance, and added pointers to it in the introduction, in the updated version of the paper. \n\n2. `Compute matching downstream experiments`\nIt would indeed be unfair to compare non-compute matched models, since pre-training on large datasets (like JFT) for longer always improves performance. Downstream experiments were all done using models pre-trained for the same amount of time, and evaluated at the same effective resolution (number of tokens) - i.e. compute matched for both training and eval cost. Specifically, *all downstream experiments used the top-rightmost points in Figure 1 (ViT-L/16 and NaViT-L/16)*. We have added a paragraph to the experimental setup to make this point clear.  Many thanks for raising this. \n\n3. `NaViT seems to benefit from training at the original aspect ratio, and at variable resolutions. However, could either of these benefits be achieved through scale and crop data augmentation?`\n\nWe do not believe aspect-ratio conservation provides an augmentation effect. It may interact with other data augmentations; e.g. inception-cropping samples a rectangular bounding box, and then squashes it back to a square. This may work better with aspect-ratio preserving models.\n\nVariable resolution and token dropping could indeed act as data augmentation; we did not conclusively demonstrate this in the paper, since in the large-scale pre-training regime we explored typically doesn’t benefit significantly from data augmentation.\n\nWe are actively exploring this in smaller-scale settings; this is however quite an undertaking. Many of the data augmentations/training techniques which enable training ViT on small-scale models, such as mixup, require some rethinking when dealing with non-square or variable resolution images. We therefore leave research on data augmentation for Native resolution vision models to future work.\n \n4. `What data augmentation is used for ViT? Additionally, I was under the impression that image stretching or shearing could be useful as a data augmentation. Could the authors please discuss this?`\n\nWe follow the original Vision Transformer paper, and follow-ups; models are trained on images with resolution 224×224, with inception crop followed by random horizontal flipping. Pre-training on JFT and large image-text datasets does not benefit much from heavy data augmentation as discussed above, and it isn’t common practice to use image stretching or shearing; this is also the case for other computer vision works with large datasets e.g. CLIP. We added a small section to the appendix  (Appendix B.4) with these details. \n\n5. `Where are the contrastive results? I might have overlooked something, but I can't find zero-shot imagenet or COCO image-text retrieval results as discussed on line 154.`\n\nApologies, we realise in hindsight these weren't clearly highlighted. A few of the experiments were performed in a contrastive setup, e.g. Figure 7, and the contrastive-pretrained models used to initialize the detection OWL-ViTs in Section 3.6, where zero-shot ImageNet accuracy was reported for both. We will make this more clear for the camera ready version.\n\n6. `Minor typos` many thanks for noticing these, we have corrected them in the uploaded pdf!\n\n7. `Honest limitations section`\n\nThanks for raising this. We have updated the limitations section to:\n* Discuss compute overhead, referencing figure 4 and related discussions.\n* Touch on the complexities of implementations/architectural changes.\n* Mention what we discussed above about the unexplored data-augmentation angle.\n* Mention some limitations relating to small-scale experiments.\n\nThanks again for your insightful review; we hope these have answered your questions, and that the updated manuscript reflects the addressed feedback."}, {"review_id": "WnCA1ThJwP", "reviewer": "Reviewer_CXkP", "summary": "This passage discusses the common practice of resizing images to a fixed resolution before processing them with computer vision models, which is not optimal. The author introduces a new model called NaViT (Native Resolution ViT) that takes advantage of flexible sequence-based modeling and allows for processing inputs of arbitrary resolutions and aspect ratios with adaptive positional embeddings. NaViT uses sequence packing and token drop which improves training efficiency for large-scale supervised and contrastive image-text pretraining. The author believes that NaViT represents a promising direction for ViTs and offers a departure from the standard input and modeling pipeline used by most computer vision models, which rely on CNN-designed approaches.", "questions": "- The authors have made an attempt to demonstrate the relative improvements over competitive counterparts in Figure 10. However, the differences between factorized position embeddings and NeRF appear to be marginal, with a range of less than $\\pm 0.2\\\\%$.\n- I would appreciate further clarification on the statement \"This enables variable aspect ratios, with resolutions of up to $R=P\\cdot \\text{maxLen}$\". Given that Pix2struct only introduces positional embeddings with the size of $[\\text{maxLen},\\text{maxLen}]$, and \"indexed with (x,y) coordinates of each patch\", I am curious how the resolution can be enlarged to $P\\cdot \\text{maxLen}$ without patch embedding in the size of $P\\times P$.\n- The authors have discussed several position encoding methods, but I remain unconvinced that any specific formulation significantly contributes to the performance improvements over the others.\n- Could the authors provide additional details on the settings of \"the total number of pixels can be resampled while preserving the aspect ratio\"? While I agree that \"for NaViT a resolution of “128” has the same area of a square 128 x 128 image, but could be 64 x 256, or 170 x 96\", I am curious as to why the \"effective resolution\" of $64 \\times 256$ is larger than $128 \\times 128$.\n- The authors have thoroughly discussed the benefits of the variable resolution, sampling strategies, and token dropping. However, none of these concepts are first proposed in this work. While the details may differ from existing works, such as the resolution sampling strategy, the fundamental idea remains the same. Therefore, the novelty of this paper appears to be limited.\n- The relationship between \"Pack n'Patch\" and \"Self-attention cost\" appears to be weak, as the packing of multiple patches from different images into a single sequence does not increase the cost compared to the original single sequence scheme. We kindly request the authors to provide further clarification on why the attention overhead should vary with the number of patches packed.", "limitations": "I found no potential negative societal impact of their work.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The authors present a simple method that significantly enhances the training efficiency of the vanilla ViT model, as evidenced by the results displayed in Figure 1. The observed improvements are noteworthy and suggest potential for practical application. \n- The authors also make a compelling argument, supported by Figure 3, that the conventional practice of resizing or padding images to a fixed size, which has been historically associated with convolutional neural networks, is flawed. Specifically, the authors demonstrate that both resizing and padding can lead to suboptimal performance and inefficiency, respectively. \n- While the overall algorithms are relatively straightforward, they are clearly depicted in Figure 2.", "weaknesses": "- The authors have presented a comprehensive set of experiments to demonstrate their results. However, I would like to raise a concern regarding the absence of comparison with the related works discussed in Section 4. This omission makes it challenging to evaluate the actual improvements over the baseline methods in a fair manner.\n\n- To address this issue, I strongly encourage the authors to provide a detailed discussion on the primary contributions of their proposed methods. It appears that the example packing technique has already been thoroughly discussed in Efficient Sequence Packing [1], and one could simply replace the word tokens with image patches to form the proposed method. Moreover, apart from the example packing, the main difference between \"Patch n'Pack\" and Pix2struct [2] seems to rely solely on the construction of positional embedding. Additionally, it is worth noting that the most recent work [3] also aims at mix-resolution tokenization. While the proposed method may have some differences from existing works, the current manuscript fails to clearly establish the novelty of this paper.\n\n### Reference:\n- [1] Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance. Submitted to ICLR2022\n- [2] Pix2struct: Screenshot parsing as pretraining for visual language understanding. Submitted to ICLR2023\n- [3] Vision Transformers with Mixed-Resolution Tokenization. CVPR2023w", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and comments.\n\nWe first address the weaknesses section, and in particular concerns around lack of novelty.\n\n*Weaknesses*\nSequence packing, variable aspect ratio, and token dropping are indeed not new concepts, and have independently been studied. This is not true for variable resolution; multiresolution models are common in the literature, typically for dense prediction tasks, but to our knowledge this is the first work to explore variable resolution to speed up training (we also note that your reference [3] was uploaded one month after the NeurIPS submission deadline).\n\nAs well as this, our major contributions are:\n* Introducing sequence packing to computer vision (Patch n’Pack), and demonstrating that it combines very well with multiple recent advancements in computer vision (token dropping and variable aspect ratio).\n* Introducing variable resolution as a way to speed up training.\n* Demonstrating that Patch n’ Pack enables a far richer design space for these computer vision techniques than was previously possible (e.g. sampling per-image resolutions and token dropping rates), and showing initial benefits from doing so.\n* Introducing positional embedding methods which demonstrably improve generalization to new resolutions.\n\nWe believe this will be of great interest to the academic community, change the de-facto non-packed method of training vision models, and meets standards of novelty for publication.\n\n*Questions*\n\n`...the differences between factorized position embeddings and NeRF appear to be marginal, with a range of less than 0.2%. `\n\n`...remain unconvinced that any specific formulation significantly contributes to the performance improvements over the others.`\n\nFigure 10-left presents the difference between positional embedding methods, and the gaps between different approaches are indeed small. \nHowever, this is for the evaluation in the “in-distribution- setup”, where the train and test resolutions are in a similar range. \nThe differences between different methods are more significant in out-of-distribution resolution evaluations (Figure 10 right); this is the main reason we propose the alternative approaches.  We added a paragraph to the positional embedding section to clarify this.\nThe benefits of our approach will hopefully also become clearer below in our response to the next question.\n\n`...further clarification on the statement \"This enables variable aspect ratios, with resolutions of up to R=P. Given that Pix2struct...`\n\nDuring training there is a maximum sequence length. As an example, say it was 9 tokens; with an example patch size of P = 10, this corresponds to a 30 x 30 image (or 10 x 90, or 20 x 40 with one padding token, etc).\n\nFor Pix2struct we would correspondingly `max_len = 9`. It uses a fixed grid of 2D embeddings of size `9 x 9`. During training, only images with a sequence length ≤ 9 are seen. This means only positional embeddings where $x \\times y \\leq 9$ are actually trained; the rest remain randomly initialized. At inference time, given a larger image, most combinations of $(x, y)$ will not be trained as any larger image will have some $x, y$ where $x\\times y > 9$. The pix2struct approach therefore cannot generalize to larger images; we have prepared [this diagram](https://i.imgur.com/8MWdkZK.png) to hopefully also help.\n\nOur factorized approach uses separate positional embeddings for X and Y coordinates. In the above example, each of these positional embeddings is of length `9`. All positional embeddings are used during training, which means when running inference on larger images, no untrained positional embeddings are used. It can e.g. generalize to a 40 x 40 image, as it saw 20 x 40, and 40 x 20 during training, even though a 40 x 40 image would never be seen during training as its sequence length of 16 exceeds the max length of 9. [This diagram](https://i.imgur.com/hjKnvSB.png) visualises this.\n\nWe hope this clarifies the setup!\n\n`...provide additional details on the settings of \"the total number of pixels can be resampled while preserving the aspect ratio...`\n\nBy effective resolution, we mean: $\\sqrt(\\mathtt{num\\textunderscore pixels})$.\nThe effective resolution of 64 x 256 is the same as 128 x 128, as they both have the same number of pixels (64 * 256 = 128 * 128 = 16384).\n\nWhat we discussed there is the difference in approaches to resizing. Typical computer vision pipelines would squash an image to square, and then resize to desired resolution $R$.\nWe instead resize the image such that it has (roughly) $R^2$ pixels. This has the same number of pixels, but retains the aspect ratio. It is important to have the same number of pixels, as this controls the number of patches and therefore the compute cost, and allows for comparable evaluations.\n\n`...Therefore, the novelty of this paper appears to be limited.` We hope this question is well addressed by the first response!\n\n`The relationship between \"Pack n'Patch\" and \"Self-attention cost\" appears to be weak, ...`\nSelf attention cost scales quadratically with sequence length $n$.\nWith $B$ images, each of sequence length $n$, the cost is roughly $O(Bn^2)$.\nIf instead we pack $k$ images per sequence, increasing sequence length to $kn$ but processing a smaller batch of size $\\frac{B}{k}$, the cost would be $O(\\frac{B}{k} \\times (kn)^2) = O(Bkn^2)$.\nThis extra factor in the cost is broadly speaking the concern around the self attention cost of Patch n’ Pack.  However, we note two things:\n* The self attention cost of a transformer model becomes an increasingly small proportion of overall costs as transformers scale up. This is what is shown in Figure 4.\n* We don’t actually have to increase the sequence length to benefit from NaViT. With the same max sequence length, NaViT can mix in images at smaller resolutions (less tokens than n) and still enjoy the increased throughput.\n\nWe hope we have addressed some of the concerns and the updated manuscript is to a satisfactory satndard!"}, {"review_id": "3c5Ubk0TMJ", "reviewer": "Reviewer_z6Gn", "summary": "This paper focuses on adapting the computer vision model to flexible usage. The authors stand from the ViT architecture and exploit its flexible sequence-based modeling to enable arbitrary resolutions and aspect ratios. The proposed NaViT could benefit the downstream tasks of object detection, image, and video classification. Evaluations on typical ViT tasks show the performance on different downstream tasks.", "questions": "Please see the suggestions in the weakness part.", "limitations": "None.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Exploiting the flexible sequence-based modeling of ViT models is interesting. This paper uses a simple but effective idea to make the image preprocessing match arbitrary resolutions and aspect ratios. The idea is motivated by convincing preliminary experiments. The performance evaluation also presents useful insights into utilizing NaViT’s property.", "weaknesses": "The architecture design of NaViT and its essential components to extract visual features could be introduced to help the readers better understand the techniques.\n\nAlthough the authors claim the proposed NaViT can be applied to different downstream tasks, the experiments are mainly based on high-level tasks of classification and detection. I am interested in NaViT’s generalization to more low-level tasks, such as super-resolution with arbitrary scales and pixel-level segmentation.\n\nThe additional overhead of introducing NaViT as the image preprocessing could be reported.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your review and critiques; we have made a few updates to the manuscript based on your feedback, and hope it helps address some of the mentioned weaknesses.\n\nWe will discuss them in more detail now:\n\n1. `The architecture design of NaViT and its essential components to extract visual features could be introduced to help the readers better understand the techniques. `\n\nAside from positional embeddings, the model architecture is functionally identical to the Vision Transformer. Given its widespread usage, we didn’t repeat ViT's architecture design and instead showed the differences introduced by our approach in Figure 2. We made this point more clear on the revised version of the paper and more explicitly directed readers to the Vision Transformer paper for further details on the architecture.\n\n2. `Although the authors claim the proposed NaViT can be applied to different downstream tasks, the experiments are mainly based on high-level tasks of classification and detection. I am interested in NaViT’s generalization to more low-level tasks, such as super-resolution with arbitrary scales and pixel-level segmentation.`\n\nGood point! In the meantime, we have explored semantic segmentation on ADE20k, finetuning the L/16 models presented in Figure 1. The benefits of NaViT transfer naturally to this setting; for example, NaViT finetuned at resolution 384 outperforms ViT finetuned at resolution 512, while finetuning twice as fast (e.g. see [this figure](https://i.imgur.com/hQK7DCq.png) for segmentation results). We have added this figure and discussion of the results to Section 3.6 of the updated manuscript.\n\n3. `The additional overhead of introducing NaViT as the image preprocessing could be reported`\n\nWe assume the reviewer refers to the cost of image preprocessing. During training, as is common, the data preprocessing occurs on CPU. Typical for deep learning models at this scale, preprocessing is significantly faster than the time for a single training step (models are not input bound), and therefore there we observed no training time overhead of using Patch 'n' Pack.\n\nAfter training is complete, the resulting model is architecturally the same as ViT, except that it performs better, and generalizes better to new image sizes. But there is no additional overhead of running the model compared to an equivalent sized ViT.\n\n\nWe hope these comments and the updated manuscript help address the mentioned concerns, and thank the reviewer again for their time!"}, {"review_id": "KQ3DPeIp4a", "reviewer": "Reviewer_SsJZ", "summary": "This paper proposed an efficien training technique for Vision Transformers, called Patch n’ Pack. Specifically, it packs multiple images of various input resolutions into a single sequence as a batch exmaple. Furthermore, based on the modified architecture, the authors proposed NaViT. By combining Patch n’ Pack and NaViT, the authors conducted extensive experiments on JFT-4B datasets as well as a few downstream tasks. Overall, the method achieves great effciency for pretraining, and the model achieves better accuracy for differnet image resolutions at inference time.", "questions": "It would be better for authors to add additional experiments for pretraining on ImageNet-1K and compare with FlexViT. Please also include memory consumption report during pretraining.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Adapting general-purpose Transformers into different input image resolutions is a fundamental research problem. Therefore, the technique that proposed in this paper is important to the community. And it works pretty well on both pretraining and downstream tasks.\n2. The experiments are comprehensive. The efficiency gain during pretraining is impressive.\n3. This paper is easy to follow. The overall presentation is of great quality.", "weaknesses": "1. Packing examples into a single sequence during training is not new in the literature. However, it's technically new for ViT training.\n2. Despite the performance, pretraining on JFT-4B is very expensive, making it difficult for the subsequent works to follow up and compare with. It would be better for the authors to include experiments of pretraining on ImageNet-1K.\n3. It is not clear how the memory cost would be under the proposed Patch n’ Pack.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for time spent reading the paper and your critique and comments; we're glad you appreciated the significance of the research problem, clarity of writing, and the strength of results.\n\nWe address here the mentioned weaknesses.\n\n1. `Packing examples into a single sequence during training is not new in the literature. However, it's technically new for ViT training. `\n\nWe agree, and have cited prior work accordingly. While this is a substantial technical change for Vision modeling, we however believe we contribute more still, by showing it unifies a number of approaches recently developed in computer vision literature (variable aspect ratio and token dropping), and enables new approaches which could not be explored before due to the restrictive need for constant sequence length per image.\n\nAlongside this, we have introduced variable resolution as a method for significantly speeding up training while producing more flexible models, and novel positional embedding schemas which enable generalization to larger resolutions. We think there's more here than just sequence packing, and believe as you pointed out there's a lot which would be of interest to the academic community.\n\n2a. `Despite the performance, pretraining on JFT-4B is very expensive, making it difficult for the subsequent works to follow up and compare with.`\n\nThis is a valid concern. We first note that other works demonstrating methodological improvements on JFT, such as the original Vision Transformer paper, opened the door to a plethora of research done at smaller scale (even experiments in the original Transformer paper were also at large scale in NLP). We believe the findings here are transferable to smaller datasets, open the door to additional innovation at smaller scales, and that this will therefore useful at many scales.\n\nThat being said, based on Figure 1, the largest relative improvements seen are at smaller training schedules, which is promising for smaller scale research. We also note that we showed NaViT's techniques are useful during downstream finetuning, which are a much lower compute setup.\n\n`It would be better for the authors to include experiments of pretraining on ImageNet-1K.`\nThat being said, we have been working on pre-training with public datasets, in order to develop a more reasonably reproducible setup and open-source models. Prior work [1] has shown that performant ViT models  pre-trained on ImageNet1k and ImageNet21k necessitate careful augmentation and regularisation. However, many of these techniques are very tailored to square images; for example, it’s unclear how mixup generalises to images of different sizes and shapes. It therefore needed more work than anticipated to adapt these techniques to NaViT.\n\nThat being said, we ran some experiments, and when training NaViT-B/16 on ImageNet-1k, we can already match ViT’s performance with less than half the compute budget; see [this figure](https://i.imgur.com/KN3EOly.png) for initial results. This is work in progress, which we aim to finish for open-sourcing code around the conference, but cannot promise that it will be ready for the camera ready deadline. Nonetheless, as discussed, learnings from JFT have been known to transfer well to new setups, and trigger further research at smaller scales. Thus we believe the paper in its current form is still of significant value to the wider academic community.\n\n[1] How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers, Steiner et Al, TMLR2022\n\n3. `It is not clear how the memory cost would be under the proposed Patch n’ Pack.`\nRegarding memory costs, we have an initial discussion on this in Figure 4, though we appreciate in hindsight this is a bit abstract.\n\nMemory-wise, NaViT generally compares favorably. Assuming identical architecture and batch size, memory cost for both approaches is controlled by the sequence length. For ViT, this is set by the resolution. For NaViT, this is controlled by the maximum resolution we want to support. Assuming we keep NaViT's maximum resolution equal to the equivalent ViT baseline, which we would do anyway, they have identical memory costs.\n\nConcretely, if we train ViT-B/16 at resolution 384 on ImageNet-1k, it would have sequence length 576. An equivalent NaViT-B/16 could train on resolutions ranging from 64 to 384; it would have the same sequence length of 576  (and therefore the same memory cost), but it would fit in expectation almost two times as many images. At the same memory cost, NaViT will fit more images (or conversely, it can fit the same batch size, at a smaller memory cost).\n\nWe believe we may have originally overstated the memory costs of NaViT, so we will clarify this in the paper.\n\nMany thanks again for time spent on this review; we hope this and the updated manuscript clarifies some aspects you raised."}]

Author rebuttal:
We thank the reviewers for their time and valuable comments. We were happy to hear that they found the paper “*very comprehensive*”, leading to “*impressive performance*” (WGZP), making a “*compelling argument*” that we need to go beyond fixed-size resolutions (CXkP), based on an “*interesting idea*” and backed up with “*useful insights*'' (z6Gn) that leads to “*impressive efficiency gains*” as a result of solving a “*fundamental research problem*” and, finally, that the paper is “*easy to follow*” due to its “*great presentation quality*” (SsJZ).

We have made multiple modifications based on the feedback, including clarification of text, fixing typos, adding semantic segmentation results, and adding/deepening discussions on asymptotic performance and limitations.


Meta-review: The paper proposes a method for training ViTs with heterogeneously-sized images. Reviewers were generally favorable on the paper, but had some questions about the experimental setup, relation to prior work, compute costs, and the applicability of the method to other tasks. The authors submitted detailed responses that addressed the majority of these concerns, and in the end all reviewers were in favor of accepting the paper.

The one outstanding concern is the relationship between the proposed method and prior work raised by Reviewer CXkP. While there is certainly similarity between the proposed method and prior methods, the AC is sympathetic to the arguments of the authors regarding novelty, and believes that this paper will have value to the community.

The authors are encouraged to update the camera-ready version of the paper in response to the suggestions of the reviewers.
Decision: Accept (poster)