Exploring Fine-Tuning for In-Context Retrieval and Efficient KV-Caching in Long-Context Language Models

TL;DR

This work investigates whether fine-tuning Long-Context Language Models (LCLMs) with Group Relative Policy Optimization (GRPO) can enable in-context retrieval that rivals or surpass retrieval-augmented generation (RAG). It introduces a data- and reward-driven RL framework to train LCLMs to selectively attend to relevant information within extremely long inputs and to be robust to KV-cache compression. Across in-domain benchmarks, certain GRPO-based rewards yield substantial improvements (up to ~+20 points) and even surpass RAG in some cases, though out-of-domain generalization is variable. Under KV-cache compression, gains are modest and task-dependent, with improvements not primarily explained by enhanced document ranking, suggesting robustness to noise rather than better retrieval signals. Overall, LCLMs with GRPO can replace RAG in certain in-domain settings, but achieving universal generalization and robust efficiency under compression requires further advances in training objectives and evaluation across diverse tasks and domains.

Abstract

With context windows of millions of tokens, Long-Context Language Models (LCLMs) can encode entire document collections, offering a strong alternative to conventional retrieval-augmented generation (RAG). However, it remains unclear whether fine-tuning strategies can improve long-context performance and translate to greater robustness under KV-cache compression techniques. In this work, we investigate which training strategies most effectively enhance LCLMs' ability to identify and use relevant information, as well as enhancing their robustness under KV-cache compression. Our experiments show substantial in-domain improvements, achieving gains of up to +20 points over the base model. However, out-of-domain generalization remains task dependent with large variance -- LCLMs excels on finance questions (+9 points), while RAG shows stronger performance on multiple-choice questions (+6 points) over the baseline models. Finally, we show that our fine-tuning approaches bring moderate improvements in robustness under KV-cache compression, with gains varying across tasks.