[Paper Review] Retrieval Augmented Classification for Long-Tail Visual Recognition

Alexander Long, Wei Yin+7 more|arXiv (Cornell University)|Feb 22, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes Retrieval Augmented Classification (RAC), a method that enhances standard image classification by integrating a retrieval module with an external memory of pre-encoded images and text. RAC achieves state-of-the-art performance on long-tail benchmarks—14.5% and 6.7% relative improvement on Places365-LT and iNaturalist-2018—by enabling the retrieval branch to implicitly learn tail-class representations, thereby freeing the base encoder to focus on head classes without fine-tuning.

ABSTRACT

We introduce Retrieval Augmented Classification (RAC), a generic approach to augmenting standard image classification pipelines with an explicit retrieval module. RAC consists of a standard base image encoder fused with a parallel retrieval branch that queries a non-parametric external memory of pre-encoded images and associated text snippets. We apply RAC to the problem of long-tail classification and demonstrate a significant improvement over previous state-of-the-art on Places365-LT and iNaturalist-2018 (14.5% and 6.7% respectively), despite using only the training datasets themselves as the external information source. We demonstrate that RAC's retrieval module, without prompting, learns a high level of accuracy on tail classes. This, in turn, frees the base encoder to focus on common classes, and improve its performance thereon. RAC represents an alternative approach to utilizing large, pretrained models without requiring fine-tuning, as well as a first step towards more effectively making use of external memory within common computer vision architectures.

Motivation & Objective

• To address the challenge of long-tail visual recognition, where rare classes suffer from insufficient training data and are overshadowed by frequent classes.
• To reduce reliance on parameter-heavy models by introducing explicit, non-parametric external memory for storing visual knowledge.
• To improve classification performance on tail classes without requiring fine-tuning of the base model.
• To explore whether retrieval-based augmentation can outperform existing consistency-based or ensemble-based approaches in long-tailed learning.

Proposed method

• RAC integrates a standard base image encoder with a parallel retrieval branch that queries a non-parametric external memory of pre-encoded images and associated text snippets.
• The retrieval module uses a fixed, pretrained image encoder (e.g., ViT) to embed images and a text encoder (e.g., BERT-like) to encode associated text, enabling cross-modal retrieval.
• During inference, the logits from the retrieval module are combined with those from the base encoder to produce final predictions.
• The method uses approximate k-NN search (via HNSW) over a large-scale index to efficiently retrieve relevant image-text pairs during training and inference.
• The external memory is constructed from the training dataset itself—images and their labels or captions—eliminating the need for external data sources.
• Joint training of the base encoder and retrieval branch is performed using a modified LACE loss, enabling end-to-end optimization without fine-tuning the retrieval components.

Experimental results

Research questions

• RQ1 Can a retrieval-based module effectively improve classification performance on long-tail visual recognition benchmarks without fine-tuning?
• RQ2 Does the retrieval module naturally learn to focus on tail classes, even without explicit prompting?
• RQ3 How does the performance of RAC scale with the size and granularity of the external memory index?
• RQ4 Can RAC outperform existing state-of-the-art methods that rely on loss reweighting or logit adjustment?
• RQ5 What is the computational overhead of integrating retrieval into standard classification pipelines?

Key findings

• RAC achieves a 14.5% relative improvement on the Places365-LT benchmark and a 6.7% relative improvement on iNaturalist-2018, surpassing prior state-of-the-art methods.
• The retrieval module learns to focus on tail classes without explicit prompting, significantly reducing the burden on the base encoder.
• Performance gains are most sensitive to label granularity—adding new classes with unique labels yields larger gains than increasing sample count with existing labels.
• The training time increases by only 1.5–2× due to the retrieval module, with the majority of overhead coming from the text encoder rather than the index lookup.
• Index size has diminishing returns beyond a certain point, as the information content of labels plateaus once most classes are represented.
• The method remains efficient at scale, with a 10M+ index searchable in real-time with negligible training overhead on a single node with 8 A100 GPUs.