PaperlayerView exact revision on arXiv

Source-linked AI summary

ML-Decoder: Scalable and Versatile Classification Head

Tal Ridnik, Gilad Sharir, Avi Ben-Cohen, Emanuel Ben-Baruch, Asaf Noy

arXiv:2111.12933v2cs.CVcs.LG

TL;DR

Classification heads must preserve useful spatial information while remaining efficient for many classes and adaptable to zero-shot prediction. ML-Decoder redesigns attention-based decoding with reduced complexity, group decoding, and flexible query use. Across single-label, multi-label, and zero-shot tasks, it reports state-of-the-art results and a better speed-accuracy trade-off than larger backbones, while its demonstrated scope remains classification tasks.

  • Problem

    Existing GAP-based heads can provide sub-optimal results and lack a natural zero-shot extension, whereas attention-based heads are often costly and infeasible for extreme classification.

  • Method

    ML-Decoder removes redundant decoder self-attention, uses group decoding with fixed queries, and supports flexible word-based queries with query augmentations.

  • Results

    ML-Decoder outperforms GAP-based heads across multi-label, single-label, and zero-shot tasks, with a better speed-accuracy trade-off than larger backbones.

  • Takeaways & Limitations

    A single classification head can serve multiple classification settings while scaling to thousands of classes and generalizing to unseen classes with word queries.

  • Takeaways & Limitations

    The paper identifies extending ML-Decoder beyond classification, including detection, video recognition, segmentation, pose estimation, and NLP-related tasks, as future work.

Abstract

from arXiv · show

In this paper, we introduce ML-Decoder, a new attention-based classification head. ML-Decoder predicts the existence of class labels via queries, and enables better utilization of spatial data compared to global average pooling. By redesigning the decoder architecture, and using a novel group-decoding scheme, ML-Decoder is highly efficient, and can scale well to thousands of classes. Compared to using a larger backbone, ML-Decoder consistently provides a better speed-accuracy trade-off. ML-Decoder is also versatile - it can be used as a drop-in replacement for various classification heads, and generalize to unseen classes when operated with word queries. Novel query augmentations further improve its generalization ability. Using ML-Decoder, we achieve state-of-the-art results on several classification tasks: on MS-COCO multi-label, we reach 91.4% mAP; on NUS-WIDE zero-shot, we reach 31.1% ZSL mAP; and on ImageNet single-label, we reach with vanilla ResNet50 backbone a new top score of 80.7%, without extra data or distillation. Public code is available at: https://github.com/Alibaba-MIIL/ML_Decoder

1. Introduction

Image classification spans single-label, multi-label, extreme-classification, and zero-shot settings, but existing heads trade off spatial effectiveness, computational cost, scalability, and query flexibility. ML-Decoder addresses these limitations with an efficient, unified attention-based head built around reduced decoder complexity and group decoding.

  • Motivation: GAP-based heads are efficient and scalable but can underuse spatial information and lack a natural zero-shot extension, while attention-based heads are costly and difficult to scale.These trade-offs motivate a classification head that combines spatial modeling with efficiency and zero-shot compatibility.
  • ML-Decoder: ML-Decoder unifies single-label, multi-label, and zero-shot classification as a common classification-head design.It is presented as a drop-in replacement for common classification heads and is evaluated across several classification tasks.
  • ML-Decoder: Removing decoder self-attention reduces query dependence from quadratic to linear, improving scalability and efficiency.The redesign removes a redundant self-attention block from the transformer-decoder architecture.
  • ML-Decoder: Group decoding uses a fixed number of queries and interpolates them to the number of classes, enabling fixed spatial pooling cost for thousands of classes.The group fully-connected block maps grouped queries to the final class space.
  • Zero-shot classification: ML-Decoder supports different query types during training and inference, enabling word-query operation and query augmentations for unseen-class generalization.The paper describes random-query and query-noise augmentations as mechanisms for encouraging generalization to unseen class queries.
  • Experiments: Comprehensive experiments evaluate ML-Decoder on MS-COCO, Open Images, NUS-WIDE, PASCAL-VOC, and ImageNet.The evaluation spans multi-label, zero-shot, and single-label classification settings.

2. Method

ML-Decoder redesigns attention-based classification heads to preserve spatial information while reducing the computational dependence on the number of classes. Its self-attention removal, group-decoding, and query flexibility support scalable classification, zero-shot learning, and broader drop-in use.

  • Baseline Classification Heads: Unlike GAP-based heads, attention-based heads retain richer spatial information for images containing multiple objects with different locations and sizes.GAP heads are efficient and class-scalable but can provide sub-optimal multi-label results and lack a direct ZSL extension.
  • Motivation: ML-Decoder addresses the scalability limits of transformer-decoder heads, whose computational cost is quadratic in the number of classes.This limitation makes transformer-decoder classification practically infeasible for datasets such as Open Images with 9,600 classes.
  • ML-Decoder Design: ML-Decoder removes redundant self-attention while preserving the transformer-decoder’s expressivity, reducing query dependence from quadratic to linear.The projection before cross-attention can transform queries as needed, making self-attention unnecessary.
  • ML-Decoder Design: Group-decoding replaces one query per class with a fixed number K of group queries, making cross-attention and feed-forward computation independent of the number of classes.A group fully-connected layer later expands group-query outputs to N logits.
  • ML-Decoder Design: The group fully-connected layer simultaneously expands each group query to N/K outputs and pools the embedding dimension into class logits.Its operation assigns each output logit to a group query and within-group output index.
  • ML-Decoder for ZSL: ML-Decoder supports fixed or learnable queries, different queries at training and inference, and word queries for zero-shot recognition.The design also extends group-decoding to ZSL and uses random-query and query-noise augmentations to encourage generalization to unseen queries.
  • Applications: ML-Decoder can serve as a drop-in replacement for GAP-based heads in tasks including single-label classification.The paper presents it as a unified head for single-label, multi-label, and zero-shot classification.

