Source-linked AI summary
FlexiViT: One Model for All Patch Sizes
Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, Filip Pavetic
TL;DR
Changing ViT patch size changes the compute-accuracy tradeoff, but standard models generally require separate training for different sizes. FlexiViT randomizes patch size during training, and usually matches or outperforms fixed-patch ViTs across downstream tasks while enabling one model to serve multiple compute budgets.
Problem
Patch size strongly affects ViT compute and predictive performance, while changing it typically requires separately trained models despite similar parametrization.
Method
FlexiViT randomizes patch size during training and adaptively resizes positional and patch embedding parameters for each size.
Results
FlexiViT usually matches and sometimes outperforms fixed-patch ViTs across a wide range of patch sizes and downstream tasks.
Takeaways & Limitations
A single FlexiViT can trade compute for predictive performance across patch sizes and reduce pretraining costs by replacing multiple scale-specific models.
Takeaways & Limitations
The authors did not explore Flexi training’s data-augmentation impact, and early experiments used a patch-size distribution that was not rerun after later findings.
Abstract
from arXiv · showhide
Vision Transformers convert images to sequences by slicing them into patches. The size of these patches controls a speed/accuracy tradeoff, with smaller patches leading to higher accuracy at greater computational cost, but changing the patch size typically requires retraining the model. In this paper, we demonstrate that simply randomizing the patch size at training time leads to a single set of weights that performs well across a wide range of patch sizes, making it possible to tailor the model to different compute budgets at deployment time. We extensively evaluate the resulting model, which we call FlexiViT, on a wide range of tasks, including classification, image-text retrieval, open-world detection, panoptic segmentation, and semantic segmentation, concluding that it usually matches, and sometimes outperforms, standard ViT models trained at a single patch size in an otherwise identical setup. Hence, FlexiViT training is a simple drop-in improvement for ViT that makes it easy to add compute-adaptive capabilities to most models relying on a ViT backbone architecture. Code and pre-trained models are available at https://github.com/google-research/big_vision
1. Introduction
Patch size is an underexamined ViT design lever that changes compute and predictive performance without changing model parametrization. FlexiViT randomizes patch size during training to support strong performance across patch sizes and tasks.
- Patchification creates ViT tokens from nonoverlapping image patches, enabling token dropping, specialized task tokens, and multimodal token mixing.
- 85.6% top-1 accuracy at 156 GFLOPs characterizes ViT-B/8, versus 79.1% at 8.6 GFLOPs for ViT-B/32, despite essentially identical parametrization.
- FlexiViT randomizes patch size during training and adaptively resizes positional and patch embedding parameters for each size.
- FlexiViT matches or outperforms fixed-patch ViTs across a wide range of patch sizes with no added cost.
- FlexiViT representations are often similar across patch sizes, especially in deeper layers, and outperform alternative architectural compute-accuracy trade-offs.
2. Related work
Related work improves ViT efficiency through token removal, variable resolution, or broader architectural flexibility. FlexiViT instead focuses on making one ViT work across trained patch sizes while retaining all tokens.
- Token-efficiency methods remove randomized, structured, or least-important tokens, whereas FlexiViT always keeps all tokens.
- Variable-resolution methods typically speed pretraining or augment self-supervised learning but operate only at a single final resolution.
- Neural Architecture Search trains one supernet from which differently shaped subnets can be extracted, usually changing multiple model dimensions.
- FlexiViT’s focus on patch size supports existing pretrained models and ViT improvements as an easy drop-in to existing training procedures.
- Matryoshka representation learning trains meaningful subvectors in output vectors and is described as complementary to FlexiViT.
3. Making ViT flexible
FlexiViT makes ViTs flexible across patch sizes by randomizing patch size during training and resizing patch and positional embeddings on the fly. It preserves or improves performance across patch sizes, with knowledge-distillation initialization providing further gains.
- Background: Patch size controls ViT sequence length and compute: smaller patches create longer, slower, more expressive inputs.For a fixed image size, self-attention scales quadratically with the number of patches.
- Making ViT flexible: Standard ViT performance rapidly degrades when inference-time patch size differs from the training patch size.The comparison uses resized patch and position embeddings to evaluate fixed-patch models at other sizes.
- Making ViT flexible: FlexiViT samples patch sizes during training and adaptively resizes patch and positional embeddings during each forward pass.The underlying learnable parameters use 32 × 32 patch and 7 × 7 position-embedding shapes, with resizing performed on the fly.
- Making ViT flexible: FlexiViT matches fixed ViTs at their training patch sizes and significantly outperforms them at other patch sizes.The model is trained in the same setting as fixed ViT-B/16 and ViT-B/30 models, except for randomized patch-size sampling.
- Making ViT flexible: PI-resize provides a principled resizing method compatible with existing pretrained ViTs without architectural changes.It aims to preserve patch embeddings when both inputs and embedding weights are resized, unlike token normalization approaches that require architecture changes.
- Making ViT flexible: Teacher initialization and distillation improve FlexiViT, while output representations remain generally aligned across grid sizes.FlexiViT is initialized from a ViT-B/8 teacher; distillation improves accuracy especially at the largest patch sizes, and CLS representations remain aligned.
4. Using pre-trained FlexiViTs
Pre-trained FlexiViTs transfer effectively across classification, retrieval, detection, and segmentation tasks. Their flexibility enables resource-efficient transfer by training cheaply with large patches and deploying with smaller patches for higher accuracy.
- Transfer tasks: Transfer evaluations cover small- and medium-scale classification, image-text retrieval, open-vocabulary detection, panoptic segmentation, and semantic segmentation.
- Results: A single pre-trained FlexiViT roughly matches two fixed-patch ViTs across diverse transfer tasks.It barely lags at large patch size and yields small or significant improvements at smaller patch size.
- Results: FlexiViT shows no significant downside compared with pretraining multiple ViTs for different patch sizes.
- Resource-efficient transfer: FlexiViT preserves flexibility after fixed-size transfer, enabling cheaper training with large patches and deployment with small patches.This saves accelerator memory and compute during transfer.
- Resource-efficient transfer: Cheap finetuning at 8×8 grid size reaches 81.8% accuracy, while evaluation at 24×24 reaches 85.3% top-1 accuracy without additional training.The larger grid gains 3.5% accuracy at no additional training cost.
5. Flexifying existing training setups
FlexiViT’s patch-size randomization can be applied during transfer to existing models and training setups. Flexible transfer generally performs best, while patch-size curricula can also improve performance per compute during pre-training.
- Flexifying transfer: Existing pre-trained models can be flexified during transfer by randomizing patch size and evaluating one resulting model at multiple sizes.
- Flexifying transfer: Flexible transfer of FlexiViT works best, while flexifying a fixed model during transfer also performs surprisingly well.
- Multimodal transfer: FlexiLiT-B/flexi performs best on average for zero-shot image-to-text retrieval, while fixed-patch transfer of FlexiViT also performs well.
- Open-vocabulary detection: Flexible OWL-ViT detection models perform close to or better than fixed-size models at any patch size during inference.For detection, the optimal patch size is not necessarily the smallest.
- Pre-training curricula: Patch-size curricula generally provide better performance per compute budget than standard training.A curriculum specifies changing probability distributions over patch sizes during training.
6. Analyzing FlexiViTs
FlexiViT changes attention patterns across patch-size scales while retaining correspondences between token representations. Its scale-specific behavior also preserves comparable texture bias, and ensembling multiple scales is usually compute-inefficient.
- Attention relevance patterns across scales: Decreasing patch size concentrates attention relevance into more numerous, smaller image areas, and attention can change significantly across scales.
- Relation of token representations across scales: Token representations remain correspondable across scales, including between a seed token and tokens produced at different patch sizes.
- Ensembling: Running a single FlexiViT at a given compute budget is nearly always better than ensembling multiple smaller-scale predictions.
- Shape or texture bias: At each patch size, FlexiViT has a texture bias similar to a ViT trained and evaluated at that same patch size.
- Model and dataset size: A FlexiViT-L trained on ImageNet-1k matches or outperforms three DeiT III models and EfficientNetV2.The paper also examines whether base model size and a 12 M-image training set are required.
7. Discussion of alternatives
The paper compares FlexiViT with alternative ways to vary ViT compute, including patch-embedding stride and model depth. Varying stride performs nearly as well, while FlexiViT offers a significantly better trade-off than depth pruning.
- Comparison setup: The alternatives are evaluated in the paper’s core setup with distillation on ImageNet-21k.
- Varying patch embedding stride: Varying patch-embedding stride uses overlapping patches to increase sequence length while keeping the intrinsic patch size fixed.Results suggest this alternative works almost as well as FlexiViT, with a slight performance gap.
- Varying model depth: Depth flexibility randomizes the number of layers by attaching a shared head to intermediate layers.Depth pruning scales linearly in depth, uses a subset of parameters, and progressively refines predictions.
- Varying model depth: FlexiViT provides a significantly better compute-accuracy tradeoff than depth pruning.
8. Conclusion
FlexiViT provides a single model for trading compute against predictive performance across patch sizes, while retaining standard ViT compatibility. The appendix finds resizing methods and several hyperparameters that preserve strong, stable performance.
- Conclusion: FlexiViT offers an efficient compute–performance tradeoff with one model across multiple patch scales and downstream tasks.The paper reports performance across transfer learning, segmentation, image-text retrieval, and open-vocabulary detection settings.
- Flexible patch-size variants: Plain resizing and PI-resizing are among the best alternatives and preserve standard ViT models, unlike normalization and untied-weight variants.Untied variants also showed slightly unstable training curves during the first half of training.
- Resizing analysis: PI-resize uses negative weights, a larger receptive field, and overall larger weights than bilinear resizing.The comparison is based on visualization of the resize matrices for upscaling and downscaling operations.
- Hyperparameter analysis: The underlying patch-embedding and position-embedding sizes have little influence on final performance within a reasonable range.The patch-embedding ablation showed no notable difference across evaluation sizes, motivating a default underlying size of 32.
- Hyperparameter analysis: Uniform patch-size sampling is preferred for simplicity because a triangular distribution was no better in the reported ablation.The triangular distribution sampled patch sizes from 16 to 30 up to three times more often than sizes outside that range.
G. FlexiLiT and FlexiCLIP results
FlexiLiT and FlexiCLIP extend patch-size randomization to image-text learning. The results indicate that FlexiCLIP can learn image and text representations jointly without patch-size randomization hindering training.
- FlexiLiT: FlexiLiT randomizes patch sizes during LiT training while following the same general setup as the fixed-patch configuration.
- FlexiCLIP: FlexiCLIP replaces the pretrained frozen backbone with a randomly initialized, unfrozen backbone, corresponding to CLIP-style training from scratch.
- FlexiCLIP: Patch-size randomization does not hinder learning image and text representations from scratch simultaneously in the FlexiCLIP setting.
H. Accelerate pre-training
A curriculum trains ViTs first with larger patches and later with the target patch size to reduce pretraining compute. The resulting FasterViT reaches standard ViT accuracy with less compute, while longer large-patch schedules trade accuracy for efficiency.
- Curriculum design: FasterViT uses a patch-size curriculum of (48, 40, 30, 24, 16) to demonstrate potential pretraining compute savings.The target evaluation patch size is 16×16, and the architecture otherwise matches standard ViT/B/16 except for variable patches and embedding layers.
- Curriculum design: The curriculum begins with a large patch size and switches to the 16×16 target using PI-resize, with optional ramp periods between distributions.Schedules allocate 50%, 75%, or 90% of training to the larger patch size; ramp periods range from 0% to 100%.
- Results: FasterViT achieves the same accuracy as standard ViTs with less compute, although the improvement is not quite significant.
- Results: Increasing the fraction of training assigned to the larger patch size improves compute at the expense of accuracy.The tradeoff is illustrated by varying the schedule parameter, including schedule = 90%.
I. Further analysis of cosine similarities between token representations across scales
The analysis compares token representations across patch-size scales using cosine similarity. Similar tokens across scales generally correspond to the same spatial location, with consistent results across grid sizes and blocks.
- Analysis setup: Cosine similarity is measured between a seed token at one scale and tokens at other scales to analyze cross-scale representation alignment.
- Findings: The most similar tokens across scales generally represent the same spatial location.Supplementary results use seed tokens from FlexiViT-B/48, FlexiViT-B/16, and FlexiViT-B/8.
- Findings: The supplementary results are consistent across additional grid sizes and transformer blocks.
J. Ensembling FlexiViT predictions across scales
Across scales, ensembling FlexiViT predictions rarely improves accuracy over selecting a single patch size at comparable computational cost. Even with low agreement between large-patch configurations, single-scale inference is generally preferable.
- All-scale ensemble: 51.7% precision@1 from ensembling all models is slightly below the 52.0% achieved at the largest grid size and smallest patch size.FlexiViT predictions are ensembled by averaging model logits.
- Pairwise ensembles: 45.8% ImageNet-21K precision@1 is achieved by both FlexiViT-B/40 and the ensemble of the two largest patch sizes.The /48 and /40 models agree on only 67.4% of examples, yet ensembling them does not improve accuracy.
- Accuracy versus computational cost: Single-scale inference is nearly always more accurate than FlexiViT ensembles at the same computational cost.The only reported exception is a two-scale ensemble whose improvement is marginal and unlikely to be statistically or practically meaningful.
- Pairwise ensembles: 52.1% versus 52.0% accuracy is obtained by the /10 + /12 ensemble and FlexiViT-B/8, respectively, at similar computational cost.The ensemble uses 186.8 GFLOPs compared with 184.5 GFLOPs for /8, and the difference is unlikely to be statistically significant or practically meaningful.