PaperlayerView exact revision on arXiv

Source-linked AI summary

DRAMA: Diverse Augmentation from Large Language Models to Smaller Dense Retrievers

Xueguang Ma, Xi Victoria Lin, Barlas Oguz, Jimmy Lin, Wen-tau Yih, Xilun Chen

arXiv:2502.18460v2cs.CLcs.IR

TL;DR

LLM retrievers can generalize effectively but incur high inference costs, while smaller retrievers may generalize poorly with limited supervised data. DRAMA combines pruned LLM backbones with diverse LLM-generated training data, and reports strong results across English and multilingual retrieval tasks. The paper identifies domain-specific supervised data as a boundary of its evaluation setup.

  • Problem

    LLM retrievers have high inference costs, while smaller retrievers often generalize poorly with limited supervised fine-tuning data.

    Source: §1 · p. 1

  • Method

    DRAMA combines pruned LLM backbones with diverse LLM-augmented data in a single-stage training framework for smaller dense retrievers.

    Source: §1 · pp. 1–2

  • Results

    DRAMA reports strong performance across English and multilingual retrieval tasks and outperforms traditional encoder-based retrievers using multi-stage contrastive learning.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    The results highlight the potential of aligning smaller retriever training with ongoing advances in LLMs.

    Source: §7 · p. 9

    The E5 fine-tuning setup excludes domain-specific tasks such as financial and medical retrieval, though broader supervised data could further improve performance across domains.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Large language models (LLMs) have demonstrated strong effectiveness and robustness while fine-tuned as dense retrievers. However, their large parameter size brings significant inference time computational challenges, including high encoding costs for large-scale corpora and increased query latency, limiting their practical deployment. While smaller retrievers offer better efficiency, they often fail to generalize effectively with limited supervised fine-tuning data. In this work, we introduce DRAMA, a training framework that leverages LLMs to train smaller generalizable dense retrievers. In particular, we adopt pruned LLMs as the backbone and train on diverse LLM-augmented data in a single-stage contrastive learning setup. Experiments show that DRAMA offers better multilingual and long-context capabilities than traditional encoder-based retrievers, and achieves strong performance across multiple tasks and languages. These highlight the potential of connecting the training of smaller retrievers with the growing advancements in LLMs, bridging the gap between efficiency and generalization.

1 Introduction

DRAMA targets the tension between LLM retrievers’ generalization and their high inference costs, and smaller retrievers’ limited generalization with little supervised data. It combines LLM-augmented training data with pruned LLM backbones to train smaller dense retrievers.

1 Introduction

  • LLM retrievers can generalize better than pre-LLM-era smaller models, but Llama3.18B incurs nearly 40× the inference cost of a BERT-based dense retriever.The stated costs include encoding large corpora and increased query latency.Source: §1 · p. 1
  • The framework studies LLMs for both augmentation data and smaller retriever backbones, aiming for generalizable, efficient dense retrievers with fewer than 1B parameters.Its augmentation study compares methods individually and in combination; its pruned decoder-only backbones are reported to retain multilingual and long-context capabilities.Source: §1 · pp. 1–2
  • The augmentation study uses the same models and corpora across methods and relies on open-source models and open-access data.It includes approaches ranging from cropped-sentence queries and corpus retrieval to LLM-generated queries and relevance judgments.Source: §1 · p. 1
  • The paper prunes decoder-only LLMs into smaller retriever backbones and uses bidirectional attention during retriever training.The proposed approach is presented as a flexible way to create smaller retrievers that retain pretrained LLM weights and remain compatible with later LLM advances.Source: §1 · pp. 1–2
  • DRAMA combines diverse LLM-augmented data with pruned LLM backbones in a single-stage framework, producing smaller retrievers that perform strongly across English and multilingual retrieval tasks.The paper reports results on BEIR, MIRACL, and multilingual MTEB tasks, outperforming traditional encoder-based retrievers that use multi-stage contrastive learning.Source: §1 · p. 2

2 Related Work

Prior work develops dense retrievers through specialized pretraining and multi-stage contrastive learning, while LLM-based ranking and data augmentation offer alternative sources of relevance signals. DRAMA studies these approaches alongside the challenge of balancing multilingual performance with English generalization.

2.1 Robust Dense Retrieval

  • Dense retrievers have been improved through customized language-modeling architectures and, in many cases, two-stage training that moves from large-scale contrastive learning to supervised contrastive learning.Other work also integrates corpus-level information into document embeddings.Source: §2.1 · p. 2

2.2 LLM for Text Ranking

  • LLM retrievers and prompted LLM rerankers have shown strong effectiveness, motivating the use of their relevance modeling during training rather than at inference.The paper frames this as shifting high inference-time costs into data augmentation.Source: §2.2 · p. 2

