Nuovo Meccanismo di Attenzione Step Riduce la Cache KV

Introduzione al Problema della Cache KV

L'uso crescente dei modelli linguistici di grandi dimensioni (LLM) e l'emergere di nuovi paradigmi nell'inferenza hanno portato in primo piano la sfida dell'inferenza efficiente su larga scala. Un collo di bottiglia significativo è la cache Key-Value (KV) all'interno dei meccanismi di attenzione tradizionali. Questa cache si espande linearmente con la dimensione del batch e la lunghezza della sequenza, diventando un vero e proprio "divoratore di memoria" che ostacola la scalabilità e l'espansione dei LLM.

Mentre varianti come MQA, GQA e MLA sono emerse per affrontare questo problema, spesso faticano a mantenere le prestazioni in condizioni di memoria ristrette o introducono complessità che pongono sfide ingegneristiche e problemi di compatibilità.

Multi-matrix Factorization Attention (MFA)

Un recente studio di ricercatori di Stepes, Tsinghua University e altre istituzioni introduce una nuova architettura di meccanismo di attenzione: Multi-matrix Factorization Attention (MFA) e la sua variante MFA-Key-Reuse (MFA-KR). Questo meccanismo riduce significativamente il costo dell'inferenza del modello linguistico, migliorando contemporaneamente le prestazioni.

MFA e MFA-KR non solo superano MLA in termini di prestazioni, ma corrispondono anche alle prestazioni tradizionali di MHA riducendo l'uso della cache KV fino al 93.7%.

MFA è progettato per la semplicità, la facile riproduzione, la bassa sensibilità agli iperparametri e la compatibilità con vari metodi di Pos-embedding.

Approccio e Analisi di MFA

Il team di ricerca ha analizzato il design generale e la capacità dei meccanismi di attenzione, identificando due dimensioni critiche relative alla capacità. Questa analisi ha portato allo sviluppo di nuovi metodi analitici e principi di progettazione.

Hanno introdotto il concetto di Generalized Multi-Head Attention (GMHA) come framework unificante per la comprensione delle diverse varianti di MHA.

Il team ha anche esplorato il calcolo e l'archiviazione di key-value da una prospettiva di inferenza ed esaminato la capacità del modello da una prospettiva di decomposizione.

Fully Parameterized Bilinear Attention (FPBA) è stata stabilita come limite superiore teorico delle prestazioni. Hanno scoperto che MHA e le sue varianti sono decomposizioni di basso rango di FPBA.

Confronto con MQA e MLA

L'analisi si è concentrata su due schemi di miglioramento rappresentativi: Multi-Query Attention (MQA) e Multi-Head Latent Attention (MLA).

• MQA utilizza una strategia di condivisione dei parametri più aggressiva, in cui tutte le teste di attenzione condividono lo stesso set di parametri key-value. Ciò riduce l'uso della memoria, ma può influire sull'espressività del modello.
• MLA introduce uno spazio latente condiviso per la compressione dei parametri, ma l'effettiva potenza espressiva è limitata dalla dimensione più piccola, il che significa che l'aumento delle dimensioni intermedie non migliora significativamente le prestazioni.

Innovazioni Chiave di MFA

Lo sviluppo di MFA è stato guidato dall'obiettivo di creare un meccanismo di attenzione che riducesse al minimo il consumo di risorse avvicinandosi ai limiti teorici di prestazione. Il design di MFA incorpora tre innovazioni chiave:

1. Aumento significativo del numero e della dimensione delle teste di attenzione per massimizzare la capacità del modello.
2. Impiego di una strategia di decomposizione a basso rango aggressiva per mantenere l'efficienza dei parametri espandendo il numero e le dimensioni delle teste di attenzione.
3. Utilizzo di un design a testa key-value singola per mantenere minimo il consumo di memoria, anche con una maggiore complessità del modello.

Misurazione e Confronto della Capacità

Per analizzare ulteriormente MFA e altri meccanismi di attenzione, il team ha introdotto due metriche chiave:

• Total Effective Rank (TER): Il prodotto del numero di teste di attenzione e del rango di fattorizzazione per testa (FRH).
• Shared Latent Subspace Dimension (SLSD): La dimensione dello spazio nascosto condiviso da tutte le teste di attenzione.

MFA raggiunge un SLSD e un TER più elevati rispetto a MQA.

Rispetto a MLA, MFA ottiene una dimensione della cache KV inferiore e un TER più elevato con budget di parametri simili, mantenendo un SLSD comparabile.

Rispetto al tradizionale MHA, MFA ha un TER più elevato, anche se il suo SLSD è inferiore.

Risultati Sperimentali

Sono stati condotti numerosi esperimenti per valutare le prestazioni della nuova architettura su scale più grandi, testando modelli che vanno da 1B a 7B di parametri e dati di addestramento da 10B a 1T.

MFA ha dimostrato capacità di scaling comparabili al tradizionale MHA, mantenendo prestazioni eccellenti anche su scale più grandi.

Mentre MFA-KR ha mostrato prestazioni leggermente inferiori, la sua tendenza di scaling si è allineata con MHA. I vantaggi in termini di risparmio di memoria di MFA e MFA-KR hanno continuato ad espandersi con le dimensioni del modello, con MFA che ha ottenuto un risparmio di memoria dell'87.5% e MFA-KR che ha ridotto l'uso di memoria al 6.25% alla scala più grande.

Studi di Ablazione

Gli studi di ablazione hanno convalidato l'efficacia di MFA e MFA-KR. I loro vantaggi in termini di prestazioni sono stati confermati anche attraverso vari metodi di codifica posizionale mainstream.

Prospettive

MFA offre miglioramenti significativi con un design semplice, affrontando efficacemente il collo di bottiglia della memoria nell'inferenza LLM senza aggiungere ulteriore complessità ingegneristica. Si integra perfettamente nell'ecosistema Transformer esistente, accelerando l'applicazione dei LLM in vari scenari.