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

1. Problemstellung

Das Verständnis der internen Aktivierungen von Vision Transformern (ViTs) ist entscheidend für die Entwicklung interpretierbarer und vertrauenswürdiger Modelle. Bisherige Ansätze wie Sparse Autoencoder (SAEs) arbeiten zwar effektiv, um menschlich interpretierbare Merkmale zu extrahieren, haben jedoch zwei wesentliche Einschränkungen:

1. Lokale Beschränkung: SAEs operieren typischerweise nur innerhalb einzelner Schichten und erfassen nicht die kreuzschichtige (cross-layer) Rechenstruktur von Transformern.
2. Fehlende Kontextualisierung: Sie können die relative Bedeutung jeder Schicht bei der Bildung der finalen Repräsentation (insbesondere der letzten Schicht) nicht angemessen abbilden.

Die Autoren fragen sich, ob es möglich ist, die komplexen MLP-Blöcke (Multi-Layer Perceptrons) in ViTs durch ein spärliches, tiefenbewusstes Modell zu ersetzen, das nicht nur die Funktion übernimmt, sondern auch eine attributionale Interpretierbarkeit über die gesamte Netzwerk-Tiefe hinweg bietet.

2. Methodik: Cross-Layer Transcoder (CLT)

Die Autoren stellen Cross-Layer Transcoder (CLTs) als neue, zuverlässige Proxy-Modelle für die MLP-Blöcke in ViTs vor.

• Architektur: CLTs verwenden ein Encoder-Decoder-Schema.

  • Encoder: Jede Schicht $i$ wandelt die Aktivierungen (nach Layer-Normalisierung 2, vor dem MLP) in einen dünnbesetzten (sparse) Code $z_i$ um. Dies geschieht durch lineare Projektion gefolgt von einer Nichtlinearität (JumpReLU, ReLU-Top-k oder Abs-Top-k).
  • Decoder: Um die Ausgabe eines MLPs in einer Ziel-Schicht $j$ zu rekonstruieren, nutzt der Decoder die spärlichen Codes aller vorhergehenden Schichten ($i \le j$).
  • Rekonstruktion: Die Ausgabe $\hat{y}_j$ wird als additive Summe der Beiträge aller vorherigen Schichten berechnet:
    $$\hat{y}_j = \sum_{i=1}^{j} z_i W_{dec}^{i \to j}$$
    Dies erzwingt eine strikt kausale Struktur (keine Information aus zukünftigen Schichten) und bildet ein dreieckiges Decoder-Architektur.

• Training: Die CLTs werden auf den Aktivierungen eines eingefrorenen Lehrmodells (Teacher Model, z.B. CLIP ViT-B/32 oder B/16) trainiert. Das Ziel ist die Minimierung des Rekonstruktionsfehlers (MSE) zwischen dem tatsächlichen MLP-Ausgang und der CLT-Rekonstruktion unter Berücksichtigung eines Sparsity-Regulierungsterms.

• Anwendung als Proxy: Nach dem Training können die MLP-Blöcke im Original-ViT durch die trainierten CLTs ersetzt werden. Dies ermöglicht zwei Dinge:

  1. Funktionaler Ersatz: Das Modell behält seine Leistungsfähigkeit bei.
  2. Interpretierbarkeit: Da die Ausgabe als Summe von Beiträgen dargestellt wird, können die Gewichte $W_{dec}^{i \to j}$ genutzt werden, um zu quantifizieren, wie stark jede Quell-Schicht $i$ zur Ziel-Schicht $j$ beiträgt.

3. Wichtige Beiträge

