Steering Sparse Autoencoder Latents to Control Dynamic Head Pruning in Vision Transformers (Student Abstract)

Titolo

Guidare i latenti degli Autoencoder Sparsi per Controllare la Potatura Dinamica delle Teste nei Vision Transformer (ViT).

1. Il Problema

I Vision Transformer (ViT) utilizzano l'attenzione multi-testa per catturare interazioni complesse tra token, ma molte di queste teste sono ridondanti, aumentando il costo computazionale senza un guadagno proporzionale nelle prestazioni.

• Limitazione attuale: Esistono framework di potatura adattiva (come AdaViT) che rimuovono dinamicamente le teste ridondanti in base all'input. Tuttavia, le politiche di decisione di questi sistemi sono spesso opache e difficili da controllare a livello latente, poiché si basano direttamente su embedding residui densi.
• Obiettivo: Rendere il processo di potatura non solo efficiente, ma anche interpretabile e controllabile a livello di caratteristiche latenti, permettendo di comprendere e guidare quali teste vengono selezionate per classi specifiche.

2. Metodologia

Gli autori propongono un nuovo framework che integra gli Autoencoder Sparsi (SAE) con la potatura dinamica delle teste nei ViT. Il processo si articola in tre fasi principali:

1. Baseline di Potatura Dinamica (AdaViT):

   • Viene utilizzato un ViT con reti di decisione leggere per ogni layer.
   • Ogni rete riceve il token CLS (che codifica il contesto globale) dallo stream residuo e produce logit di importanza per ogni testa.
   • Vengono generate maschere binarie per la potatura tramite campionamento Gumbel-Sigmoid, riducendo l'uso delle teste mantenendo l'accuratezza.

2. Addestramento dell'SAE:

   • Dopo l'addestramento del ViT, viene estratto il token CLS dello strato finale ($x$).
   • Viene addestrato un k-sparse autoencoder su questi embedding. L'SAE mappa l'embedding di input ($d=384$) in uno spazio latente più ampio ($n=3072$) e applica un'attivazione Top-K (mantenendo solo le $k=64$ dimensioni più grandi).
   • L'obiettivo è disaccoppiare le embedding dense e polisenatiche in caratteristiche latenti sparse e monosemantiche (interpretabili).

3. Guida (Steering) dei Latenti:

   • Per influenzare le decisioni di potatura, le dimensioni latenti selezionate ($S$) vengono amplificate aggiungendo un valore $\alpha$.
   • Il vettore latente modificato ($\tilde{z}$) viene ricostruito in un nuovo embedding steered ($\tilde{x}$).
   • Questo embedding steered viene reinserito nella rete di decisione per generare una nuova maschera di potatura, permettendo di osservare come cambia il comportamento di selezione delle teste.

3. Contributi Chiave

• Framework Ibrido: Integrazione innovativa tra SAE e potatura dinamica per rendere le decisioni di pruning controllabili a livello latente.
• Interpretabilità Meccanica: Dimostrazione che le caratteristiche latenti sparse possono essere mappate a sottoinsiemi specifici di teste per classi specifiche, rivelando relazioni semantiche tra le teste.
• Strategia di Guida Per-Classe: Identificazione che l'amplificazione delle attivazioni frequenti per classe è la strategia più efficace per migliorare l'efficienza senza sacrificare l'accuratezza.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su un ViT-Small pre-addestrato su ImageNet, fine-tunato su CIFAR-100.

• Baseline: Il modello potato dinamicamente rimuove il 30% delle teste mantenendo un'accuratezza del 89,79% (rispetto al 91,27% originale).
• Validazione SAE: La sostituzione dell'embedding originale con la ricostruzione SAE ha causato variazioni minime nell'accuratezza (-0,12%) e nell'uso delle teste, confermando che l'SAE preserva le informazioni essenziali.
• Effetto della Guida (Steering):
  • La strategia "Per-class frequent" (amplificare i latenti più attivi per ogni classe) riduce significativamente l'uso delle teste mantenendo o migliorando l'accuratezza.
  • Al contrario, strategie globali o casuali portano a un calo drastico dell'accuratezza all'aumentare dell'amplificazione ($\alpha$).
  • Esempi specifici:
    • Per la classe "bowl": L'accuratezza è passata dal 76% all'82%, mentre l'uso delle teste è sceso da 0,72 a 0,33, affidandosi principalmente alle teste $h2$ e $h5$.
    • Per la classe "pine tree": Accuratezza dal 79% all'84%, uso teste da 0,93 a 0,35 (teste $h2$ e $h3$).
  • Scoperta Semantica: Classi semanticamente correlate (es. "bowl" e "plate") condividono sottoinsiemi simili di teste attive, dimostrando che la guida rivela relazioni semantiche a livello di classe.

5. Significato e Conclusioni

Questo lavoro colma il divario tra efficienza computazionale e interpretabilità meccanica nei Vision Transformer.

• Dimostra che è possibile non solo ridurre il costo computazionale tramite la potatura dinamica, ma anche guidare intenzionalmente il modello verso configurazioni di teste ottimali per specifiche categorie semantiche.
• L'approccio trasforma la potatura da un processo "scatola nera" a uno strumento controllabile, dove le decisioni di pruning possono essere spiegate attraverso le caratteristiche latenti sparse.
• Lavori Futuri: Gli autori intendono estendere il framework agli strati precedenti del ViT e applicarlo a modelli foundation più grandi, andando oltre l'attuale focus sull'ultimo strato e su dataset di piccole dimensioni.