LLoCO: Aprendizaje de Contextos Largos sin Conexión

LLoCO: Learning Long Contexts Offline

April 11, 2024

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

cs.AI

Resumen

El procesamiento de contextos largos sigue siendo un desafío para los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) debido al sobrecosto computacional y de memoria cuadrático del mecanismo de autoatención y a los tamaños sustanciales de la caché KV durante la generación. Proponemos un enfoque novedoso para abordar este problema mediante el aprendizaje de contextos fuera de línea a través de la compresión de contexto y el ajuste fino eficiente en parámetros dentro del dominio. Nuestro método permite que un LLM cree una representación concisa del contexto original y recupere de manera eficiente la información relevante para responder preguntas con precisión. Introducimos LLoCO, una técnica que combina compresión de contexto, recuperación y ajuste fino eficiente en parámetros utilizando LoRA. Nuestro enfoque extiende la ventana de contexto efectiva de un modelo LLaMA2-7B de 4k tokens para manejar hasta 128k tokens. Evaluamos nuestro enfoque en varios conjuntos de datos de preguntas y respuestas de contexto largo, demostrando que LLoCO supera significativamente al aprendizaje en contexto mientras utiliza 30 veces menos tokens durante la inferencia. LLoCO logra una aceleración de hasta 7.62 veces y reduce sustancialmente el costo de responder preguntas sobre documentos largos, convirtiéndolo en una solución prometedora para el procesamiento eficiente de contextos largos. Nuestro código está disponible públicamente en 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.