Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

Dongming Jin; Jie Fu; Kechi Zhang; Wei Bi; Yifan Zhang; Zhi Jin

REVIEW 3 major objections 7 minor 1 cited by

Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

T0 review · 3 major / 7 minor · reviewed 2026-08-03 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read A Transformer can be trained so that its weights compile directly into human-readable, executable programs, without any human-written reference solution.

desk verdict A credible architecture-plus-extraction demo with exact-recovery examples, but the attention taxonomy is assumed complete and the benchmark is filtered to what PySR can express. read the letter →

arxiv 2601.05770 v3 pith:4JGIN6A2 submitted 2026-01-09 cs.LG cs.CL

Yifan Zhang , Wei Bi , Kechi Zhang , Dongming Jin , Jie Fu , Zhi Jin This is my paper · ORCID

classification cs.LGcs.CL

keywords algorithmextractionmechanisticinterpretabilitysymbolicregressiondiscretetransformerprogramsynthesistemperatureannealinghypothesistestinginterpretableAI

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper claims that a carefully constrained Transformer can encode its learned solution as a discrete computation graph in its weights, and that this graph can be read out as human-readable Python code with high fidelity. The proposed Discrete Transformer replaces dense latent vectors with a concatenative list of scalar variables, uses temperature-annealed sampling to make attention and MLP selection discrete, and then recovers the algorithm via hypothesis testing on attention patterns and symbolic regression on arithmetic modules. On a benchmark of linear, non-linear, and continuous-dynamics tasks, the extracted programs match the accuracy of an RNN-based prior extraction method on shared tasks and extend the approach to floating-point intermediate computations. The paper's core claim, taken in good faith, is that architectural design can make Transformer internals transparent enough for de novo algorithm recovery.

What carries the argument

The central mechanism is the discretized computation graph: a residual stream of scalar variables updated by concatenation, a temperature-controlled sampling function that acts as a differentiable pointer and converges to one-hot selection, and a hard-attention router whose values are raw scalars, forcing attention to become a deterministic pointer. Extraction then treats the model as a graph of router and operator nodes: hypothesis testing checks each attention head against Fixed Offset or Windowed Extrema patterns, while a genetic symbolic regression tool fits the low-dimensional sub-MLP mappings; backward traversal from a sparsified linear output head compiles these primitives into Python

What would settle it

Train the Discrete Transformer on a task whose solution requires an attention routing pattern outside the two hypothesized classes (for example, 'attend to the position of the median value in the previous k positions' or 'attend to the most recently seen position with a value equal to the current token') and check whether the extracted program still solves the task. If the model achieves low loss but hypothesis testing labels the relevant head unmatched and the extracted program fails, the claim that unmatched heads are always noise is falsified. A second direct probe is to evaluate extracted

Watch

Extended reading notes

Core claim

The paper claims that a Transformer built from a concatenative scalar residual stream, discretized reading mechanisms, hard attention routers, and bottlenecked arithmetic sub-MLPs can be trained with temperature annealing to converge to a fully discrete computation graph, and that this graph can be compiled into executable code by hypothesizing that attention heads implement one of two routing patterns—Fixed Offset and Windowed Extrema—and running symbolic regression on the low-dimensional sub-MLP mappings. Reported results include exact recovery of identities such as XOR as x_t + x_{t-1} - 2 x_t x_{t-1} and max as x_{t-1} + ReLU(x_t - x_{t-1}), plus high-fidelity programs for continuous dyn

Load-bearing premise

The claim that every functionally meaningful attention head falls into one of exactly two routing patterns—fixed offset or windowed extremum—and that any unclassifiable head is irrelevant to the final output; if a task needs any other routing pattern, the extraction silently discards a necessary component.

Editorial extensions

If this is right

If the claim holds, trained models on algorithmic tasks can yield exact, human-readable programs without any ground-truth program used in training, enabling de novo discovery of computational strategies.
The architecture extends algorithm extraction beyond discrete latent states to continuous-valued intermediate computations, covering tasks such as exponential moving average and free-fall dynamics.
Annealing dynamics show a two-phase transition where functional convergence precedes structural discretization, suggesting that learning the mapping first and then crystallizing selection is a reliable route to interpretable models.
Architectural inductive biases steer the discovered algorithm: removing MLP capacity forces attention-only solutions, and masking recent history produces mathematically equivalent higher-order recurrences, giving practitioners control over the form of the extracted program.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

