Disentangling MLP Neuron Weights in Vocabulary Space

Titolo: Disentangling MLP Neuron Weights in Vocabulary Space

Autore: Asaf Avrahamy, Yoav Gur-Arieh, Mor Geva (Tel Aviv University)
Metodo Proposto: ROTATE (Rotation-Optimized Token Alignment in weighT spacE)

---

1. Il Problema: Poliseanticità e Interpretazione dei Pesi

L'interpretazione meccanica dei modelli linguistici (LM) si scontra con il problema della poliseanticità: la maggior parte dei neuroni nei livelli MLP (Multi-Layer Perceptron) dei transformer non codifica un singolo concetto coerente, ma mescola (entanglement) molteplici concetti in un'unica rappresentazione vettoriale.
Le approcci esistenti, come gli Autoencoder Sparsi (SAE), richiedono grandi quantità di dati di attivazione (forward pass) per decomporre i neuroni in unità monosemantiche. Inoltre, questi metodi analizzano le attivazioni in uno spazio latente che non è direttamente allineato con il vocabolario del modello, rendendo difficile l'interpretazione umana immediata.
L'obiettivo di questo lavoro è sviluppare un metodo senza dati (data-free) che operi direttamente nello spazio dei pesi, disentanglando i neuroni MLP in direzioni interpretabili e allineate al vocabolario, senza richiedere passaggi in avanti sul modello.

2. Metodologia: ROTATE

Il cuore della proposta è l'osservazione statistica che i neuroni che codificano concetti coerenti e monosemantici mostrano un'alta curtosi quando i loro vettori di peso sono proiettati nello spazio del vocabolario del modello.

Concetti Chiave:

• Proiezione nel Vocabolario: Data una matrice di embedding non codificata $U$, un vettore di peso $w$ di un neurone viene proiettato come $z = wU$. Se $z$ è una distribuzione "appuntita" con code pesanti (alta curtosi), significa che il neurone promuove o sopprime fortemente un insieme sparso di token specifici, suggerendo un comportamento monosemantico.
• Ottimizzazione delle Rotazioni: ROTATE cerca di trovare una matrice di rotazione $R$ tale che il vettore ruotato $v = wR$ massimizzi la curtosi della sua proiezione nel vocabolario $z = vU$.
• Funzione di Perdita: L'ottimizzazione minimizza una funzione di perdita composta da due termini:
  1. Perdita di Curtosi ($L_{kurt}$): Massimizza la curtosi della distribuzione dei logit.
  2. Perdita di Regolarizzazione ($L_{reg}$): Penalizza la distanza coseno tra il vettore ruotato $v$ e il vettore originale $w$, assicurando che i canali scoperti rimangano fedeli al neurone originale e non convergano verso direzioni arbitrarie.
     $$L = -\lambda \cdot \log(1 + \text{Kurt}(z)) + (1 - \frac{w \cdot v}{\|w\|\|v\|})$$
• Parametrizzazione: Le rotazioni sono implementate tramite matrici di Householder, che garantiscono l'ortogonalità per costruzione.
• Algoritmo Iterativo con Mascheramento: Poiché un neurone può contenere più concetti, ROTATE esegue l'ottimizzazione in modo iterativo. Dopo aver trovato un canale (direzione ad alta curtosi), i token che contribuiscono maggiormente a quella curtosi vengono mascherati (esclusi) per forzare le iterazioni successive a scoprire nuove direzioni semantiche distinte. Questo evita che l'algoritmo converga ripetutamente sullo stesso concetto dominante.

3. Contributi Chiave

1. Osservazione Statistica: Dimostrazione empirica che l'alta curtosi nelle proiezioni vocabolari è un segnale affidabile di direzioni monosemantiche nello spazio dei pesi.
2. Metodo Data-Free: Introduzione di ROTATE, un algoritmo che disentangla i pesi dei neuroni MLP senza alcun dato di input o passaggio in avanti, rendendolo scalabile ed efficiente.
3. Canali Vocabolari (Vocabulary Channels): Recupero di direzioni sparse e interpretabili che corrispondono a concetti specifici (es. "temporal connectors", "negative numbers").
4. Superiorità rispetto agli SAE: Dimostrazione che i canali derivati dai pesi superano gli SAE (addestrati su attivazioni) sia in fedeltà che in completezza.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Gemma-2-2B-it e Llama-3.1-8B-Instruct.

A. Fedeltà (Faithfulness)

• Input-side: I canali scoperti da ROTATE attivano il neurone corrispondente in modo più affidabile rispetto ai baselines basati su SAE (Gemma Scope, Llama Scope).
  • Punteggi: ROTATE ottiene tra 0.46 e 0.71, contro 0.33–0.49 degli SAE.
• Output-side (Causalità): L'ablazione di un singolo canale riduce drasticamente la frequenza dei token associati a quel concetto (riduzione del ~90%), senza influenzare i token degli altri canali dello stesso neurone. Questo conferma che i canali sono causalmente disaccoppiati.

B. Completezza (Completeness)

• La valutazione misura quanto bene l'insieme dei canali scoperti copre lo spazio di attivazione del neurone.
• Un giudice LLM è stato in grado di associare correttamente l'input attivante alla descrizione del canale ROTATE nel 49-60% dei casi, superando significativamente gli SAE (36-49%) e il caso casuale (20%).

C. Descrizione dei Neuroni

• Aggregando le descrizioni dei singoli canali, ROTATE genera descrizioni testuali dei neuroni superiori ai metodi basati sulle attivazioni massime (MaxAct++).
• In confronti testa a testa, le descrizioni ROTATE hanno vinto il 60-69% delle volte contro MaxAct++ e il 52-65% contro MaxAct+VocabProj, specialmente su attivazioni moderate (non solo sui picchi estremi), dimostrando una migliore copertura del comportamento del neurone.

5. Significato e Implicazioni

• Scalabilità: Essendo un metodo basato solo sui pesi, ROTATE può essere applicato a qualsiasi modello LM senza la necessità di raccogliere e processare terabyte di dati di attivazione, riducendo drasticamente i costi computazionali.
• Interpretabilità Fine-Grained: Fornisce una scomposizione meccanica precisa, rivelando come un singolo neurone possa gestire concetti opposti (es. polarità positiva vs. temporale) a seconda del segnale di ingresso, grazie alla separazione dei canali di input ($w_{gate}$, $w_{in}$) e output ($w_{out}$).
• Nuovo Paradigma: Sposta il focus dall'analisi delle attivazioni (dinamica) all'analisi strutturale dei pesi (statica), suggerendo che la struttura monosemantica è intrinsecamente codificata nella geometria dei pesi e può essere estratta matematicamente.

In sintesi, ROTATE offre un mattone fondamentale per l'interpretabilità meccanica, permettendo di "leggere" direttamente i neuroni dei modelli linguistici come insiemi di concetti vocabolari discreti, senza dipendere dai dati di addestramento o inferenza.