Più valore per ogni chiave: attenzione sparsa asimmetrica per una decodifica più rapida degli LLM

Abstract degli autori

La generazione autoregressiva nei modelli linguistici di grandi dimensioni (LLM) è limitata dai requisiti di memoria e di calcolo dei meccanismi di attenzione. I metodi di attenzione sparsa riducono questo costo selezionando solo le voci ad alta probabilità della matrice di attenzione. Osserviamo che in molti di questi metodi il prodotto tra probabilità e valori diventa trascurabile, spostando il collo di bottiglia sulla fase query-key. È quindi possibile ridurre le teste key per accelerare l’inferenza, mantenendo un maggior numero di teste value per preservare la capacità con un costo aggiuntivo di decodifica limitato. Introduciamo Sparse Asymmetric Group-Query Attention (SAGA), che disaccoppia il numero di teste key e value per sfruttare questo principio, e la abbiniamo all’attenzione approximate top-N (Atop-N), un semplice metodo di attenzione sparsa progettato per studiare l’interazione tra sparsità e asimmetria nel numero di teste. Formalizziamo teoricamente i vantaggi di questa asimmetria e li convalidiamo empiricamente con misurazioni della latenza e valutazioni della qualità su modelli fino a 1,5 miliardi di parametri. Insieme, SAGA e Atop-N raggiungono accelerazioni della decodifica end-to-end superiori a $2\times$ rispetto alla nostra baseline GQA con attenzione completa, nei contesti lunghi. Sui benchmark valutati, i modelli addestrati da zero con SAGA raggiungono una qualità quasi pari a quella di varianti GQA comparabili. Per facilitarne l’adozione, introduciamo un metodo efficiente di fine-tuning che converte i modelli preaddestrati nell’architettura SAGA, consentendo agli operatori di beneficiare del nostro approccio senza costosi riaddestramenti.