Source-linked AI summary
Soft Prompt Tuning for Augmenting Dense Retrieval with Large Language Models
Zhiyuan Peng, Xuyang Wu, Qifan Wang, Yi Fang
TL;DR
Dense retrieval needs domain-specific training data, while human-written prompts used for LLM-based augmentation may produce weak queries of inconsistent quality. SPTAR learns task-specific soft prompts from limited labeled data, filters demonstration pairs, and generates weak document-query pairs for retriever training. Experiments show that it outperforms BM25 and recent LLM-based augmentation baselines across evaluated datasets and retrieval metrics, with one noted exception.
Problem
Dense retrieval is constrained by limited domain-specific training data, and existing hard prompts cannot guarantee high-quality weak queries.
Method
SPTAR tunes task-specific soft prompts on limited labeled data, selects high-quality prompt examples, and generates filtered weak document-query pairs for dense-retriever training.
Results
SPTAR consistently outperforms W/O Aug and InPars across evaluated datasets and key retrieval metrics, except ColBERT R@100, where InPars is slightly better.
Takeaways & Limitations
Soft prompt tuning and filtering provide an augmentation approach that improves task-specific dense retrieval using limited labeled data and open-source LLMs.
Takeaways & Limitations
Virtual-token representations produce varying performance across datasets, motivating deeper investigation in future work.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Dense retrieval addresses lexical mismatch by comparing dense semantic representations, but specialized domains remain limited by scarce domain-specific training data. SPTAR replaces brittle human-written prompting with task-tuned soft prompts and filtering to generate better weak training pairs, using open-source LLMs for reproducibility.
- Limited domain-specific training data constrains dense retrieval in specialized domains, while transfer learning does not benefit all models and domains equally.
- Hard prompts are difficult to craft and can underperform tuned models, motivating a learned alternative for weak-query generation.Prior augmentation methods use human-written prompts to tag unlabeled documents with weak queries.
- SPTAR tunes task-specific soft prompts on limited ground-truth data, then prompts LLMs to generate weak queries for unlabeled documents.The resulting weak document-query pairs train task-specific dense retrievers.
- A soft prompt filter selects high-quality example document-query pairs, further improving weakly tagged data and retrieval tasks.
- The study uses recent open-source LLMs and publicly available code and data to support reproducible, deterministic experiments.
2 RELATED WORK
Related work spans dense retrieval architectures, data augmentation, LLM-based retrieval, and prompt tuning. SPTAR differs by learning task-specific soft prompts and filtering demonstration pairs before generating weak retrieval data.
- 2.1 Dense Retrieval: Dense retrievers encode queries and documents as vectors for ANN search, while the paper contributes augmentation rather than a new retriever architecture.The reviewed systems include DPR, ColBERT, ANCE, TAS-B, Contriever, ReContriever, and BM25CE.
- 2.2 Data Augmentation for Dense Retrieval: Existing augmentation methods generate or construct additional document-query data through LLM query generation, representation perturbation, teacher retrieval, or distillation.InPars and PROMPTAGATOR use human-written prompts and examples to generate weak queries, while other methods use embedding or teacher-based strategies.
- 2.3 LLMs in Dense Retrieval: LLM-based retrieval work commonly uses frozen LLMs for query, relevance, hypothetical-document, or permutation generation rather than fine-tuning them.
- 2.3 LLMs in Dense Retrieval: SPTAR first learns task-specific soft prompts and then applies data filtering to improve generated weak data, unlike prior LLM augmentation methods.
- 2.4 Prompt Tuning: Prompt-tuning research learns prompt embeddings while preserving the pretrained model, with related methods including Prefix-Tuning, gisting, and other task- or context-aware prompt designs.
3 SOFT PROMPT TUNING FOR AUGMENTING DENSE RETRIEVAL
SPTAR builds task-specific dense retrievers by tuning soft prompts on limited labeled data, selecting effective examples, generating weak queries for unlabeled documents, filtering them, and training retrieval models. Its pipeline covers data preparation, prompt tuning, prompt filtering and augmentation, weak-data filtering, and dense retrieval.
- 3 SOFT PROMPT TUNING FOR AUGMENTING DENSE RETRIEVAL: SPTAR uses six modules to prepare data, tune and filter soft prompts, generate and filter weak queries, and train dense retrievers.The learned prompt and filtered weak data are passed through the augmentation and retrieval stages.
- 3.1 Data Preparation: The domain dataset contains labeled query-document pairs, while a larger collection supplies unlabeled documents for later retrieval and augmentation.After training, DR encodes the collection and searches documents by vector similarity.
- 3.1 Data Preparation: SPTAR samples small training and evaluation sets from labeled data, using evaluation data to monitor convergence and reduce overfitting.The experiments illustrate sampling 50 unique training queries and 100 evaluation queries, with other training sizes also tested.
- 3.2 Soft Prompt Tuning: Soft prompt tuning freezes the LLM parameters Φ and updates only task-specific prompt embeddings initialized from a repeated hard instruction.The prompt is prepended to document-query instances, and training uses negative log-likelihood.
- 3.3 Soft Prompt Filter: The soft prompt filter samples groups of example pairs, evaluates their loss on held-out data, and selects the group with the best evaluation metric for augmentation.Exhaustive evaluation can be impractical, so the method randomly samples candidate groups before selection.
- 3.4 Weak Data Filter: The weak-data filter retains a generated pair only when BM25 ranks its linked document among the corpus top k results.For MS MARCO and FiQA-2018, k was selected from {10, 30, 50, 70} using evaluation performance.
- 3.5 Soft Prompt Augmentor: The augmentor combines the learned soft prompt with filtered example pairs to generate weak queries for unlabeled documents, producing 100K-pair and 5K-pair datasets.These datasets are subsequently filtered before training dense retrievers.
- 3.6 Dense Retrieval: The final retrieval stage evaluates DPR, ColBERT, TAS-B, Contriever, ReContriever, and BM25-based reranking variants.The models serve as dense retrievers, bi-encoder rerankers, or cross-encoder rerankers over BM25 results.
4 EXPERIMENTAL SETUP
The experiments evaluate SPTAR across four datasets using controlled training settings, standard retrieval metrics, multiple baselines, and seven research questions. The setup compares augmentation strategies while examining prompt behavior, computational cost, filtering, and data-size effects.
- 4.1 Datasets: Experiments use four datasets from BEIR and DL2019/DL2020, with MS MARCO evaluation reported on its evaluation split.The dataset statistics are summarized in Table 2.
- 4.1 Datasets: The comparison partitions data among unaugmented models, InPars, and SPTAR, using matched evaluation and test splits for fairness.SPTAR combines limited labeled data with filtered weak data generated by its soft prompt augmentor.
- 4.2 Training Details: Soft prompt tuning uses LLaMA-7B and Vicuna-7B, while dense-retriever training follows separate hyperparameters and checkpoint selection procedures.Table 4 reports prompt-tuning settings, and ColBERT checkpoints are selected by evaluation of NDCG@10.
- 4.3 Evaluation Metrics: Evaluation uses perplexity for prompt modules and retrieval metrics including MRR@10, Recall@100, MAP, and nDCG@10 across datasets and reranking models.Recall@100 specifically evaluates BM25CE and BM25BE reranking effectiveness.
- 4.4 Baseline Methods: Baselines include BM25, Contriever, ReContriever, Without Augmentation, and InPars, with prompt and generator choices controlled for comparison.InPars uses its human-written prompt with Vicuna, whereas SPTAR uses its optimized soft prompt with LLaMA.
- 4.5 Research Questions: The study asks whether SPTAR improves retrieval over baselines and whether soft prompt tuning distills dataset knowledge into learned prompts.These questions correspond to RQ1 and RQ2.
- 4.5 Research Questions: Additional questions examine tuning cost, soft-prompt filtering, weak-data filtering, training-data size X, and example-pair count M.These questions correspond to RQ3–RQ7.
5 EXPERIMENTAL RESULTS
Across seven retrievers and re-rankers, SPTAR generally outperforms non-augmented and InPars baselines, while ablations show that soft prompts, example filtering, and weak-data filtering affect generation and retrieval quality.
- 5.1 SPTAR vs Baseline Models (RQ1): SPTAR consistently outperforms W/O Aug and InPars across datasets and retrieval metrics, except ColBERT’s R@100 on MS MARCO-related datasets; 55 of 70 improvements are significant.TAS-B achieves four of ten best results, while ColBERT accounts for eight of the 15 nonsignificant improvements.
- 5.1 SPTAR vs Baseline Models (RQ1): SPTAR improves multiple retrievers and re-rankers, but ColBERT gains are limited by token-level sensitivity to noisy generated queries, whereas DPR generalizes better.BM25CE surpasses BM25BE on 6 of 10 metrics with augmentation, reversing their small-data ordering without augmentation.
- 5.2.1 The Impact of Soft Prompt Tuning Module (RQ2).: Learned soft-prompt token distributions separate by dataset, vary with LLM and prompt length, and therefore encode task-specific information.LLaMA-7B and Vicuna-7B show greater token dispersion than GPT-2, while different prompt lengths retain substantial overlap.
- 5.2.2 The Efficiency of Soft-Prompt Tuning (RQ3).: Soft-prompt tuning updates roughly 0.003% of the LLM parameters and converges within four or five epochs on more advanced language models.The original LLM parameters remain frozen, keeping the trainable prompt size constant as the LLM grows.
- 5.2.3 The Impact of Soft Prompt Filter Module (RQ4).: Selecting better example document-query pairs reduces perplexity by 12.60%–98.59% and improves downstream DPR NDCG@10 by 3.67%–11.44%.The results show that in-context examples materially affect weak-query generation quality.
- 5.2.4 The Impact of Weak Data Filter Module (RQ5).: Weak-data filtering raises MS MARCO NDCG@10 from 0.2319 to 0.2580 at top-k=30 and FiQA-2018 from 0.2242 to 0.2404 at top-k=70.Because no correlation with NDCG@10 was found, top-k requires dataset-specific tuning.
- 5.2.4 The Impact of Weak Data Filter Module (RQ5).: Increasing soft-prompt tuning data to X=50 decreases perplexity by 99.78% and improves downstream NDCG@10 by 37.66% versus no soft-prompt tuning.Perplexity improves more readily than NDCG@10, indicating a mismatch between generation quality and downstream retrieval metrics.
- 5.2.4 The Impact of Weak Data Filter Module (RQ5).: The best number of filtered examples depends on dataset: M=2 performs best for MS MARCO, whereas M=1 performs best for FiQA-2008.The authors attribute this variation to differing dataset distributions and multiple relevant queries per document.
6 CONCLUSION AND FUTURE WORK
SPTAR addresses limited domain-specific DR data by tuning soft prompts on small labeled sets and using LLMs to generate filtered weak queries for retriever training. The authors report broad experimental effectiveness but identify limited dataset coverage, pointwise loss, and untapped multitask tuning as future concerns.
- 6 CONCLUSION AND FUTURE WORK: SPTAR tunes task-specific soft prompts on limited ground-truth data, then generates weak queries for unlabeled documents and filters examples and weak data.The resulting weak document-query pairs train domain-specific dense retrievers.
- 6 CONCLUSION AND FUTURE WORK: Experiments validate SPTAR’s effectiveness, while future work should test broader datasets, pairwise or listwise losses, and multitask soft-prompt tuning.The current pointwise loss underuses negative instances, and broad applicability remains to be examined.
A FILTERED EXAMPLE DOCUMENT-QUERY PAIRS FOR MS MARCO AND LLAMA
For MS MARCO with LLaMA, the filtered example configuration with M=2 performs better than M=1. The examples pair documents about new-driver costs and Oakland weather with task-specific queries.
- M=2 is better than M=1 for MS MARCO when LLaMA is employed.
- One example pairs information about young drivers’ insurance costs with the query “average insurance cost for new drivers.”
- Another example pairs an Oakland weather forecast with the query “weather in oakland md.”
B FILTERED EXAMPLE DOCUMENT-QUERY PAIRS FOR FIQA-2018 AND LLAMA
For FiQA-2018 with LLaMA, the filtered example configuration with M=1 performs better than M=2. The examples include agricultural-tax and payment-card documents.
- M=1 is better than M=2 for FiQA-2018 when LLaMA is employed.
- One example concerns whether an amount is taxable if a grandfather sells agricultural land.
- Another example discusses why a payment-card provider may not disclose which merchant leaked card details.