Au-delà de RAG : Compression de cache KV consciente des tâches pour un raisonnement complet sur les connaissances

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

March 6, 2025

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

cs.AI

Résumé

L'intégration de connaissances externes dans les grands modèles de langage (LLM) améliore leur utilité dans diverses applications, mais les méthodes existantes présentent des compromis. La Génération Augmentée par Récupération (RAG) extrait des preuves via une recherche de similarité, mais des informations clés peuvent se trouver en dehors des résultats les mieux classés. Les modèles à contexte long peuvent traiter plusieurs documents, mais ils sont coûteux en calcul et limités par la taille de la fenêtre contextuelle. Inspirés par les étudiants qui condensent leurs supports d'étude pour les examens à livre ouvert, nous proposons une compression de cache clé-valeur (KV) adaptée à la tâche, qui compresse les connaissances externes dans un cadre zero-shot ou few-shot. Cela permet aux LLM de raisonner efficacement sur une représentation compacte de toutes les informations pertinentes. Les expériences montrent que notre approche surpasse à la fois la RAG et les méthodes de compression indépendantes de la tâche. Sur LongBench v2, elle améliore la précision jusqu'à 7 points absolus par rapport à la RAG avec un taux de compression de 30x, tout en réduisant la latence d'inférence de 0,43s à 0,16s. Un ensemble de données synthétique met en évidence que la RAG fonctionne bien lorsque des preuves éparses suffisent, tandis que la compression adaptée à la tâche est supérieure pour les tâches nécessitant une connaissance étendue.

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.