Riassunto Tecnico: FPTQuant: Trasformazioni Preservanti la Funzione per la Quantizzazione di LLM

Definizione del Problema

I Large Language Models (LLM) comportano costi computazionali ed energetici significativi durante l'inferenza, limitando il loro dispiegamento su dispositivi edge con risorse limitate. Sebbene la quantizzazione (mappatura di valori continui su una griglia intera uniforme) sia la tecnica più efficace per migliorare l'efficienza, la quantizzazione ingenua degli LLM spesso porta a un grave degrado delle prestazioni. Ciò è dovuto principalmente agli "outlier"—attivazioni e pesi con magnitudo elevate che interrompono la distribuzione uniforme necessaria per la quantizzazione a bassi bit.

La sfida fondamentale risiede nel compromesso tra intervallo e precisione:

1. Aumentare l'intervallo per catturare gli outlier comporta una perdita di precisione per la maggior parte dei valori raggruppati intorno allo zero.
2. Mantenere la precisione tagliando (clipping) gli outlier porta a una perdita di informazioni e al collasso del modello.

I lavori precedenti hanno esplorato le Function-Preserving Transforms (FPT)—operazioni come la scalatura o le rotazioni che smorzano gli outlier senza alterare l'output del modello in assenza di quantizzazione. Tuttavia, i metodi esistenti affrontano spesso un compromesso tra espressività (la capacità di ridurre efficacemente l'errore di quantizzazione) e overhead computazionale (costo di inferenza). Molte trasformazioni espressive sono non-mergibili, richiedendo nuove operazioni nel grafo computazionale e comportando una latenza significativa.

Metodologia: FPTQuant

Il paper introduce FPTQuant, un framework che utilizza tre FPT innovative, leggere ed espressive, progettate specificamente per i transformer. La filosofia centrale è quella di sfruttare le equivarianze e le indipendenze intrinseche nelle operazioni canoniche dei transformer per progettare trasformazioni che siano sia mergibili (fuse nei pesi esistenti) che espressive.

Le tre trasformazioni primarie sono:

1. Trasformazione Pre-RoPE Mergibile per Query e Key ($T_k, \bar{T}_k$)

