Riepilogo Tecnico: TraceNAS

Definizione del Problema

Il dispiegamento dei Large Language Models (LLM) è limitato dai massicci conteggi dei parametri e dai requisiti computazionali. Sebbene il pruning strutturato offra una via per un dispiegamento efficiente rimuovendo unità architettoniche (ad es., teste di attenzione, blocchi transformer), i metodi esistenti affrontano limitazioni significative:

1. Dipendenze Locali vs. Globali: La maggior parte dei metodi di pruning valuta l'importanza di layer, teste o canali in isolamento utilizzando euristiche locali (ad es., magnitudo pesata dalle attivazioni o ricostruzione basata su Hessian). Questi approcci ignorano le complesse dipendenze strutturali globali e la propagazione degli errori indotti dal pruning attraverso il modello.
2. Costo Computazionale della Ricerca Consapevole dell'Addestramento (Training-Aware Search): I metodi che tentano di catturare le dipendenze globali spesso si affidano a una Neural Architecture Search (NAS) consapevole dell'addestramento. Questi approcci richiedono aggiornamenti del gradiente iterativi e flussi continui di token non visti per prevenire l'overfitting, rendendo il processo di ricerca computazionalmente costoso quanto l'addestramento di recupero post-pruning stesso.
3. Inadeguatezza degli Esistenti Proxy Zero-Shot: Gli attuali proxy privi di addestramento sono progettati per identificare architetture addestrabili e inizializzate casualmente concentrandosi sulla trainabilità (ad es., stabilità del gradiente, fluidità del paesaggio della perdita). Essi non riescono a valutare l'"eredità funzionale" (functional inheritance): la capacità di un modello prunato di recuperare le prestazioni preservando il paesaggio della perdita pre-addestrato.

Metodologia

TraceNAS è proposto come un framework NAS privo di addestramento progettato per ottimizzare congiuntamente la profondità e la larghezza degli LLM senza addestramento durante la ricerca.

1. Spazio di Ricerca e Codifica

Il framework tratta l'LLM pre-addestrato come una super-network. Lo spazio di ricerca $G$ è definito come una distribuzione congiunta su:

• Profondità ($d$): Una maschera discreta su $L$ blocchi transformer, che consente una riduzione della profondità non uniforme mantenendo la continuità dello stream residuo.
• Larghezza ($\kappa$): Rapporti di ritenzione dei parametri per i sub-blocchi di attenzione e MLP, consentendo una sparsità eterogenea attraverso il modello.

2. Realizzazione Architetturale In-Place

Per evitare l'overhead di memoria dell'instanziazione di migliaia di candidati, TraceNAS impiega una strategia di masking in-place:

• Masking della Larghezza: Utilizza un prodotto pesato dalle attivazioni (simile a Wanda) per generare maschere binarie per i canali dei pesi. Questa euristica $O(d^2)$ identifica i canali funzionalmente critici preservando al contempo le caratteristiche outlier.
• Masking della Profondità: Salta interi blocchi se disattivati, instradando le attivazioni attraverso le connessioni residue.
• Esecuzione: Il framework modifica temporaneamente le matrici dei pesi del modello base per calcolare i gradienti su un set di calibrazione minimo, per poi ripristinare i pesi originali, preservando il modello base per le iterazioni successive.

3. Il Proxy Zero-Shot: Correlazione della Traccia del Gradiente

L'innovazione principale è un proxy $\Phi$ invariante per scala che misura l'allineamento funzionale tra un candidato prunato ($M_{sub}$) e il base pre-addestrato ($M_{base}$).

• Tracce del Gradiente: I gradienti sono calcolati su un set di calibrazione. Per garantire la scalabilità, questi vengono proiettati in un sottospazio a basso rango (LoRA), sfruttando la bassa dimensionalità intrinseca degli LLM per ridurre l'overhead di memoria di ordini di grandezza.
• Metrica di Allineamento: Il proxy calcola il Coefficiente di Correlazione di Pearson tra le tracce del gradiente standardizzate del modello prunato e del modello base per ogni layer. Ciò garantisce l'invarianza di scala, concentrandosi sull'allineamento direzionale piuttosto che sugli spostamenti di magnitudo.
• Aggregazione Pesata dalla Sparsità: Il punteggio finale $\Phi$ è una somma pesata delle correlazioni per layer, dove i pesi corrispondono ai rapporti di ritenzione dei parametri. Questo ancora il punteggio globale nelle regioni ad alta capacità, impedendo al rumore proveniente da sub-blocchi aggressivamente prunati di distorcere la classifica.

