SubGen: Генерация токенов за сублинейное время и с использованием сублинейной памяти

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

Авторы: Amir Zandieh, Insu Han, Vahab Mirrokni, Amin Karbasi

cs.AI

Аннотация

Несмотря на значительные успехи крупных языковых моделей (LLM), их высокие требования к памяти создают трудности при их развертывании для генерации токенов в длинных контекстах. Значительный объем памяти, занимаемый декодерами LLM, обусловлен необходимостью хранения всех предыдущих токенов в модуле внимания, что является требованием, накладываемым кэшированием ключей и значений (KV). В данной работе мы сосредоточены на разработке эффективного метода сжатия KV-кэша. Эмпирические данные указывают на значительную тенденцию к кластеризации в ключевых эмбеддингах модуля внимания. Опираясь на это ключевое наблюдение, мы разработали новый метод кэширования с сублинейной сложностью, использующий онлайн-кластеризацию для ключевых токенов и онлайн-выборку по норме ell_2 для значений. В результате был создан алгоритм декодирования внимания с доказанной точностью и эффективностью, названный SubGen. Этот алгоритм не только обеспечивает сублинейный объем памяти и сублинейную временную сложность, но и устанавливает строгую границу ошибки для нашего подхода. Эмпирические оценки на задачах ответов на вопросы в длинных контекстах демонстрируют, что SubGen значительно превосходит существующие и передовые методы сжатия KV-кэша по производительности и эффективности.

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.