Title: Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Venue: NeurIPS 2023
Accepted: yes
Keywords: Vision Transformer

variable aspect ratio

flexible inference

efficient training
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

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. Wetake advantage of this with NaViT (Native Resolution ViT) which uses sequence packing during training to process inputs of arbitrary resolutions and aspect ratios. Alongside flexible model usage, we demonstrate improved training efficiency for large-scale supervised and contrastive image-text pretraining. NaViT can be efficiently transferred to standard tasks such as image and video classification, object detection, and semantic segmentation and leads to improved results on robustness and fairness benchmarks. At inference time, the input resolution flexibility can be used to smoothly navigate the test-time cost-performance trade-off. We believe that NaViT marks a departure from the standard, CNN-designed, input and modelling pipeline used by most computer vision models, and represents a promising direction for ViTs.

1 Introduction

The simple, flexible and scalable nature of the Vision Transformer (ViT) [1] has rendered it an almost ubiquitous replacement to convolution based neural networks. Underpinning this model is a simple operation: splitting an image into patches, each of which is linearly projected to a token. Typically, input images are resized to a fixed square aspect ratio and then split into a fixed number of patches.

Recent works have explored alternatives to this paradigm: FlexiViT [2] supports multiple patch sizes within one architecture, enabling smooth variation of sequence length and thus compute cost. This is achieved via random sampling of a patch size at each training step and a resizing algorithm to allow the initial convolutional embedding to support multiple patch sizes. Pix2Struct [3] introduced an alternative patching approach which preserves the aspect ratio, which is particularly useful for tasks such as chart and document understanding.

We present an alternative, NaViT. Multiple patches from different images are packed in a single sequence-termed Patch n' Pack-which enables variable resolution while preserving the aspect

10 2

10 3

Pre-training TPU chips hours

60%

65%

70%

75%

80%

Pre-training linear eval

10 2

10 3

Pre-training TPU chips hours

80%

82%

84%

86%

88%

Finetuned acc vs. pretrain cost

10 3

10 4

Inference cost (ms/image)

50%

60%

70%

80%

90%

Finetuned acc vs. inference cost

Model

ViT

Version

B/32

B/16

ratio (Figure 2). This is inspired by example packing in natural language processing, where multiple examples are packed into a single sequence to accommodate efficient training on variable length inputs.

Pooling	Representations

Image	1

Image	2

Image	3

Pad

query

key

Each	token's	receptive	field	is	restricted	to	the	tokens	within	the	same	example

Self-Attention

Pad	tokens	are	masked	out

from	attention

Pad	tokens	are	masked	out

from	pooling

Pad	examples	are	masked	out

from	loss	computation

Image1										Image2

Image3																								Image1										Image2												Image3																			Image1										Image2

Image3

Inputs

Patchify

Token	drop

Data	preprocessing

Packed

Sequence

Figure 2: Example packing enables variable resolution images with preserved aspect ratio, reducing training time, improving performance and increasing flexibility. We show here the aspects of the data preprocessing and modelling that need to be modified to support Patch n' Pack. The positionwise operations in the network, such as MLPs, residual connections, and layer normalisations, do not need to be altered. Note that this diagram only showed the differences introduced by NaViT and other parts of the model architecture is identical to the Vision Transformer [1].

Wedemonstrate that: (i) Randomly sampling resolutions at training time significantly reduces training cost. (ii) NaViT results in high performance across a wide range of resolutions, enabling smooth cost-performance trade-off at inference time, and can be adapted with less cost to new tasks, (iii) Fixed batch shapes enabled by example packing lead to new research ideas, such as aspect-ratio preserving resolution-sampling, variable token dropping rates, and adaptive computation.

These observations have major practical implications. At a fixed computational budget NaViT consistently outperforms ViT. For instance, we match the performance of the top-performing ViT with 4 × less compute (Figure 1, left). We identify the substantial increase in the number of training examples processed within the allocated compute budget as the primary contributor to the improved performance over ViT (Appendix B.5) -example packing coupled with variable resolution inputs and variable token dropping enable NaViT-L/16 to process five times more images during training (Table 2). This improved efficiency extends to the fine-tuning process (Figure 1, middle). Furthermore, by exposing NaViT to multiple resolutions during both pre-training and fine-tuning, a single model demonstrates excellent performance when evaluated on various resolutions, significantly advantaging NaViT in terms of inference cost (Figure 1, right).

NaViT's training and adaptation efficiency, and flexible inference, presents a promising avenue for Vision Transformers. Patch n' Pack empowers computer vision systems to transcend limitations imposed by current data and modeling pipelines, enabling ideas that were previously restricted by the constraints of fixed batch shapes, unlocking new possibilities for innovation and advancement.

h

=2

w

h

=

w

w

=2

h

0

1

2

3

4

5

Density

ImageNet (85.9% non-square)

h

=2

w

h

=

w

w

=2

h

Aspect ratio

LVIS (92.2% non-square)

h

=2

w

h

=

w

w

=2

h

WebLI (57.3% non-square)

2 Method

Deepneuralnetworksaretypically trained and run with batches of inputs. For efficient processing on the current hardware this implies fixed batch shapes which in turn imply fixed image sizes for computer vision applications. This coupled with architectural limitations historically associated with convolutional neural networks led to a practice of either resizing or padding images to a fixed size. Both of these have been shown to be flawed: the former harms performance and the latter is inefficient [3]. An analysis of aspect ratios in ImageNet [4], LVIS [5] and WebLI [6] as representative examples of classification, detection and web image datasets, respectively, shows that most images are typivally not square (Figure 3).

In language modelling, it is common to bypass limitations of fixed sequence lengths via example packing : tokens from multiple distinct examples are combined in one sequence, which can significantly accelerate training of language models [7]. By treating images as sequences of patches (tokens), we show that Vision Transformers [1] can benefit from the same paradigm, which we call Patch n' Pack. Using this technique ViTs can be trained on images at their 'native' resolution, and we name this approach NaViT.

2.1 Architectural changes

NaViT is built upon the original ViT, but in principle can use any ViT variant operating on a sequence of patches. To enable Patch n' Pack, we make the following architectural modifications.

Masked self attention and masked pooling. To prevent examples attending to each other, additional self-attention masks are introduced. Similarly, masked pooling on top of encoder aims to pool the token representations within each example, resulting in a single vector representation per example in the sequence. Figure 2 presents how the receptive filed of attention is controlled via masking.

Factorized & fractional positional embeddings. To handle arbitrary resolutions and aspect ratios, we revisit the position embeddings. Given square images of resolution R × R , a vanilla ViT with patch size P learns 1-D positional embeddings of length ( R / P ) 2 [1]. Linearly interpolating these embeddings is necessary to train or evaluate at higher resolution R .

Pix2struct [3] introduces learned 2D absolute positional embeddings, whereby positional embeddings of size [maxLen , maxLen] are learned, and indexed with ( x,y ) coordinates of each patch. This enables variable aspect ratios, with resolutions of up to R = P · maxLen . However, every combination of ( x,y ) coordinates must be seen during training.

