Category Feature Transformer for Semantic Segmentation

TL;DR

This paper introduces the Category Feature Transformer (CFT), a novel attention-based module for semantic segmentation that improves multi-stage feature aggregation by learning and broadcasting category embeddings, leading to state-of-the-art results.

Contribution

The paper proposes CFT, a new attention-based feature aggregation method that learns category embeddings and enhances semantic segmentation performance.

Findings

Achieves 55.1% mIoU on ADE20K dataset.

Reduces model parameters and computations.

Outperforms previous methods on benchmarks.

Abstract

Aggregation of multi-stage features has been revealed to play a significant role in semantic segmentation. Unlike previous methods employing point-wise summation or concatenation for feature aggregation, this study proposes the Category Feature Transformer (CFT) that explores the flow of category embedding and transformation among multi-stage features through the prevalent multi-head attention mechanism. CFT learns unified feature embeddings for individual semantic categories from high-level features during each aggregation process and dynamically broadcasts them to high-resolution features. Integrating the proposed CFT into a typical feature pyramid structure exhibits superior performance over a broad range of backbone networks. We conduct extensive experiments on popular semantic segmentation benchmarks. Specifically, the proposed CFT obtains a compelling 55.1% mIoU with greatly reduced model parameters and computations on the challenging ADE20K dataset.