2.3 Data Augmentation for Retriever

  • Existing augmentation methods generate synthetic queries or query-positive-negative data, but prior approaches had not been systematically compared with LLMs and corpora controlled within one framework.DRAMA compares types of LLM augmentation individually and in combination.Source: §2.3 · pp. 2–3

2.4 Multilingual Retriever

  • Multilingual retrievers can trade off multilingual retrieval performance against English generalization, while DRAMA addresses this challenge with LLM-based augmentation and pruned LLM backbones.ArcticEmbV2 is described as pursuing both capabilities using a prior training paradigm of contrastive pretraining followed by supervised fine-tuning.Source: §2.4 · p. 3

3 Methods

The methods construct query-positive-negative training data from retrieval and LLM judgments, then train smaller retrievers using contrastive learning. They also prune decoder-only LLMs to create retriever backbones at selected sizes.

3.1 Data Augmentation for Contrastive Dense Retriever Training

  • Contrastive retriever training uses a query, a relevant positive document, and hard negatives, with InfoNCE as the training loss.The paper frames augmentation as creating query-positive-hard-negative triplets and assumes access to initial supervised data, an LLM retriever, an instruction-following LLM, and a large corpus.Source: §3.1 · p. 3

3.1.1 Data Augmentation via Llama-3.18B Retriever

  • Augmentation uses cropped sentences as pseudo-queries or LLM-generated synthetic queries, then retrieves candidates to select positives and hard negatives; instruction-following LLMs can further rerank candidates.For cropped-sentence queries, the top 10 of 50 candidates are positives and the bottom 20 are hard negatives; listwise reranking uses the top reranked candidate as positive and the bottom 10 of 20 as hard negatives.Source: §3.1.1 · pp. 3–4, Figure 1 · p. 4

3.1.4 Triplet Generation from Llama-3.370B-Instruct

  • Prompted synthetic query-positive-negative triplets do not substantially improve smaller retriever training in the reported experiments.The method first prompts an instruction-following LLM to brainstorm retrieval tasks, then generate a positive and a hard negative for each query.Source: §3.1.4 · pp. 4–5

3.2 Pruning

  • Structured pruning learns parameter masks under architectural constraints, followed by continuous pretraining on the language-modeling objective to recover performance.The paper prunes Llama3.21B to 0.1B and 0.3B parameter sizes; pruning from an LLM is presented as supporting multilinguality, longer contexts, and flexible model sizing.Source: §3.2 · p. 5

4 Experiment Setup

The experiments use controlled fine-tuning data, LLM-based augmentation, and pruned retriever models. Evaluation spans generalization, multilingual retrieval, and long-context tasks against representative baselines.

4.1 Fine-tuning Data

  • Fine-tuning follows the E5 data setup, which includes general-domain retrieval data but excludes domain-specific tasks such as financial QA and scientific document retrieval.The authors use an open-source replication of the E5 fine-tuning data to balance generalization assessment with baseline comparison.Source: §4.1 · p. 5

4.2 Data Augmentation

  • The English LLM retriever is fine-tuned on MS MARCO and E5 data, while the multilingual retriever is further fine-tuned using MIRACL training data.Augmentation corpora include 25M sampled English web documents and multilingual Wikipedia and web corpora covering 19 non-English languages, segmented into chunks up to 256 tokens.Source: §4.2 · pp. 5–6

4.3 Pruning

  • The team prunes Llama3.21B into 0.1B and 0.3B models, with a maximum context length of 8,192 tokens.The first pruning stage costs about 0.5 billion tokens, and the second continuous-pretraining stage costs 26 billion tokens.Source: §4.3 · p. 6

4.4 Training

  • Smaller retrievers train on three augmentation types across English web, multilingual Wikipedia, and multilingual web data, mixed with approximately 2M E5 supervised instances.The three augmentation sources are sampled at a 2:1:1 ratio; 0.1B and 0.3B models use one positive and seven hard negatives per query, while the 1B model uses three hard negatives.Source: §4.4 · p. 6

4.5 Evaluation

  • Evaluation uses nDCG@10 across BEIR, MIRACL, multilingual MTEB subsets, and long-context benchmarks including MLDR, NeedleRetrieval, and PasskeyRetrieval.MLDR covers long-context multilingual retrieval across 13 languages.Source: §4.5 · p. 6

4.6 Baselines

  • The representative baselines include Contriever, DRAGON, E5, BGE, mE5, BGE-M3, mGTE, ArcticEmbV2, Gecko, and MistralE5.Baseline selection considers retrieval-task training-data settings described for the experiments.Source: §4.6 · p. 6

5 Results

DRAMA variants show strong English and multilingual retrieval results, while pruned LLM backbones retain multilingual and long-context performance, including on tasks outside their training conditions.

