Riepilogo Tecnico: N-vium – Trasformatore a Misto di Uscite per la Generazione Esatta Accelerata

1. Enunciato del Problema

I Modelli Linguistici di grandi dimensioni (LLM) autoregressivi standard elaborano i token in sequenza attraverso una pila di livelli di decodifica. Questo passaggio in avanti svolge due funzioni accoppiate: (1) la costruzione di rappresentazioni Chiave-Valore (KV) a ogni livello per l'attenzione futura, e (2) la proiezione dello stato nascosto finale per prevedere il token successivo. Di conseguenza, ogni token richiede un passaggio in avanti completo attraverso tutti i livelli, indipendentemente dal fatto che la distribuzione di previsione sia già convergente.

Gli approcci esistenti per migliorare l'efficienza dell'inferenza, come i metodi di Uscita Anticipata (Early Exit), tentano di disaccoppiare queste funzioni attaccando teste di previsione a livelli intermedi. Tuttavia, questi metodi trattano tipicamente le previsioni intermedie come approssimazioni della distribuzione del livello finale. Ciò porta a diversi svantaggi:

• Errori di Approssimazione: Le previsioni intermedie sono spesso più deboli rispetto a quelle a profondità completa, creando un disallineamento tra addestramento e test.
• Instradamento Euristico: Le decisioni su quando uscire spesso si basano su soglie di confidenza create manualmente che richiedono calibrazione.
• Cache KV Incomplete: Per risparmiare calcolo, i metodi di uscita anticipata spesso saltano il calcolo delle cache KV per i livelli bypassati, richiedendo copie o salti approssimativi, il che degrada la qualità.
• Degrado della Qualità: La velocità è spesso ottenuta a scapito della perplessità e della qualità della generazione.

L'articolo sostiene che, sebbene osservazioni empiriche (ad esempio, tramite la "lente dei logit") mostrino che le distribuzioni di previsione spesso si stabilizzano prima del livello finale, i metodi attuali non riescono a sfruttare questo fatto per una generazione esatta senza perdita di qualità.

2. Metodologia: N-vium

Gli autori introducono N-vium, un framework a misto di uscite che generalizza rigorosamente l'architettura standard del trasformatore. Invece di trattare le uscite anticipate come approssimazioni, N-vium ridefinisce l'output del modello come una distribuzione mista appresa su più uscite.

Architettura

• Misto di Uscite: Il modello partiziona i $L$ livelli del backbone in $N_V$ blocchi. Alla fine di ogni blocco (giunzione $k$), un router determina la probabilità di campionare il token successivo dalla distribuzione $\pi_k$ di quel livello.
• Instradamento Adattivo ai Token: Un MLP router elabora lo stato nascosto $h^{(\ell_k)}_t$ alla giunzione $k$ per produrre una probabilità $w_k$. La probabilità $p_k$ di uscire alla giunzione $k$ è calcolata come:
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  dove $w_{N_V} = 1$ per convenzione.
• Distribuzione Mista: La distribuzione finale del token successivo è una somma pesata di tutte le distribuzioni di uscita:
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• Testa Condivisa: Tutte le uscite condividono la testa di proiezione finale per il Modellamento Linguistico (LM) $W_{lm}$, con adapter MLP leggeri che trasformano gli stati nascosti intermedi per corrispondere alla geometria attesa dalla testa.

Addestramento

