PN-Mix: quantizzazione a precisione mista della cache delle key guidata dalla norma delle coppie per un’inferenza efficiente degli LLM con output lunghi

Abstract degli autori

L’inferenza con output lunghi nei modelli linguistici di grandi dimensioni è fortemente limitata dalla memoria e dalla larghezza di banda richieste dalla cache KV. La quantizzazione della cache KV attenua il problema, ma quella a pochi bit, per esempio a 2 bit, riduce drasticamente l’accuratezza del modello. I recenti metodi di quantizzazione a precisione mista della cache delle key limitano questa perdita assegnando una precisione maggiore ai canali delle key più importanti, ma ne stimano la rilevanza usando soltanto le statistiche delle key. Attraverso un’analisi delle perturbazioni dei logit di attenzione a livello delle coppie di canali RoPE, mostriamo che gli errori di quantizzazione nelle coppie di canali con prodotti elevati delle norme delle coppie di query e key contribuiscono in misura significativa alla variazione dei logit. Su questa base proponiamo PN-Mix, un metodo di quantizzazione a precisione mista della cache delle key per LLM basati su RoPE e costituiti solo da decoder. Il metodo assegna una precisione maggiore alle coppie di canali più rilevanti, individuate una sola volta nella fase di prefill e poi riutilizzate per tutta la decodifica. Nei benchmark standard di ragionamento, generazione di codice e AIME, PN-Mix, con una precisione media di 2,26 bit, raggiunge un’accuratezza media pari o superiore a quella dei precedenti metodi a precisione mista che utilizzano da 2,32 a 2,39 bit. Inoltre, nella nostra esplorazione di Pareto, le prestazioni di PN-Mix calano gradualmente fino a 1,5 bit, mentre i metodi di riferimento con cui lo confrontiamo mostrano cali più bruschi sotto i 2,25 bit.