Learning the greatest common divisor: explaining transformer predictions

Thank you for your comments, questions and suggestions. We will integrate our replies in the revised version (in a few days).

My main concern is about the significance of studying computing GCD through transformers.

The broader context of our research can be summarized as follows :

• This is a first step towards using mathematical tasks to understand and explain deep learning models, by investigating their predictions on selected tasks. This extends current approaches on explainability which mostly focus on mechanical interpretability - i.e. looking at the model weights and internal calculations. In this paper, we demonstrate the feasibility of such an approach on a non-trivial mathematical task. By leveraging our theoretical knowledge of the underlying problem, we can design train and test sets, and get a clear picture of how the model learns to solve the problem. We believe that running similar analyses on a larger set of mathematical problems would provide a lot of insight about the inner workings of deep learning models (not just transformers).
• While tasks of pure mathematics are not the main focus of LLM practitioners, they play a central role for the development of foundational models for science: large language models pre-trained on mathematical equations, instead of natural language. Our results show how transformers can be trained to perform exact calculations involving integer divisibility, a central task in integer arithmetic and number theory.
• Our results on training distributions, i.e. the fact that some distributions allow for faster learning and better out-of-distribution performance, may apply to other arithmetic tasks. In particular, log-uniform operands and outcomes could be used when fine tuning LLM, or training foundational models for physics.

Moreover, the observations are also not robust. The authors point out that the algorithm learned is deterministic but then this observations breaks down when the input distribution is changed.

Our observations on determinism and explainability, and in particular the three rules for explainability G1, G2, G3 (section 4), hold for uniform and log-uniform operands, and natural (inverse square) and log-uniform outcomes, together with the additional distributions discussed in appendix D1. They are also valid for other deep learning architectures (see our reply to reviewer 21KR). We believe this makes a strong case for their robustness.

In fact, explainability only breaks (partially) when training on uniformly distributed GCD (section 6). The G rules, then get replaced by the weaker U rules, which break again (but temporarily) as the model learns the very last GCD (appendix D3, and figure 15). We believe this limiting case actually sheds light on what the model is learning. The non-uniform distribution of outcomes is an inductive bias about the problem that the model learns to exploit (see also our reply to your question below).

Can the authors please explain why does the model only learn divisors of the base B? And, if there is a way to promote learning of other divisors?

The divisors of B are learned because divisibility by these numbers can be checked by looking at the last digits of the integer representation in base B. For instance, in base 10, a number is divisible by 2 iff its representation ends with a 0, 2, 4, 6 or 8. It is divisible by 5 iff it ends in 5 or 0, and by 20 iff it ends in 00, 20, 40, 60 or 80. This means that the model can check whether both its operands are divisible by 20, by looking at their 2 last digits (the attention mechanism allows the transformer to point at specific digits in the sequences) and comparing them to a set of memorized values (the FFN layers in the transformer allow for such memorization). For other bases, such simple rules do not exist: all digits in the number representation must be taken into account when checking for divisibility. We believe this explains our results from section 3.

In section 4, we observe that other divisors: small primes and their multiples, can be learned (i.e. grokked) when the model is trained long enough. How this phenomenon happens is not clear, but it seems to happen for all large bases, and to be accelerated by training on log-uniform distributions of operands and outcomes (section 5). In our best results, all products of primes up to 53 are learned. Observations from appendix D1 suggest that these results can be improved. Our experiments in section 6 indicate that an imbalance distribution of outcomes is required for this phenomenon to take place.

The authors mention that some non-divisible small primes are learned very late in the training stage. Did the authors see that if they continue training further, eventually other primes will be learned?

Yes. Our initial experiments with grokking, for B=1000, shows that whereas all divisors of B were learned after 84 epochs, it took 220 epochs to learn 3 and its multiples. The rightmost column in Table 4 shows, for models trained up to 1300 epochs, the number of epochs needed to learn each small divisor. For instance, for B=2023, 3 is learned after 101 epochs, 2 after 205 and 4 after 599. For B=2744, 3 is learned after 543 epochs, and 5 after 1315. We believe that longer training would help (but 1300 epochs takes over 2 weeks on the machines we used). This is suggested by our experiments with log-uniform operands, which accelerate grokking, and allow to learn all primes up to 23 with a “natural” distribution of outcomes, and 53 with a log-uniform distribution of outcomes.

