🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

Questo articolo introduce BCJR-QAT, un rilassamento differenziabile della quantizzazione dei pesi codificata a traliccio che sostituisce l'argmax non differenziabile di Viterbi con un algoritmo BCJR somma-prodotto controllato per temperatura, abilitando un addestramento consapevole della quantizzazione end-to-end che supera empiricamente la frontiera dello stato dell'arte PTQ sugli LLM.

Autori originali: Venugopalan Iyengar

Pubblicato 2026-05-12

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Venugopalan Iyengar

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di avere una biblioteca immensa di libri (un Modello Linguistico di grandi dimensioni) che desideri ridurre affinché possa essere ospitato su un laptop o un telefono standard. Per farlo, devi comprimere le "parole" (i pesi) all'interno della biblioteca.

Il documento introduce un nuovo metodo chiamato BCJR-QAT per comprimere questi libri ulteriormente, fino a raggiungere appena 2 bit per parola (il che equivale a comprimere una foto ad alta definizione in una minuscola anteprima granulosa).

Ecco la storia di come hanno fatto, utilizzando semplici analogie:

1. Il Problema: La Trappola della "Strada a Senso Unico"

In precedenza, il modo migliore per ridurre queste dimensioni era un metodo chiamato QTIP. Immagina QTIP come un bibliotecario molto intelligente che esamina una pagina di testo e dice: "Ok, sostituirò questa frase complessa con la frase semplice più vicina del mio dizionario".

Tuttavia, c'era un limite. Una volta che il bibliotecario aveva preso quella decisione, non poteva tornare indietro e cambiare idea per rendere l'intero libro più scorrevole. Era come una "strada a senso unico". Se il bibliotecario sceglieva una frase leggermente sbagliata all'inizio, il resto del libro ne risentiva e non poteva correggerlo perché la decisione era "rigida" e definitiva.

Per risolvere questo problema, i ricercatori usavano solitamente l'Addestramento (QAT), dove si permette al bibliotecario di esercitarsi e aggiustare le sue scelte. Ma ecco il punto critico: il processo decisionale del bibliotecario coinvolge un labirinto complesso (un "reticolo"). Per trovare il percorso migliore attraverso il labirinto, utilizzano una regola chiamata Viterbi, che è come un pulsante "scegli l'opzione assolutamente migliore".

Il Problema: Non puoi insegnare a un computer ad imparare attraverso un pulsante "scegli il migliore" perché quel pulsante è un vicolo cieco matematico. Se premi il pulsante, il computer non sa come aggiustare le sue scelte per migliorare; si limita a scattare all'opzione più vicina. È come cercare di sterzare un'auto permettendole solo di saltare istantaneamente alla corsia successiva: non puoi fare curve fluide.

2. La Soluzione: La Decisione "Morbida" (BCJR)

Gli autori, Venugopalan Iyengar, hanno inventato un nuovo modo per permettere al bibliotecario di esercitarsi. Invece di forzare immediatamente una decisione rigida di "scegli il migliore", hanno introdotto un concetto chiamato Temperatura ( $T$ ).

Alta Temperatura (Caldo): Immagina che il bibliotecario sia molto rilassato e indeciso. Invece di scegliere una sola frase, ne considera molte contemporaneamente, assegnando una "probabilità" a ciascuna. È come una giornata nebbiosa in cui puoi vedere diversi percorsi, non solo uno. Questa visione "morbida" è matematicamente fluida, il che significa che il computer può calcolare facilmente come aggiustare le scelte per migliorare il libro.
Bassa Temperatura (Freddo): Man mano che l'addestramento termina, il bibliotecario diventa "più freddo" e più deciso. La nebbia si dirada e si blocca sulla singola frase migliore (la decisione rigida).

Questo metodo, chiamato BCJR, trasforma la "strada a senso unico" in una strada liscia e percorribile. Il computer può ora imparare esattamente come regolare i pesi per rendere l'intero libro più scorrevole, non solo singole frasi.

3. L'Errore dell'"Eccesso"

Gli autori hanno scoperto una curiosa stranezza nel modo in cui funziona questa "temperatura".

Nella fisica tradizionale, quando cerchi la soluzione migliore, di solito inizi molto caldo (molto indeciso) per esplorare tutte le possibilità, e poi raffreddi lentamente.

La Scoperta del Documento: Per questo tipo specifico di compressione, iniziare "troppo caldo" è un disastro. Se il bibliotecario è troppo indeciso all'inizio, si perde in un quartiere cattivo (un "bacino di Voronoi peggiore") e rimane bloccato lì. Quando si raffredda e cerca di trovare il percorso migliore, è già perso in una zona negativa e non può tornare indietro.
La Soluzione: Hanno scoperto che iniziare a una temperatura moderata (né troppo calda, né troppo fredda) funziona meglio. È come iniziare un'escursione con una mappa chiara invece di vagare alla cieca nella nebbia. Saltando la fase "super calda", hanno evitato di perdersi e hanno trovato una soluzione migliore.

4. I Risultati: Una Vittoria per i Computer Piccoli

