Tevatron: An Efficient and Flexible Toolkit for Dense Retrieval

Having data ready to use is a critical preliminary step before training or encoding starts. Data access overhead and constraints could directly affect training/encoding performance. In Tevatron, we adopt the following core design: 1) text data are pre-tokenized before training or encoding happens, 2) keep tokenized data memory-mapped instead of lazy-loaded or in-memory. The former avoids overheads when running sub-word/piece level tokenizers and also reduces data traffic compared to raw text. The latter allows random data access in the training/encoding loop without consuming a large amount of physical memory.

Tevatron defines two basic raw input format templates for IR and QA context. As shown in Fig. 1, for the IR dataset (e.g. MS MARCO (Bajaj et al., 2018)), we organize a training instance into an anchor query, a list of positive target texts, and a list of negative target texts. The positive targets are usually human judged and the negative texts are usually non-relevant texts from top results of a baseline retrieval system such as BM25.

    {
       "query_id": "<query id>",
       "query":    "<query text>",
       "positive_passages": [
         {"docid": "<passage id>",
          "title": "<passage title>",
          "text":  "<passage body>"}, ...
       ],
       "negative_passages": [
         {"docid": "<passage id>",
          "title": "<passage title>",
          "text":  "<passage body>"}, ...
       ]
    }

The second format (not shown due to space limits) has an additional answers field for QA tasks (e.g. Natural Question (Kwiatkowski et al., 2019)), since the positive passages for QA dataset are usually judged by answer exact match (Chen et al., 2017; Karpukhin et al., 2020).

Users can pass raw data file pointer and processing specifications to Tevatron’s dataset class (HFTrainDataset for training, HFQueryDataset and HFCorpusDataset for encoding) which will perform fast parallel data formatting and tokenization. Processed data is internally represented as a datasets.Dataset object and is stored in Apache Arrow format which can be memory-mapped and randomly accessed by offset.

For researchers who are focusing on building new models, we make a collection of popular open-access datasets self-contained within the Tevatron toolkit. For instance, with a single line of command, one can load the training set of MS-MARCO. Under the hood, Tevatron will first download the raw data set we hosted through Huggingface²²2https://huggingface.co/tevatron. Then it will run the corresponding pre-defined pre-processing script to format and tokenize the downloaded data.

Tevatron’s model class DenseModel is a Pytorch nn.Module subclass that defines the deep neural encoder of the dense retriever. Functionally, it interfaces the underlying Transformer models and provides methods for text encoding and loss computation. Thanks to duck typing in python, DenseModel class support models in the Huggingface transformer library that return standard base model output. This means new Transformers models can be loaded into Tevatron as soon as they are available in the transformer library. On the other hand, this helps Tevatron avoid maintaining the transformer codes and reduce code reduplication. Internally, DenseModel wraps a transformer module and optionally a pooler module which controls how mapping from transformer output tensor to final representations. DenseModel class also handles loss computation during training. It implements a contrastive loss with in-batch negatives and can perform negative sharing across devices using parallel collective defined in the NCCL library.

To complete the dense retriever training setup, we introduce a DenseTrainer which implements miscellaneous training utilities. It controls basic setups such as batch size and the number of training epochs. When running on multiple GPUs, the trainer will properly set up distributed training and wrap models for gradient reduction. During training, it will asynchronously load training data to overlap computation and I/O operations. At each training step, the trainer turns a batch of loaded data into tensors and passes them to the model. In this way, the trainer glues the data sets and models together.

DenseTrainer is a subclass of Trainer in the transformers library. It inherits a collection of advanced utilities including mixed-precision training and optimizer state sharding. It is also possible to further subclass DenseTrainer to create unique training behaviors. Concretely in Tevatron, we implement a subclass GCTrainer which uses gradient caching to support large batch training on memory-limited devices (Luyu Gao and Callan, 2021).

By combining data processor, dense retrieval model, and trainer all together, Tevatron abstracts the training loop of dense retrieval model into the code block shown in Fig. 2.

# initialize model
model = DenseModel.build(model_args)
# initialize dataset
train_dataset = HFTrainDataset(data_args).process()
train_dataset = TrainDataset(data_args, train_dataset)
# initialize trainer
trainer = GCTrainer if training_args.grad_cache else DenseTrainer
trainer = trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=QPCollator(data_args),
)
# start training
trainer.train()

The retriever classes in Tevatron build dense retrieval index from text embeddings and execute search over the index. We use FAISS library (Johnson et al., 2019) as our retriever’s backend. It implements several efficient indices in C++ and exposes them through Python interfaces. For users who want the best performance, Tevatron provides a simple class BaseFaissIPRetriever which wraps a flat faiss.IndexFlatIP index for exact search. Those who want to trade-off between efficiency and effectiveness can use the more powerful FaissRetriever class. FaissRetriever takes an additional index_spec string argument in its initialization method and use faiss.index_factory method to flexibly build the specified index. Users can take advantage of this interface to build approximate search indices like HNSW (Malkov and Yashunin, 2020) or PQ (Jégou et al., 2011).

With Tevatron, we are able to replicate the training of DPR model for NQ dataset (see details in section 4.1) by a single command:

python -m tevatron.driver.train \
    --do_train \
    --dataset_name Tevatron/wikipedia-nq \
    --model_name_or_path bert-base-uncased \
    --per_device_train_batch_size 128 \
    --train_n_passages 2 \
    --num_train_epochs 40 \
    --learning_rate 1e-5 \
    --fp16 \
    --grad_cache \
    --output_dir model_nq

Besides training data, Tevatron also self-contains corresponding corpus data for each dataset. Again, we simplifies corpus encoding process into a single command:

python -m tevatron.driver.encode \
    --output_dir=temp \
    --model_name_or_path model_nq \
    --dataset_name Tevatron/wikipedia-nq-corpus \
    --encoded_save_path corpus_emb_00.pkl \
    --encode_num_shard 20 \
    --encode_shard_index 00 \
    --fp16

As encoding the entire corpus within a single process may cost large RAM usage and a long time, Tevatron support encoding the corpus by sharding. For example, the above command encodes the first 1/20 split of the entire corpus. Users can easily run multiple processes for multiple shards in parallel to speed up the encoding process.

The DPR work by Karpukhin et al. (Karpukhin et al., 2020) is among the first works that show text retrieval using learned dense representations outperforms traditional text retrieval using heuristic sparse representations (e.g. BM25) on open-domain question-answering tasks.

We evaluate the effectiveness of Tevatron by replicating the retrieval results on QA tasks (Kwiatkowski et al., 2019; Joshi et al., 2017; Rajpurkar et al., 2016; Voorhees and Tice, 2000; Berant et al., 2013) reported in original DPR work (Karpukhin et al., 2020). We compare the models trained under the "Single" setting defined in the original work where each model is trained by the corresponding individual dataset. Following the similar hyperparameters setting, we train the models with a learning rate of 1e-5 for 40 epochs with batch size 128. In Table 1, except having slightly lower accuracy than the numbers in DPR paper on SQuAD, Tevatron gives even a bit higher top-$k$ accuracy on all other four datasets. Overall, all top-$k$ accuracy results obtained via the Tevatron pipeline are at the same level of accuracy as original work. Therefore, we conclude that this is a successful replication, proving that the Tevatron pipeline is effective.

We demonstrate the efficiency of Tevatron by comparing it with the original DPR repo ³³3https://github.com/facebookresearch/DPR
To be clear, the efficiency results are based on the master branch on 2022-02-12 on three dimensions: RAM usage, GPU memory usage, and training time. The experiments are conducted on a machine with NVIDIA A100 GPUs. In both DPR-repo and Tevatron-default settings, we train the dense retriever model on 4 GPUs in distributed data-parallel mode of Pytorch. By comparing the first two rows in Table 2, we see that Tevatron is more efficient on all three dimensions than the original codebase. Concretely, Tevatron costs 3/4 less RAM, 12G less GPU memory, being 1/4 faster than training using DPR repo. This means, given the same resources, Tevatron has the potential to support larger training data, larger batch size, and faster training.

The gradient cache feature of Tevatron can further improve the GPU memory efficiency (Luyu Gao and Callan, 2021). DPR training requires batch size 128 to get the level of retrieval accuracy as reported above. With the original DPR repo, users cannot train a model with enough batch size if the GPU resource is limited, which will result in a drop in retrieval accuracy. Tevatron provides users the option to train dense retrievers using limited GPU resources but keeps the same amount of batch size for each optimization step. To illustrate this, we conduct experiments with Tevatron-GradCache on a single GPU. Tevatron-GradCache trains dense retrievers by splitting the batch of size 128 into sub-batches of size 32. Via conducting two round forward steps described in the gradient cache work (Luyu Gao and Callan, 2021), the model update step of Tevatron-GradCache is mathematically equivalent to Tevatron-default. In the experiment, it costs only 4G RAM and 15G GPU memory, to train a DPR model on NQ dataset with desired batch size. By reducing the sub-batch size, Tevatron-GradCache can save more GPU memory.

We also evaluated the performance of the training dense retriever model using the Jax backend of Tevatron on a V3-8 TPU VM. Back in the days when DPR (Karpukhin et al., 2020) first came out, it cost around a day to train dense retriever on NQ dataset using the initial DPR repo with 8$\times$ Nvidia V100-large GPUs⁴⁴4This is recorded by DPR authors in the GitHub page.. Now it is exciting to see that, such training can be done within one hour with Tevatron.

To further show the flexibility of the Tevatron toolkit across model architectures and accelerator platforms, we train multiple dense retrieval baselines on MS MARCO passage ranking task with different Transformer backbones from HuggingFace hub⁵⁵5https://huggingface.co/models. The experiments are conducted on both GPU and TPU platforms. The models are trained with a learning rate of 5e-6 with batch size 64 for 3 epochs using our self-contained dataset Tevatron/msmarco-passage.

In Table 3, we show MRR@10 for each model and training time on 4$\times$ A100 GPU and V3-8 TPU. The model backbones we choose varying across:

• •

  model size: row(1), row(2) and row(4) are models in BERT family with different size {distil, base, large}.

• •

  model type: row(4), row(5) are models in same level of size with different backbone structure {bert, roberta}.

• •

  model parameters: row(2), row(3) are same model backbone but the later one is further pre-trained from row(2), i.e. {original, fine-tuned}.

Finally, we evaluated two models trained with hard negative mining(Gao and Callan, 2021b) mined with Tevatron retriever. We craft the augmented training data by combining the hard negative passages mined using the first round dense retriever model with the original training dataset. Then we retrain the models using the hard negative augmented data. The last two rows in Table 3 show that by augmenting training data with hard negative, we can further improve the effectiveness of the dense retriever model. We also demonstrate Tevatron can replicate the state-of-the-art co-Condenser retriever(Gao and Callan, 2021b) on MS MARCO passage ranking.

The success of dense retrieval also drives research in multilingual retrieval (Asai et al., 2021; Zhang et al., 2021; Clark et al., 2020). We additionally show our Tevatron toolkit can generalize to multilingual retrieval tasks by replicating the dense retrieval baseline reported in the XOR-Retrieve task (Asai et al., 2021).