PaperlayerView exact revision on arXiv

Source-linked AI summary

OneFormer: One Transformer to Rule Universal Image Segmentation

Jitesh Jain, Jiachen Li, MangTik Chiu, Ali Hassani, Nikita Orlov, Humphrey Shi

arXiv:2211.06220v2cs.CV

TL;DR

Existing panoptic architectures do not fully unify semantic, instance, and panoptic segmentation because they require separate task-specific training for best performance. OneFormer introduces task-conditioned joint training, task tokens, and query-text contrastive learning in a single universal model. The model outperforms individually trained Mask2Former models across all three tasks and major benchmarks while reducing resource requirements.

  • Problem

    Existing panoptic architectures require individual training on semantic, instance, or panoptic segmentation to achieve their best performance, leaving truly universal segmentation unresolved.

  • Method

    OneFormer uses task-conditioned joint training, task input tokens, and query-text contrastive loss to train one transformer model across semantic, instance, and panoptic segmentation.

  • Results

    One jointly trained OneFormer model outperforms individually trained specialized Mask2Former models across all three segmentation tasks on major datasets.

  • Takeaways & Limitations

    OneFormer reduces training time, weight storage, and inference hosting requirements down to a third, making image segmentation more accessible.

Abstract

from arXiv · show

Universal Image Segmentation is not a new concept. Past attempts to unify image segmentation in the last decades include scene parsing, panoptic segmentation, and, more recently, new panoptic architectures. However, such panoptic architectures do not truly unify image segmentation because they need to be trained individually on the semantic, instance, or panoptic segmentation to achieve the best performance. Ideally, a truly universal framework should be trained only once and achieve SOTA performance across all three image segmentation tasks. To that end, we propose OneFormer, a universal image segmentation framework that unifies segmentation with a multi-task train-once design. We first propose a task-conditioned joint training strategy that enables training on ground truths of each domain (semantic, instance, and panoptic segmentation) within a single multi-task training process. Secondly, we introduce a task token to condition our model on the task at hand, making our model task-dynamic to support multi-task training and inference. Thirdly, we propose using a query-text contrastive loss during training to establish better inter-task and inter-class distinctions. Notably, our single OneFormer model outperforms specialized Mask2Former models across all three segmentation tasks on ADE20k, CityScapes, and COCO, despite the latter being trained on each of the three tasks individually with three times the resources. With new ConvNeXt and DiNAT backbones, we observe even more performance improvement. We believe OneFormer is a significant step towards making image segmentation more universal and accessible. To support further research, we open-source our code and models at https://github.com/SHI-Labs/OneFormer

1. Introduction

Existing panoptic architectures still require task-specific training, leaving universal segmentation unresolved. OneFormer addresses this with a single task-conditioned, jointly trained model that achieves state-of-the-art performance across semantic, instance, and panoptic segmentation.

  • Existing panoptic architectures require separate training for semantic, instance, and panoptic segmentation to achieve their best performance.
  • OneFormer is a universal framework trained once on one panoptic dataset to outperform existing state-of-the-art methods across all three segmentation tasks.
  • A task input token conditions the model on the active task, supporting task-guided training and task-dynamic inference.
  • Query-text contrastive loss uses task- and label-derived text to help the model learn inter-task and inter-class distinctions.
  • Task-conditioned joint training uniformly samples semantic, instance, and panoptic ground-truth domains derived from panoptic annotations.
  • OneFormer achieves state-of-the-art performance on ADE20K, Cityscapes, and COCO across all three segmentation tasks, with further gains from ConvNeXt and DiNAT backbones.

2. Related Work

Image segmentation has evolved from separate semantic and instance approaches toward panoptic and transformer-based architectures. These advances still leave a gap: the strongest unified systems require individual task training, which OneFormer targets with task-conditioned queries and contrastive learning.

  • Image Segmentation: Semantic segmentation has increasingly used transformer-based mask classification, while instance segmentation predicts a binary mask and class label for each object.
  • Image Segmentation: Panoptic segmentation combines amorphous stuff regions with distinct thing objects, but its models can remain behind specialized semantic and instance systems.
  • Universal Image Segmentation: Mask2Former established state-of-the-art performance across all three tasks but requires individual training on each task for best performance.
  • Transformer-based Architectures: OneFormer follows transformer encoder-decoder and query-based mask-classification designs while adding task-guided queries and query-text contrastive loss.
  • Transformer-based Architectures: The framework forms text descriptions from ground-truth mask classes and pads them with task-specific no-object entries.

3. Method

