Why Can't Transformers Learn Multiplication? Reverse-Engineering Reveals Long-Range Dependency Pitfalls — part 1/3

Cleaned-text code points [0, 15939). Bibliography omitted here; source.md includes it.

# Why Can’t Transformers Learn Multiplication? Reverse-Engineering Reveals Long-Range Dependency Pitfalls

###### Abstract

Language models are increasingly capable, yet still fail at a seemingly simple task of multi-digit multiplication. In this work, we study why, by reverse-engineering a model that successfully learns multiplication via implicit chain-of-thought , and report three findings: (1) Evidence of long-range structure: Logit attributions and linear probes indicate that the model encodes the necessary long-range dependencies for multi-digit multiplication. (2) Mechanism: the model encodes long-range dependencies using attention to construct a directed acyclic graph to ‘‘cache’’ and ‘‘retrieve’’ pairwise partial products. (3) Geometry: the model implements partial products in attention heads by forming Minkowski sums between pairs of digits, and digits are represented using a Fourier basis, both of which are intuitive and efficient representations that the standard fine-tuning model lacks. With these insights, we revisit the learning dynamics of standard fine-tuning and find that the model converges to a local optimum that lacks the required long-range dependencies. We further validate this understanding by introducing an auxiliary loss that predicts the ‘‘running sum’’ via a linear regression probe, which provides an inductive bias that enables the model to successfully learn multi-digit multiplication. In summary, by reverse-engineering the mechanisms of an implicit chain-of-thought model we uncover a pitfall for learning long-range dependencies in Transformers and provide an example of how the correct inductive bias can address this issue. 0 0 footnotetext: Contact: andrewlee@g.harvard.edu, smallyan@uchicago.edu 0 0 footnotetext: Code: https://github.com/ajyl/icot

## 1 Introduction

Large language models demonstrate striking capabilities across reasoning, planning, and tool use. Yet, they also fail on surprisingly simple algorithmic tasks ( Nye et al., 2021 ; Lee et al., 2023 ) . Why do Transformers excel at some tasks, but fail to learn others? One such example is multi-digit multiplication. Despite having billions of parameters, models like Llama-3.2 90B or GPT4 still fail at 4x4-digit multiplication ( Gambardella et al., 2024 ) , 1 1 1 Note that some recent proprietary models that do solve multi-digit multiplication may rely on tool-use. even when explicitly fine-tuned on the task ( Yang et al., 2023 ) . Why do Transformers fail to learn multiplication?

We study these questions by contrasting a standard fine-tuned model (SFT), which fails at multiplication, with a model trained with implicit chain-of-thought (ICoT) ( Deng et al., 2024 ; Deng et al., 2023 ) , which succeeds. ICoT works by providing explicit chain-of-thought tokens during training, but gradually removes them and thus forces the model to internalize intermediate steps in its latent states.

We partially reverse-engineer the ICoT model and uncover several insights. First, unlike the SFT model, the ICoT model learns the correct long-range structure needed for multi-digit multiplication. We provide evidence of this using logit attributions and linear regression probes. Mechanistically , the ICoT model encodes long-range dependencies by organizing its attention into a sparse, binary-tree-like graph, which (i) selects the correct digit pairs to compute partial products and (ii) “caches” these intermediate computations into earlier tokens for later retrieval. Lastly, geometrically , attention heads realize partial products as Minkowski sums of digit embeddings, and represent digits with Fourier bases, yielding a pentagonal prism structure – both of which are intuitive and efficient representations that the SFT model lacks.

With these insights, we revisit the dynamics of standard fine-tuning: under gradient descent and an auto-regressive loss, the model never learns these long-range dependencies, and thus loss plateaus on the middle digits. To confirm our understanding, we introduce a simple fix by introducing an auxiliary loss that supervises the model to predict a “running partial sum” through a lightweight linear regression probe. This provides an inductive bias to learn the proper long-range dependencies, allowing it to achieve perfect accuracy, without any supervision from chain-of-thought.

In summary, by partially reverse-engineering a network that successfully implements multi-digit multiplication, we uncover how it implements long-range dependencies, a mechanism that the unsuccessful model lacks. Our work highlights a challenge for Transformers to learn long-range dependency using gradient descent and an auto-regressive loss. While we demonstrate a task-specific inductive bias to address this issue, we anticipate generic improvements to address this limitation.

## 2 Experiment Setup, Training ICoT, Notations

