HCFormer: Unified Image Segmentation with Hierarchical Clustering

We view the attention function [58] from the clustering perspective. Let $q\in\mathbb{R}^{C\times N_{q}}$ and $k\in\mathbb{R}^{C\times N_{k}}$ be a query and a key. $N_{q}$ and $N_{k}$ are the number of tokens for the query and key, and $C$ denotes a feature dimension. Then, the attention is defined as follows:

where $\top$ denotes the transpose of a matrix and $s$ denotes a scale parameter that is usually defined as $\sqrt{C}$ [17, 58]. $\mathrm{Softmax}_{\text{row}}(\cdot)$ denotes the row-wise softmax function. The attention function is generally defined with a query, a key, and a value, but the value is omitted here for simplicity.

When $s\rightarrow 0$, eq. (1) is equivalent to the following maximization problem:

where $\langle\cdot,\cdot\rangle$ denotes the Frobenius inner product. We provide the detailed derivation in Appendix A. This maximization problem is interpreted as the clustering problem for $q$ with $k$ as cluster prototypes. With an inner product as a similarity function, $A_{nm}=1$ if $n$-th query, $q_{n}$, has the maximum similarity to $m$-th key, $k_{m}$, among all key tokens; otherwise, $A_{nm}=0$. In other words, each row of $A$ indicates the index of the cluster assigned to $q_{n}$, known as the assignment matrix. Thus, eq. (2) is an assignment problem, which is a special case or part of the clustering (e.g., agglomerative hierarchical clustering with a hierarchical level of one and one-step $K$-means clustering), and the attention in eq. (1) solves the relaxed problem of eq. (2). This insight indicates that we can accomplish clustering after downsampling using the attention in a differentiable form by corresponding the pixels in the feature maps before and after downsampling to queries and keys, respectively.

We show the computational scheme of the proposed clustering module in Fig. 2. We obtain an intermediate feature map and its downsampled feature map from a backbone model for clustering. These are fed into layer normalization [2] and a convolution layer with a kernel size of 1$\times$1, as in the transformer block [17, 58]. We define the obtained feature maps as $F^{(i)}\in\mathbb{R}^{C_{F}\times N^{(i)}}$ and $F_{d}^{(i)}\in\mathbb{R}^{C_{F}\times N_{d}^{(i)}}$, where $i\in\mathbb{N}$ denotes a downsampling factor of the spatial resolution that corresponds to the number of applied downsampling layers, and $N^{(i)}$ and $N_{d}^{(i)}$ denote the number of pixels in the feature maps before and after downsampling (i.e., $N^{(i)}=\frac{H}{2^{i}}\frac{W}{2^{i}}$ and $N^{(i)}_{d}=\frac{H}{2^{i+1}}\frac{W}{2^{i+1}}$, where $H$ and $W$ are the height and width of an input image). $C_{F}$ is the number of channels set to 128 in our experiment. Note that we assume that the downsampling halves height and width respectively, and the $\ell_{2}$-norm of the feature vector of each pixel is normalized as 1 to make the inner product the cosine similarity.

Then, the proposed clustering is defined as follows:

where $A^{(i)}\in(0,1)^{N^{(i)}\times N^{(i)}_{d}}$ is an assignment matrix. We define a scale parameter, $s^{(i)}\in\mathbb{R}$, as a trainable parameter. As already described in Sec. 2.1, when $s^{(i)}\rightarrow 0$, eq. (3) corresponds to the clustering problem for $F^{(i)}$ with $F_{d}^{(i)}$ as cluster prototypes. Thus, by computing eq. (3) at every downsampling layer, HCformer hierarchically groups pixels, as shown in Fig. 1.