Il team ha testato questo metodo su due diversi tipi di modelli:

Il Test "Proxy" (OLMoE): Hanno cercato di ottimizzare la compressione guardando solo quanto bene ogni strato del modello ricostruiva i propri dati (come verificare se una fotocopia assomiglia all'originale). Risultato: Ha fallito. Il modello è effettivamente peggiorato rispetto al vecchio metodo. Questo ha insegnato loro che fare semplicemente una "fotocopia migliore" non significa che il libro si legga meglio.
Il Test "Reale" (Llama-3.2): Hanno ottimizzato il modello per leggere e comprendere effettivamente il testo meglio (utilizzando un metodo chiamato "distillazione", dove un modello insegnante intelligente guida lo studente). Risultato: Ha funzionato!
- Su uno strato specifico del modello, il loro nuovo metodo ha battuto il vecchio metodo migliore con un margine piccolo ma significativo.
- Quando l'hanno applicato a più strati, i miglioramenti si sono sommati in modo "super-additivo" (il tutto è diventato maggiore della somma delle sue parti).

5. Il Motore Sotto il cofano

Fare questi calcoli è solitamente molto lento e richiede supercomputer costosi. Gli autori hanno anche costruito un "motore" specializzato (un kernel Triton) che esegue questi complessi calcoli incredibilmente velocemente su una singola scheda grafica consumer (come una RTX 4080). L'hanno resa 6,5 volte più veloce dei metodi standard, dimostrando che questa tecnica avanzata non ha bisogno di un data center per essere eseguita.

Riepilogo

Il documento presenta un nuovo modo per ridurre i modelli AI affinché si adattino ai dispositivi consumer.

Il Trucco: Hanno sostituito un processo decisionale rigido e non insegnabile con uno "morbido" e insegnabile che si indurisce gradualmente.
L'Insight: Non dovresti iniziare il processo di apprendimento troppo "caldo" (troppo casuale); un inizio moderato previene che il modello si perda.
Il Risultato: Se addestrato correttamente per comprendere il testo (non solo per copiare dati), questo metodo produce risultati migliori rispetto alle tecniche precedenti all'avanguardia, ed è abbastanza veloce da essere eseguito su un normale PC da gaming.

Gli autori hanno rilasciato il loro codice e i modelli addestrati in modo che altri possano provarlo.

Riepilogo Tecnico: BCJR-QAT

Enunciato del Problema

Raggiungere 2 bit per peso (bpw) o meno per i Large Language Models (LLM) su hardware consumer è un collo di bottiglia critico. Sebbene i metodi di Quantizzazione Post-Addestramento (PTQ) come QTIP abbiano stabilito una frontiera solida per la quantizzazione a 2 bit utilizzando la Quantizzazione a Codifica a Griglia (TCQ), sono limitati dal "soffitto PTQ". Spingere oltre questo limite richiede l'Addestramento Consapevole della Quantizzazione (QAT). Tuttavia, le tecniche QAT standard falliscono sui codici a griglia perché l'algoritmo di decodifica di Viterbi coinvolge un'operazione argmax non differenziabile su un numero esponenziale di percorsi. L'applicazione di un Stimatore Diretto (STE) alla griglia distrugge l'accoppiamento tra le decisioni di peso adiacenti, che è il vantaggio principale della codifica a griglia rispetto alla quantizzazione scalare.

Metodologia: BCJR-QAT

Il documento introduce BCJR-QAT, un rilassamento differenziabile della quantizzazione a codifica a griglia che sostituisce il rigido argmax di Viterbi con un'aspettativa morbida derivata dall'algoritmo BCJR (Bahl-Cocke-Jelinek-Raviv) avanti-indietro.

Meccanismo Principale

Analogia con la Meccanica Statistica: La griglia è trattata come un grafo fattoriale equivalente a una catena di spin unidimensionale di tipo Ising. L'argmax di Viterbi è visto come il limite a temperatura zero ( $T \to 0$ ) di una distribuzione di Boltzmann sui percorsi.
Parola di Codice Morbida: A una temperatura finita $T > 0$ , il metodo calcola le probabilità posteriori marginali esatte di ogni stato utilizzando l'algoritmo BCJR. L'output quantizzato è l'aspettativa di Boltzmann della parola di codice emessa:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
Questa parola di codice morbida è esattamente differenziabile rispetto all'immagine pre-ponderale e ai parametri della griglia.
Ricottura: Durante l'addestramento, la temperatura $T$ viene ricotta da un valore iniziale $T_0$ a un valore vicino allo zero $T_{end}$ . Man mano che $T \to 0$ , la parola di codice morbida converge all'output rigido di Viterbi, "cristallizzando" efficacemente l'ottimizzazione continua in una soluzione discreta.

Innovazioni Tecniche Chiave

Differenziabilità: La ricorsione avanti-indietro di BCJR (complessità $O(L \cdot S^2)$ ) fornisce gradienti esatti attraverso la griglia, preservando l'accoppiamento tra stati che i metodi STE scalari perdono.
Kernel Triton: Gli autori hanno implementato un kernel Triton fuso che combina le operazioni di raccolta, log-sum-exp e addizione in un singolo avvio. Questo riduce il passaggio avanti/indietro da centinaia di avvii di kernel a uno, ottenendo un accelerazione di 6,57× rispetto a un'implementazione di riferimento PyTorch autograd, mantenendo la parità numerica FP32.
Correzione del Programma (Skip-High-T): Il documento diagnostica una modalità di fallimento critica nei programmi di ricottura simulata standard (che iniziano a $T_0=1.0$ ). A temperature elevate, il gradiente della parola di codice morbida è "sfocato" e non informativo, causando lo spostamento dell'ottimizzatore verso un bacino di Voronoi peggiore prima che la fase di raffreddamento possa recuperare. La correzione proposta è di saltare la fase ad alta- $T$ , inizializzando a una temperatura inferiore (es. $T_0=0.3$ ) dove i gradienti sono già informativi.

Risultati Sperimentali

Il documento valuta BCJR-QAT su due architetture: OLMoE-1B-7B (MoE) e Llama-3.2-1B (Dense).

1. OLMoE-1B-7B (Risultato Negativo: Divario del Proxy)

Configurazione: QAT greedy per livello che ottimizza l'MSE di ricostruzione (un obiettivo proxy).
Esito: BCJR-QAT ha ridotto con successo l'MSE di ricostruzione per livello del 17–20% rispetto a QTIP-PTQ. Tuttavia, questo miglioramento non si è tradotto nelle prestazioni del compito finale. I modelli hanno regredito di $\sim 1.3$ di perplessità (PPL) su WikiText-2 rispetto alla baseline QTIP-PTQ.
Conclusione: Ciò conferma un "divario proxy" in cui la minimizzazione dell'errore di ricostruzione locale non garantisce prestazioni migliori nel compito finale a 2 bpw.

2. Llama-3.2-1B (Risultato Positivo: KL End-to-End)

Configurazione: QAT a singolo livello che ottimizza la distillazione KL in avanti end-to-end contro un insegnante FP16.
Esito:
- Sovrastima del Programma: L'uso di un programma ingenuo ad alta- $T$ ( $T_0=1.0$ ) ha portato a prestazioni leggermente peggiori rispetto alla PTQ (+0,005 a +0,022 PPL).
- Correzione Skip-High-T: L'uso del programma corretto ( $T_0=0.3$ ) ha eliminato la sovrastima. Su Layer 4, il modello ha raggiunto 10,1347 PPL, un miglioramento di −0,084 PPL rispetto alla baseline QTIP-PTQ (10,2189).
- Super-Additività: Combinando i livelli addestrati (Layer 4 e Layer 8), il guadagno congiunto (−0,077 PPL) ha superato la somma dei guadagni individuali (−0,062 PPL), dimostrando un "surplus di cooperazione" in cui le parole di codice addestrate con BCJR si rafforzano a vicenda.

3. Analisi Teorica

Teoria del Budget di Deriva: Gli autori derivano una condizione per uscire dal bacino di Voronoi PTQ: la deriva cumulativa del peso ( $\eta \cdot N \cdot g_{max}$ ) deve superare il raggio del bacino ( $r_{Voronoi} \approx 10^{-3}$ ). Gli esperimenti confermano che i programmi che non soddisfano questa condizione non possono uscire dalla soluzione PTQ, mentre quelli che la superano possono trovare soluzioni migliori se l'obiettivo e il programma sono corretti.

Significato e Affermazioni

Il documento afferma tre contributi principali:

Algoritmico: BCJR-QAT è il primo metodo a rendere la quantizzazione a codifica a griglia nativamente differenziabile, abilitando il QAT su TCQ senza perdere i benefici dell'accoppiamento dei percorsi.
Ingegneristico: Il kernel Triton fuso rende ciò computazionalmente fattibile su hardware consumer (singola GPU da 12 GB), rimuovendo la barriera di deployment al momento dell'inferenza per la ricerca.
Insight Empirico:
- Identifica che l'MSE di ricostruzione per livello è un proxy scarso per il QAT a 2 bit, mentre la distillazione KL end-to-end è efficace.
- Mette in discussione la saggezza convenzionale della ricottura simulata per rilassamenti discreti, mostrando che iniziare "caldi" ( $T=1.0$ ) è dannoso per BCJR-QAT perché il gradiente ad alta temperatura è dominato dal rumore. La strategia ottimale è iniziare a una temperatura in cui il gradiente è già informativo.

Il documento conclude che, sebbene i risultati a singolo livello siano promettenti, il pieno potenziale di BCJR-QAT risiede probabilmente nell'addestramento congiunto end-to-end su tutti i livelli, un compito reso accessibile dal kernel rilasciato ma che richiede hardware su scala cloud (H100/H200) per modelli superiori a 1 miliardo di parametri. Gli autori dichiarano esplicitamente di non aver eseguito questo addestramento congiunto su larga scala in questo lavoro.