PEM

PEM (Prototype-based Efficient MaskFormer) is an efficient multi-task image segmentation architecture.

• Paper: Niccolò Cavagnero, Gabriele Rosi et al. (Politecnico di Torino / Focoos AI), PEM: Prototype-based Efficient MaskFormer for Image Segmentation, CVPR 2024, arXiv:2402.19422.
• Architecture: keeps the Mask2Former meta-architecture (backbone → pixel decoder → query-based Transformer decoder, mask classification) but replaces its two expensive parts.
• Prototype-based cross-attention: the masked cross-attention becomes PEM-CA, which attends to one selected pixel per query instead of all $H_i{W}_i$ pixels.
• Pixel decoder: the deformable-attention pixel decoder becomes a fully convolutional FPN with context-based self-modulation and deformable convolutions.
• Headline result: with a ResNet50, 61.1 PQ on Cityscapes panoptic at 13.8 FPS (Mask2Former: 62.1 PQ at 4.1 FPS) and 45.5 mIoU on ADE20K semantic at 35.7 FPS.
• Standing: the fastest universal (semantic + panoptic) architecture in its comparison.

Motivation

MaskFormer-family models (Mask2Former, kMaX-DeepLab) unify semantic and panoptic segmentation under one DETR-style set-prediction objective (DETR), but their cost is dominated by (i) cross-attention between $N$ queries and high-resolution pixel features, which scales with $H_i{W}_i$, and (ii) a pixel decoder built from multi-scale deformable attention. Efficient real-time segmenters exist, but each targets a single task (BiSeNet, STDC, DDRNet, PIDNet for semantic; UPSNet, Panoptic-DeepLab, RealTimePan for panoptic). PEM aims to be both universal and fast enough for edge deployment. The closest universal-and-fast competitor is YOSO.

Architecture

Fig. 2: backbone features $F_1..F_4$ go through CSM modules and are fused top-down with deformable convolutions; ${\hat{F}}_2..{\hat{F}}_4$ feed the PEM-CA decoder blocks, ${\hat{F}}_1$ is multiplied with the refined queries to produce masks.

• Backbone — features $F_i\in {\mathbb{R}}^{H_i\times W_i\times B_i}$ at strides 4, 8, 16, 32 ($i=\mathrm{1..4}$). ResNet50 by default; STDC1/STDC2 for the faster variants.
• Pixel decoder — produces ${\hat{F}}_i\in {\mathbb{R}}^{H_i\times W_i\times C}$ ($C=128$).
• Transformer decoder — takes ${\hat{F}}_2,{\hat{F}}_3,{\hat{F}}_4$ and $N$ learnable queries $Q\in {\mathbb{R}}^{N\times C}$, returns refined queries $\hat{Q}$. Masks are $\hat{Q}\cdot {\hat{F}}_1$; class probabilities come from a linear classifier on $\hat{Q}$ (with a “no object” class), as in Mask2Former.

Prototype-based masked cross-attention (PEM-CA)

Fig. 3: the token dimension shrinks from $HW$ to $N$ after prototype selection; everything downstream is element-wise.

The premise: as training progresses, pixels belonging to one object become close in feature space, so using all of them to refine that object’s query is redundant. One representative pixel per query — its prototype — suffices.

Features and queries are projected into a shared $D$-dimensional space, giving $K\in {\mathbb{R}}^{H_i{W}_i\times D}$ and $Q\in {\mathbb{R}}^{N\times D}$. The similarity matrix and selection are

$S=K{Q}^{\mathsf{T}}\in {\mathbb{R}}^{H_i{W}_i\times N},G=\arg {\max }_{H_i{W}_i}(S+\mathcal{M}),K_p=K[G]\in {\mathbb{R}}^{N\times D}$

where the mask $\mathcal{M}$ is the Mask2Former masked-attention mask, built from the binarized mask prediction $M_{l-1}$ of the previous decoder layer resized to the feature resolution:

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