The obtained feature map from a backbone model is fed into the segmentation head used in [11] (Fig. 4). Specifically, the feature map is fed into the transformer decoder with trainable queries $Q\in\mathbb{R}^{C_{q}\times N_{m}}$, and the mask queries $\mathcal{E}_{\text{mask}}\in\mathbb{R}^{C_{m}\times N_{m}}$ are generated. The number of queries, $N_{m}$, is a hyperparameter. Note that the transformer “decoder” differs from the pixel decoder because one of its roles is to generate the mask queries, not to upsample the feature map. The feature map is also mapped into the $C_{m}$-dimensional space by a linear layer, and we define them as $\mathcal{E}_{\text{feature}}^{(i=5)}\in\mathbb{R}^{C_{m}\times N^{(i)}}$. Then, the output of the segmentation head is computed as follows:

Note that we assume that the scale $i$ is 5 because the conventional backbone has five downsampling layers. This segmentation head can be viewed as the clustering, which groups $N^{(i)}$ pixels in the feature map into $N_{m}$ clusters.²²2Eq. (4) uses the sigmoid function, not the softmax function. Thus, eq. (4) does not correspond to eq. (2). However, in the post-processing, the mask with the maximum confidence is selected as the prediction, meaning that the pipeline of the segmentation head, including the post-processing, corresponds to eq. (2). Further details can be found in Appendix E. Therefore, we also refer to $M^{(i=5)}$ as the assignment matrix. Unlike the previous work [11], we feed a low-resolution feature map into the segmentation head (eq. (4)), which contributes reduction of FLOPs in the segmentation head.

As an output of HCFormer, we obtain the assignment matrices $\{A^{(i)}\}_{i}$ obtained from the backbone and the output of the segmentation head $M^{(5)}$. We need to decode them as segmentation masks of an input image for evaluation.

One step of the decoding process is defined as a matrix multiplication between $A^{(i)}$ and $M^{(i+1)}$:

By writing down the calculations in $A^{(i)}$ explicitly, we can notice that it is equivalent to the attention function [58], although queries, keys, and values are obtained from different layers. The segmentation mask (i.e., the correspondence between pixels in an input image and the mask queries) is computed by multiplying all $\{A^{(i)}\}_{i}$ by $M^{(5)}$ as $\prod_{i}A^{(i)}M^{(5)}$.

From the hard clustering perspective, eq. (2), we can view multiplying $A^{(i)}$ as copying the cluster’s representative value to its elements. Thus, this decoding process will be accurate if each cluster is composed of pixels that are annotated with the same ground-truth label. In other words, we can evaluate the accuracy of the decoding from the hard clustering perspective by the undersegmentation error [51] that is used for assessing superpixel segmentation.

Let $N$ be the number of pixels, and then the computational costs of the proposed clustering are $O(N^{2})$, which is the same complexity as the common attention function [58] and is intractable for high-resolution images. Various studies exist on reducing complexity [38, 50, 4, 65, 60, 57, 47], and we approach it with the local attention strategy.

We divide the feature map before downsampling, $F^{(i)}$, into 2$\times$2 windows. The window size is determined by a stride of downsampling, which is assumed to be 2 in this work. Each window corresponds to the pixel in the feature map after downsampling, $F^{(i)}_{d}$. Then, the attention is computed between a pixel in the window and the corresponding pixel in $F^{(i)}_{d}$ and its surrounding eight pixels, which is the same technique used in superpixel segmentation [23, 1]. We illustrate this local attention in Fig. 3. As a result, the complexity decreases to $O(N)$.

The clustering and decoding processes with this local attention can be easily implemented using deep learning frameworks. We show an example PyTorch [45] implementation in Appendix B.

We show a baseline architecture, MaskFormer [11], and our model using the proposed clustering module in Fig. 4. The proposed model removes the pixel decoder from the baseline architecture, and instead, the decoding block, eq. (2.3), is integrated for obtaining segmentation masks. Optionally, we build a six-layer transformer encoder after the backbone, as in MaskFormer [11]. The mask queries are fed into multi-layer perceptrons and classified into $K$ classes.

