Efficient Long Context Language Model Retrieval with Compression

TL;DR

The paper tackles the computational bottleneck of long-context language model retrieval by introducing CoLoR, a retrieval-focused passage compression model trained with synthetic preference data and a dynamic length regularization term. CoLoR uses Odds Ratio Preference Optimization to rank compressed passages by retrieval efficacy while actively minimizing their length, enabling much smaller prompts for LCLMs during retrieval. Across 9 benchmark datasets, CoLoR achieves about a 6% gain in retrieval performance and reduces the input context by roughly 1.91x, outperforming both extractive/abstractive baselines and conventional retrieval methods. The approach generalizes across diverse domains and LLM backbones and offers a practical, scalable solution for efficient LCLM-based information retrieval in real-world settings.

Abstract

Long Context Language Models (LCLMs) have emerged as a new paradigm to perform Information Retrieval (IR), which enables the direct ingestion and retrieval of information by processing an entire corpus in their single context, showcasing the potential to surpass traditional sparse and dense retrieval methods. However, processing a large number of passages within in-context for retrieval is computationally expensive, and handling their representations during inference further exacerbates the processing time; thus, we aim to make LCLM retrieval more efficient and potentially more effective with passage compression. Specifically, we propose a new compression approach tailored for LCLM retrieval, which is trained to maximize the retrieval performance while minimizing the length of the compressed passages. To accomplish this, we generate the synthetic data, where compressed passages are automatically created and labeled as chosen or rejected according to their retrieval success for a given query, and we train the proposed Compression model for Long context Retrieval (CoLoR) with this data via preference optimization while adding the length regularization loss on top of it to enforce brevity. Through extensive experiments on 9 datasets, we show that CoLoR improves the retrieval performance by 6% while compressing the in-context size by a factor of 1.91. Our code is available at: https://github.com/going-doer/CoLoR.

Figures (5)

• Figure 1: Comparison of different IR approaches. (A) Dense Retrieval. To identify relevant documents to the given query, it first embeds them into the vector space and then calculates their semantic similarity. (B) LCLM Retrieval. The LCLM takes and processes the raw passages from the corpus along with the query in the input context, and identifies the relevant passages. (C) CoLoR. We compress the raw passages, and use the compressed passages alongside the query as the LCLM input for retrieval.
• Figure 2: Overview of Training Processes for CoLoR. 1. We first create the training data for CoLoR by generating multiple compressed passages from their original passages with multiple LMs. 2. The compressed passages and their associated query are used as input to the LCLM, and their retrieval performance is measured to label them as either chosen or rejected based on retrieval results. 3. CoLoR is trained using the pairs of chosen and rejected compressed passages obtained from previous steps.
• Figure 3: The trade-off of different methods, showing their compression rate (x-axis) and retrieval performance (y-axis).
• Figure 4: Results with varying the position of (compressed) passages associated with the query within the corpus, where 0% (on the x-axis) represents beginning.
• Figure 5: Results with varying the position of (compressed) passages for all datasets. Specifically, we arbitrarily adjust the positions of the gold and few-shot passages within the corpus relative to the query (0% represents the beginning). The figures at the top, middle, and bottom represent the results with CoLoR, raw passage, and document title, respectively.