SubGen: Geração de Tokens em Tempo e Memória Sublinares

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

Autores: Amir Zandieh, Insu Han, Vahab Mirrokni, Amin Karbasi

cs.AI

Resumo

Apesar do sucesso significativo dos grandes modelos de linguagem (LLMs), seus extensos requisitos de memória apresentam desafios para sua implantação na geração de tokens de contexto longo. A pegada de memória substancial dos decodificadores de LLMs surge da necessidade de armazenar todos os tokens anteriores no módulo de atenção, uma exigência imposta pelo cache de chave-valor (KV). Neste trabalho, nosso foco está no desenvolvimento de uma técnica de compressão eficiente para o cache KV. Evidências empíricas indicam uma tendência significativa de agrupamento nos embeddings de chave no módulo de atenção. Com base nessa percepção fundamental, desenvolvemos um novo método de cache com complexidade sublinear, empregando agrupamento online em tokens de chave e amostragem online ell_2 em valores. O resultado é um algoritmo de decodificação de atenção comprovadamente preciso e eficiente, denominado SubGen. Este algoritmo não apenas garante uma pegada de memória sublinear e uma complexidade de tempo sublinear, mas também estabelecemos um limite de erro rigoroso para nossa abordagem. Avaliações empíricas em tarefas de resposta a perguntas de contexto longo demonstram que o SubGen supera significativamente os métodos existentes e de última geração de compressão de cache KV em termos de desempenho e eficiência.

English

Despite the significant success of large language models (LLMs), their extensive memory requirements pose challenges for deploying them in long-context token generation. The substantial memory footprint of LLM decoders arises from the necessity to store all previous tokens in the attention module, a requirement imposed by key-value (KV) caching. In this work, our focus is on developing an efficient compression technique for the KV cache. Empirical evidence indicates a significant clustering tendency within key embeddings in the attention module. Building on this key insight, we have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online ell_2 sampling on values. The result is a provably accurate and efficient attention decoding algorithm, termed SubGen. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach. Empirical evaluations on long-context question-answering tasks demonstrate that SubGen significantly outperforms existing and state-of-the-art KV cache compression methods in terms of performance and efficiency.