Riepilogo Tecnico: STaR-Quant

Definizione del Problema

I Diffusion Large Language Models (DLLM) sono emersi come una promettente alternativa agli LLM autoregressivi, utilizzando un processo iterativo di denoising con maschera e contesto bidirezionale per generare testo. Tuttavia, la loro implementazione è ostacolata da ingenti sovraccarichi di memoria e computazionali derivanti dalle grandi dimensioni dei modelli e dal processo di denoising iterativo multi-step. Sebbene la Post-Training Quantization (PTQ) abbia avuto successo nella compressione degli LLM autoregressivi, l'applicazione diretta degli esistenti metodi PTQ ai DLLM provoca un significativo degrado delle prestazioni, in particolare sotto una quantizzazione aggressiva di pesi e attivazioni a bassi bit.

Il documento identifica due sfide specifiche, uniche per i DLLM, che causano questo fallimento:

1. Disparità di Attivazione Dipendente dallo Stato: All'interno di un singolo step di denoising, la sequenza contiene sia token mascherati (posizioni incerte da predire) che token non mascherati (contesto visibile). Questi due stati esibiscono distribuzioni di attivazione e pattern di outlier distinti. L'applicazione di una trasformazione di attivazione o di uno spazio di quantizzazione unificato a entrambi gli stati non riesce a smussare efficacemente le distribuzioni per nessuno dei due.
2. Accumulo di Errore Temporale: La generazione dei DLLM procede attraverso step iterativi in cui l'output di uno step funge da input per il successivo. Gli errori di quantizzazione, in particolare nella moltiplicazione matriciale tra la mappa di attenzione softmax e gli stati di valore, non sono isolati a un singolo step, ma si accumulano lungo la traiettoria di denoising, portando a una deviazione significativa dal percorso a precisione intera (full-precision).

Metodologia: STaR-Quant

Per affrontare queste sfide, gli autori propongono STaR-Quant, un framework PTQ progettato specificamente per i DLLM, basato sulla coerenza stato-tempo. Il framework opera su due dimensioni:

1. Trasformazione dell'Attivazione Guidata dallo Stato (SGAT)

SGAT affronta la disparità di attivazione tra token mascherati e non mascherati.

• Meccanismo: Invece di una singola trasformazione per tutti i token, SGAT decompone la dimensione nascosta in sottospazi condivisi e specifici per stato.
  • Sottospazio Condiviso ($U_{sh}$): Cattura le componenti di attivazione comuni.
  • Sottospazi Specifici per lo Stato ($U_m, U_u$): Assorbono i pattern di outlier specifici rispettivamente per i token mascherati e non mascherati.
• Implementazione: Gate di sottospazio binari indirizzano i token mascherati e non mascherati verso i rispettivi spazi di trasformazione ($P_m$ e $P_u$). Fondamentalmente, il lato dei pesi mantiene una trasformazione statica unificata ($P_e$) applicata una sola volta ai pesi. Ciò evita la duplicazione dei pesi trasformati durante l'inferenza, pur permettendo al lato delle attivazioni di adattarsi allo stato del token.
• Ottimizzazione: I parametri di trasformazione sono appresi minimizzando l'errore di quantizzazione per livello tra l'output a precisione intera e la proiezione quantizzata guidata dallo stato.

2. Compensazione dell'Attenzione Temporale (TAC)

TAC affronta l'accumulo di errori di quantificazione attraverso gli step iterativi di denoising.

• Target: Il metodo prende di mira l'output della moltiplicazione softmax-valore nel meccanismo di attenzione, un'operazione nota per essere soggetta a errori sotto quantizzazione a bassi bit.
• Meccanismo: TAC inserisce una mappatura affine leggera, a diagonale di blocco, prima della proiezione dell'output dell'attenzione. Non altera il calcolo dell'attenzione stesso, ma corregge la rappresentazione dell'attenzione quantizzata ($\hat{Z}_t$) per allinearla con la sua controparte a precisione intera ($Z^{FP}_t$).
• Stima: I parametri di compensazione (matrici per blocco $M_b$ e bias $d_b$) sono stimati in forma chiusa facendo corrispondere le statistiche del primo e secondo ordine (media e covarianza) delle rappresentazioni a precisione intera e quantizzate. Un termine di smorzamento e uno smoothing dei valori singolari vengono applicati per garantire la stabilità data la limitata quantità di dati di calibrazione.

