Segmentation

Mask2Former

The Masked-attention Mask Transformer (Cheng et al., CVPR 2022), a mask classification model whose Transformer decoder restricts cross-attention to each query's predicted mask, so that one architecture reaches strong results on panoptic, instance and semantic segmentation.

advanced

Mask2Former, the Masked-attention Mask Transformer, is an image segmentation model by Bowen Cheng, Ishan Misra, Alexander G. Schwing, Alexander Kirillov and Rohit Girdhar of Facebook AI Research and the University of Illinois at Urbana-Champaign, posted to arXiv in December 2021 (arXiv 2112.01527) and published at CVPR 2022. It is a mask classification model: instead of labelling every pixel, it predicts a fixed-size set of binary masks, each with one class label, from learned queries processed by a transformer decoder, following DETR and its direct predecessor MaskFormer. Its main change is masked attention, a form of cross-attention in which each query attends only to the pixels inside the mask it predicted at the previous decoder layer. With a multi-scale decoder and several training changes, the same architecture, loss and training procedure serve panoptic, instance and semantic segmentation, and the authors reported results on all three that matched or exceeded the specialized models of the time. The authors call this universal image segmentation, though a separate model is still trained for each task and dataset.

Problem Addressed

Per-pixel and mask classification

Semantic segmentation had been framed as per-pixel classification since the fully convolutional network: for an H×WH \times W image the model outputs a distribution pip_i over KK classes at every pixel ii and is trained with a per-pixel cross-entropy loss,

Lpixel-cls(y,ygt)=∑i=1H⋅W−log⁡pi(yigt).\mathcal{L}_{\text{pixel-cls}}(y, y^{gt}) = \sum_{i=1}^{H \cdot W} -\log p_i\big(y^{gt}_i\big).

Instance segmentation was dominated by mask classification, in which a model predicts a set of binary masks with one class each; Mask R-CNN does this from detected boxes. MaskFormer’s authors defined the general form (Cheng, Schwing & Kirillov, 2021, §3.2). The output is a set of NN probability–mask pairs z={(pi,mi)}i=1Nz = \{(p_i, m_i)\}_{i=1}^{N}, where mi∈[0,1]H×Wm_i \in [0, 1]^{H \times W} and pip_i is a distribution over the KK classes plus a “no object” label ∅\varnothing. NN need not equal KK, and several masks may carry the same class, so the same output can describe stuff regions, separate instances, or both. Per-pixel classification, by contrast, returns a fixed number of outputs and cannot return a variable number of segments, which instance-level tasks require. The authors also noted that mask classification is older than FCNs for semantic segmentation: the leading pre-FCN methods, such as O2P and SDS, classified mask proposals.

Training a mask classifier needs an assignment σ\sigma between predictions and the NgtN^{gt} ground-truth segments, padded with ∅\varnothing to size NN. MaskFormer computes it as DETR does, by bipartite matching with the Hungarian algorithm, but with a cost built from classes and masks rather than boxes, −pi(cjgt)+Lmask(mi,mjgt)-p_i(c^{gt}_j) + \mathcal{L}_{\text{mask}}(m_i, m^{gt}_j). Given the matching, the loss is

Lmask-cls(z,zgt)=∑j=1N[−log⁡pσ(j)(cjgt)+1cjgt≠∅ Lmask(mσ(j),mjgt)].\mathcal{L}_{\text{mask-cls}}(z, z^{gt}) = \sum_{j=1}^{N} \Big[ -\log p_{\sigma(j)}\big(c^{gt}_j\big) + \mathbb{1}_{c^{gt}_j \neq \varnothing}\, \mathcal{L}_{\text{mask}}\big(m_{\sigma(j)}, m^{gt}_j\big) \Big].

The DETR article covers set prediction and the matching itself.

MaskFormer

MaskFormer (NeurIPS 2021) instantiated this as three modules. A backbone produces a feature map at stride 32, and a pixel decoder, a light feature pyramid network, upsamples it to per-pixel embeddings Epixel∈RCE×H×W\mathcal{E}_{\text{pixel}} \in \mathbb{R}^{C_\mathcal{E} \times H \times W} at stride 4. A standard six-layer Transformer decoder, the same design as DETR’s, turns N=100N = 100 zero-initialized queries with learnable positional embeddings into per-segment embeddings. A linear classifier gives each pip_i, and a two-hidden-layer MLP gives a mask embedding Emask[:,i]\mathcal{E}_{\text{mask}}[:, i], so that

mi[h,w]=sigmoid⁡(Emask[:,i]⊤Epixel[:,h,w]).m_i[h, w] = \operatorname{sigmoid}\big(\mathcal{E}_{\text{mask}}[:, i]^{\top} \mathcal{E}_{\text{pixel}}[:, h, w]\big).

The masks are not forced to be mutually exclusive; the authors found a sigmoid better than a softmax across queries. The mask loss was DETR’s, a focal loss plus a dice loss.

The paper’s controlled comparison isolates the formulation. Its PerPixelBaseline+ has the same backbone, pixel decoder and Transformer decoder, but trains with per-pixel cross-entropy. On ADE20K val with a ResNet-50, it reached 41.9 mIoU and 28.3 PQSt^{\text{St}}, a region-level metric that treats every class as stuff. MaskFormer with a fixed matching of the 150 queries to the 150 classes reached 43.7 and 30.3, and with bipartite matching 44.2 and 33.4 (Table 4). The mIoU gain grew with the number of classes. On Cityscapes (19 classes, with a ResNet-101) the two tied at 78.5 mIoU, though MaskFormer still led by 2.9 PQSt^{\text{St}}, while on ADE20K-Full (847 classes) MaskFormer led by 3.5 mIoU (Table 2). The authors hypothesized that one class prediction per mask models fine-grained recognition better than per-pixel predictions. (The paper’s dataset section says 874 ADE20K-Full classes were kept, but the introduction, Table 2 and the appendix give 847.) Matching by masks rather than boxes mattered too. A MaskFormer with an added box head and box-based matching reached 43.7 PQ on COCO panoptic val, close to DETR’s 43.4, against 46.5 with mask-based matching (Table 5).

