IBM slate-30m-english-rtrvr model card

Task-specific pre-training

This stage uses the RetroMAE framework, to make our underlying LM more retrieval oriented. We initialize our base LM with slate.30m.english and continue with RetroMAE pre-training, using the data in Table 1. Our hyper-parameters are: learning rate: 2e-5, number of steps: 435000, GPUs: 8 A100 (80GB) GPUs. Note: this is our base LM for the following 2 stages.

Distillation using Unsupervised and Supervised Pairs

We use a bi-encoder framework for training an embedding model, as in Figure 1. We initialize with the RetroMAE pre-trained model, and further employ Knowledge Distillation with <query, passage> text pairs using a contrastive loss objective with in-batch negatives. Knowledge Distillation trains the student’s output probability distribution to match that of the teacher as closely as possible. In the context of retriever models, the output distribution is the similarity scores between pairs of text. Specifically, for each pair of sentences <query, passage>, the distribution of the teacher’s scores between the embeddings of query and passage, i.e., the cosine similarity between the embeddings, is distilled into the student.

The teacher used for distillation is the "slate.125m.english.rtrvr” model trained on the same data mentioned below. The teacher is formed using model fusion between two models, each of which used RetroMAE pretraining and finetuning, but differing in the finetuning data. For more details, please refer to the model card for slate.125m.english.rtrvr. The flow of knowledge transfer is shown pictorially in Figure 2.

Figure 2. Knowledge Distillation

We mine large scale pairs from various domain, as indicated in the Training Data section. Furthermore, we also include high-quality pairs for the retrieval task on the following datasets: SQuAD, Natural Questions, Specter, Stack Exchange (Title, Body) pairs, S2ORC, SearchQA, HotpotQA and Fever. Distillation hyperparameters are: learning rate:7e-4, number of steps: 400000, effective batch size: 2048, GPUs: 2 A100_80GB GPUs.

Baselines

For a fair comparison, we compare with the following baselines:

1. BM25 (a traditional model based on tf-idf).
2. ELSER (a commercial search algorithm provided by Elastic).
3. all-MiniLM-l6-v2: a popular open-source sentence transformers model. This model shares the same architecture as slate.125m.english.rtvr, with a smaller embedding dimension (384) and has been trained on more data without commercial-friendly licenses. See the Hugging Face model card (https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) for more details.
4. E5-base: a recent open-source transformer model with very good performance on the BEIR benchmark. This is a base-sized model, which has the same architecture as slate.125m.english.rtvr. [Reference: Wang et.al., 2022: Text Embeddings by Weakly-Supervised Contrastive Pre-training]. Huggingface model card (https://huggingface.co/intfloat/e5-base).
5. E5-small: a smaller model within the open source E5 family. The embedding dimension of this model matches that of all-minilm-l6-v2 (384), however it has 12 layers, and thus is larger and slightly slower. [Reference: Wang et.al., 2022: Text Embeddings by Weakly-Supervised Contrastive Pre-training]. Huggingface model card (https://huggingface.co/intfloat/e5-small).
6. BGE-base: a recent open-source transformer model with the best performance on the BEIR benchmark for the 768 embedding size (as of 01.20.2024). Huggingface model card (https://huggingface.co/BAAI/bge-base-en-v1.5).

Our Evaluation benchmark: BEIR (MTEB’s retrieval tab)

The BEIR benchmark contains of 15 open-source retrieval tasks evaluated under a zero-shot setting. BEIR focused on Diversity, including nine different retrieval tasks: Fact checking, citation prediction, duplicate question retrieval, argument retrieval, news retrieval, question answering, tweet retrieval, bio-medical IR, and entity retrieval. Further, it includes datasets from diverse text domains, datasets that cover broad topics (like Wikipedia) and specialized topics (like COVID-19 publications), different text types (news articles vs. Tweets), datasets of various sizes (3.6k - 15M documents), and datasets with different query lengths (average query length between 3 and 192 words) and document lengths (average document length between 11 and 635 words). BEIR uses the Normalized Cumulative Discount Gain (specifically, nDCG@10) metric for evaluation.

Long NQ

Long NQ is an IBM dataset designed for evaluating the full RAG pipeline, based on a subset of the NaturalQuestions dataset. The dev set has 300 answerable questions with a corpus of 178,891 passages from 2,345 Wikipedia documents. Long NQ also provides gold Wikipedia passages that are relevant for each question. During retrieval, the task is to obtain the relevant gold passage from the corpus for every question.

Results

Figure 3. Performance comparison on the BEIR benchmark (MTEB retrieval tab)

Figure 4. Performance comparison on the Long NQ dataset