Search-Adaptor: Embedding Customization for Information Retrieval (2024)

Abstract: Embeddings extracted by pre-trained Large Language Models (LLMs) have significant potential to improve information retrieval and search. Beyond the zero-shot setup in which they are being conventionally used, being able to take advantage of the information from the relevant query-corpus paired data can further boost the LLM capabilities. In this paper, we propose a novel method, Search-Adaptor, for customizing LLMs for information retrieval in an efficient and robust way. Search-Adaptor modifies the embeddings generated by pre-trained LLMs, and can be integrated with any LLM, including those only available via prediction APIs. On multiple English, multilingual, and multimodal retrieval datasets, we show consistent and significant performance benefits for Search-Adaptor – e.g., more than 5% improvements for Google Embedding APIs in nDCG@10 averaged over 14 BEIR datasets.

## introduction information retrieval is broadly considered as the task of searching for information via querying corpus database that might consist many different types of data, such as documents, webpages or logs. it has a wide range of applications across many industries, including retail, healthcare, and finance, with a significant portion of the world's economy is built on. particularly, language modeling is the key part of information retrieval as in most cases, query and corpus data are in text form. large language models (llms) have demonstrated significant achievements for a variety of text processing tasks, including question answering, summarization, and mathematical reasoning @xcite @xcite . one critical enabler for the success on these has been transforming raw text into meaningful representations that preserve semantic meanings in the representation space @xcite . for a wide range of applications, from recommendations to anomaly detection, tasks are defined as explicit operations on the learned representations. this makes the quality of the text mapping into embeddings become of paramount importance. information retrieval systems commonly utilize the text embeddings, with relevant corpora being ranked based on the similarity between queries and corpus @xcite . various llms have been proposed to extract embeddings from raw text, with the notable ones including the sentence t5 @xcite , ope-nai embedding apis (ope) and google embedding apis (gcp). however, one fundamental limitation of pre-trained llms is that they cannot utilize retrieval-specific target task data, that are in the form of positively relevant query-corpus pairs. even with a small amount, using such data for tuning is expected to significantly improve information retrieval capabilities. conventional fine-tuning @xcite can be one straightforward way of utilizing the paired query-corpus information. however, if the number of paired samples is small, tuning all the weights of a model might yield overfitting and thus poor generalization @xcite , especially in the presence of distribution shifts. in addition, it can be costly from a computational perspective as it requires large memory. there are multiple parameter-efficient tuning methods such as prompt tuning @xcite , lora @xcite , partial fine-tuning @xcite , and various adapters @xcite . these approaches only tune a subset of the parameters of llms, aiming to reduce the risks of overfitting and bringing computational gains. as a common bottleneck, all of these methods need full access to the llm's parameters to tune the model, which may not be possible for llms with api-based inference-only access. in this paper, our focus is customizing llms to obtain superior embeddings for information retrieval, particularly from the angle of how to best take advantage of the available retrieval-specific tuning data and obtain robust improvements in wide range of regimes, with a tuning method that is low cost and even applicable to llms with apibased inference-only access. fig. 1 overviews the proposed approach, search-adaptor. improving the tuning performance in this setup requires a set of innovations. we propose integration of a low capacity adapter module (to be customized for the target dataset) on top of fixed llms to modify the pre-trained embeddings. for efficient tuning, we introduce a novel differentiable ranking loss that can directly utilize the information of positive query and corpus pairs. in addition, we include multiple regularizers to improve generalization in this small data regime where without intervention, the pretrained llms would end up with catastrophic forgetting. enabled by such design approach, one major advantage of search-adaptor is that it does not require access to the parameters of the pre-trained llms -only the inference outputs of the llms are needed. commercial embedding apis that show state-of-the-art performance usually do not provide access to their model parameters. in such cases, search-adaptor can still be used to further improve those api-based embedding models, in contrast to alternative tuning methods. we demonstrate the effectiveness of search-adaptor across 14 beir datasets @xcite and 17 miracl multilingual datasets @xcite with google and openai embedding apis, applying the search-adaptor on top. in addition, we evaluate search-adaptor's performance improvements with open-source sentence t5 models @xcite . overall, search-adaptor provides significant improvements over alternatives, consistently across different datasets and models. the contributions of this paper are: ## related work pre-trained llms for zero-shot retrieval. llms to extract general text embeddings are commonly used in both academia and industry. ai solution providers like google (gcp) and openai have productionized general text embeddings that can be directly used via simple apis for zero-shot retrieval applications. in addition, many previous efforts have introduced new general text embedding models with various pre-training methods and datasets. gte @xcite proposes a multi-stage pre-training of embedding models with diverse naturally paired text datasets. e5 @xcite pre-trains the embedding models by weakly-supervised contrastive learning, utilizing consistency-based filter to generate high quality text pairs for pre-training. note that search-adaptor can be applicable on top of any pre-trained llm embedding models to customize their embeddings for superior retrieval performances. embedding customization. instead of using one unified model for zero-shot retrieval, the embeddings can be customized for each dataset or task. instruction-based embedding customization is one popular method. tart @xcite builds a retrieval system that adapts the retrieval based on the instruction. different retrieval tasks (e.g., code, question, or answer) are given as the instruction to further improve dense embedding retrieval. instructor @xcite integrates the task and domain descriptions prior to the input to fine-tune the embeddings for retrieval. however, these do not directly utilize the provided relevant query-corpus pairs. full or parameter-efficient fine-tuning (such as lora @xcite ) and (ia) 3 @xcite ) can also be considered for embedding customization. pre-trained llms can be fine-tuned with contrastive loss using positive query-corpus paired data. promptagator @xcite utilizes in-context learning to generate synthetic query-corpus pairs using a few number of original query-corpus pairs, and subsequently using those synthetic pairs to fine-tune the pre-trained llms. however, all these are only applicable when there is full access to the parameters of pre-trained llms, which is often not possible for state-of-theart commercial text embedding models. on the other hand, search-adaptor can be applied without full access to the llm parameters. ## problem formulation we formulate the retrieval problem with a given query-corpus paired dataset. assume a query set denoted as @xmath0, ..., q n } ∈ q and a corpus set denoted as @xmath1, ..., c m } ∈ c. each positive relationship between a query and corpus is represented as the triplet r @xmath2, c j , y ij ) with y @xmath3. we treat all other triplets as negative relationships (y @xmath4. the set of all query-corpus relationships is denoted as @xmath5, c j , y ij )} @xmath6:n,@xmath7:@xmath8, where r @xmath9, c j , y ij ) ∈ r|y @xmath10, c j , y ij ) ∈ r|y @xmath11. note that y ij can be either binary or continuous. query corpus llm encoder llm encoder query embedding query adaptor corpus adaptor cosine similarity adapted corpus embedding adapted query embedding recovery loss (regularizer) ranking loss (supervised loss) prediction loss (regularizer) corpus embedding query predictor recovery loss (regularizer) adapted relevance score query-corpus relevance skip-connection skip-connection predicted query embedding between the given query (q i ) and corpus (c j ) such that the predicted relationship is highly correlated with the ground truth relationship (y ij ). the scoring function f : q × @xmath12. the optimal score is the one that has the same order as the ground truth relationship for each query. ## methods fig. 2 overviews the components of search-adaptor, that are described in following sections. ## experiments we evaluate the performance of search-adaptor in multiple scenarios on numerous datasets. we demonstrate that search-adaptor is modelagnostic, applying it both on top of api-based llms (merely via access to google & openai apis) and open-sourced llms (e.g., sentence t5 @xcite ). we also demonstrate that it is data-agnostic by evaluating search-adaptor on english, multilingual and multimodal datasets. ## discussions 6.1 ablation studies search-adaptor proposes multiple innovations to improve the adaptation performance. we quantify the contributions of proposed constituents to the retrieval performance on various datasets with as st5-base as the base embedding model, with the results in ## conclusions in this paper, we focus on pushing the capabilities of llms for information retrieval and search. we propose a canonical efficient adaptation method, search-adaptor, that can also be applied to llms even with inference-only access. search-adaptor is a low-cost tuning method that brings significant and consistent improvements in retrieval performance across diverse regimes of training data size, encoder type, and corpus set. this is enabled by the judicious design of its adaptor module, along with training objectives and approaches. we have also studied the extension of search-adaptor to multimodal learning and tool use scenarios, highlighting the importance of embedding customization for such applications.