arXiv · 2605.18177
Token-Space Mask Prediction for Efficient Vision Transformer Segmentation
Abstract
Query-based Vision Transformer segmentation models typically reconstruct dense spatial feature maps to predict masks, inheriting design patterns from convolutional architectures. We show that this explicit image-space reconstruction is not required. We introduce TokenMask, a token-space mask head that computes mask logits directly from query-token affinities and performs interpolation in logit space rather than feature space. This reformulation preserves the original linear scoring mechanism while simplifying the computational structure. Across diverse ViT backbones, datasets and segmentation tasks, TokenMask consistently improves efficiency over prior approaches by reducing computational and memory requirements while maintaining competitive accuracy, leading to tangible speedups on NVIDIA Jetson AGX Orin using TensorRT FP16 inference. Overall, TokenMask yields a simpler and more deployment-friendly design for embedded vision systems.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Calvin Galagain, Martyna Poreba, François Goulette. 2026-05-18. Token-Space Mask Prediction for Efficient Vision Transformer Segmentation. https://arxiv.org/abs/2605.18177
Cite the original work for its findings. Save a collection to share your selection of sources.