Vision transformers improve image search with region-aware tokens

Region-Aware CLS Token Augmentation for Fine-Grained Image Retrieval

Computer Vision and Pattern RecognitionMachine Learning

Summary

Images are often summarized into a single code to find similar pictures, but this can miss important details, especially when looking for very similar items. The authors use special tokens from a type of AI model called vision transformers to capture more detailed parts of images automatically. By matching these detailed parts with broader image summaries, their method helps find closer matches without needing extra data about where objects are in pictures. Their work shows that using multiple detailed image codes together helps find fine differences more accurately and efficiently.

What this means in practice

  • For e-commerce product teams: Improve visual search to better distinguish similar product images for recommending precise matches to shoppers.$Commercial implications: The result enables finer image matching powering commercial visual search engines for online retail with improved accuracy.
  • For digital asset managers: Enhance large-scale photo libraries with faster, fine-grained image search using multiple localized image features.

Authors

Ian de Holanda Cavalcanti Bezerra, Vivek Trivedy, Lucas Pascotti Valem, Longin Jan Latecki

Abstract

Image retrieval methods often rely on a single global semantic descriptor extracted from an image, e.g., the [CLS] token in vision transformers. However, trying to squeeze all the semantic information of an image into a single descriptor can hurt downstream retrieval performance, especially for fine-grained retrieval tasks. In this work, we augment the semantic tokens in the newer visual transformers, the global [CLS] token and the four register tokens, with a carefully selected collection of spatial tokens, aiming to capture the spatial region representation that characterizes the contents captured in each of the semantic tokens. We leverage the DINOv2-reg model, which includes register tokens that emergently learn object and part-based representations. For each "cue" token ([CLS] and each register token), we find a "buddy" image patch token and extract an N x N patch region to produce a set of localized ROI tokens. Our approach automatically captures important regions of interest without any external bounding boxes or saliency modules, purely by matching semantic tokens with their spatial representation regions. Furthermore, we incorporate these tokens into a multi-vector retrieval framework inspired by ColBERT, enabling fine-grained matching via a per-token alignment mechanism while avoiding the large storage cost of keeping all patch embeddings. Through extensive experiments, we find that (1) register tokens encode useful fine-grained details that can complement the [CLS] token; (2) automatically pooled ROI tokens further improve fine-grained discrimination; and (3) multi-vector retrieval with a small set of tokens improves over a DINOv2-reg single-vector baseline while remaining tractable for large-scale search. The code is available at https://github.com/IdhcbIan/Augmenting_CLS_with_ROI_tokens.