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

Probleemstelling

Vision Transformers (ViTs) maken gebruik van multi-head self-attention om diverse token-interacties te modelleren. Een groot deel van deze attention-heads is echter overbodig (redundant), wat leidt tot onnodige rekenkracht zonder een evenredige prestatiewinst. Om dit aan te pakken, zijn adaptieve frameworks zoals AdaViT ontwikkeld, die dynamisch heads verwijderen op basis van invoer.

Hoewel deze dynamische pruning de efficiëntie aanzienlijk verbetert, blijft er een fundamenteel probleem bestaan: de beslissingsprocessen van deze pruning-politiek zijn ondoorzichtig (opake) en moeilijk te controleren op het latente niveau. Omdat de beslissingen gebaseerd zijn op dichte residual embeddings, ontbreekt het aan mechanische interpreteerbaarheid. Het is voor onderzoekers en ontwikkelaars lastig om te begrijpen waarom bepaalde heads worden verwijderd of om deze selectie op een gecontroleerde manier te sturen.

Methodologie

De auteurs stellen een nieuw framework voor dat Sparse Autoencoders (SAE's) integreert met dynamische head-pruning in ViTs. Het doel is om pruning-keuzes interpreteerbaar en controleerbaar te maken door gebruik te maken van de ontvlechtingseigenschappen van SAE's.

Het proces verloopt als volgt:

1. Baseline Training:

   • Er wordt een ViT getraind met een bijbehorend "decision network" per laag.
   • Dit netwerk ontvangt de CLS-token (die globale context bevat) uit de residual stream en voorspelt de belangrijkheid van elke attention-head.
   • Via Gumbel-Sigmoid sampling worden binaire masks gegenereerd om heads te verbergen (prunen) terwijl de nauwkeurigheid behouden blijft.

2. Training van de Sparse Autoencoder (SAE):

   • Na het trainen van de ViT wordt de CLS-token uit de laatste laag gebruikt als input voor een k-sparse autoencoder.
   • De SAE ontvlecht de dichte embedding ($x$) in een hoge-dimensionale, spare latente representatie ($z$).
   • De encoder selecteert alleen de $k$ grootste activeringen (Top-K), wat zorgt voor monosemantische (betekenisdrachtige) en interpreteerbare features.
   • De decoder reconstrueert de embedding ($\hat{x}$) met als doel de reconstructiefout (MSE) te minimaliseren.

3. Sturing (Steering) en Pruning:

   • Om de pruning te sturen, worden specifieke latente dimensies in de SAE versterkt (amplified) met een factor $\alpha$.
   • De versterkte latente vector wordt opnieuw gefilterd via Top-K en gereconstrueerd tot een "geleide" embedding ($\tilde{x}$).
   • Deze geleide embedding wordt teruggevoerd naar het decision network, wat leidt tot een aangepast pruning-mask.

Belangrijkste Bijdragen

• Integratie van SAE en Pruning: Het paper introduceert een innovatieve manier om SAE's te gebruiken om de interne werking van dynamische pruning te manipuleren en te begrijpen.
• Interpreteerbaarheid op Latent Niveau: In plaats van te vertrouwen op ondoorzichtige dichte embeddings, onthult het framework welke specifieke, interpreteerbare concepten (spare latents) verantwoordelijk zijn voor de selectie van heads.
• Klasse-specifieke Controle: De auteurs tonen aan dat het versterken van frequent geactiveerde latents per klasse leidt tot geoptimaliseerde, klasse-specifieke pruning-strategieën.

Resultaten

De experimenten zijn uitgevoerd op een ViT-Small model getraind op CIFAR-100 (gefine-tuned op ImageNet).

• Efficiëntie vs. Nauwkeurigheid:

  • De SAE zelf behoudt de essentiële informatie: het vervangen van de originele embedding door de gereconstrueerde versie resulteert in een verwaarloosbaar verlies in nauwkeurigheid (-0,12%).
  • Per-klasse sturing bleek de meest effectieve strategie. In tegenstelling tot globale of willekeurige sturing (die leiden tot grote nauwkeurigheidsdalingen), kon per-klasse sturing het aantal gebruikte heads aanzienlijk verminderen terwijl de nauwkeurigheid behouden bleef of zelfs verbeterde.

• Concrete Voorbeelden:

  • Voor de klasse "bowl": De nauwkeurigheid steeg van 76% naar 82%, terwijl het head-gebruik daalde van 0,72 naar 0,33. Dit werd bereikt door uitsluitend te vertrouwen op heads h2 en h5.
  • Voor de klasse "pine tree": Nauwkeurigheid steeg van 79% naar 84% met een daling in head-gebruik van 0,93 naar 0,35 (afhankelijk van h2 en h3).
  • Semantische Relaties: Interessant is dat semantisch gerelateerde klassen (zoals "bowl" en "plate") vergelijkbare subsets van heads (h2 en h5) deelden. Dit suggereert dat de SAE in staat is om semantische relaties op het niveau van attention-heads bloot te leggen.

Betekenis en Conclusie

Dit onderzoek overbrugt de kloof tussen efficiëntie en mechanische interpreteerbaarheid in Vision Transformers. Het bewijst dat dynamische head-pruning niet alleen een zwartdoos-optimalisatie hoeft te zijn, maar gecontroleerd kan worden via sparse latente features.

Door het versterken van klasse-specifieke signalen kunnen onderzoekers niet alleen de rekenkosten verlagen, maar ook begrijpen welke attention-heads essentieel zijn voor specifieke concepten. De huidige studie focust zich op de laatste laag en kleinere datasets; toekomstig werk zal het framework uitbreiden naar eerdere lagen en grotere foundation modellen.