The authors mention that explainability of the model breaks down with the change in the input distribution. Do the authors have any thoughts on why this is the case?

Explainability does not break down when the input distribution, i.e. the distribution of operands, changes (in section 5). The three rules G1, G2 and G3 from section 4 still hold for models trained on log-uniform operands, and even log-uniform outcomes. The only difference, mentioned in the last paragraph of section 5 (bottom of page 6), is that the transitions between different phases – the moment when the model is learning a new divisor – now span several epochs, during which model predictions are split between the old and the new value. Out of these transient phases, model predictions remain deterministic and explainable.

Explainability (partly) breaks when models are trained on uniform outcomes. This is described in section 6. Models trained on uniform outcomes still learn to cluster their input pairs into classes which correspond to products of divisors of the base. All elements in the class are usually predicted the same, and new classes of multiples of new primes are learned as training proceeds (see the correct GCD in figure 3). Yet, the model prediction for each class varies from epoch to epoch (Table 8): multiples of 4 are predicted sometimes as 4, sometimes as 44, but at a given time in training, all multiples are predicted the same.

We believe this is an effect of the uniform distribution of outcomes: in previous experiments, the model predicted the most common instance of each class, which turned to be its smallest member. As a result, all multiples of 4 would be predicted as 4, and 1 would be the default prediction for all “non predictable” GCD. With uniform outcomes, the most common instance randomly varies over time, and this accounts for the fluctuations observed in table 8. In other words, the distribution of outcomes in the training set provides the model with an inductive bias about divisibility, which disappears when outcomes are uniformly sampled.

Did the authors try to experiment with Chain-of-Thought and see if providing the explanation helps with the accuracy?

Chain-of-thought does not apply to this setting, as it requires a decoder-only architecture, and we use encoder-decoder. Besides, it would demand that we know in advance the algorithm that the model is going to use. In the case of GCD, providing the model with steps from the “human algorithm” (i.e. Euclid’s algorithm) would certainly not help the model learn the sieving technique it is implementing.

The authors mention they observe grokking phenomenon. Can the authors please provide more related work on grokking and transformers. And, if this is the first work observing grokking for transformers.

