Sintesi Tecnica: SnapViT – Vision Transformer elastici da modelli pre-addestrati senza riaddestramento

Problematica

I modelli di fondazione per la visione hanno raggiunto prestazioni allo stato dell'arte in compiti diversificati, ma vengono tipicamente rilasciati in un set limitato di dimensioni predeterminate (ad esempio, conteggi specifici di parametri per le famiglie DINOv3 o DeIT). Questa rigidità costringe i professionisti a fare scelte di deployment sub-ottimali quando si affrontano vincoli reali di calcolo, latenza e costi. Le soluzioni esistenti per la flessibilità dei modelli presentano limitazioni significative:

• Distillazione della Conoscenza (Knowledge Distillation): Richiede un'architettura target predeterminata e spesso si basa su dataset di pre-addestramento non pubblici, minando la robustezza.
• Inferenza Elastica: Metodi come le rappresentazioni Matryoshka o le reti Slimmable richiedono che i modelli siano progettati con strutture nidificate specifiche durante il pre-addestramento, rendendoli inapplicabili ai modelli pre-addestrati standard (off-the-shelf) già esistenti.
• Pruning Strutturato: Sebbene efficace nel ridurre memoria e calcolo, la maggior parte delle tecniche di pruning attuali è adattata a compiti o vincoli specifici e richiede tipicamente il riaddestramento, lasciando un vuoto per modelli universalmente adattabili che non necessitano di dati etichettati o di un head di classificazione.

Metodologia: SnapViT

Gli autori propongono SnapViT (Single-shot network approximation for pruned Vision Transformers), un metodo di pruning strutturato post-pre-addestramento che genera un continuum di sottoreti elastiche da un singolo modello pre-addestrato senza riaddestramento né dati etichettati. Il nucleo del metodo è un punteggio di prunability ($P$) che classifica i parametri da rimuovere in base a due componenti:

1. Approssimazione Locale dell'Hessiana (Gradienti Auto-Supervisionati)

Per catturare la sensibilità locale dei parametri senza un head di classificazione, il metodo approssima la diagonale della matrice Hessiana utilizzando gradienti auto-supervisionati. Nello specifico, impiega l'obiettivo DINO per calcolare i gradienti su un dataset di immagini. La curvatura locale $H^{(l)}$ è stimata come la norma quadratica media dei gradienti:
$$H^{(l)} \approx \frac{1}{N_D} \sum_{i=1}^{N_D} \|\nabla_\theta L_i\|^2$$
Questo fornisce una misura di base di quanto i singoli parametri siano critici all'interno dei loro blocchi specifici (ad esempio, neuroni feed-forward o teste di attenzione).

2. Stima Globale dell'Hessiana (Algoritmo Evolutivo)

Per affrontare il limite derivante dall'ignorare le dipendenze inter-blocco (elementi off-diagonal dell'Hessiana), SnapViT introduce una nuova approssimazione utilizzando l'Exponential Natural Evolution Strategy (xNES).

• Invece di calcolare esplicitamente l'intratabile Hessiana completa, xNES apprende un insieme di fattori di scala basati sulla struttura ($c$) che ricalibrano le sensibilità locali.
• L'algoritmo ottimizza una distribuzione Gaussiana multivariata sui fattori di scala. La funzione di fitness è priva di etichette (label-free), misurando la similarità del coseno tra gli embedding del modello originale e del modello prunato (dopo compressione PCA).
• Attraverso aggiornamenti di gradiente naturale, la matrice di covarianza ($\Sigma$) della strategia evolutiva evolve per approssimare l'inversa della vera Hessiana inter-blocco ($H^{(g)} \approx \alpha \Sigma^{-1}$). Ciò consente al metodo di identificare quali blocchi possono essere prunati congiuntamente senza degradare la qualità della rappresentazione.

3. Punteggio di Prunability Unificato e Generazione Elastica

Il punteggio di prunability finale è il prodotto elemento per elemento della sensibilità locale e dei fattori di scala globali:
$$P = \text{diag}(H^{(l)}) \odot (M c)$$
dove $M$ è una matrice di appartenenza che espande i fattori a livello di blocco ai singoli parametri.