To support variable aspect ratios and readily extrapolate to unseen resolutions, we introduce factorized positional embeddings, where we decompose into separate embeddings φ x and φ y of x and y coordinates. These are then summed together (alternative combination strategies explored in Section 3.4). We consider two schema: absolute embeddings, where φ ( p ) : [0 , maxLen] → R D is a function of the absolute patch index, and fractional embeddings, where φ ( r ):[0 , 1] → R D is a 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 learned embeddings φ , sinusoidal embeddings, and the learned Fourier positional embedding used by NeRF [8].

2.2 Training changes

Patch n' pack enables new techniques to be used during training of NaViT.

Continuous Token dropping. Token dropping (random omission of input patches during training) [9, 10] has been developed to accelerate training. However, typically the same proportion of tokens are dropped from all examples; packing enables continuous token dropping, whereby the token dropping rate can be varied per-image. This enables the benefits of faster throughput enabled by dropping while still seeing some complete images, reducing the train/inference discrepancy. Further, with packing, the

drop-distribution can vary throughout training, following some pre-defined schedule. In Section 3.3, we explore different schedules and the benefits of flexible token dropping.

Resolution sampling. NaViT can be trained using the original resolution of each image. Alternatively, the total number of pixels can be resampled while preserving aspect ratio. In vanilla ViT, there is a tension between greater throughput (training on smaller images), and greater performance (training on larger images, to enable high-resolution at evaluation time). Oftentimes, models are pre-trained at a smaller resolution and finetuned at a higher one [11]. NaViT is much more flexible; it allows mixedresolution training by sampling from a distribution of image sizes, while retaining each images' original aspect ratio. This allows both higher throughput and exposure to large images, yielding substantial improved performance over equivalent ViTs (in terms of models size and training duration). Section 3.2 explores different sampling strategies, and variable resolution training for pre-training and finetuning.

2.3 Efficiency of NaViT

Here we discuss some implications of Patch n' Pack on the computational efficiency of NaViT.

Self attention cost. The O ( n 2 ) cost of attention is a natural concern when packing multiple images into longer sequences. Though many works aim to remove this quadratic scaling [12, 13], we demonstrate here that as the transformer hidden dimension is scaled, the attention becomes an increasingly smaller proportion of the the overall cost, which encompasses the computation cost of the MLP as well. Figure 4 illustrates this trend, indicating a corresponding reduction in the overhead associated with packing examples. In addition to speed considerations, the memory cost of self-attention can pose a challenge for extremely long sequences. However, this challenge can be also addressed by employing memory-efficient methods [14, 15].

1024

2048

4096

6144

0%

10%

20%

30%

40%

50%

Attention overhead

1 packed

2 packed

4 packed

8 packed

Model dimension

Packing, and sequence-level padding. The final sequence lengths containing multiple examples must be fixed. We use a greedy packing approach discussed in Appendix B.3; there typically is no perfect combination of examples exactly adding up to the fixed length and padding tokens have to be used. One could for example dynamically choose the resolution or token dropping rate of the final example in a sequence to exactly fit the remaining tokens; however, we find typically less 2% of tokens are padding tokens, and thus the simple approach is sufficient.

Padding examples and the contrastive loss. Per-token losses are straightforward to implement with packed sequences. However, many computer vision models are trained with example-level losses, typically applied to a pooled representation. First, this requires modifications to the typical pooling heads to account for packing. Second, multiple pooled representations must be extracted from each sequence. Fixed batch shapes requires an assumption that, from a batch of B sequences, we extract at most B × E max pooled representations (i.e. E max examples per sequence). If a sequence contains more than E max images, the extra images will be dropped, wasting computation of the model's encoder. If a sequence has less than E max examples, then the loss will process lots of fake padding representations.

