Masked-attention Mask Transformer for Universal Image Segmentation

Idea

Based on MaskFormer,

• Apply masked attention in the decoder
• Use multi-scale high-resolution features

Architecture

Difference from MaskFormer:

Masked Attention

Constrain cross-attention within predicted mask regions, other than the whole image

From ($l$ is the layer index, and ${\text{X}}_l$ is the query features at $l$-th layer)

$${\text{X}}_l=\mathrm{softmax}({\text{Q}}_l{\text{K}}_l^T){\text{V}}_l+{\text{X}}_{l-1},{\text{X}}_l\in {\mathbb{R}}^{N\times C}$$

to

$${\text{X}}_l=\mathrm{softmax}({\mathcal{M}}_{l-1}+{\text{Q}}_l{\text{K}}_l^T){\text{V}}_l+{\text{X}}_{l-1},{\text{X}}_l\in {\mathbb{R}}^{N\times C}$$

where

$${\mathcal{M}}_{l-1}(x,y)=\{\begin{array}{ll}0,& \text{if }{\text{M}}_{l-1}(x,y)=1\\ -\infty & \text{otherwise}\end{array}$$

here ${\text{M}}_{l-1}\in {\left\{0,1\right\}}^{N\times H_l{W}_l}$ is the binarized output (with threshold at $0.5$) of the resized mask prediction of the previous $(l-1)$-th Transformer decoder layer

High-resolution features

Instead of always using the high-resolution feature map, utilize a feature pyramid which consists of both low- and high-resolution features and feed one resolution of the multi-scale feature to one Transformer decoder layer at a time.

Specifically, use the feature map produced by the pixel decoder with resolution $\frac{1}{32}$, $\frac{1}{16}$ and $\frac{1}{8}$ of the original image.