PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xiaoyan Bai, Itamar Pres, Yuntian Deng, Chenhao Tan, Stuart Shieber, Fernanda Viégas, Martin Wattenberg, Andrew Lee

arXiv:2510.00184v1cs.LGcs.AI

TL;DR

Transformers can fail at multi-digit multiplication even after task fine-tuning, raising the question of how they learn its long-range dependencies. The paper reverse-engineers a successful ICoT model and uses those findings to motivate an auxiliary loss. That loss provides an inductive bias for the dependencies and achieves perfect accuracy without chain-of-thought supervision.

  • Problem

    Even large language models can fail at 4×4-digit multiplication, including after task fine-tuning, motivating the question of why Transformers struggle with this task.

    Source: §1 · p. 1

  • Method

    The paper reverse-engineers a successful ICoT model and contrasts it with standard fine-tuning, whose training gradually removes explicit chain-of-thought tokens.

    Source: §1 · p. 1

  • Results

    A running-partial-sum auxiliary loss provides an inductive bias for long-range dependencies and enables perfect multiplication accuracy without chain-of-thought supervision.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings identify long-range dependency learning as a training pitfall and show a task-specific inductive bias that addresses it, while a generic solution remains future work.

    Source: §1 · p. 2, §8 · p. 9

    The demonstrated inductive bias is specific to multiplication; generalizing such biases to other tasks with long-range dependencies is left for future work.

    Source: §6 · p. 8

Abstract

