Learning Hierarchical Patch Splitting Policies for Faster Vision Transformers

AI summary

Overview

Research area: Computer vision / efficient Vision Transformer (ViT) architectures, specifically adaptive patch tokenization (input-level token reduction).

Technical level: Advanced — the core idea is accessible, but the method relies on reinforcement learning (PPO), a Lagrangian dual for budget control, and multi-scale patch embedding.

Scope: A single paper presenting SPLIT, a learned, budget-constrained patch-splitting policy that shortens ViT input token sequences for classification, detection, segmentation, and visual question answering.

What This Paper Is About

Standard Vision Transformers cut an image into a fixed uniform grid of patches, so the token count depends on image resolution rather than image content — every region gets the same budget, whether it is an object or empty background. Prior adaptive-patchification work (e.g., APT, Quadformer) reduces tokens by assigning coarser patches to "simple" regions, but decides where using pixel-level heuristics such as local entropy or downsampling error. This paper asks whether the decision of where to spend tokens can instead be learned from the downstream task itself, so that fine tokens land on regions the model actually finds meaningful.

Key Contributions

1. SPLIT framework. An adaptive patchification framework that uses coarse tokens as lightweight "semantic anchors" and learns which patches to split into four finer sub-patches, producing a mixed-resolution token sequence without any pixel-level heuristic.
2. Budget-constrained RL formulation. Patch refinement is cast as a constrained decision-making problem (maximize task performance subject to a token budget ratio ρ) and solved with PPO on hard sampled masks, with the token budget enforced through a Lagrange multiplier λ updated by dual descent.
3. Warm-start via patch–child dissimilarity. The decision network is initialized from deterministic pseudo-labels based on cosine dissimilarity between a parent patch embedding and its children, with a learnable threshold τ_split that tracks the same token budget — reducing RL instability.
4. Broad transfer. The same recipe adapts pretrained ViTs with minimal fine-tuning (1 epoch warm-start + 1 epoch PPO in the classification setting) and is demonstrated on ImageNet classification, COCO detection, ADE20K segmentation, and LLaVA-1.5 VQA.

Main Findings

• Classification accuracy at reduced tokens: SPLIT matches full-resolution ViT baselines while using 30% fewer tokens. At 224 resolution: SPLIT-B reaches 85.0 accuracy vs. 85.1 for ViT-B at 3635 vs. 3310 img/s (+10%); SPLIT-L reaches 87.8 (ViT-L: 87.9) at 1107 vs. 883 img/s (+25%); SPLIT-H reaches 88.2 (ViT-H: 88.3) at 533 vs. 441 img/s (+21%). Random patch selection is consistently worse in accuracy (83.7, 86.9, 87.4) despite larger speedups.
• Higher-resolution classification: At 336, SPLIT-L is 88.1 accuracy (APT-L 88.1, ViT-L 88.2) at 570 img/s (+44%); SPLIT-H is 88.4 (APT-H 88.4, ViT-H 88.5) at 281 img/s (+61%).
• Object detection (COCO, 1536×1536, EVA-02 + ViTDet head): SPLIT-B obtains 58.84 mAP / 77.80 AP50 at 4.23 img/s (+16%) vs. 58.93 / 77.85 for the baseline at 3.66 img/s, and better than APT (58.79 / 77.65 at +13%). SPLIT-L obtains 62.24 / 80.78 at +25% vs. 62.28 / 80.80 for the baseline.
• Semantic segmentation (ADE20K, EVA-02 + UperNet): At 512, SPLIT reaches 86.53 aAcc / 59.65 mIoU at 7.43 img/s (+15%) vs. 86.67 / 59.77 baseline. At 640, 86.75 / 59.99 at +15% vs. 86.83 / 60.05. SPLIT uses roughly 35% fewer tokens here vs. about 30% for APT.
• VQA (LLaVA-1.5): With the language model frozen, SPLIT-7B runs at 4.37 img/s (+18%) and reports 78.6 on VQAv2, 61.7 on GQA, 69.0 on SQA^I, 65.2 on MMB, 58.9 on MMB^C, 31.3 on MMV — above the 78.5 / 61.0 / 67.8 / 64.6 / 58.1 / 30.7 baseline — while being slightly lower on POPE (86.6 vs. 86.9) and MME (1497.3 vs. 1510.1). The 13B variant runs at 2.66 img/s (+20%) and improves SQA^I (73.4 vs. 72.7), MMB^C (64.0 vs. 63.4) and MMV (36.3 vs. 35.4) but is slightly below baseline on VQAv2 (79.6 vs. 80.0), GQA (63.1 vs. 63.2), VQA^T (60.2 vs. 61.2), POPE (86.9 vs. 87.1) and MME (1504.9 vs. 1530.6). APT is slightly faster (+22% / +23%) but falls short of baseline on most metrics.
• Semantic alignment measured directly: On ADE20K at a matched budget (~0.7 tokens), SPLIT routes fine 16×16 tokens to 0.907 of object-class patches and 0.945 of class-boundary patches, vs. 0.808 / 0.833 for APT (entropy) and 0.697 / 0.698 for random.
• Head-to-head with APT: On ViT-L at 336, SPLIT dominates APT across the accuracy-vs-token-retention Pareto frontier, with the gap widening as the budget tightens. The per-image sequence-length distributions show APT splits images fully more often; SPLIT redistributes that mass to shorter sequences.
• Overhead is small: With no sequence reduction, SPLIT adds only a 3% overhead versus vanilla ViT from multi-scale patch embedding, the decision network, and packing logic; savings from fewer tokens dominate once reduction is on.
• Adaptive λ beats fixed λ: A continuously updated λ tracks the target token budget reliably, whereas fixed λ = 0.01 requires laborious tuning and shows a larger train/validation control-error gap.
• PPO matters on dense tasks: The appendix reports the PPO phase improves detection mAP monotonically over a warm-start-only policy, from 58.06 to 58.84 at a matched budget.

