Learning the greatest common divisor: explaining transformer predictions
Authors: Francois Charton
Venue: ICLR 2024 spotlight (2024)
Keywords: mathematics
arithmetic
transformers
explainability
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Transformers can learn to predict greatest common divisors, their predictions can be fully explained, training distribution matters

Abstract: The predictions of small transformers, trained to calculate the greatest common divisor (GCD) of two positive integers, can be fully characterized by looking at model inputs and outputs.
As training proceeds, the model learns a list $\mathcal D$ of integers, products of divisors of the base used to represent integers and small primes, and predicts the largest element of $\mathcal D$ that divides both inputs. 
Training distributions impact performance. Models trained from uniform operands only learn a handful of GCD (up to $38$ GCD $\leq100$). Log-uniform operands boost performance to $73$ GCD $\leq 100$, and a log-uniform distribution of outcomes (i.e. GCD) to $91$. However, training from uniform (balanced) GCD breaks explainability.

Published as a conference paper at ICLR 2024
LEARNING THE GREATEST COMMON DIVISOR: EXPLAINING TRANSFORMER PREDICTIONS
François Charton
Meta Al
fcha mntaartommettccm
ABSTRACT
The predictions of small transformers, trained to calculate the greatest common divisor (GCD) of two positive integers, can be fully characterized by looking model inputs and outputs. As training proceeds, the model learns a list D of integers, products of divisors of the base used to represent integers and small primes, and predicts the largest element of 'D that divides both inputs. Training distributions impact performance. Models trained from uniform operands only learn handful of GCD (up to 38 GCD < 100). Log-uniform operands boost performance to GCD S 100, and log-uniform distribution of outcomes (i.e. GCD) to 91. However, training from uniform (balanced) GCD breaks explainability.
INTRODUCTION
Transformers Vaswaai et al. 2017) have been applied to problems of mathematics, both symbolic
(Lample & Charton 2019 Charton al. 2020 Shi al. (021) and numerical (Charton) (021).
Yet, they struggle with basic arithmetic (Lee et al. 2023 Nogueira et al. (02T). Large language models (LLM) can learn addition or multiplication by small prefactor, and generalize beyond
their training range when fine-tuned using scratchpad (Nye et al. 2021), chain-of-thouglle (Wei
al. 2023) or algorithmic prompting (Zhou al., 2022), but these techniques require bespoke data and do not extend to complex tasks (Dziri et al. 2023). Math transformers were also found to be brittle (Welleck et al. 202T), to fail on simple tasks (Davis) 2023). and to be hard to interpret, except in the simplest cases (Nanda al. (2023). Yet, small transformers can learn advanced calculations, such as eigen-decompoosttin (Charton 202T) and polynomial roots (Charton 2022b).
In this paper, train 4-layer transformers to compute the greatest common divisor (GCD) of two positive integers, an important operation for rational arithmetic and number theory, and observe that:
1. Transformers learn to cluster input pairs with the same GCD. All pairs of integers (a, b) with the same GCD are predicted the same.
2. Transformer predictions
be fully characterized. During training, the model learss set of integers D, and predicts, for any input pair (a, b). the largest element in D that divides a and
3. Early during training. transformers learn to predict products of divisors of the base used to represent integers. Small primes are "grokked" (Power et al.) 2022) after extended training.
4. Models trained from log-uniform operands and outcomes achieve better performance. They correctly predict up to 91 GCD < 100. Model predictions remain fully explainable.
5. An unbalanced distribution of outcomes in the training set is required for full explainability: explainability partially fails once models are trained from uniformly distributed GCD.
These results demonstrate how transformers can be trained to perform exact calculations involving integer divisibility, a central task in integer arithmetic and number theory. Beyond GCD calculations, the broader potential impact this research extends in three directions. First, it presents new approach to model explainability: fully characterizing black-box model predictions by experimenting with selected inputs and leveraging our theoretical understanding of the underlying mathematics. Sec- ond, the results on log-uniform training distributions of operands and outcomes faster learning and better performance- may extend to other arithmetic tasks, e.g. fine tuning LLM. Finally, mathemati- cal tasks play a central role for Foundational Models for Science large language models pre-traee pre-trained on mathematics, and fine-tune tuned on specific fields, such as high energy physics, computational biology astrophysics. Before they can do science, transformers must learn maths.

---

Published as a conference paper at ICLR 2024
RELATED WORK
Neural networks for arithmetic were first proposed by Siu & Roychowdhury (1992), and recurrent models by Kalchbrenner et al. (2015). Zaremba et al. (2015) and Kaiser & Sutskever (201)). Recent research mostly focuses on fine-ture toneing LLM on arithmetic tasks, to solve math word problems (Meng & Rumshisky 2019 Griffith & Kalita 2021). See Lee et al. (2023) for a summary. As an alternative, Neural Arithmetic Logical Units (Trask et al. Mistry 2023) learn exact computations that can generalize to any input, by constraining the weights of linear models to be close to 0, 1 or 11.
The difficulty of learning arithmetic tasks was discussed by many authors. Saxxon al. 2019)) benchmarking mathematical tasks, observe that number theoretic operations, like factorization, are hard. Palamas (2017) further investigates the hardness of modular arithmetic. [zzir et al. (2023) note the difficulty of extending the promising results obtained by Lee ct al. (2023) on the four operations to complex mathematical calculations or algorithms GCD and Fuclid's algorithm, here.
The role of number representation was discussed by Nogueira et al. (202)) and Charton 202
Grokking was first described by Power et al. (2022). Liu al. (2022) propose metrics to characterize it. Gromov (2023) provides an insightful analysis of grokking in feed-forward networks. Most prior work on explainability in arithmetic transformers tries to interpret model weights (Nanda et al. 2023
Zhong et al. 2023). Charton (022a) conducts similar experiments for linear algebra.
2
EXPERIMENTAL SETTINGS
GCD calculations are framed as supervised translation task. Problems (pairs of integers) are randomly sampled, represented as sequences of tokens, and used to train sequence-to-sequeee transformers translate input pairs into their GCD, by minimizing the cross-entropy between model predictions and the sequences representing correct solutions. Integers are encoded sequences of digits in base B, preceded by a sign which also serves as a separator (Table D. In base 10, the model translates (8, 12), encoded as the sequence '
1 2', into its GCD, 4, encoded as '+ 4'. The choice B is a trade-off. Small bases result in longer sequences that are harder to learn, but use small vocabulary that is easier memorize. Composite bases allow for simple tests of divisibility: in base 10, divisibility by 2, 5 and 10 is decided by looking at the rightmost token in the sequence.
Transformers with 4 layers, 512 dimensions and 8 attention heads, using Adam (Kingma & Ba 2014) are trained with a learning rate of 10-5 (no scheduling is needed) on batches of 256 examples. All inputs pairs are sampled uniformly between 1 and M 10 All data is generated on the fly: different training epochs use different examples for the train and test set. After each epoch (300,000 examples), the models are evaluated on two test sets of 100,000 examples: a natural test set of uniformly sampled pairs (a, b), and stratified test set with GCD uniformly distributed between 1 and 100. In the natural set, small GCD are more common we have P(ged(a, b) = k) मेg (Cesàro (883). The stratified set has about 1000 examples with GCD & for 1 I h < 100, and is generated by:
sampling kt uniformly between 1 and 100,
sampling and 6, uniformly between 1 and M . such that gcd(a, b) 1, using rejection sampling, adding (ka, kb) to the stratified test set.
These two test sets provide two measures of accuracy. Model accuracy, measured on the natural set, is the probability that the GCD of two random integers from to M is correctly predicted. Accuracy on the stratified test set is the number of GCD correctly predicted between and 100. The size of the problem space (10'2 possible input pairs) guarantees minimal duplication between train and test on one NVIDIA V100 GPU with 32 GB of memory. The source code com/ www.fookk.eeerrcc....oakkk
set. All experiments for these experiments can be found at
Table 1: Encoding gcd(160, 40, in base 2, 6, and 30
Base
Encoded input
Encoded output
2
1+, 1, 0, 1,0, 0, 1, 1, 0,0,0,01 [+, 1,,10, 1,,00, [,,1,
6
[t, 4, 4,
10
[+, 1,,1,6 2,2,,
[,,4,0] I+, 1,,1,101
30
[+, 5,,,,0,,,,,000
2

---

Published as a conference paper at ICLR 2024
Table 2: Numher of correct GCD under 100 and aceuracy. Best of 6 experiments.
Base
2
3
4
5
6
7
10
11
12
15
Correct GCD
7
5
7
3
19
3
13
19
Accuracy
81,6
68.9
81,4
64.0
91.5
62.5
84.7
61.8
91.5
71.7
Base
30
31
60
100
210
211
420
997
1000
1024
Correct GCD
27
2
28
13
32
1
38
14
Accuracy
94.7
61.3
95.0
84.7
95.5
61.3
96.8
61.3
84.7
81.5
3 LEARNING THE GREATEST COMMON DIVISOR BASE EXPERIMENTS
A model trained on pairs of positive integers under one million, encoded in base B = 10, correctly predicts 84.7% of the examples the natural test set, and 13 correct GCD under 100 (aceuracy on the stratified test set). Performances vary with the encoding base: from 61.8% accuracy and 2 correct GCD for base 11, to 96.8% and 38 GCD for base 420 (Table 2 The best performances are achieved for composite bases (30, 60, 210 and 420), the worst for large primes. Learning is very fast: for base 30, the model achieves 90% accuracy after 2 epochs (600,000 examples), and 93% after 6. Model size has little impact on performance (Appendix [ For base 30, 1-layer transformers with 32 dimensions (less than 300,000 parameters) achieve 93.3% accuracy. 24-layer models with 1024 dimensions (714 million parameters) achieve 93.4%. For base 31, accuracy is 61% for all models.
These variations model performance can understood by looking at model predictions. Table presents, for bases 2 and 10 and GCD up 36, the most frequent model prediction for pairs with given GCD (Pred), and its frequency in the stratified test set (%) - detailed results for 6 bases and GCD up to 100 are Appendix E.3). All frequencies are very close 100%: for every test pair with GCD k, the model makes the same prediction f(k)). In other words, the model can tell whether two input pairs have the same GCD. Correct model predictions (k) k) only happen for products of divisors of the base. In fact, all model predictions can be summarized in three rules:
(R1) Predictions are deterministic. The model predicts unique value ((k) almost all (99.9%) pairs of integers with GCD k. Predictions are correct when f(k) = k.
(R2) Correct predictions are products of primes dividing B. For base 2, they are 1, 2, 4, 8, 16, 32 and 64. For base 31, 1 and 31. For base 10, all products of elements from 2,, 4, 8, 16] and {1, 5, 25). For base 30, all products of [1, 4, 8), {1. 1. 9, 27] and {1,5, 25).
(R3) f(k) is the largest correct prediction that divides k. For instance, f(8) for base and 10, but f(55) = 5 for base 10 and f(15) = 1 for base 2.
and f(7)
These results be interpreted follows. For prime bases, such B an integer is divisible by B iff its representation ends in & zeroes. The model learns to "predict" GCD by counting the rightmost zeroes in its operands, 2g and 2g, and predicting B with min 26). This accounts for all observed results. For instance, it will correctly predict the GCD of = 8 1000g and b 12- 1100, be 2 4, and incorrectly predict the GCD of 7 111g and 1110g to be
Table 3: Model predictions and their frequencies, for GCD to 36. Correct predictions in bold face.
Base
Base 10
Base 2
Base
Base
Base
GCD
Pred
Pred
GCD
Pred
%
Pred
GCD
Pred
%
Pred
100
100
13
1
100
100
25
100
25
100
2
100
2
100
14
2
100
2
100
26
100
2
100
3
100
100
15
1
100
5
100
27
100
1
100
4
100
4
100
16
16
100
16
99.7
28
100
4
100
5
100
5
100
17
1
100
100
29
100
1
100
6
100
2
100
18
2
100
2
100
30
100
10
100
7
100
1
100
19
1
100
100
31
100
1
100
8
100
8
100
20
4
100
20
100
32
32
99,9
16
99.9
9
100
100
21
100
100
33
100
1
100
10
100
10
100
22
2
100
2
100
34
100
2
100
11
100
1
100
23
100
100
35
100
5
100
12
100
100
24
8
100
8
100
36
100
4
100
3

---

Published as a conference paper at ICLR 2024
25
30-
25
20
15
10-
Tone
120
210
100
150
600
630
730
800
50
100
150
200
250
300
Epochs
Figure 1: Correct GCD training time. Natural (Fz) distribution of GCD.
Figure 2: Correct GCD vs training
time. 5% uniform, 95% natural GCD.
1. For composite bases, such as B 10, integer is divisible by J, such that k. B", iff its rightmost digits are in {0, (.. -- The model learns to test the divisibility of its operands by comparing their n rightmost digits with the possible values, and predict the largest f that divides both operands, In practice, only divisibilities that can be tested by considering the two last digits in the representation are learned. For B 210 divisibility by learned, but divisibility by is not. For B = 420 divisibility by 16 is learned, but not by 32. The three rules also account for variations in model accuracy (computed on the natural test set) for different bases (see Appendix C
Learning GCD one prime power at a time. Learning curves have a step-like shape (Figure D. and GCD are learned in sudden batches. When the model learns a new power of a prime divisor of B, it also learns its products with already known GCD. For instance, for base 30, the model initially predicts {1, 2, 4). (1,3, 9}. {1, 5) and their products: 17 GCD under 100. A first step happens around epoch 50, when the model learns 25 and the three associated multiples 50, 75 and 100 (21 GCD), a second around epoch 220, learning 8, 24, 40 and 72, and a third at epoch 660, learning 27 and 54, for a grand total of 27 correct GCD. The three rules hold at all times during training.
Accelerating learning by balancing the distribution of GCD. The distribution of GCD verifies (age b) k) マ| (Cesàro 1883). As result, large GCD very rare in the training set, and learning them is very slow. This can be mitigated, and training accelerated, by adding a small proportion (5%) of uniformly sampled GCD the training set: for 30, the model learns 25 GCD in 30 epochs, and 27 GCD in 175, vs 250 and 660 in the original experiments (Figure 2).
6
In these experiments, models only correctly calculate GCD that are products of divisors of the base, and the best accuracies are achieved for bases divisible by many small primes, e.g. 30, 210 or 420. Still, all models learn to cluster pairs of input integers according their GCD, and output a unique prediction ff fkk) for all pairs with GCD k. This is a non-trivial result and a significant achievement.
4 LARGE COMPOSITE BASES B GROKKING SMALL PRIMES
For large bases B, non-divisors of B are sometimes learned after extended training. In one experiment with base 1000, the model predicts 13 GCD < 100 after 84 epochs: all products of {1, 2, 4, 8, 16) and 51, 25). Then, the training loss is flat during 100 epochs, and it seems that the model is no longer learning anything. But then, the model starts predicting GCD 3. with an accuracy of 0.2% at epoch 188, and 93% at epoch 193 (despite only seeing 100,000 input pairs with GCD 3 during these epochs). Multiples of are then learned, and by epoch 220, the model predicts 22 GCD: all products of {1, 2, 4, 8, 16). (1, 5, 25) and (1,3}. Model predictions still respect rules RI and R3 (Appendix E.I Table 20). and the three rules can be updated as follows:
(GI) Prediction is deterministic. All pairs with the same GCD are predicted the same, f(k). (G2) Correct predictions are products of primes divisors of B and small primes.
(G3) f(k) is the largest correct prediction that divides k.
This phenomenon is related to grokking (Power et al., 2022). Table 4 presents results for 16 large bases, with models trained up to 1300 epochs. Grokking usually sets in late during training: for bases 625 and 4000, all products divisors of B are learned and 15 epochs, but take 600 epochs for grokking (of 2 and 3) to happen. Primes and powers of primes are roughly grokked in order.
4

---

Published as a conference paper at ICLR 2024
Table 4: Predicted ged, divisors and non-divisors B. Best model of 3. For non-divisors, the epoch learned is the first epoch where model achieves 90% accuracy for this GCD.
Base
GCD predicted
Divisors predicted
Non-divisors (epoch Iearned)
625 5*
[1,5,25]
2 (634)
2017
[1]
(142), 3 (392)
2021 43.47
10
(1,43), (1,47)
2 (125), 3 (228)
2023 7.17
16
(1,7), (1,1))
3 (101). 2 (205), 4 (599)
2025 = 3 .52
(1,3,9, 81), (1,5,25)
2 (217), 4 (493), 8 (832)
2187 3
20
(1,3,9,77811
2 (86), (315) (650)
2197 = 13a
[1,13]
(62). (170), (799)
2209 = 47
{1,47)
2 (111), (260), 9 (937)
2401 7
10
(1,7,49)
2 (39), (346)
2401 = 7
14
[1,7,49)
3 (117), 2 (399), 4 (642)
2744 2 .7
30
(1.2.4.8.6.322), (1.7,49)
3 (543), 5 (1315)
3125 = 5
16
[1,5,25]
2 (46), (130), 4 (556)
3375 39.5
23
(1,3,9,77, (1,5,25)
2 (236), (319)
4000 = 2
(1,2,4,8,16,32), (1,5,25
3 (599)
4913 = 17a
17
[1.17]
2 (54), (138), (648), (873)
5000 29,5
(1.2.4.8.16,32), (1.5,25)
3 (205), 9 (886)
10000 2
22
[1,2448.16]. (15,25)
3 (211)
1. [..1 Figure retain their usual step-like shape: long periods of stagnation followed by sudden drops in the loss, and rises in accuracy, as new GCD are learned. Because it helps learn small GCD, grokking boosts model accuracy (from 63% to 91% for B 2023), but overall the number of correct GCD remains low (under 30 for all large bases).
Learning curves (Appendix
Balancing outcomes. The technique proposed in section to accelerate learning (adding small amount of uniformly distributed GCD to the training set) does apply larger bases AAppendix[E33 Table 21). However, the unbalanced distribution of GCD) can be corrected by sampling from a log- uniform distribution so that (ged b) k) instead of C as follows:
k²
Sample & between and 100, with probability P(k) C with C = Σm
x.
Sample and b uniformly from 1 to M such that gad b) 1.
Add (ak: bl)) to the training set.
A log-uniform training distribution of GCD helps the model learn new non-divisors of B for 9 bases out of 35 (Table For B = 211, primes up to are learned. For = 1 10000, 7, 9, 13 and 27 are learned, bringing the number of correct GCD to 62, best result far. For B 30, counter-intuitive situation prevails: instead of small primes, the model learns B and B +
Table 5: Log-uniform natural outcomes. Best model of trained for 700 epochs. Non-divisors in bold.
Natural
Log Log-uniform outcomes
Natural # GCD
Log-uniform outcomes
Base
# GCD
GCD
New divisors learned
Base
# GCD
New divisors learned
7
997
5
1000
22
31
9,32, 64 9
7
2017
4
2021
19
20
64
2023
11
7
3
2025
28
28
10
13
14
32
2187
20
20
11
2
2197
20
81
2209
15
10
81
2401
16
5
30
25
36
16, 29, 31
2744
29
21
31
2
3125
16
16
60
27, 32,
3375
100
13
15
32, 64
4000
25
9, 64
210
32
32
4913
17
211
18
2,3,4,5,7
5000
28
30
64
420
47
13, 49
10000
22
7, 9,9.32
625
10000
7, 13, 27. 32,
5

---

Published as a conference paper at ICLR 2024
5 LEARNING FROM LOG-UNIFORM OPERANDS
In all experiments far, all pairs in the training result, models are mostly trained from examples with large operands. 90% of operands are larger than 100,000, and small instances, like ge 9). are almost never encountered. This contrast with the way we are taught, and teach, arithmetic. We usually insist that small examples should be mastered, and sometimes memorized, before larger instances, like ged (100370,, 102372) can be tackled.
uniformly sampled between and 10 As
In this section, I sample training pairs from log-uniform distribution, by uniformly sampling real numbers I z I log M, computing ex and rounding to the nearest integer. this setting, the training set has many 1-digit as 6-digit operands. In 3% of training example, both operands are smaller than 10, and in 11% of examples, both are smaller than 100. This presents the model with many simple examples that it can memorize, just like children rote learn multiplication and addition tables. This is different from curriculum learning: the distribution of operands does not change during training. Also, the log-um oog-uniform sampling only applies to the training set (the test sets are unaffected), and it has no impact on the distribution of outcomes.
Table 6: Accuracy and correct GCD (up to 100), log-uniform operands. Best of three models, trained for 1000 epochs (300M examples). All models are tested on 100,000 pairs, uniformly distributed between and 10
Base
Accuracy
Correct GCD
Base
Accuracy
GCD
Base
Accuracy
GCD
94.4
25
98.4
2025
99.0
3
96.5
36
100
98.4
2187
98.7
4
98.4
58
210
98.5
60
2197
98.8
68
5
97.0
42
211
96.9
41
2209
98.6
65
6
96.9
39
420
98.1
2401
99.1
73
96.8
625
98.2
2744
98.9
10
97.6
48
997
98.3
64
3125
98.6
65
11
97.4
43
1000
99.1
71
3375
98.8
67
98.2
55
1024
99.0
71
4000
98.7
66
97.8
2017
98.6
4913
98.2
98.2
2021
98.6
5000
98.6
31
97.2
44
2023
98.7
65
10000
98.0
56
Training from log-uniform operands greatly improves performance (Table Accuracy for all bases between 94 and 99%, compared to 61 and 97% with uniform operands. For base 2401, the number of correct GCD is 73, best result far. For base 10, the number of correct GCD is 48 (vs 13). Learning is accelerated: for base 10, GCD 1, 2, 4 and 5 are learned as early as epoch 3, 3 and 8 by epoch 25, 7 and 9 by epoch 220 and 11 by epoch 750.
As before, large bases perform better. All models with B S 420 have an accuracy over 98% and correctly predict more than 55 GCD under 100. The divisors B are learned first, then, small powers of primes are grokked, roughly in order. After training, models have learned to predict all primes up certain value, some of their small powers, and all associated products. All primes up learned for base 2. up to 11 for base 10. up to 17 for base 100, and up to 23 for base 1024. For base 1024, 2401, and 2744, only 27 GCD are incorrectly predicted:
the 16 primes from 29 and 97, all predicted as 1, small multiples of these primes: products of 2 and 29, 31, 37, 43 and 47, predicted as and products of 3 and 29 and 31, predicted as 3.
powers small primes: 49 72, predicted as 7, and
34, predicted as 27.
small multiples of these: 98 49 2, predicted as 14.
The three rules with grokking (G1 to G3) still apply: predictions are deterministic, for a pair (a, b) with GCD k the model predicts the largest correctly predicted GCD that divides k.
Learning curves retain their step-like shape, but they are more noisy, and smoother (see Appendix E.2): transitions now span several epochs, and each new prime takes more examples to be fully learned. While the model During training, while the model leams a new divisor, rules GI and G3 are temporarily violated. During few epochs, model predictions value (c.g. between and 49 when the model is learning 49). This situation, rarely observed in previous experiments, is common with log-uniform operands.
split between the old and the new

---

Published as a conference paper at ICLR 2024
Table 7: Accuracy and correct GCD, log-uniform operands and outcomes. Best model of 3.
Base
Accuracy
Correct GCD
Base
Accuracy
GCD
Base
Accuracy
GCD
16.5
17
60
96.4
75
2025
97.9
91
93.7
51
100
97.1
78
2187
97.8
91
4
91.3
47
210
96.2
80
2197
97.6
90
92.2
58
211
95.3
67
2209
97.6
87
95.2
56
420
96.4
88
2401
97.8
89
93.0
63
625
96.0
80
2744
97.6
91
10
94.3
65
997
97.6
83
3125
97.7
91
11
94.5
57
1000
97.9
91
3375
97.6
91
12
95.0
70
1024
98.1
90
4000
97.3
90
15
95.4
62
2017
97.6
88
4913
97.1
88
30
95.8
72
2021
98.1
89
5000
97.1
89
31
94.4
64
2023
97.5
88
10000
95.2
88
Log-uniform outcomes. Balancing the distribution of GCD by making it log-uniform, as described section together with log-uniform operands, brings another large improvement in performance (Table 7) After 1000 epochs, all models with B larger than 1000 predict 87 to 91 GCD: all primes up to 53 and all composite numbers up to 100. These are our best results. They can be marginally improved by training models from inverse square root distribution of outcomes (Appendix D.I). Note the low accuracy for base 2: with log- -niuforrrm outcomes, the model fails to learn GCD 1, for lack of examples.
LEARNING FROM UNIFORM OUTCOMES
Log-uniform distributions outcomes improve model performance by reducing the imbalance between small and large GCD in the training set. It is therefore tempting push this logic further, and train models from uniform distribution of GCD and operands, i.e. sample the training set like the stratified test set from Section 2 Figure preueets presents learning curves for three models using base 10. Model accuracy (measured on the natural test set) seems to vary randomly, and the test loss is flat. Yet, the number correct GCD stable over time, and increases in steps, from 10 17, in line with the results from section 3 (13 GCD are learned). Something is learned despite the flat loss.
Accuracy
Correct GCD
Loss
20.0
4,0
80
17.5
3.5
70
15.0
3.0
50
22.5
2.5
40
10.0
2,0
30
7.5
15:
5.0
1.0
10
2.5
0.5
0.0
0.0
200
600
200
400
400
600
Epochs
Fpochs
Figure 3: Learning curves for B=10. Uniform outcomes and operands. 3 different seeds.
Table presents the most common model predictions, and their frequencies, for all GCD up to 20. At first glance, predictions seem chaotic. epoch 266, the model achieves 81% accuracy, and correctly predicts 14 GCD: 1, 2, 5, 8, 20, 32, 40, 44, 48, 50, 64, 75, 80 and 100. One epoch later, accuracy is down 6%, the model still predicts half of the correct GCD have changed! After another epoch, accuracy is 4% and the model predicts 4, 20, 25, 26, 30, 32, 40, 48, 50, 55, 64, 73, 80, 88 and 100. Again, half the correct GCD have changed.
GCD: 8, 10, 16, 40, 50, 55, 60, 64, 66, 75, 80, and 100,
As in previous experiments, frequencies are close to 100%: the model makes unique prediction f(k) for all pairs with GCD k. with the notable exception of epoch 267 where model predictions for 1, 3.... ... are split (almost evenly) between 11 and 19. Model predictions cluster by classes of GCD: all elements in class C₁ {1, 3, 7, 9, 11, 13, 17, 19} are predicted as at epoch 266, 19 at epoch 267, 73 epoch 268, and The same pattern appears for classes C2 =2, 14, 18), Ca 14)
7

---

Published as a conference paper at ICLR 2024
Table 8: Prediction for base 10 uniform operands and outcomes. Most common prediction for GCD 20, and frequency, for successive epochs. Correct predictions are bold
Epech 266 Pred
Epoch 267
Epoch 268
Epoch 269
Epoch 270
Epoch 580
Epoch 581 Pred
%
Pred
%
Pred
Pred
Pred
Pred
1
100
54
100
13
77
2
2
100
1000
26
1000
62
100
66
100
22
22
99
3
1
100
19
52
73
100
7
100
13
100
99
77
99
4
44
91
100
100
44
100
100
100
100
5
1001
100
100
5
6
2
100
66
100
26
200
62
100
66
100
22
93
22
99
7
1
100
19
62
73
100
7
100
13
100
1
99
77
99
8
99
100
100
100
8
100
88
100
88
99
53
100
13
77
10
70
70
10
100
30
99
70
100
70
100
30
100
70
100
11
I
100
19
57
73
100
7
100
13
100
98
77
99
12
44
91
100
100
44
100
4
100
100
18
13
55
100
13
100
77
14
2
100
66
100
26
100
62
100
66
100
22
92
22
99
15
5
100
55
100
55
100
55
100
5
100
5
100
5
100
16
48
97
84
99
48
16
98
48
48
17
100
19
54
100
7
100
13
100
77
100
18
2
100
66
100
26
100
62
100
66
100
22
93
22
99
19
1
100
19
53
73
100
7
100
13
100
99
77
99
20
20
100
100
20
100
20
20
100
20
and C3- (5, 15). i.e. pairs of integers both divisible by 2, 4, and that would have been predicted as 2, 4, and 5 by the base 10 model from section [3] In other words, the model learns to cluster input pairs into classes having common divisor (a product of divisors of 10), just like did in section 3 but instead of predicting the smallest (and most common) element in each class, it predict a different element at every epoch. This can be summarized into three rules with uniform outcomes:
(U1) Predictions are mostly deterministic. At a given epoch, the model usually predicts a unique value fk) for given GCD k. In rare cases, the model makes 2 or 3 predictions.
(U2) Classes of multiples of products of prime divisors of B are predicted the same. For base 10,
some classes are C1 = f1.3, 7, 9, 11, 13, 17, 19. ). C2 (2, 6, 14, 18, 22, 26, 34, 38 38..... and C3- = {5. 15, 35, 55 .....
C4 (4, 12, 24, 36, 44, 52.
(U3) For each class, the model prediction
an element of the class. Prediction varies from one epoch to the next, but the number of correct GCD is stable over time: it is the number of classes, which increases as the model learns new divisors of B.
The three rules explain the variations in the accuracy curve: since 61% of examples in the natural test set have GCD 1, accuracy jumps by 61% every time class C1 is predicted as 1. Rule U3, on the other hand, accounts for the step-shaped learning curve for correct GCD.
These results shed light on the learning process and the role of the distribution of outcomes. During training. all models, regardless of outcome distribution, learn to partition their input pairs into classes, with GCD multiples of a product of divisors of the base (or small primes when grokking happens). i.e. for base 10, multiples of 2, 4, 5, 10, 20, and default class associated to 1. The model makes unique prediction for all pairs in class. When the distribution of outcomes is unbalanced, this prediction the smallest element in the class, which happens to be the most common. When outcomes are uniformly distributed, different element of the class is predicted at every epoch, somewhat randomly: the model becomes less explainable.
Base 1000, grokking and loss of determinism. Models with base 1000, trained on uniform operands and outcomes, undergo similar learning process (see Appendix D.3) during the first 400 training epochs. Grokking sets in around epoch 200. Multiples of 11, 22, 44, 55 and 88 are learned around epoch 220, then multiples of by epoch 260 and of by epoch 400. At this point, 41 GCD are correctly predicted. Note that grokking no longer happens in order: 11 is learned before 3.
During the grokking phase, new phenomenon develops. As new primes are grokked and more classes are created, model predictions for each class become less deterministic. Instead of predicting a unique value for each class at each epoch, the model now "hesitates" between several values, and the frequency of the most common prediction goes down. By epoch 400, for the class C1, the model makes 18 different predictions with frequencies ranging from 2% to 13% (Table [5]in Appendix D.3) Model predictions no longer explainable, and the three rules are not respected.
8

---

Published as a conference paper at ICLR 2024
Interestingly, GCD continue to be learned under this new regime, starting with the largest (i.e. the smallest classes of multiples). By epoch 740, 95 GCD under 100 are correctly predicted. The worst performance is achieved for small GCD: 43, 74 and 85% correct predictions for GCD 1, 2 and 3. Appendix D.4 presents results for larger bases, where up to 99 GCD under 100 are learned.
7 DISCUSSION
Can transformers learn the greatest common divisor? With enough examples and appropriate adjustment of their training distribution, they can. Models leveraging large composite bases, and trained on log- log-uniform operands and outcomes predict over 90 of the 100 first GCD. Models trained on uniform outcomes predict 95 GCD. However, the experiments from section show the limits of naive, benchmark-based evaluations on arithmetic tasks: high accuracies (95%) can be achieved, on held-out test sets of of random examples, by models that only predict a handful of GCD.
The approach explainability presented in this paper differs from most works on the subject. Instead of looking at model parameters, I engineer experiments that reveal the algorithms that the is often repeated that transformers are incomprehensible black-boxes, that
model is implementing. sometimes confabulate and often fail in unpredictable ways. Here, model predictions can be fully characterized by a small number of rules. This is a promising direction for future research.
Experiments indicate that transformers learn sieve algorithm for computing GCD. The model first learns divisibility by products of divisors of the base, which can be tested by looking at the last digits of number, or counting its rightmost zeroes. Using these rules, the model clusters its input pairs into classes of multiples of divisors of the base, and predicts the GCD as the minimum for the class, All GCD corresponding products of divisors of B are learned this way. At the end of this phase, in base 2, the model correctly predicts 1, 2, 4, 8, 16 and 32.
As training proceeds, new prime divisors are learned (grokked) in order. They are all prime because multiples of previous divisors were learned already, i.c. the model functions like a sieve. Every time a new divisor p is learned, all existing classes are split between multiples and non-multiples of p. In base 2, once the model learns divisibility by 3, six new classes are created: multiples of 3, 6, 12, 24, 48 and 96 (splitted from 1, 2, 4, 8, 16 and 32. This accounts for the steps observed in the learning curves. A GCD is correctly predicted once all the powers of primes dividing it are learned. Eventually, all GCD will be learned this way.
Experiments with uniform outcomes suggest that unbalanced training distribution of GCD is needed for this algorithm to succeed, because it causes each class to predicted by its smallest, and most common, member (the correct GCD), and it guarantees that primes are learned in order. Interestingly, this algorithm is not related to Euclid's s algorithm. Note also that it is not specific to transformers: Appendix D.5 shows that similar results can be achieved with LSTM and GRU.
Another important finding is the role of training distributions. All models are tested on sets with uniform operands, but the best results achieved with log-uniform distribution of operands and outcomes in the training set. This may come as a surprise, since many authors observed that evaluating a model out of its training distribution has negative impact on performance. The existence of special training distributions, that allow for faster learning and more robust models (with respect to -ii-t----dttttiiionn generalization) was already observed for linear algebra (Charton) 2022a)
A log-uniform distribution of operands strikes balance between memorization and generalization, and helps models learn hard instances by memorizing easier cases. This is related to curriculum learning, but avoids catastrophic forgetting, because the training distribution never changes. These observations may apply other arithmetic tasks. On the other hand, log uniform distribution of outcomes helps learning by enforcing better representation of large GCD in the training set, classical recipe in machine learning (calssifiers are often trained on balanced datasets). The counter- intuitive result is that perfectly balanced, uniform training distribution set degrades performance by preventing the model from learning small GCD, and breaking model explainability.
Is it really grokking? Power et al. (022)) define grokking as "generalization far after overfitting." In all experiments, training and test data are generated on the fly from a very large problem space. No overfitting can happen, and the classical pattern of grokking, train accuracy dropping, and validation accuracy catching up after a long time, will not occur. The similarity with grokking lies in the sudden change in accuracy after long stagnation of the training loss.

---

Published as a conference paper at ICLR 2024
REFERENCES
Ernesto Cesàro. Question 75 (solution). Mathesis, (3):224-225) 1883.
François Charton, Linear algebra with transformers. arXiv preprint arXiv:2112.018% 2021.
François Charton, Amaury Hayat, and Guillaume Lample. Learning advanced mathematical compu- tations from examples. arXiv preprint arXiv:2006. (6462, 2020.
François Charton. What is my math transformer doing? three results on interpretability and generalization. arXiv preprint urXiv:221..00/70 2022a.
François Charton. Computing the roots of polynomials, /ttttt:/t/-- github ynomilynmmil... rootss
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv: /406. 1078, 2014.
Ernest Davis. Mathematics, word problems, common sense, and artificial intelligence, 2023.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Eitinger, 7aid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality, 2023.
Kaden Griffith and Jugal Kalita. Solving arithmetic word problems with transformers and preprocess- ing of problem text. arXiv preprint arXiv:2106.008: 2021.
Andrey Gromov. Grokking modular arithmetic, 2023.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv: 151 0,0228. 2015.
Nal Kalchbrenner, Ivo Danihelka, and Alex Graves. Grid long short-term memory. arXiv preprint arxiv:150..01526 2015.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Guillaume Lample and François Charton. Deep learning for symbolic mathematics. arXiv preprint arXhv:1912.01412. 2019.
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381, 2023..
Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. 2022.
Yuanliang Meng and Anna Rumshisky, Solving math word problems with double-decoder transformer.
arXiv preprint arXiv: /908. 10924, 2019.
Bhumika Mistry. An investigation into neural arithmetic logic modules. PhD thesis, University of Southampton, July 2023. URL FRRLhttpss: : //epa::///p/iiii int.rintts.ommm www.s..ssooo... wtt....kk/44/4555
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023.
Rodrigo Nogucira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102. 130/9, 2021.
10

---

Published as a conference paper at ICLR 2024
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:21 12.00/ 14, 2021.
Theodoros Palamas. Investigating the ability of neural networks to learn simple modular arithmetic. 2017.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization Beyond Overfitting Small Algorithmic Datasets. arXiv preprint arXiv:2201.02117, 2022.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models, 2019.
Feng Shi, Chonghan Lee, Mohammad Khairul Bashar, Nikhil Shukla, Song-Chur Zhu, and Vi- jaykrishnan Narayanan. Transformer-basee Machine Learning for Fast SAT Solvers and Logic Synthesis. arXiv preprint arXiv:2107.07116, 2021.
Kai- Kai-Yeung Siu and Vwani Roychowdhury. Optimal depth neural networks for multiplication and related problems. In S. Hanson, J. Cowan, and C. Giles (eds.), Advances in Neural Information Processing Systems, volume 5. Morgan-Kautn Kaufmann, 1992.
Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. arXiv preprint arXiv: 1808. 00508, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Chi, Quoc Le, and Denny Zhou. Chain-of-thougee prompting elicits reasoning in large language models, 2023.
Sean Welleck, Peter West, Jize Cao, and Yejin Choi. Symbolic Brittleness in Sequence Models: Systematic Generalization in Symbolic Mathematics, 2021.
Wojciech Zaremba, Tomas Mikolov, Armand Joulin, and Rob Fergus. Learning simple algorithms from examples, 2015.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023.
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:22/ 1222//.0000 1,09066, 2022.

---

Published as a conference paper at ICLR 2024
APPENDIX
A RATIONAL ARITHMETIC WITH TRANSFORMERS
In these experiments, transformers are trained to perform five arithmetic operations on positive rational numbers:
comparison: given four positive integers a, b, c and d. predict whether 8 < 5-
Integer division: given two integers and b, predict the integer [ Addition: given four integers o, b, c and d, predict the sum a
c 5 + 2- in lowest terms.
Multiplication:
given four integers a, b, e and d. predict the product & x g. in lowest terms. Simplification: given two integers a and b, predict the lowest term representation of 5. i.e. á with ged(a,b] and d ged(a,b)
For the comparison, addition and multiplication tasks, all integers a, b, and d are uniformly sampled between and M (M=1000,000 or 1,000,000).
For the simplification task, integers m, n, are uniformly sampled between 1 and M, I let a
pmn and 6 mo and the model is tasked to predict and b.
pn
ged(m..)))
For the integer division task, 3 integers I let a = pm + m and n, and the model is tasked to predict p L+J-
P are uniformly sampled between 1 and M, with m < n,
All integers are encoded as sequences of digits in base B (see section 2). Sequence to sequence transformers with layers, 512 dimensions and attention heads are trained minimize cross- entropy loss, using Adam with learning rate 10 inverse square root scheduling, linear warmup over 10, 000 optimization steps, and batch size of 256. After each epoch (300,000 examples), models tested on 100,000 random examples.
Comparison Iearned to very high accuracy, and integer division to some extent. On the other hand, the three tasks involving GCD calculations (simplification, addition and multiplication) not learned (Table (
Table 9: Rational arithmetic with transformers. Accuracy of trained models Best of 3 models, trained for 1000 to 1500 epochs.
Comparison
Integer division
Simplification M=10"
Addition
Multiplication
Base
M=10²
M=10
M=10
M=10
M=10
M=10
M=10
M=10²
M=10
10
100
100
21.2
2.4
0.14
0.02
0
0
0
30
99.9
100
14.2
2.2
0.21
0.02
0
0
31
99.9
100
14.3
2.4
0.02
0
1000
100
99.9
8.8
0.09
0.01
B MODEL SCALING FOR THE BASE EXPERIMENTS
Section 3 presents results for 4-layer transformers with 512 dimensions and 8 attention heads. In this section, experiment with very small models (down to layer and 32 dimensions), and very large ones (up to 24 layers and 1024 dimensions). Note: in Tables 10 and [1 the number of trainable parameters are indicated for base 10, they will be larger for larger bases, because larger vocabularies increase the number of parameters in the embedding and decoding layers.
Table TO
presents accuracies for models with one layer, attention heads, and 32 to 512 dimensions. These models have 3 to 100 times less parameters that the 4-layer baseline, but there is no significant change trained model accuracy for 12 different bases.
Table I
presents results for models from 6 to 24 layers, symmetric (same number of layers in the encoder and decoder), or asymmetric (using a one-layer encoder or decoder). The dimensions are 512, 640, 768 and 1024 for 6, 8, 12, and 24 layers, and the totssiionnt-ttentionn heads ratio is kept constant at 64 (i.e.there are 8, 10, 12 and 24 attention heads respectively). Again, model size has no significant impact accuracy.
12

---

Published as a conference paper at ICLR 2024
Overall, these scaling experiments suggest that trained model performance is stable over wide range of model size (300 thousands to 700 millions parameters). These results are strikingly different from what is commonly observed in Natural Language Processing, where very small transformers (under few million parameters) cannot learn, and accuracy improves with model size.
Table 10: Model accuracies for different dimensions and numbers of parameters. All models have one layer and attention heads. Parameter counts for base 10.
512 dimensions
256 dim.
128 dim.
64 dim.
32 dim.
4-layer baseline
Base
11.6M 4.0M
1.7M
0.6M
0.3M
33.7M
81.3
81.4
81.4
81.4
81.2
81.6
68.8
68.9
68.8
68.7
68.9
81.4
81.4
81.4
81.4
81.4
81.4
5
64.0
63.7
63.8
63.7
63.8
64.0
6
91.3
91.3
91.1
91.1
90.7
91.5
62.5
62.4
62.5
62.5
62.5
62.5
84.4
84.3
84.4
84.7
11
61.7
61.7
61.7
61.9
61.7
61.8
12
91.4
91.4
91.3
91.3
91.1
91.5
15
71.6
71.6
71.5
71.5
71.4
71.7
94.6
93.8
93.7
94.7
31
61.3
61.3
61.2
61.3
61.3
61.3
Table 11: Model accuracies for different depths and number of parameters (in millions). 1 and 6 layer models have 512 dimensions and 8 heads, 8-layer have 640 dimensions and 10 heads, 12-layer 768 dimensions 12 heads, 24-layer models have 1024 dimensions and heads. The largest base 2 and 3 models could not run on one 32GB GPU. All model parameters for base 10.
1/6
6/1
6/6
1/8
8/1
8/8
1/12
12/1
12/12
1/24
24/1
24/24
Base
32.5
27.3
48.3
59.1
48.4
97.1
117.1
94.7
204.8
387.4
313.3
713.8
2
81.3
81.3
81.4
81.5
81.4
81.3
81.3
81.3
81.4
81.4
3
68.7
68.8
68.7
68.8
68.9
69.0
68.9
68.8
68.8
68.8
68.6
4
81.3
81.4
81.4
81.4
81.4
81.6
81.4
81.4
81.4
81.5
81.4
81.3
5
63.8
63.8
63.7
63.8
63.6
63.7
63.7
63.7
63.6
63.9
63.7
63.6
6
91.3
91.1
91.3
91.3
91.4
91.3
91.3
91.0
91.0
91.3
91.0
90.9
7
62.6
62.6
62.4
62.5
62.4
62.6
62.5
62.4
62.4
62.4
62.3
62.2
84.3
84.2
84.4
84.7
84.4
84.5
84.4
83.4
83.4
11
61.8
61.7
61.6
61.7
61.8
61.7
62.0
61.6
61.7
61.7
61.6
61.6
12
91.4
91.3
91.3
91.4
91.5
91.4
81.4
91.2
91.2
91.4
91.3
91.2
15
71.5
71.5
71.4
71.5
71.5
71.5
71.4
71.5
71.5
71.5
70.6
71.4
30
94.6
93.4
93.5
94.7
93.6
93.6
94.7
93.6
93.6
93.5
93.4
93.4
31
61.2
61.2
61.3
61.2
61.3
61.2
61.4
61.2
61.3
61.3
61.1
C
THEORETICAL VALUES OF ACCURACY
In this section, compute a theoretical accuracy for models from section that follow the three rules, assuming that all products of prime divisors of B are correctly predicted. The distribution of the GCD random uniform positive integers verifies: P(ged(a, b) = k) न (Cesàro 1883
Therefore, = p*, with prime, theoretical model accuracy is
pd
A(p²) A(p)
72
2
x²2
i=0
if = = T1 A(p) A(g)),
if =/ A(B) = =1 11 A(p)( A(g))( A(r)), and
13

---

Published as a conference paper at ICLR 2024
Table 12: Theoretical accuracy, accuracy and number of correct GCD under 100. Best of experiments.
Base
2
3
4
5
6
7
10
11
12
15
Theoretical accuracy
81.1
68.4
81.1
63.3
90.2
62.1
88.6
61.3
90.2
80.3
Accuracy
81,6
68.9
81,4
64.0
91.5
62.5
84.7
61.8
91.5
71.7
Correct GCD
7
5
7
3
19
3
13
19
Base
30
31
60
100
210
211
420
997
1000
1024
Theoretical accuracy
94.1
60.9
94.1
88.6
96.3
60.8
96.3
60.8
88.6
81.1
Accuracy
94.7
61.3
95.0
84.7
95.5
61.3
96.8
61.3
84.7
81.5
Correct GCD
27
2
28
13
32
38
14
Table 12 compares theoretical accuracies with empirical observations. Best model performances may be higher than theory, because of sampling errors in the test set, or lower than theory when some powers of prime divisors of B have not been learned.
D ADDITIONAL EXPERIMENTS
D.1 EXPERIMENTS WITH OUTCOME DISTRIBUTIONS
The results at the end of section demonstrate that training from a log-uniform distribution GCD (P(ged k) )) improves model performance compared the natural, inverse square distribution (P(ged k) k In this section, experiment with three alternative distributions of outcomes:
"long- tail" log-uniform distribution: instead of sampling GCD between 1 and 100, they are sampled between and 200,
inverse square root distribution of outcomes: P(ged k)
V
inverse power 1.5 distribution: P(ged = k) kVE
C
Table 13: Correct GCD for different outcome distribution scaling laws. Best of 3 models, trained for 1000-1300 epochs. Log-uniform operands.
Outcome distribution scaling law
Base
k I 100
k S 200
RvE
v/
1000
71
71
91
90
91
1024
71
72
90
85
91
2017
63
64
88
87
88
2021
66
71
87
2023
65
67
85
90
2025
70
71
91
88
92
2187
66
70
91
86
91
2197
68
65
85
91
2209
65
68
85
2401
73
69
89
85
92
2744
72
72
91
88
89
3125
65
67
87
92
3375
67
87
4000
66
60
90
85
90
4913
57
60
88
90
92
5000
64
65
89
90
91
10000
56
90
Table 13 presents results for 17 bases between 1000 and 10000, for models trained with log-uniform operands and five distributions of outcomes. As observed in section 5 log-uniform distribution of outcomes achieves better performances than the natural (inverse square) distribution. The inverse power 1.5 distribution of outcomes only brings marginal improvement over the natural distribution. With log-uniform outcomes, sampling GCD up to 200 instead of 100 has a negative impact on the number of correct GCD, except for the largest bases. On the other hand, training from inverse square root distribution of outcomes improves performance for all bases. For bases, 92 GCD under 100 are predicted correctly.
14

---

Published as a conference paper at ICLR 2024
D.2 LEARNING WITH SMALLER BATCHES
A common advice, when training transformers on natural language processing tasks, is to use the largest possible batches (i.e. as many as will fit in GPU memory). Large batches have two advantages, they avoid extreme gradients by averaging them over many samples, and they accelerate training by reducing the number of optimization steps. All models in this paper were trained with batches of 256 examples. In this section, I experiment with batches of 64, training models with log-uniform operands (and various outcome distributions) for about 800 epochs.
Table 14 compares models with batches of 64 to batches of 256, trained for week (about 800 epochs for batch 64, 1300 for batch 256), on 11 different bases. For the same training time, batch size seems to have little impact performance. This suggests that models could be trained on machines with less GPU memory, at no penalty.
Table 14: Correct GCD for different batch sizes, Best 3 models, log-uniform operands. Models with batch size 64 are trained for 800 epochs, models with batch size 256 for 1300 epochs.
Inverse square outcomes
Log-uniform outcomes batch size batch size 256
Base
batch size 64 batch size 256
10
49
48
69
65
12
54
70
30
56
56
73
72
31
45
44
64
64
210
55
60
81
80
1000
70
91
2025
66
70
90
91
2401
68
73
90
89
2744
70
72
91
91
4000
67
66
90
90
10000
55
D.3 UNIFORM OPERANDS AND OUTCOMES BASE 1000
In this section, I provide detailed results for models using base 1000, and trained uniform operands and outcomes. Learning curves (Figure are similar to those for base 10 (Figure 3) during the first 200 epochs: loss curves are flat, accuracy varies wildly, and the number of correct GCD has the characteristic step-like shape observed throughout this paper. Grokking, characterized by steep drops the loss and increases in the number of correct GCD, happens between epochs 200 and 400. Then, the accuracy and the number of correct GCD, rise steadily. After 800 epochs 95 (out of 100) GCD are correctly predicted.
Accuracy
Correct GCD
Loss
70-
3.5
60-
80-
60
2.5
40
30
40
1.5
20
20
0.5
200
400
600
800
200
400
600
800
0
200
400
600
800
Epochs
Epochs
Epochs
Figure 4: Learning curves for B=1000 uniform operands and outcomes.
More precisely, by epoch 180, the model has learned to classify all examples into 14 sets: multiples of 1, 2, 4. 5, 8, 10. 16, 20, 25, 32, 40, 50. 80 and 100. At each epoch, the model selects one element in each class, which is its unique prediction for all pairs of integers with GCD the class: the rules U1 to U3 are respected.
15

---

Published as a conference paper at ICLR 2024
Grokking sets in around epoch 200, and by epoch 220, 5 new classes have been learned: multiples of 11 (11, 33, 77 and 99). 22, 44, 55 and 88, created by "splitting away" the multiples of from the classes of multiples of 1, 2, 4, 5 and 8. Because of uniform outcomes, grokking does not happen in increasing order: 11 is learned before 3. By epoch 260, multiples of 3 are learned and the model predicts 31 different outcomes (splitting 12 classes, from to 32). By epoch 400, multiples of 7 are learned, and 41 GCD are predicted.
During the grokking phase, new phenomenon develops. As new primes are grokked and more classes are created, model predictions for each class become less deterministic. Instead of predicting a unique value for each class each epoch, the model now "hesitates" between several values, and the frequency of the most common prediction goes down. By epoch 400, for the class of multiples of the model makes 18 different predictions with frequencies ranging from 2% to 13% (Table 15
Table 15: Base 1000 epoch 400 predicted values and frequencies.
GCD
1
GCD 2
GCD 3
GCD 4
GCD 5 Pred
Pred.
%
Pred
%
Pred
%
Pred
%
%
11
12
4
40
5
12
17
22
13
27
11
44
19
55
21
19
34
33
68
85
31
23
29
46
10
57
12
92
12
31
58
10
69
22
37
62
10
81
7
41
43
93
59
86
99
2
61
67
71
73
79
83
13
89
97
At this point, model predictions are neither deterministic nor interpretable, and the three rules are no longer respected. Classes have many predictions there elements, and the model begins learning individual GCD, beginning with the largest ones (i.c. the smallest classes). By epoch 740, 95 of the 100 first GCD correctly predicted, the worst performance being achieved the smallest values (GCD 1, 2 and 3, correctly predicted 43, 74 and 85% of the time).
D.4 UNIFORM OUTCOMES LARGER BASES
In these experiments, models trained for 1200 epochs, from uniform and log- log-uniform operands, and uniform outcomes. As previously, large bases achieve the best results. Models trained on uniform operands also seem perform better.
Table 16: Correct GCD with uniform outcomes. Best of 3 models, trained for 1200 epochs.
Base
Uniform operands
Log-uniform
Base
Uniform operands
Log-uniform
1000
71
94
2401
98
90
2017
98
87
2744
96
93
2023
3375
2187
99
94
4000
98
94
2209
57
90
4913
99
93
2310
96
92
10000
99
95
16

---

Published as a conference paper at ICLR 2024
D.5 EXPERIMENTS WITH DIFFERENT ARCHITECTURES
In this section, I experiment with two popular recurrent architectures: long short-term mem- ories (LSTM) [Hochreiter & Schmidhuber, 1997). and gated recurrent units (GRU) (Cho et al.] 2014). I train models with 1024 and 2048 dimensions, and four layers, uniform operands, log-un uniform operands and log-uniform operands and outcomes, for 10 different bases: 10, 30, 31, 210, 420, 1000, 2021, 2023, 2025 and 2401.
By and large, the reuslts of my experiments with transformers extend to other recurrent networks. After 500 epochs, models trained uniform operands (table 17 achieve performances similar to those obtained in sections and Composite bases like 210 and 420 achieve the best results (35 and 38 GCD), and large bases allow for grokking small primes. There is clear advantage of LSTM over GRU, or 2048 over 124 dimensions, Models trained on log-uniform operands (table [8) and outcomes (table 19) perform better, but results are lower (after comparable training time) than with trasnformers.
Table 17: Correct GCD with uniform operands. Best of models, trained for 500 epochs.
LSTM
GRU
Base
1024 dim.
2048 dim.
1024 dim.
2048 dim.
10
15
30
32
30
32
30
31
2
2
2
2
210
35
35
35
35
420
34
38
1000
22
2021
7
8
10
2023
11
11
11
2025
24
24
10
18
2401
Table 18: Correct GCD with log-uniform operands. Best of 3 models, trained for 500 epochs.
LSTM
GRU
Base
1024 dim.
2048 dim.
1024 dim.
2048 dim.
10
33
36
30
38
40
40
40
31
29
32
22
20
210
47
46
46
420
45
43
1000
53
51
46
47
2021
44
40
36
35
2023
46
48
38
39
2025
52
2401
47
41
35
33
Table 19: Correct GCD with log-uniform operands and outcomes. Best of models, trained for 450 epochs.
LSTM
GRU
Base
1024 dim.
2048 dim.
1024 dim.
2048 dim.
10
53
54
38
53
30
40
58
36
40
31
44
37
210
53
420
64
74
53
64
1000
69
73
62
69
2021
60
76
54
60
2023
53
2025
69
60
2401
55
73
50
55
17

---

Published as a conference paper at ICLR 2024
E ADDITIONAL RESULTS
E.1 GROKKING
GCD predicted
Train loss
16
14
10-
200 300
400 Epochs
500
000
700
800
100
3000
800
700
600
Epochs
Figure Learning curves base B=2023, 3 different model initializations.
Table 20: Model predictions. 1000, after 220 epochs. being learned.
GCD
Prediction
GCD
Prediction
GCD
Prediction
GCD
Prediction
GCD
Prediction
11
1
21
41
2
12
12
22
2
32
16/ 32
42
6
13
1
23
1
33
43
2
24
24
5
15
15
25
25
35
45
15
6
16
16
26
2
36
12
46
2
7
17
1
27
3
37
47
1
18
6
28
38
48
48
19
1
29
10
20
20
30
30
40
40
50
50
Table 21: Predicted GCD, natural test distribution, and 5% uniform GCD, Best model of 3.
Natural distribution
5% uniform GCD
Base
Correct GCD
Epochs
Correct GCD
Epochs
625
6
650
10
1000
22
250
1.5
560
2017
450
2021
10
550
10
600
2023
16
600
11
800
2025
28
850
18
225
2187
750
2197
11
800
775
2209
8
850
6
575
2401
14
700
14
630
2744
29
1400
650
3125
16
3375
23
400
23
475
4000
25
650
25
600
4913
17
950
575
5000
28
10000
22
250
300
18

---

Published as a conference paper at ICLR 2024
E.2
LEARNING CURVES LOG-UNIFORM OPERANDS
GCD predicted
Train loss
0.07
60
0.06
50
0.05
40
0.04
30
0.03
20
0.02
10
0.01
0
0.00
200
400
600
800
1000
0
200
400
600
800
1000
Epochs
Epochs
Figure 6: Learning curves for base B=2023. Log-uniform operands, natural outcomes, 3 different model initializations.
GCD predicted
Train loss
0.200
80
0.175
0.150
60
0.125
0.100
0.075
0,050
0.025
0
200
400
600
0.000
200
400
600
Epochs
Epochs
Figure 7: Learning curves base B=2023. Log-uniform operands, log-uniform outcomes. different
model initializations.
19

---

Published as a conference paper at ICLR 2024
E.3
DETAILED MODEL PREDICTIONS BASE EXPERIMENTS
Table 22: Predicted values for ged to 63.
Base GCD
2
31
420
Prediction
Pred.
Pred.
Pred.
Pred.
Pred.
%
100
2
100
100
100
100
100
2
100
3
1
100
100
100
100
100
3
100
4
4
100
100
100
100
100
4
100
100
100
100
100
100
99.6
7
100
100
100
100
100
7
100
8
8
100
100
100
100
100
8
100
100
100
10
100
10
100
100
10
100
11
100
100
100
100
100
1
100
100
100
100
12
100
12
99.8
100
100
100
100
100
14
2
100
100
100
100
100
14
100
15
100
100
100
15
100
100
15
99.4
16
100
100
16
100
100
100
100
18
2
100
100
100
18
100
100
18
100
19
100
100
100
100
100
1
100
20
20
100
21
100
100
100
100
21
100
22
2
100
100
100
100
100
2
100
23
100
100
100
100
100
1
100
24
25
100
100
100
100
25
99.9
26
2
100
100
100
100
100
2
100
27
100
100
100
100
100
9
100
28
28
29
100
100
100
100
100
100
30
2
100
100
10
100
30
100
100
30
99.6
31
100
100
100
100
1
100
32
98.7
99.9
100
16
100
33
100
100
100
100
100
3
100
34
2
100
100
100
100
100
2
100
100
100
100
35
100
100
100
100
100
36
37
100
100
100
100
100
100
38
2
100
100
100
100
100
2
100
99,9
40
99.9
100
99.9
100
40
99.9
41
100
100
100
100
100
1
100
42
2
100
100
100
99.9
100
42
100
44
4
100
100
100
100
100
4
100
45
100
100
100
45
100
100
45
99.8
46
100
100
100
100
100
2
100
47
100
100
100
100
48
16
100
16
100
16
99.9
24
100
100
48
99.9
49
100
100
100
100
100
7
100
100
100
100
100
50
99.6
51
100
100
100
100
3
99.8
52
4
100
100
100
100
100
4
100
53
100
100
100
100
100
1
100
100
100
18
100
55
100
100
100
100
100
5
100
56
8
100
100
99.9
100
100
56
100
57
100
100
100
100
100
3
99.9
100
59
100
100
100
100
100
60
4
100
100
20
100
60
100
100
60
99.7
61
100
100
100
100
100
1
100
63
100
100
100
100
100
63
100
20

---

Published as a conference paper at ICLR 2024
Table 23: Predicted values for ged 64 to 100.
Base GCD)
2
10
30
31
420
Prediction
Pred.
Pred.
Pred.
Pred.
55
Pred.
%
64
64
98.9
64
99.2
16
99.8
100
100
16
100
100
5
100
66
2
100
100
100
100
6
100
67
I
100
100
100
100
100
1
100
68
4
100
100
100
100
100
4
100
69
100
100
70
2
100
100
100
100
70
100
71
100
100
100
100
100
1
100
72
8
100
100
100
72
100
100
72
100
74
100
100
100
100
100
2
100
75
100
100
25
100
75
100
100
75
99.4
76
100
100
100
100
100
100
77
7
78
2
100
100
100
100
100
6
100
79
100
100
100
100
100
1
100
80
100
16
100
99.9
100
80
100
81
100
100
100
100
99.8
82
2
100
100
100
100
100
2
100
83
100
100
100
100
100
1
100
84
100
100
12
100
84
100
85
100
100
100
100
100
5
100
86
2
100
100
100
100
100
2
100
87
100
100
100
100
100
3
99.8
100
100
89
100
100
100
100
100
100
90
2
100
100
10
100
90
100
100
90
99.9
91
100
100
100
100
100
7
100
93
100
100
100
100
31
99.9
3
99.8
94
2
100
100
100
100
1
100
2
100
95
100
100
100
100
100
5
100
32
48
97
100
100
100
100
100
1
100
98
2
100
100
100
100
100
14
100
99
100
100
100
100
9
99.8
100
100
100
100
100
100
100
100
99.6
21


Reviews:
[{"review_id": "v3Ef84iDDW", "reviewer": "Reviewer_21KR", "summary": "This paper trains small transformer models (4 layers) to compute the GCD of a pair of numbers. This is framed as a sequence prediction task. The input is a pair of numbers (a, b) encoded as a sequence of numbers in a base B. The model predicts a sequence of numbers (in base B) which collectively represent the GCD of (a, b). Performance is measured by 2 metrics - accuracy in prediction GCD and number of integers between 1 to 100 which are a GCD prediction. The second metric is required according to the paper because for any given pair (a, b), the GCD is likely to be small and they also turn out to be harder to predict correctly. The authors perform a bunch of experiments in a variety of configurations which will be described below.\n\nFor the first set of experiments, the training dataset (pairs of numbers) is sampled uniformly between 1 to a million. Models trained on these datasets show pretty variable accuracy but the main trend seems to be that using larger bases of composite numbers to encode pairs work well. Models also exhibit the property of being able to predict GCDs equal to the (prime) divisor of their bases (and their small powers) pretty well. Other small prime numbers which are not divisors of the base are ‘grokked’ much later during training.\n\nNext, the authors decide to oversample pairs of small numbers to create a lop-sided training distribution (called log-uniform in the paper). Models trained on this kind of distribution show dramatically improved and consistent performance regardless choice of base.\nFinally, the authors modify the training distribution to also be uniform over the predicted GCD. This maintains good accuracy of these models but degrades the consistency with which model predictions could be explained.\n\nTo analyze the model predictions, the authors introduce ‘rules’ which they obtain after looking at the results after each stage of the above experiments. These rule help us explain model predictions in a uniform fashion.", "questions": "How is the choice of the base made? Are they also sampled from a distribution?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. The main idea is pretty straightforward and most of the concepts in the paper are well explained (see concerns discussed below).\n2. To my knowledge, computing and analyzing GCD prediction via transformers has not been done before. \n3. Very comprehensive suite of experiments and exhaustive analysis. I appreciate the authors performing such a wide range of experiments. I also really like the nice link between theoretical accuracy and practical accuracy provided in Appendix C.\n4. Claims made by the authors are clearly backed up the experiments in the paper (see some minor concerns below).", "weaknesses": "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 transformer themselves beyond the narrow context of learning GCDs. Like I mentioned earlier, the experiments are thorough and the analysis is extensive but I’m struggling to understand the value of this beyond this specific task. \n\nThe 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?\n\nHaving said all of that, I appreciate a good, rigorous experimental paper and I’m open to being convinced by the authors and reading other reviewers comments about the value of this work.\n\nOther:\n1. The authors hypothesize on Page 3 “For prime bases, such as B = 2, the three rules suggest that the model learns to predicts GCD by counting the rightmost zeros in its inputs”. Something like this seems like it should be fairly easy to confirm by visualizing the attention weights. I wonder if the authors have tried this.\n2. The authors mention that their model is ‘fully explainable’ (see abstract). I was expecting the model to output some kind of explanation for its prediction (say via attention) but what the authors actually meant was that the model predictions can be $\\textit{explained by}$ a human  because they follow certain rules. These rules keep on changing depending on the training configuration and can only be inferred after a thorough analysis of the results so I’m not sure if we can call it explainability at least in the traditional sense used in XAI. However, the authors do say in the final section “Our approach to explainability differs from most works on the subject. Instead of looking at model\nparameters, we engineer experiments that reveal the algorithms that the model is implementing”. I think a statement of this sort at the start of the paper would help in clarifying the confusion. Regardless, I’m not sure if I would call a model showing some consistent trends in the output as ‘explainable’.\n\nMinor:\n1. On page 2, when introducing the procedure to generate the stratified test set, the symbol ‘M’ is introduced without any explanation as follows “Sample a and b, uniformly between 1 and M/k , such”. I’m assuming M is the upper limit on the numbers a and b. This symbol is used throughout the paper and explained nowhere so it’s confusing.\n2. The authors use two metrics to evaluate their models - accuracy on a test set and number of correctly predicted GCD below 100. At several points in the paper, terms like “50 GCD” or “50 correct GCD” are used. I now understand that they mean that the model predicts 50 GCDs out of a 100 between 1 to 100 correctly but it was definitely confusing to me initially. I think clarifying these terms at the start can be useful.\n3. Typo - Appendix C, first line, “for models from section 3 the follow the three rules” should be “for models from section 3 that follow the three rules”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PRdiYIuI4b", "reviewer": "Reviewer_QHKu", "summary": "The paper focuses on training Transformers from scratch to predict the GCD of two numbers and explaining the algorithm used by the trained model.  It is observed that for almost all pairs of numbers for which the GCD is k, the model outputs a unique value f(k).  The model learns a set of numbers D, and for an input pair with gcd k, f(k) is the largest number in D that divides k. When numbers are represented in base B, set D contains all the products of primes dividing B. For large bases B, if training continues for very long, set D starts to include other prime numbers as well, usually learned in a monotonically increasing fashion. For example, in base 1000, set D would have numbers 1, 2, 4, 5, 10, 16, 20 and so on. If the training is continued for long, this set also starts to include 3 and its multiples with the numbers already in D. If the GCD k is one of the elements of this set, then the model output will be correct, otherwise, the model outputs the largest element of this set that divides k. The paper also investigates the role of training distribution.", "questions": "Minor suggestions/clarifications:\n\n1. On page 2, the term epochs is used for training, which suggested to me that there is a fixed training set over which the model is trained for multiple epochs. But in the last paragraph of the paper, it is mentioned that new training data is generated on the fly. Perhaps this should be clarified on page 2.\n\n2. In the abstract and introduction of the paper, there are lines of the form ``Models trained from uniform operands only learn a handful of GCD\n(up to 38 out of 100)''. This is confusing if the reader does not know that all numbers with the same GCD are mapped to the same output by the model. Without this information, it can be the case that two pairs of numbers have the same GCD but the model output is correct for only one of them. In that case, it is unclear what would it mean for the model to learn a handful of GCDs. This should be made clear in the abstract and intro.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I really enjoyed reading this paper. In most cases, the algorithms learned by Transformers do not seem interpretable. I was surprised by how structured the encoded algorithm is in the case of GCD computation. This structure is very well explored in the paper through cleverly designed experiments and the intuition behind the results is also explained well.", "weaknesses": "I don't see any substantial weakness. Perhaps one could say that 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UqgkYSVbsJ", "reviewer": "Reviewer_t8xf", "summary": "This paper studies Transformers capabilities on learning the task of computing greatest common divisor (GCD) for two numbers. First the two integers of input and the integer of the output are encoded in base $B$. Then, the model is trained with 300,000 new samples at each epoch. The authors have tried different distributions for the input/operands (uniform and log-uniform) and also the output. They have observed that in the majority of their experiments, model's output is deterministic, meaning that if $\\gcd(a,b)=k$, the output of the model is fixed depending on $k$. Further, for $k$'s that are a divisor of $B$ or small, the output is usually correct. Moreover, they show that by having a log-uniform distribution on the operands, the performance of the model is increased.", "questions": "- Q1. It's said that composite bases allow one to check for divisibility by seeing the rightmost digits. This is true, for example for powers of the base, one needs to count the number of zeros. However, the picture seems to be more complex for divisors of the base. For example, consider $B=210$. Any number ($<10^6$), can be expressed using 3 digits in this base. Now for checking divisibility by 8, one has to check the pattern of all 3 digits (and it's not as simple as being 0). On the other hand, checking divisibility for any other number also requires checking the 3 digits together. Is there anything that is making checking for 8 simpler? (Similarly, assume that $B=350$, could we expect that checking for $3, 6$ would be easier than $8$?)\n- 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 possiblility that smaller primes are learned faster than $43, 47, 43^2, \\ldots$ in this case?\n- Q3. Also given the list intuition, why 420 is performing better than 210 in Table 2?\n- 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? \n\n\n- 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 observe? Or is having small primes such as 2 and 5 playing an important role here? \n- 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. \n\n## Minor questions/remarks\n- Q7. What stopping criterion has been used? For example in Figure 4, the loss appears to be still decreasing. Maybe the combination of continued training and weight decay results in improved generalization as in grokking? \n- 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?\n- Q9. What is the positional embedding used for the model?\n- 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. \n- R2. Table 5 is a bit misleading with the current format. For example, when we see than 9 is learned for $B=2017$ we have to guess that 3 was already learned. (Similar issue can be observed for other bases.) \n- R3. For Section 5 experiments in page 6, it would helpful if training loss plot is provided.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The particular attention to the distribution of the operands is important, and it's shown that by emphasizing on small numbers, the performance of the model can be increased (resembling curriculum learning). \n- Similarly, authors have considered the distribution of the output showing that large GCDs may be slow/hard to learn as they are rare. \n- Experiments are rather extensive in several axes (e.g., number of bases, size of the models, batch-size, ...).", "weaknesses": "- The claim that Transformer predictions are fully explainable does not seem to be accurate once log-uniform operands are used. \n- 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). \n- 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.  \nSee the questions below for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ogEh4m0AIM", "reviewer": "Reviewer_zrzn", "summary": "The paper presents an exploration of using small transformers to calculate the greatest common divisor (GCD) of two positive integers. A notable aspect is that the predictions made by these models are explainable. The authors focus on how models learn a list of divisors and predict the largest element that divides both inputs. They also investigate the impact of different training distributions on model performance, including uniform operands, log-uniform operands, and balanced distributions of GCDs. They also investigate with different inputs representations in different bases.", "questions": "- Did the authors try to experiment with Chain-of-Thought and see if providing the explanation helps with the accuracy?\n- 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?\n- 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?\n- 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?\n- Do few shot prompts improve the performance?\n- 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.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Studying learning GCD through transformers is a novel problem.\n- This work raises some observations about the deterministic behavior of the model in learning GCD, and the simple shortcut algorithm it learns.\n- A lot of experiments/examples supporting the claims are provided.", "weaknesses": "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. \n\nMoreover, 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. This non-robust behavior of the observations weaken them a lot in my opinion. For example, explainability breaks down with the change in input distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)