3. Experimental Study

Experiments show that ML-Decoder preserves accuracy across query choices, reduces computational cost through query grouping, and improves the speed-accuracy trade-off across classification settings. It also generalizes to unseen classes with query augmentations and group decoding.

  • Query ablations: 88.1% mAP is achieved with learnable, fixed random, and fixed NLP-based word queries on MS-COCO.The authors use fixed queries thereafter to reduce learned parameters.
  • Classification-head comparisons: Removing self-attention preserves accuracy at the same 80-query setting while reducing computational cost.Transformer-decoder and ML-Decoder reach the same accuracy with 80 input queries.
  • Extreme classification: On Open Images, ML-Decoder group decoding raises FLOPs by only 10%–20% while significantly improving mAP over GAP.Transformer-decoder training is infeasible for the 9,600-class dataset because of out-of-memory computation.
  • Speed-accuracy comparison: ML-Decoder provides a better FLOPs-accuracy trade-off than GAP with a larger backbone on MS-COCO.The comparison covers TResNet-S, TResNet-M, and TResNet-L.
  • Zero-shot learning ablations: Both random-query and additive-noise augmentations improve unseen-class generalization, jointly increasing ZSL mAP by 1.2%.Query augmentation does not change seen-class mAP on MS-COCO or NUS-WIDE.
  • Zero-shot learning ablations: Group decoding also works for ZSL, with only a small mAP decrease compared with full decoding.The comparison is reported in the NUS-WIDE ZSL ablation.

4. Results

ML-Decoder achieves strong results across multi-label, zero-shot, and single-label benchmarks while remaining computationally practical for large class sets. The reported results include new state-of-the-art performance, although complete speed-accuracy comparisons with prior work are limited by missing computational-cost reports.

  • MS-COCO: 91.4% mAP is achieved on MS-COCO with TResNet-XL at input resolution 640.The authors report 91.1% with TResNet-L under the same resolution.
  • Evaluation limitations: A complete prior-work speed-accuracy comparison on MS-COCO is unavailable because earlier studies did not consistently report computational cost or release reproducible code.The paper provides its own FLOPs results at multiple input resolutions.
  • Additional multi-label datasets: ML-Decoder reaches 96.6% mAP on Pascal-VOC and 86.8% mAP on 9,600-class Open Images.Some attention-based methods are infeasible on Open Images because of its large class count.
  • Zero-shot classification: ML-Decoder improves the previous top NUS-WIDE ZSL result by 4.8% mAP.It also achieves top results on both ZSL and GZSL, rather than optimizing only unseen classes.
  • Single-label classification: Replacing GAP with ML-Decoder significantly improves ImageNet accuracy and gives a better speed-accuracy trade-off than using GAP with a larger backbone.The evaluation uses the A2 training configuration and compares ResNet architectures.
  • Single-label classification: ML-Decoder remains a drop-in replacement in the optimized ImageNet setting without changing or tuning training hyperparameters.The baseline vanilla ResNet50 achieves 79.7% in the cited configuration.
  • Single-label classification: 80.7% accuracy is reached with ML-Decoder and a vanilla ResNet50 after increasing training to 600 epochs.The paper reports this as the top ResNet50 result without extra data or distillation.

5. Conclusions and Future Work

The paper concludes that ML-Decoder is an efficient, versatile classification head that scales to thousands of classes and supports multi-label, single-label, and zero-shot classification. Future work will extend the approach to other vision and non-vision tasks involving spatial or structured prediction.

  • Conclusions: Removing self-attention and adding group decoding lets ML-Decoder scale to thousands of classes while improving the speed-accuracy trade-off over larger backbones.The conclusion also reports compatibility with fixed or random queries and different training and inference queries.
  • Conclusions: Word queries and query augmentations enable ML-Decoder to generalize to unseen classes.The conclusion presents this as part of its zero-shot capability.
  • Conclusions: Extensive experiments report new state-of-the-art results across multi-label, single-label, and zero-shot classification.The conclusion summarizes results across several classification tasks.
  • Future work: Future work will extend ML-Decoder to object detection, video recognition, segmentation, pose estimation, and NLP-related tasks.The authors also plan to study group decoding more generally for spatial embedding tensors.

Appendices

