Scalable LLM Reasoning Acceleration with Low-rank Distillation

1. Il Problema

I Large Language Model (LLM) stanno diventando sempre più capaci in compiti di ragionamento complesso, come la risoluzione di problemi matematici, spesso richiedendo generazioni lunghe basate su catene di pensiero (Chain-of-Thought, CoT). Tuttavia, queste generazioni lunghe comportano:

• Alta latenza e costi computazionali: Il processo di decodifica autoregressiva è computazionalmente oneroso.
• Limiti delle attuali tecniche di efficienza: Esistono metodi di inferenza efficiente basati sulla sparsità (come GRIFFIN e CATS) che riducono i calcoli nei blocchi Feed-Forward (FF) mantenendo buone prestazioni sui task linguistici standard. Tuttavia, questi metodi causano un crollo drastico delle prestazioni nei task di ragionamento (es. matematica).
  • Esempio: Applicare GRIFFIN a DeepSeek-R1-Distill-Qwen 1.5B riduce l'accuratezza su MATH-500 dal 79.40% al 42.00%.
  • Causa: Gli errori di approssimazione introdotti dalla sparsità si accumulano durante le generazioni lunghe, deviando il percorso di ragionamento.

L'obiettivo è sviluppare un metodo di inferenza efficiente che preservi sia le capacità linguistiche che quelle di ragionamento, riducendo al contempo la latenza e il numero di parametri attivi.

2. Metodologia: Caprese

Il paper propone Caprese (CAPability REcovery with Scalable Efficiency), un metodo di distillazione efficiente focalizzato sui blocchi Feed-Forward (FF) degli LLM.

Concetto Chiave

Gli autori osservano che, sebbene i metodi di compressione FF (come la selezione top-k) introducano errori, i residui (la differenza tra l'output originale e quello compresso) possiedono una struttura a basso rango. Aggiungere un piccolo strato lineare a basso rango per approssimare questi residui permette di recuperare le informazioni perse senza ripristinare l'intero calcolo.

Architettura e Procedura

1. Distillazione dei Residui:

   • Si parte da un modello efficiente compresso $\tilde{FF}(x)$.
   • Si definisce l'errore come $\|FF(x) - \tilde{FF}(x)\|^2$.
   • Si introduce un layer lineare a basso rango $xLR$ (dove $L \in \mathbb{R}^{D \times r}$ e $R \in \mathbb{R}^{r \times D}$, con $r \ll D_{FF}$) per modellare questo residuo.
   • L'obiettivo è minimizzare: $\min_{L,R} \sum \|FF(x) - \tilde{FF}(x) - xLR\|^2$.

2. Due Fasi di Addestramento:

   • Distillazione Layer-wise: Ogni strato FF viene addestrato indipendentemente per apprendere i propri residui. I pesi originali del modello sono congelati.
   • Distillazione End-to-End (E2E): I layer a basso rango appresi vengono utilizzati come inizializzazione per un addestramento congiunto su tutto il modello (escluso l'ultimo layer lineare), sempre con pesi originali congelati.

3. Efficienza Computazionale:

   • Il calcolo $xLR$ può essere eseguito in parallelo alle operazioni FF originali.
   • I parametri $L$ e $R$ possono essere concatenati alle matrici di proiezione "up" e "down" esistenti, permettendo un'integrazione pulita senza overhead significativo di latenza.
   • Costo: Aggiunge solo circa l'1% di parametri aggiuntivi (con $r=256$) e richiede solo 20.000 campioni sintetici di matematica per l'addestramento.

4. Rieselezione (Reselection):

   • Per metodi come GRIFFIN, che selezionano neuroni basandosi su statistiche di prefill, il paper introduce un meccanismo di rieselezione dinamica durante la generazione. Aggiornando le metriche di selezione man mano che i token vengono generati, si riduce ulteriormente l'accumulo di errore.

3. Contributi Chiave

• Recupero delle Capacità di Ragionamento: Caprese riesce a recuperare gran parte, se non tutta, la perdita di prestazioni nei task matematici causata da metodi di inferenza efficiente, senza danneggiare le capacità linguistiche.
• Efficienza Estrema: Riduce drasticamente i parametri attivi (circa 2 miliardi in meno per modelli come Gemma 2 9B e Llama 3.1 8B) mantenendo una latenza ridotta.
• Brevità delle Risposte: In modo inaspettato, Caprese tende a generare risposte più concise (fino all'8.5% di token in meno) rispetto ai modelli completi o ad altri metodi compressi, migliorando ulteriormente l'efficienza.
• Compatibilità: Funziona con diverse tecniche di compressione (GRIFFIN, CATS) e si integra con tecniche di scaling al test-time (generazione multipla o CoT più lunghi).

4. Risultati Sperimentali

I risultati sono stati testati su vari modelli (Llama 3, Gemma 2, DeepSeek-R1-Distill-Qwen) e task (MATH-500, GSM8K, AIME, AMC, ecc.).

• Recupero delle Prestazioni Matematiche:
  • Su DeepSeek-R1-Distill-Qwen 7B, l'applicazione di CATS riduce l'accuratezza su AMC 2023 dal 75.00% al 62.50%. Con Caprese (E2E), l'accuratezza sale al 78.13%, superando persino il modello originale completo.
  • Su DS-Qwen 1.5B, Caprese recupera l'accuratezza da 42.00% (GRIFFIN) a 60.40%, avvicinandosi al modello originale (79.40%).
• Preservazione delle Prestazioni Linguistiche:
  • Su task di linguaggio (CoQA, QASPER, XSum, CNN/DailyMail), Caprese mantiene le prestazioni invariate o leggermente migliorate, dimostrando che la distillazione sui task matematici non danneggia la generalità del modello.
• Efficienza e Latenza:
  • Su Qwen 2.5 14B, Caprese riduce la latenza end-to-end e il tempo per token (Time-to-Next-Token) di oltre il 16% rispetto al modello completo.
  • Confrontato con l'uso di LoRA su modelli compressi, Caprese è significativamente più efficiente: LoRA annulla oltre il 40% dei risparmi di tempo ottenuti dalla compressione, mentre Caprese mantiene quasi intatto il guadagno.
• Scaling al Test-Time:
  • Con la tecnica "Best-of-N" (sampling multipla), Caprese (GRIFFIN) su Llama 3.2 3B Instruct migliora il Pass@100 del 7.0% rispetto al modello originale, utilizzando il 15.8% in meno di parametri attivi.

5. Significato e Impatto

Il lavoro di Caprese rappresenta un passo fondamentale verso l'inferenza efficiente di LLM senza degradazione delle prestazioni.

• Risoluzione del Trade-off: Dimostra che è possibile ottenere l'efficienza della sparsità senza sacrificare la complessità del ragionamento, un problema che ha finora limitato l'adozione di metodi efficienti per modelli "ragionanti" (thinking models).
• Scalabilità: Il metodo è scalabile, economico da addestrare (pochi dati sintetici) e si integra facilmente nell'architettura esistente.
• Futuro: Apre la strada all'uso massiccio di modelli efficienti per compiti complessi, riducendo i costi di deployment e permettendo una scalabilità più sostenibile delle capacità di ragionamento degli LLM.

In sintesi, Caprese offre una soluzione elegante che combina compressione sparsa e distillazione a basso rango per ottenere il "meglio dei due mondi": velocità e intelligenza.