# kMaX-DeepLab: k-means Mask Transformer

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 101 of its 118 references that are in the map.
- First posted on arXiv 8 Jul 2022.

## Builds on (most cited first, 30 of 101)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Deformable DETR: Deformable Transformers for End-to-End Object Detection](https://mlsanity.com/paper/2010.04159/) (2020)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [Xception: Deep Learning with Depthwise Separable Convolutions](https://mlsanity.com/paper/1610.02357/) (2016)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Rethinking Atrous Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1706.05587/) (2017)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [ViViT: A Video Vision Transformer](https://mlsanity.com/paper/2103.15691/) (2021)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)

## Built on by (most cited first, 12 of 12)

- [PEM: Prototype-based Efficient MaskFormer for Image Segmentation](https://mlsanity.com/paper/2402.19422/) (2024)
- [Textual Query-Driven Mask Transformer for Domain Generalized Segmentation](https://mlsanity.com/paper/2407.09033/) (2024)
- [SERNet-Former: Semantic Segmentation by Efficient Residual Network with Attention-Boosting Gates and Attention-Fusion Networks](https://mlsanity.com/paper/2401.15741/) (2024)
- [Beyond Accuracy: Benchmarking Cross-Task Consistency in Unified Multimodal Models](https://mlsanity.com/paper/2604.25072/) (2026)
- [A Simple and Generalist Approach for Panoptic Segmentation](https://mlsanity.com/paper/2408.16504/) (2024)
- [High-level Feature Guided Decoding for Semantic Segmentation](https://mlsanity.com/paper/2303.08646/) (2023)
- [The revenge of BiSeNet: Efficient Multi-Task Image Segmentation](https://mlsanity.com/paper/2404.09570/) (2024)
- [Enhancing Temporal Consistency in Video Editing by Reconstructing Videos with 3D Gaussian Splatting](https://mlsanity.com/paper/2406.02541/) (2024)
- [VPNeXt -- Rethinking Dense Decoding for Plain Vision Transformer](https://mlsanity.com/paper/2502.16654/) (2025)
- [Scenes as Objects, Not Primitives: Instance-Structured 3D Tokenization from Unposed Views](https://mlsanity.com/paper/2606.29513/) (2026)
- [From Pixels to Objects: A Hierarchical Approach for Part and Object Segmentation Using Local and Global Aggregation](https://mlsanity.com/paper/2409.01353/) (2024)
- [Incorporating Feature Pyramid Tokenization and Open Vocabulary Semantic Segmentation](https://mlsanity.com/paper/2412.14145/) (2024)

## Abstract (arXiv)

The rise of transformers in vision tasks not only advances network backbone designs, but also starts a brand-new page to achieve end-to-end image recognition (e.g., object detection and panoptic segmentation). Originated from Natural Language Processing (NLP), transformer architectures, consisting of self-attention and cross-attention, effectively learn long-range interactions between elements in a sequence. However, we observe that most existing transformer-based vision models simply borrow the idea from NLP, neglecting the crucial difference between languages and images, particularly the extremely large sequence length of spatially flattened pixel features. This subsequently impedes the learning in cross-attention between pixel features and object queries. In this paper, we rethink the relationship between pixels and object queries and propose to reformulate the cross-attention learning as a clustering process. Inspired by the traditional k-means clustering algorithm, we develop a k-means Mask Xformer (kMaX-DeepLab) for segmentation tasks, which not only improves the state-of-the-art, but also enjoys a simple and elegant design. As a result, our kMaX-DeepLab achieves a new state-of-the-art performance on COCO val set with 58.0% PQ, Cityscapes val set with 68.4% PQ, 44.0% AP, and 83.5% mIoU, and ADE20K val set with 50.9% PQ and 55.2% mIoU without test-time augmentation or external dataset. We hope our work can shed some light on designing transformers tailored for vision tasks. TensorFlow code and models are available at https://github.com/google-research/deeplab2 A PyTorch re-implementation is also available at https://github.com/bytedance/kmax-deeplab