Riepilogo Tecnico: Modelli a Spazio di Stato a Campo di Densità (DF-SSM)

1. Problema

I Large Language Models (LLM), in particolare quelli con 1,3 miliardi di parametri in precisione FP16, richiedono circa 2,7 GB di memoria, superando la capacità della maggior parte dei dispositivi mobili, embedded ed edge. Sebbene la quantizzazione possa ridurre le dimensioni del modello, la quantizzazione estrema (1–4 bit) comporta tipicamente una significativa degradazione della qualità o richiede l'addestramento da zero su dataset massicci (ad es. 100B+ token). Inoltre, sebbene i Modelli a Spazio di Stato (SSM) come Mamba-2 offrano vantaggi architettonici per l'implementazione on-device (stati nascosti a dimensione fissa che eliminano la crescita lineare della KV cache), la loro compressione dei pesi rimane meno esplorata rispetto alla quantizzazione dei Transformer.

2. Metodologia: Pipeline DF-SSM

Il documento introduce i Modelli a Spazio di Stato a Campo di Densità (DF-SSM), un framework progettato per comprimere gli SSM in uno scaffold a 1 bit con correzione low-rank in int8. L'approccio consiste in tre fasi primarie:

A. Addestramento del Peso a Campo di Densità (DFW)

Gli autori impiegano la distillazione consapevole della quantizzazione (quantization-aware distillation) da un insegnante pre-addestrato in FP16 (Mamba-2 1.3B).

• Quantizzazione a 17 livelli: I pesi sono parametrizzati come valori latenti continui mappati in 17 livelli discreti tra -1 e +1 tramite un clamp lineare.
• Annealing della Quantizzazione: Un parametro di annealing $\alpha$ interpola tra pesi continui e quantizzati, aumentando linearmente da 0 a 1 durante l'addestramento. Questo evita un "salto qualitativo catastrofico" permettendo al modello di apprendere l'obiettivo di linguaggio con pesi continui prima di adattarsi ai vincoli discreti.
• Gestione del Gradiente: Viene utilizzata la stima straight-through per i gradienti. Gli autori rifiutano la mappatura sigmoidale a favore del clamping lineare per evitare l'attenuazione del gradiente ai valori estremi.

B. Scaffold Congelato con Correzione LoRA

Post-addestramento, i pesi dello scaffold quantizzato sono congelati ai loro valori deterministici arrotondati.

• Adattamento a Basso Rango (LoRA): Un modulo LoRA piccolo e addestrabile in int8 ($W_A \in \mathbb{R}^{d_{out} \times r}, W_B \in \mathbb{R}^{r \times d_{in}}$ con rango $r=16$) viene aggiunto per correggere gli errori residui.
• Architettura: LoRA è applicato a tutti i 48 layer (sia in_proj che out_proj).
• Meccanismo: La correzione LoRA non si limita a correggere gli errori di arrotondamento per singolo peso (la correlazione con il residuo di arrotondamento è ~0), bensì corregge gli errori di output pesati dall'input.
• Quantizzazione: Entrambe le matrici LoRA e le matrici di embedding/LM head sono quantizzate in int8 con scaling per canale/riga, ottenendo una perplexity (PPL) effettivamente lossless (PPL 49.1 vs 49.2 per FP16).

C. Pipeline di Inferenza Ottimizzata

Gli autori sviluppano un motore di inferenza personalizzato sia per GPU che per CPU:

• GPU (CUDA):
  • Scaffold Matmul: I pesi impacchettati a 1 bit vengono espansi on-the-fly in int8 e processati tramite cuBLAS INT8 tensor cores (usando torch._int_mm), sfruttando una riduzione dell'8x nel trasferimento di memoria e un throughput INT8 più elevato.
  • Operazioni Stateful: Kernel CUDA personalizzati gestiscono lo shift register a stato (convoluzione) e la ricorrenza SSM (riduzione a livello di warp).
  • CUDA Graphs: L'intero forward pass di 48 layer viene catturato come un grafo CUDA per eliminare l'overhead di lancio dei kernel, fondamentale per le prestazioni con batch size 1.
• CPU (AVX-512):
  • Un kernel bit-packed espande i pesi in int8 per le istruzioni VNNI (_mm512_dpbusd_epi32), ottenendo un alto throughput con 8 volte meno dati caricati rispetto ai pesi int8 standard.

D. Rappresentazione dello Stato a Campo di Densità

Per comprimere lo stato nascosto dell'SSM, gli autori mantengono un campo binario $K \times K$ per ogni elemento di stato ($K^2 \approx 4 \times d_{state}$). Il ciclo prevede normalizzazione, iniezione sigma-delta (feedback dell'errore), binarizzazione e lettura popcount per recuperare i valori continui, applicato ai confini dei chunk.

3. Risultati Chiave

Efficienza e Dimensione del Modello