MaskFormer set results on ADE20K (55.6 mIoU, Swin-L, multi-scale) and COCO panoptic (52.7 PQ). Mask2Former’s authors identified two problems. First, universal models lagged specialized ones on instance segmentation by more than 9 AP: MaskFormer with Swin-L reached 40.1 AP (measured on the panoptic model’s thing classes) against 49.5 for the specialized Swin-HTC++. Second, they were costly to train: MaskFormer took 300 epochs and fitted only one image on a 32 GB GPU, while Swin-HTC++ trained in 72 epochs.

Architecture

Mask2Former keeps MaskFormer’s meta architecture of backbone, pixel decoder and Transformer decoder, and replaces the Transformer decoder. Each query still produces a class distribution and a mask embedding, and masks are dot products with the per-pixel embeddings at stride 4. The changes are in how the queries read the image.

Backbone and pixel decoder

The paper uses ResNet-50 and -101 and Swin-T, -S, -B and -L backbones, the larger Swin models pretrained on ImageNet-22K. Mask2Former works with any pixel decoder. Its default is the multi-scale deformable attention Transformer (MSDeformAttn) from Deformable DETR (Zhu et al., 2021): six MSDeformAttn layers on the 1/8, 1/16 and 1/32 feature maps, then an upsampling layer with a lateral connection that produces the 1/4 per-pixel embedding. In the ablation it was the only pixel decoder best on all three tasks. BiFPN came close on instance segmentation (43.5 against 43.7 AP) but trailed on semantic segmentation (45.6 against 47.2 mIoU), and FaPN did well on semantic segmentation (Table 4e). The official config sets this decoder as MSDeformAttnPixelDecoder with six encoder layers on res3–res5.

Masked attention

The authors’ starting point was the observation, from studies of DETR’s slow convergence, that cross-attention over the whole image takes many epochs to learn to attend to localized object regions. They hypothesized that local features suffice to update a query, and that context can come from self-attention among the queries. Standard cross-attention with a residual path updates the NN query features Xl∈RN×CX_l \in \mathbb{R}^{N \times C} at layer ll as

Xl=softmax⁡(QlKl⊤)Vl+Xl−1,X_l = \operatorname{softmax}\big(Q_l K_l^{\top}\big) V_l + X_{l-1},

with Ql=fQ(Xl−1)Q_l = f_Q(X_{l-1}) and Kl,VlK_l, V_l linear maps of the image features at resolution Hl×WlH_l \times W_l. Masked attention adds a mask to the logits:

Xl=softmax⁡(Ml−1+QlKl⊤)Vl+Xl−1,Ml−1(x,y)={0if Ml−1(x,y)=1,−∞otherwise.X_l = \operatorname{softmax}\big(\mathcal{M}_{l-1} + Q_l K_l^{\top}\big) V_l + X_{l-1}, \qquad \mathcal{M}_{l-1}(x, y) = \begin{cases} 0 & \text{if } \mathbf{M}_{l-1}(x, y) = 1, \\ -\infty & \text{otherwise}. \end{cases}

Ml−1∈{0,1}N×HlWl\mathbf{M}_{l-1} \in \{0, 1\}^{N \times H_l W_l} is the previous layer’s mask prediction, resized to the resolution of KlK_l and binarized at 0.5. M0\mathbf{M}_0 comes from the learnable query features before the first layer (Eqs. 1–3). Each query therefore reads only the pixels it currently claims. In the official code the resize is bilinear, the mask is detached from the gradient, and a query whose mask is empty at that resolution is allowed to attend everywhere, which avoids a softmax over no entries.

High-resolution features

High-resolution features help small objects but make attention expensive. Mask2Former feeds one scale per decoder layer: the pixel decoder’s 1/32, 1/16 and 1/8 maps go to three successive layers, lowest resolution first, and the 3-layer pattern repeats LL times in a round robin. Each scale gets a sinusoidal positional embedding and a learnable scale-level embedding. The default L=3L = 3 gives 9 layers.

Decoder layer order and queries

Three further changes, which the paper groups as optimization improvements, add no computation:

  • Masked attention first. Each layer runs masked attention, then self-attention, then the feed-forward network, reversing the standard order. The query features entering the first self-attention layer carry no image signal, so self-attention there is unlikely to add information.
  • Learnable query features. DETR and MaskFormer zero-initialize the query features and learn only their positional embeddings. Mask2Former learns both, and supervises the query features directly, by predicting masks from them before the first layer. The authors found these queries act like a region proposal network: their masks already reach 50.3 class-agnostic AR@100 on COCO val2017, rising to 57.7 after layer 9 (Figure 3).
  • No dropout. The authors found dropout unnecessary and usually harmful, and removed it from the decoder.

The supplementary material sets the two decoders side by side (Table XIb). MaskFormer has 6 layers ordered self-attention, cross-attention, FFN, with dropout 0.1, features at 1/32 only, and zero-initialized query features. Mask2Former has 9 layers ordered masked attention, self-attention, FFN, with no dropout, features at {1/32, 1/16, 1/8} three times, and learnable query features. The default uses 100 queries; the Swin-L panoptic and instance models use 200.

Training

Losses

The class loss is cross-entropy over the K+1K + 1 labels with weight λcls=2.0\lambda_{\text{cls}} = 2.0 for matched predictions and 0.1 for “no object”. (In the official code, 0.1 is a class weight inside the cross-entropy, which is then multiplied by 2.0.) The mask loss replaces MaskFormer’s focal loss with binary cross-entropy and keeps the dice loss:

L=λceLce+λdiceLdice+λclsLcls,λce=λdice=5.0.\mathcal{L} = \lambda_{\text{ce}} \mathcal{L}_{\text{ce}} + \lambda_{\text{dice}} \mathcal{L}_{\text{dice}} + \lambda_{\text{cls}} \mathcal{L}_{\text{cls}}, \qquad \lambda_{\text{ce}} = \lambda_{\text{dice}} = 5.0.

In the official code the dice term for a predicted probability map m^\hat m and a target mm is a smoothed form of one minus the dice coefficient, 1−(2∑km^kmk+1)/(∑km^k+∑kmk+1)1 - (2 \sum_k \hat m_k m_k + 1) / (\sum_k \hat m_k + \sum_k m_k + 1). The matcher uses the same three terms with the same weights, with −pi(cjgt)-p_i(c^{gt}_j) in place of the cross-entropy, and solves the assignment with SciPy’s linear_sum_assignment. The losses are applied to the prediction from the learnable queries and after every decoder layer. The official COCO config therefore sets DEC_LAYERS: 10, nine decoder layers plus one for the query loss.

Point-sampled mask loss

Computing mask losses on full-resolution masks for every query dominated memory: MaskFormer fitted one image in 32 GB. PointRend (Kirillov et al., 2020) had trained its point head on NN sampled points rather than on a regular grid. It over-generates kNkN uniform candidates, keeps the βN\beta N with the most uncertain coarse predictions, and fills the remaining (1−β)N(1 - \beta) N uniformly. The Mask2Former paper cites PointRend and Implicit PointRend (Cheng, Parkhi and Kirillov’s pointly-supervised instance segmentation) as showing that a segmentation model can be trained with its mask loss on randomly sampled points. Mask2Former computes its mask losses on K=12544K = 12544 points (112×112112 \times 112), in two ways:

  • In the matching cost, the same KK uniformly random points serve all predictions and ground-truth masks, so the cost matrix is cheap.
  • In the final loss, each matched pair gets its own KK points by PointRend’s importance sampling. The official config sets OVERSAMPLE_RATIO: 3.0 and IMPORTANCE_SAMPLE_RATIO: 0.75, and the uncertainty of a point is minus the absolute value of its mask logit.

Training memory fell from 18 GB to 6 GB per image. In the ResNet-50 ablation, point sampling in the final loss alone kept accuracy (41.0 AP, 50.8 PQ, 45.9 mIoU, against 41.0, 50.3 and 45.9 with full masks). Sampling in the matching as well raised all three, to 43.7, 51.9 and 47.2 (Table 5). The paper reports the gain without explaining it.

Procedure

Panoptic and instance models are trained in Detectron2 following the updated Mask R-CNN baseline settings. The paper specifies AdamW with learning rate 10−410^{-4}, weight decay 0.05 and a 0.1 multiplier on the backbone. The learning rate drops tenfold at 0.9 and 0.95 of training. Training runs 50 epochs at batch size 16, with large-scale jittering (augmentation by random scaling from 0.1 to 2.0, then a 1024×10241024 \times 1024 crop). The Swin-L models train for 100 epochs. Semantic models follow MaskFormer’s recipe, for example 160k iterations at 512×512512 \times 512 on ADE20K, except that the 0.1 multiplier applies to Transformer backbones too and all backbones use the same learning rate and weight decay. The model zoo notes 8 V100 GPUs for most models and 16 for Swin-L.

The supplement separates the recipe from the architecture (Table XI). Retraining MaskFormer with Mask2Former’s recipe raised it from 34.0 to 37.8 AP, 46.5 to 48.2 PQ and 44.5 to 45.3 mIoU. Then swapping only the decoder, with the same ResNet-50, FPN pixel decoder and recipe, gave 41.5 AP, 50.7 PQ and 45.6 mIoU. Mask2Former converged in about 25 epochs with standard scale augmentation and almost converged in 50 with large-scale jittering, which gave higher AP when training ran longer than 25 epochs (Figure IV), against 500 epochs for DETR and 300 for MaskFormer.

Example: mask classification and masked attention on synthetic scenes

The script below trains three small models on synthetic 32×3232 \times 32 scenes. Each scene has a background (a stuff class) and two to four objects of three thing classes: disk, square and triangle. Colours are random, so the class depends only on shape, and objects may overlap and touch. All three models share a backbone, an FPN-like pixel decoder and a three-layer decoder that runs cross-attention, then self-attention, then an FFN, over 8×88 \times 8 features. They differ only in what they predict and how they attend:

  • perpixel follows MaskFormer’s PerPixelBaseline+. It has K=4K = 4 class queries whose mask embeddings, dotted with the per-pixel embeddings, act as per-pixel class logits, trained with per-pixel cross-entropy.
  • full is mask classification with 10 queries, bipartite matching by linear_sum_assignment on the cost 2(−p)+5 BCE+5 dice2(-p) + 5\,\text{BCE} + 5\,\text{dice}, the Mask2Former loss weights, and standard cross-attention.
  • masked is the same with masked attention, including the official code’s fallback for empty masks.

Losses are computed on full masks, since the images are tiny, and applied to every layer’s prediction. Semantic maps come from MaskFormer’s semantic inference. Panoptic segments come from Mask2Former’s panoptic inference with the official 0.8 thresholds; for perpixel they are the connected components of each thing class. Scoring uses mIoU and a plain PQ (IoU above 0.5, no void handling). The script also reports the share of last-layer cross-attention that each matched thing query puts on its own ground-truth object, the quantity of Mask2Former’s Figure I(b). Before training it runs checks that need no model: an all-background baseline, the PQ ceiling of a perfect per-pixel map split into connected components, and the mean object colour per class.

The script was planned with the hypotheses that mask classification would have clearly higher PQ than per-pixel classification, and that masked attention would beat full attention at the same budget. It trains 9 models (three modes, seeds 0, 1 and 2) for 10 epochs on 1,000 scenes and took about 8 minutes on a 4-thread CPU with PyTorch 2.14.1, NumPy 2.5.4 and SciPy 1.18.1. The --time flag times one epoch per model. That timing, run on 1,600 training scenes, set the size of the run (1,000 scenes and 10 epochs instead of 1,600 and 12) before it was first run.

import sys, time
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
from scipy import ndimage
from scipy.optimize import linear_sum_assignment

torch.set_num_threads(4)
S, K, NQ, D, H = 32, 4, 10, 64, 4  # image size; classes 0 bg, 1 disk, 2 square, 3 triangle; queries; width; heads
EPOCHS, SEEDS = 10, (0, 1, 2)
yy, xx = np.mgrid[:S, :S]

def shape(c, r, cy, cx):
    if c == 1: return (yy - cy) ** 2 + (xx - cx) ** 2 <= r * r
    if c == 2: return (abs(yy - cy) <= 0.85 * r) & (abs(xx - cx) <= 0.85 * r)
    return (yy >= cy - r) & (yy <= cy + 0.7 * r) & (abs(xx - cx) <= 0.6 * (yy - cy + r))

def scene(rng):
    bg = rng.uniform(0, 1, 3); img = np.tile(bg, (S, S, 1)); inst = np.zeros((S, S), int); cls = [0]
    for j in range(1, rng.integers(2, 5) + 1):
        for _ in range(100):
            c, r, (cy, cx), col = rng.integers(1, 4), rng.uniform(4, 8), rng.uniform(4, S - 4, 2), rng.uniform(0, 1, 3)
            m = shape(c, r, cy, cx); new = np.where(m, j, inst)
            if np.linalg.norm(col - bg) > 0.4 and all((new == i).sum() >= 15 for i in range(1, j + 1)): break
        inst, img[m] = new, col; cls.append(c)
    img = img + rng.normal(0, 0.05, img.shape)
    segs = [(cls[i], inst == i) for i in range(len(cls)) if (inst == i).any()]
    return img.transpose(2, 0, 1).astype(np.float32), inst, np.array(cls)[inst], segs

def dataset(n, seed):
    rng = np.random.default_rng(seed); data = [scene(rng) for _ in range(n)]
    X = torch.tensor(np.stack([d[0] for d in data]))
    T = [{"c": torch.tensor([c for c, _ in d[3]]), "m": torch.tensor(np.stack([m for _, m in d[3]])).float()} for d in data]
    return X, torch.tensor(np.stack([d[2] for d in data])), T, data

def block(i, o, s=1): return nn.Sequential(nn.Conv2d(i, o, 3, s, 1), nn.GroupNorm(8, o), nn.ReLU())

class Net(nn.Module):
    def __init__(s, mode):  # mode: "perpixel", "full" or "masked"
        super().__init__(); s.mode = mode; nq = K if mode == "perpixel" else NQ
        s.c1, s.c2 = block(3, 32), block(32, 64, 2)
        s.c3 = nn.Sequential(block(64, 64, 2), block(64, 64), block(64, 64))
        s.p3, s.l2, s.l1 = nn.Conv2d(64, D, 1), nn.Conv2d(64, D, 1), nn.Conv2d(32, D, 1)
        s.f2, s.f1, s.pix = block(D, D), block(D, D), nn.Conv2d(D, D, 1)
        s.kpos = nn.Parameter(0.02 * torch.randn(1, 64, D))
        s.qfeat, s.qpos = nn.Embedding(nq, D), nn.Embedding(nq, D)
        L = lambda f: nn.ModuleList(f() for _ in range(3))
        s.ca, s.sa = L(lambda: nn.MultiheadAttention(D, H, batch_first=True)), L(lambda: nn.MultiheadAttention(D, H, batch_first=True))
        s.ff = L(lambda: nn.Sequential(nn.Linear(D, 4 * D), nn.ReLU(), nn.Linear(4 * D, D)))
        s.n1, s.n2, s.n3 = L(lambda: nn.LayerNorm(D)), L(lambda: nn.LayerNorm(D)), L(lambda: nn.LayerNorm(D))
        s.norm, s.cls = nn.LayerNorm(D), nn.Linear(D, K + 1)
        s.memb = nn.Sequential(nn.Linear(D, D), nn.ReLU(), nn.Linear(D, D), nn.ReLU(), nn.Linear(D, D))

    def heads(s, q, E):
        q = s.norm(q)
        return s.cls(q), torch.einsum("bqd,bdhw->bqhw", s.memb(q), E)

    def forward(s, x):
        f1 = s.c1(x); f2 = s.c2(f1); f3 = s.p3(s.c3(f2))                      # 32x32, 16x16, 8x8
        p2 = s.f2(F.interpolate(f3, scale_factor=2) + s.l2(f2))
        E = s.pix(s.f1(F.interpolate(p2, scale_factor=2) + s.l1(f1)))       # per-pixel embeddings
        mem = f3.flatten(2).transpose(1, 2); B = x.shape[0]
        q, qp = s.qfeat.weight.expand(B, -1, -1), s.qpos.weight.expand(B, -1, -1)
        outs = [s.heads(q, E)]                                               # prediction from learnable queries
        for l in range(3):
            am = None
            if s.mode == "masked":                                           # Eq. (2)-(3): attend inside previous mask
                m = F.interpolate(outs[-1][1].detach(), size=(8, 8), mode="bilinear", align_corners=False)
                am = m.sigmoid().flatten(2) < 0.5
                am[am.all(-1)] = False                                       # empty mask: attend everywhere
                am = am.repeat_interleave(H, dim=0)
            a, w = s.ca[l](q + qp, mem + s.kpos, mem, attn_mask=am)          # masked / full cross-attention first
            q = s.n1[l](q + a)
            q = s.n2[l](q + s.sa[l](q + qp, q + qp, q, need_weights=False)[0])
            q = s.n3[l](q + s.ff[l](q))
            outs.append(s.heads(q, E))
        return outs, w

def costs(logit, mask, t):
    p, pm, tm = logit.softmax(-1)[:, t["c"]], mask.flatten(1), t["m"].flatten(1)
    bce = (F.softplus(pm).sum(1)[:, None] - pm @ tm.T) / pm.shape[1]       # mean BCE for every pair
    ps = pm.sigmoid(); dice = 1 - (2 * ps @ tm.T + 1) / (ps.sum(1)[:, None] + tm.sum(1)[None] + 1)
    return 2 * -p + 5 * bce + 5 * dice

def match(logit, mask, t):
    with torch.no_grad(): return linear_sum_assignment(costs(logit, mask, t).numpy())

def loss_fn(model, outs, Y, T):
    total, w = 0.0, torch.ones(K + 1); w[K] = 0.1
    for logit, mask in outs:                                                 # every layer: auxiliary loss
        if model.mode == "perpixel": total = total + F.cross_entropy(mask, Y); continue
        for b, t in enumerate(T):
            r, c = match(logit[b], mask[b], t)
            tc = torch.full((NQ,), K); tc[r] = t["c"][c]
            pm, tm = mask[b][r].flatten(1), t["m"][c].flatten(1)
            bce = F.binary_cross_entropy_with_logits(pm, tm, reduction="none").mean(1).mean()
            ps = pm.sigmoid(); dice = (1 - (2 * (ps * tm).sum(1) + 1) / (ps.sum(1) + tm.sum(1) + 1)).mean()
            total = total + (2 * F.cross_entropy(logit[b], tc, weight=w) + 5 * bce + 5 * dice) / len(T)
    return total

def components(sem):                                                         # per-pixel map -> segments
    segs = [(0, sem == 0)] if (sem == 0).any() else []
    for c in range(1, K):
        lab, n = ndimage.label(sem == c); segs += [(c, lab == i) for i in range(1, n + 1)]
    return segs

def touching(inst, sem):                                                     # object touches a same-class object?
    out = []
    for i in np.unique(inst)[1:]:
        ring = ndimage.binary_dilation(inst == i) & (inst != i) & (inst > 0)
        out.append(bool((sem[ring] == sem[inst == i][0]).any()))
    return out

def panoptic(model, logit, mask, sem):
    if model.mode == "perpixel": return components(sem)
    score, lab = logit.softmax(-1).max(-1); m = mask.sigmoid().numpy()      # Mask2Former panoptic inference
    keep = np.where((lab != K).numpy() & (score > 0.8).numpy())[0]
    if len(keep) == 0: return []
    ids = (score.numpy()[keep, None, None] * m[keep]).argmax(0); segs, stuff = [], None
    for k, q in enumerate(keep):
        area, orig = (ids == k).sum(), (m[q] >= 0.5).sum(); seg = (ids == k) & (m[q] >= 0.5)
        if area == 0 or orig == 0 or seg.sum() == 0 or area / orig < 0.8: continue
        if lab[q] == 0:                                                      # merge stuff segments of one class
            if stuff is not None: segs[stuff] = (0, segs[stuff][1] | seg); continue
            stuff = len(segs)
        segs.append((int(lab[q]), seg))
    return segs

def pq(pred, gt):
    iou_sum, tp, fp, fn = np.zeros(K), np.zeros(K), np.zeros(K), np.zeros(K)
    for P, G in zip(pred, gt):
        used = set()
        for c, g in G:
            best = 0.0
            for i, (pc, p) in enumerate(P):
                if pc == c and i not in used:
                    iou = (p & g).sum() / (p | g).sum()
                    if iou > 0.5: best, bi = iou, i
            if best > 0.5: used.add(bi); iou_sum[c] += best; tp[c] += 1
            else: fn[c] += 1
        for i, (pc, _) in enumerate(P): fp[pc] += i not in used
    ok = tp + fp + fn > 0
    return (iou_sum[ok] / (tp + 0.5 * fp + 0.5 * fn)[ok]).mean()

def miou(pred, Y):
    return np.mean([((pred == c) & (Y == c)).sum() / ((pred == c) | (Y == c)).sum() for c in range(K)])

@torch.no_grad()
def evaluate(model, X, Y, T):
    model.eval(); (outs, w) = model(X); logit, mask = outs[-1]
    if model.mode == "perpixel": sem = mask.argmax(1)
    else: sem = torch.einsum("bqc,bqhw->bchw", logit.softmax(-1)[..., :K], mask.sigmoid()).argmax(1)
    sem = sem.numpy()
    P = [panoptic(model, logit[b], mask[b], sem[b]) for b in range(len(X))]
    G = [[(int(c), m.numpy() > 0) for c, m in zip(t["c"], t["m"])] for t in T]
    fg = np.nan
    if model.mode != "perpixel":                                             # attention weight on own object
        fr = []
        for b, t in enumerate(T):
            r, c = match(logit[b], mask[b], t)
            for q, g in zip(r, c):
                if t["c"][g] > 0:
                    cell = F.avg_pool2d(t["m"][g][None, None], 4).flatten() >= 0.5
                    fr.append(w[b, q][cell].sum().item())
        fg = np.mean(fr)
    return miou(sem, Y.numpy()), pq(P, G), fg

def train(mode, seed, X, Y, T, epochs):
    torch.manual_seed(seed); model = Net(mode); opt = torch.optim.AdamW(model.parameters(), 1e-3, weight_decay=1e-4)
    g = torch.Generator().manual_seed(seed)
    for ep in range(epochs):
        model.train()
        for idx in torch.randperm(len(X), generator=g).split(32):
            outs, _ = model(X[idx]); loss = loss_fn(model, outs, Y[idx], [T[i] for i in idx])
            opt.zero_grad(); loss.backward(); opt.step()
    return model

if __name__ == "__main__":
    Xtr, Ytr, Ttr, _ = dataset(1000, 100); Xte, Yte, Tte, raw = dataset(400, 200)
    if "--time" in sys.argv:                                                 # diagnostic: one epoch per model
        for mode in ("perpixel", "full", "masked"):
            t0 = time.time(); train(mode, 0, Xtr, Ytr, Ttr, 1); print(mode, f"{time.time() - t0:.1f} s/epoch")
        sys.exit()
    # label-free difficulty checks (no training)
    G = [d[3] for d in raw]
    cc = [components(d[2]) for d in raw]
    touch = [touching(d[1], d[2]) for d in raw]
    print(f"objects per image {np.mean([len(g) - 1 for g in G]):.2f}; all-background mIoU {miou(np.zeros_like(Yte.numpy()), Yte.numpy()):.3f}")
    print(f"oracle per-pixel map + connected components: PQ {pq(cc, G):.3f}; "
          f"objects touching a same-class object {np.mean(sum(touch, [])):.3f}")
    print("mean object colour per class:", [np.round(Xte.permute(0, 2, 3, 1).numpy()[Yte.numpy() == c].mean(0), 2).tolist() for c in range(1, K)])
    res = {}
    for mode in ("perpixel", "full", "masked"):
        for seed in SEEDS:
            res[mode, seed] = evaluate(train(mode, seed, Xtr, Ytr, Ttr, EPOCHS), Xte, Yte, Tte)
            print(f"{mode:9s} seed {seed}: mIoU {res[mode, seed][0]:.3f}  PQ {res[mode, seed][1]:.3f}  attention on own object {res[mode, seed][2]:.2f}")
    for mode in ("perpixel", "full", "masked"):
        v = np.array([res[mode, s] for s in SEEDS])
        print(f"{mode:9s} range over seeds: mIoU {v[:, 0].min():.3f}-{v[:, 0].max():.3f}  PQ {v[:, 1].min():.3f}-{v[:, 1].max():.3f}  attention {v[:, 2].min():.2f}-{v[:, 2].max():.2f}")
objects per image 2.99; all-background mIoU 0.189
oracle per-pixel map + connected components: PQ 0.866; objects touching a same-class object 0.309
mean object colour per class: [[0.5, 0.5099999904632568, 0.47999998927116394], [0.5099999904632568, 0.5199999809265137, 0.46000000834465027], [0.49000000953674316, 0.47999998927116394, 0.5199999809265137]]
perpixel  seed 0: mIoU 0.572  PQ 0.406  attention on own object nan
perpixel  seed 1: mIoU 0.556  PQ 0.357  attention on own object nan
perpixel  seed 2: mIoU 0.565  PQ 0.415  attention on own object nan
full      seed 0: mIoU 0.381  PQ 0.247  attention on own object 0.11
full      seed 1: mIoU 0.393  PQ 0.301  attention on own object 0.10
full      seed 2: mIoU 0.366  PQ 0.317  attention on own object 0.16
masked    seed 0: mIoU 0.419  PQ 0.295  attention on own object 0.81
masked    seed 1: mIoU 0.428  PQ 0.295  attention on own object 0.82
masked    seed 2: mIoU 0.408  PQ 0.380  attention on own object 0.76
perpixel  range over seeds: mIoU 0.556-0.572  PQ 0.357-0.415  attention nan-nan
full      range over seeds: mIoU 0.366-0.393  PQ 0.247-0.317  attention 0.10-0.16
masked    range over seeds: mIoU 0.408-0.428  PQ 0.295-0.380  attention 0.76-0.82

The checks show a task that is neither trivial nor shortcut by colour. Predicting background everywhere gives 0.189 mIoU, and the class-wise mean colours are all close to 0.5. About 31% of objects touch another object of the same class. Even a perfect per-pixel map therefore scores only 0.866 PQ once split into connected components, because touching same-class objects merge; mask classification has no such ceiling. (nan marks the attention measure for perpixel, which has no object queries.)

The first hypothesis failed at this budget. The per-pixel model beat both mask-classification models in every seed, on mIoU and on PQ, despite its connected-components ceiling. This is the opposite of MaskFormer’s ADE20K ablation. Two things the papers report may account for it; this is reasoning, not a tested explanation. First, MaskFormer’s mIoU advantage shrank with fewer classes, down to a tie on Cityscapes’ 19 (where it still led by 2.9 PQSt^{\text{St}}), and this task has four. Second, mask transformers converge slowly: matching, per-query classification and mask prediction must all be learned before PQ rises, and 10 epochs on 1,000 images is a short schedule.

The second comparison shows the mechanism clearly. Masked attention put 0.76–0.82 of each query’s attention on its own object, against 0.10–0.16 for full cross-attention. That is the same direction as Mask2Former’s 0.59 against 0.20 on COCO; the gap is larger here, plausibly because the objects are compact and the feature map is small, though that was not tested. Masked attention also gave higher mIoU in all three seeds. Its PQ was higher in two of the three seeds and lower in seed 1, and the seed ranges overlap (0.295–0.380 against 0.247–0.317), so three seeds do not settle the PQ difference.

Inference

The outputs are the NN class distributions and masks from the last decoder layer, upsampled to the image. Post-processing follows MaskFormer and differs by task:

  • Semantic segmentation marginalizes over queries, arg⁡max⁡c∈{1,…,K}∑i=1Npi(c) mi[h,w]\arg\max_{c \in \{1, \dots, K\}} \sum_{i=1}^{N} p_i(c)\, m_i[h, w]. MaskFormer’s authors found this better than assigning each pixel to a single query.
  • Panoptic segmentation drops “no object” queries and those below a confidence threshold, then assigns each pixel to arg⁡max⁡ipi(ci) mi[h,w]\arg\max_{i} p_i(c_i)\, m_i[h, w], where cic_i is query ii’s most likely class. A segment is discarded if less than a set fraction of its own mask (mi≥0.5m_i \ge 0.5) survives the assignment, and stuff segments of one class are merged. The official COCO config sets both thresholds to 0.8.
  • Instance segmentation needs a score per mask. The class probability is multiplied by a mask confidence, the mean foreground probability over the pixels where the mask is positive. The official code ranks all query–class pairs, so one query can yield masks of several classes, and keeps the top ones; no non-maximum suppression is applied.

Because each query yields a whole segment, the model’s panoptic output can be scored for all three tasks. The paper reports APpanTh^{\text{Th}}_{\text{pan}}, the instance AP on thing classes, and mIoUpan_{\text{pan}}, the semantic mIoU after merging same-class masks, for the panoptic model.

Speed on a V100 at batch size 1, including post-processing, was 8.6 frames per second for the ResNet-50 COCO panoptic model, against 17.6 for MaskFormer, and 4.0 for Swin-L. Average FLOPs were 226G and 868G (Table 1). The authors did not use multi-scale inference for instance-level tasks, because merging instance predictions across scales would need post-processing such as NMS (supplement, §A.2).

Performance

All results below are as reported by the authors; backbones marked ImageNet-22K were pretrained on it.

COCO panoptic val2017 (133 classes, Table 1, single-scale):

Model Backbone Queries Epochs PQ PQTh^{\text{Th}} PQSt^{\text{St}} APpanTh^{\text{Th}}_{\text{pan}} mIoUpan_{\text{pan}}
DETR ResNet-50 100 500 + 25 43.4 48.2 36.3 31.1 –
MaskFormer ResNet-50 100 300 46.5 51.0 39.8 33.0 57.8
Mask2Former ResNet-50 100 50 51.9 57.7 43.0 41.7 61.7
MaskFormer Swin-L (22K) 100 300 52.7 58.5 44.0 40.1 64.8
K-Net Swin-L (22K) 100 36 54.6 60.2 46.0 – –
Mask2Former Swin-L (22K) 200 100 57.8 64.2 48.1 48.6 67.4

On COCO test-dev the Swin-L model reached 58.3 PQ, against 54.7 for the COCO challenge winner (Megvii), which used extra data and ensembles (supplement, Table II).

COCO instance segmentation val2017 (80 classes, Table 2, single-scale, train2017 only):

Model Backbone Epochs AP APS_S APM_M APL_L Boundary AP
Mask R-CNN ResNet-50 36 37.2 18.6 39.5 53.3 23.1
Mask R-CNN (large-scale jittering) ResNet-50 400 42.5 23.8 45.0 60.0 28.0
Mask2Former ResNet-50 50 43.7 23.4 47.2 64.8 30.6
Swin-HTC++ Swin-L (22K) 72 49.5 31.0 52.4 67.2 34.1
Mask2Former Swin-L (22K) 100 50.1 29.9 53.9 72.1 36.2

The authors read the 0.6 AP gain and 2.1 Boundary AP gain over Swin-HTC++ as a sign of better boundaries from high-resolution masks. The table also shows the weakness: Mask2Former trails on small objects and leads by 4.9 points on large ones. On test-dev it reached 50.5 AP with 29.1 APS_S and 71.2 APL_L, against 53.1, 36.6 and 67.7 for the challenge winner (supplement, Table IV).

ADE20K val (150 classes, Table 3): with a ResNet-50 at 512×512512 \times 512, Mask2Former reached 47.2 mIoU single-scale and 49.2 multi-scale, against 44.5 and 46.7 for MaskFormer. With Swin-L at 640×640640 \times 640 it reached 56.1 and 57.3, and with a FaPN pixel decoder 56.4 and 57.7. The best specialized model in the table, BEiT-UperNet, had 57.0 multi-scale. On the ADE20K test set, trained on train plus val with multi-scale inference, the FaPN model scored 49.72 mIoU and a challenge score of 64.76 (supplement, Table VI). ADE20K panoptic and instance results with Swin-L were 48.1 PQ and 34.9 AP (supplement, Table VIII).

Cityscapes val (19 classes, Table 6): with Swin-L, 66.6 PQ (single-scale), 43.6 APpanTh^{\text{Th}}_{\text{pan}} and 82.9 mIoUpan_{\text{pan}} from the panoptic model, and 83.3 mIoU single-scale and 84.3 multi-scale from the semantic model; the Swin-B semantic model reached 84.5 multi-scale. The comparison rows were Panoptic-DeepLab with SWideRNet at 66.4 PQ single-scale and SegFormer at 84.0 mIoU multi-scale. On Mapillary Vistas val, the Swin-L models reached 45.5 PQ (single-scale) and 64.7 mIoU (multi-scale), against 45.7 PQ for Panoptic FCN with Swin-L and 61.1 mIoU multi-scale for HMSANet (supplement, Table IX).

Ablations

The ResNet-50 ablations score COCO instance AP, COCO panoptic PQ and ADE20K mIoU together (Table 4):

  • Masked attention gave the largest gain. Replacing it with standard cross-attention dropped the model from 43.7 AP, 51.9 PQ and 47.2 mIoU to 37.8, 47.1 and 45.5 at 213G instead of 226G FLOPs. The drop was larger for the instance-level tasks than for semantic segmentation. SMCA reached 37.9, 47.2 and 46.6, and K-Net’s mask pooling, which averages features inside the mask, 43.1, 51.5 and 46.0 (Table 4c).
  • High-resolution features were next. Using only the 1/32 map cost 2.2 AP, 1.7 PQ and 1.1 mIoU. A single 1/8 map gave 44.0, 51.8 and 47.4 at 239G. Feeding all three scales to every layer gave no further gain at 247G, while the round robin nearly matched the 1/8 map (43.7, 51.9 and 47.2) at 226G (Table 4d).
  • The optimization changes together were worth 1.4 AP, 1.1 PQ and 0.9 mIoU at no cost (Table 4b). Learnable but unsupervised query features did no better than zero-initialized ones on COCO (supplement, Table Xb).
  • Number of queries. 100 queries were best for instance and semantic segmentation, and 200 slightly better for panoptic (52.2 PQ). With 1,000 queries all three tasks degraded, for example to 40.3 AP (supplement, Table Xa).

The supplement measures what masked attention changes. On COCO val, about 20% of standard cross-attention weight fell on the query’s ground-truth foreground, against about 59% for masked attention, averaged over the three scales (Figure I(b)). The authors attribute this to the softmax: its weights never reach zero, so many small weights on a large background add up. In the per-layer panoptic curves, one masked-attention layer already outperformed nine layers of standard cross-attention (Figure II).

Strengths

  • One architecture for three tasks. The same model, loss and training procedure gave state-of-the-art results, as reported in 2022, on COCO panoptic and instance segmentation and ADE20K semantic and panoptic segmentation, and results the authors call competitive on Cityscapes and Mapillary Vistas. A model trained for panoptic segmentation also yields usable instance and semantic outputs.
  • No boxes, anchors or NMS. Segments come directly from queries, with stuff and things handled alike.
  • Faster training than earlier mask transformers. 50 epochs against MaskFormer’s 300 at better accuracy, with three times less training memory from point sampling.
  • Better boundaries. The Boundary AP gain over Swin-HTC++ is larger than the mask AP gain.
  • Modular. The ablation swaps five pixel decoders and several backbones without other changes.

Limitations

  • Still trained per task. The authors’ own limitation table shows that a model trained on panoptic annotations scores slightly below the same architecture trained for instance or semantic segmentation. On COCO this was 41.7 against 43.7 AP and 61.7 against 61.5 mIoU (the semantic model is lower on this pair). On ADE20K it was 26.5 against 26.4 AP and 46.1 against 47.2 mIoU, and on Cityscapes 37.3 against 37.4 AP and 77.5 against 79.4 mIoU (Table 7). The paper’s text reads all of these as slightly worse, though the table shows the panoptic model ahead in two of the six comparisons, by 0.2 mIoU on COCO and 0.1 AP on ADE20K. Training once for all tasks was left to future work, which OneFormer took up.
  • Small objects. APS_S trails Swin-HTC++ and, with ResNet-50, the 400-epoch Mask R-CNN. The authors also note that the model does not fully exploit multi-scale features.
  • Cost. At equal backbone, the decoder and MSDeformAttn pixel decoder cost more than MaskFormer’s (226G against 181G FLOPs with ResNet-50) and halve the frame rate. The MSDeformAttn pixel decoder relies on a custom CUDA operator in the official code.
  • Masked attention trusts the previous mask. A query reads only inside the region it predicted one layer earlier. The official code’s fallback for empty masks shows that this region can vanish. From the design, it follows that a mask which misses part of the object can only be corrected through self-attention, the FFN and the next layer’s prediction; the paper does not analyse such failures.
  • Query count is a dataset-dependent hyperparameter. The best number differed between tasks, and too many queries hurt.
  • Licence. As of November 2026 the official code is under the MIT licence, with portions under other licences (Deformable DETR’s Apache-2.0), but the released model weights are under CC BY-NC 4.0, which excludes commercial use.

Historical Importance

Before 2021, semantic segmentation was per-pixel classification with context modules (DeepLab, PSPNet), and instance segmentation was box-based mask classification (Mask R-CNN). Panoptic segmentation either combined the two or, as in DETR and MaX-DeepLab, treated stuff and things alike. MaskFormer showed in 2021 that DETR-style mask classification could lead semantic segmentation too, and K-Net extended set prediction to instance segmentation. Mask2Former, published at CVPR 2022, was in its authors’ words the first architecture to outperform the best specialized models on all considered tasks and datasets. Its contribution was to make the mask-transformer recipe trainable at the cost of a specialized detector: masked attention, the multi-scale round robin and point-sampled losses.

The model became a common segmentation head on later backbones. The DINOv2 and DINOv3 papers, for example, report ADE20K results with frozen backbones inside ViT-Adapter and Mask2Former pipelines (see DINO). Its decoder was also the base of OneFormer.

Variants

  • MaskFormer (Cheng, Schwing & Kirillov, NeurIPS 2021), the predecessor described above. Mask2Former’s code is largely based on it.
  • Video Mask2Former. The official README announces support for video instance segmentation, described in a separate technical report by the same group (arXiv 2112.10764). It was not read for this article.
  • OneFormer (Jain et al., CVPR 2023) trains one model once for all three tasks. Ground truth for each task is derived from panoptic annotations, and a task is sampled uniformly for each training image. A text input, “the task is {task}”, is tokenized into a task token. The object queries are initialized from repetitions of that token, refined by a 2-layer transformer over the 1/4-scale features, and concatenated with the token itself. A query–text contrastive loss aligns the object queries with text queries built from the ground truth, such as “a photo with a car”, padded with entries that stand for “no object”. The text branch is dropped at inference. The pixel decoder (MSDeformAttn) and the masked-attention decoder follow Mask2Former, with 250 queries on ADE20K and Cityscapes and 150 on COCO. With Swin-L the single model reported 49.8 PQ, 35.9 AP and 57.0 mIoU (single-scale) on ADE20K val. The individually trained Mask2Former models in the same table had 48.7, 34.9 and 56.1; OneFormer’s table lists the Mask2Former PQ with a 0.5 confidence threshold, where Mask2Former’s own supplement gives 48.1. On Cityscapes val OneFormer reached 67.2 PQ against 66.6. On COCO val2017 it reached 57.9 PQ, on par with Mask2Former’s 57.8. A Mask2Former trained jointly without task conditioning reached 48.7 PQ, 33.7 AP and 56.2 mIoU on ADE20K, below OneFormer (OneFormer, Table 7).

Implementations

The official repository, facebookresearch/Mask2Former, is built on Detectron2 and provides configs and trained models for COCO, ADE20K, Cityscapes and Mapillary Vistas, plus video instance segmentation. The last commit is from May 2022, as of November 2026. The masked attention, matcher, point-sampled criterion and inference routines discussed above are in mask2former_transformer_decoder.py, matcher.py, criterion.py and maskformer_model.py. Hugging Face Transformers (v5.22.1) includes Mask2FormerForUniversalSegmentation with the same point-sampling parameters (train_num_points, oversample_ratio, importance_sample_ratio), and a OneFormer implementation. OneFormer’s authors released their code at github.com/SHI-Labs/OneFormer.

Related

  • Instance Segmentation

    Detecting every object of a fixed set of "thing" classes and giving each one its own pixel mask, class label and confidence score, so that two touching cars come out as two cars rather than one car-colored region.

  • Semantic Segmentation

    Labeling every pixel of an image with one of a fixed set of classes, such as road, sky, car or person, without separating individual objects; the dense-prediction counterpart of image classification, also called scene parsing.

  • Cross-Attention

    Attention in which one set of tokens asks the questions and another set supplies the answers, so that image features can read a text prompt, learned object queries can read an image, or a small set of latents can summarize a large input; the output has one row per query and the cost grows with the product of the two set sizes.

  • Mask R-CNN

    An instance segmentation model that adds a small fully convolutional mask branch and RoIAlign to the Faster R-CNN detector, predicting a box, a class and a binary mask for every detected object.

  • Panoptic Quality

    The standard score for panoptic segmentation, which matches predicted and ground-truth segments at IoU above 0.5 and combines the matched IoUs with the unmatched segments into an IoU-weighted F1, per class and averaged over classes.

  • ADE20K

    A scene-centric dataset from MIT CSAIL, first released in 2016, in which a single expert annotator labeled every object and stuff region, and many object parts, in more than 20,000 images with an open vocabulary of thousands of classes; its SceneParse150 benchmark, scored by pixel accuracy and mean IoU over 150 classes, became a standard test of semantic segmentation.

References

  1. Cheng, B., Misra, I., Schwing, A. G., Kirillov, A. & Girdhar, R. (2022). Masked-attention Mask Transformer for Universal Image Segmentation. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 1280–1289.
  2. Cheng, B., Misra, I., Schwing, A. G., Kirillov, A. & Girdhar, R. (2022). Supplementary Materials: Masked-attention Mask Transformer for Universal Image Segmentation. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), CVF open access.
  3. Cheng, B., Schwing, A. G. & Kirillov, A. (2021). Per-Pixel Classification is Not All You Need for Semantic Segmentation. Advances in Neural Information Processing Systems 34 (NeurIPS), 17864–17875.
  4. Cheng, B., Misra, I., Schwing, A. G., Kirillov, A. & Girdhar, R. Mask2Former: official implementation, README, model zoo, configs and modeling code (GitHub repository facebookresearch/Mask2Former, commit 9b0651c).
  5. Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A. & Zagoruyko, S. (2020). End-to-End Object Detection with Transformers. European Conference on Computer Vision (ECCV), Lecture Notes in Computer Science, 213–229.
  6. Kirillov, A., Wu, Y., He, K. & Girshick, R. (2020). PointRend: Image Segmentation As Rendering. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 9796–9805.
  7. Jain, J., Li, J., Chiu, M., Hassani, A., Orlov, N. & Shi, H. (2023). OneFormer: One Transformer to Rule Universal Image Segmentation. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2989–2998.
  8. Hugging Face. Transformers v5.22.1: Mask2Former and OneFormer model code (src/transformers/models/mask2former and oneformer).