The latter is an issue for contrastive learning, where loss computation scales in time and memory ∼O ( n 2 ) . To avoid this, we used the chunked contrastive loss [16], which circumvents the need to gather all data points for the softmax by performing computations on local device subsets and efficiently accumulating the necessary statistics for global softmax normalization. This enable high values of E max (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 [1], with the changes to enable packing, described in Section 2.1. In addition, we include small ViT improvements from previous works: query-key normalization and the omission of biases [17], and attention pooling [18].

We pre-train NaViT in two setups: classification training on JFT-4B [18] and contrastive languageimage training [19] on WebLI [6]. Typically, for JFT, inception crop is applied pre-training [20, 1], and in both cases, images are resized to a square (distorting aspect ratio). Unless otherwise specified,

all NaViT models are pre-trained without these operations, and preserve aspect ratio. NaViT is implemented in JAX [21] using the FLAX library [22] and built within Scenic [23].

Classification pretraining. We pre-train NaViT with supervised classification objective, using a sigmoid cross-entropy loss, following the setup of [17] on JFT-4B [18]. Visual representations are evaluated following the linear evaluation protocol used for ViT [1], where 10 examples per class are used to train a linear classifier on top of frozen representations.

Contrastive pre-training. Alongside the image model, we train a text encoder with the same architectural modifications using the contrastive image-text loss [19, 24] (details in Appendix B.2). Packing also provides efficiency improvements on the text-tower, as text sequences do not need to be padded to a fixed lengths, which is the normal setup. The contrastive models are evaluated on zero-shot ImageNet classification and COCO image-text retrieval.

Note that throughout all the pre-training experiments and the transfer to downstream experiments, both the NaViT and baseline models are trained with compute-matched setup. This implies that during pretraining, the same amount of TPU time is used and for downstream tasks, the models are evaluated at the same effective resolution (i.e., the same number of tokens). To be precise, all downstream experiments utilized the top-rightmost points in Figure 1 (ViT-L/16 and NaViT-L/16).

3.1 Improved training efficiency and performance

Figure 1 illustrates the JFT pretraining performance of different NaViT models compared to computematched ViT baselines [25]. The experimental setup details are provided in Appendix B.1. NaViT consistently surpasses ViT in performance while using the same computational budget across different compute and parameter scales; for example, the performance of the top-performing ViT can be matched by a NaViT with four times less compute. Conversely, the computationally lightest NaViT in Figure 1 is five times more cost-effective than its equivalent ViT counterpart.

TheNaViTmodelsbenefitfrompreservedaspectratios and the ability to evaluate over many resolutions, but the chief contributor here is the significant increase in the number of training examples processed by NaViT within the allocated compute budget. This is achieved through the combination of sampling multiple variable-resolution examples and token dropping, leading to variable size images that are efficiently packed into a similar sequence length as the original model. We ablate these factors below.

3.2 Benefits of variable resolution

Here, we deep-dive the benefits of mixed-resolution training. Since we preserve the native aspect ratio, when we refer to 'resolution' for NaViT, we mean 'effective resolution'. That is, images with the same area as a square image with a given resolution. For example, for NaViT a resolution of '128' has the same area of a square 128 x 128 image, but could be 64 x 256, or 170 x 96, etc., and thus has the same inference cost as regular ViT on 128 x 128 images.

Variable-resolution pre-training. Lower resolution images require fewer FLOPs to process and hence small resolutions (like 224) are used with fixed-resolution training. With fixed-resolution training, there is a trade-off between throughput and ability to process details and high-resolution images. With NaViT we can mix lower resolution images with large ones to get the best of both worlds.

Figure 5 shows a comparison between two NaViT variants trained at several different resolutions. Here, all trained for the same number of FLOPs . (1) Native aspect ratio, but fixed resolution R = R max for different chosen values of R max . (2) Variable resolution, where the resolution is distributed as R ∼U (64 ,R max ). Variable resolution models outperform models trained at only that resolution. Even in the best case for fixed resolution, where the train and evaluation resolutions are identical, variable resolution matches or outperforms fixed.

Variable-resolution finetuning. Prior works increase resolution late in pre-training or during finetuning, producing higher quality but more expensive models [1, 11]. We finetune NaViT and ViT models at different fixed resolutions, and additionally NaViT at variable resolutions. Figure 6 shows the results of fine-tuning pretrained ViT and NaViT on ImageNet-1k dataset. Performance gains during pretraining transfer well at all resolutions, but two phenomena are particularly interesting: First, NaViT finetuned with variable resolutions ("NaViT 64:512") is as good as a NaViT finetuned at a single resolution (and much better than single-resolution ViT), removing the need to pick a single downstream finetuning resolution. Second, NaViT finetuned at low resolution (64) still obtains good performance when

128

256

384

512

Max train resolution

R

76%

77%

78%

79%

Top-1 Accuracy

Best ImageNet 10shot

Variable res

U

[64,

R

]

Fixed res =

R

128

256

384

512

40%

60%

80%

native res eval

128

256

384

512

20%

40%

60%

128 res eval

128

256

384

512

65%

70%

75%

224 res eval

128

256

384

512

R

70%

75%

256 res eval

128

256

384

512

R

40%

60%

80%

384 res eval

128

256

384

512

R

20%

40%

60%

80%

512 res eval

64 128

256

384

512

Fine-tune resolution

R

55%

60%

65%

70%

75%

80%

85%

ImageNet Top-1 Accuracy

Eval at

R

ft

NaViT (one per res)

ViT (one per res)

NaViT-64:512

64 128

256

384

512

Fine-tune resolution

R

Best eval across

R

[64, 512]

NaViT (one per res)

ViT (one per res)

NaViT-64:512

evaluated at higher resolutions (Figure 6, right), enabling cheaper adaptation.This ability to perform cheap adaptation of a flexible pre-trained model corroborates findings in [2].

Resolution sampling strategies. Packing examples enables diverse resolution sampling strategies. We first consider whether to sample the target side length (average height/width, R ), or the target area (i.e. sequence length ∝ R 2 ). Sampling the side length from a uniform distribution biases towards lower sequence lengths, whereas sampling the area from a uniform distribution biases towards higher side lengths.

For each image, we sample u ∼D , where D is a distribution with support [ -1 , 1] . Werescale u linearly to [64 , 384] for sampling side-lengths or [64 2 , 384 2 ] for sampling areas. We consider four distributions D : uniform u ∼U ( -1 , 1) , truncated (to [ -1 , 1] ) standard Normal u ∼ N t (0 , 1) , and then two other Normals which bias towards lower resolution u ∼N t ( -0 . 5 , 1) and higher resolution u ∼N t (0 . 5 , 1) . The results are shown in Figure 7. Here, the best resolution resampling strategy consistently performs over the default resolution. It is consistently better to

55%

58%

60%

62%

65%

68%

ImageNet zero-shot Accuracy

normal

normal

=

0.5

normal

=0.5

uniform

distribution

63.3%

62.6%

63.4%

63.8%

63.6%

63.1%

64.3%

63.4%

Sample side length

Sample area

Native res

sample side-lengths (orange) directly as opposed to area (blue), and the best distribution is the truncated normal biasing towards lower values; both of these increase throughput by preferentially sampling smaller sequences.

3.3 Benefits of variable token dropping

Token dropping strategies.

We experimented with continuously sampled token dropping rates, and with resolution-dependent token dropping rates; both are explained in Appendix B.6.

Fig. 9a compares variable drop rates sampled from a Beta distribution to a constant drop rate, demonstrating consistent improvements from the former. Fig. 9b shows the use of a resolution dependent token dropping rate for models trained with R ∼ U (64 , 384) and dropping rates scaled between [0 . 5 -δ, 0 . 5 + δ ] ∝ R , which further improves over the beta distribution.

70.0%

80.0%

Top-1 accuracy

Fixed

Increasing

Decreasing (fast)

Decreasing (slow)

75.8%

72.2%

76.3%

76.5%

0.25

0.50

0.75

Mean token drop rate

76%

78%

80%

Top-1 Accuracy

no token drop

fixed for every image

beta:

per image

fixed

beta

=0.25

=0.4

Token drop strategy

78.0%

78.5%

79.0%

Top-1 accuracy

Resolution-dependent

Scheduled token dropping rates. Packing also enables easy

variation the token dropping rate during training. By changing the token dropping rate we can better tune the trade-off between number of images seen and information used per image, to maximize the final accuracy while keeping the total training cost constant. We varied the token dropping rate as a function of the number of images seen (details of the schedule in Appendix B.7). Fig. 8 demonstrates that further improvements are possible by reducing the token dropping rate during JFT pretraining of NaViT-B/16.

3.4 Positional embeddings

Weevaluate our factorized embeddings introduced in Section 2.1, and their design choices. We are interested in both absolute performance, and extrapolation to resolutions outside the training regime. To test this, we train NaViT-B/16 models for 200k steps on JFT, with resolutions R ∼U (160 , 352) . We evaluate performance at a range of resolutions, without modification of the embedding variables. We compare to a ViT-B/16 trained at fixed resolution 256 for the same amount of images seen, evaluated at new resolutions using standard interpolation of positional embeddings.

Figure 10a illustrates the disparity among different positional embedding methods when the model is evaluated under the 'in-distribution" setup, where the resolutions of training and testing data are within a comparable range. Conversely, Figure 10b shows the model performance across out-of-distribution resolutions, revealing notable distinctions between the approaches.

First, it is clear that the factorized approaches outperform both the baseline ViT and the Learned 2D embeddings from Pix2struct. The latter in particular struggles to generalize to higher resolution, likely because this requires an increasingly long tail of unseen ( x,y ) pairs. Factorized embeddings are best combined additively (as opposed to stacking or multiplying).

50%

55%

60%

65%

70%

75%

ImageNet 10shot accuracy

Learned 1D (ViT)

Learned 2D (Pix2struct)

Factorized (+)

Factorized (stack)

Factorized (×)

Sinusoidal

Fourier

Factorized (+)

Sinusoidal

Fourier

64.4%

71.4%

64.7%

71.3%

64.9%

71.5%

64.8%

71.5%

64.2%

71.3%

53.8%

69.3%

63.5%

71.3%

63.9%

70.8%

64.8%

71.4%

64.2%

71.5%

64

160

256

352

512

640

Evaluation resolution

0%

20%

40%

60%

80%

100%

Relative accuracy

Train

resolutions

ViT

Learned 2D (pix2struct)

Factorized (+)

Fourier (fractional)

(a)

(b)

3.5 Other aspects of NaViT's performance

Out of distribution generalization. We directly evaluate JFT-pretrained NaViT on downstream datasets, employing a label-map [26] from JFT-4B to ImageNet [4] and robustness-variants (ObjectNet [27] and ImageNet-A [28]). We compare the performance to a compute-matched ViT baseline.

10 2

10 3

Training TPU chip hours

55%

60%

65%

70%

75%

ImageNet

10 2

10 3

Training TPU chip hours

20%

30%

40%

50%

60%

70%

ImageNet-A

10 2

10 3

Training TPU chip hours

35%

40%

45%

50%

55%

ObjectNet

ViT

B/32

B/16

Figure 11 shows that NaViT compares favorably both on ImageNet as well as datasets variants that were specifically designed to test out of distribution performance. It is interesting to note that NaViT performs much better on ImageNet-A, but ViT catches up on ObjectNet, even though both these datasets contain images that have extreme aspect ratios. We believe this is due to the aspect-preserving center crop that we apply to images for ImageNet-A and ObjectNet classification (same as in [17]), which is a useful prior for ObjectNet, but less so for ImageNet-A (see Appendix G for details). If no crop is applied to images and they're instead simply resized to the image resolution expected by the model (i.e., square for ViT, and aspect preserving resize to the same number of tokens for NaViT), then the observed difference is much larger (see Figure 21 in appendix).

Calibration. In addition to the in-depth accuracy analysis, we have also quantfied the quality of the uncertainty computed by the model. In particular, for our ImageNet1K-finetuned models, we computed the expected calibration error [29] of the top prediction, as we vary the number of patches we assign per examples. We find that the calibration error remains very stable in the interval (0 . 045 , 0 . 047) as we vary the number of patches per image in the range [128 , 1024] , without any post-hoc recalibration. We provide further details in the appendix Appendix F.

Inference trade-offs. Given the flexibility of the model, there are several viable approaches how one can maximize the aggregate accuracy under a given compute budget, which we quantify in terms of the latency measured on a Cloud TPUv3 chip. In an online inference setting, choosing an inference strategy translates to an algorithm how to allocate a fixed number of tokens per example. As we show in Fig. 1, NaViT offers much better trade-offs than ViT and it shows strong diminishing returns, s.t., even relatively few patches provide highly competitive results. In Appendix D we further study a cascading approach, which both provides Pareto optimal models and gives more precise trade-off opportunities.

Fairness signal annotation. We investigate annotating images with fairness-related signals, such as those pertaining to gender and ethnicity. Prior research has shown that metrics, such as group calibration, are susceptible to labeling inaccuracy, particularly for underrepresented groups. In addition, this problem persists even when accounting for label noise during training [30]. Thus, reducing the labeling error of fairness signals improves the reliability of bias mitigation and post-hoc auditing [31]. To explore whether NaViT can help in overcoming these challenges, we train annotators on FairFace [32] and CelebA [33] datasets as linear probes (i.e. using frozen features produced by NaViT or ViT), before comparing their accuracy.

First, NaViT provides representations of higher quality that improve the accuracy of fairness signal annotation, even when dealing with square images. Original images are of size 448 × 448 in FairFace and 178 × 218 in CelebA, and we resize them to area 224 2 while preserving aspect ratios in NaViT. Despite having the same sequence length, NaViT provides a higher prediction accuracy than ViT, as shown in Figure 12 (left). We verify statistical significance using the Wilcoxon signed-rank test test [34], which reveals that the improvement in NaViT is significant with p =3 × 10 -4 .

accuracy

100

250

500

1000

Pretraining Steps (K)

93

95

97

fairface / gender

100

250

500

1000

Pretraining Steps (K)

64

68

72

accuracy

fairface / ethnicity

Model

ViT

100

250

500

1000

Pretraining Steps (K)

95

96

97

accuracy

fairface / gender

100

250

500

1000

Pretraining Steps (K)

64

68

72

accuracy

fairface / ethnicity

Square

Native

Second, we apply inception-style cropping [35] with a fixed minimum area of 50%. This changes the aspect ratio but maintains native resolution. Similar to before, we resize all cropped images to area 224 × 224 , either as square images or with native aspect ratios. Figure 12 (right) shows that native aspect ratios in NaViT improve accuracy, which is statistically significant at the 95% confidence level ( p =0 . 02 ). Appendix H contains the full set of figures for other attributes in CelebA and Fairface.

3.6 Other downstream tasks

Semantic segmentation.

We finetune NaViT to semantic segmentation on ADE20k dataset [36], following the linear decoder protocol of Segmenter [37]. We use ViT-L/16 as baseline and compare its performance with that of NaViT-L/16. Both models are pretrained on JFT-4B [18] with comparable compute budget.

We experiment with different maximum resolution R max at finetuning time: ViT takes in random square crops of resolution R max × R max while NaViT's inputs are randomly resized (with preserved aspect ratio) so that the total number of pixels is R 2 max . This way, we have the same finetuning cost for the ViT and NaViT models. Note that following common practice, in order not to alter the ground truth segmentation maps, both models are evaluated at the native resolution [37, 38, 36]. This means that for ViT, the square predictions are resized from square back to native resolution. We observe in Figure 13 that NaViT outperforms ViT when transferred to semantic segmentation with the same maximum finetuning resolution R max . Note that NaViT at R 384 beats ViT at R 512 while being twice as fast (see Appendix D). An advantage of NaViT over ViT is that it benefits from flexibility

256

384

512

800

44

47

50

53

ADE20k mIoU

Evaluated on native resolutions

ViT (square

R

2

)

R

2

)

