Opening the AI black box: program synthesis via mechanistic interpretability

TL;DR

This work introduces MIPS, an automated mechanistic interpretability framework that converts neural network–learned algorithms into Python code by training an RNN, extracting discrete representations, and applying symbolic regression. It pairs AutoML-driven simplification with a suite of normalizers and integer/Boolean autoencoders to yield transparent, verifiable programs without relying on human-written training data. In a 62-task benchmark, MIPS solves 32 tasks (including many GPT-4–challenging ones), demonstrating complementary strengths relative to large language models. The results illuminate how neural networks encode discrete information and outline a path toward scalable, trustworthy machine-learned algorithms through automated mechanistic interpretability.

Abstract

We present MIPS, a novel method for program synthesis based on automated mechanistic interpretability of neural networks trained to perform the desired task, auto-distilling the learned algorithm into Python code. We test MIPS on a benchmark of 62 algorithmic tasks that can be learned by an RNN and find it highly complementary to GPT-4: MIPS solves 32 of them, including 13 that are not solved by GPT-4 (which also solves 30). MIPS uses an integer autoencoder to convert the RNN into a finite state machine, then applies Boolean or integer symbolic regression to capture the learned algorithm. As opposed to large language models, this program synthesis technique makes no use of (and is therefore not limited by) human training data such as algorithms and code from GitHub. We discuss opportunities and challenges for scaling up this approach to make machine-learned models more interpretable and trustworthy.

Figures (6)

• Figure 1: The pipeline of our program synthesis method. MIPS relies on discovering integer representations and bit representations of hidden states, which enable regression methods to figure out the exact symbolic relations between input, hidden, and output states.
• Figure 2: These hidden structures can be turned into discrete representations. Left: the hidden states for the bitstring addition task are seen to form four clusters, corresponding to 2 bits: the output bit and the carry bit. Right: the hidden states for the Sum_Last2 task are seen to form clusters on a 2D lattice corresponding to two integers.
• Figure 3: The generated program for the addition of two binary numbers represented as bit sequences. Note that MIPS rediscovers the "ripple adder", where the variable $b$ above is the carry bit.
• Figure 4: Comparison of code generated from an RNN trained on Sum_Last5, with (top) and without (top) normalizers. The whitening normalizer provided numerical stability to the Jordan normal form normalizer, which itself simplified the recurrent portion of the program. The Toeplitz and de-biasing normalizers jointly sparsified the occurrences of $x$ in the program, and the number of terms required to compute $y$. The quantization normalizer enabled all variables to be represented as integers.
• Figure 5: We compare MIPS against program synthesis with the large language model GPT-4 Turbo, prompted with a "chain-of-thought" approach. It begins with the user providing a task, followed by the model's response, and culminates in assessing the success or failure of the generated Python code based on its accuracy in processing the provided lists.