Source-linked AI summary
Query2Label: A Simple Transformer Way to Multi-Label Classification
Shilong Liu, Lei Zhang, Xiao Yang, Hang Su, Jun Zhu
TL;DR
Multi-label classification requires managing label imbalance and extracting local features for multiple objects distributed across an image. Query2Label uses label embeddings as queries in Transformer decoders, whose cross-attention pools class-related features from a vision backbone. The framework consistently outperforms prior works across several benchmarks and establishes new state-of-the-art results.
Problem
Multi-label classification must address imbalanced one-vs-all samples and avoid losing localized object information through global feature pooling.
Method
Query2Label is a two-stage framework using a vision backbone followed by Transformer decoders that apply multi-head cross-attention with learnable label embeddings as queries.
Results
The framework consistently outperforms prior works on MS-COCO, PASCAL VOC, NUS-WIDE, and Visual Genome, establishing new state-of-the-art results.
Takeaways & Limitations
Its simple architecture and implementation provide a strong baseline for future multi-label image-classification research.
Abstract
from arXiv · showhide
This paper presents a simple and effective approach to solving the multi-label classification problem. The proposed approach leverages Transformer decoders to query the existence of a class label. The use of Transformer is rooted in the need of extracting local discriminative features adaptively for different labels, which is a strongly desired property due to the existence of multiple objects in one image. The built-in cross-attention module in the Transformer decoder offers an effective way to use label embeddings as queries to probe and pool class-related features from a feature map computed by a vision backbone for subsequent binary classifications. Compared with prior works, the new framework is simple, using standard Transformers and vision backbones, and effective, consistently outperforming all previous works on five multi-label classification data sets, including MS-COCO, PASCAL VOC, NUS-WIDE, and Visual Genome. Particularly, we establish $91.3\%$ mAP on MS-COCO. We hope its compact structure, simple implementation, and superior performance serve as a strong baseline for multi-label classification tasks and future studies. The code will be available soon at https://github.com/SlongLiu/query2labels.
1. Introduction
Query2Label addresses label imbalance and region-feature extraction in multi-label classification with a simple Transformer-decoder framework. Label queries use cross-attention to pool class-related features adaptively, while multi-head attention captures different object parts or views.
- Motivation: Multi-label classification must handle imbalanced positive and negative samples and extract discriminative features from distributed object regions.Global pooling can dilute features and make small objects difficult to identify.
- Approach: Learnable label embeddings act as queries for cross-attention, which probes and pools class-related features for subsequent binary classification.This produces adaptive, more discriminative features than globally average-pooled convolutional features.
- Approach: Multi-head attention extracts features from different object parts or views, supporting recognition under occlusions and viewpoint changes.The authors contrast this with prior single-head attention that cannot extract features by parts or views.
- Approach: Query2Label uses a two-stage framework that combines an image classification backbone with Transformer decoder layers for label-specific classification.The backbone may be a CNN or Vision Transformer, while label embeddings query object features in the decoder.
- Results: The framework uses a simplified asymmetric loss with different γ values for positive and negative samples and reports new state-of-the-art results across four named benchmarks.The benchmarks named here are MS-COCO, PASCAL VOC, NUS-WIDE, and Visual Genome.
2. Related Work
Related work addresses multi-label classification through improved losses, label-correlation modeling, region localization, and Transformer-based vision methods. Query2Label differs by using decoder cross-attention with label-specific queries to locate object features.
- Improving loss functions: Loss-based methods address sample imbalance from one-vs-rest binary classifiers using distribution-balanced or asymmetric weighting strategies.Query2Label adopts a simplified asymmetric loss without prediction-probability shifting.
- Modeling label correlations: Label-correlation methods mine concept co-occurrence as prior knowledge, often using graphs to model interactions among category-correlated features.
- Locating regions of interest: Region-localization methods seek object areas through proposals, spatial transformers, global-to-local discovery, or label-embedding attention maps.Earlier bounding boxes could be coarse and include background, while cosine-similarity attention may produce overly smooth maps.
- Transformer-based methods: Most Transformer vision methods use encoder self-attention for feature representation, whereas Query2Label uses decoder cross-attention to locate label-specific object features.DETR uses class-agnostic queries, while Query2Label assigns each query to one label class.
3. Method
Query2Label is a two-stage, backbone-agnostic framework that uses Transformer decoders to query label-specific features from spatial image representations and classify each label independently.
- Adaptive feature extraction: Cross-attention compares label queries with spatial features to produce attention maps that adaptively pool category-related features.The pooled feature is formed by linearly combining spatial features according to the attention maps before label prediction.
- Framework: Query2Label first extracts spatial image features with a vision backbone, then processes them through Transformer decoders and a linear prediction layer.The backbone may be a conventional CNN or Vision Transformer, while the second stage performs query updating, adaptive pooling, and category prediction.
- Framework: Each learnable label embedding serves as a class-specific query for multi-head cross-attention over spatial features.The queries are learned end to end and have clear semantic meanings corresponding to categories.
- Query updating: Decoder layers update queries using standard self-attention, cross-attention, and position-wise feed-forward components, with categories decoded in parallel.Attention masks are unnecessary because the model does not perform autoregressive prediction.
- Prediction: The final feature for each category is projected to a logit and passed through a sigmoid to obtain its predicted probability.Each label prediction is treated as a binary classification task with category-specific linear parameters.
- Loss function: The framework does not require a new loss function and uses a simplified asymmetric loss to address positive-negative sample imbalance.Binary cross entropy and focal loss also work, while the asymmetric loss performed best in the experiments; the default settings are γ+ = 0 and γ− = 1.
4. Experiments
The experiments evaluate Query2Label across four multi-label benchmarks using standard classification metrics and dataset-specific settings. The method achieves state-of-the-art results across MS-COCO, PASCAL VOC, NUS-WIDE, and VG500, with gains varying by resolution and dataset scale.
- Evaluation setup: Experiments report AP, mAP, OP, OR, OF1, CP, CR, and CF1 across MS-COCO, PASCAL VOC, NUS-WIDE, and Visual Genome.mAP is computed over categories; the other metrics provide overall and per-category comparisons.
- MS-COCO: Query2Label outperforms prior methods on MS-COCO medium resolution in mAP, OF1, and CF1.It exceeds ADD-GCN by 2.0%, SSGRL by 3.4%, and ASL by 0.6%.
- MS-COCO: At 640 × 640 on MS-COCO, Query2Label outperforms the best prior result and establishes a new state of the art.The high-resolution setting uses TResNetXL and removes decoder self-attention for improved training and inference efficiency.
- PASCAL VOC: On PASCAL VOC 2007, Query2Label achieves the best mAP, while the margin over ADD-GCN is small under differing resolutions and a potentially saturated metric.VOC 2012 provides a larger-margin comparison through evaluation-server results, where Query2Label leads all methods on every metric.
- NUS-WIDE: Query2Label reaches new state-of-the-art performance on NUS-WIDE, although its improvement is less significant than on MS-COCO and PASCAL VOC.The authors attribute the smaller improvement contextually to NUS-WIDE's limited image resolution.
- Visual Genome: On VG500, Query2Label outperforms all previous works despite using 512 × 512 inputs versus SSGRL's 576 × 576 inputs.The authors identify spatially adaptive feature attention as particularly important when the number of categories is large.
4.3. Ablation Study
Query2Label outperforms the baseline across small, medium, and large objects, with the largest improvement on medium objects. Small objects remain challenging despite substantial progress.
- Results on objects of different sizes: Query2Label outperforms the TResNetL baseline on small, medium, and large objects.The comparison uses MS-COCO subsets defined by object size.
- Results on objects of different sizes: The largest improvement occurs for medium objects, supporting spatially adaptive feature pooling over average pooling.Adaptive pooling helps collect information that average pooling may dilute.
- Results on objects of different sizes: Small objects remain challenging and require finer-grained image details, although Query2Label makes substantial progress.
4.4. Visualization of Attention Maps
Attention visualizations show that Query2Label approximately localizes label-specific objects, including small and medium objects. Its multi-head attention can focus on different object parts, while comparisons with TResNetL highlight gains on difficult small objects.
- Cross-attention maps: Query2Label’s cross-attention approximately locates specified objects, especially small and medium objects.The attention map is analogous to receptive-field size in the raw image.
- Multi-head attention: Different attention heads focus on different parts of a person, including the shoulder, neck, and head.Head-2 is less informative because it shows no clear focus.
- Baseline comparison: Q2L correctly classifies examples that TResNetL misclassifies when small object scale makes recognition difficult.The visualization compares Q2L attention maps with enlarged focused regions.
5. Conclusion
Query2Label is a simple multi-label classification framework built around Transformer decoders and an image backbone. Its cross-attention uses label embeddings to query class existence and pool class-related features, achieving consistent gains across established datasets.
- Conclusion: Query2Label combines Transformer decoders with an image classification backbone for multi-label classification.
- Conclusion: Transformer-decoder cross-attention uses label embeddings to query class existence and pool class-related features.
- Conclusion: The framework consistently outperforms prior works on MS-COCO, PASCAL VOC, NUS-WIDE, and Visual Genome.
- Conclusion: The authors propose Q2L as a strong baseline because of its simple architecture and outstanding performance.
A. More Implementation Details
The implementation uses standard PyTorch backbone and Transformer components with positional encodings, an 80-epoch Adam training schedule, and several regularization and augmentation choices.
- Training setup: The implementation uses official PyTorch implementations for the backbone and Transformer modules.
- Training setup: Models are trained for 80 epochs with Adam, a 1-cycle policy, and a maximal learning rate of 1e-4.
- Training setup: Regularization and augmentation include Cutout, True-weight-decay, and RandAugment.
- Model details: Following DETR, the models represent spatial positions with 2D sine and cosine encodings.
B. Metrics
The experiments report AP and mAP alongside overall and per-category precision, recall, and F1 metrics. Predictions are thresholded into positive or negative labels, making threshold choice relevant; OF1 and CF1 are primary metrics.
- The experiments report overall precision (OP), recall (OR), F1-measure (OF1), and per-category precision (CP), recall (CR), and F1-measure (CF1) beyond AP and mAP.
- For each category, the metrics use correctly predicted images, predicted images, and ground-truth images as their underlying counts.
- Predictions receive positive labels when their probabilities exceed a threshold and negative labels otherwise.
- OF1 and CF1 are the primary metrics because they jointly consider recall and precision.
C. Additional results on VOC07
The main text reports VOC07 results using ImageNet-1k pretrained backbones for fair comparison, while additional results appear in Table 8.
- The main text presents VOC07 results only with ImageNet-1k pretrained backbones.
- The ImageNet-1k backbone restriction is used to ensure a fair comparison on VOC07.
- Additional VOC07 results are listed in Table 8.
D. More Visualization Results
The paper visualizes cross-attention maps and multi-head attention maps to show how label queries focus on spatial image regions and different object parts. It also includes a VOC07 comparison table using ImageNet-22k-pretrained backbones.
- Cross-attention maps: Cross-attention visualizations compute attention between labels and image pixels, then resize the maps to raw-image dimensions for display.
- VOC07 comparison: Table 8 compares the method with prior state-of-the-art methods on VOC07 using ImageNet-22k-pretrained backbones, with all metrics reported in percent.
- Cross-attention maps: Mean-head maps in Figures 6 and 7 provide an intuitive explanation of the method’s spatial adaptive pooling.
- Multi-head attention maps: Different attention heads focus on different target parts in the multi-head visualizations for the person label.