PEM: Prototype-based Efficient MaskFormer for Image Segmentation (2402.19422v3)

Abstract: Recent transformer-based architectures have shown impressive results in the field of image segmentation. Thanks to their flexibility, they obtain outstanding performance in multiple segmentation tasks, such as semantic and panoptic, under a single unified framework. To achieve such impressive performance, these architectures employ intensive operations and require substantial computational resources, which are often not available, especially on edge devices. To fill this gap, we propose Prototype-based Efficient MaskFormer (PEM), an efficient transformer-based architecture that can operate in multiple segmentation tasks. PEM proposes a novel prototype-based cross-attention which leverages the redundancy of visual features to restrict the computation and improve the efficiency without harming the performance. In addition, PEM introduces an efficient multi-scale feature pyramid network, capable of extracting features that have high semantic content in an efficient way, thanks to the combination of deformable convolutions and context-based self-modulation. We benchmark the proposed PEM architecture on two tasks, semantic and panoptic segmentation, evaluated on two different datasets, Cityscapes and ADE20K. PEM demonstrates outstanding performance on every task and dataset, outperforming task-specific architectures while being comparable and even better than computationally-expensive baselines.

• The paper introduces PEM as a novel method that uses prototype-based masked cross-attention to reduce computational load in segmentation tasks.
• The architecture employs an efficient pixel decoder with a multi-scale FPN and deformable convolutions to capture detailed semantic features.
• PEM achieves competitive benchmark results, including a 61.1% PQ on Cityscapes, demonstrating both efficiency and high segmentation quality.

Prototype-based Efficient MaskFormer: A New Direction in Efficient Image Segmentation

Introducing PEM

In the evolving landscape of computer vision, image segmentation has been a vital area of research, contributing significantly to advancements in autonomous driving, medical image analysis, and various real-world applications requiring precise object boundaries and classifications. Recent works have demonstrated the superior capability of transformer-based architectures to address both semantic and panoptic segmentation tasks. However, these high-performance models come with a substantial computational cost, making them less feasible for deployment on devices with limited resources. This paper introduces a novel architecture, Prototype-based Efficient MaskFormer (PEM), specifically designed to address the efficiency bottleneck while maintaining, and in some cases exceeding, the segmentation performance of more computationally demanding models.

Architectural Innovations of PEM

Prototype-based Masked Cross-Attention Mechanism

PEM incorporates a prototype-based masked cross-attention mechanism, a significant part of its efficiency. Conventional transformer decoders in segmentation tasks often handle dense visual features, leading to excessive computational requirements. PEM, on the other hand, selects a single, most-representative feature (prototype) for each object descriptor, drastically reducing the complexity. This process not only improves efficiency but also aligns with the interpretability, as it highlights the most pertinent features for segmentation tasks.

Efficient Pixel Decoder

PEM introduces an efficient multi-scale Feature Pyramid Network (FPN), crucial for extracting high-resolution semantic content from images. By leveraging deformable convolutions and context-based self-modulation, PEM dynamically focuses on relevant regions and incorporates global context into the feature maps, promoting a balanced detail capture across scales. This clever design allows PEM to achieve a high semantic understanding similar to its transformer-based counterparts but at a fraction of the computational cost.

Benchmarking and Performance

PEM's performance was rigorously evaluated on two benchmark datasets: Cityscapes and ADE20K, covering semantic and panoptic segmentation tasks. Remarkably, PEM not only outperformed task-specific architectures but also showed competitive or superior results compared to established, computationally intensive models. For instance, on Cityscapes, PEM achieved a 61.1% PQ (Panoptic Quality) score, surpassing many existing models while offering significant efficiency improvements, evidenced by its higher frames per second (FPS) rate and lower Floating Point Operations Per Second (FLOPs). Its ability to maintain high performance with reduced computational demand was consistently observed across different tasks and datasets.

Implications and Future Directions

The introduction of PEM marks a significant advancement in image segmentation research, particularly in the quest for efficiency without compromising performance. This work not only demonstrates the potential of prototype-based approaches in reducing computational requirements but also sets a new precedent for future research in segmentation and other related tasks in computer vision.

PEM's architecture opens avenues for further exploration into more compact and efficient models capable of running on edge devices, thus broadening the applicability of advanced AI technologies in everyday scenarios. Further research could explore the scalability of the prototype-based approach to other vision tasks, the adaptation of PEM to different backbones or datasets, and the investigation into real-world deployment scenarios where computational resources are a critical factor.

In conclusion, the Prototype-based Efficient MaskFormer (PEM) offers a promising direction towards balancing the trade-off between efficiency and performance in image segmentation tasks. Its innovative approach not only paves the way for more sustainable AI models but also extends the reach of advanced computer vision technologies to a wider range of applications and devices.