Contributi Chiave

1. Identificazione del Problema: Gli autori identificano formalmente la disparità di attivazione dipendente dallo stato e l'accumulo di errore temporale come i principali colli di bottiglia per la quantizzazione a bassi bit dei DLLM.
2. Trasformazione dell'Attivazione Guidata dallo Stato (SGAT): Un metodo innovativo che mappa le attivazioni mascherate e non mascherate in spazi di trasformazione condivisi e specifici per stato, mantenendo una trasformazione unificata statica sul lato dei pesi, smussando efficacemente gli outlier specifici per stato senza duplicazione dei pesi.
3. Compensazione dell'Attenzione Temporale (TAC): Un metodo di compensazione affine per blocchi in forma chiusa che allinea le rappresentazioni dell'attenzione quantizzate con le controparti a precisione intera prima della proiezione dell'output, mitigando l'accumulo di errori lungo la traiettoria di denoising.
4. Validazione Empirica: Esperimenti estesi dimostrano che STaR-Quant supera costantemente i forti baseline PTQ (inclusi RTN, AWQ, QuaRot e DLLMQuant) su molteplici DLLM e task.

Risultati Sperimentali

Gli autori hanno valutato STaR-Quant su tre DLLM rappresentativi: LLaDA-8B, LLaDA-1.5-8B e Dream-7B, utilizzando la quantizzazione a 4 bit per pesi e attivazioni (W4A4).

• Performance: STaR-Quant ha ottenuto la migliore performance media su 9 benchmark (che coprono conoscenza generale, ragionamento e generazione di codice).
  • Su LLaDA-8B, ha migliorato il punteggio medio di 2,78 punti rispetto al più forte baseline (DLLMQuant++).
  • Su Dream-7B, ha migliorato il punteggio medio di 1,69 punti.
  • Sono stati osservati guadagni notevoli nei task di ragionamento (GSM8K) e generazione di codice (HumanEval), che sono sensibili agli errori di rappresentazione intermedia.
• Efficienza:
  • Accelerazione (Speedup): STaR-Quant ha ottenuto accelerazioni di inferenza da 1,65× a 1,69× rispetto al deployment in FP16.
  • Memoria: Ha ridotto l'uso di memoria GPU di picco da 3,05× a 3,14× (ad esempio, riducendo la memoria di LLaDA da ~15,9GB a ~5,2GB).
• Studi di Ablazione: Rimuovere TAC ha causato un calo significativo delle prestazioni (indicando l'importanza della correzione dell'errore temporale), mentre rimuovere SGAT ha anch'esso degradato i risultati, confermando la natura complementare della coerenza per stato e temporale.

Significato e Rivendicazioni

Il documento sostiene che STaR-Quant fornisca un necessario cambio di strategia di quantizzazione per i modelli linguistici basati sulla diffusione. Modellando esplicitamente le dimensioni dello stato (mascherato vs non mascherato) e del tempo (accumulo iterativo), il framework risolve il disallineamento tra le assunzioni convenzionali della PTQ e il processo di inferenza della diffusione.

Gli autori posizionano STaR-Quant come una soluzione pratica per il deployment dei DLLM su dispositivi con risorse limitate, offrendo un equilibrio tra alti rapporti di compressione (W4A4), significativi risparmi di memoria e accelerazioni, mantenendo livelli di accuratezza vicini ai modelli a precisione intera. Il lavoro suggerisce che la "coerenza stato-tempo" sia un principio efficace per la quantizzazione efficiente dei DLLM, sebbene gli autori notino che un'ulteriore validazione su architetture diverse e bit-width più aggressivi (ad esempio, 2-bit o 3-bit) rimanga una direzione aperta.