Differentiable Hierarchical Visual Tokenization

AI summary

Overview

Research area: Computer Vision — visual tokenization for Vision Transformers (ViTs). Technical level: Advanced (requires familiarity with ViTs, superpixels, segmentation, and model-selection statistics). Scope: This paper proposes ∂HT, an end-to-end differentiable tokenizer that replaces fixed square image patches with hierarchical, pixel-level superpixel tokens selected by information criteria, and evaluates it on classification, dense prediction, and raster-to-vector conversion.

What This Paper Is About

Vision Transformers split images into a fixed grid of square patches, which ignores object boundaries and semantic structure. The paper's goal is a tokenizer that adapts to image content at pixel granularity, can be trained end-to-end with gradients, and can still be dropped into existing pretrained ViT architectures without retraining them from scratch. The authors frame hierarchical pixel grouping as a multi-scale model-selection problem and add a differentiable feature-extraction mechanism so tokenization itself receives learning signals.

Key Contributions

1. Learnable tokens: An end-to-end differentiable tokenization method that adapts to training data and is claimed as the first fully end-to-end learnable tokenizer for ViTs, serving classification and segmentation.
2. Retrofitted tokenizers: A fine-tuning strategy that adapts pretrained ViTs to superpixel tokens at pixel-level granularity, via a reconstruction pretraining objective (Equation 7) that aligns the tokenizer encoder and linear mapping before fine-tuning the backbone.
3. Multiscale model selection: A lattice-theoretic extension of information criteria to multilevel hierarchical partitioning, letting the model pick the most informative partition without manual threshold tuning or gating.
4. Image vectorization: An out-of-the-box route from hierarchical superpixel tokenizers to raster-to-vector graphics conversion, without training specifically for that task.

Main Findings

• From-scratch classification: With identical training regimes, ∂HT ViT-B16 reached 83.2 top-1 on ImageNet1k versus 81.9 for patch ViT-B16; ∂HT ViT-S16 reached 80.0 versus 79.9 for patch ViT-S16 (Table 1, bottom).
• Comparison with other tokenizers (ViT-S16, ImageNet): ∂HT scored 80.0 (Small) and 83.2 (Base), compared with SPFormer at 81.7/82.7, SuiT at 80.9/82.1, SPiT at 75.0/80.4, DEiT at 79.9/81.8, and DEiT3 at 80.4/82.6 (Table 2).
• Retrofitting mostly holds or improves performance: DEiT3-RF ViT-B16 with ∂HT reached 83.9 ImageNet1k Acc@1 versus 82.6 for the DEiT3 ViT-B16 baseline, and 91.2 versus 90.3 on Caltech256; the small-capacity retrofitted ViT-S16 reached 80.1 versus 80.4 (Table 1, top).
• Semantic segmentation without a dense decoder: After fine-tuning with an MLP head per token instead of an upscaling decoder, ∂HT-ViT-B achieved 53.2 mIoU on ADE20k and 48.9 mIoU on COCO-Stuff164k; ∂HT-ViT-S achieved 47.1 mIoU on ADE20k (Table 3).
• Zero-shot salient segmentation: Using TokenCut, ∂HT-B reached 92.4 F_max, 79.9 IoU and 94.2 Acc@1 on ECSSD; 77.9 F_max, 64.4 IoU, 90.5 Acc@1 on DUTS; and 71.9 F_max, 58.7 IoU, 89.8 Acc@1 on DUT-OMRON (Table 4).
• Raster-to-vector conversion: ∂HT achieved 0.00178 MSE, 27.50 PSNR and 0.8541 SSIM, compared with DiffVG-AP at 0.00265 / 26.22 / 0.8494, DiffVG-MC at 0.00316 / 25.54 / 0.8287, and Adobe at 0.00712 / 21.84 / 0.7318 (Table 5), on the five examples provided by Li et al.
• Hyperparameter choice: For reconstruction quality, a Gaussian kernel with dimension d = 8 and the AICC criterion gave the best scores, with MSE 0.06 and SSIM 0.60 (Table 6).
• Component ablations: Removing the CNN projection dropped DEiT3-RF accuracy to 76.9 (down 3.2) and from-scratch to 75.5 (down 4.5); removing mask blending dropped to 71.8 (down 8.3) and 78.9 (down 1.1); removing β-mask features gave 79.7 (down 0.4) and 79.1 (down 0.9); removing kernel aggregation gave 80.0 (down 0.1) and 79.7 (down 0.3). Removing IC pruning gave 80.1 (no change) and 79.9 (down 0.1), but the paper notes it increases tokens per image and computational cost (Table 7).
• Positional embedding resolution: Increasing resolution generally helped, saturating around 24×24 for classification and 48×48 for segmentation; for ViT-B16, 48×48 gave 83.5 Acc@1 and 53.5 mIoU, versus 81.3 and 40.6 at 16×16 (Table 8).
• Scale invariance: ∂HT scaled better to higher resolutions (evaluated at 64 through 768 pixels) without modifying positional embedding resolution; at very low resolutions such as CIFAR's 32×32 the paper states canonical patch tokenization may outperform ∂HT (Figure 6).
• Efficiency trade-off: At 224 resolution with ViT-S16, patch tokenization achieved 192.3k tokens/s with 197 tokens/image (22.1M parameters) versus ∂HT at 94.2k tokens/s with 240 tokens/image (22.3M parameters); at 384 resolution, patch used 577 tokens at 112.4k tokens/s versus ∂HT's 494 tokens at 58.1k tokens/s (Table 9, measured on 8× MI250x).
• Setup details: kNN evaluation took the max score over k ∈ {10, 20, 50, 100, 150, 200}; from-scratch models use local gradient features adding +0.07 GFLOPs and +256 parameters. Training setup specifics are deferred to Appendix D and are not reported in the main text.

