🤖 machine learning

Explaining and Tuning Transformer-based LLMs in Arithmetic Tasks with Human Strategies

This paper demonstrates that transformer-based large language models can achieve improved arithmetic performance by adopting human-like learning strategies, such as task decomposition and cognitive empowerment, suggesting shared cognitive processes between LLMs and human learners while enhancing model reliability through explainable AI verification.

Original authors: Luyu Qiu, Jianing Li, Hwanhee Kim, Xiaoyong Wei, Yueyuan Zheng, Janet Hsiao, Lei Chen

Published 2026-07-21

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Luyu Qiu, Jianing Li, Hwanhee Kim, Xiaoyong Wei, Yueyuan Zheng, Janet Hsiao, Lei Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Brainy Robot and the Math Problem

Imagine a world where computers have learned to read almost every book, website, and conversation ever written. These are called Large Language Models (LLMs), and they are like super-smart robots that can write stories, answer questions, and chat with you. They are built on a special architecture called a "Transformer," which acts like a giant spotlight, helping the robot pay attention to the most important words in a sentence at the same time. You might think a robot that knows so much about language would be great at math, too. But here's the twist: while these robots are wizards with words, they often stumble over simple math problems, like adding two numbers together. It's like a person who can write a beautiful poem but can't figure out how much change they should get back at a store. This is a big deal because if we want to trust these robots with important jobs—like helping doctors or managing money—we need to know if they can do the basics correctly. Scientists are trying to figure out why these smart robots fail at simple math and, more importantly, how to teach them to get better.

The Paper's Story: Teaching Robots to Think Like Humans

In this study, the researchers decided to treat the math problem like a mystery to be solved. They didn't just throw more data at the robot and hope for the best. Instead, they asked a fascinating question: Do these AI robots learn math the same way humans do?

To find out, they took a standard "vanilla" Transformer model (a basic version of the AI) and taught it how to do integer arithmetic—adding, multiplying, and dividing whole numbers. They discovered that the AI wasn't just guessing; it was actually learning, but it was struggling with the order of operations and the complexity of the steps.

The "Lego" Breakdown
The researchers realized that just like humans break a big, scary math problem into smaller, easier steps, the AI needed to do the same. They broke down the tasks into tiny "subtasks."

For Addition: Imagine adding two long numbers. You don't just look at the whole thing; you add the last digits, then the next ones, carrying over any "extra" numbers (like carrying a 1) to the next column. The researchers found that the AI learned the easy parts (adding the last digits) very quickly, but it got confused when it had to remember and use those "carried" numbers for the next steps. It's like learning to walk before you can run.
For Multiplication: This is even trickier. When you multiply big numbers, you have to do a bunch of smaller multiplications and then add them all up. The paper found that the AI was trying to do everything at once and getting overwhelmed. However, when the researchers taught the AI to write the answer in reverse order (starting from the last digit and working backward, just like humans often do on paper), the AI's performance improved significantly for simple single-digit multiplier tasks. But here's the catch: for complex multi-digit multiplication, reversing the digits alone didn't fix the problem; the AI still struggled because the intermediate steps were too hard for a single-layer model to handle.

The "Human" Tricks
The most exciting part of the paper is that the researchers tested strategies that human teachers use to help students, and they worked on the AI, too!

Chain of Thought (CoT): Instead of asking the AI, "What is 57257 times 51422?" and expecting a magic answer, they asked it to show its work step-by-step. They told the AI to break the big multiplication into smaller, easier multiplications first. This is exactly what a human teacher tells a student: "Don't panic; just do one step at a time." When the AI used CoT combined with reversing the answer digits, its accuracy jumped significantly (reaching about 79%). To get to near-perfect scores, the researchers found they needed to combine CoT with reversing the digits and making the model "deeper."
Reversing the Output: As mentioned, for multiplication, writing the answer backward (starting from the ones place) helped the AI focus on the right numbers, but this trick has a limit. It worked well for simpler tasks but wasn't enough on its own for complex multi-digit multiplication. Interestingly, for division, reversing the digits actually made things worse. This is because division calculations naturally start from the higher-order digits (the big numbers on the left), which aligns with human experience. Trying to reverse the order for division confused the model.
Deeper Models: They also found that making the AI "deeper" (adding more layers of thinking) helped it handle the complex middle steps of division and the heavy lifting of multiplication. This gave the model the extra "brainpower" needed to keep track of intermediate results without dropping the ball.