For ZSL group decoding, queries concatenate projections of the word embeddings assigned to each label group. The approach addresses group-decoding scalability while using a shared parameter component for labels.

  • ML-Decoder uses K = N queries for MS-COCO experiments with 80 classes.
  • ZSL group decoding: ZSL group queries concatenate linear projections of all word embeddings assigned to their group.This construction adapts group decoding, where each query represents multiple labels, to word-query-based ZSL.
  • ZSL group decoding: Each group query is defined from the word embeddings of the labels assigned to that group.
  • ZSL group decoding: The group fully-connected head does not generalize well to unseen classes when using group decoding.The method therefore decomposes the group parameter matrix into label-specific word embeddings and a learned matrix shared across labels.
  • ZSL group decoding: The resulting word-embedding construction produces output logits for the ZSL group decoder and is evaluated through ablations.
  • ZSL group decoding: Figure 7 illustrates the group decoding scheme for zero-shot learning.

C. MS-COCO Training Details

The MS-COCO experiments use a specified optimization and regularization setup, with throughput measured across classification heads. ML-Decoder has a modest, class-count-independent speed reduction, unlike the transformer decoder.

  • Training setup: MS-COCO models were trained for 40 epochs with Adam, a 1-cycle policy, and a maximal learning rate of 2e-4.Regularization used Cutout, true weight decay, and auto-augment; RGB channels were scaled to [0,1].
  • Speed-accuracy evaluation: Throughput measurements used TResNet-M at 224 resolution on an Nvidia V100 with mixed precision and 80% of maximal batch size.
  • Training setup: All results were averaged over three seeds, and the TResNet-L model was version V21.
  • Speed-accuracy evaluation: 15%: ML-Decoder reduces inference speed by 15% on TResNet-M, independently of the number of classes.
  • Speed-accuracy evaluation: Transformer-decoder classification heads reduce inference speed by orders of magnitude as the number of classes increases.
  • Results: Table 10 compares MS-COCO mAP scores across different input resolutions.

F. NUS-WIDE ZSL Dataset and Training Details

NUS-WIDE is a multi-label zero-shot dataset with 925 seen labels and 81 unseen categories. The experiments use TResNet-M, 224-resolution inputs, and a full-decoding ML-Decoder baseline.

  • Dataset: NUS-WIDE contains nearly 270K images, 81 human-annotated categories, and 925 labels from Flickr user tags.
  • Dataset: The 925 Flickr-derived labels are treated as seen classes, while the 81 human-annotated categories are unseen classes.
  • Training details: NUS-WIDE experiments use cross-entropy, a TResNet-M backbone, and 224-resolution inputs.
  • Training details: The baseline ZSL ML-Decoder uses full decoding with K = N and a shared projection matrix.

G. Pascal-VOC Training Details and Results

Pascal-VOC evaluates ML-Decoder on 20-category multi-label recognition using TResNet-L at 448-pixel input resolution. The baseline uses full decoding with one query per class.

  • Dataset and setup: Pascal-VOC contains 20 object categories and averages 2.5 categories per image.
  • Dataset and setup: The dataset is divided into 5,011 trainval images and 4,952 test images.
  • Dataset and setup: Experiments use a TResNet-L backbone with input resolution 448.
  • Dataset and setup: The ML-Decoder baseline uses full decoding with K = N = 20.
  • Results: Table 11 compares ML-Decoder with known state-of-the-art models on Pascal-VOC.

H. Open-Images Training Details and Results

The Open Images experiments use a large, partially annotated dataset and a TResNet-M backbone, with reduced weights for untagged labels. The supplied passages identify comparison tables for Open Images and ImageNet classification heads, but do not provide their numerical results.

  • Dataset and training: Open Images contains 9 million training images, 41,620 validation images, and 125,436 test images.The dataset is partially annotated with human and machine-generated labels.
  • Dataset and training: Untagged Open Images labels are treated as negative labels with reduced weights.
  • Dataset and training: The Open Images model is trained for 25 epochs at 224-pixel input resolution using TResNet-M as the backbone.
  • Results: Table 12 compares ML-Decoder with known state-of-the-art results on Open Images.
  • Results: Table 13 reports ImageNet scores for classification heads and logit activations using ResNet50 and 100-group ML-Decoder decoding.

J. Comparison of ML-Decoder to State-of-the-art Models on Single-label Transfer Learning Datasets

The single-label transfer-learning comparison evaluates ML-Decoder against state-of-the-art models on CIFAR-100 and Stanford-Cars. The supplied result states that ML-Decoder ranks first on Stanford-Cars and second on CIFAR-100, while group decoding assigns multiple classes to each query.

  • Benchmark comparison: The comparison covers the single-label CIFAR-100 and Stanford-Cars datasets against known state-of-the-art models.
  • Benchmark comparison: ML-Decoder achieves first place on Stanford-Cars and second place on CIFAR-100.
  • Query augmentations: Random-query augmentation adds random queries assigned the label “noise,” while additive-noise augmentation perturbs the input queries.
  • Implementation: The group fully-connected implementation uses a loop and is described as memory-efficient during training and suitable for compile-time acceleration.
Loading 2111.12933v2…