Task, Models. We are interested in understanding the difference in a model trained with standard fine-tuning and ICoT. From experiments, we find that the simplest multi-digit multiplication in which standard fine-tuning fails but ICoT works is 4 × \times 4 digit multiplications. Similarly, the smallest architecture in which ICoT works is a 2-layer model with 4 attention heads. Thus we carefully study a 2-layer 4-head ICoT model and a standard fine-tuned model trained on 4 × \times 4 multiplication.

Training Procedures. Our ICoT setup is the same as that Deng et al. (2024) . Here we provide an informal overview of ICoT, with details in Appendix A.1 . Namely, assume two operands a = ( a 3 , a 2 , a 1 , a 0 ) , b = ( b 3 , b 2 , b 1 , b 0 ) a=(a_{3},a_{2},a_{1},a_{0}),b=(b_{3},b_{2},b_{1},b_{0}) and their product c = ( c ​ 7 ​ … ​ c 0 ) c=(c7\dots c_{0}) . Operands are written least-significant digit first, similar to other algorithmic setups ( Deng et al., 2024 ; Deng et al., 2023 ; Lee et al., 2023 ) .

For ICoT, the training data includes intermediate chain-of-thought (CoT) tokens q i q_{i} that explicitly record the step-by-step calculations. As a simple illustration, consider 12 × 34 12\times 34 . The tokens appearing between the two equal signs follow the same CoT format used in our 4 × 4 4\times 4 -digit multiplication tasks:

12 ∗ 34 = 48 ⏟ 12 ∗ 4 + 360 ⏟ 12 ∗ 30 ​ ( 408 ) ⏟ running sum = 408 \displaystyle 12*34=\underbrace{48}_{12*4}+\underbrace{360}_{12*30}\ \underbrace{(408)}_{\text{running sum}}=408

At each training epoch, a fixed number of CoT tokens are removed from the left of the chain. Concretely, the training examples at each epoch may have the following form: ( Epoch 1 ) \displaystyle(\text{Epoch 1}) a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 % % % ​ q 0 ​ … ​ q i ​ … ​ q j ​ … ​ q k ​ … ​ q τ ​ # ​ # ​ # ​ # ​ c 0 ​ … ​ c 7 \displaystyle\quad a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3}\%\%\%\ q_{0}\ldots q_{i}\ldots q_{j}\ldots q_{k}\ldots q_{\tau}\ \#\#\#\#\ c_{0}\ldots c_{7} ( Epoch 2 ) \displaystyle(\text{Epoch 2}) a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 % % % ​ q i ​ … ​ q j ​ … ​ q k ​ … ​ q τ ​ # ​ # ​ # ​ # ​ c 0 ​ … ​ c 7 \displaystyle\quad a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3}\%\%\%\ q_{i}\ldots q_{j}\ldots q_{k}\ldots q_{\tau}\ \#\#\#\#\ c_{0}\ldots c_{7} ( Epoch 3 ) \displaystyle(\text{Epoch 3}) a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 % % % ​ q j ​ … ​ q k ​ … ​ q τ ​ # ​ # ​ # ​ # ​ c 0 ​ … ​ c 7 \displaystyle\quad a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3}\%\%\%\ q_{j}\ldots q_{k}\ldots q_{\tau}\ \#\#\#\#\ c_{0}\ldots c_{7} … \displaystyle\ldots ( Epoch N ) \displaystyle(\text{Epoch N}) a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 % % % ​ # ​ # ​ # ​ # ​ c 0 ​ … ​ c 7 \displaystyle\quad a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3}\%\%\%\ \#\#\#\#\ c_{0}\ldots c_{7} where q i q_{i} are CoT tokens and % , # \%,\# are special delimiters. 2 2 2 These delimiters have no special meaning beyond matching the setup of Deng et al. (2024) . Note that after each epoch, the model sees a shorter chain by truncating some tokens, and that by the end, only the operands and final answer remain. For comparison, standard fine-tuning only trains on the operands: a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 % % % ​ # ​ # ​ # ​ # ​ c 0 ​ … ​ c 7 a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3}\%\%\%\#\#\#\#\ c_{0}\ldots c_{7} .

Interestingly, the ICoT model is able to achieve 100% accuracy on 4 × \times 4 digit multiplication, while standard fine-tuning only achieves less than 1% accuracy. Note that scaling does not help – scaling to a 12 layer 8 head model achieves the same < 1 % <1\% accuracy, and Yang et al. (2023) show that fine-tuning a 2B model still plateaus at 95% accuracy.

For more details regarding training (data format, sample size, hyperparameters), see Appendix A .

