LLoCO : Apprentissage de contextes longs hors ligne

LLoCO: Learning Long Contexts Offline

April 11, 2024

Auteurs: Sijun Tan, Xiuyu Li, Shishir Patil, Ziyang Wu, Tianjun Zhang, Kurt Keutzer, Joseph E. Gonzalez, Raluca Ada Popa

cs.AI

Résumé

Le traitement de contextes longs reste un défi pour les grands modèles de langage (LLMs) en raison de la surcharge quadratique en calcul et en mémoire du mécanisme d'auto-attention, ainsi que des tailles substantielles du cache KV pendant la génération. Nous proposons une nouvelle approche pour résoudre ce problème en apprenant les contextes hors ligne grâce à la compression de contexte et à un ajustement fin efficace en paramètres dans un domaine spécifique. Notre méthode permet à un LLM de créer une représentation concise du contexte original et de récupérer efficacement les informations pertinentes pour répondre avec précision aux questions. Nous introduisons LLoCO, une technique qui combine la compression de contexte, la récupération d'informations et l'ajustement fin efficace en paramètres en utilisant LoRA. Notre approche étend la fenêtre de contexte effective d'un modèle LLaMA2-7B de 4k tokens pour gérer jusqu'à 128k tokens. Nous évaluons notre approche sur plusieurs ensembles de données de question-réponse à contexte long, démontrant que LLoCO surpasse significativement l'apprentissage en contexte tout en utilisant 30 fois moins de tokens pendant l'inférence. LLoCO atteint une accélération jusqu'à 7,62 fois et réduit considérablement le coût de la réponse à des questions sur des documents longs, en faisant une solution prometteuse pour le traitement efficace de contextes longs. Notre code est disponible publiquement à l'adresse https://github.com/jeffreysijuntan/lloco.

English

Processing long contexts remains a challenge for large language models (LLMs) due to the quadratic computational and memory overhead of the self-attention mechanism and the substantial KV cache sizes during generation. We propose a novel approach to address this problem by learning contexts offline through context compression and in-domain parameter-efficient finetuning. Our method enables an LLM to create a concise representation of the original context and efficiently retrieve relevant information to answer questions accurately. We introduce LLoCO, a technique that combines context compression, retrieval, and parameter-efficient finetuning using LoRA. Our approach extends the effective context window of a 4k token LLaMA2-7B model to handle up to 128k tokens. We evaluate our approach on several long-context question-answering datasets, demonstrating that LLoCO significantly outperforms in-context learning while using 30times fewer tokens during inference. LLoCO achieves up to 7.62times speed-up and substantially reduces the cost of long document question answering, making it a promising solution for efficient long context processing. Our code is publicly available at https://github.com/jeffreysijuntan/lloco.