LLoCO: Aprendizado de Contextos Longos Offline

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

Resumo

O processamento de contextos longos continua sendo um desafio para modelos de linguagem de grande escala (LLMs, na sigla em inglês) devido ao custo computacional e de memória quadrático do mecanismo de autoatenção e aos tamanhos substanciais do cache de chaves-valores (KV) durante a geração. Propomos uma abordagem inovadora para resolver esse problema, aprendendo contextos offline por meio de compressão de contexto e ajuste fino eficiente em parâmetros dentro de um domínio específico. Nosso método permite que um LLM crie uma representação concisa do contexto original e recupere informações relevantes de forma eficiente para responder perguntas com precisão. Introduzimos o LLoCO, uma técnica que combina compressão de contexto, recuperação e ajuste fino eficiente em parâmetros usando LoRA. Nossa abordagem estende a janela de contexto efetiva de um modelo LLaMA2-7B de 4 mil tokens para lidar com até 128 mil tokens. Avaliamos nossa abordagem em vários conjuntos de dados de perguntas e respostas de contexto longo, demonstrando que o LLoCO supera significativamente o aprendizado em contexto (in-context learning) enquanto utiliza 30 vezes menos tokens durante a inferência. O LLoCO alcança uma aceleração de até 7,62 vezes e reduz substancialmente o custo de perguntas e respostas em documentos longos, tornando-o uma solução promissora para o processamento eficiente de contextos longos. Nosso código está disponível publicamente em 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.