pQuant: Towards Effective Low-Bit Language Models via Decoupled Linear Quantization-Aware Training

1. Il Problema: Democratizzazione dei Parametri e Limiti della Quantizzazione Estrema

Il lavoro affronta le sfide poste dalla quantizzazione dei Large Language Models (LLM) a bit estremamente bassi (sotto i 2 bit), un approccio cruciale per il deployment su dispositivi edge dove memoria e potenza di calcolo sono limitate.

Sebbene metodi recenti di Quantization-Aware Training (QAT) addestrati da zero (from scratch) abbiano mostrato progressi, gli autori identificano un collo di bottiglia fondamentale: l'effetto di "democratizzazione dei parametri" (parameter democratization).

• Fenomeno: Nei modelli 1-bit esistenti (come BitNet), la sensibilità dei parametri (la loro influenza sull'output del modello) diventa omogenea. Invece di avere un piccolo sottoinsieme di parametri "sensibili" che guidano le prestazioni, tutti i parametri vengono trattati con importanza simile.
• Conseguenza: Questa omogeneizzazione limita drasticamente l'espressività del modello e la sua scalabilità. Man mano che le dimensioni del modello aumentano, i guadagni di prestazioni sono sub-lineari e non riescono a tenere il passo con i modelli a precisione completa (FP16) o anche solo a 2 bit.
• Obiettivo: Recuperare una struttura di sensibilità differenziata, preservando i parametri critici ad alta precisione mentre si mantiene l'efficienza computazionale del resto.

2. Metodologia: pQuant e la Quantizzazione Decoupled

Per contrastare la democratizzazione dei parametri, gli autori propongono pQuant, un metodo di QAT da zero basato su una quantizzazione lineare decoupled (disaccoppiata).

Architettura Principale

pQuant sostituisce i livelli lineari standard con una struttura a due rami paralleli:

1. Ramo Dominante (1-bit): La maggior parte dei parametri viene quantizzata a 1 bit (valori +1/-1) per garantire l'efficienza computazionale massima. Questo ramo funge da "esperto condiviso" per tutti gli input.
2. Ramo Compatto ad Alta Precisione (8-bit): Un piccolo sottoinsieme di parametri (circa il 4-5% del totale) viene mantenuto in precisione 8-bit (INT8). Questo ramo è dedicato esclusivamente alla preservazione dei parametri più sensibili e informativi.

Meccanismi Chiave

• Decoupling Strutturale: A differenza dei metodi che assegnano posizioni fisse per i parametri ad alta precisione, pQuant divide strutturalmente il livello lineare. Nel Feed-Forward Network (FFN), i due rami elaborano lo stesso input e i loro output vengono sommati.
• Feature Scaling (Scalatura delle Caratteristiche): Per guidare il modello ad allocare correttamente i parametri sensibili, vengono introdotti due scalari apprendibili ($\alpha$ e $\beta$) che moltiplicano rispettivamente l'output del ramo 8-bit e del ramo 1-bit.
  • Inizialmente, $\alpha$ è impostato molto più alto di $\beta$ (es. $\alpha=2.0, \beta=0.2$).
  • Questo forza il modello a dare priorità al segnale ad alta precisione durante l'ottimizzazione, permettendo ai parametri sensibili di essere identificati e preservati dinamicamente.
• Scalabilità Efficiente (MoE): Il ramo ad alta precisione può essere esteso a più rami paralleli ($N$ esperti) attivati sparsamente. Un router leggero seleziona il ramo più adatto per ogni token (strategia Top-1). Questo permette di aumentare la capacità del modello senza aumentare linearmente il costo di inferenza, mantenendo il numero di parametri attivi costante.

Addestramento

• QAT da Zero: Il modello viene addestrato da inizializzazione casuale, non tramite fine-tuning di un modello pre-addestrato.
• Stabilità: Durante l'addestramento, i pesi FP16 sono mantenuti per il calcolo dei gradienti (usando l'estimatore Straight-Through, STE), mentre la forward pass utilizza i pesi quantizzati. I pesi FP16 vengono scartati dopo l'addestramento.
• Schedule a Due Fasi: Viene utilizzato uno schema di apprendimento con due fasi (learning rate e weight decay decrescenti) per stabilizzare l'ottimizzazione dei parametri 1-bit.