5.1 Generalization of Smaller Retrievers

  • On BEIR, DRAMA0.1B scores 56.9 nDCG@10; DRAMA0.3B scores 58.0, 0.8 points above ArcticEmb-v2-L and matching 1B Gecko.Table 1 reports benchmark effectiveness in nDCG@10.Source: §5.1 · p. 6, Table 1 · p. 7
  • On MIRACL, DRAMA0.3B is 5.5 points above ArcticEmb-v2-L; it generalizes robustly on MTEB-FR and trails ArcticEmb-v2 by under one point on MTEB-ZH.Source: Table 1 · p. 7, §5.1 · p. 7

5.2 Pruned LLM as Retriever Backbone

  • Pruned LLM backbones retain multilingual retrieval capability with less than 10B tokens of multilingual web data, unlike many 0.1B encoder retrievers based on bert-base-uncased.Source: §5.2 · p. 7
  • On MLDR, DRAMA0.1B reaches 46.8 nDCG@10 despite no MLDR training or training beyond 256 tokens, and beats M3-BGE-Dense by 2.1 points.M3-BGE-Dense used long-context contrastive pretraining but was not fine-tuned on MLDR.Source: Table 2 · p. 7, §5.2 · pp. 7–8
  • On NeedleRetrieval, DRAMA0.1B beats same-size E5-v2 variants at all input lengths; DRAMA1B leads all lengths, while DRAMA0.3B trails BGE-M3-Dense at 4096 and 8192 tokens.DRAMA0.3B outperforms BGE-M3-Dense at 512 and 32768 tokens.Source: Table 3 · p. 8, §5.2 · p. 8
  • After further fine-tuning on MLDR, DRAMA surpasses BM25 and other methods that include MLDR in their training data.Source: §5.2 · p. 8

6 Analysis and Ablation Study

The ablations examine how augmentation, backbone choice, attention and pooling, and pruning-stage language exposure affect smaller retrievers. Together, they show gains from combining augmentation and from using pruned LLM backbones with bidirectional attention, while zero-shot results vary with language exposure during pruning.

6.1 Effectiveness of Data Augmentation

  • Combining all three augmentation types achieves the highest performance; Sent plus QGen scores 53.1 versus 54.5 on BEIR with all three, at lower generation cost.Omitting LLM ranking lowers generation cost by an order of magnitude while losing around 2–3% of retrieval effectiveness.Source: Figure 2 · p. 8, §6.1 · pp. 8–9
  • Without augmentation generalization is poor; LLM-based augmentation improves BEIR performance except for triplets generated by directly prompting Llama3.370B-Instruct.The triplets comprise queries, positive documents, and negative documents.Source: §6.1 · p. 8

6.2 Effectiveness of Model Backbone

  • At 0.1B, the pruned model beats BERT by approximately 1 point on average; at 0.3B, it beats XLM-RoBERTa-Large by about 1.5 points.The 0.1B model also performs slightly better than ModernBERT while retaining multilingual support and leveraging existing LLM pretraining.Source: Table 4 · p. 9, §6.2 · p. 9

6.3 Attention and Pooling Mechanism

  • Bidirectional attention outperforms unidirectional attention, and mean pooling scores higher than last-token pooling; attention has the larger impact.Even with extensive augmented training data, unidirectional attention remains limiting; enabling bidirectional attention improves the small decoder-only model as a retriever.Source: Table 5 · p. 9, §6.3 · p. 9

6.4 Cross-lingual Generalization

  • For zero-shot German, effectiveness improves with model size; for Yoruba, the 0.3B model beats the 1B model, while unseen Polish shows a gap versus the 1B model.The authors suggest pruning-stage Yoruba exposure benefits Yoruba results and note the importance of including a language during pruning.Source: Table 6 · p. 9, §6.4 · p. 9

7 Conclusion

DRAMA integrates LLM pruning and diverse LLM data augmentation to train smaller, generalizable dense retrievers. It achieves strong performance across English and multilingual retrieval tasks.

7 Conclusion

  • DRAMA combines LLM pruning and diverse LLM data augmentation to train smaller, generalizable dense retrievers that perform strongly across English and multilingual retrieval tasks.The conclusion presents the framework as a way for smaller retriever training to improve alongside advances in LLMs.Source: §7 · p. 9

Limitations

The authors identify limits in language coverage, supervised training data, context length, training stages, and task scope, and describe these as areas for further investigation.