• Obiettivo: Il modello è addestrato per minimizzare la perdita di entropia incrociata della distribuzione mista congiunta $\pi_{mix}$, piuttosto che ottimizzare le singole teste in isolamento per corrispondere al livello finale.
• Penalità di Calcolo: Viene aggiunta una termine di perdita ausiliario $\mathcal{L}_{compute}$ per incoraggiare il router ad assegnare massa di probabilità alle uscite più anticipate, indirizzando il modello verso l'efficienza:
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  dove $\mathcal{L}_{compute}$ è la profondità del livello normalizzata attesa.
• Riscaldamento: Viene impiegata una fase di riscaldamento del router per prevenire il "crollo del router" (dove tutti i token vengono instradati all'ultimo livello) penalizzando inizialmente le deviazioni da una distribuzione di uscita uniforme.

Inferenza: Piggybacking

Un'innovazione critica è la tecnica Piggybacking per garantire il recupero esatto della cache KV senza saltare il calcolo:

1. Calcolo Differito: Quando un token esce alla giunzione $k$, il calcolo per i livelli da $k+1$ a $L$ viene differito.
2. Batching Parallelo: Man mano che i token futuri vengono elaborati, i loro passaggi in avanti raggiungono i livelli differiti. A queste profondità, gli stati nascosti dei token differiti vengono concatenati con lo stato del token corrente ed elaborati in parallelo.
3. Risultato: Questo riempie gli stati KV mancanti per tutti i token differiti in un singolo passaggio parallelo, garantendo che la cache KV sia completa, esatta e coerente con il token campionato. Nessun calcolo viene saltato; viene semplicemente riordinato da un'esecuzione sequenziale a un'esecuzione in batch parallela.

3. Contributi Chiave

1. Campionamento Misto Esatto: N-vium è il primo metodo a abilitare l'instradamento adattivo alla profondità per token con campionamento esatto dalla distribuzione addestrata, evitando gli errori di approssimazione intrinseci nei metodi standard di uscita anticipata.
2. Recupero Completo della KV: Differendo e raggruppando i calcoli dei livelli superiori, N-vium recupera cache KV complete ed esatte, una caratteristica spesso persa in altri metodi di calcolo adattivo.
3. Framework di Addestramento: Gli autori forniscono una pipeline di addestramento pratica che ottimizza direttamente l'obiettivo misto, includendo strategie per indirizzare i modelli verso l'efficienza tramite una penalità di calcolo apprendibile.
4. Scalabilità: Il metodo è validato su scale fino a 1,5 miliardi di parametri, dimostrando che l'approccio non è limitato a modelli piccoli.

4. Risultati Sperimentali

Gli autori hanno pre-addestrato modelli N-vium (in particolare "Quadrivium" con 4 uscite) sul dataset C4 utilizzando l'architettura LLaMA 2, confrontandoli con baseline dense abbinate per parametri e dati.

• Velocità vs. Perplessità:
  • Il modello più grande (1,5 miliardi di parametri, 48 livelli, larghezza nascosta 1536) ha raggiunto un accelerazione del 57,9% nel tempo reale rispetto alla baseline densa con nessun degrado della perplessità (differenza dello 0%).
  • Modelli più piccoli hanno mostrato accelerazioni costanti comprese tra il 23% e il 41% con lievi miglioramenti della perplessità (ad esempio, un modello a 36 livelli ha mostrato un'accelerazione del 36,4% con un miglioramento della perplessità del 3,7%).
• Frontiera di Pareto: Nei grafici qualità-velocità, N-vium occupa il quadrante "in alto a destra" (più veloce e qualità uguale o migliore), mentre metodi come CALM (riimplementato come CALM*) cadono nel "basso a destra" (più veloce ma qualità peggiore) e LayerSkip giace su una linea di perplessità costante al di sotto della baseline.
• Prestazioni a valle: Dopo il fine-tuning supervisionato (SFT) sul mix Tulu 3, i modelli N-vium hanno eguagliato o superato le baseline dense sui benchmark a valle (MMLU, ARC, HellaSwag, ecc.), confermando che i modelli rimangono indirizzabili e non si sovradattano all'obiettivo di entropia incrociata a scapito delle prestazioni del compito.
• Robustezza Hardware: Le misurazioni dell'accelerazione su GPU A100, RTX A6000 e RTX 3090 hanno prodotto risultati coerenti (~57%), mentre le GPU più vecchie con architettura Turing hanno mostrato guadagni inferiori a causa della mancanza di supporto nativo per bfloat16.

5. Significato e Affermazioni

L'articolo afferma che N-vium rappresenta un cambiamento di paradigma nell'efficienza dell'inferenza. A differenza della maggior parte dei metodi che migliorano i token al secondo riducendo i FLOP per token (spesso tramite approssimazione), N-vium migliora l'efficienza aumentando i FLOP al secondo. Lo ottiene riorganizzando il calcolo sequenziale in passaggi in batch paralleli, sfruttando così meglio il parallelismo hardware senza scartare alcun calcolo.

Gli autori sottolineano che N-vium è una generalizzazione rigorosa del trasformatore standard: quando i pesi di instradamento per le uscite intermedie sono zero, il modello si comporta esattamente come un trasformatore standard. Il metodo soddisfa simultaneamente quattro criteri critici, che gli autori affermano nessun altro metodo esistente soddisfi:

1. Inferenza Esatta: Campionamento dalla distribuzione addestrata senza approssimazione.
2. Recupero Completo della Cache KV: Nessun livello saltato o riempimento approssimativo della cache.
3. Instradamento Appreso: L'allocazione del calcolo per token è gestita da un router addestrato, non da euristiche.
4. Accelerazione al Decodifica: Riduce la latenza nel tempo reale durante la generazione autoregressiva.

L'articolo conclude che, sebbene l'addestramento di N-vium comporti un costo di FLOP più elevato per token (a causa della valutazione delle teste a più giunzioni), questo sovraccarico è giustificato per il deployment in produzione dove i costi di inferenza dominano, e il sovraccarico relativo diminuisce man mano che le dimensioni del modello aumentano.