4. Ricerca Evolutiva

TraceNAS impiega una strategia evolutiva per navigare nello spazio di ricerca discreto:

• Inizializzazione: Una popolazione viene generata sulla base di priorità di importanza globale.
• Evoluzione: La popolazione subisce crossover e mutazione. L'evoluzione della profondità utilizza mutazioni di bit-flip e crossover di intervalli contigui, mentre l'evoluzione della larghezza utilizza interpolazione e jitter gaussiano.
• Selezione: I candidati sono classificati esclusivamente dal proxy $\Phi$, selezionando modelli con l'alto potenziale di recupero delle prestazioni durante l'addestramento continuo post-pruning (CPT).

Contributi Chiave

1. NAS Scalabile Senza Addestramento: TraceNAS fornisce un framework unificato per il pruning congiunto di profondità e larghezza, consentendo la scoperta di architetture non uniformi senza addestramento durante la ricerca.
2. Proxy di Eredità Funzionale: Introduce un nuovo proxy basato sul gradiente e invariante per scala che valuta l'allineamento delle tracce del gradiente tra i modelli prunati e quelli pre-addestrati. Raggiunge un'alta correlazione con l'accuratezza a valle (Spearman $\rho = 0.94$, Kendall $\tau = 0.82$).
3. Efficienza: Il framework riduce l'overhead di ricerca utilizzando sottospazi del gradiente a basso rango e masking in-place, consentendo la scoperta di modelli ad alta fedeltà su una singola GPU.

Risultati Sperimentali

TraceNAS è stato valutato su Llama-2-7B, Llama-3.1-8B e Qwen-2.5-14B, effettuando il pruning a rispettivamente 2.7B, 4.6B e 8.4B parametri.

• Efficienza di Ricerca: TraceNAS ha identificato architetture ottimali in 8.5 ore di GPU (su una A100), rappresentando una riduzione di 10× nel tempo di ricerca rispetto ai baseline consapevole dell'addestramento come DarwinLM. Ha richiesto solo 98M di token di ricerca, una riduzione di 10× rispetto a DarwinLM.
• Prestazioni:
  • Llama-2-7B: TraceNAS ha raggiunto l'accuratezza media più elevata (62.81%) attraverso otto benchmark di ragionamento, superando DarwinLM (62.57%) e ShearedLLaMA (62.63%). Notevolmente, le prestazioni di ShearedLLaMA sono scese significativamente (a 60.86%) quando limitate allo stesso budget di token usato da TraceNAS.
  • Llama-3.1-8B & Qwen-2.5-14B: TraceNAS ha superato i baseline di pruning uniforme e i metodi consapevole dell'addestramento (DarwinLM, E3-Pruner) in compiti di ragionamento complessi (ad es., ARC-Challenge, BoolQ) richiedendo significativamente meno token di recupero (20B vs 94B per Minitron).
• Validazione del Proxy: Il proxy ha dimostrato un'alta stabilità attraverso diversi iperparametri (dimensione del campione, lunghezza del contesto e rango LoRA), confermando che l'allineamento della traccia del gradiente cattura efficacemente l'eredità del modello anche in condizioni vincolate.

Significato e Rivendicazioni

Il paper sostiene che TraceNAS affronta il collo di bottiglia critico del costo computazionale nel pruning strutturato. Spostando l'attenzione dalla "trainabilità" (apprendimento da zero) all' "eredità" (preservazione del paesaggio della perdita pre-addestrata), TraceNAS consente la scoperta di modelli prunati non uniformi ad alte prestazioni con un dispendio minimo di risorse.

Gli autori sottolineano che il loro metodo:

• Elimina il costo proibitivo dell'addestramento durante la ricerca, rendendo accessibile la compressione di modelli su larga scala.
• Cattura le dipendenze strutturali globali che le euristiche locali perdono, garantendo la preservazione delle complesse capacità di ragionamento.
• Riduce l'impronta di carbonio della compressionione dei modelli riducendo drasticamente le ore-GPU richieste per la ricerca dell'architettura.

Il lavoro posiziona TraceNAS come una base scalabile per il dispiegamento di LLM efficienti in ambienti con risorse limitate, mantenendo l'integrità funzionale dei modelli densi pre-addestrati.