Interpreting the Latent Structure of Operator Precedence in Language Models

AI summary

Overview

• Research area: Mechanistic interpretability of arithmetic reasoning in large language models (natural language processing).
• Technical level: Intermediate. The paper assumes familiarity with transformer components (attention blocks, MLP layers, residual stream) and interpretability tooling (logit lens, linear probes, ablation, UMAP), but the experimental logic is described clearly.
• Scope: A single-model interpretability study of how LLaMA 3.2-3B internally represents operator precedence and intermediate arithmetic results.

What This Paper Is About

Most work on LLM arithmetic looks only at prompts and final answers. This paper asks what happens inside the model when it evaluates an expression like 1 + 1 × 2 — does it compute 1 × 2 first because of mathematical precedence, or does it process the tokens left to right? The authors use interpretability tools to trace intermediate results through the model's residual stream and to locate where, and in which dimensions, operator precedence is encoded.

Key Contributions

1. Evidence that intermediate arithmetic results exist in the residual stream. Using logit lens on 4,401 correctly answered prompts, the authors find the intermediate value appears as the top logit 2,799 times (63.6%), and that this logit only becomes top-ranked after the MLP block, not after attention.
2. Causal evidence that attention layers build these representations. Layer-wise attention ablation shows that zeroing attention outputs at specific layers degrades both final-answer accuracy and the visibility of intermediate computations in the residual stream.
3. A new intervention technique: partial embedding swap. The authors swap individual embedding dimensions between operator tokens to change which precedence the model applies, then rank dimensions by their influence and find the minimal prefix of top-ranked dimensions needed to flip the model's prediction.
4. Evidence that precedence is linearly encoded and geometrically organized. A logistic regression probe predicting whether an operator is evaluated first or second reaches 100% accuracy on its test set, and UMAP projections of operator token embeddings show clustering by position and precedence after attention.

Main Findings

• Intermediate values appear late and after MLPs: Across 4,401 correctly answered prompts, the intermediate calculation surfaced as the top logit 2,799 times (roughly 63.6%), in layers 16–27. At the layer where the intermediate logit first appears, it only becomes top-ranked after the MLP block, indicating the MLP produces the intermediate computation.
• Intermediate values are linearly decodable early: The paper reports high R² scores for a linear probe predicting the intermediate value from activations, indicating that intermediate calculations are linearly encoded in the model's activations after layer 0.
• Attention ablation at layer 0 is catastrophic for accuracy: Ablating attention at layer 0 produced incorrect answers in 95.7% of cases. Layers 1, 5, and 7 also showed elevated error rates of 42.2%, 21.1%, and 24.2% respectively. Layers after 19 showed limited effects. Figure 2 is described as showing this across 19 layers, while the procedure is described as running across all 27 transformer layers.
• Ablation reduces the visibility of intermediate computations too: Attention ablation across 19 layers significantly reduced the occurrence of intermediate computations observed through the logit lens. The paper text highlights layers 5, 11, and 19 as most pronounced; the Figure 3 caption instead names layers 9, 13, and 19.
• Precedence is decodable from operator token activations after attention: A logistic regression probe trained on operator token activations before vs. after the first attention block (layer 0) achieved 100% accuracy on the test set (80%/20% train/test split), strongly indicating that operator precedence is present after attention.
• Operator embeddings cluster by position and precedence after attention: UMAP projections of operator token embeddings before and after attention in layer 0 show that operators matching in both position and precedence are clustered near each other after attention, but not before.
• Precedence can be manipulated by editing a small set of dimensions: Partial embedding swap successfully altered the model's highest logit in multiple instances. Figure 5 illustrates this for 2 + 3 * 3 = and 4 + 8 / 4 =, where swapping a subset of high-impact embedding dimensions substantially modulated the logit of the precedence-swapped (incorrect) answer, indicating precedence information is sparsely localized across specific dimensions of the residual stream activations.

Methodology in Plain English

The authors built a synthetic dataset of expressions with three operands (each a whole number from 1 to 9) and two operators drawn from mixed-precedence pairs — (+, *), (-, *), (+, /), (-, /) — so that the order of evaluation actually matters. Each operand/operator combination was rendered in six structural forms: left-parenthesized, right-parenthesized, flipped versions of each, and two no-parenthesis variants. This yielded 8,547 prompts; the 4,401 prompts the model answered correctly as its top logit were kept for analysis.