What They Found (and What They Didn't)
The researchers didn't just say "it works"; they looked inside the robot's brain to see how it worked. They used a technique called "visualization" to watch which parts of the AI's attention mechanism were lighting up. They saw that the AI had specific "attention heads" (little internal workers) that learned to do specific jobs, like "Base Multiply" or "Make Carry."

The Good News: The AI can learn these patterns. When given the right structure (like breaking tasks down, using CoT, reversing digits where appropriate, and adjusting the model depth), it performs almost perfectly.
The Bad News: Without these human-like strategies, the AI struggles. For example, on a standard 5-digit multiplication task, a basic AI might only get 13% right. But with the "human" tricks combined, it gets 100%.
The "Why": The paper suggests that the AI's failure isn't because it's "dumb," but because it's trying to solve the problem in a way that doesn't match how the math is structured. It's like trying to drive a car by pushing it from the back instead of using the steering wheel. Also, the "right" way to drive depends on the car: for multiplication, you might want to go backward, but for division, you must go forward.

The Big Picture
This paper suggests that Transformer-based AI models might share some learning patterns with human students. They both learn simple things first, then combine them to solve harder problems. By using "Explainable AI" (XAI) techniques, the researchers could see exactly where the AI was getting stuck and fix it.

The study concludes that while current AI models are amazing at language, they need a little help with math. But the good news is that we don't need to invent a new kind of robot; we just need to teach them to think like us. By using strategies like breaking problems down, showing our work, and organizing our thoughts (and knowing when to reverse the order and when not to), we can make these AI models much more reliable. This is a huge step forward for making AI safe and trustworthy, especially for the important jobs where getting the math right is the difference between success and failure. The paper doesn't claim to have solved every math problem in the universe, but it has shown a clear path forward: if we treat AI like a student who needs a good teacher, it can learn to do the math we need it to do.

Technical Summary: Explaining and Tuning Transformer-based LLMs in Arithmetic Tasks with Human Strategies

1. Problem Statement

Despite achieving state-of-the-art performance in complex Natural Language Processing (NLP) tasks, Transformer-based Large Language Models (LLMs) exhibit significant deficiencies in elementary arithmetic tasks, such as basic integer calculations. This discrepancy raises concerns regarding model reliability and safety. Unlike typical NLP tasks which allow for diverse valid outputs, arithmetic problems possess a single definitive answer where every digit is critical. Furthermore, arithmetic operations often require processing from the least significant digit (e.g., in multiplication) or involve complex intermediate steps (e.g., carries and remainders) that differ from the standard left-to-right processing of natural language. Existing Explainable AI (XAI) methods often treat Transformers as black boxes or focus on specific task instances, failing to capture the intrinsic correlations and internal mechanisms required for systematic arithmetic computation.

2. Methodology

The authors propose a comprehensive problem analysis framework grounded in cognitive science and educational psychology to inspect, explain, and improve Transformer performance on integer arithmetic. The methodology consists of three core components:

2.1 Comprehensive Problem Analysis Framework

The framework decomposes the learning process into three stages:

Subtask Decomposition: Arithmetic tasks are broken down into well-defined subtasks (e.g., Base Add, Make Carry, Use Carry) to identify which specific components are difficult for the model.
Learning Curve Analysis: The authors analyze per-digit training loss curves to observe the learning speed of specific output digits and subtasks, comparing them against human learning patterns (e.g., faster learning on simpler subtasks).
Model Visualization: Attention maps are visualized to determine how specific attention heads contribute to different subtasks and whether the model's internal mechanisms align with human specialized functioning.

2.2 Human Strategy and Capacity Transfer

Drawing from educational psychology, the study tests whether strategies effective for human learners improve LLM performance:

Computational Strategies: Reformulating arithmetic rules to match the model's processing, such as reversing the order of answer digits (processing from least significant to most significant) to align with the self-attention mechanism's ability to handle dependencies.
Chain of Thought (CoT): Decomposing complex multi-digit operations into a sequence of simpler intermediate steps (e.g., breaking $m \times m$ multiplication into $m \times u$ steps).
Cognitive Empowerment: Increasing model depth (number of layers) to enhance working memory capacity and handle complex intermediate results.

2.3 Experimental Setup

The study utilizes a single-layer decoder-only Transformer (with Multi-Head Self-Attention and Feed-Forward layers) trained on integer addition, multiplication, and division. The input consists of operands and operators, and the model predicts answer digits autoregressively. Experiments include ablation studies on attention heads, varying the number of digits (1 to 15), and testing different data formats (ordinal vs. reversed, with and without CoT).

3. Key Findings and Results

3.1 Learning Patterns and Subtask Specialization

Human-Like Learning: LLMs exhibit learning patterns similar to humans, converging faster on simpler subtasks (e.g., calculating the least significant digit without carry) and slower on complex ones requiring cascading dependencies.
Attention Head Specialization: Visualization reveals that distinct attention heads specialize in specific subtasks. For example, in addition, one head may handle "Base Add" while another handles "Use Carry."
Head Count Saturation: Performance improves significantly as the number of attention heads increases, saturating at three heads. With fewer than three heads, multiple tasks are forced into a single head, causing confusion and lower accuracy. With three or more heads, task separation is clear.

3.2 Impact of Data Format and Strategies

Reversed Digits: Reversing the output order (processing from least significant digit) significantly improves performance in multiplication tasks, allowing the model to leverage previously generated digits for current calculations. However, this strategy is less effective or detrimental for division, which naturally proceeds from the most significant digit.
Chain of Thought (CoT): CoT prompting effectively decomposes complex $m \times m$ multiplication and $m \div m$ division into simpler steps, drastically reducing prediction difficulty and improving accuracy.
Model Depth: Increasing the number of layers enhances the model's capacity to handle intermediate remainders and complex dependencies, particularly in division tasks where single-layer models struggle.

3.3 Quantitative Improvements

Multiplication: Vanilla Transformers struggle with multi-digit multiplication (e.g., dropping to ~13% accuracy for 4-digit multiplication). By combining reversed digit formats, increased depth, and CoT, accuracy reaches near 100%.
Division: Similar improvements are observed, with CoT and increased depth addressing the accumulation of errors in remainder calculations.
Ablation Studies: Removing specific attention heads (e.g., Head 2 in multiplication) causes significant loss increases, confirming the critical role of specific heads in specific subtasks (e.g., Base Multiply).

4. Key Contributions

Problem Analysis Framework: The paper proposes a novel framework combining subtask decomposition, learning curve analysis, and visualization to explain how Transformer-based LLMs handle arithmetic, moving beyond black-box evaluation.
Performance Enhancement via Human Strategies: The authors demonstrate that applying human cognitive strategies (reversed processing, CoT) and capacity enhancements (increasing depth) significantly improves arithmetic accuracy, suggesting a shared cognitive process between LLMs and human learners.
Interpretability and Mechanism Insight: The work provides a detailed, visual explanation of the internal mechanisms of Transformers in arithmetic, identifying specific attention heads responsible for distinct subtasks and validating these findings through ablation studies.

5. Significance and Claims

The paper claims that Transformer-based LLMs share cognitive processes with human learners in arithmetic, specifically in how they decompose tasks and learn subtasks sequentially. By leveraging explainable AI (XAI) verifications, the authors argue that understanding these internal mechanisms is crucial for improving model reliability.

The significance of this work lies in its potential to foster trust in LLMs for critical and high-stakes applications. By demonstrating that performance gaps in arithmetic are not inherent limitations of the architecture but rather issues of data format, task decomposition, and capacity, the paper offers a pathway to tune models for higher reliability. The findings contribute to the broader fields of model understanding and interpretability, providing a foundation for analyzing more complex tasks and architectures. The authors modestly position this work as an exploration of similarities between LLMs and human learners, supported by empirical evidence, rather than a definitive proof of human-like consciousness.