The original paper on grokking (Power et al. 2022, arxiv 2201.02177 referenced in the paper), observed the phenomenon in a transformer learning modular arithmetic – i.e. arithmetic over a finite set. The follow-up papers by Liu et al. (https://arxiv.org/abs/2205.10343) and Gromov (https://arxiv.org/abs/2301.02679) mostly focus on fully connected networks (MLP). Nanda et al. (https://arxiv.org/abs/2301.05217) study grokking in one-layer transformers. All these works focus on arithmetic over finite sets: the model is tasked to learn, or complete, a finite “addition table”. To our knowledge, our work is the first observation of grokking on a different arithmetic task.

Thank you for your comments, questions and suggestions. We will integrate our replies in the revised version (in a few days).

The claim that Transformer predictions are fully explainable does not seem to be accurate once log-uniform operands are used.

The three rules for explainability (G1, G2, G3) are hold over a large set of training distributions : uniform and log-uniform operands, natural and log-uniform outcomes (together with the additional distributions from appendix D1). We will try to clarify the last sentence in the abstract.

In fact, explainability only breaks (partially) when the training data uniformly samples GCD (section 6). The G rules, then get replaced by the weaker U rules, which break again (but temporarily) as the model learns the very last GCD (appendix D3, and figure 15). We see these last cases are limiting conditions. They suggest that an unbalanced distribution of outcomes is needed to explainability to happen.

Although, the paper usually interprets the results of experiments, it does not put forward any explanation for such results

A number of explanations are provided, but they are scattered through the paper. We will concentrate them in one subsection of the discussion. Summarizing:

We can explain the shortcuts we believe our models use to test divisibility by products of divisors of the base. When the base is a prime number, the model counts the rightmost zeroes in both of its operands, takes the minimum $k$, and predicts $B^k$. When the base is composite (or a power of a prime), the model exploits the following property: if a factor $f$ divides $B^n$, ie $B^n = kf$, an integer $m$ is divisible by $f$ if and only if its $n$ last digits in base $B$ take one of the $k$ values : $0, f, 2f, … (k-1)f$. For instance, in base 10, a number is divisible by 20 iff it ends in 0, 20, 40, 60 or 80. We believe the model memorizes those lists when k and n are small. For small bases ($B<500$), models learn all divisors of $B^2$ this way.

We also understand how the model leverages divisibility rules to predict GCD. This exploits the imbalance in the GCD distribution (i.e. the fact that, in the training set, small GCD are more common than large GCD). As training begins, the model predicts 1 (the most common value) for all input pairs $(a,b)$. When divisibility by $k$ is learned (via the shortcuts described above, or grokking), the model learns to distinguish all pairs of the form $(ka, kb)$, i.e. such that their GCD is a multiple of $k$. The model therefore predicts $k$ for all pairs of the form $(ka,kb)$, and $1$ for the others. Once a second divisor, $l$, is learned, both input classes (i.e. $(ka,kb)$ input predicted as $k$, and others, predicted as $1$), are split into four classes, predicted as $1$, $k$, $l$, and $kl$ (except if $l=k^2$, in which case we end up with 3 classes, predicted as $1$, $k$ and $k^2$). Every new prime divisor $p$, learned by the model, causes each previous class, predicted as $f$, to split into two new classes, predicted as $f$ and $pf$. This accounts for the steps observed in the training curves.

Finally, we can explain how this process leads to a correct algorithm for predicting GCD, when outcomes are not uniformly distributed. Once the divisors of the base are learned, small primes are grokked in order (so far, we do not have an explanation for grokking), and new classes are created, that correspond to the products of the prime with previously learned factors. These classes are predicted as $pf$ ($p$ the prime, $f$ the previous factor), and contain all pairs $(pfa pfb)$, divisible by $pf$. Pairs with GCD $pf$ are now correctly predicted, and multiples of $pf$ will be learned once all their prime power divisors are learned or grokked.

The only part in the process that cannot be explained is the grokking process. We leave it to future work. We will add these explanations to the discussion section.

the smaller primes are learned (or grokked) faster as they are more common in the output or really easier to learn?

The experiments from section 6 shed light on this. Non-divisors are grokked in order, and each class is predicted as its smallest element (i.e. its largest common divisor) so long the distribution of outcomes in the training set in imbalanced. This suggests that small primes are learned first, because small GCD are more common in the training sample. The distribution of outcomes provides the model with the inductive bias about what a “small number” is.

Q2. For B=2021=43×47, is there a possibility that smaller primes are learned faster than 43,47,432,…?

For B=2021, the two divisors of the base, 43 and 47, are learned long before small primes are grokked. Over three different initializations of the model, 43 and 47 always get learned around epoch 20 and 22. 2 is learned at epochs 261, 442 and 125, and 3 at epochs 450, 182 and 226. Even with log-uniform outcomes, grokking always sets in after the divisors of the base are learned.

Thanks you for your comments, and suggestions. We will integrate them in the revised version (updated in a few days).

Perhaps one could say that the implications of these results for large language models are unclear.

We agree that immediate applications to LLM are not obvious. We believe that the log-uniform sampling of operands and outcomes could be of help when fine tuning LLM on mathematical tasks (in fact, this might already be the case, since numbers scraped from the Internet tend to follow a Zipf/Benford law).

A more likely application of this research is for Foundational models for Science, large language models pre-trained on a large corpus of scientific facts (equations, calculations, scientific data), to be used in Mathematics, Physics or Chemistry. These models need to perform correct integer calculations, and our results show how they can be learned.

Finally, our main research direction is model explainability. Most current approaches focus on mechanical interpretability (i.e. understanding model weights). We propose a different direction, which focuses on elucidating model predictions for specific tasks. We believe mathematical problems are a good playground for such studies, because we can leverage our theoretical understanding of the underlying problems. Extending our approach to a larger set of mathematical problems might provide interesting insight about the algorithms learned by transformers (and other deep learning architectures).

Thank you for your comments, questions and suggestions. We will integrate our replies in the revised version (in a few days).

My main concern about this work lies with the significance of the results and observations made.

The broader context of our research can be summarized as follows:

• This is a first step towards using mathematical tasks to understand deep learning models, by investigating their predictions on selected tasks. We see this as a new approach to explainability, and demonstrate its feasibility on a non-trivial mathematical task.
• While mathematical tasks are not the main focus of LLM practitioners, they play a central role for the development of Foundational Models for Science: large language models pre-trained on mathematical equations. Our results show how transformers can be trained to perform exact calculations involving integer divisibility, a central task in integer arithmetic and number theory.
• Our results on training distributions, i.e. the fact that some distributions allow for faster learning and better out-of-distribution performance, may apply to other arithmetic tasks. In particular, log-uniform operands and outcomes could be used when fine tuning LLM, or training foundational models for physics.

Are the 3 rules task specific? Do they apply to other architectures?

We agree that the 3 rules are specific to tasks involving integer divisibility. Still, they provide insight about what our models can learn. In section 3, the model only learns divisors of the base, that can be tested by looking at a small number of digits. Divisibility by 3 or 9 (in base 10), which can be tested by summing all digits in the representation, seem harder to learn. For a transformer, we believe this suggests that the attention mechanism helps learn properties that concentrate on a few tokens in the input sequence, but struggle with non-local properties, such as the sum of all digits in an operand.

Thank you for suggesting additional experiments on architectures! We experimented with LSTM and GRU, with 1024 and 2048 dimensions, and 4 layers, for 10 different bases (from 10 to 2401). We will include our findings in the appendix of the revised paper, but initial results suggest that our observations extend to these recurrent networks.

When training on uniform operands for 260 epochs, we observe that:

• for B=420, 34 GCD below 100 are learned, products of {1,2,4,8,16},{1,3,9}, {1,5} and {1,7} (divisors of the base)
• for B=210, 31 GCD are learned: products of {1,2,4,8}, {1,3,9}, {1,5}, {1,7}
• for larger bases, grokking seems to happen as well: for B=2021, 2 and 3 are grokked, for B=2023 and 2025, 2 is grokked, for B=2401, 2 and 4 are grokked. When training on log-uniform operands, 37 GCD are learned after 150 epochs, with uniform outcomes 57 GCD under 100 are learned.

Overall, this suggests that our results extend beyond the transformer architecture.

Whether the model counts rightmost zeroes should be fairly easy to confirm by visualizing the attention weights. I wonder if the authors have tried this.

We did, and observed that some attention heads do focus on the low digits of the operands. We hesitate to jump to conclusions because:

• input sequences are short, and different heads attend to almost all tokens in the sequence,
• this pattern was observed in tasks that do not involve divisibility,
• the same phenomenon is observed in architectures without attention (see above).

Regardless, I’m not sure if I would call a model showing some consistent trends in the output as ‘explainable’.

We believe ours is a valid alternative to mechanical interpretability. Even though the low level specifics of the model are not fully understood, the 3 rules account for more than 99.5% of model predictions. We believe this qualifies as a valid explanation. Such a situation is common in physical sciences. Many macroscopic phenomena, the relation between volume, pressure and temperature in a gas, or the transformation of an egg into a fully developed individual, can be explained, even though the low level processes, the trajectories of particles in the gas or the biochemical processes of embryogenesis, are not completely elucidated. We see our approach as a new direction in explainability.

How is the choice of the base made? Are they also sampled from a distribution?

In section 3, we selected small bases (2,3,4,5,6,7), bases often used in practice (10, 100, 1000), bases with many small prime divisors (30, 60, 210, 420) and primes bases close to these values (31, 211, 997).

For section 4, we selected bases depending on their prime decomposition. In particular, we focused on bases not divisible by some small primes, in order to observe grokking. The first column of table 4 provides an explanation of sorts.

In section 5 and 6, we reused the bases from the previous sections, for the sake of comparability.