SubGen: Generación de tokens en tiempo y memoria sublineales

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

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

cs.AI

Resumen

A pesar del notable éxito de los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés), sus extensos requisitos de memoria plantean desafíos para su implementación en la generación de tokens de contexto largo. La considerable huella de memoria de los decodificadores de LLMs surge de la necesidad de almacenar todos los tokens anteriores en el módulo de atención, un requisito impuesto por el almacenamiento en caché de pares clave-valor (KV). En este trabajo, nuestro enfoque se centra en desarrollar una técnica de compresión eficiente para la caché KV. La evidencia empírica indica una tendencia significativa de agrupamiento dentro de los embeddings de clave en el módulo de atención. Basándonos en esta observación clave, hemos diseñado un nuevo método de almacenamiento en caché con complejidad sublineal, empleando agrupamiento en línea en los tokens clave y muestreo en línea de ell_2 en los valores. El resultado es un algoritmo de decodificación de atención precisa y eficiente, denominado SubGen. Este algoritmo no solo garantiza una huella de memoria sublineal y una complejidad temporal sublineal, sino que también establecemos un límite de error ajustado para nuestro enfoque. Las evaluaciones empíricas en tareas de respuesta a preguntas de contexto largo demuestran que SubGen supera significativamente a los métodos existentes y de vanguardia de compresión de caché KV en términos de rendimiento y eficiencia.

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.