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

TL;DR

Reverse-engineers a 2-layer, 4-head Transformer trained on 4×4-digit multiplication with implicit chain-of-thought (ICoT) to explain why standard fine-tuning (SFT) fails at the same task. The ICoT model reaches 100% accuracy; SFT plateaus near 1% even when scaled to 12 layers / 8 heads. The mechanism inside the ICoT model is a binary-tree-like attention DAG that caches pairwise digit products in earlier tokens for later retrieval, Minkowski-sum geometry in the attention outputs, and digits laid out on a pentagonal prism via a Fourier basis. With this understanding the authors add a single auxiliary linear-probe loss predicting the “running partial sum” — and the SFT model now learns 4×4 multiplication to 99% without any chain-of-thought supervision.

Key claims

A 2-layer, 4-head Transformer trained with ICoT reaches 100% accuracy on 4×4-digit multiplication while a matched standard-fine-tuned model achieves ~1% and digit-level accuracy plateaus near 81% even after 60 epochs; scaling SFT to 12 layers / 8 heads does not change this [§2, App. A.2].
The ICoT model encodes the long-range dependency Z_i = floor((Z_{i-1} + sum_{j+k=i} a_j·b_k) / 10); logit-attribution counterfactuals show ICoT but not SFT has the correct earlier-digit → middle-digit dependency structure [§3.1, §3.2, Fig. 2].
Linear-regression probes can decode Z_i from the ICoT hidden states at layer 2 with much lower mean-absolute-error than from the SFT model, for middle and late digits [§3.2, Fig. 3].
Mechanistically, layer-1 heads in the ICoT model attend to only a single pair of digit tokens, computing and caching pairwise partial products into the hidden state at that position; layer-2 heads then attend to those cache sites to build the per-output-digit running sum, forming a sparse binary-tree-like DAG [§3.3, Fig. 4].
When attention is sparse over two digits, the head’s output is a Minkowski sum of value-projected digit embeddings; the resulting nested-cluster geometry (global clusters indexed by one digit, local clusters of the same shape indexed by the other) follows from the covariance decomposition Σ = E[Σ_local] + Σ(E[...]) [§4.1, Eq. 5, Fig. 5].
The 10 digit embeddings (and the final hidden state, and the last-MLP weights) sit on the vertices of a pentagonal prism when projected onto their top PCs, explained by a 6-term Fourier basis with parity vector + two cos/sin pairs at frequencies 1/5 and 2/5 [§4.2, Fig. 6, App. B].
In standard fine-tuning the model first learns the easy first two and last digit, then plateaus on the middle digits — losses and gradient norms freeze at a local optimum that lacks the long-range cache-and-retrieve structure [§5, Fig. 7a].
Adding a per-timestep auxiliary MSE loss that trains a linear regression head on the layer-2 attention output to predict the running partial sum Z_i lets the SFT model reach 99% on 4×4 multiplication without any CoT supervision, and the resulting model independently rediscovers a similar attention-tree structure [§6, §6.1, Fig. 7b, Fig. 8].

Method

The setup is deliberately minimal: GPT-style 2-layer 4-head models trained from scratch on 80,800 4-by-4 multiplication examples written least-significant-digit-first. ICoT follows Deng et al. 2024 — explicit step-by-step partial-product CoT tokens are inserted between operands and answer, and at each epoch a fixed number (8) of CoT tokens is truncated from the left of the chain, so by the end the model sees only operands and final answer. Standard fine-tuning trains on operands → answer directly. Probes are linear regressions on layer-2 hidden states (post-attention, pre-MLP) trained with MSE against the analytic running partial sum Z_i = floor((Z_{i-1} + Σ_{j+k=i} a_j b_k)/10). Logit-attribution counterfactuals swap one operand digit and measure the change in answer-token logits.

The “fix without CoT” attaches a single linear regression head to each layer-2 attention head’s output and adds λ · MSE(linear_head(h_{2,t}), Z_{i(t)}) to the standard LM loss. No chain-of-thought tokens are ever shown to the model — only the prediction of the partial-sum scalar.

Results

ICoT 2L/4H: 100% test accuracy on 4×4; SFT 2L/4H: ~1% sequence accuracy, ~81% digit-level [§2, App. A.2].
SFT scaled to 12L/8H: still ~1%; Yang et al. fine-tuning a 2B model plateaus at 95% [§2 cited].
Auxiliary-running-sum SFT 2L/4H: 99% accuracy on 4×4 without any CoT tokens [§6].
Probe MAE for Z_i on middle digits is dramatically lower for ICoT than SFT (qualitative, Fig. 3).
Auxiliary-loss model independently grows an “attention tree” qualitatively similar to ICoT’s, including a single layer-2 head that attends to all the necessary digits at each timestep in a parallelogram pattern [§6.1, Fig. 8].

Why it’s interesting

The mechanistic story is precise enough to fall through into a working training intervention — that is the test mech-interp keeps trying to pass, and this is one of the cleaner instances. It complements Universal Sparse Autoencoders: Interpretable Cross-Model Concept Alignment by anchoring the interp pipeline on a circuit-level finding (attention DAG + Minkowski sum + Fourier digit basis) where the universal-SAE work targets cross-model feature alignment — same field, different primitives. It lines up with Emotion Concepts and their Function in a Large Language Model on the recurring “find the right direction / loss term, intervene on behavior” pattern, but does it at the level of an algorithmic capability rather than an alignment-relevant disposition. Practically, the result that a 2B SFT model plateaus at 95% on 4×4 multiplication while a 2-layer model with a one-line auxiliary loss reaches 99% is an existence proof against the “just scale” reflex on long-range-dependency tasks, of a piece with the broader skepticism in The Illusion of Thinking: Understanding the Strengths and Limitations of Reasoning Models via the Lens of Problem Complexity but with a constructive remedy.