Сокращение размера кеша ключ-значение трансформера с помощью внутрислойного внимания.
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
May 21, 2024
Авторы: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly
cs.AI
Аннотация
Кэш с ключами и значениями (KV) играет важную роль в ускорении декодирования для авторегрессионных моделей больших языков на основе трансформера (LLM). Однако объем памяти, необходимый для хранения кэша KV, может стать запретным при больших длинах последовательностей и больших размерах пакетов. С момента изобретения трансформера были обнаружены два из наиболее эффективных вмешательства для уменьшения размера кэша KV: Множественное Внимание Запросов (MQA) и его обобщение, Групповое Внимание Запросов (GQA). MQA и GQA изменяют конструкцию блока внимания так, что несколько головок запросов могут разделять одну головку ключа/значения, уменьшая количество различных головок ключа/значения на большой коэффициент, сохраняя при этом минимальное снижение точности. В данной статье мы показываем, что возможно пойти дальше в Множественном Внимании Запросов, также разделяя головки ключа и значения между смежными слоями, что приводит к новой конструкции внимания, которую мы называем Перекрестным Вниманием Слоев (CLA). С CLA мы обнаружили, что можно уменьшить размер кэша KV еще в 2 раза, сохраняя практически такую же точность, как у немодифицированного MQA. В экспериментах по обучению моделей с 1B и 3B параметрами с нуля мы демонстрируем, что CLA обеспечивает улучшение Парето в обмене память/точность, которое возможно с традиционным MQA, позволяя выводить более длинные последовательности и использовать большие размеры пакетов, чем это было бы возможно в противном случае.
English
Key-value (KV) caching plays an essential role in accelerating decoding for
transformer-based autoregressive large language models (LLMs). However, the
amount of memory required to store the KV cache can become prohibitive at long
sequence lengths and large batch sizes. Since the invention of the transformer,
two of the most effective interventions discovered for reducing the size of the
KV cache have been Multi-Query Attention (MQA) and its generalization,
Grouped-Query Attention (GQA). MQA and GQA both modify the design of the
attention block so that multiple query heads can share a single key/value head,
reducing the number of distinct key/value heads by a large factor while only
minimally degrading accuracy. In this paper, we show that it is possible to
take Multi-Query Attention a step further by also sharing key and value heads
between adjacent layers, yielding a new attention design we call Cross-Layer
Attention (CLA). With CLA, we find that it is possible to reduce the size of
the KV cache by another 2x while maintaining nearly the same accuracy as
unmodified MQA. In experiments training 1B- and 3B-parameter models from
scratch, we demonstrate that CLA provides a Pareto improvement over the
memory/accuracy tradeoffs which are possible with traditional MQA, enabling
inference with longer sequence lengths and larger batch sizes than would
otherwise be possible