Opening the AI black box: program synthesis via mechanistic interpretability (2402.05110v1)

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 LLMs, 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.

• The paper introduces MIPS, a method that auto-distills neural network learned algorithms into clear, executable Python code.
• It employs a multi-step process—including RNN training, finite state machine extraction, and symbolic regression—to enhance model transparency.
• MIPS solved 32 out of 62 algorithmic tasks and outperformed GPT-4 on 13, demonstrating its potential to improve AI reliability.

Mechanistic Interpretability for Program Synthesis through MIPS

Introduction and Background

The paper introduces MIPS (Mechanistic-Interpretability-based Program Synthesis), an innovative automated method designed for program synthesis. This method is rooted in the mechanistic interpretability of neural networks trained for specific algorithmic tasks. MIPS distinguishes itself by auto-distilling the learned algorithms into executable Python code, without direct reliance on human-generated training data such as algorithms and code from platforms like GitHub. This research provides a new lens through which machine-learned models can be made more interpretable and trustworthy.

Methodology

The MIPS framework involves a multi-step process that includes:

1. Neural Network Training: A black-box neural network is trained to learn an algorithm capable of performing the desired task. This paper employs a Recurrent Neural Network (RNN) to leverage its suitability for a range of algorithmic tasks.
2. Neural Network Simplification: This involves converting the neural network into a finite state machine, followed by simplification without compromising accuracy. An integer autoencoder translates the RNN into a more interpretable format, which aids in the discretization necessary for the subsequent steps.
3. Finite State Machine Extraction and Symbolic Regression: The next phases include extracting a finite state machine representation from the simplified neural network and employing symbolic regression to identify the simplest symbolic formulae that replicate the RNN's learned algorithm.

Through these steps, MIPS can distill complex, learned algorithms into Python code, making the underlying processes of neural networks clearer and potentially paving the way for advancements in interpretable AI.

Benchmark and Evaluation

MIPS was tested against a benchmark of 62 algorithmic tasks, demonstrating its capability by solving 32 tasks, including 13 that were not resolved by OpenAI's GPT-4, showcasing MIPS's complimentary nature to existing LLMs. The success of MIPS across these tasks underlines its potential to discover new algorithms autonomously, devoid of human biases or constraints found in training data. Furthermore, the methodology did not only prove effective in terms of task performance but also provided insights into how neural networks represent algorithmic knowledge, with implications for enhancing model transparency and trust.

Future Directions

The paper identifies several areas for future exploration, including extending the approach to more complex neural network architectures, addressing a broader range of data types, and scaling the method to handle larger networks. Moreover, automating formal verification of synthesized programs and exploring additional types of mechanistic simplifications represent tantalizing frontiers for research in making AI systems more decipherable and reliable.

Conclusion

In summary, the MIPS methodology introduces a novel approach to program synthesis, grounded in the mechanistic interpretability of neural networks. By converting learned algorithms into interpretable Python code, MIPS not only augments our understanding of machine learning models but also hints at a future where AI's decision-making processes are no longer opaque, contributing to the development of more transparent, verifiable, and trustworthy AI systems.