LATMiX: Learnable Affine Transformations for Microscaling Quantization of LLMs

1. Enunciato del Problema

I Large Language Models (LLM) sono sempre più distribuiti utilizzando la Quantizzazione Post-Addestramento (PTQ) per ridurre i costi di memoria e calcolo. Un formato di dati specifico che sta guadagnando terreno è il Micro-scala (MX), che partiziona i tensori in blocchi, assegnando un fattore di scala unico a ciascun blocco per gestire gli intervalli dinamici in modo più efficiente rispetto alla scala globale.

Tuttavia, combinare la quantizzazione MX con le tecniche esistenti di mitigazione degli outlier (che sono cruciali per la quantizzazione a basso numero di bit) presenta sfide significative:

• Outlier delle Attivazioni: Le attivazioni degli LLM contengono valori estremi (outlier) che dominano l'errore di quantizzazione.
• Limitazioni dei Metodi Attuali: Gli approcci esistenti utilizzano trasformazioni invertibili (tipicamente rotazioni o trasformate di Hadamard) per ridistribuire l'energia delle attivazioni e sopprimere gli outlier.
  • Trasformazioni Globali: Applicare una rotazione globale alle attivazioni prima della quantizzazione MX causa un grave degrado delle prestazioni perché la trasformazione globale è in conflitto con la struttura di scala per blocchi dell'MX.
  • Trasformazioni a Blocchi Diagonali: Lavori recenti suggeriscono di applicare trasformazioni indipendentemente all'interno di ciascun blocco MX (matrici a blocchi diagonali). Sebbene ciò eviti il disallineamento, limita la trasformazione a piccoli sottospazi isolati, impedendo una ridistribuzione efficace della massa delle attivazioni su tutto il tensore. Di conseguenza, gli outlier dominanti non possono essere sufficientemente diffusi, portando a una precisione subottimale nei regimi a basso numero di bit (ad esempio, 4-bit).

2. Metodologia: LATMiX

Gli autori propongono LATMiX (Trasformazioni Affini Apprendibili per Micro-scala), un metodo che generalizza la riduzione degli outlier a trasformazioni affini invertibili a rango completo apprendibili, ottimizzate specificamente per il formato di dati MX.

A. Analisi Teorica

Gli autori derivano un limite superiore teorico sull'errore di quantizzazione sotto la quantizzazione MX. L'analisi rivela che l'errore dipende da due fattori in competizione:

1. La Norma Spettrale della Trasformazione Inversa ($||A^{-1}||_\sigma$): Una misura di quanto la trasformazione amplifica gli errori.
2. Il Massimo Atteso delle Attivazioni per Blocco: La media dei valori assoluti massimi all'interno di ciascun blocco MX.

Il documento sostiene che le trasformazioni a blocchi diagonali (utilizzate nei lavori precedenti) sono un caso speciale che non riesce a bilanciare globalmente questi fattori. Invece, è richiesta una trasformazione affine completa (non a blocchi diagonali) per tenere conto esplicitamente sia della struttura a blocchi MX sia della distribuzione globale delle caratteristiche.

B. Architettura LATMiX

LATMiX apprende un insieme di trasformazioni affini invertibili ($T(x) = Ax + v$) che possono essere incorporate nei livelli lineari del modello, comportando zero overhead di inferenza.

• Posizioni delle Trasformazioni:
  • $T_1$: Applicata globalmente al percorso residuo (embedding di input).
  • $T_2$: Applicata localmente ai vettori di valore nell'Attention Multi-Testa (MHA) e all'interno delle Reti Feed-Forward (FFN).
  • $T_3$: Una trasformazione online a blocchi-Hadamard (mantenuta per coerenza con le baseline precedenti).
• Parametrizzazione: Per garantire l'invertibilità e abilitare l'ottimizzazione standard del deep learning, gli autori propongono due parametrizzazioni per la matrice $A$:
  1. Decomposizione LU: $A = P(L + \text{diag}(s))$, dove $L$ è unitriangolare inferiore.
  2. Decomposizione QR: $A = Q(R + \text{diag}(s))$, dove $Q$ è ortogonale (tramite esponenziale di matrice di una matrice antisimmetrica) e $R$ è triangolare superiore.
  • Viene applicata una regolarizzazione a volume preservante per garantire che il determinante rimanga vicino a 1, prevenendo l'instabilità numerica mentre permette a direzioni individuali di espandersi o contrarsi.
• Obiettivo di Ottimizzazione:
  • Loss di Distillazione: Invece di imporre un'invarianza computazionale rigorosa (che richiede matrici ortogonali), LATMiX utilizza una loss di divergenza di Kullback-Leibler (KL) tra il modello insegnante a precisione intera e il modello studente quantizzato trasformato. Ciò permette alla trasformazione di discostarsi dall'ortogonalità per adattarsi meglio alla distribuzione dei dati.
  • Regolarizzazione: Un termine per mantenere il volume della trasformazione vicino a 1.

