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

1. Il Problema

I modelli di fondazione (Foundation Models), in particolare i Vision Transformers (ViT), hanno dimostrato prestazioni eccezionali in compiti di visione artificiale. Tuttavia, rimangono "scatole nere" (opachi), il che limita la loro interpretabilità, controllabilità e affidabilità.
Le tecniche esistenti per l'interpretazione, come gli Sparse Autoencoders (SAE), operano tipicamente a livello di singolo strato. Questo approccio locale fallisce nel catturare:

• La struttura computazionale cross-layer (tra diversi livelli) dei Transformer.
• Il contributo relativo di ogni strato nella formazione della rappresentazione finale.
  Di conseguenza, le metodologie attuali offrono intuizioni limitate su come l'informazione venga trasformata attraverso la profondità della rete.

2. Metodologia: Cross-Layer Transcoders (CLT)

Gli autori introducono i Cross-Layer Transcoders (CLT) come modelli proxy affidabili, sparsi e consapevoli della profondità per sostituire i blocchi MLP (Multi-Layer Perceptron) nei ViT.

• Architettura: I CLT utilizzano uno schema encoder-decoder.

  • Encoder: Codifica le attivazioni post-attention (pre-MLP) degli strati precedenti in vettori di embedding sparsi ($z_i$). Vengono supportate diverse funzioni di non-linearità per la sparsità: JumpReLU, ReLU-Top-k e Abs-Top-k.
  • Decoder: Ricostruisce le attivazioni post-MLP di uno strato target ($y_j$) come una somma lineare dei contributi decodificati da tutti gli strati sorgente precedenti ($i \le j$).
  • Formula di Ricostruzione: $\hat{y}_j = \sum_{i=1}^{j} z_i W_{i \to j}^{dec}$.
    Questa struttura triangolare impone un vincolo causale rigoroso: uno strato target può essere ricostruito solo dalle informazioni degli strati precedenti, prevenendo la fuoriuscita di informazioni dai livelli futuri.

• Addestramento: I CLT vengono addestrati per ricostruire le attivazioni post-MLP di un ViT "maestro" (frozen) utilizzando le attivazioni pre-MLP degli strati precedenti. L'obiettivo è minimizzare l'errore di ricostruzione (MSE) mantenendo la sparsità delle feature.

• Sostituzione Funzionale: Una volta addestrati, i CLT possono sostituire i blocchi MLP originali durante l'inferenza, creando un "modello di sostituzione" che mantiene le prestazioni del modello originale ma offre una decomposizione lineare delle attivazioni.

3. Contributi Chiave

1. CLT per la Visione: È il primo studio che applica con successo i Cross-Layer Transcoders ai Vision Transformers, permettendo una ricostruzione sparsa e cross-layer delle attivazioni.
2. Sostituzione Accurata dei Blocchi MLP: I CLT possono sostituire i blocchi MLP (specialmente negli strati finali o per il token [CLS]) senza degradare le prestazioni di classificazione zero-shot, e in alcuni casi migliorandole. Questo li rende proxy affidabili per l'analisi interna.
3. Attribuzione Cross-Layer Interpretabile: La struttura lineare dei CLT genera punteggi di attribuzione che quantificano quanto ogni strato sorgente contribuisce alla rappresentazione di uno strato target.
4. Scoperta di Asimmetrie Strutturali: L'analisi rivela differenze fondamentali tra il token globale [CLS] e i token delle patch:
   • I token patch mostrano un'attribuzione fortemente diagonale-dominante (ogni strato spiega principalmente se stesso).
   • Il token [CLS] mostra un'attribuzione distribuita attraverso la profondità, aggregando informazioni da molti strati precedenti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su CLIP ViT-B/32 e ViT-B/16 utilizzando i dataset CIFAR-100, COCO e ImageNet-100.

• Fideltà della Ricostruzione: I CLT raggiungono un'alta fedeltà nella ricostruzione delle attivazioni post-MLP, con similarità coseno tra 0.92 e 0.97 e valori $R^2$ tra 0.89 e 0.95. Gli strati iniziali e finali sono ricostruiti quasi perfettamente, mentre gli strati intermedi mostrano errori leggermente superiori.
• Prestazioni di Classificazione (Zero-Shot):
  • La sostituzione dei blocchi MLP solo per il token [CLS] mantiene l'accuratezza quasi identica al modello originale (es. su ImageNet-100, l'accuratezza rimane sopra l'84% anche sostituendo tutti gli strati).
  • La sostituzione completa (tutti gli strati e tutti i token) porta a un calo delle prestazioni, specialmente per i token patch negli strati iniziali, evidenziando la difficoltà di approssimare le trasformazioni spaziali precoci.
• Validazione dell'Attribuzione (Ablazione):
  • Rimuovere lo strato sorgente con il punteggio di attribuzione più alto causa un calo significativo dell'accuratezza (es. -5.6% su ImageNet-100).
  • Mantenere solo i top-4 strati sorgenti (quelli con i punteggi più alti) è sufficiente per recuperare quasi completamente l'accuratezza originale. Questo conferma che la rappresentazione finale è concentrata in un piccolo sottoinsieme di strati dominanti.
• Visualizzazione: Un framework di recupero basato sui codici sparsi mostra che l'evoluzione semantica del token [CLS] passa da statistiche di basso livello (colore, layout) negli strati iniziali a concetti semantici ad alto livello (attività, oggetti) negli strati finali.

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo verso l'interpretabilità dei modelli di visione:

• Superamento dei Limiti Locali: Dimostra che l'analisi delle reti neurali deve considerare le dipendenze tra gli strati, non solo le trasformazioni intra-strato.
• Proxy Affidabile: I CLT offrono un metodo per analizzare i meccanismi interni dei ViT senza comprometterne le prestazioni operative, permettendo studi di causalità e attribuzione.
• Nuove Intuizioni Architetturali: La distinzione tra il comportamento "diagonale" delle patch e quello "distribuito" del [CLS] fornisce una comprensione più profonda di come i ViT integrino informazioni locali e globali.
• Futuro: Apre la strada a scoperte di concetti gerarchici, analisi di circuiti completi (inclusi i meccanismi di attenzione) e l'applicazione di questi proxy a modelli multimodali più ampi.

In sintesi, i Cross-Layer Transcoders si posizionano come un framework principato per l'interpretazione strutturata e cross-layer dei modelli di visione, trasformando rappresentazioni opache in costruzioni additive e risolvibili per strato.