Explaining Attention with Program Synthesis

Problema
Comprendere i calcoli eseguiti dalle reti neurali profonde in termini algoritmici rimane una sfida significativa nel machine learning. Gli approcci precedenti all'interpretabilità si sono ampiamente basati sull'assegnazione di un significato ai neuroni o alle feature, sia in modo top-down (tramite probe per concetti definiti dall'uomo) sia bottom-up (tramite riassunti di input/output). Sebbene potenti, questi metodi spesso non riescono a fornire una descrizione completa del calcolo neurale. Inoltre, etichettare le componenti intermedie con descrizioni in linguaggio naturale può dare origine a interpretazioni ambigue o difficili da formalizzare. Gli autori sostengono che sia necessario un mezzo che colmi il divario tra la complessità dei modelli da miliardi di parametri e le spiegazioni in linguaggio naturale—un mezzo che sia allo stesso tempo leggibile dall'uomo e formalmente verificabile.

Metodologia
Il documento propone un framework basato sulla sintesi di programmi per spiegare i calcoli delle teste di attenzione nei modelli linguistici transformer. L'obiettivo principale è trovare un programma Python eseguibile ($\pi$) che possa mappare direttamente una sequenza di testo in input ($X$) verso una matrice di attenzione ($A$), approssimando efficacemente la logica di una specifica testa di attenzione.

Il framework procede in quattro fasi distinte:

1. Estrazione della mappa di attenzione: Gli autori estraggono le matrici di attivazione dell'attenzione reali da un modello target attraverso un corpus di sequenze di addestramento (specificamente utilizzando il dataset TinyStories per la sua semplicità strutturata). Filtrano queste mappe per isolare il top 2,5% dei pesi di attenzione, al fine di concentrarsi sulle interazioni salienti tra coppie di token.
2. Sintesi del programma: Un LLM ausiliario, che funge da agente di sintesi, viene istruito con i pattern di attenzione estratti e le sequenze di input. L'agente ha il compito di generare funzioni Python eseguibili che riproducano tali pattern. L'agente di sintesi ha accesso a librerie come NumPy, spaCy e NLTK per l'elaborazione linguistica e numerica.
3. Raffinamento e Classificazione: I programmi candidati vengono validati per la sintassi e valutati rispetto ai pattern di attenzione reali utilizzando la Distanza di Jensen-Shannon (JSD). I migliori candidati sono sottoposti a un ciclo di raffinamento condizionato dal feedback, in cui l'agente di sintesi confronta i pattern reali e quelli predetti per produrre programmi rivisti.
4. Selezione: Il programma proxy ottimale per ogni testa viene selezionato massimizzando la somiglianza (misurata tramite l'Intersezione su Unione, o IoU) su sequenze di dati non utilizzate per l'addestramento (held-out).

Contributi Chiave

• Interpretabilità Programmatica: Gli autori introducono un metodo per sostituire le teste di attenzione neurali con codice simbolico ed eseguibile, offrendo un'alternativa formalmente verificabile alle descrizioni in linguaggio naturale.
• Validazione Causale tramite Sostituzione: A differenza del lavoro precedente che osserva solo le correlazioni, questo approccio valida le spiegazioni sostituendo direttamente le componenti neurali apprese con i programmi sintetizzati durante il forward pass del modello.
• Scalabilità: Il metodo è applicato a quattro diverse architetture transformer: BERT-Base, GPT-2-Small, TinyLlama-1.1B e Llama-3B, dimostrando l'applicabilità dell'approccio a diverse dimensioni e architetture di modelli (bidirezionali e causali).

Risultati

• Alto Allineamento: Una frazione sostanziale di teste di attenzione in tutti i modelli testati può essere approssimata da programmi eseguibili con alta accuratezza. I programmi con le migliori prestazioni raggiungono punteggi medi di IoU del 69% per GPT-2, 74% per TinyLlama e 79% per Llama-3B.
• Scala del Modello e Architettura: Gli autori osservano che le teste di attenzione nei modelli autoregressivi (decoder) sono generalmente più facili da adattare rispetto a quelle nei modelli bidirezionali (encoder) come BERT. Inoltre, la qualità dell'adattamento aumenta con la scala del modello, suggerendo che i modelli più grandi possano esibire una maggiore specializzazione funzionale tra le loro teste.
• Fedeltà Causale: Quando le teste di attenzione vengono sostituite con i loro programmi sintetizzati, i modelli mantengono la performance funzionale. Nello specifico, fino al 25% delle teste può essere sostituito con un aumento della perplexity di solo il 16%.
• Performance sui Task a Valle (Downstream): Sostituire dal 30% al 40% delle teste di attenzione con i loro programmi migliori non degrada significativamente le prestazioni sui benchmark di question-answering (inclusi HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA e COPA).
• Correlazione: Esiste una forte correlazione negativa (Spearman $r > 0,9$) tra la somiglianza IoU di un programma e l'aumento della perplexity in seguito alla sostituzione della testa, indicando che l'IoU è un proxy affidabile per selezionare programmi funzionalmente fedeli.

Significato e Rivendicazioni
Il documento sostiene che, anche nei modelli linguistici allo stato dell'arte, una frazione sostanziale di pattern di attenzione può essere compresa in termini simbolici. La portata principale risiede nella dimostrazione che componenti neurali addestrate possono essere sostituite da surrogati simbolici tramite codice eseguibile senza alterare sostanzialmente il comportamento del modello. Ciò suggerisce che i moderni metodi di sintesi di programmi guidati da LLM offrono un framework vitale per affrontare la questione più ampia di come operano le reti profonde.

Gli autori pongono questo lavoro come un passo verso una completa caratterizzazione simbolica dei transformer, dove il comportamento del modello può essere ragionato, modificato e testato direttamente attraverso la logica piuttosto che tramite la manipolazione dei pesi. Tuttavia, il documento rimane modesto riguardo ai propri limiti, riconoscendo che nessuna testa di attenzione di un modello è attualmente completamente caratterizzata (molte teste ottengono ancora punteggi IoU bassi) e che la libreria di programmi attuale, sebbene efficace, non è ancora esaustiva. Il lavoro stabilisce una base per la ricerca futura volta a raggiungere una completa distillazione simbolica delle architetture transformer.