Riparametrizzazione della softmax per la quantizzazione della testa di output

Abstract degli autori

Nei modelli linguistici di piccole dimensioni, i vocabolari ampi rendono le teste di output una voce di costo considerevole durante l’inferenza. Proponiamo una riparametrizzazione della softmax: un metodo applicato dopo l’addestramento che seleziona una testa di output funzionalmente equivalente prima della quantizzazione. Il metodo sottrae da ogni riga di output un multiplo scalare della media delle righe del vocabolario e seleziona il coefficiente in base alla KL di validazione, separatamente per RTN, MSE ponderato per le attivazioni e GPTQ con Hessiana completa. Questa ricerca in una dimensione comprende la testa originale e il centraggio fisso rispetto alla media, preserva la distribuzione softmax a piena precisione e lascia invariato il decoder addestrato; una correzione di rango uno gestisce i percorsi non lineari dei logit, come il soft-capping. Su sette teste, i miglioramenti con W4 si concentrano dove la quantizzazione di base altera sensibilmente le predizioni: su Phi-4-mini, la KL di AW-MSE scende da 0,936 a 0,256. I miglioramenti resistono a una calibrazione GPTQ più accurata e restano complementari alla riscalatura esatta per canale e alla quantizzazione affine. Su quattro teste e tre quantizzatori W4, i coefficienti selezionati su WikiText e poi mantenuti fissi si trasferiscono anche a C4 e OpenWebMath: superano il centraggio rispetto alla media in tutti i 18 confronti in cui il coefficiente fisso è diverso da $1$ e ne eguagliano i risultati nei restanti sei. Con W2, usato come stress test della compressione, i benefici si estendono a quasi tutta la matrice modello-quantizzatore. Un’analisi dei residui a parità di condizioni mostra che una maggiore fedeltà può accompagnarsi a un errore di ricostruzione dei logit più elevato, pur riducendo il costo del residuo ponderato secondo Fisher. Per le teste compatibili con la traslazione, la riparametrizzazione non aggiunge operazioni durante l’inferenza e preserva l’esecuzione W4 con pesi impacchettati: mantenendo il decoder in BF16, la quantizzazione della testa di output di Phi riduce del 10,8% la latenza della generazione con batch di dimensione uno rispetto alla configurazione di riferimento con la testa in BF16.