Interpreting the Latent Structure of Operator Precedence in Language Models

TL;DR

This work probes how operator precedence and intermediate arithmetic steps are represented inside a contemporary LLM (LLaMA three point two to three billion parameters) using a comprehensive interpretability toolkit. It combines logit lens, linear probes, attention/MLP ablations, partial embedding swaps, and UMAP to trace hidden computations and the encoding of precedence across a synthetic three-operand, two-operator dataset. The results show that intermediate computations reside in the residual stream, especially after MLP blocks, and that operator precedence is linearly encoded in operator embeddings after attention, with targeted embedding perturbations able to shift arithmetic outcomes. These findings advance mechanistic interpretability of arithmetic reasoning in LLMs and point to precise intervention strategies for diagnosing and improving reliability in reasoning tasks.

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.

Figures (6)

• Figure 1: (Left) shows that intermediate calculations are linearly encoded in the model’s activations after layer 0, as indicated by high $R^2$ scores. (Right) Layer-wise detection frequency of intermediate computations in the residual streams at each layer. Detection peaks around layers 18–19, suggesting that the model most strongly represents intermediate computations in the later layers of its forward pass.
• Figure 2: Shows impact of attention ablation on output accuracy across 19 layers of the model. Early layers, particularly layer 0, showed the most pronounced effect on accuracy, suggesting a key role in establishing the structural interpretation of arithmetic expressions.
• Figure 3: Shows impact of attention ablation on intermediate computation observed through Logit Lens across 4401 equations. Layers 9, 13, and 19 had the most decreases in observed intermediate values, suggesting a key role in establishing intermediate computations.
• Figure 4: Low dimensional projection of operator token embeddings before (left) and after (right) the attention block in layer 0. After attention, embeddings separate by operator position and precedence, suggesting that attention encodes operator precedence information
• Figure 5: (left) "2 + 3 * 3 = " and (right) "4 + 8 / 4 = ". The swapped-precedence-logit (blue) tracks the logit of the swapped (incorrect precedence) answer after patching, the actual-precedence-logit (green) tracks the logit of the correct answer, and the top-logit (red) shows the logit of the model's predicted token. In both examples, a subset of embedding dimensions substantially modulate the precedence-sensitive logit values, indicating that operator precedence information is sparsely localized across specific dimensions of the residual stream activations.