PaperlayerView exact revision on arXiv

Source-linked AI summary

Your ViT is Secretly an Image Segmentation Model

Tommie Kerssies, Niccolò Cavagnero, Alexander Hermans, Narges Norouzi, Giuseppe Averta, Bastian Leibe, Gijs Dubbelman, Daan de Geus

arXiv:2503.19108v1cs.CV

TL;DR

ViT segmentation commonly relies on task-specific components, raising whether sufficiently large and extensively pretrained ViTs need them. The paper introduces EoMT, which repurposes plain ViT blocks and mask annealing for segmentation. EoMT achieves competitive accuracy with substantially higher speed and simpler architecture, including 56.0 PQ at 128 FPS with ViT-L versus 54.4 PQ at 32 FPS for ViT-Adapter + M2F with ViT-B.

  • Problem

    The paper asks whether convolutional adapters and decoder modules remain necessary for state-of-the-art ViT image segmentation as model size and pre-training scale up.

  • Method

    EoMT repurposes plain ViT blocks for feature extraction, query interaction, and segmentation prediction, while mask annealing removes masked attention during inference.

  • Results

    56.0 PQ at 128 FPS with EoMT ViT-L compares with 54.4 PQ at 32 FPS for ViT-Adapter + M2F ViT-B, demonstrating higher quality and speed in this comparison.

  • Takeaways & Limitations

    Task-specific components become increasingly redundant with larger models and more pre-training, supporting investment in ViT scaling and pre-training rather than added architectural complexity.

Abstract

from arXiv · show

Vision Transformers (ViTs) have shown remarkable performance and scalability across various computer vision tasks. To apply single-scale ViTs to image segmentation, existing methods adopt a convolutional adapter to generate multi-scale features, a pixel decoder to fuse these features, and a Transformer decoder that uses the fused features to make predictions. In this paper, we show that the inductive biases introduced by these task-specific components can instead be learned by the ViT itself, given sufficiently large models and extensive pre-training. Based on these findings, we introduce the Encoder-only Mask Transformer (EoMT), which repurposes the plain ViT architecture to conduct image segmentation. With large-scale models and pre-training, EoMT obtains a segmentation accuracy similar to state-of-the-art models that use task-specific components. At the same time, EoMT is significantly faster than these methods due to its architectural simplicity, e.g., up to 4x faster with ViT-L. Across a range of model sizes, EoMT demonstrates an optimal balance between segmentation accuracy and prediction speed, suggesting that compute resources are better spent on scaling the ViT itself rather than adding architectural complexity. Code: https://www.tue-mps.org/eomt/.

1. Introduction

The paper investigates whether large, extensively pretrained ViTs still need task-specific segmentation components. It introduces EoMT, which repurposes plain ViT blocks and mask annealing to achieve competitive segmentation with substantially greater speed and simplicity.

  • Motivation: Existing ViT segmentation systems add convolutional adapters, pixel decoders, and Transformer decoders to produce and fuse multi-scale features for mask and class predictions.The adapter extracts multi-scale features, the pixel decoder fuses them, and the Transformer decoder uses object queries for predictions.
  • Motivation: Large-scale pre-training and increased ViT capacity are hypothesized to make these additional components increasingly unnecessary for extracting dense semantic information and conducting segmentation.The paper tests this hypothesis across pre-training types and model sizes.
  • EoMT: EoMT repurposes ViT blocks so they extract image features, interact with learnable object queries, and predict a mask and class label for each query.The resulting architecture differs only minimally from a plain ViT.
  • EoMT: Mask annealing enables masked attention during early training but gradually removes it, allowing masked-attention-free inference without the severe mismatch caused by simply disabling attention at inference.The strategy is designed to preserve the benefits of masked attention while improving inference efficiency.
  • Results: Removing task-specific components has only a minimal impact on segmentation accuracy with a large pretrained model such as DINOv2, while reducing computational requirements and latency.The paper argues that compute should instead support scaling the ViT and pre-training.
  • Results: 54.4 PQ at 32 FPS for ViT-Adapter + M2F with ViT-B contrasts with 56.0 PQ at 128 FPS for EoMT with ViT-L.This comparison illustrates EoMT’s stronger speed–quality balance despite using a larger model.

2. Related Work