##### Notations.

𝐡 t ℓ \mathbf{h}^{\ell}_{t} indicates the hidden states at layer ℓ \ell timestep t t . Timesteps for solution tokens c k , k = [ 0 , … , 7 ] c_{k},k=[0,\ldots,7] are notated t c k t_{c_{k}} . Att h ℓ \textsc{Att}^{\ell}_{h} ( ⋅ ) (\cdot) , MLP ℓ \textsc{MLP}^{\ell} ( ⋅ ) (\cdot) indicate the output of the attention heads or MLP blocks at layer ℓ \ell , head index h h . E , U ∈ ℝ V × d E,U\in\mathbb{R}^{V\times d} indicate (un)embedding weights.

## 3 Comparing the Mechanisms of ICoT versus SFT

### 3.1 Long-range dependencies in multi-digit multiplication

Here we discuss how one might solve multi-digit multiplication, and the required long-range dependencies needed to solve multiplication.

One approach to compute each digit, c k c_{k} , is as follows:

s k ≜ ∑ i + j = k a i ​ b j , ⏟ sum of partial products c k = ( s k + r k − 1 ) ​ mod ​ 10 , r k = ⌊ s k + r k − 1 10 ⌋ ⏟ carry , r − 1 = 0 \displaystyle s_{k}\triangleq\underbrace{\sum_{i+j=k}a_{i}b_{j},}_{\text{sum of partial products}}\quad c_{k}=(s_{k}+r_{k-1})\ \text{mod}\ 10,\quad r_{k}=\underbrace{\big\lfloor\frac{s_{k}+r_{k-1}}{10}\big\rfloor}_{\text{carry}},\quad r_{-1}=0 (1)

Note that both c k c_{k} and r k r_{k} can be expressed with an intermediary term c ^ k \hat{c}_{k} , which encapsulates both the relevant information from the partial products and the carry:

c ^ k ≜ s k + r k − 1 , c k = c ^ k ​ ( mod ​ 10 ) , r k = ⌊ c ^ k 10 ⌋ \displaystyle\hat{c}_{k}\triangleq s_{k}+r_{k-1},\qquad c_{k}=\hat{c}_{k}\;(\mathrm{mod}\;10),\qquad r_{k}=\big\lfloor\frac{\hat{c}_{k}}{10}\big\rfloor (2)

Importantly, note the long-range dependencies needed for multi-digit multiplication. Specifically, we highlight two observations: (i) To determine c k c_{k} , one must use all the partial products { a i ​ b j | i + j ≤ k } \{a_{i}b_{j}|i+j\leq k\} , since all of these terms contribute to c k c_{k} . (ii) Knowing the intermediary term c ^ k \hat{c}_{k} suffices to compute c k c_{k} and to propagate necessary information for later digits. Thus we use c ^ k \hat{c}_{k} as a probing signature (Section 3.2 ) at each timestep t c k t_{c_{k}} to check if the model is utilizing all the necessary long-range information to predict the correct tokens c k c_{k} .

In the following sections, we demonstrate how the ICoT model satisfies such long-range dependency while the standard fine-tuning model does not.

### 3.2 Evidence of Long-Range Dependencies in ICoT

We first demonstrate two lines of evidence that the ICoT model satisfies long-range dependencies in multi-digit multiplication, while the standard fine-tuning model does not.

Logit Attributions. Note from Figure 1 that digits a i , b i a_{i},b_{i} can only affect c k c_{k} terms where k ≥ i k\geq i . Also note that at timestep t c k t_{c_{k}} , the pairwise products { a i ​ b j | i + j = k } \{a_{i}b_{j}|i+j=k\} affect the final prediction c k c_{k} the most. “Earlier” pairwise products { a i ​ b j | i + j ≤ k } \{a_{i}b_{j}|i+j\leq k\} can still affect c k c_{k} , but with diminishing effects as i + j i+j gets smaller.

