Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

TL;DR

The paper introduces the Discrete Transformer, an architecture designed to bridge continuous representations and discrete symbolic logic to enable extraction of executable algorithms from Transformer models. By enforcing functional disentanglement—routing via Numerical Attention and arithmetic via Numerical MLP—and applying temperature-annealed sampling, it yields interpretable primitives that can be assembled into Python code. The authors validate on the MIPS benchmark, achieving competitive performance with RNN baselines and demonstrating the ability to recover exact arithmetic and logical forms, as well as to model continuous dynamics. They further analyze the training dynamics, show a phase transition from exploration to exploitation during annealing, and demonstrate controllability through inductive biases. This framework expands the scope of mechanistic interpretability to continuous-variable domains and offers a principled path to demonstration-free algorithm discovery.

Abstract

Algorithm extraction aims to synthesize executable programs directly from models trained on specific algorithmic tasks, enabling de novo algorithm discovery without relying on human-written code. However, extending this paradigm to Transformer is hindered by superposition, where entangled features encoded in overlapping directions obstruct the extraction of symbolic expressions. In this work, we propose the Discrete Transformer, an architecture explicitly engineered to bridge the gap between continuous representations and discrete symbolic logic. By enforcing a strict functional disentanglement, which constrains Numerical Attention to information routing and Numerical MLP to element-wise arithmetic, and employing temperature-annealed sampling, our method effectively facilitates the extraction of human-readable programs. Empirically, the Discrete Transformer not only achieves performance comparable to RNN-based baselines but crucially extends interpretability to continuous variable domains. Moreover, our analysis of the annealing process shows that the efficient discrete search undergoes a clear phase transition from exploration to exploitation. We further demonstrate that our method enables fine-grained control over synthesized programs by imposing inductive biases. Collectively, these findings establish the Discrete Transformer as a robust framework for demonstration-free algorithm discovery, offering a rigorous pathway toward Transformer interpretability.

Figures (8)

• Figure 1: Illustration of the proposed framework for extracting executable algorithms from a Discrete Transformer. (I) Discrete Search: Temperature annealing is leveraged to encourage interpretable discretization in both Numerical Attention and MLP modules. (II) Algorithm Extraction: Attention patterns are characterized via hypothesis testing (e.g., identifying token shifts), while arithmetic transformations are approximated through symbolic regression. (III) Synthesized Program: The extracted components are integrated via a linear output head to generate Python code. As shown, the framework successfully recovers the parity_last2 algorithm, correctly implementing the arithmetic XOR logic.
• Figure 2: Algorithm extraction results for the sum_last2 task. Modules are denoted by their type and indices (e.g., Attn_L0H0 represents the attention head at index 0 of layer 0). The extracted code reveals that the model utilizes specific attention heads to retrieve the previous token. Symbolic simplification (bottom) shows the mathematically simplified expression, verifying that the model correctly learns the logic $y_t = x_t + x_{t-1}$.
• Figure 3: Algorithm extraction results for the parity_last2 task. The extracted code reveals that the model utilizes specific attention heads (e.g., Attn_L0H0) to retrieve the previous token, and specific sub-MLPs (e.g., MLP_L0M0) to perform non-linear transformations. The bottom panel presents the symbolic simplification $y_t = x_t + x_{t-1} - 2 x_t x_{t-1}$, which is the algebraic formulation of the parity logic.
• Figure 4: Algorithm extraction results for maximum_prev2 and minimum_prev2 tasks. The top panel shows the raw code where sub-MLPs utilize ReLU functions to compare the current token $x_t$ with the previous token $x_{t-1}$. The bottom panel presents the simplified expressions, verifying that the model correctly reconstructs the extrema functions using the ReLU-based algebraic identities.
• Figure 5: Training dynamics exhibit a clear phase transition: the loss decreases earlier, while the pronounced drop in Discrepancy occurs slightly later, coinciding with Agreement approaching $1.0$ during temperature annealing from $10.0$ to $1.0$. The abbreviations Spring, Sum2, Diff2, Par2, and FB denote the spring, sum_last2, diff_last2, parity_last2, and freebody tasks, respectively.