Prompt-Guided Attention Head Selection for Focus-Oriented Image Retrieval

2025-04-02

AI Summary

This work addresses the challenge in focus-oriented image retrieval (FOIR) under multi-object complex scenes, where pretrained Vision Transformers (ViTs) struggle to precisely respond to user-provided visual prompts (e.g., points, bounding boxes, or masks). We propose a prompt-driven attention head selection mechanism that requires no fine-tuning and modifies neither model weights nor input images. At inference time, it dynamically matches each visual prompt against attention maps generated by individual heads, activating those most sensitive to the target region—thereby jointly capturing local object details and contextual information. As the first zero-shot, training-agnostic, plug-and-play ViT adaptation for FOIR, our method significantly improves retrieval accuracy across multiple benchmarks while preserving the original ViT’s inference efficiency. Extensive experiments validate its effectiveness, broad applicability across ViT variants, and deployment friendliness.

Technology Category

• Computer Vision: Image and Video Retrieval
• Natural Language Processing: Prompt Engineering / Prompting
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The goal of this paper is to enhance pretrained Vision Transformer (ViT) models for focus-oriented image retrieval with visual prompting. In real-world image retrieval scenarios, both query and database images often exhibit complexity, with multiple objects and intricate backgrounds. Users often want to retrieve images with specific object, which we define as the Focus-Oriented Image Retrieval (FOIR) task. While a standard image encoder can be employed to extract image features for similarity matching, it may not perform optimally in the multi-object-based FOIR task. This is because each image is represented by a single global feature vector. To overcome this, a prompt-based image retrieval solution is required. We propose an approach called Prompt-guided attention Head Selection (PHS) to leverage the head-wise potential of the multi-head attention mechanism in ViT in a promptable manner. PHS selects specific attention heads by matching their attention maps with user's visual prompts, such as a point, box, or segmentation. This empowers the model to focus on specific object of interest while preserving the surrounding visual context. Notably, PHS does not necessitate model re-training and avoids any image alteration. Experimental results show that PHS substantially improves performance on multiple datasets, offering a practical and training-free solution to enhance model performance in the FOIR task.

Problem

• Enhance ViT models for focus-oriented image retrieval with visual prompts
• Select attention heads via prompt matching for object-specific retrieval
• Improve FOIR task performance without model retraining or image alteration

Innovation

• Prompt-guided attention head selection
• Leverages multi-head attention in ViT
• Training-free focus-oriented image retrieval