so a prototype can only be picked among current foreground pixels. Selection is multi-head: features and queries are split along channels, so each channel group of a query can pick a different pixel.

Because each query now has exactly one matched token, the query–prototype interaction needs no pairwise attention (in the spirit of SwiftFormer / MobileViT V2 additive attention):

$A=(Q\odot K_p)W_A,B=\mathbf{\alpha }\odot \frac{A}{\Vert A{\Vert }_2}+K_p,Q_{out}=B{W}_{out}+Q$

with $W_A\in {\mathbb{R}}^{D\times D}$, $W_{out}\in {\mathbb{R}}^{D\times C}$, $A$ normalized across channels and scaled by a learnable $\mathbf{\alpha }\in {\mathbb{R}}^D$; the scaled interaction reweights the prototypes additively. Cross-attention cost no longer grows with the number of pixels — at Cityscapes resolution, PEM-CA is about $2\times$ faster than masked cross-attention, with the gap widening as resolution grows (Fig. 4 of the paper).

Efficient multi-scale pixel decoder

Mask2Former’s pixel decoder is an FPN plus deformable attention, which gives (i) global context, (ii) input-dependent weights and (iii) a deformable receptive field. PEM keeps all three properties with cheaper ops.

Context-based self-modulation (CSM) — squeeze-and-excitation style. Backbone features $F_i$ are projected to $F_i^{\prime }$ (channels $C$), and

${\Omega }_i=\mathrm{MLP}(\mathrm{GAP}(F_i^{\prime })),F_i^c=F_i^{\prime }\odot \sigma ({\Omega }_i)+F_i^{\prime }$

where the MLP is two $1\times 1$ convolutions. This injects a global scene representation and dynamic channel weights (properties i and ii) using only $1\times 1$ convs, normalizations and element-wise products.

Feature aggregation — top-down FPN with deformable convolutions (property iii):

