💬 NLP

Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

This paper introduces the Discrete Transformer, a novel architecture that overcomes representation entanglement in standard Transformers by injecting discreteness through temperature-annealed sampling to enable the extraction of interpretable, human-readable algorithms from continuous models.

Original authors: Yifan Zhang, Wei Bi, Kechi Zhang, Dongming Jin, Jie Fu, Zhi Jin

Published 2026-03-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yifan Zhang, Wei Bi, Kechi Zhang, Dongming Jin, Jie Fu, Zhi Jin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot chef (a Transformer) that has learned to cook thousands of recipes by tasting millions of dishes. It knows how to make a perfect cake, but if you ask it, "How did you do that?" it can't explain itself. It just says, "I mixed ingredients in a big, messy bowl in my brain."

The problem is that the robot's brain is a kitchen blender. It mixes everything together so thoroughly that you can't tell where the flour ends and the sugar begins. This is called "entanglement." Because the ingredients are blended, you can't pull out a clean, written recipe (an algorithm) from its brain.

This paper introduces a new kind of robot chef called the Discrete Transformer. Instead of a blender, this chef uses a modular assembly line.

Here is how it works, broken down into simple steps:

1. The Specialized Assembly Line

The authors redesigned the robot's brain to have two distinct, separate teams:

The Delivery Team (Numerical Attention): Their only job is to move ingredients from one spot to another. They don't mix or cook; they just say, "Take the egg from the fridge and put it on the counter."
The Mixing Team (Numerical MLP): Their only job is to do the math. They take the ingredients the Delivery Team brought and mix them (add, subtract, multiply). They don't move things around.

By forcing the robot to keep these jobs separate, the "blender" effect is gone. You can clearly see who moved what and who mixed what.

2. The "Temperature" Training

To teach this robot, the authors use a clever trick called Temperature Annealing.

Hot Start: Imagine the robot is learning in a very hot room. It's chaotic and fuzzy. It tries many different ways to move and mix ingredients. It's exploring.
Cooling Down: Slowly, they lower the temperature. As it gets colder, the robot's choices become sharper and more definite. It stops guessing and starts making firm decisions.
Frozen State: By the end, the robot is "frozen" in a perfect, rigid state. Every movement and every mix is exact. Because the robot was forced to be so precise, its internal logic is now crystal clear.

3. Reverse-Engineering the Recipe

Once the robot is frozen, the authors act like detectives to write down the recipe:

For the Delivery Team: They look at where the robot moved ingredients. Did it always grab the item from "one step back"? That's a simple rule: Take previous item.
For the Mixing Team: They look at the math the robot did. Instead of seeing a messy neural network, they use a tool (Symbolic Regression) to find the simple math equation hidden inside, like A + B or A * 2.

4. The Result: A Human-Readable Book

Finally, they stitch these clues together. Instead of a messy neural network, they produce a clean, simple Python code file that looks like a normal computer program.

Why is this a big deal?

It works on continuous numbers: Previous methods could only handle simple "on/off" switches (like 0 or 1). This new method can handle real-world numbers like temperature, speed, or money (e.g., calculating how high a ball will fall).
It finds new solutions: Because the robot isn't just copying human code, it might find a clever, efficient way to solve a problem that no human ever thought of.
It's controllable: If you tell the robot, "Don't use the mixing team, just use the delivery team," it will figure out a completely different way to solve the problem using only movement. This proves we can steer the robot's thinking.

The Bottom Line

This paper solves the mystery of the "black box" Transformer. They built a version of the AI that thinks in steps rather than a blur. By cooling it down and separating its jobs, they can look inside its brain and pull out a clear, written instruction manual that humans can read, understand, and trust. It's like turning a magic trick into a step-by-step science experiment.

1. Problem Statement

The paper addresses the challenge of Algorithm Extraction: synthesizing executable, human-readable programs directly from trained neural networks without relying on pre-existing human-written code.

The Gap: While Recurrent Neural Networks (RNNs) have shown success in this area (e.g., via the MIPS framework), extending this to the dominant Transformer architecture is difficult.
The Core Obstacle: Standard Transformers suffer from representation entanglement (specifically "superposition"), where features are encoded in overlapping, non-orthogonal directions. This polysemantic nature makes it impossible to directly recover symbolic expressions or discrete logic from the continuous, high-dimensional weight matrices.
Goal: To bridge the gap between continuous Transformer representations and discrete symbolic logic to enable the extraction of de novo algorithms.

2. Methodology: The Discrete Transformer

The authors propose a specialized architecture called the Discrete Transformer, designed to enforce functional disentanglement and bridge continuous optimization with discrete logic.

A. Architectural Design

The model is built upon the RASP (Restricted Access Sequence Processing) computational model and consists of three distinct components:

Numerical Residual Stream: Unlike standard Transformers that use dense latent vectors, the residual stream is a collection of explicit scalar variables. This prevents information interference and superposition.
Numerical Attention (Router):
- Acts strictly as an information routing operator, not a feature mixer.
- Uses Piecewise Linear Encoding (PLE) to project scalar queries/keys into a higher-dimensional space for comparison.
- Employs a hard attention mechanism via temperature-controlled sampling to select specific historical values (e.g., "copy value from $t-1$ ").
Numerical MLP (Operator):
- Dedicated solely to element-wise arithmetic and logical transformations.
- Decomposed into parallel, independent sub-modules (sub-MLPs) that perform single elementary operations on a small, fixed number of inputs (typically $k=2$ ).
- Inductive Bias: By constraining input dimensions and hidden widths, the architecture forces the network to decompose complex functions into simple arithmetic steps, making them amenable to symbolic regression.
Linear Output Head: Aggregates the discrete computational steps into a final prediction.

B. Training Mechanism: Temperature Annealing

To transition from continuous learning to discrete logic, the model uses temperature-annealed sampling:

A temperature parameter $\tau$ controls the sampling function $S(\cdot, \tau)$ used in the reading mechanisms and attention heads.
Annealing Process: Training starts with a high $\tau$ (allowing soft, probabilistic exploration) and gradually decreases $\tau \to 0$ .
Result: The model converges into a fully discrete computation graph where the solution is implicitly encoded in the weights.

C. Extraction Pipeline (From Weights to Code)

Once the model converges to a discrete state, a modular extraction pipeline recovers the algorithm:

Hypothesis Testing (for Attention): Analyzes attention patterns to identify routing logic. The authors hypothesize and verify two main patterns:
- Fixed Offset: Attending to a specific relative position (e.g., $i - \delta$ ).
- Windowed Extrema: Attending to the max/min value within a sliding window.
Symbolic Regression (for MLP): Treats each MLP sub-module as a black-box mapping $f: \mathbb{R}^k \to \mathbb{R}$ . Tools like PySR (based on genetic algorithms) are used to find exact closed-form algebraic expressions (e.g., $y = x_1 + x_2$ ) that fit the input-output data.
Global Assembly: The extracted routing patterns and arithmetic expressions are integrated via a backward traversal of the computational graph to generate a concise Python program.

3. Key Contributions

Discrete Transformer Architecture: A RASP-aligned architecture that enforces functional disentanglement (routing vs. arithmetic) and uses a scalar-based residual stream to eliminate superposition.
Modular Extraction Pipeline: A novel method combining hypothesis testing (for routing) and symbolic regression (for arithmetic) to extract interpretable programs from continuous variables, extending mechanistic interpretability beyond discrete latent approximations.
Controllable Inductive Biases: Demonstrates that architectural constraints (e.g., limiting sub-MLPs) act as strong inductive biases, allowing researchers to steer the model toward specific algorithmic paradigms (e.g., forcing a model to use attention for maximization instead of arithmetic).

4. Experimental Results

The authors evaluated the framework on a diverse benchmark including Linear Arithmetic, Non-Linear Composition, and Continuous Dynamics tasks.

Performance: The Discrete Transformer achieves performance comparable to RNN-based methods (MIPS), with Mean Squared Error (MSE) approaching zero and extraction accuracy of 1.00 on most integer-valued tasks.
Fidelity: The extracted programs are not just approximations; they recover exact algebraic identities.
- Example: For the parity_last2 task, the model extracted the exact algebraic formulation $y_t = x_t + x_{t-1} - 2x_t x_{t-1}$ , which is mathematically equivalent to the XOR operation.
- Example: For maximum_prev2, it recovered the piecewise-linear identity using ReLU: $\max(a, b) = b + \text{ReLU}(a - b)$ .
Continuous Dynamics: Unlike MIPS (which relies on discrete latent states), the Discrete Transformer natively handles floating-point variables, successfully extracting algorithms for tasks like Exponential Moving Average and Free Fall Height.
Training Dynamics: Analysis reveals a distinct phase transition where functional convergence (learning the mapping) precedes structural discretization (locking into discrete weights), ensuring a robust transition from exploration to exploitation.

5. Significance and Impact

Bridging the Gap: This work successfully overcomes the "superposition" barrier in Transformers, proving that interpretable, symbolic algorithms can be extracted from Transformer weights.
De Novo Discovery: It enables the discovery of innovative algorithms independent of human-generated code, potentially uncovering novel solutions to physical or logical problems.
Scientific Discovery: By supporting continuous variables and offering controllability via inductive biases, the framework serves as a powerful tool for scientific discovery, allowing researchers to find specific forms of algorithms that align with domain knowledge.
Interpretability: It establishes a new paradigm for Transformer interpretability, moving from analyzing "what" a model does to explicitly recovering "how" it does it in the form of executable code.