Explaining Attention with Program Synthesis

Probleemstelling
Het begrijpen van de berekeningen die door diepe neurale netwerken worden uitgevoerd in algoritmische termen blijft een belangrijke uitdaging in machine learning. Eerdere benaderingen van interpreteerbaarheid hebben grotendeels vertrouwd op het toekennen van betekenis aan neuronen of kenmerken, ofwel top-down (via probes voor door mensen gedefinieerde concepten), ofwel bottom-up (via input/output samenvattingen). Hoewel krachtig, schieten deze methoden vaak tekort in het bieden van een volledige, formele beschrijving van neurale computatie. Bov�elijk kan het labelen van intermediaire componenten met natuurlijke taalbeschrijvingen leiden tot interpretaties die ambigu of moeilijk te formaliseren zijn. De auteurs stellen dat er een behoefte is aan een medium dat de kloof overbrugt tussen de complexiteit van miljardenparametersmodellen en natuurlijke taalverklaringen—een medium dat zowel menselijk leesbaar als formeel verifieerbaar is.

Methodologie
Het artikel stelt een raamwerk voor gebaseerd op programmasynthese om de berekeningen van attention heads in transformer taalmodellen uit te leggen. Het kerndoel is het vinden van een uitvoerbaar Python-programma ($\pi$) dat een input tekstsequentie ($X$) direct kan mappen naar een attention matrix ($A$), waarmee de logica van een specifieke attention head effectief wordt benaderd.

Het raamwerk verloopt via vier afzonderlijke stappen:

1. Extractie van Attention Maps: De auteurs extraheren grondwaarheid attention activatiematrices van een doelmodel over een corpus van trainingssequenties (specifiek gebruikmakend van de TinyStories dataset vanwege de gestructureerde eenvoud). Ze filteren deze maps om de top 2,5% van de attention gewichten te isoleren om de meest relevante token-paar interacties te focussen.
2. Programmasynthese: Een hulp-Large Language Model (LLM), optredend als synthese-agent, wordt geprompt met de geëxtraheerde attention patronen en input sequenties. De agent krijgt de taak om uitvoerbare Python-functies te genereren die deze patronen reproduceren. De synthese-agent heeft toegang tot bibliotheken zoals NumPy, spaCy en NLTK voor linguïstische en numerieke verwerking.
3. Verfijning en Rangschikking: Kandidatenprogramma's worden gevalideerd op syntaxis en gescoord tegen echte attention patronen met behulp van de Jensen-Shannon Afstand (JSD). De beste kandidaten ondergaan een feedback-geconditioneerde verfijningslus waarbij de synthese-agent echte en voorspelde patronen contrasteert om herziene programma's te produceren.
4. Selectie: Het optimale proxy-programma voor elke head wordt geselecteerd door de gelijkenis te maximaliseren (gemeten via Intersection over Union, of IoU) op uitgesloten data-sequenties.

Belangrijkste Bijdragen

• Programmatische Interpreteerbaarheid: De auteurs introduceren een methode om neurale attention heads te vervangen door symbolische, uitvoerbare code, wat een formeel verifieerbaar alternatief biedt voor natuurlijke taalbeschrijvingen.
• Causale Validatie via Vervanging: In tegenstelling tot eerder werk dat alleen correlaties observeert, valideert deze aanpak verklaringen door geleerde neurale componenten direct te substitueren met gesynthetiseerde programma's tijdens de forward pass van het model.
• Schaalbaarheid: De methode wordt toegepast op vier verschillende transformer architecturen: BERT-Base, GPT-2-Small, TinyLlama-1.1B en Llama-3B, wat de toepasbaarheid van de aanpak over verschillende modelgroottes en architecturen (bidirectioneel en causaal) aantoont.

Resultaten

• Hoge Alignement: Een aanzienlijk deel van de attention heads over alle geteste modellen kan worden benaderd door uitvoerbare programma's met een hoge nauwkeurigheid. De best presterende programma's bereiken gemiddelde IoU-scores van 69% voor GPT-2, 74% voor TinyLlama en 79% voor Llama-3B.
• Modelschaal en Architectuur: De auteurs observeren dat attention heads in autoregressieve (decoder) modellen over het algemeen gemakkelijker te fitten zijn dan die in bidirectionele (encoder) modellen zoals BERT. Daarnaast neemt de kwaliteit van de fit toe met de model schaal, wat suggereert dat grotere modellen mogelijk een grotere functionele specialisatie vertonen onder hun heads.
• Causale Getrouwheid: Wanneer attention heads worden vervangen door hun gesynthetiseerde programma's, behouden de modellen hun functionele prestaties. Specifiek kunnen tot 25% van de heads worden vervangen met slechts een toename van 16% in perplexiteit.
• Downstream Prestaties: Het vervangen van 30–40% van de attention heads met hun best-fit programma's leidt niet tot een significante verslechtering van de prestaties op downstream vraag-antwoord benchmarks (inclusief HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA en COPA).
• Correlatie: Er is een sterke negatieve correlatie (Spearman $r > 0.9$) tussen de IoU-gelijkenis van een programma en de toename in perplexiteit bij head vervanging, wat aangeeft dat IoU een betrouwbare proxy is voor het selecteren van functioneel getrouwe programma's.

Betekenis en Claims
Het artikel beweert dat zelfs in state-of-the-art taalmodellen een aanzienlijk deel van de attention patronen in symbolische termen begrepen kan worden. De primaire betekenis ligt in de demonstratie dat getrainde neurale componenten kunnen worden vervangen door symbolische surrogaten via uitvoerbare code zonder het gedrag van het model substantieel te veranderen. Dit suggereert dat moderne LM-gestuurde programmasynthese-methoden een levensvatbaar raamwerk bieden voor het benaderen van de bredere vraag hoe diepe netwerken opereren.

De auteurs positioneren dit werk als een stap naar een volledige symbolische karakterisering van transformers, waarbij modelgedrag begrepen, gemodificeerd en getest kan worden via logica in plaats van gewichtsmanipulatie. Het artikel blijft echter bescheiden over de beperkingen en erkent dat geen enkel model momenteel volledig gekarakteriseerd is (veel heads behalen nog steeds een lage IoU-score) en dat de huidige programmalibraire, hoewel effectief, nog niet uitputtend is. Het werk legt een fundament voor toekomstig onderzoek gericht op het bereiken van een volledige symbolische distillatie van transformer architecturen.