# Token-Space Mask Prediction for Efficient Vision Transformer Segmentation

arXiv: https://arxiv.org/abs/2605.18177 · HTML: https://mlsanity.com/paper/2605.18177/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 18 references that are in the map.
- First posted on arXiv 18 May 2026.

## Builds on (most cited first, 15 of 15)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://mlsanity.com/paper/2107.06278/) (2021)
- [Segmenter: Transformer for Semantic Segmentation](https://mlsanity.com/paper/2105.05633/) (2021)
- [OneFormer: One Transformer to Rule Universal Image Segmentation](https://mlsanity.com/paper/2211.06220/) (2022)
- [Your ViT is Secretly an Image Segmentation Model](https://mlsanity.com/paper/2503.19108/) (2025)
- [StructToken : Rethinking Semantic Segmentation with Structural Prior](https://mlsanity.com/paper/2203.12612/) (2022)
- [Head-Free Lightweight Semantic Segmentation with Linear Transformer](https://mlsanity.com/paper/2301.04648/) (2023)
- [Vision Transformers: From Semantic Segmentation to Dense Prediction](https://mlsanity.com/paper/2207.09339/) (2022)
- [Is Semantic SLAM Ready for Embedded Systems ? A Comparative Survey](https://mlsanity.com/paper/2505.12384/) (2025)

## Abstract (arXiv)

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.