Source-linked AI summary
Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, Avital Oliver, Piotr Padlewski, Alexey Gritsenko, Mario Lučić, Neil Houlsby
TL;DR
Fixed-resolution image processing is a limiting and often suboptimal convention, especially because many images are not square. NaViT uses Patch n’ Pack to train Vision Transformers on variable-resolution, aspect-ratio-preserving inputs, improving efficiency and enabling flexible inference across resolutions. The paper reports consistent gains over ViT baselines, including matching the top-performing ViT with four times less compute, while noting constraints around packed example-level losses.
Problem
Fixed image sizes force resizing or padding, although resizing harms performance, padding is inefficient, and most representative images are not square.
Method
NaViT packs patch tokens from multiple images into sequences, using masked operations, factorized positional embeddings, and mixed-resolution training to preserve aspect ratios.
Results
NaViT consistently surpasses compute-matched ViT baselines across scales, matching the top-performing ViT with four times less compute and supporting multiple inference resolutions.
Takeaways & Limitations
Patch n’ Pack enables efficient training, cheaper adaptation to new tasks, and smooth inference cost-performance trade-offs across resolutions.
Takeaways & Limitations
Example-level losses require modified pooling and impose a limit on the number of pooled representations extracted per packed sequence.
Abstract
from arXiv · showhide
The ubiquitous and demonstrably suboptimal choice of resizing images to a fixed resolution before processing them with computer vision models has not yet been successfully challenged. However, models such as the Vision Transformer (ViT) offer flexible sequence-based modeling, and hence varying input sequence lengths. We take advantage of this with NaViT (Native Resolution ViT) which uses sequence packing during training to process inputs of arbitrary resolutions and aspect ratios. Alongside flexible model usage, we demonstrate improved training efficiency for large-scale supervised and contrastive image-text pretraining. NaViT can be efficiently transferred to standard tasks such as image and video classification, object detection, and semantic segmentation and leads to improved results on robustness and fairness benchmarks. At inference time, the input resolution flexibility can be used to smoothly navigate the test-time cost-performance trade-off. We believe that NaViT marks a departure from the standard, CNN-designed, input and modelling pipeline used by most computer vision models, and represents a promising direction for ViTs.
1 Introduction
NaViT replaces fixed-resolution image processing with packed variable-resolution sequences that preserve aspect ratio. The approach improves efficiency, supports flexible inference, and enables new training and modeling choices.
- Patch n’ Pack: NaViT packs patches from different images into one sequence, enabling variable resolutions while preserving each image’s aspect ratio.The approach is called Patch n’ Pack and adapts sequence packing from language modeling to vision transformers.
- Efficiency and flexibility: Randomly sampling resolutions during training significantly reduces training cost while maintaining flexible model usage.Fixed batch shapes also support aspect-ratio-preserving resolution sampling, variable token dropping, and adaptive computation.
- Efficiency and flexibility: A single NaViT performs across multiple resolutions, smoothly trading off inference cost and performance.Exposure to multiple resolutions during pre-training and fine-tuning supports flexible evaluation at test time.
- Efficiency and flexibility: 4× less compute matches the performance of the top-performing ViT during pre-training.NaViT consistently outperforms compute-matched ViT baselines, and its efficiency carries over to fine-tuning.
- Broader implications: Patch n’ Pack enables research directions previously constrained by fixed batch shapes, including adaptive computation and efficiency-oriented algorithms.The paper presents this flexibility as a promising direction for Vision Transformers.
2 Method
The method addresses inefficient fixed-size image pipelines by packing patch sequences from variable-resolution, non-square images. Architectural masking and positional-embedding changes, together with resolution and token-dropout strategies, support efficient training with limited packing overhead.
- Motivation: Most images in representative classification, detection, and web-image datasets are not square, making fixed resizing or padding problematic.Resizing can harm performance, whereas padding is inefficient.
- Core method: Patch n’ Pack combines tokens from multiple images into one sequence, allowing ViTs to train on images at native resolution.NaViT names the resulting approach and uses sequence packing to bypass fixed sequence-length limits.
- Architectural changes: Masked attention and masked pooling prevent packed images from interacting and produce one pooled representation per image.These modifications control each example’s attention receptive field and support example-level representations.
- Architectural changes: Factorized positional embeddings sum separate x- and y-coordinate embeddings to support variable aspect ratios and extrapolate to unseen resolutions.Fractional embeddings make parameters independent of image size but encode aspect ratio only implicitly through patch count.
- Training changes: Continuous token dropping varies the dropped-token proportion per image, combining faster throughput with exposure to some complete images.This reduces the train/inference discrepancy compared with dropping the same proportion from every example.
- Training changes: Resolution sampling mixes image sizes while preserving aspect ratio, combining lower-resolution throughput with exposure to large images.The paper reports improved performance over equivalent ViTs in model size and training duration.
- Efficiency: Packing overhead from extra attention diminishes as transformer model scale increases.The paper attributes this trend to attention becoming a smaller share of total computation as hidden dimension grows.
- Efficiency: Less than 2% of tokens are typically padding tokens under greedy fixed-length sequence packing.This makes the simple packing approach sufficient without dynamically fitting the final example.
3 Experiments
NaViT improves training efficiency and performance through mixed-resolution inputs, packing, and token dropping, while transferring effectively across resolutions, robustness, fairness, and downstream tasks.
- Improved training efficiency and performance: NaViT consistently surpasses compute-matched ViT baselines across model and parameter scales, matching the top ViT with four times less compute.The improved efficiency is primarily attributed to processing more training examples within the same compute budget.
- Benefits of variable resolution: Variable-resolution NaViT training matches or outperforms fixed-resolution training at the same computational budget, including when training and evaluation resolutions coincide.The variable-resolution setup samples resolutions uniformly between 64 and Rmax.
- Benefits of variable resolution: NaViT finetuned across variable resolutions performs as well as single-resolution NaViT and can retain good high-resolution performance after low-resolution finetuning.This reduces the need to select one downstream finetuning resolution and supports cheaper adaptation.
- Benefits of variable token dropping: Variable token-dropping strategies improve performance, including resolution-dependent rates that further improve over Beta-distributed rates.Scheduled reductions in token dropping during JFT pretraining also produce additional gains for NaViT-B/16.
- Other aspects of NaViT’s performance: NaViT compares favorably on ImageNet and out-of-distribution datasets, with particularly stronger ImageNet-A performance while ViT catches up on ObjectNet.The observed comparison depends partly on the aspect-preserving center crop used for evaluation.
- Other aspects of NaViT’s performance: NaViT maintains stable calibration as patches per image vary from 128 to 1024 and improves fairness-signal annotation accuracy over ViT.Calibration error remains between 0.045 and 0.047; fairness accuracy gains are statistically significant for both tested settings.
- Other aspects of NaViT’s performance: NaViT offers stronger inference cost-performance trade-offs, with competitive results using relatively few patches and diminishing returns as patch counts increase.The trade-off is quantified using latency on a Cloud TPUv3 chip.
- Other downstream tasks: NaViT outperforms ViT for semantic segmentation at the same maximum finetuning resolution, with R384 NaViT twice as fast as R512 ViT.NaViT also benefits from ingesting square and non-square images without changing the model pipeline.
4 Related work
Prior work addresses variable resolution, aspect ratio, and training efficiency through patch-size flexibility, multiscale representations, mixed-resolution schedules, and token dropping. Patch n’ Pack differs by incorporating mixed resolutions without complex schedules or fixed minibatch-shape restrictions.
- Flexible Vision Transformers: FlexiViT varies sequence length and compute cost by randomly sampling patch sizes during training and resizing the convolutional embedding.
- Flexible Vision Transformers: Pix2Struct preserves aspect ratios through a novel positional embedding scheme, supporting chart and document understanding.
- Multiscale Vision Transformers: Multiscale Vision Transformers use feature maps at multiple spatial scales for localization tasks such as segmentation and detection, whereas NaViT does not build hierarchical multiscale representations.
- Accelerating training with mixed resolutions: Mixed-resolution approaches commonly accelerate pretraining with fixed low resolution or staged resolution increases, as in FixRes and PaLI.
- Accelerating training with mixed resolutions: Multigrid training accelerates video modeling with coarse-to-fine hierarchical grid schedules and learning-rate scaling.
- Token dropping for improved efficiency: Patch n’ Pack enables mixed resolutions without complex schedules or training pipelines, while extending beyond earlier token-dropping strategies limited by fixed minibatch shapes.
5 Conclusions and future work
Patch n’ Pack significantly improves Vision Transformer training efficiency and allows NaViT models to operate across resolutions and adapt cheaply to new tasks. It also supports research directions such as adaptive computation and improved training and inference efficiency.
- Patch n’ Pack significantly improves training efficiency and enables NaViT models to operate at many inference resolutions and adapt cheaply to new tasks.
- Patch n’ Pack enables research on adaptive computation and algorithms for improving training and inference efficiency.
A Training details
Training details specify model configurations, resolution sampling, sequence packing, token-dropping distributions, and supervised pretraining specifications for ViT and NaViT.
- The experiments use ViT-B/32, ViT-B/16, and ViT-L/16 with a reciprocal square-root learning-rate schedule, warmup and cooldown, and separate head and body weight decay.
- ViT models use 224×224 images, whereas NaViT uniformly samples r between 64 and 256 and preserves aspect ratio while resizing to r^2 pixels.
- Text sequences use a T5 SentencePiece tokenizer, are truncated to length 24 by default, and use no text-token dropping.
- Examples are greedily packed into sequences with sufficient remaining space, then padded when no additional examples fit.
- Token-drop rates can be sampled from beta distributions by scaling a beta sample u to the maximum drop rate dmax.
- Table 2 presents the supervised-classification pretraining specifications for ViT and NaViT.
- Resolution-dependent token dropping samples rates from a truncated normal distribution whose mean scales linearly with sequence length between dmin and dmax.
- Figure 14b shows sampled drop-rate distributions for resolutions drawn uniformly between 64 and 384 under different dmin and dmax settings.
A.5 Scheduling token dropping rates
The paper studies token-dropping schedules and positional embeddings for variable-size inputs. Decreasing token dropping during training improves accuracy at fixed cost, while positional-coordinate choices trade generalization against spatial information.
- Scheduling token dropping rates: The token-dropping rate for the n-th processed image is controlled by a schedule with minimum and maximum rates and shape parameters.
- Scheduling token dropping rates: Decreasing the token-dropping rate throughout training improves ImageNet 10shot accuracy using the same pre-training resources.
- Scheduling token dropping rates: Increasing the token-dropping rate during training harms performance, whereas decreasing it improves final accuracy at fixed training cost.
- Positional embeddings: Variable input sizes require evaluating positional embeddings both within the training resolution range and on image sizes outside that distribution.
- Positional embeddings: The positional-embedding study varies whether embeddings are learned, parametric, or fixed; absolute or fractional; and factorized.
- Positional embeddings: Absolute coordinates support images up to R = P · maxLen but require extreme coordinate values to appear during training for generalization.
- Positional embeddings: Fractional coordinates normalize positions by the input dimensions, exposing extreme token coordinates during training but obscuring absolute image size.
- Positional embeddings: Coordinate embeddings may be factorized by embedding x and y independently or jointly by assigning an embedding directly to each position pair.
C Inference strategies
NaViT inference cost can be adjusted through sequence length, image resizing, model size, and cascaded allocation of tokens. Aspect-ratio-preserving resizing retains quality efficiently, while cascades provide fine-grained cost–performance trade-offs.
- Inference runtime can be tuned by changing the number of processed patches or the model size.
- Random token dropping is ineffective for reducing inference patches compared with resizing.
- Resizing preserves most ImageNet quality at 256 tokens, within 0.3%, while quality differences reach 1% at 128 tokens.Quality drops significantly below 128 tokens.
- NaViT-L/16 at 128 tokens runs as fast as NaViT-B/16 at 512 tokens, with almost a 1% quality difference.
- A cascade sends all examples through a smaller model and the hardest fraction through a larger one, producing fine-grained amortized compute–accuracy trade-offs.The hard fraction is selected using the smallest maximum probability, and aspect-ratio-preserving resizing fits the chosen token budgets.
E Calibration
The calibration experiment evaluates whether predicted uncertainties remain reliable as sequence length changes. Calibration errors were stable across the tested token scales.
- Calibration error was evaluated for an ImageNet-finetuned B-sized model across sequence lengths from 128 to 1024 tokens.The L-sized model performed similarly.
- Calibration errors remained very stable as the number of tokens per example varied.The evaluation used top-1 calibration error with ℓ1 distance and 30 buckets.
F Out of distribution evaluation
The out-of-distribution evaluation compares crop and resize preprocessing for ViT and NaViT across benchmark datasets. NaViT uses aspect-preserving resizing for both evaluation strategies.
- ViT’s Crop strategy uses an aspect-preserving central crop for ObjectNet and ImageNet-A, and square resize followed by a central crop elsewhere.
- ViT’s Resize strategy avoids cropping the images.
- NaViT uses aspect-preserving resize for both its Crop and Resize evaluation strategies.
G Fairness Signal Annotation
Fairness signal annotation results indicate benefits from NaViT representations and native aspect ratios, while the text also flags a need for caution.
- NaViT improves fairness signal annotation performance when native image resolution is used.
- Figure 20 reports that aspect-preserving NaViT resizing performs better on datasets containing images with extreme aspect ratios.
- Figure 21 summarizes higher annotator accuracy with NaViT representations and higher performance from native aspect ratios than square resizing.
H Evaluation on model-vs-human OOD datasets on different resolutions
The evaluation tests NaViT’s variable-resolution capability on the 17-dataset model-vs-human OOD benchmark across different fine-tuning resolutions. A single resolution-flexible NaViT performs roughly on par with separately fine-tuned NaViT models, while ViT’s relative performance depends on resolution.
- The benchmark evaluates OOD accuracy on 17 challenging model-vs-human datasets across different fine-tuning resolutions.
- A single NaViT trained across varying resolutions performs roughly on par with separately fine-tuned NaViT models at each test resolution.The comparison uses compute-matched pretraining and identical compute and data during fine-tuning.
- ViT is mostly worse than NaViT at lower resolutions but mostly better at higher resolutions.
- The OOD evaluation excludes the cue-conflict dataset because texture-shape conflict images lack an objective ground-truth class.
- ViT has a comparison advantage because it is fine-tuned on square images, whereas NaViT preserves aspect ratio at the same pixel count.