${\hat{F}}_i=\{\begin{array}{ll}\mathrm{DefConv}(F_4^c+\mathrm{Proj}(\mathrm{GAP}(F_4)))& i=4\\ \mathrm{DefConv}(F_i^c+\mathrm{Up}({\hat{F}}_{i+1}))& i=2,3\\ F_1^c+\mathrm{Up}({\hat{F}}_2)& i=1\end{array}$

The lowest-resolution level additionally receives a pooled scene vector, and the highest-resolution level ${\hat{F}}_1$ has no deformable conv (it is only used for the final mask product).

Training and losses

Same loss and recipe as Mask2Former: classification with binary cross-entropy, masks with BCE + Dice, Hungarian matching, deep supervision at every decoder block.

Hyperparameters

Item	Value
Optimizer	AdamW, cosine LR schedule
Learning rate	0.0007 (Cityscapes), 0.0004 (ADE20K); 0.1× multiplier on the backbone
Weight decay	0.05
Batch size	32
Iterations	90k (Cityscapes), 160k (ADE20K)
Loss weights	2.0 classification (BCE), 5.0 mask BCE, 5.0 Dice
Backbone	ResNet50, ImageNet-1k pretrained (also STDC1/STDC2)
Decoder	2 transformer decoder stages (3 PEM-CA blocks each), hidden dim 256, 8 heads, FFN expansion 8
Queries	100
Pixel decoder hidden dim	128
Train crop	Cityscapes 512×1024 (full 1024×2048 at inference); ADE20K 512×512 semantic, 640×640 panoptic (short side resized to crop at inference)

Latency/FPS: PyTorch 1.12, FP32, V100, batch size 1. Task-specific baselines on ADE20K were retrained (160k iters, AdamW, LR 0.0004, wd 0.05, crop 512×512) since they were not published there.

Results

Panoptic segmentation, ResNet50 (FPS on V100):

Method	Cityscapes PQ	FPS	ADE20K PQ	FPS
Mask2Former	62.1	4.1	39.7	19.5
YOSO	59.7	11.1	38.0	35.4
PEM	61.1	13.8	38.5	35.7

On Cityscapes PEM has 237 GFLOPs / 35.6M params vs. 519G / 44.0M for Mask2Former. It is $+1.4$ PQ over YOSO, mostly from things ($+3.3$ PQ${}^{th}$; stuff is equal at 66.1). kMaX-DeepLab (42.3 PQ on ADE20K) is more accurate but slower.

Semantic segmentation:

Method	Backbone	Cityscapes mIoU	FPS	ADE20K mIoU	FPS
Mask2Former	R50	79.4	6.6	47.2	21.5
YOSO	R50	79.4	11.1	44.7	35.3
PIDNet-L	–	80.9	31.1	40.5	65.4
PEM	STDC1	78.3	24.3	39.6	43.6
PEM	STDC2	79.0	22.0	45.0	36.3
PEM	R50	79.9	13.8	45.5	35.7

On Cityscapes only the hand-engineered task-specific PIDNet beats PEM-R50 in mIoU; on ADE20K PEM beats every task-specific model by a wide margin (+5 mIoU over PIDNet-L) — with the caveat that those baselines were retrained by the authors.

Ablations (Cityscapes panoptic)

Cross-attention variant (cumulative latency of the cross-attention modules):

Variant	PQ	PQ${}^{th}$	Latency
Cross-attention (MaskFormer)	58.4	47.4	11.4 ms
Masked CA (Mask2Former)	60.4	52.0	17.4 ms
PEM-CA, no prototypes (sum-aggregate tokens)	48.7	24.7	5.6 ms
PEM-CA, no masking	57.8	47.0	6.8 ms
PEM-CA	61.1	54.3	9.8 ms

• Tokens cannot be naively collapsed: replacing prototype selection by summation costs 13.4 PQ (almost 20 PQ${}^{th}$). Masking the selection to foreground is worth 3.3 PQ. The losses concentrate on things, i.e. prototypes matter for separating instances.
• PEM-CA beats even full masked cross-attention in PQ while being faster.

Pixel decoder (total model latency): removing CSM costs >1 PQ and ~3 PQ${}^{th}$; replacing deformable with standard convs costs 4 PQ and ~7 PQ${}^{th}$. PEM with the MaskFormer pixel decoder gets 57.8 PQ (84.7 ms), with Mask2Former’s decoder 61.4 PQ (131.6 ms) vs. PEM’s 61.1 PQ (72.4 ms) — equal accuracy at roughly twice the FPS.

Other: queries $N=50/100/200$ give 58.7/61.1/61.0 PQ (225G/237G/275G FLOPs), so 100 is the sweet spot. Performance stays robust with a single decoder stage (3 layers), giving a speed/accuracy knob; without any cross-attention layers, ADE20K PQ collapses to nearly zero, so query refinement is essential.

Limitations

• Only semantic and panoptic segmentation are evaluated — no instance segmentation and no COCO.
• Only Cityscapes and ADE20K, ResNet50/STDC backbones; no results with large or transformer backbones.
• Still slower than the best hand-engineered single-task real-time models (e.g. PIDNet) on Cityscapes, and 1 PQ below Mask2Former on Cityscapes panoptic.
• A hard argmax selects a single pixel per query per head; the ablation shows the mechanism is sensitive to the mask, so quality relies on the previous layer’s mask predictions.

See also

• Mask2Former — the architecture PEM slims down; masked attention is the mechanism PEM-CA builds on.
• Segmentation Head Comparison — where Mask2Former’s mask-classification head sits among other mask heads.
• DETR — the set-prediction / learnable-query formulation underlying all MaskFormer-family models.
• YOLACT — another prototype-based segmentation design, though its prototypes are mask bases rather than selected feature tokens.
• PANet — the class-prototype idea in few-shot segmentation, a different use of the word “prototype”.
• Attention Is All You Need — the attention operation PEM-CA replaces with element-wise interactions.