STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

Sintesi Tecnica: STARFISH

Definizione del Problema
La potatura (pruning) delle reti neurali è una tecnica standard per ridurre la dimensione del modello, l'impronta di memoria e la latenza di inferenza rimuovendo pesi, neuroni o canali. Tuttavia, una potatura aggressiva spesso porta a un degrado sostanziale dell'accuratezza predittiva. Sebbene vengano solitamente impiegati processi di recupero (healing) per riacquistare le prestazioni perse, i metodi esistenti presentano limitazioni significative. Gli approcci tradizionali richiedono spesso un intero riaddestramento o un fine-tuning sul dataset di addestramento originale, il che è computazionalmente costoso e spesso impossibile in contesti moderni in cui i dati di addestramento sono proprietari o inaccessibili. Inoltre, i metodi di recupero allo stato dell'arte spesso faticano a ripristinare l'accuratezza in condizioni di alta sparsità (ad esempio, rimuovendo >75% dei pesi), in particolare per le moderne architetture Vision Transformer (ViT).

Metodologia: STARFISH
Gli autori propongono STARFISH (faST Accuracy Recovery in pruned networks From Internal State Healing), un metodo di recupero post-potatura che opera senza accesso ai dati di addestramento originali o alle etichette. L'ipotesi centrale è che ottimizzare l'allineamento delle rappresentazioni intermedie (stati interni) tra la rete potata e l'originale rete densa sia un proxy più efficace per il ripristino dell'accuratezza rispetto all'ottimizzazione basata unicamente sulle predizioni di output finali.

L'algoritmo procede in due fasi utilizzando un piccolo set di calibrazione non etichettato ($S_{cal}$) campionato dalla distribuzione dei dati ma distinto dal set di addestramento:

1. Caching delle Rappresentazioni Target: Il modello denso originale viene eseguito su $S_{cal}$ per memorizzare (cache) gli output vettorizzati di ogni blocco computazionale (ad esempio, dopo la normalizzazione, le non-linearità e le aggiunte residue).
2. Ottimizzazione dell'Allineamento: Il modello potato ($f_{\theta \odot m}$) viene addestrato per minimizzare la perdita di allineamento tra le sue rappresentazioni intermedie e i target memorizzati. La funzione di perdita è definita come la distanza media del coseno tra gli stati nascosti del modello potato ($f^\ell_{\theta \odot m}(x)$) e gli stati nascosti del modello denso ($r^\ell_x$) attraverso tutti i blocchi $\ell$:
   $$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \odot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$$
   I pesi rimanenti del modello potato vengono aggiornati per minimizzare la perdita empirica sul set di calibrazione.

La motivazione teorica è fornita dimostrando che la distanza tra le rappresentazioni nascoste finali (input per il layer di classificazione) funge da limite superiore per la divergenza di Kullback–Leibler (KL) delle distribuzioni di output. Ciò suggerisce che l'allineamento degli stati interni vincola intrinsecamente l'errore di output.

Contributi Chiave

1. Recupero Basato sulla Rappresentazione: Il documento introduce un metodo di recupero indipendente dai tipi di potatura, dalle maschere o dalle architetture specifiche. Dimostra empiricamente che l'allineamento della rappresentazione intermedia supera gli obiettivi diretti basati sull'output (come la Cross-Entropy o il matching della divergenza KL) all'aumentare dei livelli di sparsità.
2. Efficienza ad Alta Sparsità: Il metodo è il primo a ripristinare efficacemente alte percentuali dell'accuratezza originale a livelli di sparsità estremi (ad esempio, 85% di sparsità). Raggiunge questo risultato utilizzando un minuscolo set di calibrazione (anche solo 1.000 esempi, ovvero lo 0,4% della dimensione del set di addestramento) e non richiede l'accesso ai dati di addestramento originali.
3. Modularità e Velocità: L'approccio è computazionalmente leggero, richiedendo solo pochi passaggi forward e backward su un piccolo dataset. È applicabile a schemi di potatura strutturati, semi-strutturati e non strutturati.

Risultati Sperimentali
Gli autori hanno valutato STARFISH su ImageNet-1K utilizzando varie architetture Vision Transformer (DeiT, ViT) e Convolutional Neural Network (MobileNet).

• Alta Sparsità (0,75–0,85): Su un modello DeiT-B potato con una sparsità del 75%, STARFISH ha recuperato l'82% dell'accuratezza originale del modello denso, mentre i metodi concorrenti (come CORP) hanno recuperato solo il 40%. Su un modello DeiT3-H con l'85% di sparsità, STARFISH ha mantenuto il 92,08% dell'accuratezza del modello denso, migliorando le prestazioni del modello potato di 72,36 punti percentuali.
• Moderata Sparsità (0,5): Su varie architetture DeiT e diversi scope di potatura (MLP, QK, e congiunto MLP+QK), STARFISH ha superato costantemente i metodi allo stato dell'arte. Ad esempio, in un modello DeiT-T con il 50% di potatura strutturata su MLP+QK, STARFISH ha ottenuto un mantenimento dell'81,95% dell'accuratezza densa, un miglioramento di 31,54 punti percentuali rispetto al metodo CORP.
• Confronto: STARFISH ha dimostrato prestazioni superiori rispetto a metodi come SNOWS (che utilizza l'approssimazione dell'Hessiana) e CORP (che utilizza il recupero affine a forma chiusa), particolarmente nei regimi in cui la potatura degrada severamente il modello.

Significato e Rivendicazioni
Il documento sostiene che STARFISH rappresenti un avanzamento significativo nel campo della compressione dei modelli, spostando il paradigma di recupero dal matching dell'output alla guarigione dello stato interno (internal state healing). Gli autori sottolineano che questo approccio è particolarmente vitale per i moderni modelli su larga scala dove i dati di addestramento non sono disponibili e il riaddestramento completo è impraticabile. Sfruttando il principio secondo cui le rappresentazioni interne catturano la conoscenza del modello in modo più robusto rispetto agli output finali sotto alta sparsità, STARFISH consente l'implementazione di reti potate altamente efficienti senza sacrificare le prestazioni predittive. Il metodo è presentato come una soluzione generale, "plug-and-play", che può essere applicata a qualsiasi rete potata indipendentemente dalla strategia di potatura iniziale.