If the two-pattern classification of attention heads is complete only for these benchmark tasks, the extraction method's generality is bounded: tasks requiring content-dependent routing outside Fixed Offset and Windowed Extrema (e.g., selecting a median or a value-thresholded position) would likely yield 'unmatched' heads that the pipeline would silently drop, corrupting the program. A natural tes
The architecture's RASP-style primitives suggest a reverse use: a user could specify a high-level algorithm, train the Discrete Transformer on inputs consistent with it, and check whether annealing rediscovers the intended structure or an algebraically equivalent one, making the system a probe for whether gradient descent re-discovers human-preferred algorithmic forms.
The controllability results imply a practical workflow: specify the allowed computational primitives (for instance, 'no MLP' or 'no short-range memory') and let annealing discover the simplest program in that class; a testable extension is per-module temperature schedules, which may reduce redundancy in over-parameterized settings.
The paper leaves open whether extracted programs remain exact on out-of-distribution inputs such as longer sequences or wider value ranges; running the recovered programs on such inputs would sharply test the claim that they are true algorithms rather than validation-set fits.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Best quick read: the paper is a real synthesis of existing ideas — Transformer Programs, MIPS, RASP, and PySR — and it delivers on its narrow claim: for tasks whose ground-truth rules are PySR-representable, the Discrete Transformer trains to an annealed discrete graph, and the extraction pipeline produces executable code, including continuous-dynamics tasks MIPS cannot handle. The parity_last2 and max/min examples are the real goods; the simplified expressions are algebraically exact, not merely low-RMSE fits, and the intervention experiments show the architecture can be steered to alternative algorithms. That is a genuine contribution to mechanism-extraction work.

What is new is the combination: a scalar residual stream with concatenation, hard attention restricted to routing with PLE keys and queries, sub-MLPs limited to arithmetic, and a training scheme that anneals from Gumbel-Softmax toward Sparsemax. The extraction pipeline itself — hypothesis testing for attention patterns, PySR for sub-MLPs, backward traversal — is a natural modular design. The phase-transition observation (loss drops before structural agreement) is interesting but not deeply characterized.

Soft spots, in rough order. No code or data is released, and Table 1 reports no variance despite three seeds. That matters because extraction success depends on thresholds and PySR runs. Second, the benchmark is filtered to tasks whose ground truth is expressible in PySR's operator set (Footnote 1), so exact recovery is partly by construction; the paper states this, but it caps the generality claim. Third, and most important: Section 4.1 asserts rather than proves that every functional attention head is Fixed Offset or Windowed Extrema. A content-based lookup such as 'most recent earlier position with same value as current token' falls outside both classes; Appendix B's unmatched-head defense is post-hoc and only shows that, in these tasks, the unmatched heads were redundant or negligible. The paper never characterizes the routing patterns expressible by PLE plus hard attention, so the general extraction claim is unproven. For the filtered benchmark the argument holds; for arbitrary algorithmic tasks it does not.

The paper is fairly cited and the writing is clear about its debts. It deserves a serious referee. I would want code and data plus a characterization, or at least a broader stress test, of the attention taxonomy before accepting, but the core method is sound enough to warrant that work.

Referee Report

3 major / 7 minor

Summary. The paper proposes the Discrete Transformer, a purpose-built Transformer variant designed so that a trained model can be compiled into an executable Python program. The residual stream is a concatenation of scalar variables; attention is a hard-routing operator using piecewise-linear-encoded scalar queries/keys plus relative-position bias; MLPs are restricted to low-dimensional element-wise sub-modules. A temperature-annealed Gumbel-Softmax-to-Gumbel-Sparsemax schedule drives selection weights to one-hot form. The extraction pipeline classifies each attention head as Fixed Offset or Windowed Extrema, fits each sub-MLP with PySR symbolic regression, and assembles a program by backward traversal from the linear output head. On ~25 algorithmic tasks (linear arithmetic, non-linear composition, continuous dynamics), the model and extracted programs achieve near-perfect accuracy with low RMSE, comparable to the RNN-based MIPS baseline on shared tasks; the paper also reports an annealing phase transition and demonstrates that architectural interventions controllably alter the discovered program.