Limitations

  • DRAMA’s pruning covers 20 languages, and including a language during pruning is crucial for generalization to that language.The authors suggest broader language coverage could improve performance for low-resource languages and zero-shot multilingual retrieval.Source: Limitations · p. 10
  • The supervised fine-tuning setup excludes domain-specific retrieval tasks such as financial and medical retrieval.The authors say broader supervised datasets could further improve performance across diverse domains.Source: Limitations · p. 10
  • DRAMA is trained with context length up to 256, leaving efficient integration of long-context training data open for further exploration.The authors suggest organizing training batches by context length as one possible approach.Source: Limitations · p. 10
  • DRAMA directly fine-tunes a pruned LLM in a single stage; whether multi-stage pretraining or distillation would improve effectiveness remains open.The authors identify multi-stage pretraining (Yu et al., 2024) and multi-stage distillation (Zhang et al., 2025) as possibilities.Source: Limitations · p. 10
  • DRAMA focuses on retrieval, leaving integration of supervised fine-tuning and LLM-augmented data for broader embedding, instruction-following, and reasoning-intensive tasks to future work.The passage names clustering and classification among broader text-embedding tasks (Lee et al., 2024; Wang et al., 2024c; Su et al., 2023; Asai et al., 2023; Weller et al., 2025; Shao et al., 2025).Source: Limitations · p. 10

Ethics Statement

The authors state that, to the best of their knowledge, the paper raises no ethical issues.

Ethics Statement

A Appendix

Figure 3 concerns the effectiveness of English and multilingual augmentation mixtures across three data sources, using a 0.1B backbone and one hard negative per query.

A Appendix

  • Figure 3 compares effectiveness across different English-to-multilingual data augmentation mixture ratios.Source: Figure 3 · p. 14
  • The figure covers augmentation data sources Web, mWeb, and mWiki.Source: Figure 3 · p. 14
  • The reported setup uses a 0.1B backbone and one hard negative per query.Source: Figure 3 · p. 14

A.1 Multilingual Data Balance

The appendix examines multilingual data balance and representation dimensionalities, then details pruned model configurations, training and evaluation setups, and the full evaluation tables.

A.1 Multilingual Data Balance

  • A 1:1 English-to-multilingual data balance with both mWiki and mWeb gives the best reported trade-off across in-domain multilingual effectiveness and multilingual generalization.Excluding mWeb harms multilingual generalization, while excluding mWiki lowers in-domain multilingual effectiveness; increasing English data slightly helps BEIR but significantly weakens multilingual retrieval.Source: §A.1 · p. 14

A.2 Matryoshka Representation Learning

  • Above 256 dimensions, larger DRAMA models show higher effectiveness, and representations largely retain effectiveness relative to full dimensionality; at 64 dimensions, the 0.1B model outperforms the 0.3B and 1B models.The authors note that 64 dimensions were not an MRL target for the larger models, and that target dimensionality choices matter for test-time effectiveness.Source: §A.2 · pp. 14–15, Figure 4 · p. 15

A.3 Pruned Model Configuration

  • DRAMA0.3B uses 16 layers, 1024 hidden size, 4096 intermediate size, and 16 attention heads; DRAMA0.1B uses 12 layers, 768 hidden size, 3072 intermediate size, and 12 heads.Source: §A.3 · p. 15

A.4 Detailed Training Setup

  • The pruning, augmentation, and backbone models follow their respective Llama 3.1, Llama 3.3, and Llama 3.2 community licenses; pruning and augmentation web corpora cover 20 named languages.The language list includes English, Arabic, Bengali, Spanish, Persian, Finnish, French, Hindi, Indonesian, Japanese, Korean, Russian, Swahili, Telugu, Thai, Chinese, German, Yoruba, Italian, and Portuguese.Source: §A.4 · p. 15
  • Training uses dpr-scale2 on 32 A100 GPUs for about two days, with all three variants trained for 200,000 steps and batch sizes and hard-negative counts varying by model size.DRAMA0.1B uses batch size 2048 and seven hard negatives; DRAMA0.3B uses 1024 and seven; DRAMA1B uses 256 and three.Source: §A.4 · p. 15

A.5 Detailed Evaluation Results

  • BEIR and MIRACL are evaluated with tevatron3, while MTEB-FR/ZH/DE use mteb; evaluation context limits are 512 for BEIR/MIRACL and 1024 for MTEB-FR/ZH/DE.The authors re-evaluate MTEB-FR/ZH/DE baselines for consistency because of version changes in some MTEB-FR datasets.Source: §A.5 · p. 15
  • The appendix provides full evaluation tables for BEIR, MIRACL, and MLDR.Source: §A.5 · p. 15, Table 7 · p. 16, Table 8 · p. 16
  • The appendix also provides full tables for MTEB-FR, MTEB-ZH, MTEB-DE, NeedleRetrieval, and PasskeyRetrieval.Source: §A.5 · p. 15, Table 10 · p. 16, Table 11 · p. 17

Related papers

Loading 2502.18460v2…