[Paper Review] Opening the AI black box: program synthesis via mechanistic interpretability

Eric J. Michaud, Isaac Liao+8 more|arXiv (Cornell University)|Feb 7, 2024

Evolutionary Algorithms and Applications4 citations

TL;DR

This paper introduces MIPS, a fully automated method for program synthesis that reverse-engineers neural network policies into interpretable Python code via mechanistic interpretability. By converting RNN hidden states into discrete integer or bit representations and applying symbolic regression, MIPS distills learned algorithms without relying on human-coded examples, solving 32 out of 62 algorithmic tasks—13 of which GPT-4 fails to solve—demonstrating a scalable, interpretable alternative to LLM-based code generation.

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.

Motivation & Objective

• To develop a fully automated method for distilling interpretable, symbolic programs from trained neural networks.
• To overcome the limitations of large language models in program synthesis by avoiding reliance on human-written code or examples.
• To enhance trustworthiness and interpretability of machine-learned models through automated mechanistic interpretability.
• To test whether automated interpretability can scale to practical program synthesis in algorithmic reasoning tasks.
• To provide a verifiable, symbolic alternative to black-box neural network policies in algorithmic domains.

Proposed method

• Train a recurrent neural network (RNN) with feedforward update and output functions to solve algorithmic tasks.
• Simplify the trained RNN using an integer autoencoder to extract discrete, finite-state representations of hidden states.
• Extract a finite state machine from the simplified network by identifying discrete state transitions in hidden representations.
• Apply symbolic regression (Boolean or integer) to discover exact mathematical or logical relationships between inputs, hidden states, and outputs.
• Generate human-readable Python code from the symbolic relations, ensuring correctness through formal verification with Dafny.
• Use a pipeline that avoids search-based program synthesis, instead learning the algorithm during neural network training and extracting it via interpretability.

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.

Experimental results

Research questions

• RQ1 Can mechanistic interpretability be fully automated to extract symbolic programs from neural networks?
• RQ2 Can this automated method outperform or complement large language models in program synthesis without using human-coded examples?
• RQ3 What is the performance of automated program synthesis on a benchmark of algorithmic tasks typically solved by RNNs?
• RQ4 How effective is symbolic regression in recovering the true algorithm from learned neural representations?
• RQ5 Can the synthesized programs be formally verified for correctness, ensuring reliability in safety-critical contexts?

Key findings

• MIPS successfully synthesized correct Python code for 32 out of 62 algorithmic tasks in the benchmark, demonstrating high effectiveness in program synthesis.
• Of these 32 tasks, 13 were not solved by GPT-4, indicating that MIPS complements rather than replaces LLM-based approaches.
• The method achieved high fidelity in recovering the true algorithm by converting RNN hidden states into discrete integer and bit representations.
• Formal verification using Dafny confirmed the correctness of the synthesized bit addition program, proving equivalence to the expected behavior under all valid inputs.
• The approach is entirely data-free of human-written code, relying only on input-output pairs and neural network training, making it scalable and trustworthy.
• Symbolic regression on discrete representations enabled accurate recovery of complex logic, such as modular arithmetic and iterative physics simulations.

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