Significance. If the results hold, the paper provides a concrete instance of algorithm extraction from a differentiable architecture: the model converges to a discrete computation graph, and the pipeline compiles it into executable code with independent held-out evaluation — not a circular surrogate-fit exercise. Notable strengths: extracted programs are verified on held-out data; several simplified programs are shown explicitly (Figures 2–4); the Section 6.2 interventions are falsifiable predictions that are verified (e.g., zero sub-MLPs force a windowed-extrema attention solution); and the continuous-dynamics tasks genuinely extend the MIPS discrete-latent paradigm. Offsetting these, there is no released code or data, Table 1 reports no variance despite three seeds, the benchmark is restricted to PySR-representable targets (Footnote 1), and the attention taxonomy that the pipeline relies on is an unproven two-class hypothesis. These are scope limitations on the generality of the extraction claim rather than refutations of the benchmark results.

major comments (3)

[§4.1, Eqs. (3)–(5)] The extraction pipeline classifies every attention head as Fixed Offset or Windowed Extrema and discards 'unmatched' heads as noise, supported only by two post-hoc examples in Appendix B. The two-class taxonomy is not proven complete: the attention scores in Eq. (4) are an arbitrary function of PLE-encoded scalar query/key values plus a relative-position bias, so content-based patterns other than windowed extrema — e.g., 'attend to the most recent earlier position whose value equals the current value' — are expressible in principle, would be classified as unmatched, and would be silently dropped by the §4.3 backward traversal. Footnote 1 restricts the benchmark to PySR-representable rules, so this class is never exercised. The general extraction claim thus rests on an untested completeness assumption. Concrete remedy: add a task requiring equality-based content lookup (or characterize th
[§5, Footnote 1] The benchmark includes only tasks whose ground truth lies in PySR's default operator set; the bit-palindrome rule is excluded for this reason. Since PySR is the extraction engine, every tested target is by construction inside the regression tool's searchable family. This is not circular in the narrow sense — the model is trained without reference programs and extracted programs are scored on held-out data — but the 'de novo algorithm discovery' claim is demonstrated only within that family. The continuous-dynamics tasks (exponential moving average, free fall height) are likewise PySR-representable, so the claimed 'broadening' does not leave this family. I recommend adding at least one benchmark item whose target expression is outside PySR's default operators (even if the model can implement it via its ReLU MLPs), or explicitly scoping the de novo claim.
[§5, Table 1] The headline 'comparable to MIPS' rests on Table 1, but the MIPS numbers are retrieved from Michaud et al. (2024) rather than re-run under the data protocol used here (sequence length 10, 10^6 training samples, MSE objective, token-tagging vs. language-modeling formulations as described in Appendix C). Cross-paper baselines for algorithmic benchmarks are sensitive to these choices, and the paper does not state whether task definitions, sequence lengths, and input distributions match the MIPS benchmark. Please either re-run MIPS under the same protocol or explicitly label the comparison as informal; at minimum, report the per-seed variance of the Ours columns to support parity claims.

minor comments (7)

[Table 1 / §5] The training details state that performance is averaged over three random seeds, but Table 1 reports single point values for Acc. and RMSE with no variance or range. For the headline parity claim with MIPS, the spread across seeds should be shown (or the degeneracy of the 1.00 values stated).
[§4.1–§4.3] The thresholds that operationalize both hypothesis tests ('beyond a predefined threshold', 'large fraction') and the output-head sparsity threshold ε are never reported; no null hypothesis, test statistic, or p-value is defined, so 'hypothesis testing' is in practice threshold-based classification. Report the threshold values and (ideally) a sensitivity check, since these are free parameters of the pipeline.
[§5, Figures 3–4] Please state whether the SymPy simplifications were applied to unrounded PySR fits and whether domain assumptions (e.g., binary inputs) were used. As displayed, the rounded coefficients in Figure 3 do not algebraically simplify to xt + xt−1 − 2xtxt−1, so the phrase 'exact algebraic formulation' needs qualification (exact-on-domain vs. algebraic identity).
[§5, Table 2] The magnetic task is a failure case: train loss is 1.12×10^2 (Table 2) and the extracted RMSE is ~10^1 (Table 1). The paper lists the task in the table without commenting on the failure, even though it is relevant to the 'comparable to MIPS' summary (both methods fail here). Please acknowledge it.
[§5 / Appendix E] No extracted program for the two continuous-dynamics tasks is shown anywhere (Appendix E includes only discrete/linear tasks such as sum_last, bitwise and, add mod 3, spring, gravity). Since 'broadening extraction to tasks with continuous-valued intermediate computations' is a headline contribution, one implemented example (e.g., exponential moving average) should be exhibited.
[Fig. 8 / §6.2] Some extracted expressions are not 'human-readable' in any meaningful sense: e.g., bitwise and's V5 = ((((V0−0.49)*3.28)/(V1−0.52))−3.06)*0.19 is an opaque rational interpolant that is correct only on the four binary inputs. The readability claim should be tempered for such cases, or replaced with a per-task readability assessment.
[General] No code or data repository is referenced. Given that the pipeline involves per-task grid search, unspecified thresholds, and PySR settings, Table 1 cannot be reproduced from the text alone; a public release (or at minimum a detailed hyperparameter/threshold table) should be provided.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: extraction is surrogate fitting with independent test evaluation; the attention taxonomy is a scope limitation, not a definitional loop.

full rationale

The paper's derivation chain is: (i) train a Discrete Transformer under temperature-annealed hard-attention/selection; (ii) classify each attention head by testing Fixed Offset or Windowed Extrema against the model's own attention matrices on validation data; (iii) fit each sub-MLP with PySR on collected input-output pairs; (iv) assemble by backward traversal and evaluate the resulting Python program against held-out task labels (Table 1). Each of these steps estimates a component from data and then checks the composed program independently, so no 'prediction' is forced by construction. The self-citation to Zhang et al. (2025) in Appendix C is a background remark about state-tracking and is not load-bearing. The two genuine limitations—Footnote 1 restricting the benchmark to PySR-representable ground truths and Section 4.1's unmatched-heads assumption that all non-classified heads are negligible or redundant—are explicitly acknowledged by the authors and are correctness/scope risks for the generality of extraction, not definitional circularity: the taxonomy is presented as an empirically supported hypothesis, not as a theorem derived from the target results. No equation in the paper reduces to its own inputs, and no fitted parameter is renamed as a prediction; the extracted programs are evaluated against ground-truth task outputs. Hence no significant circularity.

Assumptions & free parameters 5 free parameters · 5 assumptions · 0 invented entities

The paper introduces a new architecture but no new physical entities. The extraction results rest on hand-picked temperature schedules, unspecified hypothesis-test thresholds, a benchmark filtered to PySR-representable rules, and the prior MIPS numbers. The PySR-fitted coefficients in each extracted program are fit parameters, not independent predictions.

free parameters (5)

Temperature annealing endpoints tau1=10.0, tau2=0.1 = 10.0 -> 0.1
Chosen by hand for geometric annealing; the convergence of selection distributions to one-hot vectors depends on this schedule (Section 5, Appendix A).
Hypothesis-test thresholds for attention pattern classification = unspecified
Fixed Offset detection uses a 'predefined threshold' on attention mass concentration and Windowed Extrema uses an agreement fraction, but the threshold values are not reported (Section 4.1).
Output-head sparsity threshold epsilon = unspecified
Used to prune variables before backward traversal; its value changes which variables are retained in the extracted program (Section 4.3).
PySR fitted coefficients in extracted sub-MLP expressions = per task, e.g., parity_last2: 3.57, -0.41, -0.61, -0.56, 0.17
Each extracted program's arithmetic expressions are obtained by genetic-algorithm symbolic regression against model activations; the central 'exact recovery' claims rest on these fitted constants (Section 4.2, Figures 3-4).
Per-task architecture hyperparameters (layers, heads, sub-MLPs) = see Table 2
Selected via grid search on validation performance, and results are reported for the best checkpoint, so model selection is part of the reported success (Appendix C).

assumptions (5)

ad hoc to paper Attention heads specialize to either Fixed Offset or Windowed Extrema; unmatched heads are computational noise
Asserted in Section 4.1 and Appendix B; no formal or exhaustive characterization of possible routing patterns is given, so extraction coverage depends on this assumption.
ad hoc to paper All benchmark ground-truth functions are representable by PySR's operator set (+, -, x, /, and ReLU where added)
Footnote 1 excludes tasks whose rules cannot be expressed with these operators; this bounds the scope of the empirical claims.
domain assumption MSE training with temperature annealing converges to a fully discrete computation graph
Empirically observed in Figure 5; no convergence guarantee is provided, and the paper's own Discrepancy metric does not reach exactly zero.
domain assumption MIPS baseline results in Table 1 are accurately transcribed from Michaud et al. (2024) under comparable task formulations
The comparison relies on previously published MIPS accuracy without rerunning the baseline; task formulation differences could bias the comparison.
domain assumption Concatenation-based residual stream preserves disentangled scalar semantics without destructive interference
Equation (1) assumes that appending scalar variables to the residual stream avoids superposition; this is the design basis for the method but is not proven.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer." pith.science (2026). https://pith.science/paper/4JGIN6A2

@misc{pith2026260105770,
  author       = {Pith},
  title        = {Pith review of: Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/4JGIN6A2}},
  note         = {Machine review of arXiv:2601.05770}
}