• Generazione Single-Shot: Una volta calcolato $P$, tutti i parametri vengono classificati globalmente. Questo singolo ranking permette l'estrazione di sottoreti per qualsiasi livello di sparsità desiderato ($S \in [0, 1]$) senza ulteriore ottimizzazione.
• Ambito: Il metodo effettua il pruning sia di combinazioni riga-colonna all'interno dei blocchi feed-forward, sia di intere teste di attenzione.

Contributi Chiave

1. Strategia di Pruning Elastico: Una strategia di pruning efficiente e veloce per i Vision Transformer pre-addestrati che produce modelli elastici adattabili a qualsiasi vincolo computazionale senza riaddestramento.
2. Approssimazione Evolutiva dell'Hessiana: Una nuova strategia che utilizza un algoritmo genetico (xNES) per approssimare le componenti off-diagonal dell'Hessiana per le strutture di rete, catturando le correlazioni cross-block che le tradizionali approssimazioni diagonali trascurano.
3. Prestazioni Label-Free e Head-Agnostic: Il metodo raggiunge prestazioni allo stato dell'arte sotto un pruning significativo senza richiedere dati etichettati o un head di classificazione, generalizzando attraverso modelli di fondazione supervisionati e auto-supervisionati.

Risultati Sperimentali

Gli autori hanno valutato SnapViT su otto dataset (inclusi ImageNet-1k, Pascal VOC e vari benchmark fine-grained) utilizzando la classificazione k-nearest neighbor (k-NN), il linear probing e la segmentazione semantica lineare.

• Prestazioni vs SOTA: SnapViza costantemente eguaglia o supera i metodi allo stato dell'arte (LAMP, LLM Surgeon, FPTP, SparseGPT, SNIP Magnitude, NViT) attraverso vari livelli di sparsità.
  • Per DINOv1 ViT-B/16, il metodo effettua il pruning fino al 40% di sparsità con meno del 5% di degradazione dell'accuratezza, ottenendo un'accelerazione dell'inferenza di 1.58x.
  • Al 50% di sparsità sui modelli supervisionati AugReg, supera SNIP del 7% e FPTP del 12,3% nella classificazione lineare.
  • Per SigLIPv2 e DINOv3 (modelli di fondazione su larga scala), il metodo supera significativamente i baseline, con miglioramenti fino al 34,3% nella classificazione lineare rispetto ai secondi migliori metodi.
• Efficienza: L'intero processo di pruning per una famiglia di modelli (ad esempio, DeIT-III) richiede meno di cinque minuti su una singola GPU A100.
• Post-Processing: Il metodo è compatibile con la correzione dei pesi post-pruning (stile SparseGPT) e il fine-tuning completo. Con un singolo passaggio di correzione dei pesi, le prestazioni a sparsità estrema (ad esempio, 50% su SigLIPv2) sono ampiamente preservate.
• Scalabilità: Sebbene il pre-addestramento su larga scala (ad esempio, DINOv3, SigLIPv2) renda il pruning più complesso a causa della conoscenza rappresentativa distribuita, SnapViT rimane efficace, specialmente se combinato con la correzione dei pesi o orizzonti di ottimizzazione più lunghi (500 iterazioni xNES).

Significato e Rivendicazioni

Il paper sostiene che SnapViT colmi il divario tra la rigidità degli attuali deployment dei modelli di fondazione e la necessità di un'inferenza flessibile e adattabile alle risorse. Abilitando l'estrazione di un continuum di sottoreti da un singolo modello pre-addestrato in un unico passaggio (single-shot), consente agli utenti di personalizzare precisamente i modelli allo stato dell'arte per i loro specifici budget computazionali.

Gli autori sottolineano che il loro approccio:

• Elimina la necessità di riaddestramento, risparmiando significative risorse computazionali e tempo.
• Rimuove la dipendenza dai dati etichettati, rendendolo applicabile ai modelli di fondazione auto-supervisionati che mancano di head di classificazione.
• Cattura le correlazioni strutturali globali tramite l'approssimazione evolutiva, superando i limiti delle metriche di pruning puramente locali.
• Riduce l'impatto ambientale abbassando i requisiti di calcolo e l'uso di energia dei vision transformer, riducendo potenzialmente le emissioni di CO2 dall'inferenza.

Il lavoro si presenta come una soluzione pratica per il deployment di modelli di visione su larga scala in ambienti vincolati, mantenendo al contempo alte prestazioni attraverso una vasta gamma di compiti downstream.