of resolution during training and can ingest seamlessly square and non-square images.

Object detection.

Native-resolution training may be especially beneficial for fine-grained tasks such as object detection, which require image understanding at many different spatial scales. We use compute matched NaViT and ViT models as backbones for OWL-ViTL/14 object detectors [39], following the OWL-ViT protocol for training and evaluation. Results are presented in Table 1. The NaViT-based detector

performs significantly better on both common and unseen LVIS 'rare' classes. These experiments used shorter pre-training than the original OWL-ViT and thus reach lower absolute performance; the relative difference nonetheless suggests that NaViT produces strong representations for fine-grained vision tasks.

Video Classification. Video processing with transformers is inherently challenging due to the heterogeneity of the underlying spatio-temporal signals. In practice, cumbersome protocols have to be

established both for training and evaluation (e.g. spatial and temporal multi-cropping) [40]. NaViT alleviates some of these challenges by not only allowing training over different resolutions, but also over different temporal durations. We fine-tuned NaViT trained on JFT for Kinetics400 [41] classification by extracting three different spatio-temporal patches 'tubelets''. [42], extending the positional embedding to include the temporal dimension and initializing the embedding kernel using 'central frame embedding' [40]. We posit that NaViT is an excellent starting point for multi-scale training due to the resolution diversity at training time. We observe that NaViT-L achieves competitive performance with ViViT-L (80.4%) in approximately 6x less epochs, without multi-crop evaluation. Note that the Kinetics400 dataset used here contains less data than prior works [40].

4 Related work

Flexible Vision Transformers. FlexiViT [2] developed a novel kernel resizing approach which enables variable "resolution" via models which support multiple patch sizes. This is viewed as unifying multiple distinct models, and they study the relationship with distillation and neural architecture search. Pix2struct [3] supported variable aspect ratios with a novel positional embedding schema, and demonstrated significant efficiency and performance improvements for non-natural imagery such as chart and document understanding.

