Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

Probleemstelling

Het paper adresseert de uitdaging om uitvoerbare en interpreteerbare algoritmen direct te extraheren uit getrainde Transformer-modellen. Hoewel Large Language Models (LLMs) succesvol zijn in het genereren van code, blijft het "black-box"-karakter van standaard Transformers een belemmering voor het ontdekken van nieuwe algoritmen (de novo discovery).

De kernproblemen zijn:

1. Representatieverstrengeling (Superposition): In standaard Transformers worden features gecodeerd in overlappende, niet-orthogonale richtingen in een continue, hoogdimensionale ruimte. Dit maakt het onmogelijk om directe symbolische expressies te herleiden uit de interne representaties.
2. Lek tussen continu en discreet: Symbolische algoritmen zijn inherent discreet en spaarzaam, terwijl Transformers continue dynamiek gebruiken. Bestaande methoden voor algoritme-extractie (zoals MIPS voor RNN's) zijn beperkt tot discrete latent-variabelen en kunnen geen continue variabelen (zoals zwevende komma-getallen) effectief verwerken.

Methodologie: De Discrete Transformer

De auteurs stellen de Discrete Transformer voor, een architectuur die specifiek is ontworpen om de kloof tussen continue optimalisatie en discrete symbolische logica te overbruggen. De aanpak bestaat uit drie hoofdfasen:

1. Architectuur

De architectuur volgt het RASP-computatiemodel (Restricted Access Sequence Processing) en enforceert strikte functionele ontkoppeling:

• Numerical Residual Stream: In plaats van dichte vectorrepresentaties, wordt de residual stream gedefinieerd als een verzameling expliciete scalair variabelen. Updates gebeuren via concatenatie in plaats van optelling, wat superposition voorkomt.
• Numerical Attention (Router): Deze module fungeert uitsluitend als een routeringsoperator. Het gebruikt Piecewise Linear Encoding (PLE) om scalars naar een hogere dimensie te projecteren voor dot-product berekeningen, maar de output is een deterministische pointer (via "hard attention") die specifieke waarden uit de geschiedenis ophaalt.
• Numerical MLP (Operator): Deze module is verantwoordelijk voor element-wise rekenkundige en logische transformaties. Het is opgedeeld in parallelle, onafhankelijke sub-modules die elk een eenvoudige operatie uitvoeren. Door de input-dimensie en verborgen breedte te beperken, wordt complexiteit geforceerd tot een reeks eenvoudige stappen.
• Temperatuur-gecontroleerde Discretisatie: Om differentiabele training mogelijk te maken, wordt een temperatuurparameter ($\tau$) gebruikt in steekproefmechanismen (sampling). Tijdens training wordt $\tau$ geleidelijk verlaagd (annealing) van een hoge waarde (exploratie, continue verdeling) naar 0 (exploitatie, deterministische one-hot selectie). Hierdoor evolueert het model van een continu netwerk naar een volledig discreet computationeel graf.

2. Extractie Pipeline (Van Gewichten naar Code)

Zodra het model convergeert naar een discrete staat, wordt het algoritme geëxtraheerd via een modulaire aanpak:

• Hypothese-toetsing voor Attention: De Attention-mechanismen worden geanalyseerd op hun routeringspatronen. De auteurs testen hypothesen voor patronen zoals "Fixed Offset" (relatieve posities) en "Windowed Extrema" (inhoudsafhankelijke selectie).
• Symbolische Regressie voor MLP: Voor de Numerical MLP-submodules wordt symbolische regressie (met PySR) toegepast op de input-output paren om de exacte wiskundige expressies te vinden (bijv. $y = 2x_1 + x_2$).
• Globale Samenstelling: De geëxtraheerde primitieven (pointers en formules) worden geïntegreerd via een lineaire output head. Door achterwaartse traversie van het computationele graf en het toepassen van symbolische vereenvoudiging, wordt een leesbaar Python-programma gegenereerd.

Kernbijdragen

1. Discrete Transformer Architectuur: Een RASP-gealigneerde architectuur die functionele ontkoppeling forceert, waardoor continue optimalisatie kan worden omgezet in discrete logica.
2. Modulaire Extractie Pipeline: Een nieuwe methode die hypothese-toetsing en symbolische regressie combineert om mechanische interpretatie uit te breiden naar continue variabelen, in tegenstelling tot eerdere methoden die beperkt waren tot discrete latent-variabelen.
3. Controleerbaarheid via Inductieve Bias: Het paper toont aan dat architecturale beperkingen (zoals het aantal sub-MLP's of attention heads) gebruikt kunnen worden om het zoekruimte te sturen en specifieke vormen van algoritmen te forceren die overeenkomen met domeinkennis.

Resultaten

De auteurs evalueren het model op een diverse set van algoritmische taken, waaronder lineaire rekenkunde, niet-lineaire composities en continue dynamica.

• Prestatie: De Discrete Transformer bereikt een voorspellende prestatie (MSE) die vergelijkbaar is met of beter is dan de RNN-gebaseerde MIPS-methode.
• Extractie-Kwaliteit: De geëxtraheerde programma's lossen de meeste taken op met een nauwkeurigheid van 1.00. De Root Mean Squared Error (RMSE) tussen het originele model en het geëxtraheerde script is verwaarloosbaar laag, wat aantoont dat de extractie zeer betrouwbaar is.
• Uitbreiding naar Continue Domeinen: In tegenstelling tot MIPS, kan de Discrete Transformer taken met zwevende komma-getallen (zoals exponential moving average en free fall height) verwerken en exacte symbolische formules voor deze continue dynamica extraheren.
• Voorbeelden van Succes:
  • Parity last2: Het model leert de XOR-logica en de extractie levert de exacte algebraïsche vorm $y_t = x_t + x_{t-1} - 2x_t x_{t-1}$.
  • Max/Min prev2: Het model herkent dat ReLU-operaties gebruikt worden om conditionele logica te simuleren, wat leidt tot exacte wiskundige identiteiten voor maximum en minimum.

Betekenis en Impact

Dit werk is een belangrijke stap in het veld van mechanistische interpretbaarheid en programmasynthese.

• Transparantie: Het bewijst dat Transformer-modellen niet noodzakelijk ondoorgrondelijke "black boxes" hoeven te zijn; ze kunnen worden ontworpen om menselijk leesbare algoritmen te internaliseren.
• Nieuwe Ontdekkingen: De methode stelt onderzoekers in staat om de novo algoritmen te ontdekken die onafhankelijk zijn van menselijk geschreven code, wat potentieel kan leiden tot innovatieve oplossingen in wetenschappelijke ontdekkingen.
• Controle: Het introduceert het concept van "controleerbare algoritme-ontdekking", waarbij onderzoekers de inductieve bias van het model kunnen manipuleren om specifieke soorten oplossingen te forceren.

Samenvattend biedt de Discrete Transformer een robuust raamwerk om de interne logica van neurale netwerken te vertalen naar verifieerbare, uitvoerbare code, zelfs voor complexe taken met continue variabelen.