Unraveling Arithmetic in Large Language Models: The Role of Algebraic Structures

TL;DR

This work examines how large language models acquire arithmetic skills by internalizing algebraic structures rather than memorizing numerical patterns. It frames arithmetic in finite Abelian groups, particularly $ ext{$ ext{Z}_n$}$, and designs datasets to probe commutativity and identity through input-output relationships, while introducing operators that suppress trivial numerical leakage. The paper provides theoretical proofs that transformer attention can enforce invariance under input permutation and identity insertion, and it corroborates these ideas with extensive experiments showing high generalization to unseen inputs. The findings offer a principled view on improving LLM arithmetic by leveraging algebraic structure, with implications for designing more robust reasoning and generalization capabilities.

Abstract

Large language models (LLMs) have demonstrated remarkable mathematical capabilities, largely driven by chain-of-thought (CoT) prompting, which decomposes complex reasoning into step-by-step solutions. This approach has enabled significant advancements, as evidenced by performance on benchmarks like GSM8K and MATH. However, the mechanisms underlying LLMs' ability to perform arithmetic in a single step of CoT remain poorly understood. Existing studies debate whether LLMs encode numerical values or rely on symbolic reasoning, while others explore attention and multi-layered processing in arithmetic tasks. In this work, we propose that LLMs learn arithmetic by capturing algebraic structures, such as commutativity and identity properties. Since these structures are observable through input-output relationships, they can generalize to unseen data. We empirically demonstrate that LLMs can learn algebraic structures using a custom dataset of arithmetic problems, as well as providing theoretical evidence showing that, under specific configurations of weights and biases, the transformer-based LLMs can generate embeddings that remain invariant to both permutations of input tokens and the presence of identity elements. Our findings indicate that leveraging algebraic structures can enhance the LLMs' arithmetic capabilities, offering insights into improving their arithmetic performance.

Figures (4)

• Figure 1: Illustration of dataset for operator "$+$", $\oplus$ and $\ominus$. Notice that the same set of tokens is maintained across all operators to ensure that certain token combinations appear exclusively either in the training set or the testing set, as required.
• Figure 2: Illustration of the symbols defined for the hidden states of tokens and the variables for the attnetion layers
• Figure 3: Plots of training and testing accuracy. The first row is the training dynamics for $\mathbb{Z}_7$ given the scale of training set $K=3000$. The second row are the accuracies for $\mathbb{Z}_{7}$ (left), $\mathbb{Z}_{11}$ (middle), $\mathbb{Z}_{13}$ (right) with varying $K$ of training set.
• Figure 4: Visualization of $S^{\ell}_{\text{com}}$ and $S^{\ell}_{\text{ide}}$ where $1\leq \ell \leq 13$. The upper row displays the values of $\operatorname{S}^{\ell}_{\text{com}}(+,\ominus)$, $\operatorname{S}^{\ell}_{\text{com}}(+,\triangleleft)$, and $\operatorname{S}^{\ell}_{\text{com}}(+,\triangleright)$ and the lower row displays the values of $\operatorname{S}^{\ell}_{\text{ide}}(+,\ominus)$, $\operatorname{S}^{\ell}_{\text{ide}}(+,\triangleleft)$ and $\operatorname{S}^{\ell}_{\text{ide}}(+,\triangleright)$. The numbers in the left axis represent $K \in \{100,300,\cdots,10000\}$. For clarity, non-negative values are highlighted in green and yellow.

Theorems & Definitions (5)

• Theorem 2.1: Commutativity--Invariant to the Input Permutations
• Theorem 2.2: Identity--Invariant to the Insertion of Identity Tokens
• Remark 2.3: Non-uniqueness of Weights and Bias Assignments
• Remark 2.4: Trivial Solution of Embeddings
• Definition A.1