Multiscale Vision Transformers. Using feature maps at multiple spatial scales is a common approach to localization tasks such as segmentation and detection [43]. Many works developed Vision Transformers which do the same [44, 45], though some dispute the necessity for simple localization tasks [46]. NaViT does not build hierarchical representations using multiple scales; we believe combining the unique flexibility of our approach with the benefits of this modelling family is a promising avenue.

Accelerating training with mixed resolutions. Imagemodellingworksconsideringresolution largely focus on accelerating pretraining with a fixed, low resolution [47, 48]. FixRes [11] is a popular technique whereby resolution is increased in a final stage of pretraining, and has been used in many subsequent works [19, 6]. PaLI [6], for example, successively increases the resolution of the vision backbone during its generative training. This approach's main downside is its irreversibility: compute cannot be scaled back by reducing resolution after tuning at the higher and more expensive resolution.

Multigrid training [49] is the most closely related work. The main idea is accelerate video modelling by processing large batches with "coarse" spatiotemporal resolution early in the training, followed by a "finer" resolution later. To this end, the authors apply a hierarchical grid sampling schedule coupled with appropriate learning rate scaling. In contrast, Patch n' Pack enables effortless incorporation of mixed resolutions without complex schedules or training pipelines.

Token dropping for improved efficiency. Research initially explored random token dropping [9, 10, 50]. Follow ups demonstrated benefits from considering structured [51] or 'importance' based strategies [52, 53]. Better strategies will likely further boost performance, and Patch n' Pack sidesteps the fixed minibatch shape restriction which limited these works.

5 Conclusions and future work

We have demonstrated that Patch n' Pack-the simple application of sequence packing to vision transformers-significantly improves training efficiency. The resultant NaViT models can be applied to many resolutions at inference time, and cheaply adapted to new tasks. We discuss limitations and impacts in Appendix A, but overall Patch n' Pack enables a wide variety of research previously hindered by the need for fixed batch shapes, including adaptive computation and new algorithms for improving training and inference efficiency.

References

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR , 2021.

Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, and Filip Pavetic. Flexivit: One model for all patch sizes. In CVPR , 2023.

Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, and Kristina Toutanova. Pix2struct: Screenshot parsing as pretraining for visual language understanding, 2022.

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.

AgrimGupta, 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.

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-image model. arXiv preprint arXiv:2209.06794 , 2022.

Mario Michael Krell, Matej Kosec, Sergio P Perez, and Andrew William Fitzgibbon. Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance, 2021.

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 Systems 2020, NeurIPS 2020 , 2020.

Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang, Shih-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-24221, 2021.

Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-image pre-training via masking, 2023.

Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou. Fixing the train-test resolution discrepancy. 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 on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , 2019.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys , 55(6):1-28, 2022.

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:2011.04006 , 2020.

Markus N Rabe and Charles Staats. Self-attention does not need O ( n 2 ) memory. arXiv preprint arXiv:2112.05682 , 2021.

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memoryefficient exact attention with io-awareness. In NeurIPS , 2022.

Basil Mustafa, Josip Djolonga, and Mostafa Dehghani. On efficient losses for distributed contrastive learning. arXiv preprint , 2023.

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 Puigcerver, Utku Evci, 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´ c, Dustin Tran, Thomas Kipf, Mario Luˇ ci´ c, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning , 2023.

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR , pages 12104-12113, 2022.

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 ICML , pages 8748-8763, 2021.

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.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020.

Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab, Matthias Minderer, and Yi Tay. Scenic: A jax library for computer vision research and beyond. In CVPR , 2022.

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.

Mostafa Dehghani, Anurag Arnab, Lucas Beyer, Ashish Vaswani, and Yi Tay. The efficiency misnomer. arXiv preprint arXiv:2110.12894 , 2021.

Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR2022, New Orleans, LA, USA, June 18-24, 2022 , pages 7949-7961. IEEE, 2022.

Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. ObjectNet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. In NeurIPS , pages 9448-9458, 2019.

DanHendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In CVPR , pages 15262-15271, 2021.

Jeremy Nixon, Michael W Dusenberry, Linchuan Zhang, Ghassen Jerfel, and Dustin Tran. Measuring calibration in deep learning. In CVPRWorkshops , 2019.

Julius Adebayo, Melissa Hall, Bowen Yu, and Bobbie Chern. Quantifying and mitigating the impact of label errors on model disparity metrics. In ICLR , 2023.

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 AAAI/ACM Conference on AI, Ethics, and Society , pages 429-435, 2019.

KimmoKärkkäinen and Jungseock Joo. Fairface: Face attribute dataset for balanced race, gender, and age, 2019.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In ICCV , 2015.

Frank Wilcoxon. Individual comparisons by ranking methods . Springer, 1992.

Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In CVPR , 2016.

Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In CVPR , 2017.

Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In ICCV , 2021.

Mathilde Caron, Neil Houlsby, and Cordelia Schmid. Location-aware self-supervised transformers for semantic segmentation. arXiv preprint arXiv:2212.02400 , 2022.

Matthias Minderer, Alexey Gritsenko, Austin Stone, Maxim Neumann, Dirk Weissenborn, Alexey Dosovitskiy, Aravindh Mahendran, Anurag Arnab, Mostafa Dehghani, Zhuoran Shen, et al. Simple open-vocabulary object detection with vision transformers. arXiv preprint arXiv:2205.06230 , 2022.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. ViViT: A video vision transformer. In CVPR , 2021.

Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint arXiv:1705.06950 , 2017.

AJ Piergiovanni, Weicheng Kuo, and Anelia Angelova. Rethinking video vits: Sparse video tubes for joint image and video learning, 2022.

Olaf Ronneberger, 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 Computer-Assisted Intervention - MICCAI , 2015.

Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV , 2021.

Yanghao Li, Chao-Yuan Wu, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik, and Christoph Feichtenhofer. Mvitv2: Improved multiscale vision transformers for classification and detection. In CVPR , 2022.

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.

Hugo Touvron, Matthieu Cord, and Hervé Jégou. DeiT III: Revenge of the ViT. In ECCV , 2022.

Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin Transformer V2: Scaling Up Capacity and Resolution. CVPR , 2022.

Chao-Yuan Wu, Ross B. Girshick, Kaiming He, Christoph Feichtenhofer, and Philipp Krähenbühl. A multigrid method for efficiently training video models. In CVPR , pages 150-159. Computer Vision Foundation / IEEE, 2020.

KaimingHe,XinleiChen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In CVPR , pages 15979-15988. IEEE, 2022.

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.

Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. CoRR , abs/2210.09461, 2022.

HongxuYin, Arash Vahdat, Jose M. Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In CVPR , 2022.

Samira Abnar, Mostafa Dehghani, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training. arXiv preprint arXiv:2110.02095 , 2021.

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 a unified text-to-text transformer. arXiv preprint arXiv:1910.10683 , 2019.

Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In EMNLP , pages 66-71, November 2018.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Foad Hamidi, Morgan Klaus Scheuerman, and Stacy M Branham. Gender recognition or gender reductionism? the social implications of embedded gender recognition systems. In Proceedings of the 2018 chi conference on human factors in computing systems , pages 1-13, 2018.

OsKeyes. The misgendering machines: Trans/hci implications of automatic gender recognition. Proceedings of the ACM on human-computer interaction , 2(CSCW):1-22, 2018.

Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. In NeurIPS , pages 23885-23899, 2021.

