Paper: Learning the greatest common divisor: explaining transformer predictions (ICLR-2024)

Reviewer comment (Novelty):
My main concern is about the significance of studying computing GCD through transformers. The observations are all limited to this particular task and it is not clear how broadly applicable they are. Or, if there is anything to learn from these observations and apply it elsewhere.

Author response:
*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.

Reviewer comment (Evaluation):
The observations are also not robust. The authors point out that the algorithm learned is deterministic but then this observation breaks down when the input distribution is changed. This non-robust behavior of the observations weakens them a lot in my opinion.

Author response:
*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).

Reviewer comment (Evaluation):
For example, explainability breaks down with the change in input distribution.

Author response:
*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.

Reviewer comment (Experiments):
Did the authors try to experiment with Chain-of-Thought and see if providing the explanation helps with the accuracy?

Author response:
*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.

Reviewer comment (Theory):
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?

Author response:
*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.

Reviewer comment (Experiments):
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?

Author response:
*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.

Reviewer comment (Novelty):
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.

Author response:
*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.

Reviewer comment (Theory):
The claim that Transformer predictions are fully explainable does not seem to be accurate once log-uniform operands are used.

Author response:
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.

Reviewer comment (Theory):
Although the paper usually interprets the results of experiments, it does not put forward any explanation for such results (for example, defining and justifying the shortcuts Transformers may take).

Author response:
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.

Reviewer comment (Evaluation):
Some of the report rules might be consequences of other factors (and they may not be robust as a result). For example, the smaller primes are learned (or grokked) faster as they are more common in the output or really easier to learn? Although these problems can be clarified. See the questions below for more details.

Author response:
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.

Reviewer comment (Theory):
Q2. At the beginning of the paper there is the idea that divisors of \(B\) are easily learned. Is this also the case when \(B\) is for example, the product of two large primes, e.g., \(B=2021=43\times 47\)? Is there a possibility that smaller primes are learned faster than \(43, 47, 43^2, \ldots\) in this case?

Author response:
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.

Reviewer comment (Evaluation):
Q3. Also given the list intuition, why is 420 performing better than 210 in Table 2?

Author response:
This is because 420 is divisible by 4, and not 210. Table 19 in appendix E2 indicates that for B=420, 8 and 16 (which can be tested on the last two digits) are correctly predicted, but not 32 (which requires three digits). For B=210, only 2 and 4 (and their multiples) are correctly predicted. 

As a result, in B=420, models will correctly predict GCD that are multiples of 8 and 16 (and other divisors of the base, e.g. 24, 48, 40, 80…), but B=210 cannot. This explains the different in performance.

Reviewer comment (Evaluation):
Q4. Further, the list intuition does not seem to be applicable anymore in Table 6 (e.g., \(2401=7^4\) is having the best performance). What could be the reason for this?

Author response:
The list intuition is specific to divisors of the base. For 2401, 7 and 49 are learned this way, but all other divisors are learned by grokking. Performance therefore depends on how fast the model groks the small prime, a quantity that varies a lot with model initialisation (figure 5 in appendix E1 shows learning curves for 3 different model initializations, for B=2023). 

In other words, while it is clear that large base learn faster, randomness in model initialization is probably the main reason 2401 achieves better results than other large bases.

Reviewer comment (Experiments):
Q5. In claim U2, one can see that 21 is wrongly classified into \(C_4\). As a result I wonder, for other bases that are product of larger primes, e.g., again, \(2021=43\times 47\), can this phenomenon still be observed? Or is having small primes such as 2 and 5 playing an important role here?

Author response:
The 21 is a typo. After checking, it is 24. Thank you for spotting it. We did not find exceptions to the classification rules U1, U2 and U3.

Reviewer comment (Evaluation):
Q6. In the line before Table 5, we see that for \(B=30\), \(B-1, B+1\) are learned faster than other primes. I guess this might be due to the fact that \((x_3x_2x_1)_B \pmod{B-1} = x_3 \times B^2 + x_2 \times B + x_1 \pmod{B-1} =  x_3 + x_2  + x_1 \pmod{B-1}\) which may be easier to learn due to symmetry between \(x_3, x_2, x_1\). Similar explanations can be given for any divisor of \(B-1\) and \(B+1\). If that's the case, this factor should also be considered when we see some primes are grokked.

Author response:
For B-1 and B+1, simple rules exist for all bases: the sum of the digits should be a small multiple of B-1 (because $B^k = 1 \mod B-1$), and the alternating difference $x_1-x_2+x_3-x_4+..$. should be a small multiple of B+1 (because $B^k = (-1)^k \mod B+1$). Similar rules exist for other divisors: e.g. the divisibility by 3 rule in base 10. However, all these rules involve all digits in the representation, and imply that the model must have learned to add or subtract. 
This result for B=30, also observed in one experiment for base 31, suggests that such divisibility rules can indeed be learned, but that the usual “learn divisors of B, then grok small primes” is easier.

Reviewer comment (Novelty):
Q8. The addition and multiplication of rationals mentioned in the intro and studied in the appendix require the output to be simplified. So I wonder, is there any particular difference or additional complexity comparing to the simplification task itself?

Author response:
Multiplication is more complex, in the sense that the two pairs of operands must be multiplied before the result is simplified. 

Addition is even more complex, because the operands must first be multiplied so that they have a common denominator, then numerators need to be added, and finally the result must be simplified.

This said, our results on fraction comparison, which involve computing ad-bc when comparing a/b and c/d, suggest that transformers can handle addition and multiplication.

Reviewer comment (Reproducibility):
Q9. What is the positional embedding used for the model?

Author response:
We use absolute encodings of positions, positional embeddings are learned during training.

Reviewer comment (Writing):
R1. My understanding is that the samples at each batch are freshly generated. I think it would be beneficial if this is communicated more clearly and earlier to the reader.

Author response:
We will clarify this

Reviewer comment (Presentation):
R2. Table 5 is a bit misleading with the current format. For example, when we see that 9 is learned for \(B=2017\) we have to guess that 3 was already learned. (Similar issue can be observed for other bases.)

Author response:
The new divisors learned are “new” compared to the natural outcome distribution. In the B=1000 case, 3 was learned with the natural distribution of outcomes. We will clarify this.

Reviewer comment (Evaluation):
The implications of these results for large language models are unclear. However, we barely know anything about the internal workings of Transformer models (the architecture behind LLMs) and this paper makes a small (and interesting!) step in enhancing our understanding.

Author response:
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).

Reviewer comment (Novelty):
My main concern about this work lies with the significance of the results and observations made. The authors train a transformer to predict the GCD which seems to work fairly well with some tricks in picking the right dataset. However, I’m not convinced about why this result would be of significant interest to the wider community and what it says about the representational power of transformers themselves beyond the narrow context of learning GCDs. While the experiments are thorough and the analysis is extensive, I’m struggling to understand the value of this beyond this specific task.

Author response:
*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.

Reviewer comment (Experiments):
The authors present ‘3 rules’ which supposedly explain model predictions and while they seem technically correct, they do seem pretty specific to the task and don’t seem to point to something specific about transformers themselves. For example, do we get these same rules if we swap the transformer architecture to something else? Is there something special about the architecture?

Author response:
*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.

Reviewer comment (Reproducibility):
How is the choice of the base made? Are they also sampled from a distribution?

Author response:
*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.