Related work established ViT architectures, large-scale visual pre-training, and unified Mask Transformer pipelines for segmentation. This paper builds on that context by pursuing a minimalistic ViT-based alternative to more complex segmentation architectures.

  • Image segmentation: Image segmentation assigns pixel-level segments using masks and class labels, with semantic, instance, and panoptic segmentation differing in their treatment of classes and object instances.Panoptic segmentation combines instance-level segments for things with class-level segments for stuff.
  • Image segmentation: Mask Transformer frameworks unified semantic, instance, and panoptic segmentation through learnable object queries representing either stuff classes or thing instances.These frameworks support a shared architecture and training pipeline across the three tasks.
  • Vision Transformers: ViTs divide images into fixed-size patches, embed them as tokens, and process them with Transformer blocks rather than the multiscale downsampling structure typical of CNNs.The passage contrasts fixed-resolution ViT processing with CNN downscaling steps.
  • Vision Transformers: Multi-scale features and local processing are commonly considered beneficial for segmentation, motivating Transformer variants that add local attention or token downsampling.Such variants depart from the plain ViT architecture and can limit access to large-scale pre-training advances.
  • Large-scale visual pre-training: Large-scale visual pre-training has progressed from supervised ImageNet initialization toward weakly supervised and self-supervised learning on massive datasets.The passage identifies vision foundation models using masked image modeling as a recent direction.

3. Towards Encoder-only Mask Transformer

The section traces how conventional ViT segmentation components are progressively removed, culminating in EoMT, which repurposes plain ViT blocks and learned queries for segmentation.

  • 3.1. Preliminaries: ViTs divide images into fixed-size patches, project them into tokens, and process those tokens through Transformer blocks.Final patch tokens can be reordered into spatial features whose resolution is determined by the patch size.
  • 3.1. Preliminaries: Standard ViT segmentation systems add adapters for convolutional biases and multi-scale features, followed by pixel and Transformer decoders for predictions.The adapter extracts multi-scale features, the pixel decoder fuses them, and the Transformer decoder uses learned queries to predict masks and classes.
  • 3.1. Preliminaries: Mask2Former uses masked cross-attention so each query attends only to the image region corresponding to its predicted segmentation mask.Intermediate masks and class labels are predicted before cross-attention, and the masks constrain subsequent query-to-image attention.
  • 3.2. Removing task-specific components: The authors remove the adapter, pixel decoder, and multi-scale processing, retaining a single-scale Transformer decoder that attends to ViT output features.The simplified pyramid is passed directly to the decoder, then only F_vit = F_16 is used for query attention while upscaling produces high-resolution masks.
  • 3.3. Encoder-only Mask Transformer: EoMT removes the dedicated decoder by concatenating learned queries with patch tokens after L_1 blocks and jointly processing them through the final L_2 ViT blocks.The same self-attention operation jointly provides query interaction and query-to-patch information transfer, followed by class and mask prediction.
  • 3.3. Encoder-only Mask Transformer: During training, intermediate masks constrain query-to-patch self-attention, while the model progressively removes this masking during training.This preserves masked-attention behavior during learning while enabling the final architecture to operate without a dedicated decoder or masked attention at inference.

4. Experiments

Experiments evaluate EoMT across segmentation tasks, pre-training regimes, model sizes, and efficiency settings. Results show that simpler ViT-based designs retain competitive accuracy while improving speed.

  • Stepwise component removal: EoMT reduces PQ from 57.1 to 56.0 while increasing prediction speed by 4.4× relative to ViT-Adapter + Mask2Former.Removing the adapter, pixel decoder, and multi-scale features reduces performance by only 0.4 PQ and makes the model 2.2× faster.
  • Pre-training: DINOv2 and EVA-02 pre-training limit EoMT’s PQ gap to 1.1 and 1.2, compared with 3.9 and 6.1 for ImageNet-21K and ImageNet-1K.These results support the reported link between large-scale pre-training and reduced need for task-specific components.
  • Model size: The PQ gap between EoMT and ViT-Adapter + Mask2Former narrows from 5.8 for ViT-S to 0.7 for ViT-g as model size increases.EoMT improves its relative performance with larger models and strong pre-training while remaining faster at all tested sizes.
  • Efficiency trade-off: EoMT with ViT-L achieves PQ 56.0 at 128 FPS, versus PQ 54.4 at 32 FPS for ViT-Adapter + M2F with ViT-B.This comparison gives EoMT a better PQ-versus-FPS trade-off despite using a larger model.
  • Cross-task evaluation: Across benchmarks, EoMT is comparable in mIoU while up to 4.4× faster for semantic segmentation and achieves 48.8 AP at 30 FPS versus 47.6 AP at 29 FPS for instance segmentation.For panoptic segmentation on COCO, EoMT is on par with state-of-the-art methods and up to 2.1× faster.
  • Mask annealing: Mask annealing enables removing masked attention during inference while roughly maintaining PQ and more than doubling inference speed.Training with masked attention and disabling it only at inference causes a severe performance drop, whereas annealing phases masking out during training.
  • ViT compatibility: ALGM increases EoMT throughput by up to 31% without affecting mIoU, whereas ViT-Adapter + M2F gains no throughput from token merging.The contrast is attributed to EoMT’s proximity to the plain ViT and the adapter’s token-unmerging overhead.

5. Conclusion

