From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Riepilogo Tecnico: Da Livelli a Sotto-moduli (SUBFIT)

Definizione del Problema
La compressione post-training dei Large Language Models (LLM) mira a ridurre i costi di deployment rimuovendo componenti architettonici. Sebbene i metodi che eliminano interi livelli o blocchi contigui siano efficaci, essi interrompono lo stream residuo, causando spesso un brusco degrado della qualità. I recenti approcci basati sulla "sostituzione" affrontano questo problema sostituendo i componenti rimossi con moduli leggeri e adattati. Tuttavia, i metodi esistenti condividono due vincoli progettuali restrittivi:

1. Granularità a livello di intero livello: Trattano l'intero livello Transformer come l'unità atomica di rimozione.
2. Selezione contigua: Rimuovono i componenti in intervalli di profondità consecutivi.

Gli autori sostengono che questi vincoli siano eccessivamente restrittivi. La ridondanza nei transformer pre-addestrati non è confinata in regioni contigue, né è distribuita uniformemente tra gli output di Attention e FeedForward (FFN). Di conseguenza, sono necessarie strategie diverse per approssimare diversi tipi di sotto-moduli, e i componenti rimovibili non devono necessariamente raggrupparsi in intervalli di profondità contigui.

Metodologia: SUBFIT
Il documento introduce SUBFIT (Submodule-level Fitted residual replacement), un metodo di compressione post-training che opera alla granularità di singoli sotto-moduli di Attention e FFN. Il processo prevede due fasi principali:

1. Selezione Non Contigua:

   • SUBFIT seleziona un numero target di sotto-moduli di Attention e FFN indipendentemente attraverso la profondità del modello basandosi su un budget di sparsità $s$.
   • La selezione è non contigua; i moduli sono scelti indipendentemente tramite punteggio senza formare un blocco consecutivo.
   • Strategia di Punteggio:
     • Attention: I moduli sono classificati tramite un "Impact score" calcolato sulle attivazioni di calibrazione, che misura quanto un sotto-modulo ruota lo stream residuo e la sua magnitudo relativa.
     • FFN: Per i modelli base, viene utilizzato un punteggio consapevole della sostituzione che combina gli effetti di rimozione e l'errore surrogato. Per i modelli istruiti (instruction-tuned), viene utilizzato un punteggio basato solo sul coseno per evitare il rumore derivante dalle magnitudo residue scarsamente calibrate.
   • Ordinamento: I sotto-moduli di Attention vengono selezionati e sostituiti per primi. I sotto-moduli FFN vengono quindi valutati e selezionati sul modello parzialmente compresso per tenere conto dello spostamento dello stream residuo indotto dalla sostituzione dell'Attention.

2. Sostituzione del Residuo Adattata (Fitted Residual Replacement):

   • Ogni sotto-modulo selezionato è sostituito da un surrogato parametrico leggero $\hat{F}_\ell$ che approssima il contributo residuo originale del sotto-modulo ($y = F_\ell(x)$) piuttosto che il suo calcolo interno.
   • Struttura del Surrogato: Il surrogato utilizza una correzione affine (gate elemento per elemento, bias e centratura) più un termine di interazione strutturato definito da un operatore a basso rango $U_\ell = V_\ell^\top W_\ell$.
   • Parametrizzazione Asimmetrica:
     • Attention: Ogni sotto-modulo di Attention rimosso riceve la propria base a basso rango ($V_\ell$) e i propri parametri indipendenti, riflettendo ruoli eterogenei attraverso le profondità.
     • FFN: Una base di input condivisa $V$ è utilizzata attraverso tutti i livelli FFN selezionati. Solo i fattori di output ($W_\ell$) e i termini affini rimangono specifici per ogni livello. Questa condivisione dei parametri tra i livelli riduce significativamente il costo di deployment.
   • Fitting: I parametri sono recuperati analiticamente tramite fitting in forma chiusa sulle statistiche di calibrazione (minimi quadrati regolarizzati con ridge ed eigen-decomposition), senza richiedere ottimizzazione basata su gradienti o fine-tuning dell'intero modello.

Contributi Chiave

1. Identificazione dei Vincoli: Gli autori identificano la granularità a livello di intero livello e la selezione contigua come fattori limitanti nella compressione basata sulla sostituzione esistente.
2. Framework SUBFIT: Un nuovo metodo che esegue la selezione e la sostituzione congiunta a livello di sotto-modulo (Attention e FFN) con selezione non contigua.
3. Design Asimmetrico: Un'architettura specifica in cui i bypass di Attention utilizzano basi indipendenti, mentre i bypass FFN condividono una base di input comune per minimizzare l'overhead dei parametri.
4. Valutazione Completa: Test estesi su 10 LLM (5 base, 5 istruiti) e 5 livelli di sparsità (dal 12,5% al 37,5%), dimostrando trade-off migliorati rispetto a quattro baseline (Streamline, ReplaceMe).

Risultati Sperimentali

• Perplexity (Modelli Base): SUBFIT raggiunge il minor degrado di perplexity in 9 scenari su 10 di densità/sparsità del modello. Il divario rispetto alla baseline più forte si amplia significativamente sotto compressione aggressiva (ad esempio, al 25% di sparsità, SUBFIT mostra un degrado di perplexity di 2,42× rispetto a 4,34× della migliore baseline).
• Accuratezza sui Task (Modelli Istruiti): Al 25% di sparsità, SUBFIT mantiene il 84,6% dell'accuratezza downstream densa in media, superando la baseline più forte (81,6%). È l'unico metodo a mantenere oltre l'80% di accuratezza al 25% di sparsità e sopra il 73% al 37,5%.
• Stabilità: SUBFIT esibisce una varianza significativamente inferiore nel degrado della perplexity rispetto alle baseline, che mostrano instabilità (ad esempio, la deviazione standard cresce di ordini di grandezza) ad alta sparsità.
• Efficienza: Il metodo produce miglioramenti misurabili nella velocità di inferenza (1,18×–1,40× nel time-to-first-token) e risparmi nella KV-cache (proporzionali alla sparsità) senza richiedere kernel sparsi o hardware personalizzato.
• Studi di Ablazione:
  • Compensazione Residua: La sostituzione degli FFN è il fattore dominante per la stabilità; sostituire sia l'Attention che l'FFN produce i risultati migliori.
  • Base Condivisa: La base FFF condivisa raggiunge prestazioni comparabili a basi per livello, ma con significativamente meno parametri.
  • Sensibilità al Rango: La qualità FFN migliora fino al rango 4096, mentre l'Attention satura al rango 256, giustificando le impostazioni di rango asimmetriche.

Significato e Rivendicazioni
Il documento afferma che la granularità è un asse poco esplorato nella compressione post-training. Passando dalla sostituzione a livello di livello a quella a livello di sotto-modulo, SUBFIT recupera più comportamento del modello denso rispetto alla sostituzione contigua di interi livelli a parità di budget di parametri. I risultati suggeriscono che il recupero non contiguo e specifico per sotto-modulo fornisce un trade-off aggregato qualità-efficienza più forte, particolarmente sotto compressione aggressiva. Il metodo è progettato per essere economico da implementare (fitting in forma chiusa) ed è compatibile con gli stack di inferenza esistenti. Gli autori suggeriscono con modestia che questo approccio possa generalizzarsi ad altre unità di rimozione, come le teste di attenzione o i singoli neuroni, ma non rivendicano una generalizzazione immediata a tali unità specifiche in questo lavoro.