We directly test for these relationships in our ICoT and SFT models using logit attributions. Namely, given an input sample orig := a 0 ​ a 1 ​ a 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 \textsc{orig}:=a_{0}a_{1}a_{2}a_{3}*b_{0}b_{1}b_{2}b_{3} , we measure the logits of the model’s predictions for c 0 − 7 : logit c k ​ ( orig ) c_{0-7}:\text{logit}_{c_{k}}(\textsc{orig}) . We then randomly swap out one of the operand digits at timestep t t (e.g., a ~ 2 \tilde{a}_{2} ) to construct a counterfactual input counter t = a 0 ​ a 1 ​ a ~ 2 ​ a 3 ∗ b 0 ​ b 1 ​ b 2 ​ b 3 \textsc{counter}_{t}=a_{0}a_{1}\tilde{a}_{2}a_{3}*b_{0}b_{1}b_{2}b_{3} and measure the change in logits: Δ t , k = logit c k ​ ( orig ) − logit c k ​ ( counter t ) \Delta_{t,k}=\text{logit}_{c_{k}}(\textsc{orig})-\text{logit}_{c_{k}}(\textsc{counter}_{t}) Thus Δ t , k \Delta_{t,k} measures the effect that digit at timestep t t has on the prediction of token c k c_{k} .

We use 1,000 samples for each ( t , k ) (t,k) pair and show the results in Figure 2 . Note that for SFT, the model does not see the correct dependencies between earlier tokens to middle tokens, while the ICoT model does, suggesting that the model has indeed learned the correct long-range dependencies.

Probing for 𝐜 ^ k \mathbf{\hat{c}}_{k} . Note from Figure 1 and Equation 2 that the long-range dependencies can be captured by an intermediate term, c ^ k \hat{c}_{k} . We test for whether c ^ k \hat{c}_{k} information can be decoded from the hidden states of the models using linear regression probes. Namely, at each timestep t c k t_{c_{k}} we predict for c ^ k \hat{c}_{k} by training a single vector 𝐰 k ∈ ℝ d \mathbf{w}_{k}\in\mathbb{R}{}^{d} such that 𝐰 k ​ 𝐡 t c ​ k 2 . mid = c ^ k \mathbf{w}_{k}\mathbf{h}_{t_{ck}}^{2.\text{mid}}=\hat{c}_{k} using a MSE loss, where 𝐡 2 . mid \mathbf{h}^{2.\text{mid}} is the hidden state at layer 2 after attention heads, before MLPs.

Figure 3 reports the mean absolute error from probing for c ^ k \hat{c}_{k} for middle and late digits, k = 2 , … , 6 k=2,\ldots,6 . Note that the accuracy from the ICoT model is much higher than that of SFT, further suggesting that the ICoT model has learned the correct long-range dependencies while SFT has not.

### 3.3 Encoding Long-Range Dependencies via Attention Trees

How does the ICoT model compute long-range dependencies? Here we describe how the model’s attention patterns induce a shallow directed acyclic graph, akin to a binary expression tree, in order to encode long-range dependencies.

Namely, in the first layer, across all timesteps t > 5 t>5 , 3 3 3 Note that only after timestep 5, both a a and b b tokens appear in the context. each attention head only attends to a pair of digit tokens, { a i , b j } \{a_{i},b_{j}\} (Figure 4 , left). This allows the model to produce the pairwise product a i ​ b j a_{i}b_{j} (see Section 4.1 for how attention heads represent pairwise products), but also allows the model to cache the product a i ​ b j a_{i}b_{j} in the hidden state of layer 1 at timestep t t (i.e., 𝐡 t 1 \mathbf{h}^{1}_{t} ). Put differently, product pairs { a i ​ b j } i , j ∈ { 0 , … ​ 4 } \{a_{i}b_{j}\}_{i,j\in\{0,\ldots 4\}} are “cached” in the first layer across different timesteps ( 𝐡 t 1 , t < t c k \mathbf{h}^{1}_{t},t<t_{c_{k}} ).

At later timesteps t ≥ t c k t\geq t_{c_{k}} , when the model predicts solution tokens c k c_{k} , this allows the second layer attention heads to attend to a small set of previous cache sites , i.e., where the appropriate pairs of products a i ​ b j , i + j = k a_{i}b_{j},i+j=k are stored from earlier timesteps.

Example: Figure 4 depicts the attention patterns when the model predicts c 2 c_{2} , given input “ a 0 ​ … ​ 3 ∗ b 0 ​ … ​ 3 = c 0 ​ c 1 a_{0\ldots 3}*b_{0\ldots 3}=c_{0}c_{1} ”. These attention maps are averaged from 1,000 samples from a held out test set. The necessary terms to compute c 2 c_{2} are a 2 ​ b 0 , a 1 ​ b 1 , a 0 ​ b 2 a_{2}b_{0},a_{1}b_{1},a_{0}b_{2} , and c ^ 1 \hat{c}_{1} (which in turn requires a 1 ​ b 0 , a 0 ​ b 1 , a 0 ​ b 0 a_{1}b_{0},a_{0}b_{1},a_{0}b_{0} ).