• Contesto: I moderni LLM utilizzano le Rotary Positional Embeddings (RoPE), che impediscono la semplice fusione di trasformazioni nei pesi di query ($W_q$) e key ($W_k$) perché RoPE non commuta con trasformazioni lineari arbitrarie.
• Meccanismo: Gli autori introducono una coppia di trasformazioni applicate prima di RoPE. $T_k$ è applicata alle key, e $\bar{T}_k$ (un'approssimazione inversa) è applicata alle query. Queste consistono in matrici di rotazione $2 \times 2$ scalate.
• Proprietà: Grazie alla struttura a blocchi diagonali di RoPE, queste specifiche matrici di rotazione/scalatura commutano con l'operazione RoPE. Ciò consente alle trasformazioni di essere fuse direttamente nei pesi $W_q$ e $W_k$ post-training, comportando zero overhead di inferenza.
• Applicabilità: Valido per modelli con RoPE (es. serie Llama) e modelli senza embedding posizionali.

2. Trasformazione Multihead Value Mergibile ($T_v, \bar{T}_v$)

• Contesto: Nel meccanismo di attenzione, la moltiplicazione di matrici batch (BMM) opera sui valori ($V$) e sulla proiezione di output ($W_o$). La dimensione $d$ (dimensione del valore) è indipendente dal calcolo della probabilità di attenzione.
• Meccanismo: Una trasformazione invertibile $T_v$ viene applicata ai pesi di proiezione dei valori ($W_v$), e la sua inversa $\bar{T}_v$ viene applicata ai pesi di proiezione di output ($W_o$).
• Proprietà: Poiché i diversi head di attenzione sono indipendenti, una matrice invertibile unica può essere ottimizzata per ogni head. Questa trasformazione viene fusa in $W_v$ e $W_o$.
• Vantaggio: Offre un'alta espressività (matrici invertibili complete per head) pur rimanendo mergibile, a differenza dei precedenti metodi basati su rotazione che spesso condividevano una singola rotazione tra gli head o non erano mergibili.

3. Pseudo-Dynamic Residual Scaling ($S_n$)

• Contesto: Le connessioni residue nei transformer spesso bypassano i livelli di normalizzazione, portando a token con scale molto diverse. Ciò causa outlier nei blocchi successivi (es. l'output di SwiGLU).
• Meccanismo: Gli autori propongono di spostare l'RMSNorm affinché si applichi anche allo stream residuo. Un fattore di scalatura per token $S_n$ (basato sulla radice quadrata della media dei quadrati del residuo) viene applicato all'input residuo. Per preservare l'output, gli output dei blocchi di attenzione e MLP ($Y_n$) vengono riscalati dallo stesso $S_n$ prima di essere aggiunti al residuo.
• Proprietà: La moltiplicazione di matrici e i livelli lineari senza bias commutano con la scalatura per token. Ciò consente di applicare la scalatura in profondità all'interno dei blocchi di attenzione e MLP per ridurre gli outlier intra-blocco.
• Implementazione: Le scale sono calcolate ricorsivamente usando l'output normalizzato del blocco precedente. In pratica, la divisione finale per la scala è non necessaria poiché l'LM head inizia con un RMSNorm.

Strategia di Ottimizzazione

FPTQuant impiega un processo di ottimizzazione in due fasi:

1. Ottimizzazione Locale: Le trasformazioni sono ottimizzate localmente e indipendentemente per minimizzare la norma $L_p$ (specificamente $L_4$) dei pesi fusi. Ciò riduce gli outlier e fornisce un'inizializzazione stabile.
2. Addestramento End-to-End (Student-Teacher): Un modello studente quantizzato (con trasformazioni) viene addestrato per approssimare un modello insegnante a piena precisione utilizzando la perdita di Jensen-Shannon Divergence (JSD).
   • Distinzione Cruciale: A differenza dei metodi che utilizzano la perdita di predizione del next-token, l'approccio student-teacher evita l'overfitting sul dataset di addestramento (es. Wikitext) confrontando la distribuzione di probabilità completa dell'insegnante. Ciò migliora la generalizzazione sui task downstream.

Contributi Chiave

1. Nuove FPT: Introduzione di tre trasformazioni specifiche (Pre-RoPE, Value e Residual Scaling) che sono progettate per essere sia altamente espressive che mergibili, eliminando l'overhead di inferenza.
2. Quantizzazione Statica INT4: Dimostrazione che FPTQuant consente la quantizzazione statica INT4 (pesi, attivazioni, KV cache) con un overhead minimo, ottenendo un'accelerazione fino a 3.9× rispetto a FP16 nei compiti di prefill.
3. Compromesso Accuratezza-Velocità: FPTQuant raggiunge prestazioni pari o superiori ai metodi SOTA (come FlatQuant e SpinQuant), pur essendo significativamente più veloce (fino al 29% più veloce di FlatQuant) e non richiedendo kernel personalizzati.

Risultati Sperimentali

Gli autori hanno valutato FPTQuant su vari modelli (Llama-2/3, Ministral, Qwen) attraverso molteplici configurazioni di quantizzazione (W4A4KV4, W4A8KV8, ecc.).

• Velocità: FPTQuant ottiene costantemente un'accelerazione di prefill di 2.8×–3.9× rispetto a FP16. È più veloce di SpinQuant e costantemente più veloce di FlatQuant (accelerazione relativa del 15–29%).
• Accuratezza:
  • In scenari di deployment "difficili" (quantizzando tutte le attivazioni intermedie), FPTQuant supera i baseline come SpinQuant e FlatQuant.
  • Nei setting standard (Linear+KV), FPTQuant performa comparabilmente a FlatQuant, che è attualmente lo SOTA per accuratezza.
  • Nella impegnativa configurazione W4A4KV4, FPTQuant mostra una perplexity competitiva e punteggi di ragionamento zero-shot, spesso superando SpinQuant e QuaRot.
• Stabilità: L'analisi di sensibilità mostra che FPTQuant è function-preserving anche sotto aggiornamenti di parametri rumorosi, mentre metodi come FlatQuant (che si affidano a inversioni in virgola mobile) e OSTQuant (non-function-preserving) mostrano instabilità o degradazione con il rumore.
• Costo di Addestramento: L'addestramento di FPTQuant è efficiente. Sebbene sia leggermente più lento di RTN-opt, è generalmente più veloce di FlatQuant perché le sue trasformazioni sono mergibili, evitando i passaggi forward online richiesti dalle trasformazioni non-mergibili.

Significato e Rivendicazioni

Il paper sostiene che FPTQuant rappresenti un passo avanti significativo per rendere gli LLM utilizzabili nel deployment edge. Sfruttando esplicitamente le proprietà matematiche delle operazioni dei transformer (equivarianze), gli autori hanno creato trasformazioni che non sacrificano la velocità di inferenza per l'accuratezza.

• Impatto Pratico: FPTQuant consente l'uso di bit-width più piccoli (INT4) senza l'elevato costo computazionale della quantizzazione dinamica o delle trasformazioni non-mergibili. Ciò riduce l'impronta di memoria e il consumo energetico degli LLM, potenzialmente democratizzando l'accesso all'IA su smartphone e dispositivi edge.
• Contributo Metodologico: Il lavoro evidenzia come la comprensione delle proprietà di commutazione delle operazioni esistenti sia la chiave per progettare strategie di quantizzazione efficaci. Argomenta contro l'uso della perdita di next-token prediction per l'addestramento delle FPT a causa dei rischi di overfitting, promuovendo invece l'addestramento student-teacher con matching delle distribuzioni.
• Limitazioni: Gli autori notano con modestia che, sebbene l'approccio sia robusto tra diverse famiglie di modelli, non è garantito che si traduca allo stesso modo per tutti gli LLM (es. Mixture-of-Experts). Inoltre, le attuali limitazioni hardware (mancanza di supporto nativo INT4 su molti GPU) significano che il benchmarking si basa sull'emulazione, e i pieni benefici potrebbero realizzarsi solo su hardware futuri o acceleratori specifici per l'edge.