Soft Prompt Tuning for Augmenting Dense Retrieval with Large Language Models

TL;DR

This work tackles domain-specific data scarcity in dense retrieval by introducing Soft Prompt Tuning for Augmenting Dense Retrieval (SPTAR). It learns a task-specific soft prompt on limited labeled data and uses LLMs to tag unlabeled documents with weak queries, generating ample weak document-query pairs to train domain-tailored dense retrievers. A soft prompt filter and a BM25-based weak data filter refine the prompt guidance and the generated data, respectively, yielding substantial improvements across seven retrievers on BEIR-style benchmarks. The results show that SPTAR consistently outperforms unsupervised baselines and prior LLM-based augmentation methods, with efficiency gains due to the small number of trainable parameters and rapid convergence. This approach provides a reproducible, data-efficient path to leverage LLMs for enhancing DR in low-data regimes, with open-source tooling and datasets enabling broader adoption and future extensions.

Abstract

Dense retrieval (DR) converts queries and documents into dense embeddings and measures the similarity between queries and documents in vector space. One of the challenges in DR is the lack of domain-specific training data. While DR models can learn from large-scale public datasets like MS MARCO through transfer learning, evidence shows that not all DR models and domains can benefit from transfer learning equally. Recently, some researchers have resorted to large language models (LLMs) to improve the zero-shot and few-shot DR models. However, the hard prompts or human-written prompts utilized in these works cannot guarantee the good quality of generated weak queries. To tackle this, we propose soft prompt tuning for augmenting DR (SPTAR): For each task, we leverage soft prompt-tuning to optimize a task-specific soft prompt on limited ground truth data and then prompt the LLMs to tag unlabeled documents with weak queries, yielding enough weak document-query pairs to train task-specific dense retrievers. We design a filter to select high-quality example document-query pairs in the prompt to further improve the quality of weak tagged queries. To the best of our knowledge, there is no prior work utilizing soft prompt tuning to augment DR models. The experiments demonstrate that SPTAR outperforms the unsupervised baselines BM25 and the recently proposed LLMs-based augmentation method for DR.

Figures (6)

• Figure 1: The pipeline of the proposed Soft Prompt Tuning for Augmenting dense Retrieval (SPTAR).
• Figure 2: The main architecture of the proposed SPTAR: a) The same LLM is shared by soft prompt tuning module, soft prompt filter module and soft prompt augmentor module; b) soft prompt tuning module fixs the LLM's original parameters $\Phi$ and only fine-tune the parameters of soft prompt's embedding layer $\theta$ on the sampled small dataset (Section \ref{['sec: data_preparation']}); c) soft prompt filter module fixs the learned parameters $\theta^{\ast}$, and for each group of sampled example document-query pairs, computes the loss on evaluation dataset. The group of example document-query pairs with the smallest loss will be utilized in the soft prompt augmentor module; d) with the learned parameters $\theta^{\ast}$ and a group of filtered example document-query pairs, the soft prompt augmentor module iterates over the unlabeled document dataset $D_{unlabeled}$ to generate weak queries.
• Figure 3: T-SNE embedding visualization of soft prompt's virtual tokens: a) soft prompt’s virtual tokens with different datasets; b) soft prompt’s virtual tokens with different LLMs; c) virtual tokens of soft prompt with different lengths.
• Figure 4: SPTAR-DPR NDCG@10 scores with different top-$k$ of weak data filter. SPTAR-DPR is trained on $S_{train}^{50}+S_{eval}^{100}+F_{k}(W_{large})$ (Section \ref{['sec: dataset']}). Results are obtained on LLaMA-7B. For MS MARCO and FiQA-2018, $M=2$ and $M=1$ respectively.
• Figure 5: Evaluation of SPTAR-DPR with different $X$ compared with W/O (Section \ref{['sec: dataset']}). SPTAR-DPR is trained on $S_{train}^{X}+S_{eval}^{100}+W_{small}$ and tested on $D_{test}$. Results are obtained on LLaMA-7B and MS MARCO.