Task-Adaptive Embedding Refinement via Test-time LLM Guidance

TL;DR

This paper introduces a test-time LLM-guided query refinement method that enhances embedding models for zero-shot search and classification, significantly improving their effectiveness in challenging tasks.

Contribution

It presents a novel approach that refines query embeddings in real time using LLM feedback, expanding the practical deployment of embedding models in complex scenarios.

Findings

Up to +25% improvement in search and classification tasks

Refined queries lead to better ranking and clearer separation in embedding space

Method enhances embedding model utility without costly LLM pipelines

Abstract

We explore the effectiveness of an LLM-guided query refinement paradigm for extending the usability of embedding models to challenging zero-shot search and classification tasks. Our approach refines the embedding representation of a user query using feedback from a generative LLM on a small set of documents, enabling embeddings to adapt in real time to the target task. We conduct extensive experiments with state-of-the-art text embedding models across a diverse set of challenging search and classification benchmarks. Empirical results indicate that LLM-guided query refinement yields consistent gains across all models and datasets, with relative improvements of up to +25% in literature search, intent detection, key-point matching, and nuanced query-instruction following. The refined queries improve ranking quality and induce clearer binary separation across the corpus, enabling the embedding space to better reflect the nuanced, task-specific constraints of each ad-hoc user query. Importantly, this expands the range of practical settings in which embedding models can be effectively deployed, making them a compelling alternative when costly LLM pipelines are not viable at corpus-scale. We release our experimental code for reproducibility, at https://github.com/IBM/task-aware-embedding-refinement.