Efficient Masked Attention Transformer for Few-Shot Classification and Segmentation

Few-shot classification and segmentation (FS-CS) focuses on jointly performing multi-label classification and multi-class segmentation using few annotated examples. Although the current state of the art (SOTA) achieves high accuracy in both tasks, it struggles with small objects. To overcome this, we propose the Efficient Masked Attention Transformer (EMAT), which improves classification and segmentation accuracy, especially for small objects. EMAT introduces three modifications: a novel memory-efficient masked attention mechanism, a learnable downscaling strategy, and parameter-efficiency enhancements. EMAT outperforms all FS-CS methods on the PASCAL-5$^i$ and COCO-20$^i$ datasets, using at least four times fewer trainable parameters. Moreover, as the current FS-CS evaluation setting discards available annotations, despite their costly collection, we introduce two novel evaluation settings that consider these annotations to better reflect practical scenarios.

💡 Research Summary

The paper tackles the challenging problem of few‑shot classification and segmentation (FS‑CS), where a model must simultaneously predict a multi‑label classification vector and a multi‑class segmentation mask for a query image given only a few annotated support examples. While the current state‑of‑the‑art method, the Classification‑Segmentation Transformer (CST), achieves high overall accuracy, it suffers from poor performance on small objects because it relies on a memory‑intensive masked‑attention mechanism that forces aggressive down‑sampling of the support feature maps. To address these limitations, the authors propose the Efficient Masked Attention Transformer (EMAT), which introduces three key innovations:

1. Memory‑Efficient Masked Attention – Instead of processing all entries of the correlation tensor, EMAT explicitly removes entries that are masked out by the support mask. A custom element‑wise masking operator (⊘) discards key and value vectors corresponding to zero‑mask positions before the softmax and weighted sum, dramatically reducing memory consumption. This enables the use of much higher‑resolution correlation tokens (support spatial dimensions increased from 12×12 to 20×20 in the first attention layer and from 3×3 to 10×10 in the second) without exceeding GPU memory limits.

2. Learnable Down‑Scaling – CST uses only average pooling to shrink the support token map, which limits flexibility. EMAT replaces this with a lightweight, learnable down‑scaling module that combines small 2‑D and 3‑D convolutions with pooling. The support query matrix Q is split into image tokens Q_i and a class token Q_c; Q_i undergoes a 3‑D convolution and reshaping, while Q_c passes through a 2‑D convolution. The resulting tensors are concatenated to form a down‑scaled query Q_d. In the second attention layer, Q_i is further collapsed to a single token via 3‑D average pooling before concatenation. This hybrid design removes the need for large pooling kernels while preserving spatial detail.

3. Parameter‑Efficiency Enhancements – Few‑shot learning is prone to over‑fitting when models contain many parameters. EMAT reduces channel widths throughout the two‑layer transformer (64→32 and 32→16 channels) and the task‑specific heads (128→32 and 64→16 in CST), cutting the total number of trainable parameters to roughly one‑quarter of CST’s count.

The overall pipeline mirrors CST: a frozen ViT‑S backbone pretrained with DINOv2 extracts patch tokens from support and query images, a class token is appended, and cosine similarity produces a correlation tensor C. The two‑layer transformer processes C with the new masked‑attention and down‑scaling, after which separate heads output the multi‑label classification scores and the multi‑class segmentation mask. Training is performed on the 1‑way‑1‑shot setting using a combined loss L = λ L_clf + L_seg. During inference on an N‑way‑K‑shot task, each class is treated as an independent 1‑way‑K‑shot problem; logits are averaged over the K support examples, thresholded at δ = 0.5, and masks are generated accordingly. Importantly, EMAT can emit empty masks when a class is absent from the query image, a capability lacking in many prior FS‑S methods.

Beyond architectural advances, the authors critique the prevailing FS‑CS evaluation protocol, which discards any additional annotations present in support images, thereby wasting valuable label information. They propose two new evaluation settings:

• Partially Augmented Setting – retains all annotations belonging to the selected support classes, even if a support image is used for a different class in the episode.
• Fully Augmented Setting – retains every annotation present in each support image, regardless of whether the corresponding class is part of the episode’s support set.

These settings better reflect real‑world scenarios where all collected annotations are available and align with the generalized few‑shot learning (GFSL) paradigm.

Experiments on the widely used PASCAL‑5i (20 classes) and COCO‑20i (80 classes) benchmarks demonstrate that EMAT outperforms all prior FS‑CS methods in both mean Intersection‑over‑Union (mIoU) and mean Average Precision (mAP). The gains are especially pronounced for small objects (occupying <15 % of the image area), confirming the effectiveness of the high‑resolution correlation processing. Despite using roughly four times fewer trainable parameters, EMAT’s inference time remains comparable to CST, thanks to the efficient masking strategy.

In summary, EMAT delivers a memory‑, compute‑, and parameter‑efficient transformer architecture that markedly improves small‑object performance in few‑shot classification and segmentation, while also introducing more realistic evaluation protocols that fully exploit available annotations. This makes EMAT a compelling, practical solution for real‑world applications where data is scarce and objects of interest may be tiny.