Oltre RAG: Compressione Task-Aware della KV Cache per il Ragionamento su Conoscenza Completa

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

March 6, 2025

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

cs.AI

Abstract

L'integrazione di conoscenze esterne nei grandi modelli linguistici (LLM) ne migliora l'utilità in diverse applicazioni, ma i metodi esistenti presentano compromessi. La Generazione Aumentata da Recupero (RAG) recupera prove tramite ricerca di similarità, ma informazioni chiave potrebbero non rientrare nei risultati più rilevanti. I modelli a contesto lungo possono elaborare più documenti, ma sono computazionalmente costosi e limitati dalla dimensione della finestra di contesto. Ispirati dagli studenti che condensano il materiale di studio per esami a libro aperto, proponiamo una compressione task-aware della cache chiave-valore (KV), che comprime le conoscenze esterne in un'impostazione zero-shot o few-shot. Ciò consente ai LLM di ragionare in modo efficiente su una rappresentazione compatta di tutte le informazioni rilevanti. Gli esperimenti dimostrano che il nostro approccio supera sia la RAG che i metodi di compressione task-agnostici. Su LongBench v2, migliora l'accuratezza fino a 7 punti assoluti rispetto alla RAG con un tasso di compressione di 30x, riducendo al contempo la latenza di inferenza da 0,43s a 0,16s. Un dataset sintetico evidenzia che la RAG funziona bene quando prove sparse sono sufficienti, mentre la compressione task-aware è superiore per compiti che richiedono una conoscenza ampia.

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.