Accelerating Vision Transformers with Adaptive Patch Sizes

AI summary

Overview

Research area: Efficient computer vision, specifically accelerating Vision Transformers (ViTs) through content-aware, variable-size image patchification.

Technical level: Intermediate. The paper assumes familiarity with ViT patch embedding, self-attention cost scaling, sequence packing, and fine-tuning workflows, but it explains its core mechanism clearly enough for readers who know the basics of transformer vision backbones.

One-sentence scope: The paper introduces Adaptive Patch Transformers (APT), a method that assigns different patch sizes within a single image based on local entropy, and reports speedups on ImageNet classification, visual question answering, object detection, and semantic segmentation.

What This Paper Is About

Standard Vision Transformers split every image into equal-sized patches, so high-resolution inputs produce very long token sequences even when large parts of the image are visually simple and redundant. The authors propose APT, which measures entropy at multiple scales and assigns large patches to homogeneous regions and small patches to detailed regions, reducing the number of input tokens while trying to preserve downstream accuracy. The goal is wall-clock speedup in both training and inference, rather than theoretical FLOP reduction alone.

Key Contributions

1. Adaptive Patch Transformer (APT): A content-aware patchification scheme that uses multiple patch sizes within one image, allocating larger patches to low-entropy regions and smaller patches to complex ones, reducing token count.

2. Preservation of pretrained accuracy: APT is reported to preserve the accuracy of standard pretrained models across resolutions and model scales, and to match prior performance on dense prediction and vision-language tasks.

3. Cheap adaptation via a zero-initialized MLP: A zero-initialized MLP (inspired by ControlNet) combines resized large-patch embeddings with aggregated sub-patch embeddings, allowing APT to be applied to an already fine-tuned ViT and to converge in as little as one epoch.

4. Extension to downstream and window-attention tasks: APT is extended to varying input sizes via sequence packing and positional encoding interpolation, and adapted to object detection and segmentation, including models that use window attention.

Main Findings

• Throughput gains on classification: The abstract reports a 40% throughput increase on ViT-L and 50% on ViT-H while maintaining downstream performance. In the 1-epoch fine-tuning table, APT reaches +7% (ViT-B 224/16), +21% (ViT-B 384/16), +12% (ViT-L 224/14), +33% (ViT-L 336/14), +23% (ViT-H 224/14), and +50% (ViT-H 336/14) speedups.

• Full fine-tuning on ImageNet: Using the MAE training recipe, ViT-B at 384/16 goes from 11.6h to 9.0h (+29%), ViT-L at 336/14 from 15.9h to 9.9h (+61%), and ViT-L at 448/14 from 31.4h to 16.9h (+86%), with accuracy of 84.2, 86.1, and 86.3 respectively versus baselines of 84.2, 86.1, and 86.4.

• Accuracy matching after one epoch: Table 2 shows APT matching or slightly trailing the original network after one additional epoch, e.g. ViT-B 384/16 at 86.1 versus 86.1 baseline and ViT-H 336/14 at 88.4 versus 88.5 baseline.

• Gains grow with resolution and model size: The paper states that speedups increase dramatically at higher resolutions and larger scales, which it attributes partly to attention computation dominating more of the per-iteration training time.

• Best trade-off against ablations: The paper reports APT offers a better speed-accuracy trade-off than random masking or resizing-only, which serve as stronger versions of FLIP and Quadformer respectively.

• Layer-level comparison: In accuracy-versus-throughput comparisons on ViT-L/14 at 224×224 and ViT-H/14 at 336×336, APT is reported to outperform layer-level merging methods (EViT, ToMe, PPT, DTEM) including re-implemented "Advanced" versions using FlashAttention.

• Visual QA: On LLaVA-1.5, APT increases throughput by 22% on the 7B model and 23% on the 13B model, matching or exceeding baseline across VQA-v2, GQA, SQA-I, VQA-T, POPE, MME, MMB, MMB-C, and MMV. The paper states the throughput increase is 23% and that APT exceeds original performance on a range of vision-language benchmarks.