OneFormer unifies semantic, instance, and panoptic segmentation through task-conditioned joint training, task-guided queries, and query-text contrastive learning. Its single jointly trained model is designed to address task-dependent query meanings while retaining a shared architecture.

  • 3.1. Task Conditioned Joint Training: OneFormer jointly trains semantic, instance, and panoptic segmentation by uniformly sampling tasks and deriving task-specific labels from one set of panoptic annotations.This task-conditioned strategy uses a single annotation source while producing task-specific training labels.
  • 3.1. Task Conditioned Joint Training: Task input tokens of the form “the task is {task}” condition OneFormer on the segmentation task during training and inference.The token is mapped to a task representation used to condition the architecture.
  • 3.2. Query Representations: Semantic, instance, and panoptic tasks require different query meanings, motivating task-guided query initialization and concatenation with the task token.The queries must represent class-level masks, thing instances, or both, depending on the task.
  • 3.2. Query Representations: OneFormer forms text queries from padded ground-truth-derived text and image queries from task-token-guided transformer updates.Text queries represent image segments, while object queries provide image-based representations used by the decoder.
  • 3.3. Task Guided Contrastive Queries: The query-text contrastive loss aligns object and text queries to learn inter-task and inter-class distinctions and reduce category misclassifications.The loss contains object-to-text and text-to-object components, with a learnable temperature scaling the similarity logits.
  • 3.4. Other Architecture Components: The architecture uses multi-scale deformable pixel decoding, a transformer decoder, class prediction, mask generation, and classification, binary cross-entropy, and dice losses.Decoder queries interact with multi-scale features through masked cross-attention, self-attention, and feed-forward updates.

4. Experiments

OneFormer is evaluated across three datasets and all three segmentation tasks, using one jointly trained model and task-conditioned components. It outperforms specialized Mask2Former models on major benchmarks, while ablations show benefits from task conditioning, text context, contrastive learning, and joint training.

  • Main Results: OneFormer achieves state-of-the-art performance across semantic, instance, and panoptic segmentation on ADE20K, Cityscapes, and COCO with a single jointly trained model.
  • Main Results: OneFormer outperforms Mask2Former on Cityscapes, improving PQ by +0.6% and AP by +1.9% with the Swin-L backbone.
  • Main Results: OneFormer reaches 68.5% PQ with ConvNeXt-L and 46.7% AP with ConvNeXt-XL on Cityscapes.
  • Ablation Studies: The task-conditioned architecture improves AP, while learnable text context improves PQ; removing the task token causes a −2.7% AP drop.
  • Ablation Studies: OneFormer’s joint-training design improves PQ, AP, and mIoU over a jointly trained Mask2Former baseline by +1.1%, +2.2%, and +0.8%, respectively.
  • Ablation Studies: Query-text contrastive loss reduces category confusions between visually similar classes, producing more accurate regions than Mask2Former in qualitative examples.

5. Conclusion

The conclusion presents OneFormer as a train-once transformer framework that unifies semantic, instance, and panoptic segmentation. The authors report stronger performance than individually trained specialized models while reducing deployment resources.

  • OneFormer unifies semantic, instance, and panoptic segmentation with one architecture, one model, and training on one dataset.
  • The single jointly trained model outperforms individually trained specialized Mask2Former models across all three segmentation tasks and major datasets.
  • Training time, weight storage, and inference hosting requirements can be reduced to a third, making image segmentation more accessible.

A. Implementation Details

Implementation uses multi-scale deformable attention and reports task-specific evaluation metrics, alongside selected training and ablation settings. The appendix experiments identify query count, text-context count, and contrastive-loss weight choices.

  • Implementation: The pixel decoder uses six MSDeformAttn modules over 1/8, 1/16, and 1/32-resolution feature maps, producing 1/4-resolution features with hidden dimension 256.
  • Training Settings: COCO training uses AdamW with base learning rate 0.0001, weight decay 0.05, 100 epochs, and 1024×1024 fixed-size crops after scale augmentation.
  • Ablation Settings: N = 150 performs best in the COCO query-count ablation with the ResNet-50 backbone.
  • Ablation Settings: Nctx = 16 performs best in the learnable text-context ablation.
  • Ablation Settings: λQ↔Qtext = 0.5 gives the best performance in the contrastive-loss weight ablation.

B. Additional Ablations

Additional ablations examine query count and contrastive-loss weighting across datasets and backbones. The best query count varies with dataset and backbone, while the supplied passage does not state the best loss weight numerically.

  • Number of Queries: N = 150 performs best on COCO with the ResNet-50 backbone after 50 training epochs.
  • Number of Queries: With Swin-L, N = 250 is best on ADE20K and Cityscapes, whereas N = 150 is best on COCO.
  • Number of Queries: With the smaller ResNet-50 backbone, N = 150 is optimal on ADE20K.
  • Contrastive Loss: The experiments also ablate contrastive-loss weight on COCO using a ResNet-50 backbone.