The paper concludes that scaling ViT model size and pre-training makes task-specific segmentation components increasingly redundant. EoMT consequently offers a simpler, faster, and scalable segmentation design.

  • Task-specific segmentation components become increasingly redundant as ViT model size and pre-training are scaled up.
  • EoMT uses a plain ViT to deliver high accuracy and speed with a significantly simpler design than existing segmentation models.
  • The authors state that compute resources should be directed toward scaling the ViT and pre-training rather than adding architectural complexity.

A.1. Models

The appendix specifies the ViT backbones, pretrained model weights, query counts, and block configurations used for the experiments.

  • Pretrained model adaptation: Pretrained patch embeddings and positional embeddings are resized to support 16 × 16 patches and different input sizes.
  • Queries: Panoptic and instance models use K = 200 queries, while semantic models use K = 100 queries.
  • Block configuration: EoMT adds query tokens only to the final L2 ViT blocks, with L2 set to 3 for ViT-S and ViT-B, 4 for ViT-L, and 5 for ViT-g.

A.2. Training

Training follows Mask2Former’s loss formulation and augmentation setup, with task-specific schedules and pretrained-model initialization details.

  • Augmentation: Training uses random flipping, scale jittering, padding, cropping, and selected color jittering, with jitter ranges varying by segmentation task.
  • Model and implementation setup: The implementation uses Mask2Former-compatible class and model specifications, including ViT-Adapter + Mask2Former configurations and pretrained timm models.
  • Loss function: The total loss combines cross-entropy for class logits with binary-cross entropy and Dice losses for mask logits.The loss terms are weighted by λ_bce, λ_dice, and λ_ce.
  • Optimization: A two-stage linear learning-rate warm-up first initializes randomly initialized parameters for 500 iterations, then pretrained parameters for 1000 iterations.

A.3. Evaluation

The evaluation uses task-specific inference procedures for segmentation and measures throughput consistently across models. ALGM improves EoMT throughput during inference without additional training.

  • Inference setup: Panoptic and instance segmentation use padded square-image inference, while semantic segmentation uses sliding-window crops.The preprocessing differs by task to accommodate the model’s fixed input format.
  • Efficiency measurements: FPS is measured on identical hardware across existing works and the authors’ models, while FLOPs come from papers or are calculated directly.Compilation and precision settings introduce specific measurement caveats for some baselines.
  • Token merging: ALGM adaptively merges tokens according to image complexity and reports throughput in images per second using batch size 32.For batching, the minimum mergeable-token count across each batch is applied to all images.
  • Token merging: Inference-only ALGM increases EoMT throughput with no additional training required.The improvement comes from applying token merging during inference and processing image batches.

B. Detailed Experimental Analysis

The detailed experiments examine component removal, mask annealing, query-processing depth, and qualitative confidence behavior. Results support efficient configurations while preserving segmentation quality.

  • Component analysis: Panoptic and instance analyses separately report thing/stuff PQ and large/medium/small-object AP after removing task-specific components.These breakdowns reveal how component removal affects different segmentation categories and object sizes.
  • Mask annealing: Mask annealing is effective for both EoMT and the ViT-Adapter + M2F baseline.The strategy is therefore not limited to the encoder-only architecture.
  • Query-processing depth: ViT-L achieves its highest PQ around L2 = 4, while prediction speed is not significantly affected by changing L2.EoMT performance remains stable across the tested query-processing configurations.
  • OOD confidence: EoMT assigns low confidence to the full OOD object, whereas ViT-Adapter + M2F does so only partially in the qualitative comparison.The visualization uses darker colors for lower confidence and evaluates Cityscapes-trained models on BRAVO.
  • Mask annealing: 56.2 and 57.1 PQ are obtained with masking during training and inference, compared with 53.2 and 54.0 without masking.Mask annealing reduces PQ only slightly, to 56.0 and 56.8, while removing inference masking.

C. Out-of-distribution Confidence Estimation

EoMT is evaluated for distinguishing novel OOD objects from in-distribution regions using confidence scores. DINOv2-based EoMT models outperform comparison models in this setting.

  • OOD generalization: DINOv2-based models such as EoMT significantly outperform non-ViT Swin models in OOD generalization despite similar in-distribution performance.The comparison is discussed for models evaluated across OOD and ID behavior.
  • Confidence metric: OOD regions are novel object classes absent from the training data, and AUPRCOOD measures whether models assign them lower confidence.Higher AUPRCOOD reflects better separation of unseen objects from ID regions.
  • Results: 89.7 AUPRCOOD for EoMT exceeds 68.7 for ViT-Adapter + M2F and 56.8 for Swin + M2F.The visualization likewise shows consistently low confidence across the OOD object for EoMT.

D. Qualitative Examples

The qualitative examples compare EoMT with ViT-Adapter + M2F on COCO panoptic segmentation using DINOv2-g and 1280 × 1280 inputs.

  • Qualitative comparison: Figure C visualizes qualitative COCO panoptic-segmentation predictions from ViT-Adapter + M2F and EoMT.The examples use DINOv2-g with a 1280 × 1280 input size.
Loading 2503.19108v1…