Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT
This addresses retrieval challenges for domains with long documents where chunking is ineffective, though it appears incremental as it builds on existing architectures like Monarch Mixer.
The paper tackles the problem of poor retrieval performance with long documents by introducing LoCoV1, a 12-task benchmark for evaluating long-context retrieval, and M2-BERT, an 80M parameter encoder that scales to 32K tokens. The result shows M2-BERT outperforms Transformer-based models by at least 23.3 points on LoCoV1 with 90x fewer parameters.
Retrieval pipelines-an integral component of many machine learning systems-perform poorly in domains where documents are long (e.g., 10K tokens or more) and where identifying the relevant document requires synthesizing information across the entire text. Developing long-context retrieval encoders suitable for these domains raises three challenges: (1) how to evaluate long-context retrieval performance, (2) how to pretrain a base language model to represent both short contexts (corresponding to queries) and long contexts (corresponding to documents), and (3) how to fine-tune this model for retrieval under the batch size limitations imposed by GPU memory constraints. To address these challenges, we first introduce LoCoV1, a novel 12 task benchmark constructed to measure long-context retrieval where chunking is not possible or not effective. We next present the M2-BERT retrieval encoder, an 80M parameter state-space encoder model built from the Monarch Mixer architecture, capable of scaling to documents up to 32K tokens long. We describe a pretraining data mixture which allows this encoder to process both short and long context sequences, and a finetuning approach that adapts this base model to retrieval with only single-sample batches. Finally, we validate the M2-BERT retrieval encoder on LoCoV1, finding that it outperforms competitive Transformer-based models by at least 23.3 points, despite containing upwards of 90x fewer parameters.