GradPruner: Gradient-Guided Layer Pruning Enabling Efficient Fine-Tuning and Inference for LLMs

Riepilogo Tecnico: GradPruner

Definizione del Problema

Il fine-tuning di Large Language Models (LLM) su dati downstream è spesso computazionalmente costoso e dispendioso in termini di tempo. Sebbene esistano metodi di pruning strutturato per migliorare l'efficienza dell'inferenza, essi richiedono tipicamente un ulteriore addestramento, distillazione della conoscenza o ricerche di strutture complesse, il che mina l'obiettivo di un fine-tuning efficiente. Inoltre, gli approcci di pruning esistenti spesso si affidano a dati di calibrazione per valutare l'importanza dei parametri; tuttavia, gli LLM mostrano frequentemente prestazioni subottimali su dati specifici del dominio, portando a valutazioni di importanza distorte.

Gli autori identificano tre sfide chiave:

1. Sviluppare un metodo per misurare l'importanza dei parametri adattato a specifici dati downstream senza aumentare la memoria o il tempo di addestramento.
2. Massimizzare il pruning dei parametri preservando il più possibile la struttura originale del modello.
3. Supportare sia il fine-tuning completo (FFT) che i metodi di Parameter-Efficient Fine-Tuning (PEFT) come LoRA.

Metodologia: GradPruner

GradPruner è un framework di pruning dei layer guidato dal gradiente, progettato per migliorare l'efficienza sia dell'addestramento che dell'inferenza. Il metodo opera in tre fasi primarie:

1. Valutazione dell'Importanza dei Parametri (Matrice IGIA)

Ispirandosi all'osservazione che la perdita diminuisce bruscamente durante i primi step del fine-tuning (indicando un apprendimento rapido dei task downstream), GradPruner utilizza la Matrice di Accumulo delle Informazioni del Gradiente Iniziale (IGIA-Matrix).

• Processo: Invece del training completo, il modello subisce una breve fase di fine-tuning LoRA (circa l'1% dei totali step di addestramento).
• Simulazione del Gradiente: Per un layer lineare con peso $W$ e adapter LoRA $W_A$ e $W_B$, il metodo simula il gradiente di $W$ calcolando il prodotto dei gradienti LoRA: $\nabla W_{sim} = \nabla W_B \cdot \nabla W_A$.
• Accumulo: La IGIA-Matrix ($F_W$) è calcolata mediando i quadrati dei gradienti simulati sugli $t$ step iniziali:
  $$F_W = \frac{1}{t} \sum_{i=1}^{t} (\nabla W_{sim}^{(i)})^2$$
• Scoring del Layer: Il punteggio di importanza per un determinato layer $j$ è derivato sommando i valori della IGIA-Matrix di tutti i layer lineari all'interno di quel blocco.

2. Pruning dei Layer

In base ai punteggi di importanza calcolati, il metodo identifica e rimuove i layer meno importanti. Gli autori notano che il pruning fino al 30% dei layer ha un impatto minimo sull'accuratezza, ma un pruning aggressivo oltre questo punto causa una significativa degradazione delle prestazioni.

3. Merging dei Layer

Per consentire rapporti di pruning superiori al 30% senza perdita di accuratezza, GradPruner introduce una strategia di Layer Merging. Invece di scartare completamente i layer prunati, essi vengono fusi con i layer rimanenti.

• Sparsificazione: I layer prunati vengono prima sparsificati utilizzando la IGIA-Matrix, mantenendo solo il $p\%$ dei parametri (in base alla magnitudo) e azzerando il resto.
• Merging basato sul Segno: Per prevenire l'interferenza causata da segni di parametri contrastanti (dove valori positivi e negativi si annullano a vicenda), il metodo fonde solo elementi con lo stesso segno. Nello specifico, un parametro in un layer prunato viene fuso con il parametro corrispondente in un layer mantenuto solo se i loro segni coincidono. Se i segni differiscono, il parametro prunato viene scartato.

Contributi Chiave

1. Framework GradPruner: La proposta di un metodo che calcola la IGIA-Matrix utilizzando i gradienti iniziali del fine-tuning LoRA per valutare l'importanza dei parametri, eliminando la necessità di estesa calibrazione o di un intero ri-addestramento.
2. Layer Merging Sensibile al Segno: Una nuova tecnica di merging che sparsifica i layer prunati e li combina con i layer mantenuti sulla base della coerenza del segno, permettendo rapporti di pruning più elevati (fino al 40%) mantenendo l'accuratezza.
3. Validazione Empirica: Esperimenti estesi che dimostrano come GradPruner raggiunga una riduzione dei parametri del 40% con una diminuzione dell'accuratezza nei task downstream di appena lo 0,99% su benchmark medici, finanziari e generali.

Risultati Sperimentali

Gli autori hanno valutato GradPruner su due LLM (Llama3.1-8B e Mistral-7B) attraverso otto dataset downstream (tra cui PubMedQA, MedMCQA, BillSum, FinGPT, HellaSwag, WinoGrande, ARC e PIQA).

• Accuratezza: GradPruner ha ottenuto una riduzione media dei parametri del 40% con una diminuzione dell'accuratezza trascurabile dello 0,99% rispetto ai modelli densi. In diversi casi, il modello Llama3.1-8B prunato ha superato il modello più piccolo e non prunato Llama3.2-3B.
• Efficienza: Rispetto ai modelli densi, GradPriment ha ridotto il tempo di addestramento di circa il 36% e la memoria di addestramento del 34%. Il tempo di inferenza e la memoria sono stati ridotti rispettivamente di circa il 39% e il 39%.
• Confronto: GradPruner ha costantemente superato altri baseline di pruning strutturato (LLMPruner, LaCo, MINITRON, SAT, APT) sia in termini di accuratezza che di metriche di efficienza.
• Studi di Ablazione:
  • Rimuovendo la fase di layer merging, si è verificata una significativa caduta dell'accuratezza quando il pruning superava il 30% dei layer.
  • Sostituendo il merging basato sul segno con la media ponderata, l'accuratezza è risultata inferiore, confermando l'importanza della coerenza del segno.
  • Il metodo è rimasto robusto anche quando i dati di addestramento sono stati ridotti al 10% della dimensione originale, sebbene le prestazioni abbiano fluttuato leggermente a volumi di dati estremamente bassi (1%).

Significato e Rivendicazioni

L'articolo sostiene che GradPruner affronti il compromesso tra efficienza del fine-tuning e prestazioni del modello. Sfruttando la rapida dinamica di apprendimento nelle prime fasi del fine-tuning, esso identifica i parametri critici senza l'overhead di un addestramento completo o di una complessa distillazione. Gli autori affermano che la loro strategia di merging basata sul segno mitiga efficacementamente l'interferenza tipicamente associata al pruning aggressivo, consentendo una sostanziale riduzione dei parametri (40%) preservando la capacità del modello sui task specializzati downstream. Il lavoro suggerisce che il pruning a livello di layer combinato con un merging intelligente sia una via percorribile per il deployment di LLM efficienti e adattati al dominio.