from arXiv · show

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 \emph{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.

1 INTRODUCTION

The paper investigates why Transformers struggle with multi-digit multiplication by reverse-engineering a successful implicit chain-of-thought model. It identifies its long-range dependency mechanism and representations, then uses those findings to motivate an auxiliary loss that enables perfect accuracy without chain-of-thought supervision.

1 INTRODUCTION

  • The study asks why models can fail at multi-digit multiplication despite broader capabilities, including after task fine-tuning.The introduction notes that Llama-3.2 90B and GPT4 fail at 4×4-digit multiplication.Source: §1 · p. 1
  • The ICoT model learns the long-range structure needed for multiplication, supported by logit attributions and linear regression probes.The standard fine-tuned model lacks this structure.Source: §1 · p. 1
  • Attention forms a sparse, binary-tree-like graph that selects digit pairs for partial products and caches them in earlier tokens for later retrieval.This mechanism organizes computations across timesteps.Source: §1 · p. 1
  • Attention heads form partial products as Minkowski sums of digit embeddings, while Fourier bases represent digits; the SFT model lacks these representations.The paper characterizes these representations as intuitive and efficient.Source: §1 · p. 2
  • A running-partial-sum auxiliary loss gives the model an inductive bias for long-range dependencies and achieves perfect multiplication accuracy without chain-of-thought supervision.Standard fine-tuning otherwise plateaus on middle digits without learning those dependencies.Source: §1 · p. 2

2 EXPERIMENT SETUP, TRAINING ICOT, NOTATIONS

The experiments compare a small ICoT model with standard fine-tuning on 4×4-digit multiplication. ICoT training gradually removes explicit intermediate steps, and the resulting model greatly outperforms standard fine-tuning on the task.

2 EXPERIMENT SETUP, TRAINING ICOT, NOTATIONS

  • The study compares a 2-layer, 4-head ICoT model with standard fine-tuning on 4×4-digit multiplication.The authors identify these as the smallest task and architecture settings where standard fine-tuning fails and ICoT succeeds.Source: §2 · p. 2
  • Operands are written least-significant digit first, and the illustrated chain includes a running sum between partial calculations.The example shows 12×4 and 12×30 contributing to 408.Source: §2 · p. 2
  • ICoT training supplies explicit chain-of-thought calculation tokens, then removes tokens over successive epochs until only operands and the final answer remain.Standard fine-tuning trains on operands and the final answer without those intermediate tokens.Source: §2 · p. 2
  • ICoT reaches 100% accuracy on 4×4-digit multiplication, while standard fine-tuning achieves less than 1%.The passage also reports that scaling to a 12-layer, 8-head model still achieves less than 1%.Source: §2 · p. 2
  • An intermediate value ĉ_i can carry information from which both the solution digit c_i and carry r_i can be derived.Source: Figure 1 · p. 3

3 COMPARING THE MECHANISMS OF ICOT VERSUS SFT

The section defines the long-range information required to compute each multiplication digit, then tests whether ICoT and standard fine-tuning encode it. The evidence favors ICoT, whose attention patterns form a tree-like cache-and-retrieve mechanism for partial products.

3.1 LONG-RANGE DEPENDENCIES IN MULTI-DIGIT MULTIPLICATION

  • Computing c_k requires all partial products a_i b_j where i + j ≤ k, and ĉ_k summarizes information sufficient to derive c_k and propagate information to later digits.Source: §3.1 · p. 3

3.2 EVIDENCE OF LONG-RANGE DEPENDENCIES IN ICOT

  • Logit-attribution and ĉ_k probing results indicate that ICoT captures long-range dependencies that standard fine-tuning misses.The logit analysis tests operand-digit effects on output logits; the probes test whether ĉ_k is decodable from hidden states.Source: §3.2 · p. 4, Figure 2 · p. 4, Figure 3 · p. 4

3.3 ENCODING LONG-RANGE DEPENDENCIES VIA ATTENTION TREES

  • For c2, reconstructed attention paths retrieve a2b0, a1b1, and a0b2, along with earlier products needed for ĉ_1.Those earlier products include a1b0, a0b1, and a0b0.Source: §3.3 · p. 5
  • The authors report that corresponding attention trees can be reconstructed for all output digits c0 through c7.Source: §3.3 · p. 5
  • For each output digit, attention forms a shallow binary-tree-like graph: first-layer heads compute and cache digit-pair products, which later second-layer heads retrieve.The graph is spread across timesteps and selects the relevant pairs for each output.Source: §3.3 · pp. 4–5

4 FEATURE GEOMETRY OF ICOT

ICoT attention outputs organize digit pairs into nested Minkowski-sum geometries, while Fourier structure gives its digit representations a pentagonal-prism form unlike SFT.

4.1 DIGIT-WISE MULTIPLICATIONS AS MINKOWSKI SUMS

  • Attention-head outputs form Minkowski sums of paired digit embeddings, with nested global and local clusters corresponding to ai and bj.The global clusters correspond to ai, and the nested local clusters have the same geometry for bj.Source: §4.1 · p. 6, Figure 5 · p. 6

4.2 EMBEDDING DIGITS ON A PENTAGONAL PRISM VIA FOURIER BASES

  • ICoT embeddings, final hidden states, and last-MLP weights are well reconstructed from a small set of Fourier basis functions.The final hidden layer can be reconstructed via six terms.Source: §4.2 · pp. 6–7
  • ICoT's ten digits form a pentagonal prism in the final hidden layer, while SFT hidden states show no obvious pattern.The first principal component separates even and odd digits, and the two pentagons are stacked along that component.Source: §4.2 · pp. 6–7, Figure 6 · p. 7

5 PITFALLS OF LEARNING: LACK OF LONG-RANGE DEPENDENCY

Standard fine-tuning learns some output digits but plateaus on the middle digits, consistent with a local optimum lacking the needed long-range dependencies; auxiliary-loss training learns the middle digits.

5 PITFALLS OF LEARNING: LACK OF LONG-RANGE DEPENDENCY

  • SFT learns c0, c1, and c7 first and eventually c2, but c3–c6 remain unlearned with plateaued losses, suggesting a local optimum without needed long-range dependencies.Source: §5 · p. 7
  • The same failure pattern appears in a 12-layer, 8-head model, so scaling up does not resolve the issue.Source: §5 · p. 7
  • Standard fine-tuning plateaus on middle digits, whereas auxiliary-loss training allows the model to learn them.Both methods learn c0, c1, and c7 first.Source: Figure 7 · p. 8

6 LEARNING MULTIPLICATION WITHOUT ICOT

The section tests whether auxiliary supervision of an intermediate running sum can help a Transformer learn multiplication without implicit chain-of-thought. The 2-layer model reaches 99% accuracy on 4×4 multiplication, and its attention shows a pattern similar to—but not necessarily identical to—the ICoT mechanism.

6 LEARNING MULTIPLICATION WITHOUT ICOT

  • Two second-layer attention-head outputs feed linear regression heads trained with MSE to predict the accumulated sum at each timestep.The auxiliary loss is added to the standard language-modeling loss.Source: §6 · p. 8
  • The task’s inductive bias lets a 2-layer model learn 4×4 multiplication with 99% accuracy, while a larger 12-layer model still fails at multiplication under standard fine-tuning.Source: §6 · p. 8
  • The model learns the early and last output digits first, then works inward, learning c5 last.The reported order is c0, c1, c7; then c2, c3, c4, c6; and finally c5.Source: §6 · p. 8
  • The suggested inductive bias is specific to multiplication; more general inductive biases for long-range tasks are left for future work.The authors describe the task-specific bias as an example of overcoming a long-range dependency pitfall.Source: §6 · p. 8

6.1 DOES THE MODEL WITH AUXILIARY LOSS LEARN THE SAME MECHANISM AS ICOT?

  • The auxiliary-loss model forms a sparse attention tree and also has a head that attends to all necessary pairwise digits simultaneously.The attention pattern suggests a similar, but not necessarily exact, mechanism to ICoT.Source: §6.1 · p. 9, Figure 8 · p. 9

7 RELATED WORK

Prior work studies Transformers in controlled settings, often using arithmetic tasks to characterize their behavior. The paper also relates implicit chain-of-thought to process supervision because it gradually removes reasoning tokens during training.

Studying Transformers with Arithmetic Tasks

  • Arithmetic tasks have been used to study grokking, Fourier-based digit representations, heuristics, and length generalization in Transformers.The paper positions its study as an investigation of why Transformers fail to learn multi-digit multiplication.Source: §7 · p. 9

Process Supervision

  • Process supervision provides feedback on intermediate reasoning steps, with prior examples reporting results on GSM8K and MATH.The cited examples include human feedback and automated step-wise rewards.Source: §7 · p. 9
  • ICoT is presented as latent-space process supervision: chain-of-thought tokens are gradually removed so the model internalizes the reasoning procedure.The authors use ICoT's multiplication success to study why Transformers fail.Source: §7 · p. 9

8 CONCLUSION

The paper reverse-engineers an implicit chain-of-thought model to study how Transformers learn the long-range dependencies needed for multi-digit multiplication. It identifies a limitation of standard training and presents a task-specific inductive bias as an example of addressing it.

8 CONCLUSION

  • Reverse-engineering the ICoT model reveals that it learns the long-range dependencies needed for multi-digit multiplication.The paper studies this mechanism to understand why Transformers fail on the task.Source: §8 · p. 9
  • Gradient descent with an auto-regressive loss does not encourage Transformers to learn the right long-range dependencies, according to the paper's conclusion.Source: §8 · p. 9
  • The paper offers a simple task-specific inductive bias as an example and leaves generic solutions for future work.Source: §8 · p. 9

A TRAINING DETAILS

The appendix section introduces additional details about the training procedures of the models.

A TRAINING DETAILS

A.1 ICOT TRAINING

The experiments contrast implicit chain-of-thought training, which succeeds at multiplication, with standard fine-tuning, which plateaus at low accuracy. The ICoT training sequence gradually removes explicit intermediate tokens.

A.1 ICOT TRAINING

  • ICoT training gradually removes explicit chain-of-thought tokens until only the operands and final answer remain.In the described schedule, tokens are truncated across epochs, with eight removed at each stage in the experiment.Source: §A.1 · p. 11

A.2 STANDARD FINE-TUNING

  • ICoT reaches 100% test accuracy, while standard fine-tuning remains near 1% train and validation accuracy and larger SFT reaches only 1% accuracy.The 2-layer SFT model reaches about 81% digit-level accuracy; the 12-layer, 8-head model reaches 80%.Source: §A.1 · p. 11, §A.2 · p. 12

A.1 ICOT TRAINING

  • The ICoT experiment trains a 2-layer, 4-head model from scratch on 80,800 samples and converges after 13 epochs.Validation and test sets each contain 1,000 held-out samples.Source: §A.1 · p. 11

B FOURIER STRUCTURE IN MODEL’S WEIGHTS, ACTIVATIONS

The section examines Fourier structure in the ICoT model and reports that even a larger model fails to learn the right long-range dependencies under standard fine-tuning. It also describes attention patterns from which ICoT solution-token trees can be derived.

B FOURIER STRUCTURE IN MODEL’S WEIGHTS, ACTIVATIONS

  • Embedding weights, final MLP output weights, and final hidden states are analyzed as digit-indexed vectors that lie in a low-dimensional trigonometric subspace.Each vector’s ten entries correspond to digits 0–9.Source: B FOURIER STRUCTURE IN MODEL’S… · p. 12
  • Six-dimensional Fourier fits have median R2 values of 0.85–0.99, while adding k = 3, 4 terms yields perfect R2 fits.The table reports median R2 of Fourier fits over digits n = 0 . . . 9.Source: B FOURIER STRUCTURE IN MODEL’S… · p. 12, Table 1 · p. 13

C PER TOKEN GRADIENTS AND LOSSES: 12-LAYER MODEL

D ATTENTION PATTERNS OF ALL MODELS

  • Attention patterns across the three models provide trees for each ICoT solution token.The paper presents patterns for all digits and derives attention trees for ICoT.Source: D ATTENTION PATTERNS OF ALL MODELS · p. 13
Loading 2510.00184v1…