Dissecting Multiplication in Transformers: Insights into LLMs

TL;DR

The paper tackles why transformers struggle with arithmetic, focusing on integer multiplication by training a vanilla transformer on n-digit multiplication and analyzing its internal computations. It reveals that multiplication decomposes into parallel subtasks and that the model learns these subtasks in a digit-wise, sequential order, with carry propagation as a key bottleneck. By reversing answer digits, increasing model depth, and adopting a curriculum that favors simple samples, the authors achieve substantial accuracy gains (up to about 99.9% on 5-digit multiplication with a small transformer) and provide an explainability pipeline that clarifies attention-head roles and intermediate steps. The work advances interpretability and reliability of Transformer-based models for structured tasks and offers concrete strategies for boosting arithmetic performance in LLMs, with potential implications for broader algorithmic reasoning tasks and safe deployment.

Abstract

Transformer-based large language models have achieved remarkable performance across various natural language processing tasks. However, they often struggle with seemingly easy tasks like arithmetic despite their vast capabilities. This stark disparity raise human's concerns about their safe and ethical use, hinder their widespread adoption.In this paper, we focus on a typical arithmetic task, integer multiplication, to explore and explain the imperfection of transformers in this domain. We provide comprehensive analysis of a vanilla transformer trained to perform n-digit integer multiplication. Our observations indicate that the model decomposes multiplication task into multiple parallel subtasks, sequentially optimizing each subtask for each digit to complete the final multiplication. Based on observation and analysis, we infer the reasons of transformers deficiencies in multiplication tasks lies in their difficulty in calculating successive carryovers and caching intermediate results, and confirmed this inference through experiments. Guided by these findings, we propose improvements to enhance transformers performance on multiplication tasks. These enhancements are validated through rigorous testing and mathematical modeling, not only enhance transformer's interpretability, but also improve its performance, e.g., we achieve over 99.9% accuracy on 5-digit integer multiplication with a tiny transformer, outperform LLMs GPT-4. Our method contributes to the broader fields of model understanding and interpretability, paving the way for analyzing more complex tasks and Transformer models. This work underscores the importance of explainable AI, helping to build trust in large language models and promoting their adoption in critical applications.

Figures (5)

• Figure 1: The decomposed steps of (a) addition, (b) multi-digit and unit-digit ($m\times u$) multiplication, and (c) multi-digit and multi-digit ($m\times m$) multiplication.
• Figure 2: Illustrations of (a) the overall per-digit loss curve, and (b-f) per-digit loss curve for each subtask.
• Figure 3: Attention map of ordinal and reversed answer digit format. To predict answer digit, the multiple attention heads are responsible for different tasks and combine the information in subsequent MLP layers.
• Figure 4: The overlap of per-digit product with different multiplier format. The darker the color, the more overlapping digits there are.
• Figure 5: Accuracy (%) of ordinal and reversed transformer trained with different proportion of simple samples.