They then asked four questions. First, does the intermediate answer (for example, 3 × 3 = 9 inside 2 + 3 * 3 =) ever show up if you project each layer's residual stream through the unembedding matrix — the logit lens — and look at the top 10 tokens? Second, can a simple linear probe predict the intermediate value from layer activations even when it does not appear in the top tokens? Third, if you zero out the attention output at one layer at a time (27 layers, applied before the MLP block, before the residual addition), does the model's accuracy or the visibility of the intermediate value break? Fourth, where in the embedding space is precedence stored?

For the last question they introduced partial embedding swap: take an expression like 3 + 4 * 5 =, which correctly evaluates to 23 but evaluates to 35 left-to-right. Swap one dimension at a time between the hidden representations of the + and * operator tokens, measure how much the logit of 35 rises, and rank the dimensions. Then swap growing prefixes of the top-ranked dimensions to find the minimal subset needed to push 35 to the top prediction. This is complemented by a logistic regression probe trained on operator token activations before and after layer 0's attention block to classify whether an operator is evaluated first or second, and by UMAP projections of operator token embeddings labeled by position, operator name, and evaluation order.

Why This Matters

Impact on research. Prior interpretability work on arithmetic largely traced circuits for final predictions or studied prompting strategies. This paper shifts attention to the internal handling of order of operations — a structural property of the input rather than a numerical one — and reports that this property is linearly decodable and sparsely editable. The partial embedding swap technique is a reusable intervention for causal testing of latent features, and the layer-0 attention result gives circuit-level researchers a concrete target to investigate further in a small model that fails often enough to be a useful study subject.

Real-world applications.

• Diagnosing and debugging arithmetic errors in deployed assistants and calculators built on smaller open models, where the paper notes failures are most prominent.
• Targeted model editing: identifying a small set of dimensions that control a behavior suggests cheaper, more surgical interventions than full fine-tuning.
• Evaluation design: checking whether a model applies precedence, rather than only whether the final number is right, since the related-work discussion notes models can produce correct answers through flawed internal logic.
• Trust and monitoring: probing internal representations for structural understanding (does the model internally apply the right grouping?) without needing to expose it through output prompts.

Industry relevance. The paper uses LLaMA 3.2-3B, an open-source instruction-tuned model, and all techniques — logit lens, probes, ablation, activation patching, UMAP — are standard and comparatively cheap. This makes the findings directly applicable to teams deploying or auditing small models on tool-use and calculation pipelines, where precedence errors silently corrupt downstream results.

Future Directions

• The paper reports results only for LLaMA 3.2-3B; whether the layer-16–27 intermediate-value signature, the layer-0 attention dependence, and the operator embedding clustering replicate in other model families and sizes is not reported.
• The discrepancy between the text (layers 5, 11, 19) and the Figure 3 caption (layers 9, 13, 19) for the ablation effect on intermediate computations invites closer layer-level localization.
• The dataset was restricted to positive whole numbers in all intermediate steps; behaviors on negative, fractional, or larger-number expressions remain open.
• The paper does not report the minimal number of swapped dimensions found by the cumulative patching algorithm (Algorithm 2) in numeric terms, nor a systematic success rate for partial embedding swap beyond "multiple instances," leaving the reliability and generality of the intervention to be quantified.
• Whether precedence representations can be used to improve arithmetic performance — for example through targeted steering rather than only altering predictions — is raised as an implication but not tested here.

Target Audience

Mechanistic interpretability researchers and graduate students studying arithmetic reasoning in transformers; NLP practitioners working with small open-source models on numerical tasks; and safety or evaluation engineers who want internal, not just output-level, checks on whether a model reasons about structure correctly. Readers need working knowledge of transformer internals and familiarity with tools like logit lens and activation patching to follow the experimental design.

Authors’ abstract

Large Language Models (LLMs) have demonstrated impressive reasoning capabilities but continue to struggle with arithmetic tasks. Prior works largely focus on outputs or prompting strategies, leaving the open question of the internal structure through which models do arithmetic computation. In this work, we investigate whether LLMs encode operator precedence in their internal representations via the open-source instruction-tuned LLaMA 3.2-3B model. We constructed a dataset of arithmetic expressions with three operands and two operators, varying the order and placement of parentheses. Using this dataset, we trace whether intermediate results appear in the residual stream of the instruction-tuned LLaMA 3.2-3B model. We apply interpretability techniques such as logit lens, linear classification probes, and UMAP geometric visualization. Our results show that intermediate computations are present in the residual stream, particularly after MLP blocks. We also find that the model linearly encodes precedence in each operator's embeddings post attention layer. We introduce partial embedding swap, a technique that modifies operator precedence by exchanging high-impact embedding dimensions between operators.