• Compressione: Il modello risultante è di 278 MB, una riduzione di 9,7× rispetto al docente FP16 da 2,7 GB.
• Composizione: 56% Scaffold (1-bit), 37% Embedding/LM Head (int8), 4% LoRA (int8), 3% Altro (FP16/FP32).
• Costo di Addestramento: La distillazione richiede solo 32M di token e 6 ore su una singola GPU A100, rispetto ai 150B di token necessari per addestrare BitMamba-2 da zero.

Prestazioni di Inferenza

• Accelerazione GPU: Raggiunge un'inferenza 21,4× più veloce su GPU (batch=1) rispetto all'implementazione di riferimento ottimizzata mamba-ssm. Con batch=512, il throughput si avvicina al limite di banda HBM dell'A100 (~5.081 tok/s).
• Accelerazione CPU: Raggiunge un incremento di velocità di 2,9× su CPU (Intel Xeon, AVX-512) con un uso di RAM di 567 MB.

Prestazioni sui Task Downstream

DF-SSM mantiene le prestazioni entro 2–4 punti percentuali di BitMamba-2 (un modello a 1,58-bit addestrato da zero su 150B di token) e conserva il 78–95% delle prestazioni del docente FP16:

• BoolQ: 60,8% (Docente: 64,2%)
• PIQA: 67,1% (Docente: 73,2%)
• WinoGrande: 54,7% (Docente: 60,9%)
• ARC-easy: 50,2% (Docente: 64,1%)
• HellaSwag: 41,4% (Docente: 59,9%) – Segnalato come l'outlier, probabilmente dovuto a una quantità insufficiente di token di distillazione per il ragionamento multi-step.

4. Risultati di Interpretabilità

Il documento presenta la prima analisi sistematica dell'organizzazione della conoscenza fattuale negli SSM, rivelando tre fasi distinte di elaborazione:

1. Classificazione dell'Intento (Layer 0–3): Opera in uno spazio astratto senza allineamento al vocabolario. La classificazione è sintattica (guidata dalla struttura del template, ad es. "La capitale di...") piuttosto che semantica. Le categorie con template identici si raggruppano immediatamente; le categorie con template misti no.
2. Recupero della Conoscenza (Layer 25–35): Le associazioni fattuali si localizzano in una finestra di 5 layer (con picco a L32–L36). Esperimenti di intervento causale mostrano che la sostituzione degli stati nascosti in questa finestra inverte le predizioni fattuali. La conoscenza occupa un sottospazio a bassa dimensione (10–16 dimensioni) dello stato da 2048 dimensioni.
3. Formattazione dell'Output (Layer 36–47): La struttura delle categorie si dissolve mentre il modello converge verso rappresentazioni di output condivise, preparandosi a emettere i token.

Osservazione Chiave: Il modello esibisce strutture di conoscenza interna ben organizzate (clustering pulito, localizzazione causale) nonostante una debole forza fattuale (le risposte corrette raggiungono raramente il rank top-1). Ciò suggerisce che la struttura rappresentativa può precedere la forza fattuale.

5. Significato e Rivendicazioni

Il documento rivendica i seguenti contributi:

• Efficienza: Una pipeline completa per comprimere un SSM da 1,3B in un modello binario implementabile in 6 ore con una quantità minima di dati di distillazione (32M di token).
• Prestazioni: Un'accelerazione dell'inferenza di 21,4× su GPU e 2,9× su CPU mantenendo prestazioni vicine al docente, guidata principalmente dall'ottimizzazione della larghezza di banda della memoria e dall'utilizzo dei tensor core INT8.
• Architettura: Un nuovo modello a precisione mista (scaffold a 1-bit + LoRA in int8) che evita i cali qualitativi dei modelli a quantizzazione uniforme estrema.
• Interpretabilità: La scoperta di un modello di elaborazione a tre fasi negli SSM e il fatto che la classificazione nei primi layer sia sintattica e guidata dai template, operando in uno spazio non allineato al vocabolario.
• Implicazione: I risultati suggeriscono che la distillazione può trasferire efficientemente la struttura organizzativa delle rappresentazioni della conoscenza anche quando il contenuto fattuale è solo parzialmente mantenuto, implicando che ulteriori token di addestramento potrebbero riempire le strutture esistenti piuttosto che riorganizzarle.

Limitazioni Ammesse:

• Il degrado della PPL (49,2 vs 14,3) indica che c'è margine di miglioramento tramite un addestramento più lungo o ranghi LoRA più elevati.
• I risultati di interpretabilità sono specifici per un singolo modello distillato (Mamba-2 1,3B); la generalizzazione ad altre scale o architetture addestrate da zero non è verificata.
• L'etichetta "1-bit" si applica allo scaffold (56% dei pesi); gli embedding e il LoRA rimangono in int8, in linea con altri metodi SSM binari.