Title: Thicker and Quicker: The Jumbo Token for Fast Plain Vision Transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: efficient deep learning, computer vision, vision transformers

Abstract: ViTs are general and accurate, and address many tasks, but ViTs are slow, and
are not always practical when efficiency is key. Existing methods for faster ViTs
design hybrid non-ViT architectures, losing generality, or shrink their tokens,
sacrificing accuracy. While many non-ViT architectures are both fast and accurate,
they cannot flexibly process other input shapes, pre-train by SOTA self-supervised
learning, reduce computation by dropping tokens, and more like ViTs can. We
make ViTs faster by reducing patch token width while increasing global token
width by adding a new Jumbo token. Our wider Jumbo token is processed by its
own wider FFN to increase model capacity. Yet our Jumbo FFN is efficient: it
processes a single token, for speed, and its parameters are shared across all layers,
for memory. Crucially, our Jumbo is attention-only and non-hierarchical, like a
plain ViT, so it is simple, scalable, flexible, and compatible with ViT methods new
and old. Jumbo improves over ViT baselines with Registers from Nano to Large
scales while maintaining speed/throughput on ImageNet-1K (↑0.1−13%). Jumbo
also improves MAE pre-training (↑4.9% linear probing on ImageNet-1K), test-time
adaptation (↑5.2% on ImageNet-C), and time series modeling. Our Jumbo models
even achieve better speed-accuracy trade-offs than specialized non-ViT compute-efficient models, while maintaining plain-ViT compatibility for practicality.

TL;DR: We add a new wider "Jumbo" token to ViTs to improve accuracy and efficiency by adding model capacity in just the right way while preserving generality.