Token-Space Mask Prediction for Efficient Vision Transformer Segmentation

• URL: http://arxiv.org/abs/2605.18177v1
• Date: Mon, 18 May 2026 10:20:23 GMT
• Title: Token-Space Mask Prediction for Efficient Vision Transformer Segmentation
• Abstract summary: TokenMask is a token-space mask head that computes masktoken logits directly from querytoken affinities.<n>Across diverse ViT backbones, datasets and segmentation tasks, TokenMask consistently improves efficiency.
• Score: 3.090546888821788
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• GenMask: Adapting DiT for Segmentation via Direct Mask Generation [81.54526445834294]
  We argue that instead of indirect adaptation, segmentation tasks should be trained directly in a generative manner.<n>We introduce timesteps sampling strategy for binary masks that emphasizes extreme noise levels for segmentation and moderate noise for image generation.<n>We present GenMask, a DiT trains to generate black-and-white segmentation masks as well as colorful images in RGB space under the original generative objective.
  arXiv  Detail & Related papers  (2026-03-25T03:52:05Z)
• AFFMAE: Scalable and Efficient Vision Pretraining for Desktop Graphics Cards [8.026364090097951]
  Masked Autoencoders (MAE) reduce computation by encoding only visible tokens.<n>We introduce AFFMAE, a masking-friendly hierarchical pretraining framework built on adaptive, off-grid token merging.
  arXiv  Detail & Related papers  (2026-02-18T07:58:47Z)
• The Missing Point in Vision Transformers for Universal Image Segmentation [17.571552686063335]
  We introduce ViT-P, a two-stage segmentation framework that decouples mask generation from classification.<n>ViT-P serves as a pre-training-free adapter, allowing the integration of various pre-trained vision transformers.<n>Experiments across COCO, ADE20K, and Cityscapes datasets validate the effectiveness of ViT-P.
  arXiv  Detail & Related papers  (2025-05-26T10:29:13Z)
• High-Quality Mask Tuning Matters for Open-Vocabulary Segmentation [109.19165503929992]
  We present MaskCLIP++, which uses ground-truth masks instead of generated masks to enhance the mask classification capability of CLIP.<n>After low-cost fine-tuning, MaskCLIP++ significantly improves the mask classification performance on multi-domain datasets.<n>We achieve performance improvements of +1.7, +2.3, +2.1, +3.1, and +0.3 mIoU on the A-847, PC-459, A-150, PC-59, and PAS-20 datasets.
  arXiv  Detail & Related papers  (2024-12-16T05:44:45Z)
• ColorMAE: Exploring data-independent masking strategies in Masked AutoEncoders [53.3185750528969]
  Masked AutoEncoders (MAE) have emerged as a robust self-supervised framework. We introduce a data-independent method, termed ColorMAE, which generates different binary mask patterns by filtering random noise. We demonstrate our strategy's superiority in downstream tasks compared to random masking.
  arXiv  Detail & Related papers  (2024-07-17T22:04:00Z)
• Mask Propagation for Efficient Video Semantic Segmentation [63.09523058489429]
  Video Semantic baseline degradation (VSS) involves assigning a semantic label to each pixel in a video sequence. We propose an efficient mask propagation framework for VSS, called SSSS. Our framework reduces up to 4x FLOPs compared to the per-frame Mask2Former with only up to 2% mIoU on the Cityscapes validation set.
  arXiv  Detail & Related papers  (2023-10-29T09:55:28Z)
• Pre-training with Random Orthogonal Projection Image Modeling [32.667183132025094]
  Masked Image Modeling (MIM) is a powerful self-supervised strategy for visual pre-training without the use of labels. We propose an Image Modeling framework based on Random Orthogonal Projection Image Modeling (ROPIM) ROPIM reduces spatially-wise token information under guaranteed bound on the noise variance and can be considered as masking entire spatial image area under locally varying masking degrees.
  arXiv  Detail & Related papers  (2023-10-28T15:42:07Z)
• MixMask: Revisiting Masking Strategy for Siamese ConvNets [23.946791390657875]
  This work introduces a novel filling-based masking approach, termed textbfMixMask. The proposed method replaces erased areas with content from a different image, effectively countering the information depletion seen in traditional masking methods. We empirically validate our framework's enhanced performance in areas such as linear probing, semi-supervised and supervised finetuning, object detection and segmentation.
  arXiv  Detail & Related papers  (2022-10-20T17:54:03Z)
• SegViT: Semantic Segmentation with Plain Vision Transformers [91.50075506561598]
  We explore the capability of plain Vision Transformers (ViTs) for semantic segmentation. We propose the Attention-to-Mask (ATM) module, in which similarity maps between a set of learnable class tokens and the spatial feature maps are transferred to the segmentation masks. Experiments show that our proposed SegVit using the ATM module outperforms its counterparts using the plain ViT backbone.
  arXiv  Detail & Related papers  (2022-10-12T00:30:26Z)
• SODAR: Segmenting Objects by DynamicallyAggregating Neighboring Mask Representations [90.8752454643737]
  Recent state-of-the-art one-stage instance segmentation model SOLO divides the input image into a grid and directly predicts per grid cell object masks with fully-convolutional networks. We observe SOLO generates similar masks for an object at nearby grid cells, and these neighboring predictions can complement each other as some may better segment certain object part. Motivated by the observed gap, we develop a novel learning-based aggregation method that improves upon SOLO by leveraging the rich neighboring information.
  arXiv  Detail & Related papers  (2022-02-15T13:53:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.