Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

Algorithm extraction aims to synthesize executable programs directly from models trained on algorithmic tasks, enabling de novo recovery of executable mechanisms from weights without relying on human-written target programs. However, applying this paradigm to Transformer is complicated by representation entanglement (e.g., superposition), where features encoded in overlapping directions substantially hinder the recovery of symbolic expressions. We propose the Discrete Transformer, an architecture explicitly designed to bridge the gap between continuous representations and discrete symbolic logic. By injecting discreteness through temperature-annealed sampling, our framework effectively leverages hypothesis testing and symbolic regression to extract human-readable programs. Empirically, the Discrete Transformer achieves performance comparable to the RNN-based MIPS baseline on shared discrete tasks, while broadening extraction to tasks with continuous-valued intermediate computations. Finally, we show that architectural inductive biases provide fine-grained control over synthesized programs, establishing the Discrete Transformer as a controllable testbed for algorithm extraction and Transformer interpretability.

Modern AI models can often solve algorithmic problems, but the rules they learn are usually hidden inside millions of numerical parameters and are difficult for humans to inspect. This work introduces the Discrete Transformer, a Transformer architecture designed so that its learned computations can be converted back into readable program code. The key idea is to make the model gradually move from flexible continuous computation toward discrete, interpretable choices during training. After training, the model’s attention modules can be understood as routing information between positions, while its arithmetic modules can be translated into symbolic formulas. These pieces are then assembled into executable Python programs that reproduce the model’s behavior. Experiments show that the method can recover clear algorithms for arithmetic, logical, and physical-dynamics tasks, including cases involving continuous variables. This provides a step toward AI systems whose learned reasoning can be inspected, verified, and potentially reused as human-understandable algorithms.