The flexibility of the downsampling is important for our clustering module because the downsampled pixels are the cluster prototypes that should have representative values of clusters. However, unlike transformer-based backbones, the receptive field of CNN-based backbones is limited even if the deeper model is used [41, 63]. Thus, the CNN-based backbones may not sample effective pixels for the clustering module. To alleviate this problem, we replace downsampling layers to which the proposed clustering module is attached with DCNv2 [74] with a stride of two, which deforms kernel shapes and modulates weight values. This replacement is only adopted for a CNN-based backbone. We verify its effect in Appendix F.2.

Since FCNs [39] have been proposed, deep neural networks have become a de facto standard approach for image segmentation. As main segmentation tasks, semantic, instance, and panoptic segmentation are studied.

Semantic segmentation is often defined as per-pixel classification tasks, and various FCN-based methods have been proposed [72, 6, 3, 7, 8, 63]. Instance segmentation has the object detection perspective. Thus, many methods [19, 15, 34, 14, 31] use an object proposal module as the task-specific module, which is used for instance discrimination. Panoptic segmentation has been proposed in  [26], which is a task combining semantic and instance segmentation. In early work, panoptic segmentation is approached with two separated modules for generating semantic and instance masks and then fusing them [26, 25, 9, 64]. To simplify the framework, Panoptic FCN [32] unifies the modules by using dynamic kernels. Panoptic FCN generates kernels from a feature map and produces segmentation masks by cross-correlation between the kernels and feature maps. As a similar approach, cross-attention with the transformer decoder is adopted in other panoptic segmentation methods [59, 5, 11, 10, 33, 70]. Such approaches also simplify the panoptic segmentation pipeline.

While many studies investigate task-specific modules and architectures, they cannot be applied to other tasks. Thus, recent work investigates the unified architectures [11, 71, 10]. These methods use cross-attention-based approaches that can generate segmentation masks regardless of the task definition. They have demonstrated effectiveness in the various segmentation tasks and achieved comparable or better results than the specialized models. Our study also focuses on such unified architectures and builds our model based on MaskFormer [11]. Note that, as seen in eq. (4), the segmentation head used in the unified models can be viewed as a clustering module; namely, the unified models would also be clustering-based frameworks. However, they predict segmentation masks directly from per-pixel features (i.e., not a hierarchical clustering method), and these methods do not allow us to visualize and evaluate the intermediate results.

Several studies define the segmentation tasks as a clustering problem [27, 44, 24, 16, 28, 70], especially in the proposal-free instance segmentation methods. The main focus of these methods is to learn an effective representation for clustering, and their approach is not hierarchical.

Many studies on image segmentation focus on the segmentation head or the pixel decoder. In particular, various decoder architectures have been proposed, thereby improving segmentation accuracy [36, 46, 3, 49, 22, 30, 62, 8]. The decoder would make the segmentation problem complex: models using the decoder have to solve both the upsampling and pixel classification problems internally. Such models suffer from an error in the upsampling process of the pixel decoder. But, it is difficult to know whether the prediction error is due to upsampling or classification because upsampling is performed in the high-dimensional feature space.

Our decoding process can be viewed as an attention-based decoder, as shown in eq. (2.3), and there are several methods using the attention or similar modules for the pixel decoder [30, 29, 35, 66, 22]. However, existing methods do not have a clustering perspective and do not allow us to interpret, visualize, and evaluate the attention map as the clustering results. In contrast, since HCFormer is built on the clustering perspective, we can interpret and visualize intermediate results as the clustering results and evaluate their accuracy by comparing ground-truth labels.

There are some methods that do not use the pixel decoder. For example, FCN-32s [39], the simplest variant of FCNs, does not use the trainable decoder, although the generated masks are low-resolution and the segmentation accuracy is somewhat low. Several methods [6, 69, 72, 7] use dilated convolution [69, 6] to keep the resolution in the backbone, which expands the convolution kernel by inserting holes between its consecutive elements and replacing the stride with the dilation rate. The methods using dilated convolution improve the segmentation accuracy at the expense of FLOPs.