A Broader impacts & limitations

Broader Impacts

NaViT enables training of vision transformers on variable size inputs, which has a profound impact on advancing adaptive computation research. By training models to handle various input size, we can explore adaptive computation techniques that dynamically adjust the computational resources based on the specific requirements of a given input. This flexibility opens up new avenues for implementing ideas that aim at adjusting allocation of compute and improving efficiency in vision tasks per input. Furthermore, NaViT computational efficiency unlocks the potential for scaling up pre-training of vision models. With the ability to handle different resolutions, models can effectively tackle more complex and diverse visual data, allowing for the development of larger and more powerful vision models.

Limitations

There is a wide range of applications that could benefit from a model capable of processing inputs of different resolutions, including OCR or document understanding with the use of vision models. Although we highlight the considerable advantages of employing NaViT and provide a comprehensive analysis of common computer vision tasks, we did not specifically investigate the benefits of NaViT in these particular applications. We consider this area as a priority for future research and follow-up work.

Implementation complexities The modelling changes shown in Figure 1 can complicate ViT; in particular, packing-aware changes to the loss function may be complex for more complicated usecases.

Computational overhead Done carelessly, significant extra cost can be introduced due to sequence packing (Discussed in Figure 4). NaViT can however be run with the same sequence length as a corresponding ViT model (e.g. ViT at resolution 384 with sequence length 576 vs NaViT at resolutions 64:384 with sequence length 576; the latter will fit on average ∼ 2 × as many images).

Data augmentation aspects were not explored in this work; firstly, it is possible that token dropping and variable resolution can also act as data augmentations, which was not exhaustively studied here. Secondly, common data augmentations such as mixup and various random cropping schema assume square, constant resolution images, and analogues supporting variable aspect ratio and size were not developed here.

Scale ; experiments were performed in a large-dataset regime, and future work should validate these results at smaller scale. This however necessitates research into data augmentation techniques as mentioned above.

B Training details

B.1 Classification pretraining

The experiments and ablations in the paper are with ViT-B/32, ViT-B/16, and ViT-L/16. We use a reciprocal square-root learning rate schedule, with linear warmup and cooldown, with a maximum value of 8 e -4 , and phases. We follow [18, 54] and use a higher weight decay of 3 . 0 on the head compared to the body's weight decay of 0 . 03 during upstream training to improve transfer to downstream tasks. For our experiments, we evaluated both NaViT and ViT models using configurations B/32, B/16, and L/16. Each ViT model was trained with varying compute budgets, with cooling down at different stages of training. We trained a corresponding NaViT model for each ViT size and computational budget, allowing us to perform 'compute-matched' comparisons [25].

Table 2 presents the pretraining specifications for both ViT and NaViT models. During the pretraining phase, ViT models were trained using images of size 224 × 224 . In contrast, NaViT models uniformly sampled a value, denoted as r , between 64 and 256 and resized the image to have a total of r 2 pixels while preserving the aspect ratio. Although training NaViT models on native resolutions is possible, we empirically discovered that sampling a resolution provides greater control over maximizing the number of examples observed during pretaining within a fixed computational budget while maintaining performance across different resolutions in downstream tasks. Additionally, by controlling the resolution, we can ensure efficient packing by tuning the sequence length and limit padding to less than 2%.

B.2 Contrastive pretraining

Weuse the 32000 token T5 [55] sentencepiece [56] tokenizer. By default, text sequences are truncated to a maximum length of 24. No token dropping is used for text. Models are trained under the same optimization regime as the classification models, but with a learning rate of 3 × 10 -3 . Weight decay of 1 × 10 -6 is applied consistently to all kernels in the model (no change for projection heads). By default, image side-lengths are sampled ∼U (64 ,R max ) , and no other image augmentations are applied.

B.3 Packing algorithm

Packing of examples into sequences is done alongside batching. A simple greedy approach is used which adds examples to the first sequence with enough remaining space. Once no more examples can fit, sequences are filled with padding tokens, yielding the fixed sequence lengths needed for batched operations. Such simple packing algorithm can lead to a significant padding, depending on the distribution of length of inputs. There are several methods to address such limitations, like bin packing [7], which allows minimizing the padding. Here, in NaViT, since controlling the resolutions we sample, we can ensure efficient packing by tuning the sequence length and limit padding to less than 2%.

B.4 Pre-processing and augmentation

For both classification and contrastive pre-training, for the baseline models, we follow the original Vision Transformer paper [1], and follow-ups; models are trained on images with resolution 224 × 224 , with inception crop followed by random horizontal flipping. For NaViT, we have no preprocessing besides resizing to the sampled resolution and token dropping.

10 2

10 3

Training TPU chips hours

60%

65%

70%

75%

80%

ImageNet1K-10shot

10 11

10 12

Tokens seen

10 9

10 10

Images seen

Model

ViT

Version

B/32

B/16

B.5 Sample efficiency with NaViT

Based on the experiments we ran, we believe a lot of the observed performance benefits are due to speeding up training (i.e., seeing more images in a given compute budget). This is evident in Figure 1, where relative performance gains diminish with longer training. We can also show this by plotting the number of images seen on the x-axis:

If compute cost was not a concern, for a given number of images seen, it is best to train on purely large resolution, with no token dropping. Asymptotically we therefore expect the token dropping and mixed resolutions to eventually hurt NaViT in unlimited compute regime. This assumes a large enough dataset that overfitting is not an issue. Eventually one may enter a regime where the token dropping and variable resolution is a useful data augmentation to prevent overfitting, which may alter results.

B.6 Sampling token dropping rates

Sampling with a beta distribution We use a parameterisation based on the mean d µ and standard deviation σ . Weaim to sample dropout rate d ∈ [0 . 0 ,d max ] , with some mean d µ .

Accordingly, we sample u ∈ [0 , 1] ∼B ( α,β ) and set drop rate d = u × d max . α and β are set such that the mean of u is u µ = d µ d max . The maximum supported variance for a beta distribution of mean u µ is u µ (1 -u µ ) ; we pick by default a variance σ 2 =0 . 3 u µ (1 -u µ ) , which we found to work well in practice. The resultant distributions of token dropouts for different settings of d µ and d max are shown in Figure 15a.

25%

50%

75%

100%

Token drop rate

0%

5%

10%

15%

20%

25%

Percent

=0.10

p

<0.1 =60.3%

=0.25

p

<0.1 =11.5%

=0.35

p

<0.1 =1.3%

Max drop rate = 40%

25%

50%

75%

100%

Token drop rate

=0.10

p

<0.1 =66.9%

=0.25

p

<0.1 =27.3%

=0.50

p

<0.1 =2.3%

=0.65

p

<0.1 =0.1%

Max drop rate = 70%

25%

50%

75%

100%

Token drop rate

=0.10

p

<0.1 =70.4%

=0.25

p

<0.1 =33.5%

=0.50

p

<0.1 =5.9%

=0.70

p

<0.1 =0.9%

=0.85

p

<0.1 =0.1%

Max drop rate = 90%

25%

50%

75% 100%

Token drop rate

0%

1%

2%

3%

4%

5%

Percent

d

[0.1, 0.9]

d

[0.25, 0.75]

d

[0.4, 0.6]

