MLKV: Teste Multi-Livello a Coppie Chiave-Valore per la Decodifica Efficiente in Memoria dei Trasformatori

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

June 13, 2024

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

cs.AI

Abstract

L'inferenza auto-regressiva dei transformer trae grande beneficio dalla memorizzazione delle chiavi-valori (KV), ma può portare a significativi colli di bottiglia nella memoria man mano che le dimensioni del modello, la dimensione del batch e la lunghezza delle sequenze crescono su larga scala. Introduciamo la condivisione multi-strato delle chiavi-valori (MLKV), un approccio innovativo che estende la condivisione delle KV attraverso i livelli del transformer per ridurre ulteriormente l'uso della memoria rispetto a quanto possibile con l'attenzione multi-query (MQA) e l'attenzione a query raggruppate (GQA). Le valutazioni su vari benchmark NLP e metriche di inferenza utilizzando varianti riaddestrate di Pythia-160M dimostrano che MLKV riduce significativamente l'uso della memoria con una minima perdita di prestazioni, riducendo la dimensione della cache KV fino a un fattore di 6x rispetto a MQA. Questi risultati evidenziano il potenziale di MLKV per un'implementazione efficiente dei modelli transformer su larga scala. Forniamo il codice all'indirizzo 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