Can Cross-Layer Transcoders Replace Vision Transformer Activations? An Interpretable Perspective on Vision

Probleemstelling

Het begrijpen van de interne activaties van Vision Transformers (ViTs) is cruciaal voor het bouwen van interpreteerbare en betrouwbare modellen. Bestaande methoden, zoals Sparse Autoencoders (SAEs), worden gebruikt om menselijk interpreteerbare functies te extraheren, maar hebben twee belangrijke beperkingen:

1. Lokale werking: SAEs opereren op individuele lagen en missen de cross-layer computationele structuur van Transformers.
2. Gebrek aan context: Ze kunnen niet vaststellen wat de relatieve bijdrage van elke laag is aan de uiteindelijke representatie van de laatste laag.

Dit leidt tot een beperkt inzicht in hoe informatie door de diepte van het netwerk wordt getransformeerd. Het paper stelt de vraag of Cross-Layer Transcoders (CLTs), die eerder succesvol zijn toegepast in taalmodellen (LLMs), ook effectief kunnen worden ingezet voor visuele modellen om zowel functionaliteit als interpretatie te verbeteren.

Methodologie: Cross-Layer Transcoders (CLTs)

De auteurs introduceren CLTs als betrouwbare, sparse en diepte-bewuste proxy-modellen voor de MLP-blokken (Multi-Layer Perceptrons) in ViTs.

• Architectuur: Een CLT gebruikt een encoder-decoder schema.

  • Encoder: Elke laag $i$ leert een lineaire encoder om de pre-MLP activaties ($x_i$) te mappen naar een sparse feature vector ($z_i$). Er worden drie niet-lineaire functies gebruikt voor sparsiteit: JumpReLU, ReLU-Top-k, en Abs-Top-k.
  • Decoder: Voor een doellag $j$ wordt de post-MLP activatie ($\hat{y}_j$) gereconstrueerd als een som van de sparse codes van alle voorgaande lagen ($i \leq j$):
    $$\hat{y}_j = \sum_{i=1}^{j} z_i W_{dec}^{i \to j}$$
    Dit creëert een triangulaire decoder-architectuur die een strikt causale structuur garandeert (geen informatielekken van toekomstige lagen).

• Training: De CLTs worden getraind om de post-MLP activaties van een bevroren "leraar" ViT (bijv. CLIP ViT-B/32 of B/16) te reconstrueren. Het trainingsdoel is het minimaliseren van de reconstructiefout (MSE) gecombineerd met een sparsiteitsstraf.

• Vervanging (Replacement): Tijdens inferentie kunnen de originele MLP-blokken worden vervangen door de CLT-berekeningen. Dit resulteert in een "vervangingmodel" dat fungeert als een interpreteerbare proxy voor de originele ViT zonder de prestaties op downstream-taken (zoals zero-shot classificatie) te compromitteren.

• Attributie: Omdat de reconstructie lineair is, kan de bijdrage van elke bronlaag $i$ aan een doellag $j$ worden kwantitatief bepaald via projectie-scores. Dit maakt het mogelijk om te zien welke lagen het meest verantwoordelijk zijn voor de uiteindelijke representatie.

Belangrijkste Bijdragen

1. Toepassing van CLTs in Vision: Het paper onderzoekt voor het eerst de effectiviteit van Cross-Layer Transcoders specifiek voor Vision Transformers, waarbij ze in staat blijken om post-MLP activaties te reconstrueren vanuit eerdere lagen.
2. Functionele Vervanging: CLTs kunnen MLP-blokken vervangen, vooral in latere lagen of voor [CLS]-tokens, terwijl de zero-shot classificatieprestaties behouden blijven of zelfs licht verbeteren. Dit maakt ze tot een betrouwbare proxy voor analyse.
3. Interpreteerbare Cross-Layer Attributie: De methode levert projectie-gebaseerde bijdragescijfers op. Deze tonen een opvallend asymmetrisch patroon:
   • Patch-tokens: Vertonen een sterk diagonaal-dominante attributie (elke laag verklaart voornamelijk zijn eigen output).
   • [CLS]-token: Haalt credit breed uit de diepte van het netwerk, waarbij informatie uit vele voorgaande lagen wordt geaggregeerd.
4. Geloofwaardige Attributie: Ablatie-experimenten tonen aan dat de uiteindelijke representatie grotendeels wordt bepaald door een kleine subset van dominante bronlagen. Het behouden van slechts de top-4 lagen herstelt de originele nauwkeurigheid, terwijl het verwijderen van de hoogst-scorende laag de prestaties aanzienlijk verslechtert.

Resultaten

De methoden zijn getest op CLIP ViT-B/32 en ViT-B/16 modellen met datasets CIFAR-100, COCO en ImageNet-100.

• Reconstructiekwaliteit: CLTs bereiken hoge reconstructiefideliteit. De gemiddelde cosine-similariteit ligt tussen 0,92 en 0,97, en de verklaarde variantie ($R^2$) tussen 0,89 en 0,95. ViT-B/16 is over het algemeen makkelijker te benaderen dan ViT-B/32 (waarschijnlijk door fijnere patch-granulariteit).
• Classificatieprestaties bij Vervanging:
  • [CLS]-tokens: Vervanging van alle MLP-lagen door CLTs voor alleen de [CLS]-token resulteert in prestaties die bijna identiek zijn aan het baseline-model (bijv. 61,74% vs 61,65% op CIFAR-100).
  • Patch-tokens: Vervanging van alle lagen voor patch-tokens leidt tot een daling in nauwkeurigheid, vooral bij vroege lagen. Echter, vervanging van alleen de laatste lagen (bijv. 10→11) behoudt de nauwkeurigheid.
  • ViT-B/16 toont over het algemeen meer robustheid tegen vervanging dan ViT-B/32.
• Interpretatie: De heatmaps van de bijdragescijfers bevestigen dat [CLS]-tokens semantische signalen over de diepte integreren, terwijl patch-tokens lokaal opereren. Ablatie-studies bevestigen dat de attributie-scores "faithful" (geloofwaardig) zijn: het verwijderen van belangrijke lagen degradeert de prestaties, terwijl het behouden van een kleine subset de prestaties behoudt.

Significantie en Toekomst

Dit werk positioneert CLTs als een principieel kader voor gestructureerde, cross-layer interpretatie van visuele modellen zonder in te leveren op prestaties.

• Inzicht: Het onthult fundamentele architecturale verschillen tussen hoe [CLS]-tokens en patch-tokens informatie verwerken in ViTs.
• Toekomstige Richtingen: De auteurs suggereren drie richtingen:
  1. Hiërarchische conceptontdekking (hoe laag-niveau features samenkomen tot semantische abstracties).
  2. Uitbreiding van CLTs om ook attention-heads te modelleren voor een volledig circuit-analyse.
  3. Toepassing op bredere visuele taken en Multimodal Large Language Models (MLLMs).

Kortom, het paper bewijst dat Cross-Layer Transcoders niet alleen een functionele vervanging kunnen zijn voor ViT-blokken, maar ook een krachtig instrument bieden om de "black box" van Vision Transformers te openen en te begrijpen.