LoopVQ: quantizzazione della cache KV a meno di 1 bit per i transformer iterativi

Abstract degli autori

I transformer iterativi sono un’alternativa efficiente in termini di parametri ai transformer standard per la modellazione del linguaggio. Riapplicando gli stessi layer in più passaggi del ciclo, i recenti modelli linguistici iterativi raggiungono prestazioni pari a quelle di modelli non iterativi di dimensioni diverse volte superiori. Sebbene questa architettura riduca efficacemente la memoria occupata dai parametri, ogni passaggio del ciclo memorizza comunque le proprie chiavi e i propri valori, perciò la cache key-value (KV) cresce linearmente con il numero di passaggi e diventa il principale collo di bottiglia in termini di memoria. Osserviamo che la cache KV dei modelli iterativi è altamente ridondante lungo l’asse del ciclo e che i diversi passaggi presentano sensibilità alla quantizzazione molto diverse. Sulla base di queste osservazioni, proponiamo $\textbf{LoopVQ}$, un metodo di quantizzazione vettoriale della cache KV per transformer iterativi, che codifica congiuntamente i canali di più passaggi del ciclo e mette a punto i codebook in modo che corrispondano all’output del modello a precisione piena. A 0,5 bit per elemento, LoopVQ supera i metodi esistenti di quantizzazione vettoriale della cache KV per gli LLM standard nei compiti generali, di ragionamento a catena di pensiero su sequenze lunghe e di recupero con contesto lungo, con i modelli Ouro e Huginn. Il suo design è ortogonale a quello di questi metodi e, in combinazione con TaSQ, conserva gran parte della precisione BF16 anche a 0,25 bit per elemento. Abbiamo inoltre implementato LoopVQ in SGLang, dove riduce quasi della metà la latenza di decodifica con un contesto di 32K token e aumenta il throughput massimo di serving fino a $4.3\times$ e $6.1\times$ rispettivamente sui modelli Ouro da 1.4B e 2.6B, rispetto alla baseline BF16.