• Object detection: On COCO at 1536×1536 with an EVA-02 backbone and ViTDet-style head, APT reports +14% speedup for EVA-02-B (mAP 58.79 vs 58.93; AP50 77.65 vs 77.85) and +30% for EVA-02-L (mAP 62.07 vs 62.28; AP50 80.64 vs 80.80), reducing 30% of input tokens.

• Semantic segmentation: On ADE20K with EVA-02-L and UperNet, APT gives +11% speedup at both 512 and 640 resolution while reducing 28~32% of input tokens. At 512: aAcc 86.68 vs 86.67, mIoU 59.70 vs 59.77. At 640: aAcc 86.82 vs 86.83, mIoU 60.01 vs 60.05.

• Overhead measured, not free: With no sequence reduction, APT is about 10% slower (for example, ViT-B 224/16 drops from 3310 to 3090 images/s). With τ=5.75, the ZeroConv operation takes about 10% of end-to-end model latency, but token removal yields a net 33% speedup.

• Zero-initialization ablation: Zero-initialized connection gives 87.98 without training and 88.13 with training, compared to residual (87.40 / 87.52) and non-zero initialization (87.50 / 87.81), against a base of 88.15.

• Threshold behavior: Default thresholds are τ32 = 5.75 and τ64 = 4.0. Object detection used a lower threshold of 2, which the authors attribute to that task's reliance on precise edge localization. The paper notes accuracy slowly decreases as the threshold rises, then drops significantly once useful information is blurred.

• Sequence length distribution: Token counts per image are concentrated near the maximum and slowly tail off, stopping at around 30% of the maximum value.

• Qualitative behavior: APT assigns the largest (64×64) patches to dark backgrounds, blue sky, and blurry backdrops, the second largest (32×32) to other simple regions, and the smallest patches to faces and in-focus objects. The authors note APT is agnostic to downstream goals, so a texturally simple pink wall would still receive coarse patches even if a user asked about its color.

Methodology in Plain English

Deciding patch sizes. APT defines a fixed number of patch scales. For a base patch size p, patches in level i have size 2^i·p. In the illustrated example with S=3 and p=16, the candidate sizes are 16×16, 32×32, and 64×64. The image is first divided at the coarsest scale, and the Shannon entropy of each patch's binned pixel-intensity distribution is computed using H(P) = −Σ p_i log₂ p_i. Patches whose entropy falls below a tunable threshold τ_i are retained at that scale. The process repeats at finer scales until the smallest patch size is reached, and all remaining regions are assigned the smallest patches. Patches follow a quadtree-like structure on a regular grid.

Embedding mixed patch sizes. Rather than resizing every patch or training separate embedding layers per size, APT resizes each large patch to the base size p×p and passes it through the standard patch embedding layer, while also splitting the original large patch into p×p sub-patches, embedding each of those, and aggregating them with successive Conv2d downsampling layers. The aggregated high-resolution embedding is combined with the resized embedding through a zero-initialized MLP, which lets the model gradually incorporate fine detail without initially degrading the pretrained network's behavior.

Handling variable token counts. Because token counts differ per image, APT uses sequence packing: tokens from a batch are concatenated into one sequence, and a block-diagonal attention mask ensures tokens only attend within their own example. This is supported natively by FlashAttention and xFormers and the paper states it adds no overhead to the network. After the transformer, the sequence is split back into per-example subsequences for class-token extraction or pooling.

Positional encodings. Following NaViT, each base p×p grid is assigned an initial positional encoding, and encodings for larger patches are obtained by interpolation. A 2p patch uses a grid of H/(2p) × W/(2p), a 4p patch uses H/(4p) × W/(4p), and so on.

Dense prediction adaptation. Since APT output cannot be reshaped into a rectangular feature map, tokens representing larger patches are repeated 2^(2i) times to reconstruct a fully differentiable feature map usable by transposed convolutions. For window attention, the image is divided into windows that are multiples of the largest patch size, and each window then holds a variable number of tokens, with attention applied within windows via sequence packing and masks.

