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

TL;DR

This paper investigates fine-tuning strategies to improve long-context language models' ability to retrieve relevant information and maintain robustness under KV-cache compression, showing significant in-domain gains but variable out-of-domain performance.

Contribution

The study systematically evaluates fine-tuning methods for enhancing long-context models' retrieval accuracy and robustness, highlighting their effectiveness and limitations across different tasks.

Findings

In-domain performance improved by up to +20 points.

Out-of-domain results vary significantly by task.

Moderate robustness gains under KV-cache compression.

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.