Source-linked AI summary
Trove: A Flexible Toolkit for Dense Retrieval
Reza Esfandiarpoor, Max Zuo, Stephen H. Bach
TL;DR
Dense-retrieval experiments require substantial engineering for data management, model customization, and distributed evaluation. Trove provides a configurable toolkit for on-the-fly data processing and customizable retrieval pipelines; it reduces data-preparation memory use and supports inference scaling across nodes without overhead.
Problem
Method
Trove combines on-the-fly retrieval-data processing with customizable modeling and a unified evaluation and hard-negative-mining interface.
Source: §1 · p. 2Results
Trove reduces data-preparation memory use by 2.6x, while inference time decreases linearly with added nodes and incurs no overhead.
Source: §4.2 · p. 6Takeaways & Limitations
Trove reduces engineering effort in retrieval research while allowing users to customize or replace modeling components.
Source: §1 · p. 2, §5 · p. 6FastResultHeapq is not a full heapq implementation and only mimics functionality for tracking top-k documents per query.
Source: §4.1 · p. 5
Abstract
from arXiv · showhide
We introduce Trove, an easy-to-use open-source retrieval toolkit that simplifies research experiments without sacrificing flexibility or speed. For the first time, we introduce efficient data management features that load and process (filter, select, transform, and combine) retrieval datasets on the fly, with just a few lines of code. This gives users the flexibility to easily experiment with different dataset configurations without the need to compute and store multiple copies of large datasets. Trove is highly customizable: in addition to many built-in options, it allows users to freely modify existing components or replace them entirely with user-defined objects. It also provides a low-code and unified pipeline for evaluation and hard negative mining, which supports multi-node execution without any code changes. Trove's data management features reduce memory consumption by a factor of 2.6. Moreover, Trove's easy-to-use inference pipeline incurs no overhead, and inference times decrease linearly with the number of available nodes. Most importantly, we demonstrate how Trove simplifies retrieval experiments and allows for arbitrary customizations, thus facilitating exploratory research.
1 Introduction
Trove addresses the engineering burden of dense-retrieval experiments with on-the-fly data management, customizable modeling, and a unified multi-node evaluation and hard-negative-mining interface.
1 Introduction
- Existing retrieval toolkits rely on large pre-processed data files, limit modeling customization, and make multi-node evaluation difficult or labor-intensive.These constraints add engineering effort to exploratory experiments.Source: §1 · p. 1
- Trove processes retrieval data on the fly and supports filtering, transforming, and combining data without maintaining large pre-processed files for each experiment.Its design is intended to make data operations fast and memory-efficient.Source: §1 · pp. 1–2
- Trove provides a unified interface for evaluation and hard negative mining that supports multi-node and multi-GPU inference without additional code.The introduction identifies distributed evaluation as a retrieval-specific engineering challenge.Source: §1 · pp. 1–2
- Trove exposes modeling components for users to modify or replace with arbitrary modules while remaining compatible with the Hugging Face transformers ecosystem.The toolkit also offers built-in customization options and documented guides and examples.Source: §1 · p. 2
2 Background and Challenges
Dense-retrieval research spans varied training approaches, but retrieval’s data structure and shared computation make data management, modeling, and distributed evaluation challenging.
2 Background and Challenges
- Existing dense-retrieval work includes methods using mined hard negatives, in-batch negatives, and synthetic training data.The cited literature also includes large decoder-only models, parameter-efficient fine-tuning, and retrieval instructions.Source: §2 · p. 2
- Retrieval datasets require maintaining large pre-processed files with duplicate data, and these files make data changes difficult to track with version control.Creating MS MARCO training samples from raw files requires substantial memory and extra code.Source: §2 · p. 2
- Retrieval toolkits restrict modeling customization by wrapping transformer encoders, and component interactions can prevent training with graduated relevance labels.Tevatron is cited as an example where overwriting the loss does not enable graduated relevance labels.Source: §2 · p. 2
- IR evaluation is computationally demanding and multi-node execution is not straightforward: SentenceTransformers and MTEB are single-node only, while Tevatron requires separate jobs to encode dataset shards, retrieve documents, and calculate metrics.Source: §2 · p. 3
3 System Design
Trove organizes experiments around configuration objects and exposes its main pipeline components for customization or replacement.
3.1 Workflow
- Users configure raw query, corpus, and qrel processing, dataset details, and model settings before instantiating a retriever and a trainer or evaluator.Configuration objects can be created from command-line arguments to support diverse experiments.Source: §3.1 · p. 3
3 System Design
- Users can customize or replace main pipeline modules without changing Trove’s source code or substantially changing their experiment workflow.The workflow centers on configuration objects and components for data, models, training, and evaluation.Source: §3.1 · p. 3
3.2 Data Management
Trove manages retrieval data through an ID-based representation and configurable dataset classes, with caching and callbacks supporting efficient, reproducible processing.
3.2 Data Management
- Preparing retrieval instances from large query, corpus, and qrel files is memory-intensive, while merging millions of qrel records can slow complex pipelines.The passage notes MS MARCO has 500K queries and 8M documents.Source: §3.2 · p. 3
3.2.1 Internal Representation
- MaterializedQRel groups qrel triplets by query ID and stores query, corpus, and qrel records in memory-mapped Arrow tables indexed by ID.It loads only the records needed for each training instance, reducing memory consumption.Source: §3.2.1 · p. 3
3.2.2 User-facing Classes
- MultiLevelDataset and BinaryDataset combine MaterializedQRel instances, allowing separate preprocessing of data sources and version-controllable dataset pipelines.Configuration objects specify the loading and processing steps.Source: §3.2.2 · p. 4, §3.2.1 · p. 4
- EncodingDataset and a memory-mapped embedding cache support lazy loading of cached vectors, while RetrievalCollator tokenizes and batches retrieval examples.The cache returns a stored embedding instead of raw text when one is available.Source: §3.2.2 · p. 4
3.2.3 Additional Optimizations
- Callbacks let users customize loading, query and passage formatting, and qrel filtering; cached artifacts and atomic writes support repeatable, reliable dataset generation.Datasets are fast after the first run, and fingerprinting tracks changes.Source: §3.2.3 · p. 4
3.3 Modeling
Trove separates retrieval modeling into retriever, encoder, and loss components that users can customize independently. Its built-in options and subclass interfaces support both standard and user-defined retrieval behavior.
3.3 Modeling
- Trove divides modeling into retriever, encoder, and loss components, each independently customizable.Users can use arbitrary nn.Module objects as encoders or subclass PretrainedRetriever to implement custom retrieval logic.Source: §3.3 · p. 4
- PretrainedRetriever supports Hugging Face encoders, common pooling and normalization, LoRA adapters, and quantization.BiEncoderRetriever provides dual-encoder retrieval with cross-device in-batch negatives.Source: §3.3 · p. 4
- Users can add encoder wrappers and loss functions as subclasses, then select them through retriever options.Trove includes InfoNCE and KL Divergence losses.Source: §3.3 · p. 4
3.4 Training
Trove supports approximate IR metrics during training and uses Hugging Face's Trainer module with minimal changes. Figure 3 is titled as a training example with mined hard negatives.
3.4 Training
- Trove uses Hugging Face's Trainer module for training, with minimal changes.The components are compatible with the Hugging Face transformers ecosystem.Source: §3.4 · p. 4
- IRMetrics computes approximate IR metrics during training for small MultiLevelDataset instances.It can be used as the compute_metric callback by ranking a small number of annotated documents for each development query.Source: §3.4 · p. 4
- Figure 3 is titled “Training with Mined Hard Negatives.”Source: Figure 3 · p. 5
3.5 Inference
Trove's RetrievalEvaluator provides a unified interface for retrieval evaluation and hard-negative mining. Its inference supports distributed execution and GPU-accelerated top-k tracking.
3.5 Inference
- RetrievalEvaluator provides a simple interface for evaluation and hard-negative mining, and supports logging to Wandb or other trackers.Other experiment trackers can be integrated through callback functions.Source: §3.5 · p. 5
- Distributed inference runs the same script across nodes and GPUs, with shard sizes adjusted to GPU throughput.Faster devices receive more samples to avoid stalling them behind slower GPUs.Source: §3.5 · p. 5
- FastResultHeapq is 16x and 600x faster than Python heapq for online and cached embeddings, respectively.It uses fast matrix operations and GPU acceleration to track top-k documents for each query.Source: §3.5 · p. 5
4 Demonstration
Trove demonstrates flexible retrieval-data configuration and component customization, then benchmarks data-management and inference efficiency. The reported results include reduced memory use, linear inference scaling across nodes, and faster top-k tracking than Python heapq.
4.1 Flexibility and Ease of Use
- The mined-hard-negative training example supports multi-node/GPU training, standard pooling and normalization, LoRA adapters, and quantization.The authors describe the code as requiring little engineering effort for common training setups.Source: §4.1 · p. 5
- Replacing dataset configuration lines combines multi-level synthetic data, annotated positives, and mined hard negatives.The example filters positive labels and randomly selects two hard negatives per query.Source: §4.1 · p. 5
- In SyCL, the example implements Wasserstein distance as a loss and selects it with the --loss=ws option.Source: §4.1 · p. 6
4.2 Efficiency
- 2.6x less memory is required to prepare MS MARCO training data with Trove than with the naive baseline.Combining synthetic and real data adds 0.73 GB of memory for 2M passages, versus 2.45 GB for the naive approach; on 8 GPUs, naive loading uses 90 GB of RAM.Source: §4.2 · p. 6, Table 1 · p. 6
- After the first run, internal caching makes data available almost instantaneously; short time to first sample matters for debugging and interactive development.The appendix reports time to first sample as the time required to load and process data.Source: §4.2 · p. 6
- Inference time decreases linearly as nodes increase, without code changes to the script.Table 2 reports inference time in HH:MM for different node counts.Source: §4.2 · p. 6, Table 2 · p. 6
- More than 600x faster top-k tracking is reported for FastResultHeapq than Python heapq in the online setup.The online setup embeds a batch of 256 documents and compares it with queries on the fly; Python heapq can take up to 130 hours as query counts grow.Source: §4.2 · p. 6, Table 3 · p. 6
- 16x to 21x faster top-k tracking is reported for Trove with cached embeddings and large GPU batches.The comparison uses batches of 40,960 documents; in practice Python heapq speedup is often limited by disk I/O with simple caching mechanisms.Source: §4.2 · p. 6, Table 3 · p. 6
5 Conclusion
Trove streamlines dense-retrieval research with on-the-fly data management, flexible modeling, and unified evaluation and hard-negative-mining support. Its components can be configured, modified, or replaced, helping researchers experiment with new ideas.
5 Conclusion
- Trove loads and processes retrieval data on the fly, eliminating the need for large pre-processed data files while keeping memory use small.Source: §5 · p. 6
- Its unified interface supports evaluation and hard negative mining with multi-node inference and no extra code.Source: §5 · p. 6
- Researchers can configure, modify, or replace Trove components to experiment quickly and freely with new ideas.Source: §5 · pp. 6–7
A Limitations
Trove prioritizes a simple, freely customizable codebase for researchers, so it does not include every retrieval method out of the box. The authors also describe custom alternatives to industry-standard tools, including a search container that can track scores for arbitrary documents.
A Limitations
- Trove does not support every retrieval-related method out of the box, but users can implement arbitrary methods in scripts.Source: A Limitations · p. 8
- The authors sometimes avoid industry standards with custom solutions; they cite large one-use search indexes as inefficient in their case.Source: A Limitations · p. 8
- Trove can track similarity scores for arbitrary documents, including documents outside the top-k results.Source: A Limitations · p. 9
- Table 4 reports time to first sample in seconds for first and subsequent runs.Source: Table 4 · p. 9
B Model Customization
Trove supports model customization through encoder wrappers that users can adapt and select, including for input formatting and pooling. Users can also supply an encoder object directly.
B Model Customization
- Users can modify query and passage formatting methods to add instructions, then select the modified wrapper in their scripts.Source: B Model Customization · p. 9
- Users can modify the default encoder wrapper to implement different pooling operations and select the wrapper through its encoder-class option.Source: B Model Customization · p. 9
- Users can subclass PretrainedEncoder wrappers to control model loading, saving, and embedding calculation, and access the wrappers through configuration options.Source: B Model Customization · p. 9
- Users can instantiate BiEncoderRetriever with any nn.Module object as its encoder.Source: B Model Customization · p. 9