Experimental setup. ImageNet experiments used 8× NVIDIA A100s; detection, segmentation, and VQA used 8× NVIDIA RTX A6000. Models were built with the timm library, and MAE pretrained checkpoints were used for full fine-tuning (ViT-B for 100 epochs, ViT-L for 50, batch size 1024, base learning rate 1.5e-3, layer decay 0.75). Short fine-tuning used 1 epoch, layer decay 0.99, learning rate 1e-6, and no augmentations. Downstream tasks were fine-tuned for 5% of the iterations used in each model's original scheme with the learning rate reduced by a factor of 10. Baselines included random masking, resizing-only, vanilla timm ViT, and layer-level merging methods (EViT, ToMe, PPT, DTEM), with "advanced" FlashAttention-compatible re-implementations except for PPT, which depends on attention scores.

Why This Matters

Impact on research. The paper argues that input-level patch reduction is more efficient and reliable than layer-level token merging, and that most prior token-reduction work either uses a fixed reduction ratio that mismatches image complexity or introduces padding and irregular shapes that negate real speedups. By reporting wall-clock and throughput numbers for both training and inference, it shifts the evaluation emphasis away from theoretical FLOP counts.

Real-world applications:

• Vision-language assistants: APT reduces visual tokens in LLaVA-1.5 backbones, accelerating both the vision encoder and cross-attention layers in models that answer questions about images.
• High-resolution object detection: APT is demonstrated at 1536×1536 resolution on COCO with EVA-02 and window attention, relevant to applications such as surveillance, autonomous driving, and industrial inspection.
• Semantic segmentation: Demonstrated on ADE20K at 512 and 640 resolution with UperNet, relevant to medical imaging, satellite analysis, and scene parsing.
• Large-scale image classification backbones: APT can be applied to an existing fine-tuned ViT with one epoch of fine-tuning, lowering the cost of upgrading deployed backbones.

Industry relevance. The method targets the practical cost of training and serving large ViT backbones. Claims of 40% throughput improvement on ViT-L and 50% on ViT-H, plus reduced wall-clock fine-tuning time (for example, 31.4h to 16.9h for ViT-L at 448/14), directly affect GPU-hour budgets. Compatibility with FlashAttention and sequence packing matters because those are standard in production inference stacks, whereas the paper notes several competing methods are not natively compatible with FlashAttention and end up slower than a FlashAttention-enabled vanilla ViT.

Future Directions

• Learned or less hand-crafted patch allocation: The limitations section states APT relies on a hand-crafted heuristic to determine patch sizes, which may not align with downstream user preferences.
• Removing the tuned threshold hyperparameter: The authors note the empirically-tuned threshold adds friction to adoption on downstream tasks, and object detection already needed a different value (2) than the default (5.75 / 4.0).
• Extending to image generation: The paper says APT currently does not support image generation, which operates at very high resolution with large models, and calls it an ideal future application.
• Task-aware patchification: The qualitative results show APT ignores the user's goal (for example, assigning coarse patches to a texturally simple background a user might ask about), raising the question of how to make patch allocation aware of downstream intent.

Target Audience

Researchers and engineers working on efficient transformer architectures, ViT training and inference optimization, and token reduction methods. It is also relevant to practitioners deploying high-resolution vision systems for detection and segmentation, and to those working on vision-language models who want to cut visual token counts without retraining a backbone from scratch. Readers should already understand ViT patch embedding, attention complexity, and standard fine-tuning procedures.

Authors’ abstract

Vision Transformers (ViTs) partition input images into uniformly sized patches regardless of their content, resulting in long input sequence lengths for high-resolution images. We present Adaptive Patch Transformers (APT), which addresses this by using multiple different patch sizes within the same image. APT reduces the total number of input tokens by allocating larger patch sizes in more homogeneous areas and smaller patches in more complex ones. APT achieves a drastic speedup in ViT inference and training, increasing throughput by 40% on ViT-L and 50% on ViT-H while maintaining downstream performance, and can be applied to a previously fine-tuned ViT, converging in as little as 1 epoch. It also significantly reduces training and inference time without loss of performance in high-resolution dense visual tasks, achieving up to 30\% faster training and inference in visual QA, object detection, and semantic segmentation.