read the original abstract

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.

Figures

Figures reproduced from arXiv: 2601.05770 by the authors.

**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, while arithmetic transformations are approximated through symbolic regression. (III… view at source ↗

**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 yt = xt … view at source ↗

**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 0.1. The abbreviations Spring, Sum2, Diff2, Par2, and FB denote the spring, sum last2, diff last2, parity last2, and freebody tasks, respectively. the model can switch algorithmic paradigms,… view at source ↗

Figures from the paper (5 more)

**Figure 4.** 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 xt with the previous token xt−1. The bottom panel presents the simplified expressions, verifying…

**Figure 6.** Figure 6: Robustness landscape: program complexity (Top) and convergence loss (Bottom) on the sum last2 task. Excessive model capacity often increases the length of the extracted programs due to redundancy, while exerting only a minor effect on the MSE loss. 14 [PITH_FULL_IMAGE…

**Figure 7.** Figure 7: Intervened synthesized programs revealing alternative logical pathways. Left: When MLP-based arithmetic is prohibited, the model solves maximum prev2 and minimum prev2 by shifting to a pure attention mechanism. The extracted code shows explicit Windowed Extrema attenti…

**Figure 8.** Figure 8: Synthesized programs for sum last (Top Left), bitwise and (Bottom Left), and add mod 3 (Right). 16 [PITH_FULL_IMAGE:figures/full_fig_p016_8.png]

**Figure 9.** Figure 9: Synthesized programs for freebody (Top Left), spring (Bottom Left), and gravity (Right). 17 [PITH_FULL_IMAGE:figures/full_fig_p017_9.png]

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

Discovering Interpretable Algorithms by Decompiling Transformers to RASP
cs.LG 2026-02 conditional novelty 7.0 of 10

Trained length-generalizing transformers on small algorithmic and formal-language tasks can be automatically decompiled into short, interpretable RASP-style programs, while non-generalizing models generally cannot.

Reference graph

Works this paper leans on

2 extracted references · 2 linked inside Pith · cited by 1 Pith paper

[2020]

Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L

URL https://arxiv.org/abs/2006.1 1287. Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L. Sparse Autoencoders find highly interpretable features in Language Models, 2023. URL https:// arxiv.org/abs/2309.08600. Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Gross...

arXiv 2006
[2024]

unmatched heads

URL https://arxiv.org/abs/2401.1 4196. Jang, E., Gu, S., and Poole, B. Categorical reparame- terization with Gumbel-Softmax, 2017. URL https: //arxiv.org/abs/1611.01144. Lai-Dang, Q.-V ., Kang, T., and Son, S. Adaptive Trans- former programs: Bridging the gap between performance and interpretability in Transformers. In Yue, Y ., Garg, A., Peng, N., Sha, F...

arXiv 2017

Pith tools

Reviewed August 3, 2026 · model on record in the stance chip above.

[1] [2020]

Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L

URL https://arxiv.org/abs/2006.1 1287. Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L. Sparse Autoencoders find highly interpretable features in Language Models, 2023. URL https:// arxiv.org/abs/2309.08600. Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Gross...

arXiv 2006

[2] [2024]

unmatched heads

URL https://arxiv.org/abs/2401.1 4196. Jang, E., Gu, S., and Poole, B. Categorical reparame- terization with Gumbel-Softmax, 2017. URL https: //arxiv.org/abs/1611.01144. Lai-Dang, Q.-V ., Kang, T., and Son, S. Adaptive Trans- former programs: Bridging the gap between performance and interpretability in Transformers. In Yue, Y ., Garg, A., Peng, N., Sha, F...

arXiv 2017