Hierarchical clustering approaches are sometimes used to solve image segmentation. For example, some previous methods [13, 18, 21, 55, 67, 68, 54, 53, 56, 48] group pixels into small segments using superpixel segmentation [1, 52, 23] as pre-processing; the obtained segments are then merged or classified to obtain the desired cluster. Such a hierarchical approach often reduces computational costs and improves segmentation accuracy compared to directly clustering or classifying pixels.

A recently proposed method, called GroupViT [66], is a bottom-up hierarchical clustering method. However, GroupViT is specialized in vision transformers [17] and semantic segmentation with text supervision. In contrast, our method can be incorporated into almost all multi-scale backbone models, such as ResNet [20] and Swin Transformer [38], and used for arbitrary image segmentation tasks.

We implement our model on the Mask2Former author’s implementation.³³3https://github.com/facebookresearch/Mask2Former We use the transformer decoder with the masked attention [10], and the number of decoder layers is set to 8. Note that we do not use multi-scale feature maps in the transformer decoder because HCFormer does not have the pixel decoder. Thus, HCFormer is built on top of MaskFormer rather than on top of Mask2Former since the transformer decoder used in HCFormer and MaskFormer is independent of the pixel decoder. We describe the detailed difference between transformer decoders of HCFormer and Mask2Former in Appendix C.

We train HCFormer with almost the same protocol as in Mask2Former [10]; we use the binary cross-entropy loss and the dice loss [42] as the mask loss: $\mathcal{L}_{\text{mask}}=\lambda_{\text{ce}}\mathcal{L}_{\text{ce}}+\lambda_{\text{dice}}\mathcal{L}_{\text{dice}}$. The training loss combines mask loss, classification loss, and an additional regularization term for $s^{(i)}$ in eq. (3): $\mathcal{L}_{\text{mask}}+\lambda_{\text{cls}}\mathcal{L}_{\text{cls}}+\lambda_{\text{reg}}\mathcal{L}_{\text{reg}}$, where $\mathcal{L}_{\text{cls}}$ is a cross-entropy loss and $\mathcal{L}_{\text{reg}}=\sum_{i}|s^{(i)}|$. $\lambda_{\text{reg}}$ is set to $0.1$. The regularization enforces the proposed attention-based clustering, eq. (3), to be the hard clustering, eq. (2). Following [10], we set $\lambda_{\text{ce}}=5.0$, $\lambda_{\text{dice}}=5.0$, and $\lambda_{\text{cls}}=2.0$. Another training protocol is the same as for Mask2Former [10] (e.g., optimizer, its hyperparameters, and the number of training iterations). The details can be found in Appendix D.

We use the same post-processing as in [11]: we multiply class confidence and mask confidence and use the output as the confidence score. Then, the mask with the maximum confidence score is selected as the predicted mask. The details can be found in Appendix E.

The clustering module defined in eq. (3) is incorporated into every downsampling layer, except for those in the stem block of ResNet [20] and the patch embedding layer in Swin Transformer [38]. Thus, the hierarchical level is 3 (i.e., $\{A^{(2)},A^{(3)},A^{(4)}\}$ are computed). The relation between the hierarchical level and downsampling layers using the clustering module is shown in Fig. II in the appendix.

As baseline methods, we choose the recently proposed methods for unifying segmentation tasks, MaskFormer [11], K-Net [71], and Mask2Former [10], and specialized models for each task [32, 70, 63, 19, 61]. We compare our method to the baselines with several backbones, ResNet-50 [20], Swin-S [38], and Swin-L [38]. Note that ResNet-50 and Swin-S were pretrained with ImageNet-1K, and Swin-L was pretrained with ImageNet-22K. The training procedure and evaluation metrics are following  [10]. We trained models three times and reported medians.

