Além do RAG: Compressão de Cache KV Consciente da Tarefa para Raciocínio de Conhecimento Abrangente

Beyond RAG: Task-Aware KV Cache Compression for Comprehensive Knowledge Reasoning

March 6, 2025

Autores: Giulio Corallo, Orion Weller, Fabio Petroni, Paolo Papotti

cs.AI

Resumo

A incorporação de conhecimento externo em modelos de linguagem de grande escala (LLMs) amplia sua utilidade em diversas aplicações, mas os métodos existentes apresentam compensações. A Geração Aumentada por Recuperação (RAG) busca evidências por meio de pesquisa de similaridade, mas informações-chave podem ficar fora dos resultados mais bem classificados. Modelos de contexto longo podem processar múltiplos documentos, mas são computacionalmente caros e limitados pelo tamanho da janela de contexto. Inspirados por estudantes que condensam material de estudo para provas de livro aberto, propomos a compressão de cache chave-valor (KV) consciente da tarefa, que comprime conhecimento externo em uma configuração zero-shot ou few-shot. Isso permite que os LLMs raciocinem de forma eficiente sobre uma representação compactada de todas as informações relevantes. Experimentos mostram que nossa abordagem supera tanto a RAG quanto métodos de compressão agnósticos à tarefa. No LongBench v2, ela melhora a precisão em até 7 pontos absolutos em relação à RAG com uma taxa de compressão de 30x, enquanto reduz a latência de inferência de 0,43s para 0,16s. Um conjunto de dados sintético destaca que a RAG desempenha bem quando evidências esparsas são suficientes, enquanto a compressão consciente da tarefa é superior para tarefas de conhecimento amplo.

English

Incorporating external knowledge in large language models (LLMs) enhances their utility across diverse applications, but existing methods have trade-offs. Retrieval-Augmented Generation (RAG) fetches evidence via similarity search, but key information may fall outside top ranked results. Long-context models can process multiple documents but are computationally expensive and limited by context window size. Inspired by students condensing study material for open-book exams, we propose task-aware key-value (KV) cache compression, which compresses external knowledge in a zero- or few-shot setup. This enables LLMs to reason efficiently over a compacted representation of all relevant information. Experiments show our approach outperforms both RAG and task-agnostic compression methods. On LongBench v2, it improves accuracy by up to 7 absolute points over RAG with a 30x compression rate, while reducing inference latency from 0.43s to 0.16s. A synthetic dataset highlights that RAG performs well when sparse evidence suffices, whereas task-aware compression is superior for broad knowledge tasks.