3. Risultati Sperimentali

Gli autori hanno valutato pQuant su modelli di diverse dimensioni (da 300M a 2.6B parametri) confrontandoli con BitNet (1-bit), BitNet1.58 (2-bit) e LLaMA-2 (FP16).

• Prestazioni Superiori:
  • pQuant riduce l'incertezza (Perplexity) del 32% rispetto ai migliori baseline 1-bit.
  • A scala 1.3B, pQuant (1.35-bit) supera BitNet (1-bit) e si avvicina notevolmente a BitNet1.58 (2-bit), pur utilizzando meno bit per peso.
  • Il modello pQuant da 2.6B supera in accuratezza media il modello LLaMA-2 da 1.3B (FP16) e supera significativamente i metodi di quantizzazione post-training (PTQ) applicati a modelli 7B.
• Scalabilità:
  • Aumentando il numero di rami ad alta precisione ($N$), pQuant segue una legge di scalatura molto più vicina a quella dei modelli FP16 rispetto ai modelli 1-bit puri.
  • Con $N=8$, pQuant raggiunge prestazioni quasi pari a LLaMA-2 FP16.
• Efficienza di Memoria e Inferenza:
  • pQuant riduce l'impronta di memoria del 92% rispetto a LLaMA-2 FP16 e del 31% rispetto a BitNet1.58.
  • Grazie all'attivazione sparsa e all'uso di operazioni bitwise, pQuant è 1.6 volte più veloce di BitNet1.58 a parità di budget di parametri attivi e 82% più veloce di LLaMA-2 FP16.
• Analisi della Sensibilità:
  • Le analisi confermano che pQuant ripristina una distribuzione di sensibilità differenziata, con picchi di sensibilità chiaramente identificati nel ramo 8-bit, risolvendo il problema della democratizzazione osservato in BitNet.

4. Contributi Chiave

1. Identificazione del Collo di Bottiglia: Dimostrazione empirica che la "democratizzazione dei parametri" (omogeneizzazione della sensibilità) è il limite principale dei modelli 1-bit attuali.
2. Architettura Decoupled: Progettazione di un livello lineare ibrido (1-bit + 8-bit) che separa strutturalmente l'efficienza dalla capacità espressiva.
3. Guida Dinamica: Introduzione dello Feature Scaling per guidare l'allocazione dinamica dei parametri sensibili verso il ramo ad alta precisione, senza bisogno di maschere predefinite.
4. Scalabilità MoE: Estensione del ramo ad alta precisione in un modulo Mixture of Experts (MoE) sparsamente attivato, permettendo di scalare la capacità del modello mantenendo costi di inferenza costanti.

5. Significato e Impatto

Il lavoro pQuant rappresenta un passo avanti significativo verso l'effettiva utilizzazione pratica degli LLM su dispositivi edge.

• Superamento dei Limiti Attuali: Dimostra che è possibile ottenere prestazioni vicine ai modelli FP16 utilizzando pesi mediamente inferiori a 1.5 bit, superando i limiti di accuratezza dei metodi precedenti.
• Efficienza Hardware: La riduzione drastica della larghezza di banda della memoria e l'uso di operazioni bitwise rendono pQuant ideale per l'esecuzione su CPU e GPU di consumo, dove la memoria è spesso il collo di bottiglia principale.
• Nuova Direttiva di Ricerca: Suggerisce che la chiave per la quantizzazione estrema non è solo comprimere i pesi, ma preservare strategicamente la struttura di sensibilità del modello attraverso architetture ibride e dinamiche.

In sintesi, pQuant offre una soluzione robusta per addestrare modelli linguistici estremamente efficienti senza sacrificare le prestazioni, aprendo la strada a LLM potenti ed economici per il deployment su larga scala.