Methodology in Plain English

The method has four stages. First, a lightweight CNN encoder projects each pixel into a learned d-dimensional feature space, giving a similarity measure between pixels. Second, starting from individual pixels, the algorithm repeatedly pairs each region with its most similar neighbor and merges them into connected superpixels, doing this repeatedly to build a full hierarchy from pixels up to the whole image. Third, instead of a hand-tuned threshold, the method scores partitions using information criteria — formulas that balance how well a partition fits the image against how complex it is — and picks the best partition; the complexity estimate is derived from the lattice structure of partition space and is inversely related to the number of connected edges inside each region. Fourth, to keep gradients flowing, the tokenizer replaces the mean value of each superpixel region in the original image with a learnable estimate (a "mean-injection" trick), and token features are extracted with a foreground/background mask blending scheme plus a learnable shared background feature β to avoid sparse token representations. Pretrained models can be retrofitted by first pretraining the encoder and linear mapping with a reconstruction loss so that the tokenizer reconstructs the original image, then fine-tuning the backbone.

Why This Matters

Impact on research: The paper reframes tokenization as an adaptive, learnable component rather than a fixed preprocessing step, and claims modular backward compatibility so pretrained ViTs can be retargeted to new tokenization schemes instead of retrained. It also connects superpixel segmentation to information-criterion model selection, and shows tokens transfer to dense tasks and vectorization without task-specific tokenizer training.

Real-world applications (as suggested by the paper's framing):

• Raster-to-vector graphics conversion and image tracing, using the superpixel partitions with potrace.
• Semantic segmentation for scene understanding, including a decoder-free MLP classification of each superpixel.
• Zero-shot salient object detection, where tokens are used directly without post-processing or specialized training.
• Domains cited as promising: medical imaging, video understanding, and document-focused vision-language tasks where fixed patches misalign with text and layout.

Industry relevance: Adaptive tokenization changes the compute-versus-tokens trade-off — the paper reports lower throughput than canonical ViTs at 224 resolution but a narrower gap at 384 resolution, where ∂HT uses fewer tokens than patch tokenization. The retrofitting path matters for organizations holding pretrained models, and the quadratic attention cost in video transformers makes redundancy management commercially relevant.

Future Directions

• Replacing the implicit kernel-weighted aggregation with explicit learnable aggregation, such as graph neural networks, while preserving symmetry and positive semi-definiteness for well-defined edge contractions.
• Extending differentiable tokenization to self-supervised learning, including masked image modeling and DINO-style training.
• Applying adaptive tokenization to vision-language models and document-focused tasks where fixed patches align poorly with heterogeneous text and layout.
• Spatiotemporal superpixel tokenization for video transformers, where quadratic attention makes redundancy management critical.

Target Audience

ViT architecture researchers, efficiency and tokenization specialists, and practitioners who want to retrofit pretrained vision models with content-adaptive tokens. It also suits readers working on superpixel segmentation, dense prediction, image vectorization, and multimodal or video architectures. Because the method relies on information-criterion model selection and lattice-theoretic derivations, readers should be comfortable with statistics and transformer internals.

Authors’ abstract

Vision Transformers rely on fixed patch tokens that ignore the spatial and semantic structure of images. In this work, we introduce an end-to-end differentiable tokenizer that adapts to image content with pixel-level granularity while remaining backward-compatible with existing architectures for retrofitting pretrained models. Our method uses hierarchical model selection with information criteria to provide competitive performance in both image-level classification and dense-prediction tasks, and even supports out-of-the-box raster-to-vector conversion.