C. Individual Training

OneFormer remains effective when trained individually or jointly, while joint training gives it an advantage over Mask2Former. The individual-training setup conditions the task token and restricts supervision to the corresponding task labels.

  • The individual-training analysis uses ADE20K with an R50 backbone and a fixed 160k-iteration, batch-size-16 protocol.
  • OneFormer outperforms Mask2Former under every individual and joint training strategy.The comparison uses R50 on ADE20K for 160k iterations with batch size 16.
  • Joint training gives OneFormer the highest PQ, AP, and mIoU scores, whereas Mask2Former suffers a significant performance drop.
  • Individual training sets the task-token value to panoptic, semantic, or instance for the corresponding training samples.Only the matching ground-truth labels are used for each individual-training strategy.

D. Analysis on the Task-Dynamic Nature of OneFormer

Changing the task token changes OneFormer’s outputs to reflect instance, semantic, or panoptic segmentation. Quantitative and qualitative results show distinct handling of “thing” objects and “stuff” regions across task settings.

  • Setting the task token to instance makes OneFormer suppress “stuff” masks while preserving distinct “thing” instances.On Cityscapes, PQSt drops to 0.0% and PQTh changes by only −0.2% relative to panoptic inference.
  • Setting the task token to semantic produces amorphous “stuff” masks without separate masks for individual “thing” objects.The semantic setting shows sizable drops in PQ, PQTh, and AP while PQSt improves by +0.2%.
  • Qualitative examples show task-dependent errors: instance inference can misclassify “stuff,” while semantic inference merges distinct objects into one amorphous mask.

E. Comparison to SOTA Methods at System-Level for Image Segmentation

Across ADE20K, Cityscapes, and COCO, a single OneFormer model outperforms Mask2Former on panoptic, instance, and semantic segmentation without additional system-level training techniques or data.

  • OneFormer outperforms Mask2Former on all three segmentation tasks across ADE20K, Cityscapes, and COCO.The comparison uses the same OneFormer models referenced in the main text and does not apply additional system-level training techniques or data.

E.1. SOTA Systems on ADE20K val

OneFormer reaches state-of-the-art results on ADE20K across segmentation tasks and backbones, while the evaluation notes a COCO annotation discrepancy requiring adjusted instance annotations.

  • SOTA Systems on ADE20K val: 37.8% AP makes Swin-L OneFormer state of the art for instance segmentation without extra training data.
  • SOTA Systems on ADE20K val: 51.5% PQ makes DiNAT-L OneFormer state of the art for panoptic segmentation without extra training data.
  • SOTA Systems on ADE20K val: The ADE20K individual-versus-joint comparison uses R50 models trained for 160k iterations, and reports variance across training strategies.
  • SOTA Systems on ADE20K val: 58.3% single-scale mIoU and 58.8% multi-scale mIoU are reported as new state-of-the-art scores without extra training data.
  • SOTA Systems on ADE20K val: 68.5% PQ is achieved by ConvNeXt-L OneFormer with single-scale inference, while ConvNeXt-XL OneFormer reaches 46.7% AP.
  • SOTA Systems on ADE20K val: DiNAT-L OneFormer reaches 68.1% mIoU, while Swin-L OneFormer achieves a 64.4% PQTh score without extra training data.
  • SOTA Systems on ADE20K val: The COCO evaluation found serious discrepancies between ground-truth panoptic and instance annotations.For fair comparison, instance annotations were generated from panoptic annotations during evaluation.

F. Analysis on Discrepancy between Instance and Panoptic Annotations in COCO

COCO’s separately developed instance and panoptic annotations contain substantial object-level discrepancies, complicating unified segmentation training and evaluation. OneFormer therefore evaluates instance AP using instance ground truths derived from panoptic annotations.

  • Annotation discrepancies: COCO’s separately developed panoptic and instance annotations contain significant discrepancies, unlike Cityscapes and ADE20K annotations.The discrepancies arise because COCO maintains separate annotation sets, whereas the other datasets combine semantic and instance annotations to generate panoptic annotations.
  • Train-set examples: COCO instance annotations merge distinct objects, including ties with people and dogs with boats, while panoptic annotations separate them.
  • Validation-set examples: COCO instance annotations omit multiple people and motorcycles, while panoptic annotations include them or merge background groups into one mask.
  • Impact: These discrepancies create a significant barrier to developing and evaluating a unified image segmentation model.
  • Evaluation choice: OneFormer evaluates instance AP using instance ground truths derived from panoptic annotations because its training uses only panoptic annotations.
Loading 2211.06220v2…