We first show the evaluation results for panoptic segmentation on the COCO validation dataset [37] in Tab. 1. HCFormer outperforms the unified architectures, MaskFormer and K-Net, and Panoptic FCN for all metrics with fewer parameters. The additional transformer encoder (i.e., HCFormer+) boosts the segmentation accuracy for the relatively small backbone models and outperforms CMT-Deeplab with fewer parameters and FLOPs. Although the transformer decoder used in HCFormer is different from that used in MaskFormer in this comparison, we verify HCFormer still outperforms MaskFormer even when MaskFormer’s transformer decoder is applied (see Appendix F.3).

Mask2Former shows the best accuracy, though it sacrifices the FLOPs. According to [10], using multi-scale features in the transformer decoder improves 1.7 PQ for Mask2Former with the ResNet-50 backbone (i.e., PQ of Mask2Former without the multi-scale feature maps is 50.2, which is the same as PQ of HCFormer+ with the ResNet-50 backbone in Tab. 1). Thus, we believe the gap in PQ between HCFormer and Mask2Former is due to the design of the transformer. In other words, PQ of HCFormer is comparable to that of Mask2Former in the same setting, but HCFormer shows lower FLOPs and has the interpretability. The detailed analysis is described in Appendix C.

Tab. 2 shows the average precision on instance segmentation tasks. HCFormer also outperforms MaskFormer and the specialized model, Mask R-CNN and SOLOv2. Notably, HCFormer significantly improves AP${}^{\text{L}}$ from MaskFormer, which indicates HCFormer can generate well-aligned masks for large objects. For recovering the large objects by the pixel decoder, the model needs to capture the long-range dependence, which may be difficult for the simple pixel decoder. Thus, the conventional pipelines need a well-design decoder to recover the large objects, such as that used in Mask2Former. However, since HCFormer groups pixels locally and hierarchically, it may not need long-range dependence to capture the large objects, compared to the conventional methods. As a result, HCFormer outperforms MaskFormer, especially in terms of AP${}^{\text{L}}$.

Tabs. 3 and 4 show the results on semantic segmentation. On ADE20k, HCFormer outperforms MaskFormer and the specialized model, SegFormer, and Mask2Former shows the best mIoU. However, on Cityscapes [12] (Tab. 4), mIoU of HCFormer is worse than that of SegFormer, although HCFormer improves mIoU from MaskFormer. Unlike COCO and ADE20K, the Cityscapes dataset contains only urban street scenes, and there are many thin or small objects, such as poles, signs, and traffic lights. MaskFormer and HCFormer do not have specialized modules to capture such small and thin objects; hence such objects would be missed in an intermediate layer.

For further analysis, we visualize intermediate clustering results and a predicted mask using HCFormer with Swin-L in Fig. 5. From the visualization of undersegmentation error [51], we find that the model groups pixels well except for the boundaries in intermediate clustering, but it misclassifies clusters for the small objects in this image (e.g., poles and persons in the center). HCFormer cannot extract semantically discriminative features for such small objects, although the obtained features are discriminative in terms of clustering. Thus, to solve errors for the small objects, we may improve the transformer decoder or stack more layers on the backbone for the coarsest scale. We believe that this analysis plays one of the roles in interpretability, and conventional models do not allow for this type of analysis. We present further results and analysis on other datasets in Appendix F.4.

To investigate the effect of the hierarchical level and the more effective architecture, we evaluate PQ, FLOPs, and the number of parameters of our model with various hierarchical levels and different numbers of transformer decoder layers. We use the COCO dataset and the ResNet-50 backbone for evaluation. The additional ablation study can be found in Appendix F.3.

The results are shown in Tabs. 5 and 6. The higher the hierarchical level, the higher the resolution of the predicted masks, increasing PQ and FLOPs. The hierarchical level of 2 would be preferred regarding the balance between accuracy and computational costs in practice, though we adopt the hierarchical level of 3 in Sec. 4.2. In particular, PQ in our model is still comparable to that of MaskFormer [11] even when the hierarchical level is 2.

Our model does not work well with only one decoder layer, as with other models using the transformer decoder [11, 10, 5]. Regarding the balance between accuracy and computational costs, four or eight layers would be suitable for our model.