C. Meccanismo di Incorporamento (Folding)

Un contributo critico è la capacità di incorporare (fold) le trasformazioni affini apprese ($A$ e $v$) direttamente nelle matrici dei pesi e nei bias dell'LLM.

• La trasformazione inversa è incorporata nelle proiezioni di input (K, Q, V) del meccanismo di attenzione.
• La trasformazione in avanti è incorporata nelle proiezioni di output dei livelli di attenzione e FFN.
• Ciò garantisce che durante l'inferenza, il modello si comporti esattamente come un modello quantizzato standard senza costi computazionali aggiuntivi.

3. Contributi Chiave

1. Insight Teorico: Fornito un limite teorico che mostra come l'errore di quantizzazione nei formati MX dipenda sia dalla distribuzione globale delle caratteristiche sia dalla struttura a blocchi, dimostrando che le trasformazioni a blocchi diagonali sono insufficienti per una riduzione ottimale dell'errore.
2. Algoritmo LATMiX: Introdotto un metodo per apprendere trasformazioni affini invertibili generali (non limitate alle rotazioni) utilizzando parametrizzazioni LU/QR e loss di distillazione.
3. Deployment a Zero-Overhead: Dimostrato che queste complesse trasformazioni apprese possono essere matematicamente incorporate nei livelli lineari esistenti, mantenendo la velocità di inferenza e l'ingombro di memoria dei modelli quantizzati standard.
4. Valutazione Completa: Validato il metodo su molteplici dimensioni di modello (da 1B a 32B) e architetture (Llama, Qwen) su benchmark diversificati.

4. Risultati Sperimentali

Gli autori hanno valutato LATMiX su sette benchmark di ragionamento zero-shot (ARC, HellaSwag, WinoGrande, PIQA, BoolQ, OpenBookQA) e sul compito di perplessità WikiText2.

• Miglioramenti delle Prestazioni: LATMiX ha costantemente superato le baseline solide, tra cui:
  • Quantizzazione MX standard (RTN, GPTQ).
  • Metodi basati su rotazione (QuaRot, SpinQuant).
  • Metodi di rotazione a blocchi diagonali (MR-GPTQ, OSTQuant).
  • Altri metodi basati su trasformazioni (FlatQuant).
• Recupero della Precisione: Nella quantizzazione MXFP4 (4-bit), LATMiX ha raggiunto fino a un 4% di miglioramento nel recupero medio della precisione rispetto alle migliori baseline. Ad esempio, su Llama-3.2-1B, LATMiX-LU ha raggiunto un recupero del 93,88% contro il 93,07% di MR-GPTQ.
• Robustezza: Il metodo ha mostrato miglioramenti coerenti tra diverse dimensioni di modello (da 1B a 32B) e granularità di quantizzazione (MXFP4 e MXINT4).
• Studi di Ablazione:
  • Completo vs. A Blocchi: Le trasformazioni affini complete (a livello di tensore) hanno superato significativamente quelle a blocchi diagonali.
  • Affine vs. Ortogonale: Apprendere trasformazioni affini generali (con bias) è stato superiore all'apprendimento di sole matrici ortogonali (rotazione).
  • Funzione di Loss: La loss di distillazione KL si è dimostrata più efficace per la generalizzazione zero-shot rispetto alle loss MSE o Cross-Entropy.

5. Significato

• Colmare il Divario: LATMiX colma con successo il divario tra la necessità di una soppressione aggressiva degli outlier (tramite trasformazioni globali) e i vincoli strutturali dei moderni formati MX supportati dall'hardware.
• Efficienza Hardware: Incorporando le trasformazioni, LATMiX permette l'uso di potenti operazioni matematiche apprese senza sacrificare i guadagni di efficienza di inferenza forniti dalla quantizzazione MX.
• Direzione Futura: Il lavoro suggerisce che limitare le trasformazioni a rotazioni o matrici a blocchi diagonali è subottimale. Apre la porta all'apprendimento di trasformazioni più espressive e dipendenti dai dati che rispettano i formati di dati specifici dell'hardware AI di prossima generazione.

In sintesi, LATMiX dimostra che analizzando teoricamente l'interazione tra distribuzioni di attivazione e strutture a blocchi MX, e apprendendo trasformazioni affini flessibili, è possibile raggiungere una quantizzazione quasi senza perdita degli LLM a larghezze di bit estremamente basse (4-bit) senza incorrere in overhead di inferenza.