MLKV: Cabeças de Chave-Valor Multicamadas para Decodificação Eficiente em Memória de Transformers

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

June 13, 2024

Autores: Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

cs.AI

Resumo

A inferência auto-regressiva de transformers se beneficia significativamente do cache de Chave-Valor (KV), mas pode levar a grandes gargalos de memória à medida que o tamanho do modelo, o tamanho do lote e o comprimento da sequência aumentam em escala. Introduzimos o compartilhamento de Chave-Valor Multi-Camadas (MLKV), uma abordagem inovadora que estende o compartilhamento de KV através das camadas do transformer para reduzir o uso de memória além do que era possível com a Atenção Multi-Consulta (MQA) e a Atenção de Consulta Agrupada (GQA). Avaliações em vários benchmarks de PLN e métricas de inferência usando variantes do Pythia-160M retreinadas demonstram que o MLKV reduz significativamente o uso de memória com perda mínima de desempenho, diminuindo o tamanho do cache KV em até 6 vezes em comparação com o MQA. Esses resultados destacam o potencial do MLKV para a implantação eficiente de modelos transformers em escala. Disponibilizamos o código em https://github.com/zaydzuhri/pythia-mlkv.

English

Auto-regressive inference of transformers benefit greatly from Key-Value (KV) caching, but can lead to major memory bottlenecks as model size, batch size, and sequence length grow at scale. We introduce Multi-Layer Key-Value (MLKV) sharing, a novel approach extending KV sharing across transformer layers to reduce memory usage beyond what was possible with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Evaluations on various NLP benchmarks and inference metrics using uptrained Pythia-160M variants demonstrate that MLKV significantly reduces memory usage with minimal performance loss, reducing KV cache size down to a factor of 6x compared to MQA. These results highlight MLKV's potential for efficient deployment of transformer models at scale. We provide code at https://github.com/zaydzuhri/pythia-mlkv