K² · Artificial intelligence
AI rewrites a chatbot's inner machinery as readable Python code
Amiri Hayes, Belinda Li, Jacob Andreas
3 authors · cs.LG, cs.AI
Like explaining it at the dinner table.
Inside every AI language model sits a swarm of tiny decision-makers called attention heads. Each one decides which earlier words a sentence should "look back at" when predicting the next word — but they do this through a wall of numbers nobody can read. This paper turns some of those number-walls into plain Python programs you can actually open and inspect.
Here's the trick. For one attention head, you feed it example sentences and record exactly which words it focused on — a grid showing how much each word paid attention to every other word. Then you hand a summary of that grid to a separate AI and ask it: write Python code that produces this same focus pattern from the raw text. The AI guesses many candidate programs. You keep the ones that best predict the head's behavior on fresh, unseen sentences.
The results: under 1,000 such programs recreate the attention patterns of three different models (GPT-2, TinyLlama, Llama-3B), matching the real focus more than 75% of the time. Swap out a quarter of the heads with their code stand-ins, and the model gets only 16% worse at predicting text — and stays just as good at answering questions.
But 75% overlap means a sizable chunk of each head's behavior still escapes the code, and a 16% accuracy drop is real, not free.
Why you should care: When an AI gives a baffling or biased answer, "it's a black box" is the usual shrug. Code you can read for a quarter of its parts is a crack of daylight into that box.
arXiv preprint — these findings haven’t been peer-reviewed yet. Treat them as early results, not settled science.