Methodology in Plain English

The image is first chopped into large patches. A small "decision network" then looks at each patch and its four would-be smaller children and asks: keep this patch as one token, or split it into four? This is repeated over several stages, so the final sequence mixes big, medium, and small tokens.

The decision network is a lightweight module: it uses the parent patch embedding as a query and the four child embeddings as keys/values in a cross-attention layer, concatenates the pooled summary with the parent, and passes it through a two-layer MLP to output a split probability. During training a decision is sampled from that probability; at inference the rule is simply "split if probability > 0.5."

Because choosing patches is a discrete, non-differentiable operation, the authors do not backpropagate through it. Instead they treat it as a one-step reinforcement learning problem and use PPO. For each image they sample several masks, score each by a combined cost that mixes downstream error with the fraction of tokens used (weighted by λ), and give a mask a positive advantage if it beats the other masks sampled for the same image. A separate dual-descent update raises λ when the policy exceeds the token budget and lowers it when the policy is under budget — making λ the single knob controlling how aggressively the model splits.

To avoid unstable training from a random start, the decision network is first warm-started with pseudo-labels: split a patch if any of its children is sufficiently dissimilar to the parent in embedding space (cosine dissimilarity above a learnable threshold). The backbone is fine-tuned during this warm-start and then frozen while the policy is optimized with PPO.

Why This Matters

The paper reframes an efficiency problem as an alignment problem: instead of asking "which pixels look complicated?" it asks "which regions does the model's own objective say matter?" That is a conceptual shift for the adaptive-tokenization literature, and the measured fine-token recall numbers (0.907 vs. 0.808 for object patches) give a concrete, quantitative way to test the claim rather than only showing qualitative heatmaps. It also shows that input-level tokenization can be learned with policy gradients, avoiding the biased straight-through relaxations that earlier mixed-scale tokenizers used.

Real-world applications:

• High-resolution detection and autonomy: COCO experiments at 1536×1536 with window attention map directly onto settings where resolution is expensive but localization accuracy is safety-critical.
• Vision-language assistants: Applying SPLIT to LLaVA-1.5's vision backbone improves throughput while holding or improving most VQA scores, relevant for serving multimodal models at scale.
• Medical and satellite imaging: Large, mostly-homogeneous images with small regions of interest are the regime where content-dependent tokenization saves the most compute; the paper's design targets exactly this, though it does not evaluate these domains.
• Edge and on-device inference: A 3% overhead when no reduction is applied, and speedups of +10% to +61% when reduction is on, make the cost model easy to reason about for deployment.

Industry relevance: the method attaches to pretrained checkpoints from timm, EVA-02, and LLaVA-1.5 with short fine-tuning runs (1 epoch warm-start plus 1 epoch PPO for classification, 5% of total iterations for detection and segmentation), which lowers the adoption cost for teams that already have ViT backbones in production.

Future Directions

• Ratio-conditioned policies. The reduction ratio is fixed at training time; changing the desired speedup requires a short retraining run. Conditioning the policy on the target ratio at inference would remove this constraint, as the authors note in their limitations.
• Generative models. The authors state that image generation operates at much higher resolutions and larger model scales — the regime where the speedup would be largest — but this is untested.
• Task-agnostic / self-supervised policies. Training SPLIT with self-supervised objectives to produce a policy reusable across tasks, and characterizing how a policy trained on one task behaves on another, is explicitly left open.
• Extending beyond two-to-three splitting stages and comparing decision-network cost models — the runtime breakdown attributes SPLIT's edge over APT at similar token reduction partly to lower decision-making overhead versus entropy-based multi-scale scoring when extended past two patch sizes, which invites further work on cheap importance scoring.

Target Audience

Researchers and engineers working on efficient transformer architectures, adaptive tokenization, or token merging/pruning; practitioners deploying ViTs at high resolution or serving multimodal models who need throughput gains without retraining from scratch; and reinforcement-learning researchers interested in constrained policy optimization applied to architectural decisions. Readers should be comfortable with ViT basics and standard RL notation; the paper is not an introductory survey.

Authors’ abstract

Vision Transformers (ViTs) are bottlenecked by the length of their input token sequence. Standard patchification fixes this length with a single patch size applied across the image, regardless of regional content. Prior work on adaptive patch sizing reduces token count by assigning coarser patches to visually simple regions, but relies on noisy pixel-level heuristics such as local entropy to make patchification decisions. We argue that this compute allocation should be grounded in the model's interpretation of task relevant information, not in raw pixel statistics. We introduce Split Policy Learned Image Tokenization which learns to allocate token resolution based on what the model finds meaningful, producing shorter token sequences that are better matched to task-relevant content. SPLIT matches the accuracy of full-resolution ViT classification baselines while using 30% fewer tokens. The same gains hold on dense prediction: SPLIT achieves competitive accuracy on COCO and ADE20K, demonstrating that semantically aligned patchification generalises well beyond classification.