Sampling resolution-dependent dropping rates Given input data with sequence lengths ranging from s min to s max , we sample dropout rate d from a truncated normal distribution d ∼ N trunc ( µ, 0 . 02) , where samples more than two standard deviations away from µ are rejected.

The mean of this distribution µ is set according to the minimum and maximum token dropping rates d min and d max , and simply scales linearly with the sequence length s (such that s = s min has µ = d min and s = s max has µ = d max .

Figure 15b shows example distributions of sampled drop rates given inputs with resolution R ∼ U (64 , 384) , and different values of d min and d max .

0

200 k

400 k

Total training steps

0.25

0.50

0.75

Average drop rate

0

200 k

400 k

Total training steps

0

1

2

Total images seen

1e9

0

200 k

400 k

Total training steps

65%

70%

75%

ImageNet 10shot acc

Fixed

= 1

2

=

N

10

,

=

N

2

=

N

10

,

=

N

2

=

N

5

,

= 2

N

3

B.7 Scheduling token dropping rates

We experiment with a token dropping schedule which varies with total number of images seen. In particular, the rate applied for the n -th processed image during training is given by:

where σ represents the sigmoid function; ρ min , ρ max control the minimum and maximum dropping rate applied; and µ and τ control the shape of the schedule. We experimented with both increasing ( τ > 0 ) and decreasing ( τ < 0 ) schedules. In all cases we set ρ min =0 . 2 and ρ max =0 . 8 . Figure 16 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 on image sizes outside of the training distribution. Results and discussion of these experiments can be found in Section 3.4.

Broadly, we 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.

(0.0, 0.0)

(2, 3)

(0, 0)

(6, 4)

0

2

0

3

0

6

0

4

(0.0, 0.0)

(1.0, 1.0)

(0.0, 0.0)

(1.0, 1.0)

0.0

1.0

0.0

1.0

0.0

1.0

0.0

1.0

(0.0, 0.0)

(0.75, 1.0)

(0.0, 0.0)

(1.0, .71)

0.0

0.75

0.0

1.0

0.0

1.0

0.0

.71

Absolute coordinates

Fractional coordinates (non aspect ratio preserving)

Fractional coordinates (aspect ratio preserving)

Absolute and fractional coordinates Anatural way of indexing token within an image is to select a priori a 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 in this way allow models to consume images with resolutions up to R = P · maxLen . However, when learned absolute coordinate embeddings are considered, extreme values of x and y must also be observed during training, which necessitates training on images with varied aspect ratios and limits models generalisation.

Toalleviate 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 x 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 x and y should be embedded independently or jointly. In case of independent embedding, the two coordinates x and y are embedded independently, and their embeddings are combined via addition or by stacking. For joint embeddings and embedding for each position ( x,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 we followed [57] and used sinusoidal positional embeddings, and learned embeddings were implemented as in [1].

For parametric positional embeddings we followed [8] and used Fourier embeddings. Specifically, coordinates ( x,y ) were mapped using a 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 3 and Figure 10a.

In sinusoidal and factorised embeddings experiments with fractional coordinates fractional coordinate embeddings were obtained from absolute coordinate embeddings via bilinear interpolation.

D Inference strategies

Weperformed 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.

Wefirstly looked at how model quality changes in respect to decreasing area of the image compared to native resolution, presented in Figure 18a. We observed that on 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 18b.

Please note that this highly depends on the native resolution of the images in the dataset - e.g. dataset with twice as big images than ImageNet can probably be safely resized to 20% of area.

Abetter way to quantify the performance is by giving a constant compute budget corresponding to number of patches. Figure 19a shows that resizing the image (while preserving aspect ratio) to 256 tokens retains most of the quality (within 0.3%). This corresponds to 256x256 area (given patch size of 16). At 128 tokens (181x181 area) the quality difference reaches 1% and drops significantly after that.

0.25

0.50

0.75

1.00

1.25

1.50

Image area

85

86

87

88

Imagenet accuracy

NaViT-L/16

NaViT-B/16

(a)

0.25

0.50

0.75

1.00

1.25

1.50

Image area

86.5

87.0

87.5

88.0

88.5

Imagenet accuracy

resize

drop tokens

(b)

128

256

384

512

640

768

896

1024

sequence length

87.6

87.8

88.0

88.2

88.4

88.6

Imagenet accuracy

NaViT-L

NaViT-L(no upscale)

(c)

64128

256

384

512

640

768

896

1024

sequence length

84

85

86

87

88

Imagenet accuracy

6x reduction

1%

5x reduction

1%

NaViT-L/16

NaViT-B/16

(a)

5

10

15

20

25

average runtime in ms per image

86.5

87.0

87.5

88.0

Imagenet accuracy

128

256

384

512

640

768

896

1024

256

384

512

640

768

896 1024

NaViT-L/16

NaViT-B/16

(b)

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 18c.

Figure 19b presents the runtime of NaViT-L/16 and NaViT-B/16 for different sequence lengths. We can see that NaViT-L/16 at sequence length 128 is as fast as NaViT-B/16 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 is in particular interesting for bulk inference workloads, where one can amortize over large datasets and where only the total inference time matters.

0

5

10

15

20

25

Average inference time in milliseconds

85.5

86.0

86.5

87.0

87.5

88.0

88.5

ImageNet1K-Accuracy

128

256

1024

Cascade (ViT-B)

ViT-B/16

Cascade (ViT-L)

ViT-L/16

Pareto frontier

To evaluate the feasibility of this approach, we consider two sequence lengths n 1 and n 2 with respective inference times t 1 and t 2 . Then, we (i) send all examples though the model with n 1 tokens, and send only the α ∈ (0 , 1) -fraction deemed hardest (those with the smallest maximum probability) to the model with n 2 tokens. To have an input almost exactly fit into n 1 or n 2 tokens we perform and aspect ratio preserving resize. Hence, the total amortized inference time per-example is t i + αt 2 , while the accuracy obtained by combining the accuracy of the first model on the 1 -α most-confident fraction of the data, and the performance of the more expensive model on the remaining data. By considering several pairs of models and varying α we obtain the plot in Figure 20. As we can see this strategy is indeed useful and provides not only the best performing models at given compute budgets, but because α is a real parameter one can obtain very fine-grained trade-offs.

F Calibration

To evaluate behaviour of the predicted uncertainties with scale, we compute the calibration error of a -B sized ImageNet-finetuned 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 up to 1. As we varied the sequence of tokens per example between 128 and 1024, we obtained very stable calibration errors (top-1, using glyph[lscript] 1 and 30 buckets, i.e., the settings from [29]), which we present in Table 4.

G Out of distribution evaluation

For ViT, we apply the 'Crop' strategy from [17], namely an aspect-preserving crop of the central 75% of the image for ObjectNet and ImageNet-A, and square resize followed by a 87.5% central crop for the other datasets. We also apply a simple 'Resize' strategy that does not crop the images. For NaViT, both the 'Crop' and the 'Resize' strategy do an aspect preserving resize of the target images.

10 12

10 13

Training TPU chip hours

55%

60%

65%

70%

75%

ImageNet

10 12

10 13

Training TPU chip hours

10%

20%

30%

40%

50%

60%

ImageNet-A

10 12

10 13

Training TPU chip hours

15%

20%

25%

30%

35%

40%

45%

50%

ObjectNet

ViT

B/32

B/16

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 metrics, 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 [30]. Thus, reducing the labeling error of fairness signals has the potential of improving the reliability of bias mitigation and post-hoc auditing [31]. Nevertheless,

accuracy accuracy

ImageNet

ImageNet-A

ObjectNet

ViT

ViT

ViT

Compute

Crop

B/32

7

.

9

×

10

1

54.4

63.5

14.7

32.7

32.9

41.2

1

.

9

×

10

2

60.2

65.6

22.0

37.0

38.0

43.5

3

.

9

×

10

2

63.7

67.2

26.6

40.1

41.7

45.1

7

.

9

×

10

2

65.5

68.2

30.7

42.5

44.2

45.7

B/16

2

.

6

×

10

2

66.4

68.5

37.3

52.3

47.2

48.4

6

.

6

×

10

2

61.7

66.0

27.0

44.2

41.3

45.9

1

.

3

×

10

3

70.3

71.2

48.8

57.0

52.8

50.7

2

.

6

×

10

3

68.7

70.1

43.5

55.1

50.5

49.8

5

.

4

×

10

2

70.7

73.6

51.5

65.5

53.3

55.0

1

.

3

×

10

3

66.4

71.1

39.2

58.6

47.6

52.1

2

.

7

×

10

3

73.9

75.1

60.4

68.9

57.7

57.9

5

.

4

×

10

3

73.0

74.6

57.6

67.9

56.2

57.1

Resize

B/32

7

.

9

×

10

1

51.7

64.0

12.6

26.7

15.9

31.6

1

.

9

×

10

2

57.9

66.2

17.2

30.1

20.0

33.8

3

.

9

×

10

2

61.6

67.4

20.7

32.6

23.5

34.6

7

.

9

×

10

2

63.6

68.6

22.8

35.0

25.3

36.1

B/16

2

.

6

×

10

2

65.4

69.2

27.4

43.9

28.7

38.7

6

.

6

×

10

2

60.4

66.5

20.4

36.8

23.6

35.3

1

.

3

×

10

3

69.5

72.5

36.2

51.5

34.2

42.1

2

.

6

×

10

3

67.7

71.0

32.5

48.6

32.0

40.5

5

.

4

×

10

2

70.0

73.9

38.5

59.9

35.3

45.6

1

.

3

×

10

3

65.1

71.5

28.2

51.6

29.3

41.5

2

.

7

×

10

3

73.2

76.0

47.3

65.5

39.8

48.8

5

.

4

×

10

3

72.3

75.3

44.3

64.1

38.8

48.2

100

250

500

1000

Pretraining Steps (K)

98

99

celeba / gender

100

250

500

1000

Pretraining Steps (K)

93

95

97

accuracy

fairface / gender

100

250

500

1000

Pretraining Steps (K)

54

57

60

accuracy

fairface / age

100

250

500

1000

Pretraining Steps (K)

64

68

72

accuracy

fairface / ethnicity

Model

ViT

100

250

500

1000

Pretraining Steps (K)

98

99

celeba / gender

100

250

500

1000

Pretraining Steps (K)

95

96

97

accuracy

fairface / gender

100

250

500

1000

Pretraining Steps (K)

55

57

59

accuracy

fairface / age

100

250

500

1000

Pretraining Steps (K)

64

68

72

accuracy

fairface / ethnicity

Square

Native

we emphasize that while NaViT improves the annotation accuracy in these tasks, care must be taken in such situations since classifiers can be inaccurate and lead to a broad categorization of people that misidentifies real identities. We encourage readers to delve into the comprehensive work outlining such potential risks, e.g. [58, 59], for further insight. In assessing the technical capabilities of NaViT, our intent is not to promote or encourage their application in inappropriate contexts. Rather, our objective is only to illuminate these technical findings for scenarios where they may be considered beneficial, such as when measuring the level of diversity in a dataset or auditing/mitigating biases in predictive models. We strongly advocate for responsible AI use, maintaining that the benefits of technological advancements should not overshadow the importance of user safety and privacy. AI tools, including

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 NaViT, 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 × 224 pixel window!). We investigate how the ability to cope with variable resolutions affects performance on 'model-vs-human', a benchmark of 17 challenging datasets [60]. 1 For this purpose, we replicate the setup from Figure 6, but instead of evaluating ImageNet accuracy, we evaluate OOD accuracy on the model-vs-human benchmark.

64

128

224

384

512

Fine-tune resolution

Rft

85

90

95

Accuracy (%)

colour

NaViT (one per res)

ViT (one per res)

NaViT-64:512

64

128

224

384

512

Fine-tune resolution

Rft

85

90

95

Accuracy (%)

false-colour

64

128

224

384

512

Fine-tune resolution

Rft

60

80

Accuracy (%)

power-equalisation

64

128

224

384

512

Fine-tune resolution

Rft

40

60

80

Accuracy (%)

rotation

64

128

224

384

512

Fine-tune resolution

Rft

40

50

Accuracy (%)

uniform-noise

64

128

224

384

512

Fine-tune resolution

Rft

60

80

Accuracy (%)

contrast

64

128

224

384

512

Fine-tune resolution

Rft

50

60

Accuracy (%)

low-pass

64

128

224

384

512

Fine-tune resolution

Rft

20

40

60

Accuracy (%)

high-pass

64

128

224

384

512

Fine-tune resolution

Rft

60

80

Accuracy (%)

phase-scrambling

64 128

224

384

512

Fine-tune resolution

Rft

60

65

Accuracy (%)

eidolonI

64

128

224

384

512

Fine-tune resolution

Rft

55.0

57.5

60.0

62.5

Accuracy (%)

eidolonII

64

128

224

384

512

Fine-tune resolution

Rft

60

65

Accuracy (%)

eidolonIII

64

128

224

384

512

Fine-tune resolution

Rft

60

70

80

90

Accuracy (%)

sketch

64

128

224

384

512

Fine-tune resolution

Rft

40

50

60

Accuracy (%)

stylized

64

128

224

384

512

Fine-tune resolution

Rft

55

60

65

70

Accuracy (%)

silhouette

64

128

224

384

512

Fine-tune resolution

Rft

20

40

60

Accuracy (%)

edge

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 finetuning 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-matched 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 NaViT model trained on varying resolutions (red) performs roughly on par with fine-tuning one NaViT model per test resolution (blue).

The ViT baseline (orange) is mostly worse than NaViT models for lower resolutions and mostly better for higher resolutions. It may be worth noting that ViT models have a 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 aspect ratio) that have the same number of pixels, while not necessarily being square.

Peer reviews:
comments: 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.

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.
2. The experiments are comprehensive. The efficiency gain during pretraining is impressive.
3. 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.
2. 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.
3. It is not clear how the memory cost would be under the proposed Patch n’ Pack.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-28

comments: 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.

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.

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.

The additional overhead of introducing NaViT as the image preprocessing could be reported.

QUESTIONS:
Please see the suggestions in the weakness part.

LIMITATIONS:
None.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nI found no potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOverall, 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.\n\nWEAKNESSES:\nAccording 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.\n\nQUESTIONS:\nSee weaknesses section, thank you.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}