PACE: codifica periodica e allineata ai canali degli stati ricorrenti a quattro bit

Abstract degli autori

I modelli linguistici ad attenzione lineare conservano il contesto in uno stato ricorrente di dimensioni fisse, di cui un sistema di inferenza memorizza una copia per ogni sequenza che decodifica contemporaneamente. Memorizzare lo stato con una precisione inferiore consente quindi a più sequenze di condividere un acceleratore. I quantizzatori di stato esistenti arrotondano lo stato dopo ogni token generato e riferiscono che a quattro bit si perde precisione, a meno di aumentare il numero di bit o sottoporre il modello a fine-tuning.\par Presentiamo PACE, un quantizzatore di stato a quattro bit applicabile a posteriori. PACE arrotonda lo stato solo una volta ogni pochi token generati, come i kernel suddivisi in blocchi fanno già una volta per blocco durante il prefill, e conserva i token generati nel frattempo in un piccolo buffer. Arrotondare meno spesso riduce l’errore che si accumula nello stato. Inoltre, osserviamo che i valori più elevati dello stato gated-delta si concentrano in colonne, perché alcuni canali dell’aggiornamento scritto a ogni token restano elevati per molti token. PACE usa quindi una stessa scala di quantizzazione per ogni colonna. Mantiene inoltre a otto bit le teste che conservano più a lungo l’errore di arrotondamento. Su quattro modelli ibridi da 4B a 35B parametri, PACE con teste miste a quattro e otto bit ottiene risultati competitivi rispetto a DAMP a otto bit sui tre modelli più grandi, usando circa due terzi della memoria richiesta da DAMP; PACE con tutte le teste a quattro bit ottiene punteggi superiori a Q-Mamba a parità di memoria; e l’arrotondamento ogni pochi token preserva l’accelerazione della decodifica offerta dal buffer.