1. CLTs für Vision: Die Arbeit untersucht erstmals die Wirksamkeit von Cross-Layer Transcodern in Vision Transformern und ermöglicht eine spärliche, kreuzschichtige Rekonstruktion von MLP-Aktivierungen.
2. Präziser Ersatz von MLPs: CLTs können MLP-Blöcke, insbesondere in späteren Schichten oder für den globalen [CLS]-Token, ersetzen, ohne die Zero-Shot-Klassifizierungsleistung zu beeinträchtigen (in einigen Fällen sogar Verbesserung).
3. Interpretierbarkeit durch Kreuzschicht-Zuordnung: Die Projektions-basierten Beitragswerte zeigen eine signifikante Asymmetrie:
   • Patch-Tokens: Zeigen eine stark diagonale Dominanz (jede Schicht erklärt primär ihren eigenen Output).
   • [CLS]-Token: Zeigt eine breite Verteilung über die Tiefe; die finale Repräsentation ist eine Aggregation von semantischen Signalen aus vielen vorherigen Schichten.
4. Vertrauenswürdige Attribution: Ablationsstudien belegen, dass die finalen Repräsentationen stark von einer kleinen Teilmenge dominanter Quell-Schichten abhängen. Das Entfernen der am höchsten bewerteten Schicht verschlechtert die Leistung drastisch, während das Beibehalten der Top-4-Schichten die ursprüngliche Genauigkeit weitgehend wiederherstellt.

4. Ergebnisse

Die Autoren trainierten CLTs auf CLIP ViT-B/32 und ViT-B/16 mit den Datensätzen CIFAR-100, COCO und ImageNet-100.

• Rekonstruktionsqualität: CLTs erreichen eine hohe Rekonstruktionsgenauigkeit (Cosine-Similarity zwischen 0,92 und 0,97, $R^2$ zwischen 0,89 und 0,95). ViT-B/16 lässt sich aufgrund kleinerer Patches und einfacherer pro-Token-Aktivierungen besser approximieren als ViT-B/32.
• Funktionaler Ersatz (Zero-Shot Accuracy):
  • Der Ersatz von MLPs nur für den [CLS]-Token über alle Schichten hinweg führt zu einer Genauigkeit, die der des Baseline-Modells entspricht oder diese sogar leicht übertrifft.
  • Der vollständige Ersatz aller Schichten und aller Tokens (Patches + CLS) führt zu einem Leistungsabfall, insbesondere bei frühen Schichten und Patch-Tokens.
  • Der Ersatz der letzten Schichten (z.B. 10→11) ist für alle Token-Typen sicher und verlustfrei.
• Interpretierbarkeit:
  • Die Analyse der Beitragswerte zeigt, dass der [CLS]-Token Informationen aus dem gesamten Netzwerk integriert, während Patch-Tokens lokalere Transformationen bevorzugen.
  • Ablationsexperimente bestätigen, dass die von CLTs identifizierten „wichtigsten" Schichten tatsächlich funktional kritisch sind. Das Entfernen der Top-1-Schicht senkt die Genauigkeit signifikant (z.B. von 80,54% auf 74,94% auf ImageNet-100), während das Beibehalten der Top-4-Schichten die Leistung fast vollständig erhält.

5. Bedeutung und Fazit

Diese Arbeit demonstriert, dass Cross-Layer Transcoder eine vielversprechende Alternative zu herkömmlichen Methoden wie SAEs oder per-layer Transcodern darstellen. Sie bieten einen prinzipiellen Rahmen für strukturierte, kreuzschichtige Interpretierbarkeit in Vision-Modellen, ohne die Leistungsfähigkeit der Modelle zu beeinträchtigen.

Die Haupterkenntnis ist, dass die finale Repräsentation eines ViT nicht isoliert in der letzten Schicht entsteht, sondern als aggregiertes Ergebnis über die gesamte Tiefe des Netzwerks, wobei der [CLS]-Token diese Integration übernimmt. CLTs machen diese interne Dynamik sichtbar und quantifizierbar, was neue Wege für das Verständnis, die Kontrolle und die Verifizierung von Foundation Models im visuellen Bereich eröffnet. Zukünftige Arbeiten könnten diese Methode auf Aufmerksamkeitsköpfe (Attention Heads) erweitern und auf multimodale Large Language Models anwenden.